Update utils/database.py

utils/database.py

@@ -11,7 +11,8 @@ from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, Base
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
 from langchain.agents import initialize_agent
-from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, # utils/database.py
+from langchain_core.runnables import RunnablePassthrough
 
 def create_connection(db_file):
     try:
@@ -104,17 +105,13 @@ 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.5,
             model_name="gpt-4",
-            api_key=os.environ.get("OPENAI_API_KEY"),
+            api_key=os.environ.get("OPENAI_API_KEY")
         )
 
         # Create chat memory
@@ -124,37 +121,45 @@ def initialize_qa_system(vector_store):
             k=5
         )
 
-        # Create the base QA chain
-        qa = ConversationalRetrievalChain.from_llm(
-            llm=llm,
-            retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
-            chain_type="stuff",
-        )
-
-        # Define the tools
-        tools = [
-            Tool(
-                name="RFP_Knowledge_Base",
-                func=qa.run,
-                description="Use this tool to analyze RFP documents and answer questions about their content."
-            )
-        ]
-
         # Create the prompt template
         prompt = ChatPromptTemplate.from_messages([
             ("system", "You are a helpful assistant analyzing RFP documents."),
             MessagesPlaceholder(variable_name="chat_history"),
-            ("human", "{input}"),
-            MessagesPlaceholder(variable_name="agent_scratchpad"),
+            ("human", "{question}"),
+            MessagesPlaceholder(variable_name="agent_scratchpad")
         ])
 
-        # Create the agent
-        agent = create_openai_tools_agent(llm, tools, prompt)
+        # Create the RAG chain with lambda function
+        rag_chain = (
+            {
+                "context": lambda x: vector_store.as_retriever().get_relevant_documents(x["question"]),
+                "question": RunnablePassthrough(),
+                "chat_history": lambda x: memory.chat_memory.messages
+            }
+            | prompt 
+            | llm
+        )
 
         # Create the agent executor
         agent_executor = AgentExecutor(
-            agent=agent,
-            tools=tools,
+            agent=create_openai_tools_agent(
+                llm=llm,
+                tools=[
+                    Tool(
+                        name="RFP_Knowledge_Base",
+                        func=rag_chain.invoke,
+                        description="Use this tool to analyze RFP documents and answer questions about their content."
+                    )
+                ],
+                prompt=prompt
+            ),
+            tools=[
+                Tool(
+                    name="RFP_Knowledge_Base",
+                    func=rag_chain.invoke,
+                    description="Use this tool to analyze RFP documents and answer questions about their content."
+                )
+            ],
             memory=memory,
             verbose=True,
             handle_parsing_errors=True