Update utils/database.py

utils/database.py

@@ -1,18 +1,15 @@
 # utils/database.py
-import streamlit as st
-import sqlite3
-from sqlite3 import Error
-from datetime import datetime
-#from langchain.memory import ConversationBufferMemory
-from langchain.chat_models import ChatOpenAI
-import os
 from langchain.memory import ConversationBufferWindowMemory
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chat_models import ChatOpenAI
-from langchain.agents import initialize_agent
-from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder # utils/database.py
+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
 
 def create_connection(db_file):
     try:
@@ -105,6 +102,19 @@ def format_chat_history(messages: list[BaseMessage]) -> list[dict]:
             formatted.append({"role": "system", "content": msg.content})
     return formatted
 
+# utils/database.py
+from langchain.memory import ConversationBufferWindowMemory
+from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+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
+
 def initialize_qa_system(vector_store):
     """Initialize QA system with proper chat handling"""
     try:
@@ -121,48 +131,44 @@ def initialize_qa_system(vector_store):
             k=5
         )
 
-        # Create the prompt template
+        # Create the prompt template with the correct variable names
         prompt = ChatPromptTemplate.from_messages([
             ("system", "You are a helpful assistant analyzing RFP documents."),
             MessagesPlaceholder(variable_name="chat_history"),
-            ("human", "{question}"),
+            ("user", "{input}\nContext: {context}"),
             MessagesPlaceholder(variable_name="agent_scratchpad")
         ])
 
-        # Create the RAG chain with lambda function
-        rag_chain = (
+        # Create retriever function
+        retriever = vector_store.as_retriever(search_kwargs={"k": 2})
+
+        # Create the RAG pipeline
+        rag_pipe = (
             {
-                "context": lambda x: vector_store.as_retriever().get_relevant_documents(x["question"]),
-                "question": RunnablePassthrough(),
-                "chat_history": lambda x: memory.chat_memory.messages
+                "context": lambda x: retriever.get_relevant_documents(x["input"]),
+                "input": lambda x: x["input"],
+                "chat_history": lambda x: memory.chat_memory.messages,
+                "agent_scratchpad": lambda x: format_to_tool_messages(x["intermediate_steps"])
             }
-            | prompt 
-            | llm
+            | prompt
+            | llm.bind(stop=["\nHuman:"])
+            | OpenAIToolsAgentOutputParser()
         )
 
         # Create the agent executor
         agent_executor = AgentExecutor(
-            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
-            ),
+            agent=rag_pipe,
             tools=[
                 Tool(
                     name="RFP_Knowledge_Base",
-                    func=rag_chain.invoke,
+                    func=lambda x: retriever.get_relevant_documents(x),
                     description="Use this tool to analyze RFP documents and answer questions about their content."
                 )
             ],
             memory=memory,
             verbose=True,
-            handle_parsing_errors=True
+            handle_parsing_errors=True,
+            return_intermediate_steps=True
         )
 
         return agent_executor