Update app.py

app.py

@@ -1,5 +1,6 @@
 from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
 import datetime
+from difflib import get_close_matches
 import requests
 import pytz
 import yaml
@@ -18,6 +19,128 @@ def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return
     """
     return "What magic will you build ?"
 
+MIN_YEAR = 2023
+
+def sanitize_name(name: str) -> str | None:
+    """Only allow letters, hyphens, apostrophes and spaces (covers names like O'Ward, Häkkinen)."""
+    if not re.match(r"^[a-zA-ZÀ-ÿ\s'\-]{2,50}$", name):
+        return None
+    return name.strip()
+
+def sanitize_year(year: int) -> int | None:
+    """Only allow integers within a valid range."""
+    if not isinstance(year, int) or isinstance(year, bool):
+        return None
+    if year < MIN_YEAR or year > datetime.now().year:
+        return None
+    return year
+
+@tool
+def driver_number(name: str, year: int) -> str:
+    """A tool that give the number of a F1 driver for a specific session  
+    Args:
+        name: the name of the driver
+        year: the year for which we want the number from 2023
+    """
+
+    # Sanitize inputs before anything else
+    clean_name = sanitize_name(name)
+    if clean_name is None:
+        return "Invalid driver name. Only letters, spaces, hyphens and apostrophes are allowed (2-100 characters)."
+
+    clean_year = sanitize_year(year)
+    if clean_year is None:
+        return f"Invalid year. Please provide a year between {MIN_YEAR} and {datetime.now().year}."
+
+    try:
+        # Step 1: Get any session_key for the given year
+        sessions_resp = requests.get(
+            "https://api.openf1.org/v1/sessions",
+            params={"year": clean_year},
+            timeout=10
+        )
+
+        if sessions_resp.status_code != 200:
+            return f"Error fetching sessions: HTTP {sessions_resp.status_code}"
+
+        sessions = sessions_resp.json()
+
+        if not sessions:
+            return f"No sessions found for year {clean_year}."
+
+        session_key = sessions[0]["session_key"]
+
+        # Validate session_key is an integer before using it
+        if not isinstance(session_key, int):
+            return "Unexpected session data returned from API."
+
+        # Step 2: Search for the driver by last name within that session
+        drivers_resp = requests.get(
+            "https://api.openf1.org/v1/drivers",
+            params={"session_key": session_key, "last_name": clean_name},
+            timeout=10
+        )
+
+        if drivers_resp.status_code != 200:
+            return f"Error fetching drivers: HTTP {drivers_resp.status_code}"
+
+        drivers = drivers_resp.json()
+
+        # Fallback 1: try full name exact match
+        if not drivers:
+            drivers_resp = requests.get(
+                "https://api.openf1.org/v1/drivers",
+                params={"session_key": session_key, "full_name": clean_name.upper()},
+                timeout=10
+            )
+            drivers = drivers_resp.json() if drivers_resp.status_code == 200 else []
+
+        # Fallback 2: fuzzy match against all drivers in that session
+        if not drivers:
+            all_drivers_resp = requests.get(
+                "https://api.openf1.org/v1/drivers",
+                params={"session_key": session_key},
+                timeout=10
+            )
+
+            if all_drivers_resp.status_code == 200:
+                all_drivers = all_drivers_resp.json()
+
+                # Build a mapping of searchable names -> driver object
+                name_map = {}
+                for d in all_drivers:
+                    name_map[d.get("last_name", "").lower()] = d
+                    name_map[d.get("full_name", "").lower()] = d
+                    name_map[d.get("broadcast_name", "").lower()] = d
+                    name_map[d.get("name_acronym", "").lower()] = d
+
+                matches = get_close_matches(
+                    clean_name.lower(),
+                    name_map.keys(),
+                    n=1,
+                    cutoff=0.6
+                )
+
+                if matches:
+                    drivers = [name_map[matches[0]]]
+
+        if not drivers:
+            return f"No driver found with name '{clean_name}' in {clean_year}."
+
+        # Validate the returned driver_number is actually an integer
+        driver_num = drivers[0].get("driver_number")
+        if not isinstance(driver_num, int):
+            return "Unexpected driver data returned from API."
+
+        return str(driver_num)
+
+    except requests.exceptions.Timeout:
+        return "Request timed out. The OpenF1 API may be slow or unavailable."
+    except requests.exceptions.ConnectionError:
+        return "Could not connect to the OpenF1 API. Check network connectivity."
+    except Exception as e:
+        return f"Unexpected error: {str(e)}"
+
 @tool
 def get_current_time_in_timezone(timezone: str) -> str:
     """A tool that fetches the current local time in a specified timezone.