Upload 2 files

app.py

@@ -2,12 +2,14 @@ import os
 from pathlib import Path
 
 from langchain.chains import ConversationalRetrievalChain
-from langchain.embeddings import OpenAIEmbeddings
 from langchain.vectorstores import Chroma
-from langchain.llms.openai import OpenAIChat
+from langchain.llms.openai import OpenAIChat, OpenAI
 from langchain.document_loaders import PyPDFLoader, WebBaseLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import LLMChainExtractor
+from langchain_experimental.text_splitter import SemanticChunker
 
 import streamlit as st
 
@@ -29,7 +31,7 @@ def load_documents():
 
 
 def split_documents(documents):
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=256)
+    text_splitter = SemanticChunker(OpenAIEmbeddings())
     texts = text_splitter.split_documents(documents)
     return texts
 
@@ -41,7 +43,10 @@ def embeddings_on_local_vectordb(texts):
         persist_directory=LOCAL_VECTOR_STORE_DIR.as_posix(),
     )
     vectordb.persist()
-    retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+    retriever = ContextualCompressionRetriever(
+        base_compressor=LLMChainExtractor.from_llm(OpenAI(temperature=0)),
+        base_retriever=vectordb.as_retriever(search_kwargs={"k": 3}, search_type="mmr"),
+    )
     return retriever
 
 
@@ -51,10 +56,11 @@ def query_llm(retriever, query):
         retriever=retriever,
         return_source_documents=True,
     )
+    relevant_docs = retriever.get_relevant_documents(query)
     result = qa_chain({"question": query, "chat_history": st.session_state.messages})
     result = result["answer"]
     st.session_state.messages.append((query, result))
-    return result
+    return relevant_docs, result
 
 
 def input_fields():
@@ -77,6 +83,8 @@ def boot():
     st.title("Enigma Chatbot")
     input_fields()
     st.sidebar.button("Submit Documents", on_click=process_documents)
+    st.sidebar.write("---")
+    st.sidebar.write("References made during the chat will appear here")
     if "messages" not in st.session_state:
         st.session_state.messages = []
     for message in st.session_state.messages:
@@ -84,7 +92,9 @@ def boot():
         st.chat_message("ai").write(message[1])
     if query := st.chat_input():
         st.chat_message("human").write(query)
-        response = query_llm(st.session_state.retriever, query)
+        references, response = query_llm(st.session_state.retriever, query)
+        for doc in references:
+            st.sidebar.info(f"Page {doc.metadata['page']}\n\n{doc.page_content}")
         st.chat_message("ai").write(response)
 
 

requirements.txt

@@ -1,4 +1,5 @@
 openai==0.28
 langchain==0.1.1
 pypdf==4.0.0
-chromadb==0.4.22
+chromadb==0.4.22
+langchain-experimental==0.0.49