RunnableWithMessageHistory and DynamoDBChatMessageHistory

[Pairrot-Lore]

Hi all,

I'm currently working on a Langserve API and need help integrating chat history with AWS DynamoDB using RunnableWithMessageHistory and DynamoDBChatMessageHistory. Despite my efforts, I’m struggling to combine the necessary chains and successfully store chat history in DynamoDB.

Could anyone share examples, documentation, or advice on the best practices to accomplish this integration?

The error that I am currently having is: Error in RootListenersTracer.on_chain_end callback: KeyError('question')

Below is the relevant code snippet for reference:

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "Use this context to answer the question {context}"), 
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ]
)

class InputChat(BaseModel):
    """Input for the chat endpoint."""

    # The field extra defines a chat widget.
    # As of 2024-02-05, this chat widget is not fully supported.
    # It's included in documentation to show how it should be specified, but
    # will not work until the widget is fully supported for history persistence
    # on the backend.
    human_input: str = Field(
        ...,
        description="The human input to the chat system.",
        extra={"widget": {"type": "chat", "input": "human_input"}},
    )

# Create a chatbot Question & Answer chain from the retriever
def format_docs(docs):
    formatted_docs = "\n\n".join(f"{doc.page_content}\n\nMetadata: {doc.metadata}" for doc in docs)

rag_chain_from_docs = (
    RunnablePassthrough.assign(
        context=(lambda x: format_docs(x["context"])[0]))
    | prompt
    | model
)

rag_chain_with_source = RunnableParallel(
    {"context": retriever , "question": RunnablePassthrough(), "history": lambda x: x["history"]}
).assign(answer=rag_chain_from_docs)

chain_with_history = RunnableWithMessageHistory(
    rag_chain_with_source,
    lambda session_id: DynamoDBChatMessageHistory(
        table_name="p2024001SessionTable", session_id=session_id, boto3_session=boto3.Session(region_name='us-east-1')
    ),
    input_messages_key="question",
    history_messages_key="history",
).with_types(input_type=InputChat)

# This is where we configure the session id
config = {"configurable": {"session_id": "{session_id}"}}

add_routes(
    app,
    chain_with_history.with_config(),
    path="/knowledge",
    config_keys=["configurable"]
)