chat questions and answer retains.

app.py

@@ -1,5 +1,12 @@
 from rag_pipeline import process_file, answer_query
 from pydantic import BaseModel
+from typing import List, Dict, Optional
+from datetime import datetime
+
+class ChatMessage(BaseModel):
+    question: str
+    answer: str
+    timestamp: datetime
 
 class QueryRequest(BaseModel):
     file_id: str
@@ -22,7 +29,6 @@ load_dotenv()
 CHROMA_DIR = "./chroma_db"
 embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
 
-
 app = FastAPI()
 BASE_DIR = "files"
 app.add_middleware(
@@ -32,6 +38,10 @@ app.add_middleware(
     allow_headers=["*"],
 )
 file_store = {}
+
+# Add chat memory store
+chat_memory: Dict[str, List[ChatMessage]] = {}
+
 @app.get("/test")
 async def test():
     return {"message": "hello world!"}
@@ -104,9 +114,25 @@ async def query_endpoint(request = Body(...)):
             for context in contexts
         )
 
+        # Add chat history to context if it exists
+        if file_id in chat_memory and chat_memory[file_id]:
+            chat_history = "\n\nPrevious Conversation:\n"
+            for msg in chat_memory[file_id][-3:]:  # Include last 3 exchanges
+                chat_history += f"Q: {msg.question}\nA: {msg.answer}\n\n"
+            formatted_context = chat_history + formatted_context
+
         # Get the answer using the enhanced context
         answer = answer_query(question, formatted_context, explain_like_5)
         
+        # Store the Q&A in chat memory
+        if file_id not in chat_memory:
+            chat_memory[file_id] = []
+        chat_memory[file_id].append(ChatMessage(
+            question=question,
+            answer=answer,
+            timestamp=datetime.now()
+        ))
+        
         return {
             "answer": answer,
             "context_used": formatted_context  # Optionally return context for debugging
@@ -146,6 +172,10 @@ async def delete_file(file_id: str):
                 except Exception as e:
                     print(f"Error deleting file {filename}: {str(e)}")
 
+        # 3. Clear chat memory for this file
+        if file_id in chat_memory:
+            del chat_memory[file_id]
+
         if not matching_files and not os.path.exists(chroma_path):
             raise HTTPException(
                 status_code=404,
@@ -153,7 +183,7 @@ async def delete_file(file_id: str):
             )
 
         return {
-            "message": "File and its embeddings deleted successfully",
+            "message": "File, embeddings, and chat history deleted successfully",
             "deleted_files": matching_files,
             "embeddings_deleted": os.path.exists(chroma_path)
         }

rag_pipeline.py

@@ -91,18 +91,19 @@ def answer_query(question, context, explain_like_5=False):
         context = "\n\n".join(str(c) for c in context)
     
     system_prompt = (
-        "You are a helpful assistant answering user queries based STRICTLY on the provided document chunks.\n"
+        "You are a helpful assistant answering user queries based on the provided document chunks and conversation history.\n"
         "IMPORTANT RULES:\n"
-        "1. ONLY use information from the given context. Do not use any external knowledge.\n"
-        "2. If the answer cannot be fully derived from the context, say 'I cannot answer this question based on the provided context.'\n"
+        "1. Use information from both the document context and previous conversation history.\n"
+        "2. If the answer cannot be fully derived from the context or conversation history, say 'I cannot answer this question based on the provided information.'\n"
         "3. If you're unsure about any part of the answer, acknowledge the uncertainty.\n"
-        "4. Do not make assumptions beyond what's explicitly stated in the context.\n"
-        "5. Quote relevant parts of the context to support your answers when possible."
+        "4. Do not make assumptions beyond what's explicitly stated in the context or conversation history.\n"
+        "5. Quote relevant parts of the context to support your answers when possible.\n"
+        "6. When referencing previous conversation, be clear about which information came from where."
     )
     
     if explain_like_5:
-        system_prompt += "\nExplain the answer in a simple way, like you're talking to a 5-year-old, but still only use information from the context."
-    print("Context:", context)
+        system_prompt += "\nExplain the answer in a simple way, like you're talking to a 5-year-old, but still only use information from the context and conversation history."
+    
     try:
         # Send to LLM with formatted prompt
         response = client.chat.completions.create(
@@ -110,10 +111,10 @@ def answer_query(question, context, explain_like_5=False):
             messages=[
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": (
-                    f"Context:\n{context}\n\n"
+                    f"Context and Conversation History:\n{context}\n\n"
                     f"Question: {question}\n\n"
-                    "Remember to answer ONLY based on the information provided in the context above. "
-                    "If you cannot find the answer in the context, say so explicitly."
+                    "Remember to answer based on both the document context and conversation history above. "
+                    "If you cannot find the answer in either, say so explicitly."
                 )}
             ],
             temperature=0.3  # Lower temperature for more focused answers