Use RunnableWithMessageHistory with structured output

[segoranov]

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

from pydantic.v1 import BaseModel
from langchain_openai import AzureChatOpenAI
from typing import Optional
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_mongodb import MongoDBChatMessageHistory


class MyStructuredOutput(BaseModel):
    long_explanation: str
    short_explanation: Optional[str] = None


def interact_with_llm(user_question: str) -> MyStructuredOutput:
    chat_llm_with_history = _get_chat_llm_with_history()
    config = {"configurable": {"session_id": "some-test-session-id"}}
    return chat_llm_with_history.invoke({"question": user_question}, config=config)


def _get_chat_llm_with_history():
    chat_llm = AzureChatOpenAI().with_structured_output(MyStructuredOutput)
    prompt_template: ChatPromptTemplate = _get_prompt_template()
    runnable = prompt_template | chat_llm

    return RunnableWithMessageHistory(runnable,
                                      _get_message_history,
                                      input_messages_key="question",
                                      history_messages_key="history",
                                      )


def _get_message_history(session_id: str):
    return MongoDBChatMessageHistory(
        session_id=session_id,
        connection_string="my-mongo-connection-string",
        database_name="my-db-name",
        collection_name="chat_histories",
    )


def _get_prompt_template() -> ChatPromptTemplate:
    return ChatPromptTemplate.from_messages(
        [
            ("system", "You are a knowledgeable AI assistant good at math."),
            MessagesPlaceholder(variable_name="history"),
            ("human", "{question}"),
        ]
    )


if __name__ == "__main__":
    interact_with_llm("Explain to me theory of relativity")

Description

I'm trying to use the langchain library to implement chat with history support using structured output (with RunnableWithMessageHistory)

I expect langchain to persist in MongoDB the human messages and the AI responses, but no history is being saved and this error is observed in the logs:

WARNING [langchain_core.callbacks.manager] [manager.py:282] [trace_id=9fd20617ccf590787cdd9ad1b0ba12dd span_id=03557bfbc0db6bf9 resource.service.name=service.name trace_sampled=True] - Error in RootListenersTracer.on_chain_end callback: ValueError()

I tried using one of the following keys as output_messages_key parameter for the RunnableWithMessageHistory object: long_explanation, short_explanation, MyStructuredOutput, but all of the attempts resulted with KeyError like WARNING [langchain_core.callbacks.manager] [manager.py:282] [trace_id="some_id" span_id="some_id" resource.service.name=service.name trace_sampled=True] - Error in RootListenersTracer.on_chain_end callback: KeyError('long_explanation'). On top of this I want the whole structured output to be stored in history, not just a specific part (key) of it.

So is it possible to use at all structured output with RunnableWithMessageHistory or should I implement history management by myself?

System Info

langchain==0.1.14
langchain-community==0.0.38
langchain-core==0.1.52
langchain-mongodb==0.1.6
langchain-openai==0.1.1
langchain-text-splitters==0.0.2

[dosubot[bot]]

Hello, @segoranov! I'm here to assist you with any bugs, questions, or contributions you might have. Let's work together to resolve the issue at hand.

Yes, it is possible to use structured output with RunnableWithMessageHistory in the LangChain library. The RunnableWithMessageHistory class is designed to manage chat message history for another Runnable, and it supports structured input and output formats.

Here is an example of how you can use RunnableWithMessageHistory with structured output:

from pydantic.v1 import BaseModel
from langchain_openai import AzureChatOpenAI
from typing import Optional
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_mongodb import MongoDBChatMessageHistory


class MyStructuredOutput(BaseModel):
    long_explanation: str
    short_explanation: Optional[str] = None


def interact_with_llm(user_question: str) -> MyStructuredOutput:
    chat_llm_with_history = _get_chat_llm_with_history()
    config = {"configurable": {"session_id": "some-test-session-id"}}
    return chat_llm_with_history.invoke({"question": user_question}, config=config)


def _get_chat_llm_with_history():
    chat_llm = AzureChatOpenAI().with_structured_output(MyStructuredOutput)
    prompt_template: ChatPromptTemplate = _get_prompt_template()
    runnable = prompt_template | chat_llm

    return RunnableWithMessageHistory(
        runnable,
        _get_message_history,
        input_messages_key="question",
        history_messages_key="history",
        output_messages_key=None  # Ensure the whole structured output is stored
    )


def _get_message_history(session_id: str):
    return MongoDBChatMessageHistory(
        session_id=session_id,
        connection_string="my-mongo-connection-string",
        database_name="my-db-name",
        collection_name="chat_histories",
    )


def _get_prompt_template() -> ChatPromptTemplate:
    return ChatPromptTemplate.from_messages(
        [
            ("system", "You are a knowledgeable AI assistant good at math."),
            MessagesPlaceholder(variable_name="history"),
            ("human", "{question}"),
        ]
    )


if __name__ == "__main__":
    interact_with_llm("Explain to me theory of relativity")

In this example, the output_messages_key parameter is set to None, which ensures that the entire structured output is stored in the history. This should resolve the KeyError you encountered when specifying individual keys.

Additionally, ensure that your MongoDB setup is correctly configured and that the connection string, database name, and collection name are accurate. This will help in persisting the chat history as expected.

By following this approach, you can effectively use structured output with RunnableWithMessageHistory and store the entire structured output in MongoDB ^[1].

To continue talking to Dosu, mention @dosu.

[segoranov]

@dosu same ValueError() is received and output_messages_key is None by default anyway.

[dosubot[bot]]

To resolve the ValueError() and ensure the history is saved correctly in MongoDB when using structured output with RunnableWithMessageHistory, you need to ensure that the input and output values are correctly formatted and that the session history is properly managed.

Here is an updated example of how to set up and use RunnableWithMessageHistory with MongoDB:

from pydantic.v1 import BaseModel
from langchain_openai import AzureChatOpenAI
from typing import Optional, List
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_mongodb import MongoDBChatMessageHistory
from langchain_core.messages import BaseMessage, AIMessage, HumanMessage, message_to_dict, messages_from_dict
from pymongo import MongoClient, errors
import json
import logging

logger = logging.getLogger(__name__)

class MyStructuredOutput(BaseModel):
    long_explanation: str
    short_explanation: Optional[str] = None

class MongoDBChatMessageHistory(BaseChatMessageHistory):
    """Chat message history that stores history in MongoDB."""

    def __init__(self, connection_string: str, session_id: str, database_name: str = "chat_history", collection_name: str = "message_store", create_index: bool = True):
        self.connection_string = connection_string
        self.session_id = session_id
        self.database_name = database_name
        self.collection_name = collection_name

        try:
            self.client: MongoClient = MongoClient(connection_string)
        except errors.ConnectionFailure as error:
            logger.error(error)

        self.db = self.client[database_name]
        self.collection = self.db[collection_name]
        if create_index:
            self.collection.create_index("SessionId")

    @property
    def messages(self) -> List[BaseMessage]:
        """Retrieve the messages from MongoDB"""
        try:
            cursor = self.collection.find({"SessionId": self.session_id})
        except errors.OperationFailure as error:
            logger.error(error)
            return []

        items = [json.loads(document["History"]) for document in cursor]
        messages = messages_from_dict(items)
        return messages

    def add_message(self, message: BaseMessage) -> None:
        """Append the message to the record in MongoDB"""
        try:
            self.collection.insert_one(
                {
                    "SessionId": self.session_id,
                    "History": json.dumps(message_to_dict(message)),
                }
            )
        except errors.WriteError as err:
            logger.error(err)

    def clear(self) -> None:
        """Clear session memory from MongoDB"""
        try:
            self.collection.delete_many({"SessionId": self.session_id})
        except errors.WriteError as err:
            logger.error(err)

def get_by_session_id(session_id: str) -> BaseChatMessageHistory:
    return MongoDBChatMessageHistory(
        connection_string="your_mongodb_connection_string",
        session_id=session_id
    )

def interact_with_llm(user_question: str) -> MyStructuredOutput:
    chat_llm_with_history = _get_chat_llm_with_history()
    config = {"configurable": {"session_id": "some-test-session-id"}}
    return chat_llm_with_history.invoke({"question": user_question}, config=config)

def _get_chat_llm_with_history():
    chat_llm = AzureChatOpenAI().with_structured_output(MyStructuredOutput)
    prompt_template: ChatPromptTemplate = _get_prompt_template()
    runnable = prompt_template | chat_llm

    return RunnableWithMessageHistory(
        runnable,
        get_by_session_id,
        input_messages_key="question",
        history_messages_key="history",
        output_messages_key=None  # Ensure the whole structured output is stored
    )

def _get_prompt_template() -> ChatPromptTemplate:
    return ChatPromptTemplate.from_messages(
        [
            ("system", "You are a knowledgeable AI assistant good at math."),
            MessagesPlaceholder(variable_name="history"),
            ("human", "{question}"),
        ]
    )

if __name__ == "__main__":
    interact_with_llm("Explain to me theory of relativity")

In this example, the get_by_session_id function manages the session history using MongoDB, and the RunnableWithMessageHistory is configured to use this function. The invoke method is called with the appropriate input and configuration, ensuring that the history is correctly saved and managed in MongoDB ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[hschawe]

I'm facing the same problem, did you find a solution for this? Or did you have to move away from RunnableWithMessageHistory?

[eyurtsev]

You can't use it with with_structured_output unfortunately. RunnableWithMessageHistroy expects messages on the output of the wrapped runnable. When you use with_structured_output the output is not an AIMessage, but either a dict or a a pydantic object.

Your options are:

1. Swap to using the bind_tools interface instead -- it'll involve a bit more manual effort
2. Use with_structured_output with include_raw=False . See API reference, https://python.langchain.com/api_reference/anthropic/chat_models/langchain_anthropic.chat_models.ChatAnthropic.html#langchain_anthropic.chat_models.ChatAnthropic.with_structured_output

In both cases, you'll may need to restructure the LCEL chain a bit to make sure that RunnableWithMessageHistories is wrapping a chain that is returning an AIMessage rather than any other type. This might be challenging for folks not as comfortable with LCEL, I'd recomend swapping to langgraph for adding chat history to the conversation it'll be more obvious how to do this.

[Yairama]

But is worth to make somethin like this?:

class QueryReviewer(BaseModel):
    ...

lsto = llm.with_structured_output(QueryReviewer, include_raw=False)
result = lsto.invoke("There is some prompt")
result = AIMessage (content=str(result), artifact=result.dict())

# then pass result to the history

Thanks for supporting with the questions!

[mlnrt]

@segoranov @Yairama @eyurtsev , There is one way to make it work out of the box.
First you instruct to include the raw output (which you might want anyway, in case the Model is unable to generate a response in the desired structure): llm.with_structured_output(MyStructuredOutput, include_raw=True)
This will generate an output in the format

{
    "raw": AIMessage(
        content=[
            {
                "type": "tool_use",
                "name": "ChatBotResponse",
                "input": <the answer according to your specified data structure>,
                "id": "****"
            }
        ], 
        additional_kwargs={}, 
        response_metadata={}, 
        id="", 
        tool_calls=[], 
        usage_metadata={}
    ),
    "parsed": ChatBotResponse(<the answer according to your specified data structure>),
    "parsing_error": None
}

So what you can do is set the output_messages_key="raw" when you create the RunnableWithMessageHistory, because then you have an AIMessage() that it can use.

RunnableWithMessageHistory(
        chain,
        ...
        input_messages_key="question",
        history_messages_key="history",
        output_messages_key="raw",
   )