tools/flight_status.py

from typing import Any, Optional
from smolagents.tools import Tool
import requests

class FlightStatus(Tool):
    name = "flight_status"
    description = "Fetches the status of a given flight based on the flight number using AviationStack API."
    inputs = {'flight': {'type': 'string', 'description': 'Flight Code to fetch the information.'}}
    output_type = "any"

    def forward(self,flight: str) -> dict:
        """
        Fetches the status of a flight by querying the AviationStack API.
        Returns a dictionary containing flight status or error details.
        """
        #flight_number = self.inputs.get('flight')
        flight_number = flight
        api_key = "6ffd5cb6e366ed0cf24a1caa38b17c31"

        # Ensure flight number and API key are provided
        if not flight_number or not api_key:
            return {'error': 'Missing flight number or API key'}

        base_url = "http://api.aviationstack.com/v1/flights"
        
        # Parameters for the API request
        params = {
            'access_key': api_key,
            'flight_iata': flight_number  # Flight number in IATA format (e.g., "AA100")
        }

        try:
            # Send a GET request to the AviationStack API
            response = requests.get(base_url, params=params)
            
            #Check if the request was successful
            if response.status_code == 200:
                data = response.json()

                # If no flight data is found
                if not data.get("data"):
                    return {'error': f"No information available for flight {flight_number}."}

                flight_data = data['data'][0]  # Get the first flight result

                # Extract relevant flight information
                flight_info = {
                    'flight_number': flight_data['flight']['iata'],
                    'airline': flight_data['airline']['name'],
                    'status': flight_data['flight_status'],
                    'departure': {
                        'airport': flight_data['departure']['airport'],
                        'estimated': flight_data['departure'].get('estimated', 'N/A'),
                        'actual': flight_data['departure'].get('estimated_runway', 'N/A'),
                    },
                    'arrival': {
                        'airport': flight_data['arrival']['airport'],
                        'estimated': flight_data['arrival'].get('estimated', 'N/A'),
                        'actual': flight_data['arrival'].get('estimated_runway', 'N/A'),
                    },
                }

                return flight_info
            
            else:
                return {'error': f"Error fetching data for {flight_number}. Status code: {response.status_code}"}
        
        except Exception as e:
            return {'error': f"An error occurred: {str(e)}"}

    def __init__(self, *args, **kwargs):
        self.is_initialized = False  # Marking the tool as not initialized