Source code for langchain_community.llms.amazon_api_gateway

from typing import Any, Dict, List, Mapping, Optional

import requests
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.pydantic_v1 import Extra

from langchain_community.llms.utils import enforce_stop_tokens

[docs]class ContentHandlerAmazonAPIGateway: """Adapter to prepare the inputs from Langchain to a format that LLM model expects. It also provides helper function to extract the generated text from the model response."""
[docs] @classmethod def transform_input( cls, prompt: str, model_kwargs: Dict[str, Any] ) -> Dict[str, Any]: return {"inputs": prompt, "parameters": model_kwargs}
[docs] @classmethod def transform_output(cls, response: Any) -> str: return response.json()[0]["generated_text"]
[docs]class AmazonAPIGateway(LLM): """Amazon API Gateway to access LLM models hosted on AWS.""" api_url: str """API Gateway URL""" headers: Optional[Dict] = None """API Gateway HTTP Headers to send, e.g. for authentication""" model_kwargs: Optional[Dict] = None """Keyword arguments to pass to the model.""" content_handler: ContentHandlerAmazonAPIGateway = ContentHandlerAmazonAPIGateway() """The content handler class that provides an input and output transform functions to handle formats between LLM and the endpoint. """ class Config: """Configuration for this pydantic object.""" extra = Extra.forbid @property def _identifying_params(self) -> Mapping[str, Any]: """Get the identifying parameters.""" _model_kwargs = self.model_kwargs or {} return { **{"api_url": self.api_url, "headers": self.headers}, **{"model_kwargs": _model_kwargs}, } @property def _llm_type(self) -> str: """Return type of llm.""" return "amazon_api_gateway" def _call( self, prompt: str, stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any, ) -> str: """Call out to Amazon API Gateway model. Args: prompt: The prompt to pass into the model. stop: Optional list of stop words to use when generating. Returns: The string generated by the model. Example: .. code-block:: python response = se("Tell me a joke.") """ _model_kwargs = self.model_kwargs or {} payload = self.content_handler.transform_input(prompt, _model_kwargs) try: response = self.api_url, headers=self.headers, json=payload, ) text = self.content_handler.transform_output(response) except Exception as error: raise ValueError(f"Error raised by the service: {error}") if stop is not None: text = enforce_stop_tokens(text, stop) return text