memory per session

app.py

@@ -229,8 +229,16 @@ def create_agent():
     
     logger.info("Starting agent creation process...")
     
-def create_agent(llm, db_connection):
-    """Create and return a SQL database agent with conversation memory."""
+def create_agent(llm, db_connection, *, run_test: bool = True):
+    """Create and return a SQL database agent with conversation memory.
+
+    Args:
+        llm: Language model instance
+        db_connection: SQLDatabase connection
+        run_test: If True, executes a small test query which may add entries
+                  to agent memory. Disable for per-session agents to start
+                  with a clean memory.
+    """
     if not llm:
         error_msg = "Cannot create agent: LLM is not available"
         logger.error(error_msg)
@@ -244,10 +252,14 @@ def create_agent(llm, db_connection):
     try:
         logger.info("Creating SQL agent with memory...")
         
-        # Create conversation memory
+        # Create conversation memory (configurable window via env MEMORY_K)
+        try:
+            memory_k = int(os.getenv("MEMORY_K", "8"))
+        except Exception:
+            memory_k = 8
         memory = ConversationBufferWindowMemory(
             memory_key="chat_history",
-            k=5,  # Keep last 5 message exchanges in memory
+            k=memory_k,
             return_messages=True,
             output_key="output"
         )
@@ -271,15 +283,16 @@ def create_agent(llm, db_connection):
             return_intermediate_steps=True  # Important for memory to work properly
         )
         
-        # Test the agent with a simple query
-        logger.info("Testing agent with a simple query...")
-        try:
-            test_query = "SELECT 1"
-            test_result = agent.run(test_query)
-            logger.info(f"Agent test query successful: {str(test_result)[:200]}...")
-        except Exception as e:
-            logger.warning(f"Agent test query failed (this might be expected): {str(e)}")
-            # Continue even if test fails, as it might be due to model limitations
+        # if run_test:
+        #     # Test the agent with a simple query
+        #     logger.info("Testing agent with a simple query...")
+        #     try:
+        #         test_query = "SELECT 1"
+        #         test_result = agent.run(test_query)
+        #         logger.info(f"Agent test query successful: {str(test_result)[:200]}...")
+        #     except Exception as e:
+        #         logger.warning(f"Agent test query failed (this might be expected): {str(e)}")
+        #         # Continue even if test fails, as it might be due to model limitations
         
         logger.info("SQL agent created successfully")
         return agent, ""
@@ -497,7 +510,7 @@ def convert_to_messages_format(chat_history):
     
     return messages
 
-async def stream_agent_response(question: str, chat_history: List[List[str]]) -> Tuple[str, Optional["go.Figure"]]:
+async def stream_agent_response(question: str, chat_history: List[List[str]], session_agent=None) -> Tuple[str, Optional["go.Figure"]]:
     """Procesa la pregunta del usuario y devuelve la respuesta del agente con memoria de conversación."""
     global agent  # Make sure we can modify the agent's memory
     
@@ -517,7 +530,8 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
     user_message = HumanMessage(content=question)
     messages.append(user_message)
     
-    if not agent:
+    active_agent = session_agent if session_agent is not None else agent
+    if not active_agent:
         error_msg = (
             "## ⚠️ Error: Agente no inicializado\n\n"
             "No se pudo inicializar el agente de base de datos. Por favor, verifica que:\n"
@@ -531,11 +545,11 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
     # Update the agent's memory with the full conversation history
     try:
         # Rebuild agent memory from chat history pairs
-        if hasattr(agent, 'memory') and agent.memory is not None:
-            agent.memory.clear()
+        if hasattr(active_agent, 'memory') and active_agent.memory is not None:
+            active_agent.memory.clear()
             for i in range(0, len(messages)-1, 2):  # (user, assistant)
                 if i+1 < len(messages):
-                    agent.memory.save_context(
+                    active_agent.memory.save_context(
                         {"input": messages[i].content},
                         {"output": messages[i+1].content}
                     )
@@ -550,7 +564,7 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
         # Execute the agent with proper error handling
         try:
             # Let the agent use its memory; don't pass raw chat_history
-            response = await agent.ainvoke({"input": question})
+            response = await active_agent.ainvoke({"input": question})
             logger.info(f"Agent response type: {type(response)}")
             logger.info(f"Agent response content: {str(response)[:500]}...")
             
@@ -640,7 +654,7 @@ async def stream_agent_response(question: str, chat_history: List[List[str]]) ->
                             "Devuelve SOLO la consulta SQL en un bloque ```sql``` para responder a: "
                             f"{question}. No incluyas explicación ni texto adicional."
                         )
-                        sql_only_resp = await agent.ainvoke({"input": sql_only_prompt})
+                        sql_only_resp = await active_agent.ainvoke({"input": sql_only_prompt})
                         sql_only_text = str(sql_only_resp)
                         sql_query2 = extract_sql_query(sql_only_text)
                         if sql_query2 and looks_like_sql(sql_query2):
@@ -866,6 +880,9 @@ def create_ui():
         if not env_ok:
             gr.Warning("⚠️ " + env_message)
         
+        # Create session-scoped state for the agent with memory
+        session_agent_state = gr.State(value=None)
+
         # Create the chat interface
         with gr.Row():
             chatbot = gr.Chatbot(
@@ -975,12 +992,28 @@ def create_ui():
         # Hidden component for streaming output
         streaming_output_display = gr.Textbox(visible=False)
         
-        return demo, chatbot, chart_display, question_input, submit_button, streaming_output_display
+        return (
+            demo,
+            chatbot,
+            chart_display,
+            question_input,
+            submit_button,
+            streaming_output_display,
+            session_agent_state,
+        )
 
 def create_application():
     """Create and configure the Gradio application."""
     # Create the UI components
-    demo, chatbot, chart_display, question_input, submit_button, streaming_output_display = create_ui()
+    (
+        demo,
+        chatbot,
+        chart_display,
+        question_input,
+        submit_button,
+        streaming_output_display,
+        session_agent_state,
+    ) = create_ui()
     
     def user_message(user_input: str, chat_history: List[Dict[str, str]]) -> Tuple[str, List[Dict[str, str]]]:
         """Add user message to chat history (messages format) and clear input."""
@@ -997,20 +1030,36 @@ def create_application():
 
         return "", chat_history
     
-    async def bot_response(chat_history: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Optional[go.Figure]]:
-        """Generate bot response for messages-format chat history and return optional chart figure."""
+    async def bot_response(
+        chat_history: List[Dict[str, str]],
+        session_agent,
+    ) -> Tuple[List[Dict[str, str]], Optional[go.Figure], Any]:
+        """Generate bot response using a session-scoped agent with memory."""
         if not chat_history:
-            return chat_history, None
+            return chat_history, None, session_agent
 
         # Ensure last message is a user turn awaiting assistant reply
         last = chat_history[-1]
         if not isinstance(last, dict) or last.get("role") != "user" or not last.get("content"):
-            return chat_history, None
+            return chat_history, None, session_agent
 
         try:
             question = last["content"]
             logger.info(f"Processing question: {question}")
 
+            # Ensure we have a session-specific agent with memory
+            local_agent = session_agent
+            if local_agent is None:
+                try:
+                    local_agent, _ = create_agent(llm, db_connection, run_test=False)
+                    session_agent = local_agent
+                    logger.info("Created new session agent with memory")
+                except Exception as e:
+                    logger.error(f"Could not create session agent: {e}")
+                    # Fallback to global agent if available
+                    session_agent = None
+                    local_agent = None
+
             # Convert prior messages to pair history for stream_agent_response()
             pair_history: List[List[str]] = []
             i = 0
@@ -1029,7 +1078,11 @@ def create_application():
                     i += 1
 
             # Call the agent for this new user question
-            assistant_message, chart_fig = await stream_agent_response(question, pair_history)
+            assistant_message, chart_fig = await stream_agent_response(
+                question,
+                pair_history,
+                session_agent=session_agent,
+            )
 
             # Append assistant message back into messages history
             chat_history.append({"role": "assistant", "content": assistant_message})
@@ -1066,14 +1119,14 @@ def create_application():
                         )
 
             logger.info("Response generation complete")
-            return chat_history, chart_fig
+            return chat_history, chart_fig, session_agent
 
         except Exception as e:
             error_msg = f"## ❌ Error\n\nError al procesar la solicitud:\n\n```\n{str(e)}\n```"
             logger.error(error_msg, exc_info=True)
             # Ensure we add an assistant error message for the UI
             chat_history.append({"role": "assistant", "content": error_msg})
-            return chat_history, None
+            return chat_history, None, session_agent
     
     # Event handlers
     with demo:
@@ -1085,8 +1138,8 @@ def create_application():
             queue=True
         ).then(
             fn=bot_response,
-            inputs=[chatbot],
-            outputs=[chatbot, chart_display],
+            inputs=[chatbot, session_agent_state],
+            outputs=[chatbot, chart_display, session_agent_state],
             api_name="ask"
         )
         
@@ -1098,8 +1151,8 @@ def create_application():
             queue=True
         ).then(
             fn=bot_response,
-            inputs=[chatbot],
-            outputs=[chatbot, chart_display]
+            inputs=[chatbot, session_agent_state],
+            outputs=[chatbot, chart_display, session_agent_state]
         )
     
     return demo