app.py

import os
from smolagents import CodeAgent, HfApiModel, tool
import datetime
import pytz
import gradio as gr
from geopy.geocoders import Nominatim
from timezonefinder import TimezoneFinder
import requests
from langdetect import detect

@tool
def get_timezone_by_city(city: str) -> str:
    """Fetches the timezone for a given city using geolocation.
    
    Args:
        city: A string representing the city name, optionally with region or country (e.g., 'Бураево', 'Sydney').
    
    Returns:
        A string with the timezone (e.g., 'Europe/Moscow') or an error message.
    """
    if not isinstance(city, str):
        return f"Error: Expected a string for city, got: {type(city)}"
    try:
        geolocator = Nominatim(user_agent="smolagents_bot")
        location = geolocator.geocode(city, language="en")
        if not location:
            city_name = city.split(",")[0].strip().title()
            location = geolocator.geocode(city_name, language="en")
            if not location:
                return f"Error: City '{city}' not found. Try a different spelling or provide more context."
        tf = TimezoneFinder()
        timezone = tf.timezone_at(lat=location.latitude, lng=location.longitude)
        if not timezone:
            return f"Error: Timezone not found for city '{city}'."
        return timezone
    except Exception as e:
        return f"Error fetching timezone for city '{city}': {str(e)}"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Fetches the current local time in a specified timezone.
    
    Args:
        timezone: A string representing a valid timezone (e.g., 'Europe/Moscow').
    
    Returns:
        A string with the current time (e.g., '2025-05-03 12:00:00') or an error message.
    """
    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 local_time
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"

@tool
def get_air_quality(city: str, lang: str = "en") -> str:
    """Fetches air quality data for a given city using OpenWeatherMap API.
    
    Args:
        city: A string representing the city name (e.g., 'Бирск', 'Sydney').
        lang: A string representing the language code (e.g., 'en', 'ru').
    
    Returns:
        A string with air quality index (AQI) and pollutant levels or an error message.
    """
    if not isinstance(city, str):
        return f"Error: Expected a string for city, got: {type(city)}"
    try:
        geolocator = Nominatim(user_agent="smolagents_bot")
        location = geolocator.geocode(city, language="en")
        if not location:
            city_name = city.split(",")[0].strip().title()
            location = geolocator.geocode(city_name, language="en")
            if not location:
                return f"Error: City '{city}' not found for air quality data."
        
        api_key = os.getenv("OPENWEATHERMAP_API_KEY", "93f518fda8d24cf89aee7050c26b27e9")
        url = f"http://api.openweathermap.org/data/2.5/air_pollution?lat={location.latitude}&lon={location.longitude}&appid={api_key}"
        response = requests.get(url)
        if response.status_code != 200:
            return f"Error fetching air quality for '{city}': API returned status {response.status_code}"
        
        data = response.json()
        aqi = data["list"][0]["main"]["aqi"]
        components = data["list"][0]["components"]

        translations = {
            "en": {
                "aqi_desc": {1: "Good", 2: "Fair", 3: "Moderate", 4: "Poor", 5: "Very Poor"},
                "air_quality": "Air quality",
                "pm2_5": "Fine particles (PM2.5)",
                "pm10": "Coarse particles (PM10)",
                "co": "Carbon monoxide (CO)"
            },
            "ru": {
                "aqi_desc": {1: "Хорошее", 2: "Удовлетворительное", 3: "Среднее", 4: "Плохое", 5: "Очень плохое"},
                "air_quality": "Качество воздуха",
                "pm2_5": "Мелкие частицы (PM2.5)",
                "pm10": "Крупные частицы (PM10)",
                "co": "Угарный газ (CO)"
            }
        }
        lang = lang if lang in translations else "en"
        t = translations[lang]
        aqi_emoji = "🌱" if aqi <= 2 else "😷" if aqi <= 4 else "☣️"
        
        return (
            f"😷 {t['air_quality']}: AQI {aqi} ({t['aqi_desc'].get(aqi, 'Unknown')}) {aqi_emoji}\n"
            f"{t['pm2_5']}: {components['pm2_5']} µg/m³\n"
            f"{t['pm10']}: {components['pm10']} µg/m³\n"
            f"{t['co']}: {components['co']} µg/m³"
        )
    except Exception as e:
        return f"Error fetching air quality for '{city}': {str(e)}"

@tool
def get_weather(city: str, lang: str = "en") -> str:
    """Fetches weather data for a given city using OpenWeatherMap API.
    
    Args:
        city: A string representing the city name (e.g., 'Бирск', 'Sydney').
        lang: A string representing the language code (e.g., 'en', 'ru').
    
    Returns:
        A string with weather conditions, temperature, and humidity or an error message.
    """
    if not isinstance(city, str):
        return f"Error: Expected a string for city, got: {type(city)}"
    try:
        geolocator = Nominatim(user_agent="smolagents_bot")
        location = geolocator.geocode(city, language="en")
        if not location:
            city_name = city.split(",")[0].strip().title()
            location = geolocator.geocode(city_name, language="en")
            if not location:
                return f"Error: City '{city}' not found for weather data."
        
        api_key = os.getenv("OPENWEATHERMAP_API_KEY", "93f518fda8d24cf89aee7050c26b27e9")
        url = f"http://api.openweathermap.org/data/2.5/weather?lat={location.latitude}&lon={location.longitude}&appid={api_key}&units=metric"
        response = requests.get(url)
        if response.status_code != 200:
            return f"Error fetching weather for '{city}': API returned status {response.status_code}"
        
        data = response.json()
        temp = data["main"]["temp"]
        weather = data["weather"][0]["main"]
        humidity = data["main"]["humidity"]

        translations = {
            "en": {
                "weather": "Weather",
                "temp": "Temperature",
                "humidity": "Humidity",
                "conditions": {
                    "Clear": "Clear ☀️",
                    "Clouds": "Cloudy ☁️",
                    "Rain": "Rain 🌧",
                    "Snow": "Snow ❄️",
                    "Thunderstorm": "Thunderstorm ⛈",
                    "Drizzle": "Drizzle 🌦",
                    "Mist": "Mist 🌫"
                }
            },
            "ru": {
                "weather": "Погода",
                "temp": "Температура",
                "humidity": "Влажность",
                "conditions": {
                    "Clear": "Ясно ☀️",
                    "Clouds": "Облачно ☁️",
                    "Rain": "Дождь 🌧",
                    "Snow": "Снег ❄️",
                    "Thunderstorm": "Гроза ⛈",
                    "Drizzle": "Морось 🌦",
                    "Mist": "Туман 🌫"
                }
            }
        }
        lang = lang if lang in translations else "en"
        t = translations[lang]
        condition = t["conditions"].get(weather, weather)
        
        return (
            f"🌡️ {t['weather']}: {condition}\n"
            f"{t['temp']}: {temp}°C\n"
            f"{t['humidity']}: {humidity}%"
        )
    except Exception as e:
        return f"Error fetching weather for '{city}': {str(e)}"

@tool
def final_answer(answer: str) -> str:
    """Returns the final answer.
    
    Args:
        answer: A string containing the final answer.
    
    Returns:
        The input answer 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. "
        "If the query contains a city name (e.g., 'Бураево', 'Sydney', 'न्यू यॉर्क') or phrases like 'сколько времени', 'what time', "
        "or a city in any context (e.g., 'Что там в Бирске?'), return current time, air quality, and weather for that city. "
        "Use ONLY the provided functions: `get_timezone_by_city`, `get_current_time_in_timezone`, `get_air_quality`, `get_weather`, `final_answer`. "
        "Detect query language with `detect` and pass it to `get_air_quality` and `get_weather` for translated output. "
        "Return ONLY the result via `final_answer()` in a code block:\n"
        "```py\n"
        "final_answer('YOUR ANSWER HERE')\n"
        "```\n"
        "Do NOT include 'thoughts', 'code', explanations, extra imports, or invalid dates. "
        "Do NOT define new functions or modify existing ones. "
        "If the city is not found or query is unclear, return an error via `final_answer()`. "
        "NEVER generate code without ```py``` block or with invalid syntax."
    ),
    "default": "Response: {{question}}",
    "planning": {
        "initial_plan": (
            "Analyze query: {{question}}. If it contains a city, get time, air quality, and weather using provided functions. "
            "Return code:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
        "update_plan_pre_messages": (
            "Review query: {{question}}. Adjust plan. Return code:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
        "update_plan_post_messages": (
            "Review query: {{question}} and results. Adjust plan. Return code:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
    },
    "managed_agent": {
        "execute": (
            "Execute task: {{question}}. If it contains a city or phrases like 'сколько времени', use provided functions. "
            "Return code:\n"
            "```py\n"
            "from langdetect import detect\n"
            "query = '{{question}}'.strip()\n"
            "city = query.split(' в ')[-1].strip().title() if ' в ' in query else query.title()\n"
            "lang = detect(query) if query else 'en'\n"
            "timezone = get_timezone_by_city(city)\n"
            "if timezone.startswith('Error'):\n"
            "    final_answer(timezone)\n"
            "else:\n"
            "    time_result = get_current_time_in_timezone(timezone)\n"
            "    air_result = get_air_quality(city, lang)\n"
            "    weather_result = get_weather(city, lang)\n"
            "    final_answer(f'🕒 Время в {city}: {time_result}\\n' + air_result + '\\n' + weather_result)\n"
            "```"
        ),
        "report": (
            "Summarize task: {{question}}. Return code:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
        "task": (
            "Define task: {{question}}. Return code:\n"
            "```py\n"
            "final_answer('YOUR TASK DEFINITION HERE')\n"
            "```"
        ),
    },
    "final_answer": {
        "pre_messages": (
            "Prepare response for: {{question}}. Return code:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
        "template": "Final response: {{answer}}",
        "post_messages": (
            "Review response for: {{question}}. Return code:\n"
            "```py\n"
            "final_answer('YOUR ANSWER HERE')\n"
            "```"
        ),
    },
}

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

def process_input(user_input):
    try:
        response = agent.run(user_input)
        return response if response else "Error: No response received from the agent."
    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, air quality, weather, and more!"
).launch()