Update app.py

app.py

@@ -5,6 +5,7 @@ import pytz
 import gradio as gr
 from geopy.geocoders import Nominatim
 from timezonefinder import TimezoneFinder
+import requests
 
 @tool
 def get_timezone_by_city(city: str) -> str:
@@ -18,10 +19,8 @@ def get_timezone_by_city(city: str) -> 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()
             location = geolocator.geocode(city_name, language="en")
             if not location:
@@ -49,6 +48,44 @@ def get_current_time_in_timezone(timezone: str) -> str:
     except Exception as e:
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
+@tool
+def get_air_quality(city: str) -> str:
+    """Tool to fetch air quality data for a given city using OpenWeatherMap API.
+    Args:
+        city: A string representing the city name (e.g., 'Бирск', 'New York').
+    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()
+            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"]
+        
+        aqi_description = {1: "Good", 2: "Fair", 3: "Moderate", 4: "Poor", 5: "Very Poor"}
+        return (
+            f"Air quality in {city}: AQI {aqi} ({aqi_description.get(aqi, 'Unknown')})\n"
+            f"PM2.5: {components['pm2_5']} µg/m³, PM10: {components['pm10']} µg/m³\n"
+            f"CO: {components['co']} µg/m³, NO2: {components['no2']} µg/m³, SO2: {components['so2']} µg/m³, O3: {components['o3']} µg/m³"
+        )
+    except Exception as e:
+        return f"Error fetching air quality for '{city}': {str(e)}"
+
 @tool
 def final_answer(answer: str) -> str:
     """Returns the final answer.
@@ -71,22 +108,24 @@ prompt_templates = {
         "You are a helpful assistant that responds in the same language as the user's query. "
         "Detect the language and return only the result, without explanations, images, or repeating instructions. "
         "If the query contains a city name (e.g., 'Бураево', 'Осиновка', 'Fosnavåg', 'न्यू यॉर्क') or phrases like 'сколько времени', 'what time', 'quelle heure', "
-        "use `get_timezone_by_city` to find the timezone, then `get_current_time_in_timezone` to get the time. "
+        "or if a city is mentioned in any context (e.g., 'Что там в Бирске?'), treat it as a request for the current time and air quality in that city. "
+        "Use `get_timezone_by_city` to find the timezone, `get_current_time_in_timezone` for the time, and `get_air_quality` for air quality data. "
+        "Combine results in a single response. "
         "Always return the response via `final_answer()` in a code block like this:\n"
         "```py\n"
         "final_answer('YOUR ANSWER HERE')\n"
         "```\n"
-        "Execute tools in sequence: first call get_timezone_by_city(city), then get_current_time_in_timezone(timezone). "
+        "Execute tools in sequence: first call get_timezone_by_city(city), then get_current_time_in_timezone(timezone), then get_air_quality(city). "
         "Do NOT include tool calls, formatting like 'Current time in X: ', or invalid dates (e.g., '2023-03-22') inside final_answer(). "
-        "If the city is not found or the query is unclear, return an error message in the same language via final_answer(). "
+        "If the city is not found or the query is unclear (e.g., no city mentioned), return an error message in the same language via final_answer(). "
         "NEVER generate code without the ```py``` block or with invalid syntax. "
-        "Do NOT ask for clarifications; return an error if the query is ambiguous."
+        "Do NOT ask for clarifications or generate responses outside the ```py``` block."
     ),
     "default": "Response: {{question}}",
     "planning": {
         "initial_plan": (
             "Analyze the query: {{question}}. If it contains a city name or phrases like 'сколько времени', "
-            "use `get_timezone_by_city` and `get_current_time_in_timezone`. "
+            "use `get_timezone_by_city`, `get_current_time_in_timezone`, and `get_air_quality`. "
             "Write code in a code block:\n"
             "```py\n"
             "final_answer('YOUR ANSWER HERE')\n"
@@ -109,16 +148,18 @@ prompt_templates = {
     },
     "managed_agent": {
         "execute": (
-            "Execute the task: {{question}}. If it contains a city name or phrases like 'сколько времени', use `get_timezone_by_city` and `get_current_time_in_timezone`. "
+            "Execute the task: {{question}}. If it contains a city name or phrases like 'сколько времени', use `get_timezone_by_city`, `get_current_time_in_timezone`, and `get_air_quality`. "
             "Write code in a code block:\n"
             "```py\n"
-            "city = '{{question}}'.strip()\n"
+            "query = '{{question}}'.strip()\n"
+            "city = query.split(' в ')[-1].strip() if ' в ' in query else query\n"
             "timezone = get_timezone_by_city(city)\n"
-            "if not timezone.startswith('Error'):\n"
-            "    result = get_current_time_in_timezone(timezone)\n"
-            "    final_answer(result)\n"
-            "else:\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)\n"
+            "    final_answer(time_result + '\\n' + air_result)\n"
             "```"
         ),
         "report": (
@@ -157,7 +198,7 @@ prompt_templates = {
 
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, get_current_time_in_timezone, get_timezone_by_city],
+    tools=[final_answer, get_current_time_in_timezone, get_timezone_by_city, get_air_quality],
     max_steps=3,
     verbosity_level=2,
     prompt_templates=prompt_templates,
@@ -175,5 +216,5 @@ gr.Interface(
     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!"
+    description="Hello! What city or town are you from? I can tell you the time, air quality, and more!"
 ).launch()