langchain/libs/community/langchain_community/agent_toolkits/nasa/toolkit.py
Eugene Yurtsev 25fbe356b4
community[patch]: upgrade to recent version of mypy (#21616)
This PR upgrades community to a recent version of mypy. It inserts type:
ignore on all existing failures.
2024-05-13 14:55:07 -04:00

59 lines
1.9 KiB
Python

from typing import Dict, List
from langchain_core.tools import BaseToolkit
from langchain_community.tools import BaseTool
from langchain_community.tools.nasa.prompt import (
NASA_CAPTIONS_PROMPT,
NASA_MANIFEST_PROMPT,
NASA_METADATA_PROMPT,
NASA_SEARCH_PROMPT,
)
from langchain_community.tools.nasa.tool import NasaAction
from langchain_community.utilities.nasa import NasaAPIWrapper
class NasaToolkit(BaseToolkit):
"""Nasa Toolkit."""
tools: List[BaseTool] = []
@classmethod
def from_nasa_api_wrapper(cls, nasa_api_wrapper: NasaAPIWrapper) -> "NasaToolkit":
operations: List[Dict] = [
{
"mode": "search_media",
"name": "Search NASA Image and Video Library media",
"description": NASA_SEARCH_PROMPT,
},
{
"mode": "get_media_metadata_manifest",
"name": "Get NASA Image and Video Library media metadata manifest",
"description": NASA_MANIFEST_PROMPT,
},
{
"mode": "get_media_metadata_location",
"name": "Get NASA Image and Video Library media metadata location",
"description": NASA_METADATA_PROMPT,
},
{
"mode": "get_video_captions_location",
"name": "Get NASA Image and Video Library video captions location",
"description": NASA_CAPTIONS_PROMPT,
},
]
tools = [
NasaAction(
name=action["name"],
description=action["description"],
mode=action["mode"],
api_wrapper=nasa_api_wrapper,
)
for action in operations
]
return cls(tools=tools) # type: ignore[arg-type]
def get_tools(self) -> List[BaseTool]:
"""Get the tools in the toolkit."""
return self.tools