Update utils/database.py

utils/database.py

@@ -11,6 +11,7 @@ from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, Base
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate
+from langchain.agents import initialize_agent
 
 def create_connection(db_file):
     try:
@@ -90,14 +91,6 @@ 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"""
@@ -111,47 +104,54 @@ def format_chat_history(messages: list[BaseMessage]) -> list[dict]:
             formatted.append({"role": "system", "content": msg.content})
     return formatted
 
+
+
 def initialize_qa_system(vector_store):
     """Initialize QA system with proper chat handling"""
     try:
         llm = ChatOpenAI(
-            temperature=0,
+            temperature=0.5,
             model_name="gpt-4",
             api_key=os.environ.get("OPENAI_API_KEY"),
         )
 
-        # Create a custom prompt template
-        prompt = ChatPromptTemplate.from_messages([
-            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"),
-            HumanMessagePromptTemplate.from_template("{input}"),
-        ])
-
-        # Initialize memory with proper configuration
+        # Create chat memory
         memory = ConversationBufferWindowMemory(
             memory_key="chat_history",
             return_messages=True,
-            k=5  # Keep last 5 interactions
+            k=5
         )
 
-        qa_chain = ConversationalRetrievalChain.from_llm(
+        # Create retrieval QA chain
+        qa = ConversationalRetrievalChain.from_llm(
             llm=llm,
             retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
-            memory=memory,
+            chain_type="stuff",
+        )
+
+        # Create tool for the agent
+        qa_tool = Tool(
+            name='Knowledge Base',
+            func=qa.run,
+            description='use this tool when answering questions about the RFP documents'
+        )
+
+        # Initialize agent
+        agent = initialize_agent(
+            agent='chat-conversational-react-description',
+            tools=[qa_tool],
+            llm=llm,
             verbose=True,
-            return_source_documents=True,
-            combine_docs_chain_kwargs={"prompt": prompt}
+            max_iterations=3,
+            early_stopping_method='generate',
+            memory=memory,
         )
 
-        return qa_chain
+        return agent
 
     except Exception as e:
         st.error(f"Error initializing QA system: {e}")
         return None
-
 def initialize_faiss(embeddings, documents, document_names):
     """Initialize FAISS vector store"""
     try: