Use a RouterQueryEngine as a chat

[KirianCaumes]

I'm working on project that has a RouterQueryEngine based on PDF files and SQLite DB.

I have multiple QueryEngineTool based on both BaseQueryEngine and NLSQLTableQueryEngine.

If I want to ask a question on my RouterQueryEngine via the query function it works perfectly.

But now, I would like to use the same system but as a chat where this chat keep memory of the previous messages.

Is this something possible? How can I achieve it?

I tried multiple things (with OpenAIAgent, ChatMemoryBuffer, and so on), but I'm really stuck on this point 😕

Here my (simplified) code bellow:

from __future__ import annotations
from llama_index.core.query_engine import NLSQLTableQueryEngine
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core import SQLDatabase
from llama_index.llms.azure_openai import AzureOpenAI
from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector
from llama_index.core import Settings
from dotenv import load_dotenv
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.agent.openai import OpenAIAgent

files = []

files_base_path = "./data/downloaded"

if __name__ == '__main__':
    load_dotenv()

    # Init LLM and store it globally
    Settings.llm = AzureOpenAI()

    # Init EMBED and store it globally
    Settings.embed_model = AzureOpenAIEmbedding()

    # List of tables that will be used
    sql_tables = [table_name]

    # SQL query engine to generate SQL Query
    sql_query_engine = NLSQLTableQueryEngine(
        sql_database=SQLDatabase(
            engine=engine,
            include_tables=sql_tables
        ),
        tables=sql_tables,
    )

    # Vector stores
    vector_stores: list[VectorStoreIndex] = []

    # Load downloaded file
    documents = SimpleDirectoryReader(
        input_dir=files_base_path
    ).load_data()

    # For each documents found, generate a Vector store
    for document in documents:
        # Create and automatically save vector
        vector_store = VectorStoreIndex.from_documents(
            documents=[document],
        )

        vector_stores.append(vector_store)

    # Vector Query Engine from Vector Stores
    vector_query_engines = [
        index.as_query_engine()
        for index in vector_stores
    ]

    # Vector Query Engine Tool from Vector Query Engine
    vector_query_engine_tools = [
        QueryEngineTool.from_defaults(
            query_engine=query_engine,
            name=f"..._engine",
            description=f"...",
        )
        for file, query_engine in zip(files, vector_query_engines)
    ]

    # SQL Query Engine Tool from SQL Query Engine
    sql_query_engine_tool = QueryEngineTool(
        query_engine=sql_query_engine,
        metadata=ToolMetadata(
            name="sql_query_engine",
            description="...",
        ),
    )

    # Router query engine based of all query engine tools
    router_query_engine = RouterQueryEngine(
        selector=LLMSingleSelector.from_defaults(),
        query_engine_tools=(
            [sql_query_engine_tool] +
            vector_query_engine_tools
        ),
        verbose=True,
    )

    while True:
        question = input("> ")
        if question == "exit":
            break
        response = router_query_engine.query(question) # TODO question here
        print(str(response))

[dosubot[bot]]

Hey @KirianCaumes, back for more adventures in coding, I see! 🚀

Based on your requirements, it seems like you want to maintain a history of previous messages in your chat system. This can be achieved by using a memory buffer that keeps track of the conversation history.

In your current setup, you're using the RouterQueryEngine to handle queries. However, to maintain a chat history, you'll need to use a chat engine that supports memory, such as the SimpleChatEngine class in the llama_index/core/chat_engine/simple.py file.

Here's a simplified example of how you can modify your code to use SimpleChatEngine:

from llama_index.core.chat_engine.simple import SimpleChatEngine
from llama_index.core.memory import ChatMemoryBuffer

# Initialize memory buffer
memory = ChatMemoryBuffer.from_defaults()

# Initialize chat engine
chat_engine = SimpleChatEngine(llm=Settings.llm, memory=memory)

while True:
    question = input("> ")
    if question == "exit":
        break
    response = chat_engine.chat(question)
    print(str(response))

In this example, ChatMemoryBuffer is used to store the chat history. The SimpleChatEngine is initialized with this memory buffer and your LLM instance. When you call the chat method, it automatically updates the memory with the user's message and the AI's response. You can access the chat history at any point using the chat_history property of the SimpleChatEngine instance.

Please note that this is a simplified example and you might need to adjust it according to your project's structure and requirements. For more details on how SimpleChatEngine and ChatMemoryBuffer work, you can refer to the llama_index/core/chat_engine/simple.py and llama_index/core/memory.py files in the LlamaIndex repository.

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

Sources

• llama-index-core/llama_index/core/chat_engine/simple.py
• llama-index-legacy/llama_index/legacy/query_engine/router_query_engine.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.

[KirianCaumes]

@dosu-bot As you explained it, if I use SimpleChatEngine and ChatMemoryBuffer, I will not be able to use any knowledge base (my RouterQueryEngine in that case). It is mentionned in the documentation : "the SimpleChatEngine does not make use of a knowledge base, whereas all others make use of a query engine over knowledge base." (sourc : https://docs.llamaindex.ai/en/stable/module_guides/deploying/chat_engines/modules.html#module-guides).

Is there a way to combine RouterQueryEngine with any chat engine ?

[dosubot[bot]]

Based on the information you've provided, it seems you're trying to integrate a chat engine with the RouterQueryEngine in the LlamaIndex framework. As you've correctly pointed out, the SimpleChatEngine and ChatMemoryBuffer are not compatible with the RouterQueryEngine as per the documentation.

However, there are other chat engines in the LlamaIndex framework that you might consider using. These include CondensePlusContextChatEngine, CondenseQuestionChatEngine, and ContextChatEngine. These engines are designed to handle more complex chat scenarios and might be compatible with the RouterQueryEngine.

To integrate the RouterQueryEngine with one of these chat engines, you would typically need to:

1. Initialize the RouterQueryEngine and make it available within the chat engine. This could be done by passing an instance of the RouterQueryEngine to the chat engine's constructor or making it accessible through a service context or global settings.

2. Modify the chat engine's methods to use the RouterQueryEngine for processing queries or retrieving relevant information. This could involve modifying methods such as _retrieve_context or _condense_question to utilize the RouterQueryEngine.

3. Adjust the chat flow in methods like chat, stream_chat, achat, and astream_chat to accommodate the RouterQueryEngine.

4. Ensure that any necessary configuration or settings required by the RouterQueryEngine are properly set up.

5. Thoroughly test the chat engine to ensure that the integration works as expected.

Please note that the exact steps might vary depending on the specific functionalities and interfaces of the RouterQueryEngine and the chosen chat engine.

As for your current code, you might need to modify the while loop at the end to integrate the chat engine. Instead of directly querying the RouterQueryEngine, you would need to pass the user's input to the chat engine, which would then use the RouterQueryEngine to process the query and generate a response.

I hope this helps! If you have any further questions or need more specific guidance, please let me know.

Sources

• llama-index-core/llama_index/core/chat_engine/init.py
• llama-index-core/llama_index/core/chat_engine/condense_plus_context.py
• llama-index-core/llama_index/core/query_engine/router_query_engine.py
• llama-index-core/llama_index/core/memory/chat_memory_buffer.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.

[aminolroayaXgen]

@dosu-bot can you write the code about what you mean?

[bolaka1111]

Hey @KirianCaumes! I am also checking out on this requirement wherein I need the RouterQueryEngine keep track of the chat history. Or if there is a way to switch between tools within the Chat Engine modes. Let me know if you were able to find any solution on this yet!

[KirianCaumes]

Hey! Sadly no, I haven't found a solution, and my research is on pause for the moment 😕

[parallelepipede]

Hey @KirianCaumes

I found an alternative solution. I am not sure it can be applied to your problem nonetheless.

Instead of creating a RouterQueryEngine with chat mode, I create different retrievertools and a routerRetriever. Then I apply this new retriever to a ContextChatEngine object.

Hope it can help

[KirianCaumes]

Hey @parallelepipede

I think your right, I just tried something quickly and it seems to be working, thank you! 🤩

I'm definitely going to dig into this solution, but here the code I tried so far (I suppose you have something that looks like this too):

base_retrievers = [
    index.as_retriever()
    for index in vector_stores
]

retriever_tools = [
    RetrieverTool.from_defaults(
        retriever=base_retriever,
        name="Something",
        description=f"A description",
    )
    for base_retriever in base_retrievers
]

retriever = RouterRetriever(
    selector=LLMSingleSelector.from_defaults(),
    retriever_tools=retriever_tools,
)

context_chat_engine = ContextChatEngine.from_defaults(
    retriever=retriever
)

response = context_chat_engine.chat("A question...")

print(str(response))

[parallelepipede]

However It doesn't take into account the chat history to select the correct choice with the LLMSingleSelector. There is need to add previous questions for the LLMSingleSelector or outputParser for example.

[Rumeysakeskin]

Hello @KirianCaumes ,
Your solution worked! For reference, here is my code:

from llama_index.core import SummaryIndex
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    SimpleKeywordTableIndex,
)

# load documents
documents = SimpleDirectoryReader("./data/").load_data()
from llama_index.core import Settings

# initialize settings (set chunk size)
Settings.chunk_size = 1024
nodes = Settings.node_parser.get_nodes_from_documents(documents)

# initialize storage context (by default it's in-memory)
summary_index = SummaryIndex(nodes, storage_context=storage_context)
vector_index = VectorStoreIndex(nodes, storage_context=storage_context)
keyword_index = SimpleKeywordTableIndex(nodes, storage_context=storage_context)

list_retriever = summary_index.as_retriever()
vector_retriever = vector_index.as_retriever()
keyword_retriever = keyword_index.as_retriever()


from llama_index.core.tools import RetrieverTool

list_tool = RetrieverTool.from_defaults(
    retriever=list_retriever,
    description=(
        "Will retrieve all context"
    ),
)
vector_tool = RetrieverTool.from_defaults(
    retriever=vector_retriever,
    description=(
        "Useful for retrieving specific context"
    ),
)
keyword_tool = RetrieverTool.from_defaults(
    retriever=keyword_retriever,
    description=(
        "Useful for retrieving specific context "
    ),
)


from llama_index.core.selectors import LLMSingleSelector, LLMMultiSelector
from llama_index.core.selectors import (
    PydanticMultiSelector,
    PydanticSingleSelector,
)
from llama_index.core.retrievers import RouterRetriever
from llama_index.core.response.notebook_utils import display_source_node

retriever = RouterRetriever(
    selector=LLMSingleSelector.from_defaults(llm=llm),
    retriever_tools=[
        list_tool,
        vector_tool,
    ],
)
from llama_index.core.chat_engine import ContextChatEngine

context_chat_engine = ContextChatEngine.from_defaults(
    retriever=retriever
)
response = context_chat_engine.stream_chat("Question...")
for token in response.response_gen:
    print(token, end="")

Sources:

• https://docs.llamaindex.ai/en/stable/module_guides/querying/router/?h=routerretriever#using-as-a-retriever
• https://docs.llamaindex.ai/en/stable/examples/chat_engine/chat_engine_context/?h=contextchatengine
• https://docs.llamaindex.ai/en/stable/examples/retrievers/router_retriever/?h=router

[aiworld12345]

Thanks @Rumeysakeskin , this is helpful but using "as_retriever()" somehow reduced my accuracy of information retrieval.

[aiworld12345]

In the following function i am not able to save the conversation history, also the prompt template is not working , i get answers outside of the context given . can anyone help me out this .

def assistant(verbose: bool = True):
Settings.llm = OpenAI(model="gpt-3.5-turbo")
summary_storage = StorageContext.from_defaults(persist_dir="storage/summaryindex")
qna_storage = StorageContext.from_defaults(persist_dir="storage/meeting_index")
embed_model = HuggingFaceEmbedding(model_name="model")

summary_index = load_index_from_storage(summary_storage)
vector_index = load_index_from_storage(qna_storage,embed_model=embed_model)

#Adding prompt template
template = (
"We have provided context information below. \n"
"---------------------\n"
"{context_str}"
"\n---------------------\n"
"Given this information and not prior knowledge, please answer the question: {query_str}\n"
"You must not give answers to any questions that are not related to the {context_str} .\n"
"If the user asks out of context then you must tell that you do not know the answer.\n"
)
query_template = PromptTemplate(template)


summary_query_engine = summary_index.as_query_engine(query_template=query_template,
    response_mode="tree_summarize",
    use_async=True,
)
retrival_query_engine = vector_index.as_query_engine(query_template=query_template,
similarity_top_k=3,
node_postprocessors=[
    MetadataReplacementPostProcessor(target_metadata_key="window")
],

)

summary_tool = QueryEngineTool.from_defaults(
    query_engine=summary_query_engine,
    description="Useful for summarization questions related to the document."
    "Do not give any information that is not given in the context"
)
vector_tool = QueryEngineTool.from_defaults(
    query_engine=retrival_query_engine,
    description="Useful for retrieving specific information from context given in document."
    "Do not give any information that is not given in the context"
)

engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=[summary_tool, vector_tool],
    verbose=verbose
)
return  engine

query_engine = assistant()
user_input= input("Q: ")
response = query_engine.query(user_input)
print(response)