andreped/ReferenceBot

1from typing import List2 3from langchain.chains.qa_with_sources import load_qa_with_sources_chain4from langchain.chat_models.base import BaseChatModel5from langchain.docstore.document import Document6from pydantic import BaseModel7 8from knowledge_gpt.core.embedding import FolderIndex9from knowledge_gpt.core.prompts import STUFF_PROMPT10 11 12class AnswerWithSources(BaseModel):13    answer: str14    sources: List[Document]15 16 17def query_folder(18    query: str,19    folder_index: FolderIndex,20    llm: BaseChatModel,21    return_all: bool = False,22) -> AnswerWithSources:23    """Queries a folder index for an answer.24 25    Args:26        query (str): The query to search for.27        folder_index (FolderIndex): The folder index to search.28        return_all (bool): Whether to return all the documents from the embedding or29        just the sources for the answer.30        model (str): The model to use for the answer generation.31        **model_kwargs (Any): Keyword arguments for the model.32 33    Returns:34        AnswerWithSources: The answer and the source documents.35    """36 37    chain = load_qa_with_sources_chain(38        llm=llm,39        chain_type="stuff",40        prompt=STUFF_PROMPT,41    )42 43    relevant_docs = folder_index.index.similarity_search(query, k=5)44    result = chain({"input_documents": relevant_docs, "question": query}, return_only_outputs=True)45    sources = relevant_docs46 47    if not return_all:48        sources = get_sources(result["output_text"], folder_index)49 50    answer = result["output_text"].split("SOURCES: ")[0]51 52    return AnswerWithSources(answer=answer, sources=sources)53 54 55def get_sources(answer: str, folder_index: FolderIndex) -> List[Document]:56    """Retrieves the docs that were used to answer the question the generated answer."""57 58    source_keys = [s for s in answer.split("SOURCES: ")[-1].split(", ")]59 60    source_docs = []61    for file in folder_index.files:62        for doc in file.docs:63            if doc.metadata["source"] in source_keys:64                source_docs.append(doc)65    return source_docs66