Error when using SagemakerEndpoint

[massi-ang]

When running the following code (taken from the docs):

from langchain_aws import SagemakerEndpoint
from langchain_community.llms.sagemaker_endpoint import LLMContentHandler

class ContentHandler(LLMContentHandler):
    content_type = "application/json"
    accepts = "application/json"

    def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
        input_str = json.dumps({"inputs": prompt, "parameters": 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_handler = ContentHandler()

llm=SagemakerEndpoint(
    endpoint_name="<endpoint>",
    client=sm,
    model_kwargs={"temperature": 1e-10},
    content_handler=content_handler
)

the following error is thrown:

{
	"name": "ValidationError",
	"message": "1 validation error for SagemakerEndpoint
content_handler
  Input should be an instance of LLMContentHandler [type=is_instance_of, input_value=<__main__.ContentHandler object at 0x7f665a6e7250>, input_type=ContentHandler]
    For further information visit https://errors.pydantic.dev/2.10/v/is_instance_of",
	"stack": "---------------------------------------------------------------------------
ValidationError                           Traceback (most recent call last)
/home/sagemaker-user/s1/trainining.ipynb Cell 21 line 2
     <a href='vscode-notebook-cell://auonzkaukrqfcia.studio.us-east-1.sagemaker.aws/home/sagemaker-user/s1/trainining.ipynb#X26sdnNjb2RlLXJlbW90ZQ%3D%3D?line=14'>15</a>         return response_json[0][\"generated_text\"]
     <a href='vscode-notebook-cell://auonzkaukrqfcia.studio.us-east-1.sagemaker.aws/home/sagemaker-user/s1/trainining.ipynb#X26sdnNjb2RlLXJlbW90ZQ%3D%3D?line=17'>18</a> content_handler = ContentHandler()
---> <a href='vscode-notebook-cell://auonzkaukrqfcia.studio.us-east-1.sagemaker.aws/home/sagemaker-user/s1/trainining.ipynb#X26sdnNjb2RlLXJlbW90ZQ%3D%3D?line=19'>20</a> llm=SagemakerEndpoint(
     <a href='vscode-notebook-cell://auonzkaukrqfcia.studio.us-east-1.sagemaker.aws/home/sagemaker-user/s1/trainining.ipynb#X26sdnNjb2RlLXJlbW90ZQ%3D%3D?line=20'>21</a>     endpoint_name=\"<endpoint>\",
     <a href='vscode-notebook-cell://auonzkaukrqfcia.studio.us-east-1.sagemaker.aws/home/sagemaker-user/s1/trainining.ipynb#X26sdnNjb2RlLXJlbW90ZQ%3D%3D?line=21'>22</a>     client=sm,
     <a href='vscode-notebook-cell://auonzkaukrqfcia.studio.us-east-1.sagemaker.aws/home/sagemaker-user/s1/trainining.ipynb#X26sdnNjb2RlLXJlbW90ZQ%3D%3D?line=22'>23</a>     model_kwargs={\"temperature\": 1e-10},
     <a href='vscode-notebook-cell://auonzkaukrqfcia.studio.us-east-1.sagemaker.aws/home/sagemaker-user/s1/trainining.ipynb#X26sdnNjb2RlLXJlbW90ZQ%3D%3D?line=23'>24</a>     content_handler=content_handler
     <a href='vscode-notebook-cell://auonzkaukrqfcia.studio.us-east-1.sagemaker.aws/home/sagemaker-user/s1/trainining.ipynb#X26sdnNjb2RlLXJlbW90ZQ%3D%3D?line=24'>25</a> )

File /opt/conda/lib/python3.11/site-packages/langchain_core/load/serializable.py:125, in Serializable.__init__(self, *args, **kwargs)
    123 def __init__(self, *args: Any, **kwargs: Any) -> None:
    124     \"\"\"\"\"\"
--> 125     super().__init__(*args, **kwargs)

File /opt/conda/lib/python3.11/site-packages/pydantic/main.py:214, in BaseModel.__init__(self, **data)
    212 # `__tracebackhide__` tells pytest and some other tools to omit this function from tracebacks
    213 __tracebackhide__ = True
--> 214 validated_self = self.__pydantic_validator__.validate_python(data, self_instance=self)
    215 if self is not validated_self:
    216     warnings.warn(
    217         'A custom validator is returning a value other than `self`.\
'
    218         \"Returning anything other than `self` from a top level model validator isn't supported when validating via `__init__`.\
\"
    219         'See the `model_validator` docs (https://docs.pydantic.dev/latest/concepts/validators/#model-validators) for more details.',
    220         stacklevel=2,
    221     )

ValidationError: 1 validation error for SagemakerEndpoint
content_handler
  Input should be an instance of LLMContentHandler [type=is_instance_of, input_value=<__main__.ContentHandler object at 0x7f665a6e7250>, input_type=ContentHandler]
    For further information visit https://errors.pydantic.dev/2.10/v/is_instance_of"
}

When changing from langchain_aws import SagemakerEndpoint to from langchain_community.llms import SagemakerEndpoint it works but gives an error that the class is deprecated and one should use langchain_aws

[michaelnchin]

Langchain-aws's SagemakerEndpoint requires an instance of the LLMContentHandler class defined in the same package (and vice versa for langchain-community's implementation).

This import should work:

from langchain_aws import SagemakerEndpoint
from langchain_aws.llms.sagemaker_endpoint import LLMContentHandler