Update utils/database.py

utils/database.py

@@ -7,10 +7,10 @@ from datetime import datetime
 from langchain.chat_models import ChatOpenAI
 import os
 from langchain.memory import ConversationBufferWindowMemory
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
-from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate
 import os
 
 def create_connection(db_file):
@@ -91,6 +91,27 @@ def insert_document(conn, doc_name, doc_content):
         st.error(f"Error inserting document: {e}")
     return False
 
+# utils/database.py
+from langchain.memory import ConversationBufferWindowMemory
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
+from langchain.chains import ConversationalRetrievalChain
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate
+import os
+import streamlit as st
+
+def format_chat_history(messages: list[BaseMessage]) -> list[dict]:
+    """Convert chat history to the format expected by langchain"""
+    formatted = []
+    for msg in messages:
+        if isinstance(msg, HumanMessage):
+            formatted.append({"role": "user", "content": msg.content})
+        elif isinstance(msg, AIMessage):
+            formatted.append({"role": "assistant", "content": msg.content})
+        elif isinstance(msg, SystemMessage):
+            formatted.append({"role": "system", "content": msg.content})
+    return formatted
+
 def initialize_qa_system(vector_store):
     """Initialize QA system with proper chat handling"""
     try:
@@ -100,12 +121,14 @@ def initialize_qa_system(vector_store):
             api_key=os.environ.get("OPENAI_API_KEY"),
         )
 
-        # Create the proper prompt template
+        # Create a custom prompt template
         prompt = ChatPromptTemplate.from_messages([
-            ("system", "You are an AI assistant analyzing RFP documents. Provide clear and concise answers based on the document content."),
+            SystemMessage(content=(
+                "You are an AI assistant analyzing RFP documents. "
+                "Provide clear and concise answers based on the document content."
+            )),
             MessagesPlaceholder(variable_name="chat_history"),
-            ("user", MessagesPlaceholder(variable_name="input")),
-            MessagesPlaceholder(variable_name="agent_scratchpad")
+            HumanMessagePromptTemplate.from_template("{input}"),
         ])
 
         # Initialize memory with proper configuration
@@ -119,9 +142,9 @@ def initialize_qa_system(vector_store):
             llm=llm,
             retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
             memory=memory,
-            combine_docs_chain_kwargs={"prompt": prompt},
             verbose=True,
-            return_source_documents=True
+            return_source_documents=True,
+            combine_docs_chain_kwargs={"prompt": prompt}
         )
 
         return qa_chain