Source code for langchain.llms.sagemaker_endpoint
"""Sagemaker InvokeEndpoint API."""
from abc import abstractmethod
from typing import Any, Dict, Generic, List, Mapping, Optional, TypeVar, Union
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM
from langchain.llms.utils import enforce_stop_tokens
from langchain.pydantic_v1 import Extra, root_validator
INPUT_TYPE = TypeVar("INPUT_TYPE", bound=Union[str, List[str]])
OUTPUT_TYPE = TypeVar("OUTPUT_TYPE", bound=Union[str, List[List[float]]])
[docs]class ContentHandlerBase(Generic[INPUT_TYPE, OUTPUT_TYPE]):
    """A handler class to transform input from LLM to a
    format that SageMaker endpoint expects.
    Similarly, the class handles transforming output from the
    SageMaker endpoint to a format that LLM class expects.
    """
    """
    Example:
        .. code-block:: python
            class ContentHandler(ContentHandlerBase):
                content_type = "application/json"
                accepts = "application/json"
                def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
                    input_str = json.dumps({prompt: prompt, **model_kwargs})
                    return input_str.encode('utf-8')
                
                def transform_output(self, output: bytes) -> str:
                    response_json = json.loads(output.read().decode("utf-8"))
                    return response_json[0]["generated_text"]
    """
    content_type: Optional[str] = "text/plain"
    """The MIME type of the input data passed to endpoint"""
    accepts: Optional[str] = "text/plain"
    """The MIME type of the response data returned from endpoint"""
[docs]    @abstractmethod