get_args.py

import json
import openai

def get_args(query_user, OPENAI_KEY):
    """
    Extracts necessary parameters for flight search from user query using OpenAI API.
    
    This function takes a user query and an OpenAI key as inputs, sends the user query to 
    OpenAI API to extract necessary parameters including the number of adults, departure date,
    return date, destination location code, and origin location code for flight search.

    Parameters:
    query_user (str): User query to be sent to OpenAI API.
    openai_key (str): OpenAI key to authenticate with the API.

    Returns:
    num_adults (int): Number of adults for the flight.
    departureDate (str): Departure date in the format YYYY-MM-DD.
    returnDate (str): Return date in the format YYYY-MM-DD.
    destinationLocationCode (str): IATA code for the destination location.
    originLocationCode (str): IATA code for the origin location.
    TypeofflightReuqest (str): Returns the type of flight e,g cheapest, fastest etc.
    """

    function_call = [
    {
      "name": "search_for_flights",
      "description": "Requests flight data from Amadeus API and writes to SQLite database",
      "parameters": {
        "type": "object",
        "properties": {
            "num_adults":{
                "type":"integer",
                "description": '''Based on the query, respond with the number of adults'''
            },
            "departureDate": {
                "type":"string",
                "description": '''Based on the query, respond with the Departure Date. Dates are specified in the ISO 8601 YYYY-MM-DD format. '''
            },
            "returnDate": {
                "type":"string",
                "description": '''Based on the query, respond with the Return Date. Dates are specified in the ISO 8601 YYYY-MM-DD format. '''
            },
            "destinationLocationCode":{
                "type":"string",
                "description": '''Based on the query, respond with an airport IATA code from the city which the traveler is going. E.g CDG for Charles de Gaulle Airport'''
            },
          "originLocationCode": {
            "type": "string",
            "description": '''Based on the query, respond with an airport IATA code from the city which the traveler will depart from. E.g CDG for Charles de Gaulle Airport'''
          },

         "TypeofflightReuqest": {
            "type": "string",
            "description": '''Based on the query, respond with the type of flight the user is requesting E.g cheapest, shortest, fastest, least stops etc.'''
          },

        },
        "required": ["destinationLocationCode", "originLocationCode", "departureDate", "returnDate", "num_adults", "TypeofflightReuqest"]
      }
    }
    ]
    
    openai.api_key = OPENAI_KEY

    message = openai.ChatCompletion.create(
        model="gpt-4-0613",
        messages=[{"role": "user", "content": query_user}],
        functions = function_call,
        function_call = 'auto',
        temperature=0
    )
    response_message = message["choices"][0]["message"]["function_call"]["arguments"]

    parsed_data = json.loads(response_message)

    # Accessing variables
    num_adults = parsed_data['num_adults']
    departureDate = parsed_data['departureDate']
    returnDate = parsed_data['returnDate']
    destinationLocationCode = parsed_data['destinationLocationCode']
    originLocationCode = parsed_data['originLocationCode']
    TypeofflightReuqest = parsed_data['TypeofflightReuqest']

    
    print("Number of Adults: ", num_adults)
    print("Departure Date: ", departureDate)
    print("Return Date: ", returnDate)
    print("Destination Location Code: ", destinationLocationCode)
    print("Origin Location Code: ", originLocationCode)
    print("Origin Location Code: ", TypeofflightReuqest)


    return num_adults, departureDate, returnDate, destinationLocationCode, originLocationCode, TypeofflightReuqest