Router with a Query Pipeline

[sumitbindra]

Hello,

I have an existing query pipeline built over tabular data using llama index text to sql modules.

I want to add a router that can decide between the data and source document for this data (say a data dictionary) to answer user questions.

Current examples all user query engines but can I use query pipeline itself as one of the selector items?

Thanks.

[dosubot[bot]]

Yes, you can use the query pipeline itself as one of the selector items in a router that decides between querying the data and querying the source document for this data. Here's an example of how to construct query pipelines for different purposes and use a router component to dynamically choose between these pipelines based on the query:

from llama_index.core.query_pipeline import RouterComponent, QueryPipeline

# Define your existing query pipeline for tabular data
tabular_data_pipeline = QueryPipeline(chain=[your_tabular_data_query_engine])

# Define another query pipeline for the source document (e.g., data dictionary)
source_document_pipeline = QueryPipeline(chain=[your_source_document_query_engine])

# Define the choices for the router
choices = [
    "This tool answers specific questions about the tabular data",
    "This tool answers questions about the source document (data dictionary)",
]

# Create the router component
router_c = RouterComponent(
    selector=your_selector_function,
    choices=choices,
    components=[tabular_data_pipeline, source_document_pipeline],
    verbose=True,
)

# Create the top-level query pipeline with the router component
qp = QueryPipeline(chain=[router_c], verbose=True)

In this setup, tabular_data_pipeline and source_document_pipeline are query pipelines used as selector items in the RouterComponent. The router decides which pipeline to use based on the query, ensuring that the appropriate query pipeline is chosen dynamically ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

I assume then I can also do a mix of a query pipeline and a query engine?

[dosubot[bot]]

Yes, you can mix a query pipeline and a query engine within the same router setup. Here is an example of how to achieve this:

from llama_index.core.query_pipeline import QueryPipeline, RouterComponent
from llama_index.llms.openai import OpenAI
from llama_index.core import Document, VectorStoreIndex, SummaryIndex
from llama_index.core.response_synthesizers import TreeSummarize
from llama_index.core.schema import NodeWithScore, TextNode
from llama_index.core import PromptTemplate
from llama_index.core.selectors import LLMSingleSelector

# Define HyDE template
hyde_str = """\
Please write a passage to answer the question: {query_str}

Try to include as many key details as possible.

Passage: """
hyde_prompt = PromptTemplate(hyde_str)

# Define LLM
llm = OpenAI(model="gpt-3.5-turbo")

# Define synthesizer
summarizer = TreeSummarize(llm=llm)

# Load documents
reader = SimpleDirectoryReader(input_files=["pg_essay.txt"])
documents = reader.load_data()

# Define vector retriever
vector_index = VectorStoreIndex.from_documents(documents)
vector_query_engine = vector_index.as_query_engine(similarity_top_k=2)

# Define summary query prompts + retrievers
summary_index = SummaryIndex.from_documents(documents)
summary_qrewrite_str = """\
Here's a question:
{query_str}

You are responsible for feeding the question to an agent that given context will try to answer the question.
The context may or may not be relevant. Rewrite the question to highlight the fact that
only some pieces of context (or none) maybe be relevant.
"""
summary_qrewrite_prompt = PromptTemplate(summary_qrewrite_str)
summary_query_engine = summary_index.as_query_engine()

# Define selector
selector = LLMSingleSelector.from_defaults()

# Define query pipelines
vector_chain = QueryPipeline(chain=[vector_query_engine])
summary_chain = QueryPipeline(
    chain=[summary_qrewrite_prompt, llm, summary_query_engine], verbose=True
)

choices = [
    "This tool answers specific questions about the document (not summary questions across the document)",
    "This tool answers summary questions about the document (not specific questions)",
]

router_c = RouterComponent(
    selector=selector,
    choices=choices,
    components=[vector_chain, summary_chain],
    verbose=True,
)

# Top-level pipeline
qp = QueryPipeline(chain=[router_c], verbose=True)

# Try out queries
response = qp.run("What did the author do during his time in YC?")
print(str(response))

response = qp.run("What is a summary of this document?")
print(str(response))

This setup demonstrates how to mix a query pipeline and a query engine within the same router setup, allowing you to dynamically choose the appropriate query pipeline based on the query ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

[dosubot[bot]]

To incorporate a multi-step query engine (run_multi) with your setup, you need to instantiate the MultiStepQueryEngine class with the appropriate parameters. Here's an example of how you can do this:

from llama_index.core.query_engine import MultiStepQueryEngine
from llama_index.legacy.indices.query.query_transform.base import StepDecomposeQueryTransform
from llama_index.legacy.response_synthesizers import BaseSynthesizer, get_response_synthesizer
from llama_index.legacy.core.base_query_engine import BaseQueryEngine
from llama_index.legacy.schema import QueryBundle

# Assuming you have instances of the required components
base_query_engine = BaseQueryEngine()  # Replace with your actual base query engine instance
query_transform = StepDecomposeQueryTransform()  # Replace with your actual query transform instance
response_synthesizer = BaseSynthesizer()  # Replace with your actual response synthesizer instance

# Create the MultiStepQueryEngine instance
multi_step_query_engine = MultiStepQueryEngine(
    query_engine=base_query_engine,
    query_transform=query_transform,
    response_synthesizer=response_synthesizer,
    num_steps=3,  # Number of steps to run the multi-step query
    early_stopping=True,  # Whether to stop early if the stop function returns True
    index_summary="Your index summary here",  # Provide your index summary
    stop_fn=None  # Optional: Provide a custom stop function if needed
)

# Now you can use the multi_step_query_engine to run multi-step queries
query_bundle = QueryBundle(query_str="Your query here")  # Replace with your actual query bundle
response = multi_step_query_engine.query(query_bundle)
print(response)

You can integrate this multi_step_query_engine into your existing router setup by adding it as one of the components in the RouterComponent. Here's how you can modify your existing setup to include the multi-step query engine:

# Define your existing query pipeline for tabular data
tabular_data_pipeline = QueryPipeline(chain=[your_tabular_data_query_engine])

# Define another query pipeline for the source document (e.g., data dictionary)
source_document_pipeline = QueryPipeline(chain=[your_source_document_query_engine])

# Define the multi-step query engine
multi_step_query_engine = MultiStepQueryEngine(
    query_engine=base_query_engine,
    query_transform=query_transform,
    response_synthesizer=response_synthesizer,
    num_steps=3,
    early_stopping=True,
    index_summary="Your index summary here",
    stop_fn=None
)

# Define the choices for the router
choices = [
    "This tool answers specific questions about the tabular data",
    "This tool answers questions about the source document (data dictionary)",
    "This tool runs a multi-step query"
]

# Create the router component
router_c = RouterComponent(
    selector=your_selector_function,
    choices=choices,
    components=[tabular_data_pipeline, source_document_pipeline, multi_step_query_engine],
    verbose=True,
)

# Create the top-level query pipeline with the router component
qp = QueryPipeline(chain=[router_c], verbose=True)

# Try out queries
response = qp.run("Your query here")
print(str(response))

This setup allows you to dynamically choose between the tabular data pipeline, the source document pipeline, and the multi-step query engine based on the query ^[1][2].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

how do i incorporate a multi step query pipeline (run_multi) with this?

[dosubot[bot]]

To incorporate a multi-step query pipeline (run_multi) into your setup, you need to instantiate the MultiStepQueryEngine class and integrate it into your router component. Here's how you can do this:

1. Define the Multi-Step Query Engine: Create an instance of MultiStepQueryEngine with the necessary parameters.

2. Integrate the Multi-Step Query Engine into the Router: Add the multi-step query engine as one of the components in the RouterComponent.

Here's an example:

from llama_index.core.query_pipeline import QueryPipeline, RouterComponent
from llama_index.core.query_engine import MultiStepQueryEngine
from llama_index.legacy.indices.query.query_transform.base import StepDecomposeQueryTransform
from llama_index.legacy.response_synthesizers import BaseSynthesizer, get_response_synthesizer
from llama_index.legacy.core.base_query_engine import BaseQueryEngine
from llama_index.legacy.schema import QueryBundle
from llama_index.core.selectors import LLMSingleSelector

# Define your existing query pipeline for tabular data
tabular_data_pipeline = QueryPipeline(chain=[your_tabular_data_query_engine])

# Define another query pipeline for the source document (e.g., data dictionary)
source_document_pipeline = QueryPipeline(chain=[your_source_document_query_engine])

# Define the multi-step query engine
base_query_engine = BaseQueryEngine()  # Replace with your actual base query engine instance
query_transform = StepDecomposeQueryTransform()  # Replace with your actual query transform instance
response_synthesizer = BaseSynthesizer()  # Replace with your actual response synthesizer instance

multi_step_query_engine = MultiStepQueryEngine(
    query_engine=base_query_engine,
    query_transform=query_transform,
    response_synthesizer=response_synthesizer,
    num_steps=3,  # Number of steps to run the multi-step query
    early_stopping=True,  # Whether to stop early if the stop function returns True
    index_summary="Your index summary here",  # Provide your index summary
    stop_fn=None  # Optional: Provide a custom stop function if needed
)

# Define the choices for the router
choices = [
    "This tool answers specific questions about the tabular data",
    "This tool answers questions about the source document (data dictionary)",
    "This tool runs a multi-step query"
]

# Define the selector
selector = LLMSingleSelector.from_defaults()

# Create the router component
router_c = RouterComponent(
    selector=selector,
    choices=choices,
    components=[tabular_data_pipeline, source_document_pipeline, multi_step_query_engine],
    verbose=True,
)

# Create the top-level query pipeline with the router component
qp = QueryPipeline(chain=[router_c], verbose=True)

# Try out queries
query_bundle = QueryBundle(query_str="Your query here")  # Replace with your actual query
response = qp.run(query_bundle)
print(str(response))

This setup allows you to dynamically choose between the tabular data pipeline, the source document pipeline, and the multi-step query engine based on the query ^[1][2].

To continue talking to Dosu, mention @dosu.