fixing agent session state

app.py

@@ -131,14 +131,17 @@ WEB_TOOLS = [
     TextInspectorTool(model, text_limit),
 ]
 
-agent = CodeAgent(
-    model=model,
-    tools=[visualizer] + WEB_TOOLS,
-    max_steps=5,
-    verbosity_level=2,
-    additional_authorized_imports=AUTHORIZED_IMPORTS,
-    planning_interval=4,
-)
+# Agent creation in a factory function
+def create_agent():
+    """Creates a fresh agent instance for each session"""
+    return CodeAgent(
+        model=model,
+        tools=[visualizer] + WEB_TOOLS,
+        max_steps=5,
+        verbosity_level=2,
+        additional_authorized_imports=AUTHORIZED_IMPORTS,
+        planning_interval=4,
+    )
 
 document_inspection_tool = TextInspectorTool(model, 20000)
 
@@ -198,17 +201,23 @@ def stream_to_gradio(
 class GradioUI:
     """A one-line interface to launch your agent in Gradio"""
 
-    def __init__(self, agent, file_upload_folder: str | None = None):
-        self.agent = agent
+    def __init__(self, file_upload_folder: str | None = None):
+        
         self.file_upload_folder = file_upload_folder
         if self.file_upload_folder is not None:
             if not os.path.exists(file_upload_folder):
                 os.mkdir(file_upload_folder)
 
-    def interact_with_agent(self, prompt, messages):
+    def interact_with_agent(self, prompt, messages, session_state):
+        # Get or create session-specific agent
+        if 'agent' not in session_state:
+            session_state['agent'] = create_agent()
+            
         messages.append(gr.ChatMessage(role="user", content=prompt))
         yield messages
-        for msg in stream_to_gradio(self.agent, task=prompt, reset_agent_memory=False):
+
+        # Use session's agent instance
+        for msg in stream_to_gradio(session_state['agent'], task=prompt, reset_agent_memory=False):
             messages.append(msg)
             yield messages
         yield messages
@@ -279,6 +288,8 @@ OpenAI just published [Deep Research](https://openai.com/index/introducing-deep-
 However, their agent has a huge downside: it's not open. So we've started a 24-hour rush to replicate and open-source it. Our resulting [open-Deep-Research agent](https://github.com/huggingface/smolagents/tree/main/examples/open_deep_research) took the #1 rank of any open submission on the GAIA leaderboard! ✨
 
 You can try a simplified version below. 👇""")
+            # Add session state to store session-specific data
+            session_state = gr.State({})  # Initialize empty state for each session
             stored_messages = gr.State([])
             file_uploads_log = gr.State([])
             chatbot = gr.Chatbot(
@@ -305,8 +316,12 @@ You can try a simplified version below. 👇""")
                 self.log_user_message,
                 [text_input, file_uploads_log],
                 [stored_messages, text_input],
-            ).then(self.interact_with_agent, [stored_messages, chatbot], [chatbot])
+            ).then(self.interact_with_agent,
+                # Include session_state in function calls
+                [stored_messages, chatbot, session_state],
+                [chatbot]
+            )
 
         demo.launch(debug=True, share=True, **kwargs)
 
-GradioUI(agent).launch()
+GradioUI().launch()