mirror of
https://github.com/hwchase17/langchain
synced 2024-11-04 06:00:26 +00:00
4eda647fdd
Previously, if this did not find a mypy cache then it wouldnt run this makes it always run adding mypy ignore comments with existing uncaught issues to unblock other prs --------- Co-authored-by: Erick Friis <erick@langchain.dev> Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com>
62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
from typing import Any, Dict, Optional, Type
|
|
|
|
from langchain_core.callbacks import CallbackManagerForToolRun
|
|
from langchain_core.language_models import BaseLanguageModel
|
|
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
|
|
|
|
from langchain_community.chat_models import ChatOpenAI
|
|
from langchain_community.tools.amadeus.base import AmadeusBaseTool
|
|
|
|
|
|
class ClosestAirportSchema(BaseModel):
|
|
"""Schema for the AmadeusClosestAirport tool."""
|
|
|
|
location: str = Field(
|
|
description=(
|
|
" The location for which you would like to find the nearest airport "
|
|
" along with optional details such as country, state, region, or "
|
|
" province, allowing for easy processing and identification of "
|
|
" the closest airport. Examples of the format are the following:\n"
|
|
" Cali, Colombia\n "
|
|
" Lincoln, Nebraska, United States\n"
|
|
" New York, United States\n"
|
|
" Sydney, New South Wales, Australia\n"
|
|
" Rome, Lazio, Italy\n"
|
|
" Toronto, Ontario, Canada\n"
|
|
)
|
|
)
|
|
|
|
|
|
class AmadeusClosestAirport(AmadeusBaseTool):
|
|
"""Tool for finding the closest airport to a particular location."""
|
|
|
|
name: str = "closest_airport"
|
|
description: str = (
|
|
"Use this tool to find the closest airport to a particular location."
|
|
)
|
|
args_schema: Type[ClosestAirportSchema] = ClosestAirportSchema
|
|
|
|
llm: Optional[BaseLanguageModel] = Field(default=None)
|
|
"""Tool's llm used for calculating the closest airport. Defaults to `ChatOpenAI`."""
|
|
|
|
@root_validator(pre=True)
|
|
def set_llm(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
|
if not values.get("llm"):
|
|
# For backward-compatibility
|
|
values["llm"] = ChatOpenAI(temperature=0)
|
|
return values
|
|
|
|
def _run(
|
|
self,
|
|
location: str,
|
|
run_manager: Optional[CallbackManagerForToolRun] = None,
|
|
) -> str:
|
|
content = (
|
|
f" What is the nearest airport to {location}? Please respond with the "
|
|
" airport's International Air Transport Association (IATA) Location "
|
|
' Identifier in the following JSON format. JSON: "iataCode": "IATA '
|
|
' Location Identifier" '
|
|
)
|
|
|
|
return self.llm.invoke(content) # type: ignore[union-attr]
|