ChatSagemakerEndpoint not working. Anyone gotten it to work?

[SuryaKari]

I am trying to use a Sagemaker Endpoint as a chat LLM and have started to explore ChatSagemakerEndpoint. But it is confusing and is not working. I got the ChatBedrockConverse working but not this. I have attached my code here

import openlit
from langgraph.graph import StateGraph
from langchain_aws.llms import SagemakerEndpoint
from langchain_aws.llms.sagemaker_endpoint import LLMContentHandler
import json
import boto3
from langchain_aws.llms.bedrock import BedrockLLM
from contextlib import asynccontextmanager
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent
from langchain_core.language_models import BaseChatModel
from langchain_core.outputs import ChatResult, ChatGeneration
from langchain_aws.chat_models import ChatBedrock
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.graph import StateGraph
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    HumanMessage,
    SystemMessage,
    merge_message_runs,
)
from langchain_aws import ChatBedrockConverse
from langchain_aws.chat_models.sagemaker_endpoint import ChatSagemakerEndpoint, ChatModelContentHandler
from typing import Any, Dict, Tuple, Type, List


# Configure AWS credentials directly in the code
aws_access_key_id = "mock"
aws_secret_access_key = "mock"
aws_region = "us-east-2"

# Create a session with explicit credentials
session = boto3.Session(
    aws_access_key_id=aws_access_key_id,
    aws_secret_access_key=aws_secret_access_key,
    region_name=aws_region
)

class DeepSeekChatHandler(ChatModelContentHandler):
    content_type = "application/json"
    accepts = "application/json"

    def transform_input(self, messages, model_kwargs):
        """Transform messages to DeepSeek format"""
        # DeepSeek expects a specific format with roles
        print(messages)
        input_data = {
            "inputs": [[{"role": msg['role'], "content": msg['content']} for msg in messages]],
            "parameters": {
                "temperature": model_kwargs.get("temperature", 0.5),
                "max_new_tokens": model_kwargs.get("max_new_tokens", 700),
                "top_p": model_kwargs.get("top_p", 0.9)
            }
        }
        return json.dumps(input_data).encode("utf-8")

    def transform_output(self, output_bytes) -> BaseMessage:
        """Transform DeepSeek output to AIMessage"""
        response_json = json.loads(output_bytes.read().decode("utf-8"))
        return AIMessage(content=response_json[0]["generation"]["content"])

chat_content_handler = DeepSeekChatHandler()

# Initialize chat model
sagemaker_llm_chat = ChatSagemakerEndpoint(
    endpoint_name="jumpstart-dft-deepseek-llm-r1-disti-20250302-011118",
    region_name=aws_region,
    content_handler=DeepSeekChatHandler(),
    client=sagemaker_client,
    model_kwargs={
        "temperature": 0.5,
        "max_new_tokens": 700,
        "top_p": 0.9
    }
)


messages = [
    SystemMessage(content="You are an AWS expert."),
    HumanMessage(content="What is Sagemaker endpoints?"),
]

service_response = sagemaker_llm_chat.invoke(messages)

print(service_response)