diff --git a/README.md b/README.md index 2e83c447..c9f3ac27 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,9 @@ The MCP server provides two main tools: | `airlines` | list | Filter by airline codes (e.g., ['BA', 'AA']) | | `sort_by` | string | CHEAPEST, DURATION, DEPARTURE_TIME, or ARRIVAL_TIME | | `passengers` | int | Number of adult passengers | +| `children` | int | Number of child passengers | +| `infants_in_seat` | int | Number of infants traveling in their own seat | +| `infants_on_lap` | int | Number of lap infants | #### `search_dates` Parameters @@ -82,6 +85,9 @@ The MCP server provides two main tools: | `airlines` | list | Filter by airline codes (e.g., ['BA', 'AA']) | | `sort_by_price` | bool | Sort results by price (lowest first) | | `passengers` | int | Number of adult passengers | +| `children` | int | Number of child passengers | +| `infants_in_seat` | int | Number of infants traveling in their own seat | +| `infants_on_lap` | int | Number of lap infants | ## Quick Start @@ -185,6 +191,10 @@ fli dates JFK LHR \ | `--class, -c` | Cabin class | `ECONOMY`, `BUSINESS` | | `--stops, -s` | Maximum stops | `NON_STOP`, `ONE_STOP` | | `--sort, -o` | Sort results by | `CHEAPEST`, `DURATION` | +| `--passengers` | Number of adult passengers | `2` | +| `--children` | Number of children | `2` | +| `--infants-in-seat` | Number of infants in seat | `1` | +| `--infants-on-lap` | Number of infants on lap | `1` | | `--format` | Output format | `text`, `json` | #### Dates Command (`fli dates`) @@ -200,6 +210,10 @@ fli dates JFK LHR \ | `--stops, -s` | Maximum stops | `NON_STOP`, `ONE_STOP` | | `--time` | Departure time window | `6-20` | | `--sort` | Sort by price | (flag) | +| `--passengers` | Number of adult passengers | `2` | +| `--children` | Number of children | `2` | +| `--infants-in-seat` | Number of infants in seat | `1` | +| `--infants-on-lap` | Number of infants on lap | `1` | | `--[day]` | Day filters | `--monday`, `--friday` | | `--format` | Output format | `text`, `json` | diff --git a/docs/api/models.md b/docs/api/models.md index d026b3b3..1acd2795 100644 --- a/docs/api/models.md +++ b/docs/api/models.md @@ -88,6 +88,13 @@ Type of trip for flight search. Configuration for passenger counts. +Fields: + +- `adults` +- `children` +- `infants_in_seat` +- `infants_on_lap` + ::: fli.models.google_flights.PassengerInfo ### TimeRestrictions diff --git a/docs/examples/advanced.md b/docs/examples/advanced.md index cc26acf5..63caf823 100644 --- a/docs/examples/advanced.md +++ b/docs/examples/advanced.md @@ -35,6 +35,7 @@ filters = FlightSearchFilters( passenger_info=PassengerInfo( adults=2, children=1, + infants_in_seat=1, infants_on_lap=1 ), flight_segments=[ @@ -280,6 +281,7 @@ filters = FlightSearchFilters( passenger_info=PassengerInfo( adults=2, children=1, + infants_in_seat=1, infants_on_lap=1 ), flight_segments=[outbound, return_flight], diff --git a/docs/guides/mcp.md b/docs/guides/mcp.md index 27857dea..834d7726 100644 --- a/docs/guides/mcp.md +++ b/docs/guides/mcp.md @@ -78,6 +78,9 @@ Search for flights between two airports on a specific date. | `airlines` | list | No | null | Filter by airline codes (e.g., ['BA', 'AA']) | | `sort_by` | string | No | CHEAPEST | CHEAPEST, DURATION, DEPARTURE_TIME, or ARRIVAL_TIME | | `passengers` | int | No | 1 | Number of adult passengers | +| `children` | int | No | 0 | Number of child passengers | +| `infants_in_seat` | int | No | 0 | Number of infants traveling in their own seat | +| `infants_on_lap` | int | No | 0 | Number of lap infants | **Example Response:** @@ -126,6 +129,9 @@ Find the cheapest travel dates between two airports within a date range. | `airlines` | list | No | null | Filter by airline codes (e.g., ['BA', 'AA']) | | `sort_by_price` | bool | No | false | Sort results by price (lowest first) | | `passengers` | int | No | 1 | Number of adult passengers | +| `children` | int | No | 0 | Number of child passengers | +| `infants_in_seat` | int | No | 0 | Number of infants traveling in their own seat | +| `infants_on_lap` | int | No | 0 | Number of lap infants | **Example Response:** diff --git a/examples/README.md b/examples/README.md index a8bb6bac..92f94b9d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -211,7 +211,7 @@ from fli.models import PassengerInfo PassengerInfo(adults=1) # Family with children -PassengerInfo(adults=2, children=2, infants_on_lap=1) +PassengerInfo(adults=2, children=2, infants_in_seat=1, infants_on_lap=1) ``` ## Support diff --git a/examples/complex_flight_search.py b/examples/complex_flight_search.py index d352b083..806a1c00 100644 --- a/examples/complex_flight_search.py +++ b/examples/complex_flight_search.py @@ -25,7 +25,12 @@ def main(): # Create detailed filters filters = FlightSearchFilters( trip_type=TripType.ONE_WAY, - passenger_info=PassengerInfo(adults=2, children=1, infants_on_lap=1), + passenger_info=PassengerInfo( + adults=2, + children=1, + infants_in_seat=1, + infants_on_lap=1, + ), flight_segments=[ FlightSegment( departure_airport=[[Airport.JFK, 0]], diff --git a/examples/complex_round_trip_validation.py b/examples/complex_round_trip_validation.py index b7ae551e..a7628f3a 100644 --- a/examples/complex_round_trip_validation.py +++ b/examples/complex_round_trip_validation.py @@ -80,7 +80,12 @@ def main(): # Create filters with complex requirements filters = FlightSearchFilters( trip_type=TripType.ROUND_TRIP, - passenger_info=PassengerInfo(adults=2, children=1, infants_on_lap=1), + passenger_info=PassengerInfo( + adults=2, + children=1, + infants_in_seat=1, + infants_on_lap=1, + ), flight_segments=[outbound, return_flight], stops=MaxStops.ONE_STOP_OR_FEWER, seat_type=SeatType.BUSINESS, diff --git a/fli/cli/commands/dates.py b/fli/cli/commands/dates.py index baf53216..f8d29cf4 100644 --- a/fli/cli/commands/dates.py +++ b/fli/cli/commands/dates.py @@ -80,6 +80,38 @@ def dates( help="Trip duration in days", ), ] = 3, + passengers: Annotated[ + int, + typer.Option( + "--passengers", + help="Number of adult passengers", + min=1, + ), + ] = 1, + children: Annotated[ + int, + typer.Option( + "--children", + help="Number of children", + min=0, + ), + ] = 0, + infants_in_seat: Annotated[ + int, + typer.Option( + "--infants-in-seat", + help="Number of infants in seat", + min=0, + ), + ] = 0, + infants_on_lap: Annotated[ + int, + typer.Option( + "--infants-on-lap", + help="Number of infants on lap", + min=0, + ), + ] = 0, airlines: Annotated[ list[str] | None, typer.Option( @@ -226,6 +258,10 @@ def dates( "end_date": end_date, "trip_duration": trip_duration, "is_round_trip": is_round_trip, + "passengers": passengers, + "children": children, + "infants_in_seat": infants_in_seat, + "infants_on_lap": infants_on_lap, "cabin_class": seat_type.name, "max_stops": stops.name, "departure_window": ( @@ -260,7 +296,12 @@ def dates( # Create search filters filters = DateSearchFilters( trip_type=trip_type, - passenger_info=PassengerInfo(adults=1), + passenger_info=PassengerInfo( + adults=passengers, + children=children, + infants_in_seat=infants_in_seat, + infants_on_lap=infants_on_lap, + ), flight_segments=segments, stops=stops, seat_type=seat_type, @@ -320,6 +361,10 @@ def dates( "end_date": end_date, "trip_duration": trip_duration, "is_round_trip": is_round_trip, + "passengers": passengers, + "children": children, + "infants_in_seat": infants_in_seat, + "infants_on_lap": infants_on_lap, "cabin_class": cabin_class, "max_stops": max_stops, "departure_window": ( diff --git a/fli/cli/commands/flights.py b/fli/cli/commands/flights.py index ca87f84a..5bfa8490 100644 --- a/fli/cli/commands/flights.py +++ b/fli/cli/commands/flights.py @@ -36,6 +36,10 @@ def _search_flights_core( departure_date: str, return_date: str | None = None, departure_window: str | tuple[int, int] | None = None, + passengers: int = 1, + children: int = 0, + infants_in_seat: int = 0, + infants_on_lap: int = 0, airlines: list[str] | None = None, cabin_class: str = "ECONOMY", max_stops: str = "ANY", @@ -50,6 +54,10 @@ def _search_flights_core( "departure_date": departure_date, "return_date": return_date, "departure_window": None, + "passengers": passengers, + "children": children, + "infants_in_seat": infants_in_seat, + "infants_on_lap": infants_on_lap, "airlines": [airline.upper() for airline in airlines] if airlines else None, "cabin_class": cabin_class.upper(), "max_stops": max_stops.upper(), @@ -96,7 +104,12 @@ def _search_flights_core( # Create search filters filters = FlightSearchFilters( trip_type=trip_type, - passenger_info=PassengerInfo(adults=1), + passenger_info=PassengerInfo( + adults=passengers, + children=children, + infants_in_seat=infants_in_seat, + infants_on_lap=infants_on_lap, + ), flight_segments=segments, stops=stops, seat_type=seat_type, @@ -188,6 +201,38 @@ def flights( help="Departure time window in 24h format (e.g., 6-20)", ), ] = None, + passengers: Annotated[ + int, + typer.Option( + "--passengers", + help="Number of adult passengers", + min=1, + ), + ] = 1, + children: Annotated[ + int, + typer.Option( + "--children", + help="Number of children", + min=0, + ), + ] = 0, + infants_in_seat: Annotated[ + int, + typer.Option( + "--infants-in-seat", + help="Number of infants in seat", + min=0, + ), + ] = 0, + infants_on_lap: Annotated[ + int, + typer.Option( + "--infants-on-lap", + help="Number of infants on lap", + min=0, + ), + ] = 0, airlines: Annotated[ list[str] | None, typer.Option( @@ -251,6 +296,10 @@ def flights( departure_date=departure_date, return_date=return_date, departure_window=departure_window, + passengers=passengers, + children=children, + infants_in_seat=infants_in_seat, + infants_on_lap=infants_on_lap, airlines=airlines, cabin_class=cabin_class, max_stops=max_stops, diff --git a/fli/mcp/server.py b/fli/mcp/server.py index 5f449d03..b1a533fc 100755 --- a/fli/mcp/server.py +++ b/fli/mcp/server.py @@ -240,6 +240,9 @@ class FlightSearchParams(BaseModel): ge=1, description="Number of adult passengers", ) + children: int = Field(0, ge=0, description="Number of children") + infants_in_seat: int = Field(0, ge=0, description="Number of infants in seat") + infants_on_lap: int = Field(0, ge=0, description="Number of infants on lap") exclude_basic_economy: bool = Field( False, description="Exclude basic economy fares from results" ) @@ -275,6 +278,9 @@ class DateSearchParams(BaseModel): ge=1, description="Number of adult passengers", ) + children: int = Field(0, ge=0, description="Number of children") + infants_in_seat: int = Field(0, ge=0, description="Number of infants in seat") + infants_on_lap: int = Field(0, ge=0, description="Number of infants on lap") # ============================================================================= @@ -359,7 +365,12 @@ def _execute_flight_search(params: FlightSearchParams) -> dict[str, Any]: # Create search 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, @@ -425,7 +436,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, @@ -516,6 +532,18 @@ def search_flights( int | None, Field(description="Number of adult passengers", ge=1), ] = None, + children: Annotated[ + int, + Field(description="Number of children", ge=0), + ] = 0, + infants_in_seat: Annotated[ + int, + Field(description="Number of infants in seat", ge=0), + ] = 0, + infants_on_lap: Annotated[ + int, + Field(description="Number of infants on lap", ge=0), + ] = 0, exclude_basic_economy: Annotated[ bool, Field(description="Exclude basic economy fares from results"), @@ -538,6 +566,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, ) return _execute_flight_search(params) @@ -595,6 +626,18 @@ def search_dates( int | None, Field(description="Number of adult passengers", ge=1), ] = None, + children: Annotated[ + int, + Field(description="Number of children", ge=0), + ] = 0, + infants_in_seat: Annotated[ + int, + Field(description="Number of infants in seat", ge=0), + ] = 0, + infants_on_lap: Annotated[ + int, + Field(description="Number of infants on lap", ge=0), + ] = 0, ) -> dict[str, Any]: """Find the cheapest travel dates between two airports within a date range. @@ -615,6 +658,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, ) return _execute_date_search(params) diff --git a/tests/cli/test_dates.py b/tests/cli/test_dates.py index 22ba94d9..3020aa9d 100644 --- a/tests/cli/test_dates.py +++ b/tests/cli/test_dates.py @@ -135,6 +135,41 @@ def test_dates_with_time(runner, mock_search_dates, mock_console): mock_search_dates.search.assert_called_once() +def test_dates_with_family_passenger_counts(runner, mock_search_dates, mock_console): + """Test dates search passes family passenger counts into filters.""" + mock_search_dates.search.return_value = [ + DatePrice( + date=(datetime.now() + timedelta(days=1),), + price=299.99, + ), + ] + + result = runner.invoke( + app, + [ + "dates", + "JFK", + "LAX", + "--passengers", + "2", + "--children", + "2", + "--infants-in-seat", + "1", + "--infants-on-lap", + "1", + ], + ) + + assert result.exit_code == 0 + mock_search_dates.search.assert_called_once() + args, _ = mock_search_dates.search.call_args + assert args[0].passenger_info.adults == 2 + assert args[0].passenger_info.children == 2 + assert args[0].passenger_info.infants_in_seat == 1 + assert args[0].passenger_info.infants_on_lap == 1 + + def test_dates_with_sort(runner, mock_search_dates, mock_console): """Test dates search with sort option.""" mock_search_dates.search.return_value = [ diff --git a/tests/cli/test_flights.py b/tests/cli/test_flights.py index 3418a9c5..9c2fcfce 100644 --- a/tests/cli/test_flights.py +++ b/tests/cli/test_flights.py @@ -94,6 +94,35 @@ def test_flights_with_stops(runner, mock_search_flights, mock_console): mock_search_flights.search.assert_called_once() +def test_flights_with_family_passenger_counts(runner, mock_search_flights, mock_console): + """Test flights search passes family passenger counts into filters.""" + result = runner.invoke( + app, + [ + "flights", + "JFK", + "LAX", + datetime.now().strftime("%Y-%m-%d"), + "--passengers", + "2", + "--children", + "2", + "--infants-in-seat", + "1", + "--infants-on-lap", + "1", + ], + ) + + assert result.exit_code == 0 + mock_search_flights.search.assert_called_once() + args, _ = mock_search_flights.search.call_args + assert args[0].passenger_info.adults == 2 + assert args[0].passenger_info.children == 2 + assert args[0].passenger_info.infants_in_seat == 1 + assert args[0].passenger_info.infants_on_lap == 1 + + def test_flights_invalid_airport(runner, mock_search_flights, mock_console): """Test flights search with invalid airport code.""" result = runner.invoke( diff --git a/tests/mcp/test_mcp_server.py b/tests/mcp/test_mcp_server.py index d0fe5555..29edcaa7 100644 --- a/tests/mcp/test_mcp_server.py +++ b/tests/mcp/test_mcp_server.py @@ -1,6 +1,7 @@ """Test MCP server functionality.""" from datetime import datetime, timedelta +from unittest.mock import MagicMock from fli.mcp.server import ( DateSearchParams, @@ -211,6 +212,53 @@ def test_invalid_airline_code(self): assert "airline" in result["error"].lower() assert result["flights"] == [] + def test_search_flights_passes_family_passenger_counts(self, monkeypatch): + """Test flight search tool passes family passenger counts into filters.""" + mock_client = MagicMock() + mock_client.search.return_value = [] + monkeypatch.setattr("fli.mcp.server.SearchFlights", lambda: mock_client) + + result = search_flights( + origin="JFK", + destination="LHR", + departure_date=get_future_date(30), + passengers=2, + children=1, + infants_in_seat=1, + infants_on_lap=1, + ) + + assert result["success"] is True + args, _ = mock_client.search.call_args + assert args[0].passenger_info.adults == 2 + assert args[0].passenger_info.children == 1 + assert args[0].passenger_info.infants_in_seat == 1 + assert args[0].passenger_info.infants_on_lap == 1 + + def test_search_dates_passes_family_passenger_counts(self, monkeypatch): + """Test date search tool passes family passenger counts into filters.""" + mock_client = MagicMock() + mock_client.search.return_value = [] + monkeypatch.setattr("fli.mcp.server.SearchDates", lambda: mock_client) + + result = search_dates( + origin="JFK", + destination="LHR", + start_date=get_future_date(30), + end_date=get_future_date(60), + passengers=2, + children=1, + infants_in_seat=1, + infants_on_lap=1, + ) + + assert result["success"] is True + args, _ = mock_client.search.call_args + assert args[0].passenger_info.adults == 2 + assert args[0].passenger_info.children == 1 + assert args[0].passenger_info.infants_in_seat == 1 + assert args[0].passenger_info.infants_on_lap == 1 + def test_flight_search_params_validation(self): """Test FlightSearchParams validation.""" future_date = get_future_date(30) @@ -221,6 +269,9 @@ def test_flight_search_params_validation(self): assert params.cabin_class == "ECONOMY" # default assert params.max_stops == "ANY" # default assert params.sort_by == "CHEAPEST" # default + assert params.children == 0 + assert params.infants_in_seat == 0 + assert params.infants_on_lap == 0 def test_date_search_params_validation(self): """Test DateSearchParams validation.""" @@ -241,3 +292,6 @@ def test_date_search_params_validation(self): assert params.cabin_class == "ECONOMY" # default assert params.max_stops == "ANY" # default assert params.sort_by_price is False # default + assert params.children == 0 + assert params.infants_in_seat == 0 + assert params.infants_on_lap == 0 diff --git a/tests/models/test_date_search_filters.py b/tests/models/test_date_search_filters.py index 73922512..4c7ad111 100644 --- a/tests/models/test_date_search_filters.py +++ b/tests/models/test_date_search_filters.py @@ -175,7 +175,7 @@ def get_future_date(days: int = 30) -> str: "name": "Test 3: Flight Search Data", "search": DateSearchFilters( passenger_info=PassengerInfo( - adults=2, + adults=3, children=3, infants_in_seat=0, infants_on_lap=1, @@ -215,7 +215,7 @@ def get_future_date(days: int = 30) -> str: None, [], 1, - [2, 3, 1, 0], + [3, 3, 1, 0], [None, 900], None, None, diff --git a/tests/models/test_flight_search_filters.py b/tests/models/test_flight_search_filters.py index 2a57e40b..501d382d 100644 --- a/tests/models/test_flight_search_filters.py +++ b/tests/models/test_flight_search_filters.py @@ -190,7 +190,7 @@ def get_future_date(days: int = 30) -> str: "name": "Test 3: Flight Search Data", "search": FlightSearchFilters( passenger_info=PassengerInfo( - adults=2, + adults=3, children=3, infants_in_seat=0, infants_on_lap=1, @@ -228,7 +228,7 @@ def get_future_date(days: int = 30) -> str: None, [], 1, - [2, 3, 1, 0], + [3, 3, 1, 0], [None, 900], None, None,