diff --git a/fli/mcp/server.py b/fli/mcp/server.py index a4e286fa..3f280308 100755 --- a/fli/mcp/server.py +++ b/fli/mcp/server.py @@ -132,6 +132,11 @@ class FlightSearchParams(BaseModel): ge=1, description="Number of adult passengers", ) + children: int = Field(0, ge=0, description="Number of children (ages 2-11)") + infants_in_seat: int = Field( + 0, ge=0, description="Number of infants (under 2) occupying their own seat" + ) + infants_on_lap: int = Field(0, ge=0, description="Number of lap infants (under 2, no seat)") exclude_basic_economy: bool = Field( False, description="Exclude basic economy fares from results" ) @@ -218,6 +223,11 @@ class DateSearchParams(BaseModel): ge=1, description="Number of adult passengers", ) + children: int = Field(0, ge=0, description="Number of children (ages 2-11)") + infants_in_seat: int = Field( + 0, ge=0, description="Number of infants (under 2) occupying their own seat" + ) + infants_on_lap: int = Field(0, ge=0, description="Number of lap infants (under 2, no seat)") currency: str | None = Field( None, description=( @@ -581,7 +591,12 @@ def _build_flight_filters( filters = FlightSearchFilters( trip_type=trip_type, - passenger_info=PassengerInfo(adults=params.passengers), + passenger_info=PassengerInfo( + adults=params.passengers, + children=params.children, + infants_in_seat=params.infants_in_seat, + infants_on_lap=params.infants_on_lap, + ), flight_segments=segments, stops=max_stops, seat_type=cabin_class, @@ -803,7 +818,12 @@ def _execute_date_search(params: DateSearchParams) -> dict[str, Any]: # Create search filters filters = DateSearchFilters( trip_type=trip_type, - passenger_info=PassengerInfo(adults=params.passengers), + passenger_info=PassengerInfo( + adults=params.passengers, + children=params.children, + infants_in_seat=params.infants_in_seat, + infants_on_lap=params.infants_on_lap, + ), flight_segments=segments, stops=max_stops, seat_type=cabin_class, @@ -920,6 +940,18 @@ def search_flights( int | None, Field(description="Number of adult passengers", ge=1), ] = None, + children: Annotated[ + int, + Field(description="Number of children (ages 2-11)", ge=0), + ] = 0, + infants_in_seat: Annotated[ + int, + Field(description="Number of infants (under 2) occupying their own seat", ge=0), + ] = 0, + infants_on_lap: Annotated[ + int, + Field(description="Number of lap infants (under 2, no seat)", ge=0), + ] = 0, exclude_basic_economy: Annotated[ bool, Field(description="Exclude basic economy fares from results"), @@ -995,6 +1027,9 @@ def search_flights( max_stops=max_stops, sort_by=sort_by, passengers=passengers or CONFIG.default_passengers, + children=children, + infants_in_seat=infants_in_seat, + infants_on_lap=infants_on_lap, exclude_basic_economy=exclude_basic_economy, emissions=emissions, checked_bags=checked_bags, @@ -1073,6 +1108,18 @@ def search_dates( int | None, Field(description="Number of adult passengers", ge=1), ] = None, + children: Annotated[ + int, + Field(description="Number of children (ages 2-11)", ge=0), + ] = 0, + infants_in_seat: Annotated[ + int, + Field(description="Number of infants (under 2) occupying their own seat", ge=0), + ] = 0, + infants_on_lap: Annotated[ + int, + Field(description="Number of lap infants (under 2, no seat)", ge=0), + ] = 0, currency: Annotated[ str | None, Field(description="ISO 4217 currency code (USD, EUR, GBP, JPY...) for prices."), @@ -1125,6 +1172,9 @@ def search_dates( departure_window=effective_departure_window, sort_by_price=sort_by_price, passengers=passengers or CONFIG.default_passengers, + children=children, + infants_in_seat=infants_in_seat, + infants_on_lap=infants_on_lap, currency=currency, language=language, country=country, @@ -1186,6 +1236,18 @@ def get_booking_options( int | None, Field(description="Number of adult passengers", ge=1), ] = None, + children: Annotated[ + int, + Field(description="Number of children (ages 2-11)", ge=0), + ] = 0, + infants_in_seat: Annotated[ + int, + Field(description="Number of infants (under 2) occupying their own seat", ge=0), + ] = 0, + infants_on_lap: Annotated[ + int, + Field(description="Number of lap infants (under 2, no seat)", ge=0), + ] = 0, airlines: Annotated[ list[str] | None, Field(description="Filter by airline IATA codes (e.g., ['BA', 'AA'])"), @@ -1276,6 +1338,9 @@ def get_booking_options( max_stops=max_stops, sort_by=sort_by, passengers=passengers or CONFIG.default_passengers, + children=children, + infants_in_seat=infants_in_seat, + infants_on_lap=infants_on_lap, airlines=airlines, exclude_basic_economy=exclude_basic_economy, emissions=emissions, diff --git a/tests/mcp/test_new_filters.py b/tests/mcp/test_new_filters.py index b83c5c21..8478d2fb 100644 --- a/tests/mcp/test_new_filters.py +++ b/tests/mcp/test_new_filters.py @@ -19,6 +19,7 @@ from fli.mcp.server import ( DateSearchParams, FlightSearchParams, + _get_booking_options_from_params, _search_dates_from_params, _search_flights_from_params, ) @@ -233,3 +234,77 @@ def test_layover_restrictions_built(self, captured_dates): assert filters.layover_restrictions is not None assert filters.layover_restrictions.min_duration == 120 assert filters.layover_restrictions.max_duration == 600 + + +# --------------------------------------------------------------------------- +# passenger mix — adults / children / infants flow into PassengerInfo +# --------------------------------------------------------------------------- + + +class TestPassengerMix: + def test_flight_search_passenger_mix(self, captured_search): + params = FlightSearchParams( + origin="OPO", + destination="HKG", + departure_date=_future(30), + passengers=2, + children=1, + infants_in_seat=1, + ) + _search_flights_from_params(params) + info = captured_search["filters"].passenger_info + assert info.adults == 2 + assert info.children == 1 + assert info.infants_in_seat == 1 + assert info.infants_on_lap == 0 + + def test_flight_search_defaults_to_adults_only(self, captured_search): + params = FlightSearchParams( + origin="JFK", + destination="LAX", + departure_date=_future(30), + passengers=1, + ) + _search_flights_from_params(params) + info = captured_search["filters"].passenger_info + assert (info.children, info.infants_in_seat, info.infants_on_lap) == (0, 0, 0) + + def test_date_search_passenger_mix(self, captured_dates): + params = DateSearchParams( + origin="OPO", + destination="HKG", + start_date=_future(30), + end_date=_future(60), + passengers=2, + children=1, + infants_on_lap=1, + ) + _search_dates_from_params(params) + info = captured_dates["filters"].passenger_info + assert info.adults == 2 + assert info.children == 1 + assert info.infants_on_lap == 1 + assert info.infants_in_seat == 0 + + def test_booking_options_passenger_mix(self, captured_search): + # ``get_booking_options`` builds its own ``FlightSearchParams`` and + # re-runs the search via ``_build_flight_filters`` rather than + # delegating to ``_search_flights_from_params``. Cover that separate + # construction site so a passenger-field typo there is caught too. + # The patched ``search`` returns ``[]``, so the booking path stops at + # its empty-result branch after the filter is captured. + params = FlightSearchParams( + origin="OPO", + destination="HKG", + departure_date=_future(30), + passengers=2, + children=1, + infants_in_seat=1, + infants_on_lap=1, + ) + _get_booking_options_from_params(params) + info = captured_search["filters"].passenger_info + assert info.adults == 2 + assert info.children == 1 + assert info.infants_in_seat == 1 + assert info.infants_on_lap == 1