Update GRADIO_UI.py

GRADIO_UI.py

@@ -132,14 +132,13 @@ def pull_messages_from_step(step_log: MemoryStep):
                         metadata={"title": "📝 Scraped Text"}
                     )
 
-def stream_to_gradio(agent, task: str, api_key: str = None, reset_agent_memory: bool = False, additional_args: Optional[dict] = None):
+def stream_to_gradio(initialize_agent, task: str, api_key: str = None, reset_agent_memory: bool = False, additional_args: Optional[dict] = None):
     if not _is_package_available("gradio"):
         raise ModuleNotFoundError("Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`")
     import gradio as gr
 
-    if api_key:
-        os.environ["GOOGLE_API_KEY"] = api_key
-        logger.debug("Using user-provided GOOGLE_API_KEY")
+    logger.debug(f"Received api_key: {'****' if api_key else 'None'}")
+    agent = initialize_agent(api_key)
 
     total_input_tokens = 0
     total_output_tokens = 0
@@ -172,16 +171,16 @@ def stream_to_gradio(agent, task: str, api_key: str = None, reset_agent_memory:
         yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
 
 class GradioUI:
-    def __init__(self, agent: MultiStepAgent):
+    def __init__(self, initialize_agent):
         if not _is_package_available("gradio"):
             raise ModuleNotFoundError("Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`")
-        self.agent = agent
+        self.initialize_agent = initialize_agent
 
     def interact_with_agent(self, prompt, api_key, messages):
         import gradio as gr
         messages.append(gr.ChatMessage(role="user", content=prompt))
         yield messages
-        for msg in stream_to_gradio(self.agent, task=prompt, api_key=api_key, reset_agent_memory=False):
+        for msg in stream_to_gradio(self.initialize_agent, task=prompt, api_key=api_key, reset_agent_memory=False):
             messages.append(msg)
             yield messages
         yield messages
@@ -207,7 +206,6 @@ class GradioUI:
         """
 
         with gr.Blocks(fill_height=True, css=css) as demo:
-            stored_messages = gr.State([])
             gr.Markdown("**Note**: Please provide your own Gemini API key below. The default key may run out of quota.")
             api_key_input = gr.Textbox(
                 lines=1, label="Gemini API Key (optional)", placeholder="Enter your Gemini API key here", type="password"
@@ -217,7 +215,7 @@ class GradioUI:
                 avatar_images=(None, "./icon.png"), scale=1, height=600
             )
             text_input = gr.Textbox(
-                lines=1, label="Enter URL and request (e.g., navigate to https://en.wikipedia.org/wiki/Nvidia, scrape the infobox table, fill a search form)"
+                lines=1, label="Enter URL and request (e.g., navigate to https://en.wikipedia.org/wiki/Nvidia, scrape the infobox table)"
             )
             text_input.submit(self.interact_with_agent, [text_input, api_key_input, stored_messages], [chatbot])