Integrate AviationStack API with Langchain's Agent Tools

[Jacky3003]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

We would like to integrate Aviationstack's API into LangChain as a tool to query information about different airlines. The idea here is to chain this tool with other related tools in LangChain such as Google Places and Amadeus by querying via airlines to connect information together. Additionally, a toolkit will be made to help give the agent an action to choose based on the prompt, similar to many other toolkits in LangChain, an example being Steam’s API. For now, this implementation will cover only a limited amount of use cases with the option to expand on more features later.

A potential example can be seen below:

from langchain_core.prompts import PromptTemplate
from langchain.agents import create_react_agent
from langchain_openai import OpenAI
from langchain_community.agent_toolkits.aviationstack.toolkit import AviationstackToolkit
from langchain_commuity.utilities.aviationstack import AviationstackAPIWrapper

template = """Find any flight to book where the name of the airline is Philippine Airlines."""
prompt_to_use = PromptTemplate.from_template(template)

chosen_llm = OpenAI()

aviationstack_wrapper = AviationstackAPIWrapper()
tools_to_use = AviationstackToolkit.from_aviationstack_api_wrapper(aviationstack_wrapper)

agent = create_react_agent(
    llm=chosen_llm,
    tools=tools_to_use,
    prompt=prompt_to_use
)

result = agent.invoke("Find a flight to book for me please.")
print(result)

And the output could be the following (TBD):

Here are flights that can be booked under Phillipine Airlines:
1. Flight date is 2024-10-15 and status is scheduled
- Departs at terminal 1 from Ninoy Aquino International
- Arrives at terminal 1 at Seoul (Incheon)
- Flight iata is PR466
...

Motivation

LangChain already offers flight data in the form of Amadeus’ API, which includes information such as seats, bookings, hotels, transfers, etc. Aviationstack is able to fill in gaps with more tools related to aviation intelligence, such as real time and historical flight data, as well as information about aircrafts, allowings for more aviation analysis amongst different airlines. Overall, this would help expand the catalog for the aviation space in LangChain.

Proposal (If applicable)

This implementation will be limited to the endpoints that are only allowed on the free plan of Aviationstack. There are other endpoints that are also beyond the scope of what this feature tries to achieve, these endpoints can serve as a stretch goal. The following endpoints as a starter will be implemented: Flights and Airlines. Note that depending on the complexity of the implementation, some of these endpoints may have to be omitted as well.

The following files will need to be added:

- libs/community/langchain_community/tools/aviationstack/__init__.py
- libs/community/langchain_community/tools/aviationstack/prompt.py
- libs/community/langchain_community/tools/aviationstack/tool.py
- libs/community/langchain_community/utilities/aviationstack.py
- libs/community/tests/integration_tests/utilities/test_aviationstack_api.py
- libs/community/langchain_community/agent_toolkits/aviationstack/__init__.py
- libs/community/langchain_community/agent_toolkits/aviationstack/toolkit.py

The following files will need to be modified:

- libs/community/tests/unit_tests/tools/test_imports.py
- libs/community/langchain_community/tools/__init__.py
- libs/langchain/langchain/tools/__init__.py
- libs/community/langchain_community/agent_toolkits/__init__.py
- libs/langchain/langchain/agents/agent_toolkits/__init__.py
- libs/community/langchain_community/utilities/__init__.py
- libs/community/langchain_community/agent_toolkits/load_tools.py

Sample code for tool.py:

"""Tool for Aviationstack API"""

from langchain.tools import BaseTool
from typing import Optional
from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_community.utilities.aviationstack import AviationstackAPIWrapper

class AviationstackQueryRun(BaseTool):

    api_wrapper: AviationstackAPIWrapper
    mode: str
    name: str = "aviationstack"
    description: str = (
        "Wrapper for Aviationstack API"
        "Used to get information about the aviation industry."
        "Input should be an airline name."
    )

    def _run(
        self, 
        query: str, 
        mode: str,
        run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        """Uses the Aviationstack API tool."""
        return self.api_wrapper.run(self.mode, query)

Sample code for the utility wrapper:

import requests

from langchain_core.utils import get_from_dict_or_env
from typing import Dict, Optional, Any
from pydantic import BaseModel, ConfigDict, model_validator

AVIATIONSTACK_API_URLS = (
    "https://api.aviationstack.com/v1/airlines",
    "https://api.aviationstack.com/v1/flights"
)
AVIATIONSTACK_TIMEOUT = 5000

class AviationstackAPIWrapper(BaseModel):
    """Wrapper for Aviationstack's API

    Follow the instructions below for how to use the API:

    1. Head to https://aviationstack.com/, sign up for an
    account, and get an API key from your dashboard.
    2. (Optional) you can upgrade you API plan here: 
    https://aviationstack.com/product if you would like to
    increase the pagination limit past 100.
    3. Save your API key to an environment variable called
    AVIATIONSTACK_API_KEY, and optionally set an environment
    variable called AVIATIONSTACK_LIMIT to the pagination
    limit you want (be sure that your plan is basic or above
    if you set it above 100), defaults to 100.

    """

    aviationstack_api_key: Optional[str] = None
    aviationstack_pagination_limit = 100

    model_config = ConfigDict(
        extra="forbid",
    )

    @model_validator(mode="before")
    @classmethod
    def validate_environment(cls, values: Dict) -> Any:
        #TODO
        """Validate that the api key exists, and set the limit if needed"""

        aviationstack_api_key = get_from_dict_or_env(
            values, "AVIATIONSTACK_API_KEY", "aviationstack_api_key"
        )
        aviationstack_limit = get_from_dict_or_env(
            values, "AVIATIONSTACK_LIMIT", "aviationstack_limit"
        )

        values["aviationstack_api_key"] = aviationstack_api_key
        values["aviationstack_limit"] = aviationstack_limit

        return values

    def run(self, mode: str, query:str) -> str:
        """Call the API for Aviationstack depending on the mode."""
        #TODO
        response = None
        # Handle API request based on the mode.
        if mode == "get_airline":
            response = requests.get(AVIATIONSTACK_API_URLS[0])
        if mode == "get_flights":
            response = requests.get(AVIATIONSTACK_API_URLS[1])
        return self.format_response(response)

    def format_response(self, response: requests.Response):
        #TODO
        # Process the response in a readable format.
        pass

Sample toolkit:

from typing import List

from langchain_core.tools.base import BaseToolkit
from langchain_core.tools import BaseTool

from langchain_community.tools.aviationstack.tool import AviationstackQueryRun
from langchain_community.utilities.aviationstack import AviationstackAPIWrapper

# More prompts to be added as needed
from langchain_community.tools.aviationstack.prompt import (
    AVIATIONSTACK_GET_AIRLINE_INFORMATION
)

class AviationstackToolkit(BaseToolkit):

    tools: List[BaseTool] = []

    @classmethod
    def create_toolkit_from_api_wrapper(
        cls, aviationstack_wrapper: AviationstackAPIWrapper
    ):
        # Add operations as needed per new prompt
        operations: List[dict] = [
            {
                "mode": "get_airline",
                "name": "Get Airline Information",
                "description": AVIATIONSTACK_GET_AIRLINE_INFORMATION
            }
        ]
        tools = [
                AviationstackQueryRun(
                name=operation["name"],
                mode=operation["mode"],
                description=operation["description"],
                api_wrapper=AviationstackAPIWrapper
            ) for operation in operations
        ]
        return cls(tools=tools)

Sample Prompt:

AVIATIONSTACK_GET_AIRLINE_INFORMATION ="""
    This tool is a wrapper for Aviationstack's API used for getting
    information about a specific airline. The input to this tool
    will be the name of an airline. As an example, if you want to
    search for the airline "American Airlines", you would pass the
    name "American Airlines" as input. The input is then used to
    filter the response from Aviationstack's airlines API here:
    https://api.aviationstack.com/v1/airlines in order to give
    information regarding the inputted airline name.
"""

Next Steps

• Finalize extra prompts where needed in prompt.py.
• Add docstrings to relevant methods and classes where needed.
• Implement code for testing framework.
• Create custom response formatters for both the flights and airlines endpoint.

[rd4cake]

.