How to configure different retrievals with RunnableWithMessageHistory in a RAG pipeline

[strivereep]

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

QA_SYSTEM_PROMPT = """You are assistant. Answer the questions based on the following context.
{context}

Human Input: {human_input}
"""

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)


vectorstore1 = Chroma.from_texts(
    ["cats like fish", "dogs like sticks"], embedding=OpenAIEmbeddings()
)

vectorstore2 = Chroma.from_texts(["x_n+1=a * xn * (1-xn)"], embedding=OpenAIEmbeddings())


# need help here
def get_retriever(option: str):
    if option == 'index1':
        return vectorstore1.as_retriever()
    else:
        return vectorstore2.as_retriever()


qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", QA_SYSTEM_PROMPT),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{human_input}"),
    ]
)


context = itemgetter("human_input") | get_retriever | format_docs

chain = RunnablePassthrough.assign(context=context) | qa_prompt | ChatOpenAI()


# User input
class InputChat(BaseModel):
    """Input for the chat endpoint."""
    human_input: str = Field(
        ...,
        description="The human input to the chat system.",
        extra={
            "widget": {
                "type": "chat",
                "input": "human_input"
            }
        },
    )
    option: str = Field(
        ...,
        description="The option to select from the vector store.",
    )


chain_with_history = RunnableWithMessageHistory(
    chain,
    lambda session_id: RedisChatMessageHistory(session_id, url=REDIS_URL),
    input_messages_key="human_input",
    history_messages_key="history",
).with_types(input_type=InputChat)

res = chain_with_history.invoke({"human_input": "Hello!", "option": "index1"}, {"configurable": {"session_id": session_id}})
print(res)

Description

I have created chain with history with redis as primary source of memory for storing the chat histories. Everythin is working fine but now I need to pass extra params other than human_input, i.e. say option. And based upon the option params, I want to switch the retriever i.e. if user passes index1, then use vectorstore1 else vectorstore2.
Not sure if this is better way to do that, any other approaches will be fine.

System Info

langchain===0.1.5
langchain-community===0.0.17
langchain_openai===0.0.5
langchain_core=== 0.1.18
langserve===0.0.41

[dosubot[bot]]

🤖

Hey @strivereep! Great to see you diving into more adventures with LangChain. How's the coding journey going?

Based on your code and requirements, it seems like you want to switch between different retrievers based on an 'option' parameter. This is indeed possible in LangChain.

You can modify your get_retriever function to accept the 'option' parameter from the input dictionary. Here's how you can do it:

def get_retriever(input_dict: dict):
    option = input_dict.get('option', 'default_option')
    if option == 'index1':
        return vectorstore1.as_retriever()
    else:
        return vectorstore2.as_retriever()

In this modified function, get_retriever now accepts an input dictionary and retrieves the 'option' parameter from it. If 'option' is not provided, it defaults to 'default_option'.

Then, you can modify your context to pass the entire input dictionary to get_retriever:

context = itemgetter("human_input") | (lambda x: {"human_input": x, "option": option}) | get_retriever | format_docs

In this modified context, the lambda function creates a new dictionary that includes both 'human_input' and 'option', and passes this dictionary to get_retriever.

Please note that this solution assumes that the 'option' parameter is always provided in the input dictionary when invoking chain_with_history. If 'option' might not be provided, you should handle this case in your get_retriever function.

I hope this helps! If you have any further questions, feel free to ask.

Sources

• libs/langchain/langchain/chains/history_aware_retriever.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.

[strivereep]

how to pass the option as parameters in the chain_with_history?
Is this the format: chain_with_history.invoke({"human_input": "Hello!", "option": "index1"}, {"configurable": {"session_id": session_id}})

And this snippet din't work:
context = itemgetter("human_input") | (lambda x: {"human_input": x, "option": option}) | get_retriever | format_docs

[Zaeem6100]

how you resolved this @strivereep