feat: remove cache, add context expander

.gitignore

@@ -1,3 +1,4 @@
 /__pycache__
 /temp
-/models
+/models
+/chroma

app.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import streamlit as st
 from llm import load_llm, response_generator
 from vector_store import load_vector_store, process_pdf
@@ -9,20 +10,26 @@ from uuid import uuid4
 repo_id = "Qwen/Qwen2.5-3B-Instruct-GGUF"
 filename = "qwen2.5-3b-instruct-q5_k_m.gguf"
 
+
 llm = load_llm(repo_id, filename)
+vector_store = load_vector_store()
+
 
 st.title("PDF QA")
 # Initialize chat history
 if "messages" not in st.session_state:
+    vector_store.reset_collection()
+    if os.path.exists("./temp"):
+        shutil.rmtree("./temp")
     st.session_state.messages = []
 
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         if message["role"] == "user":
-            st.markdown(message["content"])
+            st.write(message["content"])
         else:
-            st.code(message["content"])
+            st.write(message["content"])
 
 # Accept user input
 if prompt := st.chat_input("What is up?"):
@@ -34,13 +41,12 @@ if prompt := st.chat_input("What is up?"):
 
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
-        vector_store = load_vector_store()
-        retriever = vector_store.as_retriever()
-        docs = retriever.get_relevant_documents(prompt)
-
+        retriever = vector_store.as_retriever(search_kwargs={"k": 3})
         response = response_generator(llm, st.session_state.messages, prompt, retriever)
 
         st.markdown(response["answer"])
+        with st.expander("See context"):
+            st.write(response["context"])
 
     # Add assistant response to chat history
     st.session_state.messages.append(
@@ -54,7 +60,7 @@ with st.sidebar:
         "Choose a PDF file", accept_multiple_files=True, type="pdf"
     )
     if uploaded_files is not None:
-        vector_store = load_vector_store()
+        st.session_state.uploaded_pdf = True
         for uploaded_file in uploaded_files:
             temp_dir = "./temp"
             if not os.path.exists(temp_dir):

llm.py

@@ -7,6 +7,10 @@ from langchain.chains import create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import ChatPromptTemplate
 
+from langchain_core.globals import set_debug
+
+set_debug(True)
+
 
 @st.cache_resource()
 def load_llm(repo_id, filename):
@@ -29,6 +33,8 @@ def load_llm(repo_id, filename):
         n_threads=4,
         n_threads_batch=4,
         n_ctx=8000,
+        max_tokens=128,
+        # stop=["."],
     )
     print(f"{repo_id} loaded successfully. ✅")
     return llm
@@ -36,26 +42,45 @@ def load_llm(repo_id, filename):
 
 # Streamed response emulator
 def response_generator(llm, messages, question, retriever):
+    # System prompt setting up context for the assistant
     system_prompt = (
-        "You are an assistant for question-answering tasks. "
-        "Use the following pieces of retrieved context to answer "
-        "the question. If you don't know the answer, say that you "
-        "don't know. Use three sentences maximum and keep the "
-        "answer concise."
+        "<|im_start|>system\n"
+        "You are an AI assistant specializing in question-answering tasks. "
+        "Utilize the provided context and past conversation to answer "
+        "the current question. If the answer is unknown, clearly state that you "
+        "don't know. Keep responses concise and direct."
         "\n\n"
-        "{context}"
+        "Context: {context}"
+        "\n<|im_end|>"
     )
 
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            ("system", system_prompt),
-            ("user", "{input}"),
-        ]
-    )
+    # Prepare message history
+    message_history = [("system", system_prompt)]
+
+    # Append conversation history to messages
+    for message in messages:
+        if message["role"] == "user":
+            message_history.append(
+                ("user", "<|im_start|>user\n" + message["content"] + "\n<|im_end|>")
+            )
+        elif message["role"] == "assistant":
+            message_history.append(
+                (
+                    "assistant",
+                    "<|im_start|>assistant\n" + message["content"] + "\n<|im_end|>",
+                )
+            )
+
+    message_history.append(("assistant", "<|im_start|>assistant\n"))
+
+    # Create prompt template with full message history
+    prompt = ChatPromptTemplate.from_messages(message_history)
 
+    # Instantiate chains for document retrieval and question answering
     question_answer_chain = create_stuff_documents_chain(llm, prompt)
     rag_chain = create_retrieval_chain(retriever, question_answer_chain)
 
-    results = rag_chain.invoke({"input": question})
+    # Invoke RAG (retrieval-augmented generation) chain with current input
+    results = rag_chain.invoke({"input": question}, verbose=True)
 
     return results

vector_store.py

@@ -1,7 +1,8 @@
 import streamlit as st
-import chromadb
+
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_chroma import Chroma
+from langchain_community.vectorstores import InMemoryVectorStore
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 
@@ -16,17 +17,14 @@ def load_embedding_model(model):
     return model
 
 
+@st.cache_resource()
 def load_vector_store():
-    """
-    Loads a simple vector store
-    I didn't use @st.cache because I want to
-    load vector store on every page load
-    """
-    model = load_embedding_model("sentence-transformers/all-MiniLM-L6-v2")
-    chromadb.api.client.SharedSystemClient.clear_system_cache()
+    model = load_embedding_model("sentence-transformers/all-mpnet-base-v2")
+
     vector_store = Chroma(
         collection_name="main_store",
         embedding_function=model,
+        persist_directory="./chroma",
     )
     return vector_store