Update utils/database.py

utils/database.py

@@ -1646,18 +1646,19 @@ def initialize_qa_system(vector_store):
         dict: QA system chain or None if initialization fails.
     """
     try:
-        llm = ChatOpenAI(
+          llm = ChatOpenAI(
             temperature=0.5,
             model_name="gpt-4",
-            max_tokens=4000,  # Explicitly set max tokens
+            max_tokens=4000,
             api_key=os.environ.get("OPENAI_API_KEY")
         )
 
-        # Optimize retriever settings
+        # Optimize retriever settings and add source tracking
         retriever = vector_store.as_retriever(
             search_kwargs={
                 "k": 3,  # Retrieve fewer, more relevant chunks
-                "fetch_k": 5  # Consider more candidates before selecting top k
+                "fetch_k": 5,  # Consider more candidates before selecting top k
+                "include_metadata": True  # Enable source tracking
             }
         )
 
@@ -1691,7 +1692,7 @@ Accuracy: Double-check all information for accuracy and completeness before prov
             MessagesPlaceholder(variable_name="chat_history"),
             ("human", "{input}\n\nContext: {context}")
         ])
-
+        
         def get_chat_history(inputs):
             chat_history = inputs.get("chat_history", [])
             if not isinstance(chat_history, list):
@@ -1708,8 +1709,8 @@ Accuracy: Double-check all information for accuracy and completeness before prov
 
         chain = (
             {
-                "context": get_context,
-                "chat_history": get_chat_history,
+                "context": lambda x: get_context_with_sources(retriever, x["input"]),
+                "chat_history": lambda x: format_chat_history(x["chat_history"]),
                 "input": lambda x: x["input"]
             }
             | prompt