from smolagents import CodeAgent, InferenceClientModel, tool
import gradio as gr

# ----------------------
# TOOL
# ----------------------
@tool
def my_custom_tool(arg1: str) -> str:
    """
    A simple tool that echoes user input.

    Args:
        arg1 (str): The text input provided by the user.

    Returns:
        str: The same text prefixed with 'Echo:'.
    """
    return f"Echo: {arg1}"

# ----------------------
# MODEL (FIXED PROPERLY)
# ----------------------
model = InferenceClientModel(
    model_id="HuggingFaceH4/zephyr-7b-beta",
    temperature=0.2
)

# ----------------------
# AGENT
# ----------------------
agent = CodeAgent(
    model=model,
    tools=[my_custom_tool]
)

# ----------------------
# FUNCTION
# ----------------------
def run_agent(user_input):
    try:
        return str(agent.run(user_input))
    except Exception as e:
        return f"Error: {str(e)}"

# ----------------------
# UI
# ----------------------
demo = gr.Interface(
    fn=run_agent,
    inputs=gr.Textbox(lines=5, placeholder="Ask something..."),
    outputs=gr.Textbox(),
    title="SmolAgents Demo",
    description="Working agent with tool"
)

# ----------------------
# LAUNCH
# ----------------------
if __name__ == "__main__":
    demo.launch()