Update utils/database.py

utils/database.py

@@ -130,30 +130,23 @@ def initialize_qa_system(vector_store):
             api_key=os.environ.get("OPENAI_API_KEY")
         )
 
-        # Define the template strings
-        system_template = "You are a helpful assistant analyzing RFP documents. Use the provided context to answer questions."
-        human_template = "Question: {input}\nContext: {context}"
-
-        # Create the prompt template using the simpler format
-        chat_prompt = ChatPromptTemplate.from_messages([
-            ("system", system_template),
-            ("human", human_template)
+        # Create the prompt template
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", "You are a helpful assistant analyzing RFP documents."),
+            MessagesPlaceholder(variable_name="chat_history"),
+            ("human", "{input}\nContext: {context}")
         ])
 
         # Create retriever function
         retriever = vector_store.as_retriever(search_kwargs={"k": 2})
 
-        # Create the chain
-        qa_chain = (
-            {
-                "context": lambda x: retriever.get_relevant_documents(x["input"]),
-                "input": lambda x: x["input"]
-            }
-            | chat_prompt
-            | llm
-        )
+        # Create the chain with proper chat history handling
+        chain = RunnablePassthrough.assign(
+            context=lambda x: "\n".join(doc.page_content for doc in retriever.get_relevant_documents(x["input"])),
+            chat_history=lambda x: x.get("chat_history", [])
+        ) | prompt | llm
 
-        return qa_chain
+        return chain
 
     except Exception as e:
         st.error(f"Error initializing QA system: {e}")