Update app.py

app.py

@@ -3,6 +3,31 @@ from smolagents import CodeAgent, HfApiModel, tool
 import datetime
 import pytz
 import gradio as gr
+from geopy.geocoders import Nominatim
+from timezonefinder import TimezoneFinder
+
+@tool
+def get_timezone_by_city(city: str) -> str:
+    """Tool to fetch the timezone for a given city using geolocation.
+    Args:
+        city: A string representing the city name (e.g., 'Бураево', 'Tønsberg').
+    Returns:
+        A string with the timezone (e.g., 'Europe/Samara') 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:
+            return f"Error: City '{city}' not found."
+        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:
@@ -39,30 +64,31 @@ except Exception as 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. "
-        "If the query contains a city name (e.g., 'Birsk', 'Tønsberg') or is explicitly about time (e.g., 'What time is it in Moscow?'), "
-        "ALWAYS use the `get_current_time_in_timezone` tool with the appropriate timezone (e.g., 'Europe/Moscow' for Moscow, 'Asia/Yekaterinburg' for Birsk, 'Europe/Oslo' for Tønsberg). "
+        "Detect the language and return only the result, without explanations. "
+        "If the query contains a city name (e.g., 'Birsk', 'Tønsberg', 'Бураево') or phrases like 'сколько времени', 'what time', 'quelle heure', "
+        "ALWAYS use the `get_timezone_by_city` tool to find the timezone, then use `get_current_time_in_timezone` to get the time. "
         "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"
-        "Execute the tool first, then pass the result as a string to final_answer(). "
-        "Do NOT include tool calls (e.g., get_current_time_in_timezone()) inside final_answer(). "
-        "If the query is unclear or does not specify a city, return a clarification request in the same language via final_answer()."
+        "Execute the tools in sequence (get_timezone_by_city, then get_current_time_in_timezone), then pass the result as a string to final_answer(). "
+        "Do NOT include tool calls 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()."
     ),
     "default": "Response: {{question}}",
     "planning": {
         "initial_plan": (
-            "Analyze the query: {{question}}. If it contains a city name (e.g., 'Birsk', 'Tønsberg') or is about time, use `get_current_time_in_timezone` with the correct timezone (e.g., 'Asia/Yekaterinburg' for Birsk, 'Europe/Oslo' for Tønsberg). "
-            "Write code to execute the tool and return the result via final_answer() in a code block:\n"
+            "Analyze the query: {{question}}. If it contains a city name (e.g., 'Birsk', 'Бураево') or phrases like 'сколько времени', "
+            "use `get_timezone_by_city` to find the timezone, then `get_current_time_in_timezone` to get the time. "
+            "Write code to execute the tools 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. "
+            "If the plan needs adjustment (e.g., incorrect city or timezone), update it. "
             "Write corrected code in a code block:\n"
             "```py\n"
             "final_answer('YOUR ANSWER HERE')\n"
@@ -79,11 +105,16 @@ prompt_templates = {
     },
     "managed_agent": {
         "execute": (
-            "Execute the task: {{question}}. If it contains a city name or is about time, use `get_current_time_in_timezone`. "
+            "Execute the task: {{question}}. If it contains a city name or phrases like 'сколько времени', use `get_timezone_by_city` and `get_current_time_in_timezone`. "
             "Write code in a code block:\n"
             "```py\n"
-            "result = get_current_time_in_timezone('TIMEZONE')\n"
-            "final_answer(result)\n"
+            "city = '{{question}}'.lower().split()[-1]  # Extract last word as city\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"
+            "    final_answer(timezone)  # Return error message\n"
             "```"
         ),
         "report": (
@@ -95,7 +126,7 @@ prompt_templates = {
         ),
         "task": (
             "Define the task based on the query: {{question}}. "
-            "If it contains a city name, use `get_current_time_in_timezone`. "
+            "If it contains a city name or phrases like 'сколько времени', use `get_timezone_by_city` and `get_current_time_in_timezone`. "
             "Return the task definition via final_answer() in a code block:\n"
             "```py\n"
             "final_answer('YOUR TASK DEFINITION HERE')\n"
@@ -125,8 +156,8 @@ prompt_templates = {
 
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, get_current_time_in_timezone],
-    max_steps=2,
+    tools=[final_answer, get_current_time_in_timezone, get_timezone_by_city],
+    max_steps=3,
     verbosity_level=2,
     prompt_templates=prompt_templates,
 )