MCP Server for Elicitation

mcp_server.py

# server.py
from fastmcp import FastMCP, Context
from dataclasses import dataclass
import httpx
import os
from datetime import datetime, timedelta
import pytz

mcp = FastMCP("Flight Status Server")

@dataclass
class FlightInfo:
    flight_number: str

@mcp.tool
async def get_flight_status(ctx: Context) -> str:
    """Get live flight status information by collecting flight number interactively."""
    
    # Elicit flight number from user
    result = await ctx.elicit(
        message="Please provide the flight number you want to check",
        response_type=FlightInfo
    )
    
    if result.action == "decline":
        return "Flight number not provided"
    elif result.action == "cancel":
        return "Operation cancelled"
    elif result.action != "accept":
        return "Invalid response"
    
    flight_number = result.data.flight_number
    
    try:
        # Fetch flight data
        flight_details = await _fetch_flight_status(flight_number)
        
        # Format response
        return f"""
Flight Status for {flight_details['flight']}:
• Route: {flight_details['source']} → {flight_details['destination']}
• Departure: {flight_details['depart_time']}
• Arrival: {flight_details['arrival_time']}
• Status: {flight_details['status']}
        """.strip()
        
    except Exception as e:
        return f"Error fetching flight status: {str(e)}"

async def _fetch_flight_status(flight: str) -> dict:
    """Internal function to fetch flight status from FlightAware API."""
    AEROAPI_BASE_URL = "https://aeroapi.flightaware.com/aeroapi"
    AEROAPI_KEY = os.getenv("AERO_API_KEY")
    
    if not AEROAPI_KEY:
        raise ValueError("AERO_API_KEY is not set in environment variables")
    
    def _clean_flight_id(flight_id):
        if "flight_id=" in flight_id:
            return flight_id.split("flight_id=")[1]
        return flight_id
    
    def _utc_to_local(utc_date_str, local_timezone_str):
        utc_datetime = datetime.strptime(utc_date_str, '%Y-%m-%dT%H:%M:%SZ').replace(tzinfo=pytz.utc)
        local_timezone = pytz.timezone(local_timezone_str)
        local_datetime = utc_datetime.astimezone(local_timezone)
        return local_datetime.strftime('%Y-%m-%d %H:%M:%S')
    
    # Prepare API request
    flight_id = _clean_flight_id(flight)
    start_date = datetime.now().date().strftime('%Y-%m-%d')
    end_date = (datetime.now().date() + timedelta(days=1)).strftime('%Y-%m-%d')
    api_resource = f"/flights/{flight_id}?start={start_date}&end={end_date}"
    
    # Make API call
    async with httpx.AsyncClient(headers={"x-apikey": AEROAPI_KEY}) as client:
        response = await client.get(f"{AEROAPI_BASE_URL}{api_resource}")
        response.raise_for_status()
        flight_data = response.json()['flights'][0]
    
    # Determine best available times
    dep_key = ('estimated_out' if flight_data.get('estimated_out') else
               'actual_out' if flight_data.get('actual_out') else
               'scheduled_out')
    
    arr_key = ('estimated_in' if flight_data.get('estimated_in') else
               'actual_in' if flight_data.get('actual_in') else
               'scheduled_in')
    
    # Build response
    flight_details = {
        'flight': flight,
        'source': flight_data['origin']['city'],
        'destination': flight_data['destination']['city'],
        'depart_time': _utc_to_local(flight_data[dep_key], flight_data['origin']['timezone']),
        'arrival_time': _utc_to_local(flight_data[arr_key], flight_data['destination']['timezone']),
        'status': flight_data['status']
    }
    
    return flight_details

if __name__ == "__main__":
    mcp.run(transport="streamable-http", host="127.0.0.1", port=8080, path="/mcp")