find_flights.py

# run SQLDatabase chain
def find_flights(query, llm, db):
    """
    Executes a search for flights using a language model and a SQL database toolkit.
    
    Parameters:
    query (str): The query to be executed, typically a natural language description of the flights to find.
    llm (LanguageModel): The language model used to process the query and generate SQL commands.
    db (Database): The database object where the flight data is stored and from which data will be retrieved.
    
    Returns:
    Response: The response from the agent executor's run method, typically containing the search results or an error message.
    """
        
    toolkit = SQLDatabaseToolkit(db=db, llm=llm)
    agent_executor = create_sql_agent(
        llm=llm,
        toolkit=toolkit,
        verbose=True,
        agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    )
    
    return agent_executor.run(query)