Update app.py

app.py

@@ -1,36 +1,31 @@
-from smolagents import CodeAgent, OpenAIServerModel, load_tool, tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, OpenAIServerModel, load_tool, tool
 import datetime
 import pytz
 import yaml
 from tools.final_answer import FinalAnswerTool
 import gradio as gr
 
-# ----- Tools -----
+#  Tool to create a personalized greeting repeated a specific number of times.
 @tool
 def my_custom_tool(arg1: str, arg2: int) -> str:
     """
-    Creates a personalized greeting repeated multiple times.
-
+    A tool that creates a personalized greeting repeated multiple times.
+    
     Args:
-        arg1 (str): The name or phrase to greet.
-        arg2 (int): The number of times to repeat the greeting.
-
-    Returns:
-        str: The repeated greeting.
+        arg1: the name or phrase to greet
+        arg2: the number of times to repeat the greeting
     """
     greeting = f"Hello, {arg1}!"
     return " ".join([greeting] * arg2)
 
+# Tool to get the current local time in your chosen timezone.
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """
-    Fetches the current local time in a specified timezone.
-
+    A tool that fetches the current local time in a specified timezone.
+    
     Args:
-        timezone (str): A valid timezone string (e.g., 'America/New_York').
-
-    Returns:
-        str: The current time in that timezone or an error message.
+        timezone: A string representing a valid timezone (e.g., 'America/New_York').
     """
     try:
         tz = pytz.timezone(timezone)
@@ -41,24 +36,14 @@ def get_current_time_in_timezone(timezone: str) -> str:
 
 final_answer = FinalAnswerTool()
 
-# Load prompt templates
+# Import tool from Hub
+image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
+
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
 
-# Load additional tools from Hugging Face Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
-
-# ----- Agent Initialization -----
+# ----- Initialize agent with API key -----
 def initialize_agent(api_key: str):
-    """
-    Initializes the CodeAgent with the user-provided OpenAI API key.
-
-    Args:
-        api_key (str): OpenAI API key provided by the user.
-
-    Returns:
-        CodeAgent: The initialized agent ready to use.
-    """
     model = OpenAIServerModel(
         max_tokens=2096,
         temperature=0.5,
@@ -78,34 +63,41 @@ def initialize_agent(api_key: str):
         description=None,
         prompt_templates=prompt_templates
     )
-
     return agent
 
-# ----- Gradio Query Function -----
-def agent_response(api_key: str, user_query: str):
-    """
-    Responds to a user query using the CodeAgent initialized with the API key.
-
-    Args:
-        api_key (str): OpenAI API key provided by the user.
-        user_query (str): The user's query for the agent.
-
-    Returns:
-        str: Agent's response.
-    """
-    agent = initialize_agent(api_key)
-    return agent.run(user_query)  # Run the agent on the user's query
-
-# ----- Gradio Interface -----
-iface = gr.Interface(
-    fn=agent_response,
-    inputs=[
-        gr.Textbox(label="Enter your OpenAI API Key", type="password"),
-        gr.Textbox(label="Your Query", placeholder="Type your question here...")
-    ],
-    outputs=gr.Textbox(label="Agent Response"),
-    title="Custom SmolAgent App",
-    description="Enter your OpenAI API key and ask a question to the agent."
-)
-
-iface.launch()
+# ----- Gradio functions -----
+# Store agent globally after API key input
+AGENT = None
+
+def set_api_key(api_key: str):
+    global AGENT
+    AGENT = initialize_agent(api_key)
+    return "Agent initialized! You can now enter your queries below."
+
+def ask_agent(user_query: str):
+    global AGENT
+    if AGENT is None:
+        return "Please enter your OpenAI API key first."
+    response = AGENT.run(user_query)
+    return response
+
+# ----- Gradio interface -----
+with gr.Blocks() as demo:
+    gr.Markdown("## Enter your OpenAI API key and start querying the agent")
+    
+    with gr.Row():
+        api_key_input = gr.Textbox(label="OpenAI API Key", placeholder="sk-...", type="password")
+        api_key_button = gr.Button("Initialize Agent")
+    
+    init_output = gr.Textbox(label="Status", interactive=False)
+
+    with gr.Row():
+        user_query_input = gr.Textbox(label="Enter your query here", placeholder="Ask something...")
+        submit_query_button = gr.Button("Submit Query")
+    
+    query_output = gr.Textbox(label="Agent Response", interactive=False)
+
+    api_key_button.click(set_api_key, inputs=api_key_input, outputs=init_output)
+    submit_query_button.click(ask_agent, inputs=user_query_input, outputs=query_output)
+
+demo.launch()