flight.py

from typing import Any

import httpx
from mcp.server.fastmcp import FastMCP
from dotenv import load_dotenv
import os

load_dotenv()
SERPAPI_KEY = os.getenv("SERPAPI_KEY")

mcp = FastMCP("Flight Search Tool", log_level="ERROR")

async def fetch_serp_data(params: dict[str, Any]) -> dict[str, Any] | None:
    """Make a request to the Serp API and return the response data."""
    async with httpx.AsyncClient(timeout=30.0) as client:
        try:
            
            params["api_key"] = SERPAPI_KEY
            params["engine"] = "google_flights"
            response = await client.get("https://serpapi.com/search.json", params=params)
            response.raise_for_status()
            return response.json()
        except Exception as error:
            return {"error": str(error)}
        
def format_flight_entry(flight_data: dict) -> str:
    """Extracts relevant info from a single flight result."""
    # SerpApi nests flight legs inside a 'flights' list
    legs = flight_data.get("flights", [])
    if not legs:
        return "Flight details unavailable."
    
    first_leg = legs[0]
    last_leg = legs[-1]
    
    airline = first_leg.get("airline")
    price = flight_data.get("price", "N/A")
    duration = flight_data.get("total_duration", "Unknown")
    
    departure = f"{first_leg['departure_airport']['name']} ({first_leg['departure_airport']['time']})"
    arrival = f"{last_leg['arrival_airport']['name']} ({last_leg['arrival_airport']['time']})"
    
    return f"Airline: {airline} | {price} | Duration: {duration} mins\n   From: {departure}\n   To: {arrival}\n"
        
@mcp.tool()
async def search_flights(origin: str, destination: str, outbound_date: str, return_date: str = None) -> str:
    """
    Search for flights between two airports using IATA codes.
    
    Args:
        origin: 3-letter IATA code (e.g., 'JFK', 'SFO')
        destination: 3-letter IATA code (e.g., 'LHR', 'CDG')
        outbound_date: Departure date in YYYY-MM-DD format
        return_date: Optional return date in YYYY-MM-DD format
    """
    params = {
        "departure_id": origin,
        "arrival_id": destination,
        "outbound_date": outbound_date
    }

    if return_date:
        params["return_date"] = return_date
        params["type"] = "1"
    else:
        params["type"] = "2"
    
    data = await fetch_serp_data(params)
    
    if not data or "error" in data:
        return f"Error fetching flight data. {data.get('error', 'Unknown error')}"
    
    results = data.get("best_flights", [])
    if not results:
        return "No flights found for the given criteria."
    
    #Top 5 results formatted
    formatted_results = [format_flight_entry(flight) for flight in results[:5]]
    
    return "\n---\n".join(formatted_results)

def main():
    mcp.run(transport="stdio")

if __name__ == "__main__":
    main()