app.py

# app.py
import gradio as gr
import os

# (1) Import everything you need from smolagents:
#     – CodeAgent, the various Tool classes, and InferenceClientModel
#     – If your agent definition lives in a file like simple_agent.py in this directory,
#       you can import it directly. For example, if you have a `SimpleAgentBuilder`
#       function or class that returns a ready‐to‐run CodeAgent, import it here.
#
# Option A: If you defined your agent in simple_agent.py, you could do:
#     from simple_agent import build_agent
#
#     def build_agent() -> CodeAgent:
#         # your code that sets up CodeAgent(...) with tools=[…], etc.
#         return agent
#
# Option B: Re‐copy the agent definition directly into app.py (less DRY, but self‐contained).
#
# Below is an example of “Option B,” which simply re‐defines the same tools + agent
# that you used in simple_agent.py. Adjust names/imports if yours differ.

from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, FinalAnswerTool, InferenceClientModel, Tool

# ---------------  (Re-define or import your custom tools)  ---------------

class SuperheroPartyThemeTool(Tool):
    name = "superhero_party_theme_generator"
    description = """
    Suggests creative superhero (or villain) themed party ideas based on a category.
    Returns a string with a unique party theme idea.
    """
    inputs = {
        "category": {
            "type": "string",
            "description": "The type of superhero party (e.g. 'classic heroes', 'villain masquerade', 'futuristic Gotham')."
        }
    }
    output_type = "string"

    def forward(self, category: str) -> str:
        themes = {
            "classic heroes": "Justice League Gala: Guests come dressed as their favorite DC heroes with themed cocktails like 'The Kryptonite Punch'.",
            "villain masquerade": "Gotham Rogues' Ball: A mysterious masquerade where guests dress as classic Batman villains.",
            "futuristic Gotham": "Neo-Gotham Night: A cyberpunk-style party inspired by Batman Beyond, with neon decorations and futuristic gadgets."
        }
        return themes.get(category.lower(), "Themed party idea not found. Try 'classic heroes', 'villain masquerade', or 'futuristic Gotham'.")

class SuggestMenuTool(Tool):
    name = "suggest_menu"
    description = "Suggests a menu based on the type of occasion."
    inputs = {
        "occasion": {
            "type": "string",
            "description": "The type of occasion for the party (e.g. 'casual', 'formal', 'superhero', etc.)."
        }
    }
    output_type = "string"

    def forward(self, occasion: str) -> str:
        if occasion.lower() == "casual":
            return "Pizza, snacks, and drinks."
        elif occasion.lower() == "formal":
            return "3-course dinner with wine and dessert."
        elif occasion.lower() == "superhero":
            return "Buffet with high-energy and healthy food."
        else:
            return "Custom menu for the butler."

class CateringServiceTool(Tool):
    name = "catering_service_tool"
    description = "Returns the highest-rated catering service in Gotham City."
    inputs = {
        "query": {
            "type": "string",
            "description": "A search term for finding catering services."
        }
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        services = {
            "Gotham Catering Co.": 4.9,
            "Wayne Manor Catering": 4.8,
            "Gotham City Events": 4.7,
        }
        best_service = max(services, key=services.get)
        return best_service

# ---------------  (Instantiate your CodeAgent exactly as in simple_agent.py)  ---------------
def build_agent() -> CodeAgent:
    """
    Construct and return the same CodeAgent you originally pushed to the Hub.
    Make sure ‘tools=[ … ]’ and the model (InferenceClientModel) lines match exactly.
    """
    agent = CodeAgent(
        tools=[
            DuckDuckGoSearchTool(),
            VisitWebpageTool(),
            SuggestMenuTool(),
            CateringServiceTool(),
            SuperheroPartyThemeTool(),
            FinalAnswerTool()
        ],
        model=InferenceClientModel(),
        max_steps=10,
        verbosity_level=2
    )
    return agent

# ---------------  (Wrap the agent in a Gradio interface)  ---------------
agent_instance = build_agent()

def run_agent(prompt: str) -> str:
    try:
        return agent_instance.run(prompt)
    except Exception as e:
        import traceback
        return f"⚠️ An error occurred:\n\n{traceback.format_exc()}"


iface = gr.Interface(
    fn=run_agent,
    inputs=gr.Textbox(
        lines=2,
        placeholder="Type something like:\n" 
                    "\"Give me the best playlist for a 'villain masquerade' party at Wayne Manor\""
    ),
    outputs=gr.Textbox(label="Agent Response"),
    title="Simple SmolAgents Agent",
    description=(
        "Enter a natural‐language prompt for the SmolAgents CodeAgent. "
        "Under the hood, this is the same agent you pushed to the Hub!"
    ),
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()