Using SagemakerEndpoint together with RunnableWithMessageHistory with custom prompt template creates broken prompt

[bojan2501]

I am testing SagemakerEndpoint(Llama3 Instruct) with RunnableWithMessageHistory. This endpoint is using specfic template and the generated template from RunnableWithMessageHistory is broken as the messages from History dont follow formatting.
Code example:

prompt_data = """
        <|begin_of_text|>
        <|start_header_id|>system<|end_header_id|>
        {system}
        <|eot_id|>
        {chat_history}
        <|start_header_id|>user<|end_header_id|>
        {input}
        <|eot_id|>
        <|start_header_id|>assistant<|end_header_id|>
        """
prompt_template = PromptTemplate(template=prompt_data)

history = StreamlitChatMessageHistory(key="chat_history")

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"]

llm = SagemakerEndpoint(...)

chain = prompt_template | llm
wrapped_chain = RunnableWithMessageHistory(chain,
                                           lambda session_id: history,
                                           history_messages_key="chat_history")

response = wrapped_chain.invoke({"input": "What can you do?", "system": "You are a good assistant."}, config)

The output of the call when invoked is:

        <|begin_of_text|>
        <|start_header_id|>system<|end_header_id|>
        You are a good assistant.
        <|eot_id|>
        [AIMessage(content='How can I help you?')]
        <|start_header_id|>user<|end_header_id|>
        What can you do?
        <|eot_id|>
        <|start_header_id|>assistant<|end_header_id|>

And as you can see the issue is that messages from history are not formatted.
Is there a way to format history messages?
I tried:

1. See can I do it in when providing function which returns history. Not able to do it.
2. Do the formatting in ContentHandler transform_input but it is a string there which is mixed bag as you can see.
3. Create template using ChatPromptTemplate.from _messages and then try again 2. Again string which is not the best to handle as it can contain a lot of bad things.
4. Checked can I reuse langchain_aws.chat_models.bedrock.convert_messages_to_prompt_llama3 and this failed.

Any suggestion or explanation about this issue is more then welcome.

All best
Bojan