Update app.py

app.py

@@ -1,18 +1,33 @@
 import os
-from typing import Optional
+from typing import Optional, List, Dict, Any
 import gradio as gr
 import requests
 from smolagents import CodeAgent, Tool
 from smolagents.models import HfApiModel
 from smolagents.monitoring import LogLevel
 from gradio import ChatMessage
+import logging
 from functools import lru_cache
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 DEFAULT_MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"
 HF_API_TOKEN = os.getenv("HF_TOKEN")
 
+# Tool descriptions for the UI
+TOOL_DESCRIPTIONS = {
+    "Hub Collections": "Add tool collections from Hugging Face Hub.",
+    "Spaces": "Add tools from Hugging Face Spaces.",
+}
+
 @lru_cache(maxsize=128)
-def search_spaces(query, limit=1):
+def search_spaces(query: str, limit: int = 1) -> Optional[Dict[str, str]]:
+    """
+    Search for Hugging Face Spaces using the API.
+    Returns the first result or None if no results.
+    """
     try:
         url = f"https://huggingface.co/api/spaces?search={query}&limit={limit}"
         response = requests.get(url, headers={"Authorization": f"Bearer {HF_API_TOKEN}"})
@@ -21,10 +36,14 @@ def search_spaces(query, limit=1):
         if not spaces:
             return None
         return extract_space_info(spaces[0])
-    except requests.RequestException:
+    except requests.RequestException as e:
+        logger.error(f"Error searching spaces: {e}")
         return None
 
-def extract_space_info(space):
+def extract_space_info(space: Dict[str, Any]) -> Dict[str, str]:
+    """
+    Extracts space information from the API response.
+    """
     space_id = space["id"]
     title = space_id.split("/")[-1]
     description = f"Tool from {space_id}"
@@ -38,27 +57,33 @@ def extract_space_info(space):
         description = space["cardData"]["description"]
     return {"id": space_id, "title": title, "description": description}
 
-def get_space_metadata(space_id):
+def get_space_metadata(space_id: str) -> Optional[Dict[str, str]]:
+    """
+    Get metadata for a specific Hugging Face Space.
+    """
     try:
         url = f"https://huggingface.co/api/spaces/{space_id}"
         response = requests.get(url, headers={"Authorization": f"Bearer {HF_API_TOKEN}"})
         response.raise_for_status()
         space = response.json()
         return extract_space_info(space)
-    except requests.RequestException:
+    except requests.RequestException as e:
+        logger.error(f"Error getting space metadata: {e}")
         return None
 
-def create_agent(model_name, space_tools=None):
+def create_agent(model_name: str, space_tools: Optional[List[Dict[str, str]]] = None) -> Optional[CodeAgent]:
+    """
+    Create a CodeAgent with the specified model and tools.
+    """
     if not space_tools:
         space_tools = []
     try:
         tools = [
             Tool.from_space(
-                t["id"],
-                name=t.get("name", t["id"]),
-                description=t.get("description", ""),
-            )
-            for t in space_tools
+                tool_info["id"],
+                name=tool_info.get("name", tool_info["id"]),
+                description=tool_info.get("description", ""),
+            ) for tool_info in space_tools
         ]
         model = HfApiModel(model_id=model_name, token=HF_API_TOKEN)
         agent = CodeAgent(
@@ -67,39 +92,46 @@ def create_agent(model_name, space_tools=None):
             additional_authorized_imports=["PIL", "requests"],
             verbosity_level=LogLevel.DEBUG,
         )
+        logger.info(f"Agent created successfully with {len(tools)} tools")
         return agent
-    except:
-        try:
-            fallback_model = HfApiModel(
-                model_id="Qwen/Qwen2.5-Coder-7B-Instruct", token=HF_API_TOKEN
-            )
-            agent = CodeAgent(
-                tools=tools,
-                model=fallback_model,
-                additional_authorized_imports=["PIL", "requests"],
-                verbosity_level=LogLevel.DEBUG,
-            )
-            return agent
-        except:
-            return None
+    except Exception as e:
+        logger.error(f"Error creating agent: {e}")
+        return create_fallback_agent(tools)
 
-def on_search_spaces(query):
+def create_fallback_agent(tools: List[Tool]) -> Optional[CodeAgent]:
+    """
+    Create a fallback CodeAgent if the primary model fails.
+    """
+    try:
+        logger.info("Trying fallback model...")
+        fallback_model = HfApiModel(model_id="Qwen/Qwen2.5-Coder-7B-Instruct", token=HF_API_TOKEN)
+        agent = CodeAgent(
+            tools=tools,
+            model=fallback_model,
+            additional_authorized_imports=["PIL", "requests"],
+            verbosity_level=LogLevel.DEBUG,
+        )
+        logger.info("Agent created successfully with fallback model")
+        return agent
+    except Exception as e:
+        logger.error(f"Error creating agent with fallback model: {e}")
+        return None
+
+# Event handler functions
+def on_search_spaces(query: str) -> tuple:
     if not query:
         return "Please enter a search term.", "", "", ""
     try:
         space_info = search_spaces(query)
         if space_info is None:
             return "No spaces found.", "", "", ""
-        results_md = (
-            f"### Search Results:\n- ID: `{space_info['id']}`\n"
-            f"- Title: {space_info['title']}\n"
-            f"- Description: {space_info['description']}\n"
-        )
+        results_md = f"### Search Results:\n- ID: `{space_info['id']}`\n- Title: {space_info['title']}\n- Description: {space_info['description']}\n"
         return results_md, space_info["id"], space_info["title"], space_info["description"]
     except Exception as e:
+        logger.error(f"Error in search: {e}")
         return f"Error: {str(e)}", "", "", ""
 
-def on_validate_space(space_id):
+def on_validate_space(space_id: str) -> tuple:
     if not space_id:
         return "Please enter a space ID or search term.", "", ""
     try:
@@ -108,20 +140,15 @@ def on_validate_space(space_id):
             space_info = search_spaces(space_id)
             if space_info is None:
                 return f"No spaces found for '{space_id}'.", "", ""
-            result_md = (
-                f"### Found Space via Search:\n- ID: `{space_info['id']}`\n"
-                f"- Title: {space_info['title']}\n- Description: {space_info['description']}\n"
-            )
+            result_md = f"### Found Space via Search:\n- ID: `{space_info['id']}`\n- Title: {space_info['title']}\n- Description: {space_info['description']}\n"
             return result_md, space_info["title"], space_info["description"]
-        result_md = (
-            f"### Space Validated Successfully:\n- ID: `{space_info['id']}`\n"
-            f"- Title: {space_info['title']}\n- Description: {space_info['description']}\n"
-        )
+        result_md = f"### Space Validated Successfully:\n- ID: `{space_info['id']}`\n- Title: {space_info['title']}\n- Description: {space_info['description']}\n"
         return result_md, space_info["title"], space_info["description"]
     except Exception as e:
+        logger.error(f"Error validating space: {e}")
         return f"Error: {str(e)}", "", ""
 
-def on_add_tool(space_id, space_name, space_description, current_tools):
+def on_add_tool(space_id: str, space_name: str, space_description: str, current_tools: List[Dict[str, str]]) -> tuple:
     if not space_id:
         return current_tools, "Please enter a space ID."
     for tool in current_tools:
@@ -135,31 +162,30 @@ def on_add_tool(space_id, space_name, space_description, current_tools):
     updated_tools = current_tools + [new_tool]
     tools_md = "### Added Tools:\n"
     for i, tool in enumerate(updated_tools, 1):
-        tools_md += (
-            f"{i}. **{tool['name']}** (`{tool['id']}`)\n   {tool['description']}\n\n"
-        )
+        tools_md += f"{i}. **{tool['name']}** (`{tool['id']}`)\n   {tool['description']}\n\n"
     return updated_tools, tools_md
 
-def on_create_agent(model, space_tools):
+def on_create_agent(model: str, space_tools: List[Dict[str, str]]) -> tuple:
     if not space_tools:
         return None, [], "", "Please add at least one tool before creating an agent.", "No agent created yet."
     try:
         agent = create_agent(model, space_tools)
         if agent is None:
             return None, [], "", "Failed to create agent. Please try again with different tools or model.", "No agent created yet."
-        tools_str = ", ".join([f"{t['name']} ({t['id']})" for t in space_tools])
+        tools_str = ", ".join([f"{tool['name']} ({tool['id']})" for tool in space_tools])
         agent_status = update_agent_status(agent)
         return agent, [], "", f"✅ Agent created successfully with {model}!\nTools: {tools_str}", agent_status
     except Exception as e:
+        logger.error(f"Error creating agent: {e}")
         return None, [], "", f"Error creating agent: {str(e)}", "No agent created yet."
 
-def add_user_message(message, chat_history):
+def add_user_message(message: str, chat_history: List[ChatMessage]) -> tuple:
     if not message:
         return "", chat_history
     chat_history = chat_history + [ChatMessage(role="user", content=message)]
     return message, chat_history
 
-def stream_to_gradio(agent, task: str, reset_agent_memory: bool = False, additional_args: Optional[dict] = None):
+def stream_to_gradio(agent: CodeAgent, task: str, reset_agent_memory: bool = False, additional_args: Optional[dict] = None):
     from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
     from smolagents.agent_types import AgentAudio, AgentImage, AgentText
     for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
@@ -176,7 +202,7 @@ def stream_to_gradio(agent, task: str, reset_agent_memory: bool = False, additio
     else:
         yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
 
-def stream_agent_response(agent, message, chat_history):
+def stream_agent_response(agent: CodeAgent, message: str, chat_history: List[ChatMessage]):
     if not message or agent is None:
         return chat_history
     yield chat_history
@@ -189,15 +215,18 @@ def stream_agent_response(agent, message, chat_history):
         chat_history = chat_history + [ChatMessage(role="assistant", content=error_msg)]
         yield chat_history
 
-def on_clear(agent=None):
+def on_clear(agent: Optional[CodeAgent] = None) -> tuple:
     return agent, [], "", "Agent cleared. Create a new one to continue.", "", gr.update(interactive=False)
 
-def update_agent_status(agent):
+def update_agent_status(agent: Optional[CodeAgent]) -> str:
     if agent is None:
         return "No agent created yet. Add a Space tool to get started."
     tools = agent.tools if hasattr(agent, "tools") else []
-    return f"Agent ready with {len(tools)} tools"
+    tool_count = len(tools)
+    status = f"Agent ready with {tool_count} tools"
+    return status
 
+# Create the Gradio app
 with gr.Blocks(title="AI Agent Builder") as app:
     gr.Markdown("# AI Agent Builder with smolagents")
     gr.Markdown("Build your own AI agent by selecting tools from Hugging Face Spaces.")