feat: 02_10

app/app.py

@@ -12,12 +12,11 @@ from chainlit.types import AskFileResponse
 
 import chromadb
 from chromadb.config import Settings
-from langchain.chains import LLMChain
+from langchain.chains import LLMChain, RetrievalQAWithSourcesChain
 from langchain.chat_models import ChatOpenAI
 from langchain.document_loaders import PDFPlumberLoader
 from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.prompts import ChatPromptTemplate
-from langchain.schema import Document, StrOutputParser
+from langchain.schema import Document
 from langchain.schema.embeddings import Embeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
@@ -144,24 +143,28 @@ async def on_chat_start():
     msg.content = f"`{file.name}` loaded. You can now ask questions!"
     await msg.update()
 
+    ##########################################################################
+    # Exercise 1a:
+    # Now we have search engine setup, our Chat with PDF application can do
+    # RAG architecture pattern. Please use the appropriate RetrievalQA Chain
+    # from Langchain. 
+    # 
+    # Remember, we would want to set the model temperature to
+    # 0 to ensure model outputs do not vary across runs, and we would want to
+    # also return sources to our answers.
+    ##########################################################################
     model = ChatOpenAI(
         model="gpt-3.5-turbo-16k-0613",
+        temperature=0,
         streaming=True
     )
 
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            (
-                "system",
-                "You are Chainlit GPT, a helpful assistant.",
-            ),
-            (
-                "human",
-                "{question}"
-            ),
-        ]
+    chain = RetrievalQAWithSourcesChain.from_chain_type(
+        llm=model,
+        chain_type="stuff",
+        retriever=search_engine.as_retriever(max_tokens_limit=4097),
     )
-    chain = LLMChain(llm=model, prompt=prompt, output_parser=StrOutputParser())
+    ##########################################################################
 
     # We are saving the chain in user_session, so we do not have to rebuild
     # it every single time.
@@ -172,11 +175,51 @@ async def on_chat_start():
 async def main(message: cl.Message):
 
     # Let's load the chain from user_session
-    chain = cl.user_session.get("chain")  # type: LLMChain
+    chain = cl.user_session.get("chain")  # type: RetrievalQAWithSourcesChain
 
-    response = await chain.arun(
-        question=message.content, callbacks=[cl.LangchainCallbackHandler()]
+    response = await chain.acall(
+        message.content,
+        callbacks=[cl.AsyncLangchainCallbackHandler(stream_final_answer=True)]
     )
-
-    await cl.Message(content=response).send()
-
+    answer = response["answer"]
+    sources = response["sources"].strip()
+    source_elements = []
+
+    # Get the documents from the user session
+    docs = cl.user_session.get("docs")
+    metadatas = [doc.metadata for doc in docs]
+    all_sources = [m["source"] for m in metadatas]
+
+    ##########################################################################
+    # Exercise 1a:
+    # Now we have search engine setup, our Chat with PDF application can do
+    # RAG architecture pattern. Please use the appropriate RetrievalQA Chain
+    # from Langchain. 
+    # 
+    # Remember, we would want to set the model temperature to
+    # 0 to ensure model outputs do not vary across runs, and we would want to
+    # also return sources to our answers.
+    ##########################################################################
+    # Adding sources to the answer
+    if sources:
+        found_sources = []
+
+        # Add the sources to the message
+        for source in sources.split(","):
+            source_name = source.strip().replace(".", "")
+            # Get the index of the source
+            try:
+                index = all_sources.index(source_name)
+            except ValueError:
+                continue
+            text = docs[index].page_content
+            found_sources.append(source_name)
+            # Create the text element referenced in the message
+            source_elements.append(cl.Text(content=text, name=source_name))
+
+        if found_sources:
+            answer += f"\nSources: {', '.join(found_sources)}"
+        else:
+            answer += "\nNo sources found"
+
+    await cl.Message(content=answer, elements=source_elements).send()