RunnableWithMessageHistory to output source documents in RAG

[alexmondaini]

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

vectorstore = FAISS.from_texts(
    ["harrison worked at kensho"], embedding=OpenAIEmbeddings(api_key=)
    )
    retriever = vectorstore.as_retriever()

    ### Contextualize question ###
    contextualize_q_system_prompt = """Given a chat history and the latest user question \
    which might reference context in the chat history, formulate a standalone question \
    which can be understood without the chat history. Do NOT answer the question, \
    just reformulate it if needed and otherwise return it as is."""
    contextualize_q_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", contextualize_q_system_prompt),
            MessagesPlaceholder("history"),
            ("human", "{input}"),
        ]
    )
    history_aware_retriever = create_history_aware_retriever(
        llm, retriever, contextualize_q_prompt
    )


    ### Answer question ###
    qa_system_prompt = """You are an assistant for question-answering tasks. \
    Use the following pieces of retrieved context to answer the question. \
    If you don't know the answer, just say that you don't know. \
    Use three sentences maximum and keep the answer concise and outpout the source document under the context key for which you based your answer.\

    {context}"""
    qa_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", qa_system_prompt),
            MessagesPlaceholder("history"),
            ("human", "{input}"),
        ]
    )
    question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)

    rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

    return {
        'input': lambda x: str(x)
    } | RunnableWithMessageHistory(
        rag_chain,  # type: ignore
        get_session_history=history_factory,
        input_messages_key="input",
        history_messages_key="history",
        output_messages_key="answer"
    )

Description

I would like to output the source document as part of the answer when using RunnableWithMessageHistory, from the class signature I noticed only one key as in output_messages_key can be defined, so how would I be ale to also return the source document in this case ?

System Info

langchain @ file:///home/conda/feedstock_root/build_artifacts/langchain_1711012871747/work
langchain-cli==0.0.21
langchain-community @ file:///home/conda/feedstock_root/build_artifacts/langchain-community_1711008261310/work
langchain-core @ file:///home/conda/feedstock_root/build_artifacts/langchain-core_1711659762823/work
langchain-openai @ file:///home/conda/feedstock_root/build_artifacts/langchain-openai_1710403027641/work
langchain-text-splitters @ file:///home/conda/feedstock_root/build_artifacts/langchain-text-splitters_1709389732771/work
langchainhub==0.1.15

[dosubot[bot]]

Hey @alexmondaini, great to see you diving deeper into LangChain! How's everything going on your end?

To modify your implementation of RunnableWithMessageHistory in LangChain to output both the answer and the source document, you can utilize the flexibility of returning a dictionary as output from the wrapped Runnable. Since the RunnableWithMessageHistory class allows for the output to be a dictionary with a key for a BaseMessage or sequence of BaseMessages, you can extend this functionality to include both the answer and the source document in the output dictionary.

Here's a strategy to achieve this:

1. Modify the Wrapped Runnable's Output: Ensure that the Runnable wrapped by RunnableWithMessageHistory returns a dictionary containing both the answer and the source document. The keys for the answer and the source document should be distinct.

2. Specify the output_messages_key Appropriately: When initializing RunnableWithMessageHistory, specify the output_messages_key to point to the key in the output dictionary that contains the answer. This is necessary because the _get_output_messages method uses the output_messages_key to extract the relevant part of the output for further processing.

3. Access the Source Document: Since the source document is also part of the output dictionary, you can access it directly from the output of RunnableWithMessageHistory.invoke() or within any subsequent processing steps.

Here is an illustrative example based on the provided context:

# Step 1: Modify the wrapped Runnable to return a dictionary with both answer and source document
class CustomRunnable(Runnable):
    async def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> dict:
        # Your logic here to generate answer and source document
        answer = "This is the answer."
        source_document = "This is the source document."
        return {"answer": answer, "source_document": source_document}

# Initialize the CustomRunnable
custom_runnable = CustomRunnable()

# Step 2: Wrap the CustomRunnable with RunnableWithMessageHistory
# Specify 'answer' as the output_messages_key since it contains the answer part of the output
runnable_with_history = RunnableWithMessageHistory(
    custom_runnable,
    get_session_history=get_by_session_id,  # Assuming get_by_session_id is defined as in the context
    output_messages_key="answer",  # Key in the output dict that contains the answer
)

# Invoke the runnable_with_history and access both the answer and the source document
output = runnable_with_history.invoke(input_data, config)
answer = output["answer"]  # This is extracted and processed by _get_output_messages
source_document = output["source_document"]  # Directly access the source document from the output

In this example, CustomRunnable is a placeholder for your actual Runnable implementation that needs to return a dictionary with both the answer and the source document. The key point is to structure the output of the wrapped Runnable to include both pieces of information and then access them as needed after invoking RunnableWithMessageHistory.

For more details, you can refer to the RunnableWithMessageHistory class in the LangChain codebase.

Sources

• libs/core/langchain_core/runnables/history.py
• libs/core/tests/unit_tests/runnables/test_history.py

About Dosu

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[DhavalThkkar]

@alexmondaini Were you able to generate sources? I don't think the above implementation from dosubot helps

[alexmondaini]

Hi @DhavalThkkar , I was able to generate sources by applying concatenation of strings from the dictionary lookup entries as in:

answer = output["answer"]  # This is extracted and processed by _get_output_messages
source_document = output["source_document"]  # Directly access the source document from the output

However the only catch here is that this will also return documents for a greetings message in a chatbot, which makes it odd when a user says hello for example and the bot outputs source documents for it.

[niruka1994]

        chain = (
                {"context": itemgetter("question") | ensemble_retriever | format_docs,
                 "question": itemgetter("question"), "history": itemgetter("history")}
                | prompt
                | llm
        )

        chain_with_memory = (
            RunnableWithMessageHistory(
                chain,
                get_session_history,
                input_messages_key="question",
                history_messages_key="history",
            )
        )

        response = chain_with_memory.invoke(
            {"question": user_input}, cfg
        )

        msg = response.content

My source is the same as above, how did you apply source_document?

[sharathsadari]

Hey!

i am also trying to get the source document. Did you find a way to do it?

[kyle-skillins-regn]

@alexmondaini What did you end up converting to the custom runnable? Or, where did you get your rag_chain set up to feed into the CustomRunnable?