codekingpro/portable-devtools

1"""Chain that makes API calls and summarizes the responses to answer a question."""2 3from __future__ import annotations4 5import json6from typing import Any, Dict, List, NamedTuple, Optional, cast7 8from langchain_classic.chains.api.openapi.requests_chain import APIRequesterChain9from langchain_classic.chains.api.openapi.response_chain import APIResponderChain10from langchain_classic.chains.base import Chain11from langchain_classic.chains.llm import LLMChain12from langchain_core.callbacks import CallbackManagerForChainRun, Callbacks13from langchain_core.language_models import BaseLanguageModel14from pydantic import BaseModel, Field15from requests import Response16 17from langchain_community.tools.openapi.utils.api_models import APIOperation18from langchain_community.utilities.requests import Requests19 20 21class _ParamMapping(NamedTuple):22    """Mapping from parameter name to parameter value."""23 24    query_params: List[str]25    body_params: List[str]26    path_params: List[str]27 28 29class OpenAPIEndpointChain(Chain, BaseModel):30    """Chain interacts with an OpenAPI endpoint using natural language."""31 32    api_request_chain: LLMChain33    api_response_chain: Optional[LLMChain] = None34    api_operation: APIOperation35    requests: Requests = Field(exclude=True, default_factory=Requests)36    param_mapping: _ParamMapping = Field(alias="param_mapping")37    return_intermediate_steps: bool = False38    instructions_key: str = "instructions"  #: :meta private:39    output_key: str = "output"  #: :meta private:40    max_text_length: Optional[int] = Field(ge=0)  #: :meta private:41 42    @property43    def input_keys(self) -> List[str]:44        """Expect input key.45 46        :meta private:47        """48        return [self.instructions_key]49 50    @property51    def output_keys(self) -> List[str]:52        """Expect output key.53 54        :meta private:55        """56        if not self.return_intermediate_steps:57            return [self.output_key]58        else:59            return [self.output_key, "intermediate_steps"]60 61    def _construct_path(self, args: Dict[str, str]) -> str:62        """Construct the path from the deserialized input."""63        path = self.api_operation.base_url + self.api_operation.path64        for param in self.param_mapping.path_params:65            path = path.replace(f"{{{param}}}", str(args.pop(param, "")))66        return path67 68    def _extract_query_params(self, args: Dict[str, str]) -> Dict[str, str]:69        """Extract the query params from the deserialized input."""70        query_params = {}71        for param in self.param_mapping.query_params:72            if param in args:73                query_params[param] = args.pop(param)74        return query_params75 76    def _extract_body_params(self, args: Dict[str, str]) -> Optional[Dict[str, str]]:77        """Extract the request body params from the deserialized input."""78        body_params = None79        if self.param_mapping.body_params:80            body_params = {}81            for param in self.param_mapping.body_params:82                if param in args:83                    body_params[param] = args.pop(param)84        return body_params85 86    def deserialize_json_input(self, serialized_args: str) -> dict:87        """Use the serialized typescript dictionary.88 89        Resolve the path, query params dict, and optional requestBody dict.90        """91        args: dict = json.loads(serialized_args)92        path = self._construct_path(args)93        body_params = self._extract_body_params(args)94        query_params = self._extract_query_params(args)95        return {96            "url": path,97            "data": body_params,98            "params": query_params,99        }100 101    def _get_output(self, output: str, intermediate_steps: dict) -> dict:102        """Return the output from the API call."""103        if self.return_intermediate_steps:104            return {105                self.output_key: output,106                "intermediate_steps": intermediate_steps,107            }108        else:109            return {self.output_key: output}110 111    def _call(112        self,113        inputs: Dict[str, Any],114        run_manager: Optional[CallbackManagerForChainRun] = None,115    ) -> Dict[str, str]:116        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()117        intermediate_steps = {}118        instructions = inputs[self.instructions_key]119        instructions = instructions[: self.max_text_length]120        _api_arguments = self.api_request_chain.predict_and_parse(121            instructions=instructions, callbacks=_run_manager.get_child()122        )123        api_arguments = cast(str, _api_arguments)124        intermediate_steps["request_args"] = api_arguments125        _run_manager.on_text(126            api_arguments, color="green", end="\n", verbose=self.verbose127        )128        if api_arguments.startswith("ERROR"):129            return self._get_output(api_arguments, intermediate_steps)130        elif api_arguments.startswith("MESSAGE:"):131            return self._get_output(132                api_arguments[len("MESSAGE:") :], intermediate_steps133            )134        try:135            request_args = self.deserialize_json_input(api_arguments)136            method = getattr(self.requests, self.api_operation.method.value)137            api_response: Response = method(**request_args)138            if api_response.status_code != 200:139                method_str = str(self.api_operation.method.value)140                response_text = (141                    f"{api_response.status_code}: {api_response.reason}"142                    + f"\nFor {method_str.upper()}  {request_args['url']}\n"143                    + f"Called with args: {request_args['params']}"144                )145            else:146                response_text = api_response.text147        except Exception as e:148            response_text = f"Error with message {str(e)}"149        response_text = response_text[: self.max_text_length]150        intermediate_steps["response_text"] = response_text151        _run_manager.on_text(152            response_text, color="blue", end="\n", verbose=self.verbose153        )154        if self.api_response_chain is not None:155            _answer = self.api_response_chain.predict_and_parse(156                response=response_text,157                instructions=instructions,158                callbacks=_run_manager.get_child(),159            )160            answer = cast(str, _answer)161            _run_manager.on_text(answer, color="yellow", end="\n", verbose=self.verbose)162            return self._get_output(answer, intermediate_steps)163        else:164            return self._get_output(response_text, intermediate_steps)165 166    @classmethod167    def from_url_and_method(168        cls,169        spec_url: str,170        path: str,171        method: str,172        llm: BaseLanguageModel,173        requests: Optional[Requests] = None,174        return_intermediate_steps: bool = False,175        **kwargs: Any,176        # TODO: Handle async177    ) -> "OpenAPIEndpointChain":178        """Create an OpenAPIEndpoint from a spec at the specified url."""179        operation = APIOperation.from_openapi_url(spec_url, path, method)180        return cls.from_api_operation(181            operation,182            requests=requests,183            llm=llm,184            return_intermediate_steps=return_intermediate_steps,185            **kwargs,186        )187 188    @classmethod189    def from_api_operation(190        cls,191        operation: APIOperation,192        llm: BaseLanguageModel,193        requests: Optional[Requests] = None,194        verbose: bool = False,195        return_intermediate_steps: bool = False,196        raw_response: bool = False,197        callbacks: Callbacks = None,198        **kwargs: Any,199        # TODO: Handle async200    ) -> "OpenAPIEndpointChain":201        """Create an OpenAPIEndpointChain from an operation and a spec."""202        param_mapping = _ParamMapping(203            query_params=operation.query_params,204            body_params=operation.body_params,205            path_params=operation.path_params,206        )207        requests_chain = APIRequesterChain.from_llm_and_typescript(208            llm,209            typescript_definition=operation.to_typescript(),210            verbose=verbose,211            callbacks=callbacks,212        )213        if raw_response:214            response_chain = None215        else:216            response_chain = APIResponderChain.from_llm(217                llm, verbose=verbose, callbacks=callbacks218            )219        _requests = requests or Requests()220        return cls(221            api_request_chain=requests_chain,222            api_response_chain=response_chain,223            api_operation=operation,224            requests=_requests,225            param_mapping=param_mapping,226            verbose=verbose,227            return_intermediate_steps=return_intermediate_steps,228            callbacks=callbacks,229            **kwargs,230        )231