langchain/libs/community/langchain_community/agent_toolkits/nla/toolkit.py

from __future__ import annotations

from typing import Any, List, Optional, Sequence

from langchain_core.language_models import BaseLanguageModel
from langchain_core.tools import BaseTool
from langchain_core.tools.base import BaseToolkit
from pydantic import Field

from langchain_community.agent_toolkits.nla.tool import NLATool
from langchain_community.tools.openapi.utils.openapi_utils import OpenAPISpec
from langchain_community.tools.plugin import AIPlugin
from langchain_community.utilities.requests import Requests


class NLAToolkit(BaseToolkit):
    """Natural Language API Toolkit.

    *Security Note*: This toolkit creates tools that enable making calls
        to an Open API compliant API.

        The tools created by this toolkit may be able to make GET, POST,
        PATCH, PUT, DELETE requests to any of the exposed endpoints on
        the API.

        Control access to who can use this toolkit.

        See https://python.langchain.com/docs/security for more information.
    """

    nla_tools: Sequence[NLATool] = Field(...)
    """List of API Endpoint Tools."""

    def get_tools(self) -> List[BaseTool]:
        """Get the tools for all the API operations."""
        return list(self.nla_tools)

    @staticmethod
    def _get_http_operation_tools(
        llm: BaseLanguageModel,
        spec: OpenAPISpec,
        requests: Optional[Requests] = None,
        verbose: bool = False,
        **kwargs: Any,
    ) -> List[NLATool]:
        """Get the tools for all the API operations."""
        if not spec.paths:
            return []
        http_operation_tools = []
        for path in spec.paths:
            for method in spec.get_methods_for_path(path):
                endpoint_tool = NLATool.from_llm_and_method(
                    llm=llm,
                    path=path,
                    method=method,
                    spec=spec,
                    requests=requests,
                    verbose=verbose,
                    **kwargs,
                )
                http_operation_tools.append(endpoint_tool)
        return http_operation_tools

    @classmethod
    def from_llm_and_spec(
        cls,
        llm: BaseLanguageModel,
        spec: OpenAPISpec,
        requests: Optional[Requests] = None,
        verbose: bool = False,
        **kwargs: Any,
    ) -> NLAToolkit:
        """Instantiate the toolkit by creating tools for each operation.

        Args:
            llm: The language model to use.
            spec: The OpenAPI spec.
            requests: Optional requests object. Default is None.
            verbose: Whether to print verbose output. Default is False.
            kwargs: Additional arguments.

        Returns:
            The toolkit.
        """
        http_operation_tools = cls._get_http_operation_tools(
            llm=llm, spec=spec, requests=requests, verbose=verbose, **kwargs
        )
        return cls(nla_tools=http_operation_tools)

    @classmethod
    def from_llm_and_url(
        cls,
        llm: BaseLanguageModel,
        open_api_url: str,
        requests: Optional[Requests] = None,
        verbose: bool = False,
        **kwargs: Any,
    ) -> NLAToolkit:
        """Instantiate the toolkit from an OpenAPI Spec URL.

        Args:
            llm: The language model to use.
            open_api_url: The URL of the OpenAPI spec.
            requests: Optional requests object. Default is None.
            verbose: Whether to print verbose output. Default is False.
            kwargs: Additional arguments.

        Returns:
            The toolkit.
        """

        spec = OpenAPISpec.from_url(open_api_url)
        return cls.from_llm_and_spec(
            llm=llm, spec=spec, requests=requests, verbose=verbose, **kwargs
        )

    @classmethod
    def from_llm_and_ai_plugin(
        cls,
        llm: BaseLanguageModel,
        ai_plugin: AIPlugin,
        requests: Optional[Requests] = None,
        verbose: bool = False,
        **kwargs: Any,
    ) -> NLAToolkit:
        """Instantiate the toolkit from an OpenAPI Spec URL"""
        spec = OpenAPISpec.from_url(ai_plugin.api.url)
        # TODO: Merge optional Auth information with the `requests` argument
        return cls.from_llm_and_spec(
            llm=llm,
            spec=spec,
            requests=requests,
            verbose=verbose,
            **kwargs,
        )

    @classmethod
    def from_llm_and_ai_plugin_url(
        cls,
        llm: BaseLanguageModel,
        ai_plugin_url: str,
        requests: Optional[Requests] = None,
        verbose: bool = False,
        **kwargs: Any,
    ) -> NLAToolkit:
        """Instantiate the toolkit from an OpenAPI Spec URL"""
        plugin = AIPlugin.from_url(ai_plugin_url)
        return cls.from_llm_and_ai_plugin(
            llm=llm, ai_plugin=plugin, requests=requests, verbose=verbose, **kwargs
        )