Spaces:
Sleeping
Sleeping
File size: 2,408 Bytes
e71c4e6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
from typing import Any, List
from canonical_demo_memory.core.debug import FakeChatModel
from canonical_demo_memory.core.embedding import FolderIndex
from canonical_demo_memory.core.prompts import STUFF_PROMPT
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
from langchain.chat_models import ChatOpenAI
from langchain.docstore.document import Document
from pydantic import BaseModel
class AnswerWithSources(BaseModel):
answer: str
sources: List[Document]
def query_folder(
query: str,
folder_index: FolderIndex,
return_all: bool = False,
model: str = "openai",
**model_kwargs: Any,
) -> AnswerWithSources:
"""Queries a folder index for an answer.
Args:
query (str): The query to search for.
folder_index (FolderIndex): The folder index to search.
return_all (bool): Whether to return all the documents from the embedding or
just the sources for the answer.
model (str): The model to use for the answer generation.
**model_kwargs (Any): Keyword arguments for the model.
Returns:
AnswerWithSources: The answer and the source documents.
"""
supported_models = {
"openai": ChatOpenAI,
"debug": FakeChatModel,
}
if model in supported_models:
llm = supported_models[model](**model_kwargs)
else:
raise ValueError(f"Model {model} not supported.")
chain = load_qa_with_sources_chain(
llm=llm,
chain_type="stuff",
prompt=STUFF_PROMPT,
)
relevant_docs = folder_index.index.similarity_search(query, k=5)
result = chain(
{"input_documents": relevant_docs, "question": query}, return_only_outputs=True
)
sources = relevant_docs
if not return_all:
sources = get_sources(result["output_text"], folder_index)
answer = result["output_text"].split("SOURCES: ")[0]
return AnswerWithSources(answer=answer, sources=sources)
def get_sources(answer: str, folder_index: FolderIndex) -> List[Document]:
"""Retrieves the docs that were used to answer the question the generated answer."""
source_keys = [s for s in answer.split("SOURCES: ")[-1].split(", ")]
source_docs = []
for file in folder_index.files:
for doc in file.docs:
if doc.metadata["source"] in source_keys:
source_docs.append(doc)
return source_docs
|