src/tools/flixbus_tools.py

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()