Update utils/database.py

utils/database.py

@@ -1,18 +1,20 @@
 # utils/database.py
-from langchain.memory import ConversationBufferWindowMemory
+# Update the imports first
+from langchain_community.chat_models import ChatOpenAI
 from langchain_core.messages import (
     HumanMessage,
     AIMessage,
     SystemMessage,
-    BaseMessage  # Added this import
+    BaseMessage
 )
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.runnables import RunnablePassthrough
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
 from langchain.agents import AgentExecutor, Tool, create_openai_tools_agent
-from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.agents.format_scratchpad.tools import format_to_tool_messages
 from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
-from langchain_core.runnables import RunnablePassthrough
+
 import os
 import streamlit as st
 import sqlite3
@@ -130,21 +132,35 @@ def initialize_qa_system(vector_store):
             api_key=os.environ.get("OPENAI_API_KEY")
         )
 
-        # Create the prompt template
+        # Create retriever function
+        retriever = vector_store.as_retriever(search_kwargs={"k": 2})
+
+        # Create a simpler prompt template
         prompt = ChatPromptTemplate.from_messages([
             ("system", "You are a helpful assistant analyzing RFP documents."),
             MessagesPlaceholder(variable_name="chat_history"),
-            ("human", "{input}\nContext: {context}")
+            ("human", "{question}")
         ])
 
-        # Create retriever function
-        retriever = vector_store.as_retriever(search_kwargs={"k": 2})
-
-        # 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
+        def get_chat_history(inputs):
+            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):
+            docs = retriever.get_relevant_documents(inputs["question"])
+            return "\n".join(doc.page_content for doc in docs)
+
+        chain = (
+            {
+                "question": lambda x: x["input"],
+                "chat_history": get_chat_history,
+                "context": get_context
+            }
+            | prompt
+            | llm
+        )
 
         return chain