Creating SQL agent with memory.

app.py

@@ -13,6 +13,7 @@ try:
     from langchain_community.utilities import SQLDatabase
     from langchain_google_genai import ChatGoogleGenerativeAI
     from langchain.agents.agent_types import AgentType
+    from langchain.memory import ConversationBufferWindowMemory
     from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
     import pymysql
     from dotenv import load_dotenv
@@ -151,7 +152,7 @@ def create_agent():
     logger.info("Starting agent creation process...")
     
 def create_agent(llm, db_connection):
-    """Create and return a SQL database agent."""
+    """Create and return a SQL database agent with conversation memory."""
     if not llm:
         error_msg = "Cannot create agent: LLM is not available"
         logger.error(error_msg)
@@ -163,7 +164,15 @@ def create_agent(llm, db_connection):
         return None, error_msg
         
     try:
-        logger.info("Creating SQL agent...")
+        logger.info("Creating SQL agent with memory...")
+        
+        # Create conversation memory
+        memory = ConversationBufferWindowMemory(
+            memory_key="chat_history",
+            k=5,  # Keep last 5 message exchanges in memory
+            return_messages=True,
+            output_key="output"
+        )
         
         # Create the database toolkit with additional configuration
         toolkit = SQLDatabaseToolkit(
@@ -171,7 +180,7 @@ def create_agent(llm, db_connection):
             llm=llm
         )
         
-        # Create the agent with more detailed configuration
+        # Create the agent with memory and more detailed configuration
         agent = create_sql_agent(
             llm=llm,
             toolkit=toolkit,
@@ -179,7 +188,9 @@ def create_agent(llm, db_connection):
             verbose=True,
             handle_parsing_errors=True,  # Better error handling for parsing
             max_iterations=10,  # Limit the number of iterations
-            early_stopping_method="generate"  # Stop early if the agent is stuck
+            early_stopping_method="generate",  # Stop early if the agent is stuck
+            memory=memory,  # Add memory to the agent
+            return_intermediate_steps=True  # Important for memory to work properly
         )
         
         # Test the agent with a simple query
@@ -328,17 +339,22 @@ def convert_to_messages_format(chat_history):
     return messages
 
 async def stream_agent_response(question: str, chat_history: List) -> List[Dict]:
-    """Procesa la pregunta del usuario y devuelve la respuesta del agente."""
+    """Procesa la pregunta del usuario y devuelve la respuesta del agente con memoria de conversación."""
     # Initialize response
     response_text = ""
     messages = []
     
-    # Add previous chat history
+    # Add previous chat history in the correct format for the agent
     if chat_history:
-        messages.extend(chat_history)
+        # Convert chat history to the format expected by the agent's memory
+        for msg in chat_history:
+            if msg["role"] == "user":
+                messages.append(HumanMessage(content=msg["content"]))
+            elif msg["role"] == "assistant":
+                messages.append(AIMessage(content=msg["content"]))
     
-    # Add user's question
-    user_message = {"role": "user", "content": question}
+    # Add current user's question
+    user_message = HumanMessage(content=question)
     messages.append(user_message)
     
     if not agent:
@@ -351,7 +367,6 @@ async def stream_agent_response(question: str, chat_history: List) -> List[Dict]
             f"Error: {agent_error}"
         )
         assistant_message = {"role": "assistant", "content": error_msg}
-        messages.append(assistant_message)
         return [assistant_message]
     
     try: