Update utils/database.py

utils/database.py

@@ -337,15 +337,17 @@ def display_vector_store_info():
         st.error(traceback.format_exc())
         
 def initialize_qa_system(vector_store):
-    """Initialize QA system with clean response formatting."""
+    """Initialize QA system with proper chat handling."""
     try:
         llm = ChatOpenAI(
             temperature=0.5,
             model_name="gpt-4",
             api_key=os.environ.get("OPENAI_API_KEY")
         )
-# Create retriever function
+        
+        # Create retriever function
         retriever = vector_store.as_retriever(search_kwargs={"k": 2})
+        
         # Create a template that enforces clean formatting
         prompt = ChatPromptTemplate.from_messages([
             ("system", """You are a helpful assistant analyzing RFP documents. 
@@ -362,6 +364,22 @@ def initialize_qa_system(vector_store):
             ("human", "{input}\n\nContext: {context}")
         ])
 
+        def get_chat_history(inputs):
+            """Get formatted chat history."""
+            chat_history = inputs.get("chat_history", [])
+            if not isinstance(chat_history, list):
+                return []
+            return [msg for msg in chat_history if isinstance(msg, BaseMessage)]
+
+        def get_context(inputs):
+            """Get formatted context from documents."""
+            docs = retriever.get_relevant_documents(inputs["input"])
+            context_parts = []
+            for doc in docs:
+                source = doc.metadata.get('source', 'Unknown source')
+                context_parts.append(f"\nFrom {source}:\n{doc.page_content}")
+            return "\n".join(context_parts)
+
         def format_response(response_text):
             """Clean up the response formatting."""
             # Remove technical metadata
@@ -375,20 +393,10 @@ def initialize_qa_system(vector_store):
             response_text = response_text.replace('\\n', '\n')
             
             return response_text
-        def get_context(inputs):
-            docs = retriever.get_relevant_documents(inputs["input"])
-            context_parts = []
-            for doc in docs:
-                source = doc.metadata.get('source', 'Unknown source')
-                context_parts.append(f"\nFrom {source}:\n{doc.page_content}")
-            return "\n".join(context_parts)
 
         chain = (
             {
-                "context": lambda x: "\n".join(
-                    f"\nFrom {doc.metadata.get('source', 'Unknown')}:\n{doc.page_content}"
-                    for doc in retriever.get_relevant_documents(x["input"])
-                ),
+                "context": get_context,
                 "chat_history": get_chat_history,
                 "input": lambda x: x["input"]
             }