I want to create history aware chatbot that uses self-query to retrieve data from chromadb with custom context every time it retrieves the data. When asking question related to a previous one (stored in history store {}) no retrieving should happen, when answer can’t be find in history the retriever should get the data and generate the response.

I have the following code

retriever = SelfQueryRetriever.from_llm(
    llm,
    vectorstore,
    document_content_description,
    metadata_field_info,
    search_kwargs={"k": 5}
)

def format_metadata(result):
    metadata_context = "n".join([
        f"id: {doc.metadata['id']}, category: {doc.metadata['category']}, subcategory: {doc.metadata['subcategory']}, name: {doc.metadata['name']}"
        for doc in result
    ])
    return metadata_context

def self_query(query_text):
    model = llm
    try:
        result = retriever.invoke(query_text)
        logger.info("Data retrieved from database.")
        if len(result) == 0:
            logger.info(f"Unable to find matching results.")
            return "Unable to find matching results."
    except Exception as e:
        return ({'Retrieval failed: ': str(e)})
    return result

class InMemoryHistory(BaseChatMessageHistory, BaseModel):
    """In memory implementation of chat message history."""

    messages: List[BaseMessage] = Field(default_factory=list)

    def add_messages(self, messages: List[BaseMessage]) -> None:
        """Add a list of messages to the store"""
        self.messages.extend(messages)

    def clear(self) -> None:
        self.messages = []

def get_answer(query_text, metadata_context):

    def get_session_history(
        user_id: str, conversation_id: str
    ) -> BaseChatMessageHistory:
        if (user_id, conversation_id) not in store:
            store[(user_id, conversation_id)] = InMemoryHistory()
        return store[(user_id, conversation_id)]

    prompt = ChatPromptTemplate.from_messages([
        (PROMPT_TEMPLATE),
        ("system", "{context}"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}"),
    ])

    contextualize_q_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", contextualize_q_system_prompt),
            ("system", "{context}"),
            MessagesPlaceholder("chat_history"),
            ("human", "{input}"),
        ]
    )

    history_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)

    question_answer_chain = create_stuff_documents_chain(llm, prompt)
    rag_chain = create_retrieval_chain(history_retriever, question_answer_chain)

    with_message_history = RunnableWithMessageHistory(
        rag_chain,
        get_session_history=get_session_history,
        input_messages_key="input",
        history_messages_key="history",
        history_factory_config=[
            ConfigurableFieldSpec(
                id="user_id",
                annotation=str,
                name="User ID",
                description="Unique identifier for the user.",
                default="",
                is_shared=True,
            ),
            ConfigurableFieldSpec(
                id="conversation_id",
                annotation=str,
                name="Conversation ID",
                description="Unique identifier for the conversation.",
                default="",
                is_shared=True,
            ),
        ],
    )

    response = with_message_history.invoke(
        {"context": itemgetter(metadata_context), "input": query_text},
        config={"configurable": {"user_id": "123", "conversation_id": "1"}}
    )
    return response


query_text = "I am looking for id 4."
query_text_1 = "What did I ask you about?"
query_text_2 = "I am looking for id 10."

result = self_query(query_text)
metadata_context = format_metadata(result)

print(get_answer(query_text, metadata_context))
print(get_answer(query_text_1, metadata_context))
print(get_answer(query_text_2, metadata_context))

Retrieving of data works well, the thing is that the context is not formatted and does not have the form of the the metadata_context and the answer generated is wrong. I also tried the following option:

def get_answer(query_text, metadata_context):

    def get_session_history(
        user_id: str, conversation_id: str
    ) -> BaseChatMessageHistory:
        if (user_id, conversation_id) not in store:
            store[(user_id, conversation_id)] = InMemoryHistory()
        return store[(user_id, conversation_id)]

    prompt = ChatPromptTemplate.from_messages([
        (PROMPT_TEMPLATE),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ])

    chain = prompt | llm

    with_message_history = RunnableWithMessageHistory(
        chain,
        get_session_history=get_session_history,
        input_messages_key="question",
        history_messages_key="history",
        history_factory_config=[
            ConfigurableFieldSpec(
                id="user_id",
                annotation=str,
                name="User ID",
                description="Unique identifier for the user.",
                default="",
                is_shared=True,
            ),
            ConfigurableFieldSpec(
                id="conversation_id",
                annotation=str,
                name="Conversation ID",
                description="Unique identifier for the conversation.",
                default="",
                is_shared=True,
            ),
        ],
    )

    response = with_message_history.invoke(
        {"context": metadata_context, "question": query_text},
        config={"configurable": {"user_id": "123", "conversation_id": "1"}}
    )
    return response

query_text = "I am looking for id 4."
query_text_1 = "What did I ask you about?"
query_text_2 = "I am looking for id 10."

result = self_query(query_text)
metadata_context = format_metadata(result)

print(get_answer(query_text, metadata_context))
print(get_answer(query_text_1, metadata_context))
print(get_answer(query_text_2, metadata_context))

Which gets the right first answer, respond to the second one based on the stored history but when asked third question it does not retrieve the data again so that is why I would like to implement history_aware_retriever.