Formating messages from History when using RunnableWithMessageHistory with PromptTemplate

[bojan2501]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

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

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)

Description

I am trying to use custom prompt format with RunnableWithMessageHistory.
For History I am using StreamlitChatMessageHistory.
And I cant find any info how to format messages from History for custom prompt.
Model in Sagemaker is Llama3.1 Instruct. Therefore custom prompt is needed.

This is prompt just before sending to LLM(Sagemaker):
<|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 [AIMessage(content='How can I help you?')] needs to be formatted in specific way, like prompt template is set.
Are there any resources with examples how to format output from History in RunnableWithMessageHistory?

All best
Bojan

System Info

langchain==0.2.14
langchain-aws==0.1.17
langchain-community==0.2.12
langchain-core==0.2.34
streamlit==1.37.1

[dosubot[bot]]

Hello, @bojan2501! I'm here to help you with your issue. I'm a bot that can assist you with bug fixes, answer your questions, and guide you on how to become a contributor. Let's work together on resolving your concern.

To format messages from StreamlitChatMessageHistory for a custom prompt in RunnableWithMessageHistory, you need to create a function that formats the chat history according to your prompt template. Here's an example of how you can achieve this:

1. Define a function to format the chat history:

def format_chat_history(chat_history):
    formatted_history = ""
    for message in chat_history.messages:
        if isinstance(message, HumanMessage):
            role = "user"
        elif isinstance(message, AIMessage):
            role = "assistant"
        else:
            role = "unknown"
        formatted_history += f"<|start_header_id|>{role}<|end_header_id|>\n{message.content}\n<|eot_id|>\n"
    return formatted_history

2. Integrate this function into your chain:

from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.messages import HumanMessage, AIMessage

# Define your prompt template
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)

# Initialize the message history
history = StreamlitChatMessageHistory(key="chat_messages")

# Define your LLM (SagemakerEndpoint in this case)
llm = SagemakerEndpoint(...)

# Create the chain
chain = prompt_template | llm

# Wrap the chain with RunnableWithMessageHistory
wrapped_chain = RunnableWithMessageHistory(
    chain,
    lambda session_id: history,
    history_messages_key="chat_history"
)

# Format the chat history and invoke the chain
formatted_history = format_chat_history(history)
response = wrapped_chain.invoke(
    {"input": "What can you do?", "system": "You are a good assistant.", "chat_history": formatted_history},
    config
)

print(response)

This example demonstrates how to format the chat history and integrate it into your chain using RunnableWithMessageHistory. The format_chat_history function ensures that the messages are formatted according to your custom prompt template before being sent to the LLM.

For more detailed examples and guides, you can refer to the LangChain documentation on message history and the LangChain tutorial on building a chatbot ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[bojan2501]

@dosu
Unfortunately it did not help. Upon debugging I can see that function format_chat_history is not even called.
What could be wrong?

[dosubot[bot]]

The issue might be that the format_chat_history function is not being called within the RunnableWithMessageHistory setup. Ensure that the function is correctly integrated and invoked within the chain. Here's a revised approach:

1. Ensure the format_chat_history function is defined and accessible:

def format_chat_history(chat_history):
    formatted_history = ""
    for message in chat_history.messages:
        if isinstance(message, HumanMessage):
            role = "user"
        elif isinstance(message, AIMessage):
            role = "assistant"
        else:
            role = "unknown"
        formatted_history += f"<|start_header_id|>{role}<|end_header_id|>\n{message.content}\n<|eot_id|>\n"
    return formatted_history

2. Modify the RunnableWithMessageHistory to use the format_chat_history function:

from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.messages import HumanMessage, AIMessage
from langchain_community.llms import SagemakerEndpoint

# Define your prompt template
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)

# Initialize the message history
history = StreamlitChatMessageHistory(key="chat_messages")

# Define your LLM (SagemakerEndpoint in this case)
llm = SagemakerEndpoint(...)

# Create the chain
chain = prompt_template | llm

# Wrap the chain with RunnableWithMessageHistory
wrapped_chain = RunnableWithMessageHistory(
    chain,
    lambda session_id: history,
    history_messages_key="chat_history"
)

# Ensure the chat history is formatted before invoking the chain
def invoke_with_formatted_history(input_data, config):
    formatted_history = format_chat_history(history)
    input_data["chat_history"] = formatted_history
    return wrapped_chain.invoke(input_data, config)

# Invoke the chain with formatted history
response = invoke_with_formatted_history(
    {"input": "What can you do?", "system": "You are a good assistant."},
    config
)

print(response)

This approach ensures that the format_chat_history function is explicitly called to format the chat history before invoking the chain. This should resolve the issue of the function not being called.

To continue talking to Dosu, mention @dosu.

[bojan2501]

@dosu
The function was called but the prompt did not update.

[bojan2501]

@dosu
You are not using prompt_data template.

[dosubot[bot]]

To specify the format of history messages in the prompt_data template when using RunnableWithMessageHistory, you can use the MessagesPlaceholder in your ChatPromptTemplate. Here is an example that incorporates your prompt_data template:

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from langchain_community.llms import SagemakerEndpoint

# Define your prompt template
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 = ChatPromptTemplate.from_messages([
    ("system", "{system}"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
])

# Initialize the message history
history = StreamlitChatMessageHistory(key="chat_messages")

# Define your LLM (SagemakerEndpoint in this case)
llm = SagemakerEndpoint(...)

# Create the chain
chain = prompt_template | llm

# Wrap the chain with RunnableWithMessageHistory
wrapped_chain = RunnableWithMessageHistory(
    chain,
    lambda session_id: history,
    input_messages_key="input",
    history_messages_key="chat_history"
)

# Invoke the chain with formatted history
response = wrapped_chain.invoke(
    {"input": "What can you do?", "system": "You are a good assistant."},
    config={"configurable": {"session_id": "1"}}
)

print(response)

In this example:

• MessagesPlaceholder(variable_name="chat_history") is used to specify where the history messages should be inserted in the prompt_data template.
• input_messages_key="input" specifies that the latest input message is under the key "input".
• history_messages_key="chat_history" specifies that historical messages should be added under the key "chat_history".

This setup ensures that the history messages are correctly formatted and integrated into the prompt template ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[bojan2501]

@dosu.
This is not working, prompt_data is not used in your example.
For anyone else I tried update prompt in ContentHandler.transform_input() but it failed also as the promt at this point is string.

[dosubot[bot]]

To ensure that the prompt_data template is correctly used and updated when using RunnableWithMessageHistory, you need to make sure that the prompt object passed to the create_history_aware_retriever function includes input as one of its input variables. This is validated in the function with the following check:

if "input" not in prompt.input_variables:
    raise ValueError(
        "Expected `input` to be a prompt variable, "
        f"but got {prompt.input_variables}"
    )

This ensures that the prompt template is correctly set up to handle the input variable. Additionally, the RunnableBranch is configured to handle cases with and without chat_history appropriately, ensuring that the prompt is used correctly in generating the search query when chat_history is present.

Here is the relevant part of the code that sets up the RunnableBranch:

retrieve_documents: RetrieverOutputLike = RunnableBranch(
    (
        # Both empty string and empty list evaluate to False
        lambda x: not x.get("chat_history", False),
        # If no chat history, then we just pass input to retriever
        (lambda x: x["input"]) | retriever,
    ),
    # If chat history, then we pass inputs to LLM chain, then to retriever
    prompt | llm | StrOutputParser() | retriever,
).with_config(run_name="chat_retriever_chain")

This configuration ensures that the prompt is used to generate a search query when chat_history is present, and the query is then passed to the retriever.

Additionally, LangChain provides a RunnableWithMessageHistory class that can handle the process of appending the input variable to the chat history automatically. Here’s how you can set it up:

1. Define the Prompt Template:

   from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
   
   prompt = ChatPromptTemplate.from_messages(
       [
           ("system", "You are a helpful assistant. Answer all questions to the best of your ability."),
           MessagesPlaceholder(variable_name="messages"),
       ]
   )

2. Create the Chain:

   from langchain_openai import ChatOpenAI
   
   chat = ChatOpenAI(model="gpt-3.5-turbo")
   chain = prompt | chat

3. Wrap the Chain with RunnableWithMessageHistory:

   from langchain_core.runnables.history import RunnableWithMessageHistory
   from langchain_core.chat_history import InMemoryChatMessageHistory
   
   store = {}
   
   def get_session_history(session_id: str):
       if session_id not in store:
           store[session_id] = InMemoryChatMessageHistory()
       return store[session_id]
   
   chain_with_message_history = RunnableWithMessageHistory(
       chain,
       get_session_history,
       input_messages_key="messages",
   )

4. Invoke the Chain:

   from langchain_core.messages import HumanMessage
   
   config = {"configurable": {"session_id": "abc2"}}
   
   response = chain_with_message_history.invoke(
       {"messages": [HumanMessage(content="Translate this sentence from English to French: I love programming.")]},
       config=config,
   )

By following these steps, you ensure that the prompt_data template is correctly used and updated, leveraging the RunnableWithMessageHistory class to manage the chat history automatically.

Additional Considerations

When constructing RunnableWithMessageHistory, you need to pass in a get_session_history function that takes a session_id and returns a BaseChatMessageHistory object. This allows you to maintain several conversations/threads with the same chain at the same time. Here’s an example of how to set this up:

from langchain_core.chat_history import SQLChatMessageHistory

def get_session_history(session_id):
    return SQLChatMessageHistory(session_id, "sqlite:///memory.db")

You can also customize the configuration parameters by which you track message histories by passing in a list of ConfigurableFieldSpec objects to the history_factory_config parameter. This allows for more complex setups involving multiple identifiers like user_id and conversation_id.

By properly setting up the RunnableWithMessageHistory and ensuring the prompt template includes the necessary variables, you can effectively manage and update the prompt_data template in LangChain ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[bojan2501]

I give up. There is no solution for now.