app.py

import os
from smolagents import CodeAgent, HfApiModel, tool
import datetime
import pytz
import gradio as gr

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Tool to fetch the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'Europe/Moscow').
    """
    if not isinstance(timezone, str):
        return f"Error: Expected a string for timezone, got: {type(timezone)}"
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"Current time in {timezone}: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

@tool
def final_answer(answer: str) -> str:
    """Returns the final answer.
    Args:
        answer: The final answer as a string.
    """
    return answer

try:
    model = HfApiModel(
        max_tokens=1000,
        temperature=0.5,
        model_id='mistralai/Mixtral-8x7B-Instruct-v0.1',
    )
except Exception as e:
    print(f"Error initializing model: {str(e)}")

prompt_templates = {
    "system_prompt": (
        "You are a helpful assistant that responds in the same language as the user's query. "
        "Detect the language of the query and respond with the result only, without extra explanations. "
        "For time-related questions (e.g., 'What time is it in [place]?') or city names (e.g., 'Birsk'), use the `get_current_time_in_timezone` tool "
        "with the appropriate timezone (e.g., 'Europe/Moscow' for Moscow, 'Asia/Yekaterinburg' for Birsk). "
        "NEVER redefine tools—use only the provided ones. "
        "Always return the response via `final_answer()` in a code block like this:\n"
        "```py\n"
        "final_answer('YOUR ANSWER HERE')\n"
        "```\n"
        "Do NOT include tool calls (e.g., get_current_time_in_timezone()) inside final_answer(). "
        "Execute the tool first, then pass the result as a string to final_answer(). "
        "If the query is unclear, ask for clarification in the same language via final_answer()."
    ),
    "default": "Response: {{question}}",
    "planning": {
        "initial_plan": (
            "Analyze the query: {{question}}. Determine if a tool is needed. "
            "If the query is about time or mentions a city (e.g., 'Birsk'), use `get_current_time_in_timezone` with the correct timezone (e.g., 'Asia/Yekaterinburg' for Birsk). "
            "Write code to execute the tool and return the result via final_answer() in a code block:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
        "update_plan_pre_messages": (
            "Review the query: {{question}} and previous steps. "
            "If the plan needs adjustment (e.g., wrong timezone), update it. "
            "Write corrected code in a code block:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
        "update_plan_post_messages": (
            "Review the query: {{question}} and execution results. "
            "If the result is incorrect or incomplete, adjust the plan and code. "
            "Return the result via final_answer() in a code block:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        )
    },
    "managed_agent": {
        "execute": (
            "Execute the task: {{question}}. Use tools as specified in the plan. "
            "Write code in a code block:\n"
            "```py\n"
            "result = get_current_time_in_timezone('TIMEZONE')\n"
            "final_answer(result)\n"
            "```"
        ),
        "report": (
            "Summarize the execution of the task: {{question}}. "
            "Return the result via final_answer() in a code block:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
        "task": (
            "Define the task based on the query: {{question}}. "
            "Identify the required tools and steps. "
            "Return the task definition via final_answer() in a code block:\n"
            "```py\n"
            "final_answer('YOUR TASK DEFINITION HERE')\n"
            "```"
        )
    },
    "final_answer": {
        "pre_messages": (
            "Prepare the final response for the query: {{question}}. "
            "Ensure the answer is clear and in the same language as the query. "
            "Return the response via final_answer() in a code block:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
        "template": "Final response: {{answer}}",
        "post_messages": (
            "Review the formatted response for the query: {{question}}. "
            "Ensure the response is accurate and complete. "
            "Return the final response via final_answer() in a code block:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        )
    }
}

agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone],
    max_steps=2,
    verbosity_level=2,
    prompt_templates=prompt_templates,
)

def process_input(user_input):
    try:
        response = agent.run(user_input)
        return response
    except Exception as e:
        return f"Error: {str(e)}"

gr.Interface(
    fn=process_input,
    inputs="text",
    outputs="text",
    title="Helpful Assistant",
    description="Hello! What city or town are you from? I can tell you the time, weather, air quality, and more!"
).launch()