import os import gradio as gr import json from typing import Optional from src.common.logger import get_logger from src.common.flixbus_client import CityInfo, FlixBusClient from src.common.utils import format_docstring from src.tools.schemas.flixbus_tools_schema import ( GetCityInfoResilt, GetTripeOptionsListResult, SearchCitiesInRadiusResult, TripeOption, ) logger = get_logger(__name__) ##### get_city_info ##### async def get_city_info(city_name: str = "") -> Optional[GetCityInfoResilt]: """ Resolves a city name to its Flixbus UUID. """ client = FlixBusClient() try: suggestions = await client.get_city_suggestions(q=city_name) if not suggestions: return None city = suggestions[0] return GetCityInfoResilt( id=city.id, name=city.name, country=city.country, latitude=city.location.lat, longitude=city.location.lon, ) except Exception as e: logger.error(f"Error resolving city ID: {e}") return None finally: await client.close() @gr.mcp.resource( "ui://widget/get_city_info_widget.html", mime_type="text/html+skybridge" ) def get_city_info_resource(): with open( os.path.join(os.path.dirname(__file__), "widgets/get_city_info_widget.html"), "r", ) as f: return f.read() @format_docstring(template_args={"return_schema": GetCityInfoResilt._model_json_schema}) @gr.mcp.tool( name="get_city_info_tool", _meta={ "openai/outputTemplate": "ui://widget/get_city_info_widget.html", "openai/resultCanProduceWidget": True, "openai/widgetAccessible": True, }, ) async def get_city_info_tool( city_name: str = "", ) -> str: """ Resolves a city name to its Flixbus UUID. Args: city_name: The name of the city to resolve. Returns: {return_schema} """ city_info = await get_city_info(city_name) if city_info is None: return json.dumps( {"error": f"City '{city_name}' not found in Flixbus database."} ) return city_info.model_dump_json() ##### search_cities_in_radius ##### async def search_cities_in_radius( lat: float = 0, lon: float = 0, radius_km: int = 300, limit: int = 20 ) -> SearchCitiesInRadiusResult: """ Finds Flixbus cities within a bounding box roughly corresponding to the radius. Note: This is an approximation using a bounding box. """ client = FlixBusClient() try: # 1 degree of latitude is approx 111km lat_delta = radius_km / 111 # 1 degree of longitude is approx 111km * cos(lat) - simplifying to 111km for rough box lon_delta = radius_km / 111 top_left_lat = lat + lat_delta top_left_lon = lon - lon_delta bottom_right_lat = lat - lat_delta bottom_right_lon = lon + lon_delta response = await client.get_cms_cities( top_left_lat=top_left_lat, top_left_lon=top_left_lon, bottom_right_lat=bottom_right_lat, bottom_right_lon=bottom_right_lon, limit=limit, ) cities: list[GetCityInfoResilt] = [] for city in response.result: cities.append( GetCityInfoResilt( name=city.name, id=city.uuid, # Using UUID as ID latitude=city.location.lat, longitude=city.location.lon, country=city.country, ) ) return SearchCitiesInRadiusResult(cities=cities) finally: await client.close() @gr.mcp.resource( "ui://widget/search_cities_in_radius_widget.html", mime_type="text/html+skybridge" ) def search_cities_in_radius_resource(): """ Returns the HTML content of the search cities in radius widget. """ with open( os.path.join( os.path.dirname(__file__), "widgets/search_cities_in_radius_widget.html" ), "r", ) as f: return f.read() @format_docstring( template_args={"return_schema": SearchCitiesInRadiusResult._model_json_schema} ) @gr.mcp.tool( name="search_cities_in_radius_tool", _meta={ "openai/outputTemplate": "ui://widget/search_cities_in_radius_widget.html", "openai/resultCanProduceWidget": True, "openai/widgetAccessible": True, }, ) async def search_cities_in_radius_tool( lat: float = 0, lon: float = 0, radius_km: int = 300, limit: int = 20 ) -> str: """ Finds Flixbus cities within a bounding box roughly corresponding to the radius. Note: This is an approximation using a bounding box. Args: lat: Latitude of the center point. lon: Longitude of the center point. radius_km: Search radius in kilometers, default is 300km. limit: Maximum number of cities to return, default is 20. Returns: {return_schema} """ result = await search_cities_in_radius(lat, lon, radius_km, limit) return result.model_dump_json() ##### get_trip_options_list ##### async def get_tripe_options_list( from_city_id: str = "", to_city_id: str = "", date: str = "" ) -> GetTripeOptionsListResult: """ Searches for the trips between two cities. """ client = FlixBusClient() shop_url = client.get_shop_url_for_search_trips( from_city_id=from_city_id, to_city_id=to_city_id, departure_date=date ) empty_result = GetTripeOptionsListResult( from_city=CityInfo(name="-", slug="-", id=from_city_id), to_city=CityInfo(name="-", slug="-", id=to_city_id), options=[], best_option=None, shop_url=shop_url, ) try: response = await client.search_trips( from_city_id=from_city_id, to_city_id=to_city_id, departure_date=date ) if not response.trips: return empty_result options: list[TripeOption] = [] best_option: Optional[TripeOption] = None tripe_options = response.trips[0].results for _, trip_option in tripe_options.items(): if trip_option.status != "available": continue option = TripeOption( price=trip_option.price.total, departure_station_id=trip_option.departure.station_id, arrival_station_id=trip_option.arrival.station_id, departure_time=trip_option.departure.date.strftime( "%Y-%m-%dT%H:%M:%S%z" ), arrival_time=trip_option.arrival.date.strftime("%Y-%m-%dT%H:%M:%S%z"), duration_hours=trip_option.duration.hours, available_seats=trip_option.available.seats, ) options.append(option) if best_option is None or option.price < best_option.price: best_option = option from_city = response.cities[from_city_id] to_city = response.cities[to_city_id] return GetTripeOptionsListResult( options=options, best_option=best_option, shop_url=shop_url, from_city=from_city, to_city=to_city, ) except Exception as e: logger.error(f"Error finding trip: {e}") return empty_result finally: await client.close() @gr.mcp.resource( "ui://widget/get_tripe_options_list_widget.html", mime_type="text/html+skybridge" ) def get_tripe_options_list_resource(): """ Returns the HTML content of the get trip options list widget. """ with open( os.path.join( os.path.dirname(__file__), "widgets/get_tripe_options_list_widget.html" ), "r", ) as f: return f.read() @format_docstring( template_args={"return_schema": GetTripeOptionsListResult._model_json_schema} ) @gr.mcp.tool( name="get_tripe_options_list_tool", _meta={ "openai/outputTemplate": "ui://widget/get_tripe_options_list_widget.html", "openai/resultCanProduceWidget": True, "openai/widgetAccessible": True, }, ) async def get_tripe_options_list_tool( from_city_id: str = "", to_city_id: str = "", date: str = "" ) -> str: """ Searches for the trips between two cities. Args: from_city_id: Flixbus UUID of the departure city. to_city_id: Flixbus UUID of the destination city. date: Departure date in format DD.MM.YYYY. Returns: {return_schema} """ result = await get_tripe_options_list(from_city_id, to_city_id, date) return result.model_dump_json()