restarting

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+resources/design-by-fire.pdf filter=lfs diff=lfs merge=lfs -text
+resources/lets-talk.pdf filter=lfs diff=lfs merge=lfs -text
+resources/progit.pdf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

README.md

@@ -5,7 +5,7 @@ colorFrom: green
 colorTo: green
 sdk: gradio
 sdk_version: 3.40.1
-app_file: app.py
+app_file: src/app.py
 pinned: false
 ---
 

resources/design-by-fire.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cf8176ae0f4873ca6547f22ecad1fcc366b170488782087a3be48801721eba0
+size 1353204

resources/lets-talk.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1545cbf074f8e363200e32f1da24b0319163bfafd2871078c8e750a32b02098b
+size 4156635

resources/progit.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dea5f1cce14aabd2d3e7246fd4d0e6fe632c13561060bb47eee069b7f257289a
+size 18915172

src/app.py

@@ -0,0 +1,138 @@
+
+import os
+from threading import Lock
+from typing import Any, Dict, Optional, Tuple
+
+import gradio as gr
+from langchain.chains import ConversationalRetrievalChain
+from langchain.chat_models import ChatOpenAI
+from langchain.memory import ConversationBufferMemory
+from langchain.prompts.chat import (ChatPromptTemplate,
+                                    HumanMessagePromptTemplate,
+                                    SystemMessagePromptTemplate)
+
+from .core.chunking import chunk_file
+from .core.embedding import embed_files
+from .core.parsing import read_file
+
+VECTOR_STORE = "faiss"
+MODEL = "openai"
+EMBEDDING = "openai"
+MODEL = "gpt-3.5-turbo-16k"
+K = 5
+USE_VERBOSE = True
+API_KEY = os.environ["OPENAI_API_KEY"]
+system_template = """
+Use the context below to answer questions. You must only use the Context to answer questions. If I ask you about 'the book' or 'this book' or similar references, then answer using the Context. If you cannot find the answer from the Context below, you must respond with
+"I'm sorry, but I can't find the answer to your question in the book, 'Design by Fire,' by Emily Elizabeth Schlickman and Brett Milligan." All answers must be in English unless you are explicitly asked to translate to a different language.
+----------------
+{context}
+{chat_history}
+"""
+
+# Create the chat prompt templates
+messages = [
+  SystemMessagePromptTemplate.from_template(system_template),
+  HumanMessagePromptTemplate.from_template("{question}")
+]
+qa_prompt = ChatPromptTemplate.from_messages(messages)
+
+class AnswerConversationBufferMemory(ConversationBufferMemory):
+  def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
+    return super(AnswerConversationBufferMemory, self).save_context(inputs,{'response': outputs['answer']})
+
+def getretriever():
+  with open("./resources/design-by-fire.pdf", 'rb') as uploaded_file:
+    try:
+      file = read_file(uploaded_file)
+    except Exception as e:
+      print(e)
+
+  chunked_file = chunk_file(file, chunk_size=512, chunk_overlap=0)
+  folder_index = embed_files(
+    files=[chunked_file],
+    embedding=EMBEDDING,
+    vector_store=VECTOR_STORE,
+    openai_api_key=API_KEY,
+  )
+  return folder_index.index.as_retriever(verbose=True, search_type="similarity", search_kwargs={"k": K})
+
+retriever = getretriever()
+
+def getanswer(chain, question, history):
+  if hasattr(chain, "value"):
+    chain = chain.value
+  if hasattr(history, "value"):
+    history = history.value
+  if hasattr(question, "value"):
+    question = question.value
+
+  history = history or []
+  lock = Lock()
+  lock.acquire()
+  try:
+    output = chain({"question": question})
+    output = output["answer"]
+    history.append((question, output))
+  except Exception as e:
+    raise e
+  finally:
+    lock.release()
+  return history, history, gr.update(value="")
+
+def load_chain(inputs = None):
+  llm = ChatOpenAI(
+        openai_api_key=API_KEY,
+        model_name=MODEL,
+        verbose=True)
+  chain = ConversationalRetrievalChain.from_llm(
+    llm,
+    retriever=retriever,
+    return_source_documents=USE_VERBOSE,
+    memory=AnswerConversationBufferMemory(memory_key="chat_history", return_messages=True),
+    verbose=USE_VERBOSE,
+    combine_docs_chain_kwargs={"prompt": qa_prompt})
+  return chain
+
+CSS ="""
+.contain { display: flex; flex-direction: column; }
+.gradio-container { height: 100vh !important; }
+#component-0 { height: 100%; }
+#chatbot { flex-grow: 1; overflow: auto;}
+"""
+
+with gr.Blocks() as block:
+  with gr.Row():
+    with gr.Column(scale=0.75):
+      with gr.Row():
+        gr.Markdown("<h1>Design by Fire</h1>")
+      with gr.Row():
+        gr.Markdown("by Emily Elizabeth Schlickman and Brett Milligan")
+      chatbot = gr.Chatbot(elem_id="chatbot").style(height=600)
+
+      with gr.Row():
+          message = gr.Textbox(
+              label="",
+              placeholder="Design by Fire",
+              lines=1,
+          )
+      with gr.Row():
+          submit = gr.Button(value="Send", variant="primary", scale=1)
+
+      state = gr.State()
+      chain_state = gr.State(load_chain)
+
+      submit.click(getanswer, inputs=[chain_state, message, state], outputs=[chatbot, state, message])
+      message.submit(getanswer, inputs=[chain_state, message, state], outputs=[chatbot, state, message])
+
+    with gr.Column(scale=0.25):
+      with gr.Row():
+        gr.Markdown("<h1><center>Suggestions</center></h1>")
+      ex1 = gr.Button(value="What are the main factors and trends discussed in the book that contribute to the changing behavior of wildfires?", variant="primary")
+      ex1.click(getanswer, inputs=[chain_state, ex1, state], outputs=[chatbot, state, message])
+      ex2 = gr.Button(value="How does the book explore the relationship between fire and different landscapes, such as wilderness and urban areas?", variant="primary")
+      ex2.click(getanswer, inputs=[chain_state, ex2, state], outputs=[chatbot, state, message])
+      ex3 = gr.Button(value="What are the three approaches to designing with fire that the book presents?", variant="primary")
+      ex3.click(getanswer, inputs=[chain_state, ex3, state], outputs=[chatbot, state, message])
+
+block.launch(debug=True)

src/components/faq.py

@@ -0,0 +1,46 @@
+# flake8: noqa
+import streamlit as st
+
+
+def faq():
+    st.markdown(
+        """
+# FAQ
+## How does KnowledgeGPT work?
+When you upload a document, it will be divided into smaller chunks 
+and stored in a special type of database called a vector index 
+that allows for semantic search and retrieval.
+
+When you ask a question, KnowledgeGPT will search through the
+document chunks and find the most relevant ones using the vector index.
+Then, it will use GPT3 to generate a final answer.
+
+## Is my data safe?
+Yes, your data is safe. KnowledgeGPT does not store your documents or
+questions. All uploaded data is deleted after you close the browser tab.
+
+## Why does it take so long to index my document?
+If you are using a free OpenAI API key, it will take a while to index
+your document. This is because the free API key has strict [rate limits](https://platform.openai.com/docs/guides/rate-limits/overview).
+To speed up the indexing process, you can use a paid API key.
+
+## What do the numbers mean under each source?
+For a PDF document, you will see a citation number like this: 3-12. 
+The first number is the page number and the second number is 
+the chunk number on that page. For DOCS and TXT documents, 
+the first number is set to 1 and the second number is the chunk number.
+
+## Are the answers 100% accurate?
+No, the answers are not 100% accurate. KnowledgeGPT uses GPT-3 to generate
+answers. GPT-3 is a powerful language model, but it sometimes makes mistakes 
+and is prone to hallucinations. Also, KnowledgeGPT uses semantic search
+to find the most relevant chunks and does not see the entire document,
+which means that it may not be able to find all the relevant information and
+may not be able to answer all questions (especially summary-type questions
+or questions that require a lot of context from the document).
+
+But for most use cases, KnowledgeGPT is very accurate and can answer
+most questions. Always check with the sources to make sure that the answers
+are correct.
+"""
+    )

src/core/caching.py

@@ -0,0 +1,32 @@
+import canonical_demo_memory.core.chunking as chunking
+import canonical_demo_memory.core.embedding as embedding
+import canonical_demo_memory.core.parsing as parsing
+import streamlit as st
+from canonical_demo_memory.core.parsing import File
+from streamlit.runtime.caching.hashing import HashFuncsDict
+
+
+def file_hash_func(file: File) -> str:
+    """Get a unique hash for a file"""
+    return file.id
+
+
+@st.cache_data(show_spinner=False)
+def bootstrap_caching():
+    """Patch module functions with caching"""
+
+    # Get all substypes of File from module
+    file_subtypes = [
+        cls
+        for cls in vars(parsing).values()
+        if isinstance(cls, type) and issubclass(cls, File) and cls != File
+    ]
+    file_hash_funcs: HashFuncsDict = {cls: file_hash_func for cls in file_subtypes}
+
+    parsing.read_file = st.cache_data(show_spinner=False)(parsing.read_file)
+    chunking.chunk_file = st.cache_data(show_spinner=False, hash_funcs=file_hash_funcs)(
+        chunking.chunk_file
+    )
+    embedding.embed_files = st.cache_data(
+        show_spinner=False, hash_funcs=file_hash_funcs
+    )(embedding.embed_files)

src/core/chunking.py

@@ -0,0 +1,61 @@
+from langchain.docstore.document import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+from .parsing import File
+
+
+def chunk_sentences(sentences, chunk_size=512):
+  sents = []
+  current_sent = ""
+
+  for sentence in sentences:
+    # If adding the next sentence doesn't exceed the chunk_size,
+    # we add the sentence to the current chunk.
+    if len(current_sent) + len(sentence) <= chunk_size:
+      current_sent += " " + sentence
+    else:
+      # If adding the sentence would make the chunk too long,
+      # we add the current_sent chunk to the list of chunks and start a new chunk.
+      sents.append(current_sent)
+      current_sent = sentence
+
+  # After going through all the sentences, there may be a chunk that hasn't yet been added to the list.
+  # We add it now:
+  if current_sent:
+    sents.append(current_sent)
+
+  return sents
+
+def chunk_file(
+    file: File, chunk_size: int, chunk_overlap: int = 0, model_name="gpt-3.5-turbo"
+) -> File:
+    """Chunks each document in a file into smaller documents
+    according to the specified chunk size and overlap
+    where the size is determined by the number of token for the specified model.
+    """
+
+    # split each document into chunks
+    chunked_docs = []
+    for doc in file.docs:
+        text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+            model_name=model_name,
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+        )
+
+        chunks = text_splitter.split_text(doc.page_content)
+
+        for i, chunk in enumerate(chunks):
+            doc = Document(
+                page_content=chunk,
+                metadata={
+                    "page": doc.metadata.get("page", 1),
+                    "chunk": i + 1,
+                    "source": f"{doc.metadata.get('page', 1)}-{i + 1}",
+                },
+            )
+            chunked_docs.append(doc)
+
+    chunked_file = file.copy()
+    chunked_file.docs = chunked_docs
+    return chunked_file

src/core/debug.py

@@ -0,0 +1,49 @@
+from langchain.vectorstores import VectorStore
+from typing import Iterable, List, Any
+from langchain.docstore.document import Document
+from langchain.embeddings.base import Embeddings
+from langchain.embeddings.fake import FakeEmbeddings as FakeEmbeddingsBase
+from langchain.chat_models.fake import FakeListChatModel
+from typing import Optional
+
+
+class FakeChatModel(FakeListChatModel):
+    def __init__(self, **kwargs):
+        responses = ["The answer is 42. SOURCES: 1, 2, 3, 4"]
+        super().__init__(responses=responses, **kwargs)
+
+
+class FakeEmbeddings(FakeEmbeddingsBase):
+    def __init__(self, **kwargs):
+        super().__init__(size=4, **kwargs)
+
+
+class FakeVectorStore(VectorStore):
+    """Fake vector store for testing purposes."""
+
+    def __init__(self, texts: List[str]):
+        self.texts: List[str] = texts
+
+    def add_texts(
+        self, texts: Iterable[str], metadatas: List[dict] | None = None, **kwargs: Any
+    ) -> List[str]:
+        self.texts.extend(texts)
+        return self.texts
+
+    @classmethod
+    def from_texts(
+        cls,
+        texts: List[str],
+        embedding: Embeddings,
+        metadatas: Optional[List[dict]] = None,
+        **kwargs: Any,
+    ) -> "FakeVectorStore":
+        return cls(texts=list(texts))
+
+    def similarity_search(
+        self, query: str, k: int = 4, **kwargs: Any
+    ) -> List[Document]:
+        return [
+            Document(page_content=text, metadata={"source": f"{i+1}-{1}"})
+            for i, text in enumerate(self.texts)
+        ]

src/core/embedding.py

@@ -0,0 +1,76 @@
+from typing import List, Type
+
+from langchain.docstore.document import Document
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.embeddings.base import Embeddings
+from langchain.vectorstores import VectorStore
+from langchain.vectorstores.faiss import FAISS
+
+from .debug import FakeEmbeddings, FakeVectorStore
+from .parsing import File
+
+
+class FolderIndex:
+    """Index for a collection of files (a folder)"""
+
+    def __init__(self, files: List[File], index: VectorStore):
+        self.name: str = "default"
+        self.files = files
+        self.index: VectorStore = index
+
+    @staticmethod
+    def _combine_files(files: List[File]) -> List[Document]:
+        """Combines all the documents in a list of files into a single list."""
+
+        all_texts = []
+        for file in files:
+            for doc in file.docs:
+                doc.metadata["file_name"] = file.name
+                doc.metadata["file_id"] = file.id
+                all_texts.append(doc)
+
+        return all_texts
+
+    @classmethod
+    def from_files(
+        cls, files: List[File], embeddings: Embeddings, vector_store: Type[VectorStore]
+    ) -> "FolderIndex":
+        """Creates an index from files."""
+
+        all_docs = cls._combine_files(files)
+
+        index = vector_store.from_documents(
+            documents=all_docs,
+            embedding=embeddings,
+        )
+
+        return cls(files=files, index=index)
+
+
+def embed_files(
+    files: List[File], embedding: str, vector_store: str, **kwargs
+) -> FolderIndex:
+    """Embeds a collection of files and stores them in a FolderIndex."""
+
+    supported_embeddings: dict[str, Type[Embeddings]] = {
+        "openai": OpenAIEmbeddings,
+        "debug": FakeEmbeddings,
+    }
+    supported_vector_stores: dict[str, Type[VectorStore]] = {
+        "faiss": FAISS,
+        "debug": FakeVectorStore,
+    }
+
+    if embedding in supported_embeddings:
+        _embeddings = supported_embeddings[embedding](**kwargs)
+    else:
+        raise NotImplementedError(f"Embedding {embedding} not supported.")
+
+    if vector_store in supported_vector_stores:
+        _vector_store = supported_vector_stores[vector_store]
+    else:
+        raise NotImplementedError(f"Vector store {vector_store} not supported.")
+
+    return FolderIndex.from_files(
+        files=files, embeddings=_embeddings, vector_store=_vector_store
+    )

src/core/parsing.py

@@ -0,0 +1,105 @@
+import re
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from hashlib import md5
+from io import BytesIO
+from typing import Any, List, Optional
+
+import docx2txt
+import fitz
+from langchain.docstore.document import Document
+
+
+
+class File(ABC):
+    """Represents an uploaded file comprised of Documents"""
+
+    def __init__(
+        self,
+        name: str,
+        id: str,
+        metadata: Optional[dict[str, Any]] = None,
+        docs: Optional[List[Document]] = None,
+    ):
+        self.name = name
+        self.id = id
+        self.metadata = metadata or {}
+        self.docs = docs or []
+
+    @classmethod
+    @abstractmethod
+    def from_bytes(cls, file: BytesIO) -> "File":
+        """Creates a File from a BytesIO object"""
+
+    def __repr__(self) -> str:
+        return (
+            f"File(name={self.name}, id={self.id},"
+            " metadata={self.metadata}, docs={self.docs})"
+        )
+
+    def __str__(self) -> str:
+        return f"File(name={self.name}, id={self.id}, metadata={self.metadata})"
+
+    def copy(self) -> "File":
+        """Create a deep copy of this File"""
+        return self.__class__(
+            name=self.name,
+            id=self.id,
+            metadata=deepcopy(self.metadata),
+            docs=deepcopy(self.docs),
+        )
+
+
+def strip_consecutive_newlines(text: str) -> str:
+    """Strips consecutive newlines from a string
+    possibly with whitespace in between
+    """
+    return re.sub(r"\s*\n\s*", "\n", text)
+
+
+class DocxFile(File):
+    @classmethod
+    def from_bytes(cls, file: BytesIO) -> "DocxFile":
+        text = docx2txt.process(file)
+        text = strip_consecutive_newlines(text)
+        doc = Document(page_content=text.strip())
+        return cls(name=file.name, id=md5(file.read()).hexdigest(), docs=[doc])
+
+
+class PdfFile(File):
+    @classmethod
+    def from_bytes(cls, file: BytesIO) -> "PdfFile":
+        pdf = fitz.open(stream=file.read(), filetype="pdf")  # type: ignore
+        docs = []
+        for i, page in enumerate(pdf):
+            text = page.get_text(sort=True)
+            text = strip_consecutive_newlines(text)
+            doc = Document(page_content=text.strip())
+            doc.metadata["page"] = i + 1
+            docs.append(doc)
+        # file.read() mutates the file object, which can affect caching
+        # so we need to reset the file pointer to the beginning
+        file.seek(0)
+        return cls(name=file.name, id=md5(file.read()).hexdigest(), docs=docs)
+
+
+class TxtFile(File):
+    @classmethod
+    def from_bytes(cls, file: BytesIO) -> "TxtFile":
+        text = file.read().decode("utf-8")
+        text = strip_consecutive_newlines(text)
+        file.seek(0)
+        doc = Document(page_content=text.strip())
+        return cls(name=file.name, id=md5(file.read()).hexdigest(), docs=[doc])
+
+
+def read_file(file: BytesIO) -> File:
+    """Reads an uploaded file and returns a File object"""
+    if file.name.lower().endswith(".docx"):
+        return DocxFile.from_bytes(file)
+    elif file.name.lower().endswith(".pdf"):
+        return PdfFile.from_bytes(file)
+    elif file.name.lower().endswith(".txt"):
+        return TxtFile.from_bytes(file)
+    else:
+        raise NotImplementedError(f"File type {file.name.split('.')[-1]} not supported")

src/core/prompts.py

@@ -0,0 +1,31 @@
+# flake8: noqa
+from langchain.prompts import PromptTemplate
+
+## Use a shorter template to reduce the number of tokens in the prompt
+template = """Create a final answer to the given questions using the provided document excerpts(in no particular order) as references. ALWAYS include a "SOURCES" section in your answer including only the minimal set of sources needed to answer the question. If you are unable to answer the question, simply state that you do not know. Do not attempt to fabricate an answer and leave the SOURCES section empty.
+
+---------
+
+QUESTION: What  is the purpose of ARPA-H?
+=========
+Content: More support for patients and families. \n\nTo get there, I call on Congress to fund ARPA-H, the Advanced Research Projects Agency for Health. \n\nIt’s based on DARPA—the Defense Department project that led to the Internet, GPS, and so much more.  \n\nARPA-H will have a singular purpose—to drive breakthroughs in cancer, Alzheimer’s, diabetes, and more.
+Source: 1-32
+Content: While we’re at it, let’s make sure every American can get the health care they need. \n\nWe’ve already made historic investments in health care. \n\nWe’ve made it easier for Americans to get the care they need, when they need it. \n\nWe’ve made it easier for Americans to get the treatments they need, when they need them. \n\nWe’ve made it easier for Americans to get the medications they need, when they need them.
+Source: 1-33
+Content: The V.A. is pioneering new ways of linking toxic exposures to disease, already helping  veterans get the care they deserve. \n\nWe need to extend that same care to all Americans. \n\nThat’s why I’m calling on Congress to pass legislation that would establish a national registry of toxic exposures, and provide health care and financial assistance to those affected.
+Source: 1-30
+=========
+FINAL ANSWER: The purpose of ARPA-H is to drive breakthroughs in cancer, Alzheimer’s, diabetes, and more.
+SOURCES: 1-32
+
+---------
+
+QUESTION: {question}
+=========
+{summaries}
+=========
+FINAL ANSWER:"""
+
+STUFF_PROMPT = PromptTemplate(
+    template=template, input_variables=["summaries", "question"]
+)

src/core/qa.py

@@ -0,0 +1,78 @@
+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