query pipelines with router #15852

debarun-ghosh17 · 2024-09-05T11:44:10Z

debarun-ghosh17
Sep 5, 2024

I am trying to extend the Query Pipeline for Advanced Text-to-SQL guide here (https://docs.llamaindex.ai/en/stable/examples/pipeline/query_pipeline_sql/).
One query pipeline is the text to sql. While the other is a RAG based pipeline.
vector_tool_qp = QP(
modules={
"input": InputComponent(),
"table_retriever": obj_retriever,
"vector_tool_retriever": vector_tool_retriever,
"response_synthesis_prompt": response_synthesis_prompt,
"response_synthesis_llm": llm,
},
verbose=True,
)
Both are working fine individually. I want to combine them with a router. But while setting up, I am getting the following error

ValueError Traceback (most recent call last)
Cell In[182], line 8
1 from llama_index.core.query_pipeline import RouterComponent
3 choices = [
4 "This tool answers specific and mathematical questions about the document (not summary questions across the document)",
5 "This tool answers summary questions about the document (not specific questions)",
6 ]
----> 8 router_c = RouterComponent(
9 selector=selector,
10 choices=choices,
11 components=[vector_tool_qp],
12 verbose=True,
13 )
14 # top-level pipeline
15 router_qp = QueryPipeline(chain=[router_c], verbose=True)

File /opt/conda/lib/python3.10/site-packages/llama_index/core/query_pipeline/components/router.py:111, in RouterComponent.init(self, selector, choices, components, verbose)
109 # validate component has one input key
110 if len(new_component.free_req_input_keys) != 1:
--> 111 raise ValueError("Expected one required input key")
112 query_keys.append(next(iter(new_component.free_req_input_keys)))
113 new_components.append(new_component)

ValueError: Expected one required input key

Answered by dosubot[bot]

Sep 6, 2024

To resolve the error and set the input keys correctly, you should use the InputKeys class to define the required and optional keys for each component. Here’s how you can do it:

Define the Input and Output Keys for Each Pipeline: Use the InputKeys class to specify the required and optional keys.
Wrap Each Pipeline in a Component: Create a wrapper component for each pipeline that conforms to the single input key requirement.
Initialize the RouterComponent: Pass the wrapped components to the RouterComponent.

Here’s an example:

from llama_index.core.query_pipeline import RouterComponent, QueryPipeline, InputComponent
from llama_index.core.base_selector import BaseSelector
from llama_…

View full answer

@debarun-ghosh17 · 2024-09-05T11:46:27Z

dosubot[bot]
bot Sep 5, 2024

Hey there, @debarun-ghosh17! I'm here to help you with any bugs, questions, or contributions you have in mind. Let's tackle this issue together!

I found a similar unsolved discussion that might be relevant: Router with a Query Pipeline ^[1].

To combine a text-to-SQL query pipeline and a RAG-based pipeline using a RouterComponent without encountering the "Expected one required input key" error, ensure that each component in the RouterComponent has exactly one required input key. Here's a step-by-step approach:

Define the Selector: Create a selector that will decide which component to route the query to.
Ensure Single Input Key for Each Component: Make sure that both the text-to-SQL query pipeline and the RAG-based pipeline components have exactly one required input key.
Initialize the RouterComponent: Pass the selector, choices, and components to the RouterComponent.

Here is an example code snippet to illustrate this:

from llama_index.legacy.core.query_pipeline.components.router import RouterComponent
from llama_index.legacy.core.base_selector import BaseSelector
from llama_index.legacy.core.query_pipeline.query_component import QueryComponent, InputKeys, OutputKeys

# Define your selector (assuming you have a custom selector implementation)
class MySelector(BaseSelector):
    def select(self, choices, query):
        # Custom logic to select the appropriate component
        pass

# Define your text-to-SQL query component
class TextToSQLComponent(QueryComponent):
    @property
    def input_keys(self):
        return InputKeys.from_keys({"query"})

    @property
    def output_keys(self):
        return OutputKeys.from_keys({"output"})

    def _run_component(self, **kwargs):
        # Implementation of the text-to-SQL logic
        pass

# Define your RAG-based query component
class RAGComponent(QueryComponent):
    @property
    def input_keys(self):
        return InputKeys.from_keys({"query"})

    @property
    def output_keys(self):
        return OutputKeys.from_keys({"output"})

    def _run_component(self, **kwargs):
        # Implementation of the RAG-based logic
        pass

# Initialize the components
text_to_sql_component = TextToSQLComponent()
rag_component = RAGComponent()

# Initialize the selector
selector = MySelector()

# Initialize the RouterComponent
router_component = RouterComponent(
    selector=selector,
    choices=["text_to_sql", "rag"],
    components=[text_to_sql_component, rag_component],
    verbose=True
)

# Now you can use the router_component to route queries
result = router_component.run_component(query="Your query here")

In this example:

TextToSQLComponent and RAGComponent are defined with exactly one required input key (query).
MySelector is a custom selector that decides which component to use based on the query.
RouterComponent is initialized with the selector, choices, and components.

This setup ensures that the RouterComponent will not encounter the "Expected one required input key" error ^[2]^[3].

To continue talking to Dosu, mention @dosu.

10 replies

debarun-ghosh17 Sep 5, 2024
Author

I have two query pipelines, When feeding them to router component I am getting the expected one required input key error

dosubot[bot] bot Sep 5, 2024

To resolve the "Expected one required input key" error when feeding two query pipelines into a RouterComponent, you need to ensure that each query pipeline has exactly one required input key. Here’s how you can achieve this:

Define the Input and Output Keys for Each Pipeline: Ensure that each pipeline has a single required input key.
Wrap Each Pipeline in a Component: Create a wrapper component for each pipeline that conforms to the single input key requirement.
Initialize the RouterComponent: Pass the wrapped components to the RouterComponent.

Here’s an example:

from llama_index.core.query_pipeline import RouterComponent, QueryPipeline, InputComponent
from llama_index.core.base_selector import BaseSelector

# Define your selector (assuming you have a custom selector implementation)
class MySelector(BaseSelector):
    def select(self, choices, query):
        # Custom logic to select the appropriate component
        pass

# Define wrapper components for each pipeline
class TextToSQLPipelineComponent(QueryPipeline):
    @property
    def input_keys(self):
        return {"query_str"}

    @property
    def output_keys(self):
        return {"sql_query"}

class RAGPipelineComponent(QueryPipeline):
    @property
    def input_keys(self):
        return {"query_str"}

    @property
    def output_keys(self):
        return {"response"}

# Create instances of the pipelines
text_to_sql_pipeline = TextToSQLPipelineComponent(
    modules={
        "input": InputComponent(),
        "text_to_sql": text_to_sql_component,
    },
    verbose=False,
)

rag_pipeline = RAGPipelineComponent(
    modules={
        "input": InputComponent(),
        "retriever": retriever,
        "rag": rag_component,
    },
    verbose=False,
)

# Link the components in their respective pipelines
text_to_sql_pipeline.add_link("input", "text_to_sql", src_key="query_str", dest_key="query_str")
rag_pipeline.add_link("input", "retriever", src_key="query_str", dest_key="query_str")
rag_pipeline.add_link("retriever", "rag", dest_key="retrieved_docs")
rag_pipeline.add_link("input", "rag", src_key="query_str", dest_key="query_str")

# Initialize the selector
selector = MySelector()

# Initialize the RouterComponent
router_component = RouterComponent(
    selector=selector,
    choices=["text_to_sql", "rag"],
    components=[text_to_sql_pipeline, rag_pipeline],
    verbose=True
)

# Now you can use the router_component to route queries
result = router_component.run_component(query="Your query here")

In this setup:

TextToSQLPipelineComponent and RAGPipelineComponent are wrapper components that ensure each pipeline has exactly one required input key (query_str).
The RouterComponent is initialized with these wrapped components, avoiding the "Expected one required input key" error ^[1].

To continue talking to Dosu, mention @dosu.

debarun-ghosh17 Sep 6, 2024
Author

'set' object has no attribute 'required_keys'. Seems like the expected format for input keys is InputKeys(required_keys=set(), optional_keys=set())

dosubot[bot] bot Sep 6, 2024

To resolve the error and set the input keys correctly, you should use the InputKeys class to define the required and optional keys for each component. Here’s how you can do it:

Define the Input and Output Keys for Each Pipeline: Use the InputKeys class to specify the required and optional keys.
Wrap Each Pipeline in a Component: Create a wrapper component for each pipeline that conforms to the single input key requirement.
Initialize the RouterComponent: Pass the wrapped components to the RouterComponent.

Here’s an example:

from llama_index.core.query_pipeline import RouterComponent, QueryPipeline, InputComponent
from llama_index.core.base_selector import BaseSelector
from llama_index.core.query_pipeline.query import InputKeys, OutputKeys

# Define your selector (assuming you have a custom selector implementation)
class MySelector(BaseSelector):
    def select(self, choices, query):
        # Custom logic to select the appropriate component
        pass

# Define wrapper components for each pipeline
class TextToSQLPipelineComponent(QueryPipeline):
    @property
    def input_keys(self):
        return InputKeys.from_keys(required_keys={"query_str"})

    @property
    def output_keys(self):
        return OutputKeys.from_keys(required_keys={"sql_query"})

class RAGPipelineComponent(QueryPipeline):
    @property
    def input_keys(self):
        return InputKeys.from_keys(required_keys={"query_str"})

    @property
    def output_keys(self):
        return OutputKeys.from_keys(required_keys={"response"})

# Create instances of the pipelines
text_to_sql_pipeline = TextToSQLPipelineComponent(
    modules={
        "input": InputComponent(),
        "text_to_sql": text_to_sql_component,
    },
    verbose=False,
)

rag_pipeline = RAGPipelineComponent(
    modules={
        "input": InputComponent(),
        "retriever": retriever,
        "rag": rag_component,
    },
    verbose=False,
)

# Link the components in their respective pipelines
text_to_sql_pipeline.add_link("input", "text_to_sql", src_key="query_str", dest_key="query_str")
rag_pipeline.add_link("input", "retriever", src_key="query_str", dest_key="query_str")
rag_pipeline.add_link("retriever", "rag", dest_key="retrieved_docs")
rag_pipeline.add_link("input", "rag", src_key="query_str", dest_key="query_str")

# Initialize the selector
selector = MySelector()

# Initialize the RouterComponent
router_component = RouterComponent(
    selector=selector,
    choices=["text_to_sql", "rag"],
    components=[text_to_sql_pipeline, rag_pipeline],
    verbose=True
)

# Now you can use the router_component to route queries
result = router_component.run_component(query="Your query here")

In this setup:

TextToSQLPipelineComponent and RAGPipelineComponent use the InputKeys class to define the required input key (query_str).
The RouterComponent is initialized with these wrapped components, avoiding the "Expected one required input key" error ^[1].

To continue talking to Dosu, mention @dosu.

Answer selected by debarun-ghosh17

debarun-ghosh17 Sep 6, 2024
Author

visualize the final router pipeline as a directed graph

dosubot[bot] bot Sep 6, 2024

To visualize the final router pipeline as a directed graph, you can use the pyvis library to create an interactive graph visualization. Here’s an example of how to do this:

Install the pyvis library:
```
pip install pyvis
```

Create and visualize the directed graph:

from pyvis.network import Network
from llama_index.core.query_pipeline import QueryPipeline

# Assuming you have your router pipeline defined as `router_pipeline`
# For example:
# router_pipeline = QueryPipeline(chain=[router_component], verbose=True)

# Create a Network object
net = Network(notebook=True, cdn_resources="in_line", directed=True)

# Convert the pipeline's DAG to a NetworkX graph and add it to the Network object
net.from_nx(router_pipeline.dag)

# Visualize the graph
net.show("router_pipeline.html")

This code snippet will generate an HTML file (router_pipeline.html) that contains an interactive visualization of your router pipeline as a directed graph. You can open this file in a web browser to explore the graph.

This approach leverages the pyvis library to create a visual representation of the pipeline's directed acyclic graph (DAG), making it easier to understand the structure and flow of your query pipeline ^[1]^[2].

To continue talking to Dosu, mention @dosu.

debarun-ghosh17 Sep 9, 2024
Author

If one of the router components fails to get an answer or returns an error, can we direct the query to the alternative pipeline

dosubot[bot] bot Sep 9, 2024

Yes, if one of the router components fails to get an answer or returns an error, you can direct the query to an alternative pipeline. The RouterComponent class is designed to route queries to different query components based on a selector. If a component fails, you can handle the error and redirect the query to another component.

Here is an example of how you might modify the _run_component method to handle errors and redirect the query to an alternative component:

def _run_component(self, **kwargs: Any) -> Any:
    """Run component."""
    # for the output selection, run the corresponding component, aggregate into list
    sel_output = self.selector.select(self.choices, kwargs["query"])
    # assume one selection
    if len(sel_output.selections) != 1:
        raise ValueError("Expected one selection")
    component = self.components[sel_output.ind]
    log_str = f"Selecting component {sel_output.ind}: " f"{sel_output.reason}."
    if self.verbose:
        print_text(log_str + "\n", color="pink")
    try:
        # run component
        return component.run_component(
            **{self._query_keys[sel_output.ind]: kwargs["query"]}
        )
    except Exception as e:
        # Handle error and redirect to alternative component
        alternative_component = self.get_alternative_component(sel_output.ind)
        if alternative_component:
            return alternative_component.run_component(
                **{self._query_keys[sel_output.ind]: kwargs["query"]}
            )
        else:
            raise e

def get_alternative_component(self, failed_index: int) -> Optional[QueryComponent]:
    """Get an alternative component in case of failure."""
    # Logic to select an alternative component
    # For simplicity, return the next component in the list
    for i, component in enumerate(self.components):
        if i != failed_index:
            return component
    return None

This code attempts to run the selected component and, if it fails, redirects the query to an alternative component. The get_alternative_component method provides a simple way to select an alternative component. You can customize this logic based on your specific requirements ^[1]^[2].

To continue talking to Dosu, mention @dosu.

debarun-ghosh17 · 2024-09-12T05:01:29Z

debarun-ghosh17
Sep 12, 2024
Author

2 replies

dosubot[bot] bot Sep 12, 2024

To add a retry loop to the sql_retriever in your text2sql pipeline, you can modify the retrieve_with_metadata method in the SQLRetriever class. This method will attempt to execute the SQL query up to a specified number of retries. If the query fails, it will send the query, SQL, and error message back to the LLM to generate a new SQL query. After the maximum number of retries, it will exit gracefully with an exit message.

Here is an example implementation:

def retrieve_with_metadata(
        self, str_or_query_bundle: QueryType, max_retries: int = 3
    ) -> Tuple[List[NodeWithScore], Dict]:
    """Retrieve with metadata and retry logic."""
    if isinstance(str_or_query_bundle, str):
        query_bundle = QueryBundle(str_or_query_bundle)
    else:
        query_bundle = str_or_query_bundle
    table_desc_str = self._get_table_context(query_bundle)
    logger.info(f"> Table desc str: {table_desc_str}")
    if self._verbose:
        print(f"> Table desc str: {table_desc_str}")

    retries = 0
    while retries < max_retries:
        try:
            response_str = self._llm.predict(
                self._text_to_sql_prompt,
                query_str=query_bundle.query_str,
                schema=table_desc_str,
                dialect=self._sql_database.dialect,
            )

            sql_query_str = self._sql_parser.parse_response_to_sql(
                response_str, query_bundle
            )
            logger.debug(f"> Predicted SQL query: {sql_query_str}")
            if self._verbose:
                print(f"> Predicted SQL query: {sql_query_str}")

            if self._sql_only:
                sql_only_node = TextNode(text=f"{sql_query_str}")
                retrieved_nodes = [NodeWithScore(node=sql_only_node)]
                metadata = {"result": sql_query_str}
            else:
                retrieved_nodes, metadata = self._sql_retriever.retrieve_with_metadata(
                    sql_query_str
                )
            return retrieved_nodes, {"sql_query": sql_query_str, **metadata}
        except Exception as e:
            logger.error(f"SQL query failed: {e}")
            retries += 1
            if retries >= max_retries:
                exit_message = f"Failed to generate a valid SQL query after {max_retries} attempts."
                logger.error(exit_message)
                err_node = TextNode(text=exit_message)
                return [NodeWithScore(node=err_node)], {"error": str(e)}

    return [], {"error": "Unexpected error"}

Explanation:

Initialization:
- The method initializes the QueryBundle and retrieves the table context.
Retry Loop:
- A while loop is used to attempt the SQL query up to max_retries times.
- If an exception occurs during the SQL query execution, it logs the error and retries.
Error Handling:
- If the maximum number of retries is reached, it exits gracefully with an exit message and returns an error node.
Successful Execution:
- If the SQL query is successful, it returns the retrieved nodes and metadata.

This implementation ensures that the sql_retriever will retry generating and executing the SQL query up to the specified number of retries, and it will handle errors gracefully if the maximum number of retries is reached ^[1]^[2].

To continue talking to Dosu, mention @dosu.

debarun-ghosh17 Sep 12, 2024
Author

@dosu How would you add retry loop so that if sql_retriiever comes with an error. We go back to llm, with the query+sql+error to regenrate anoter one. The pipeline setup is as follows:
qp = QP(
modules={
"input": InputComponent(),
"table_retriever": obj_retriever,
"table_output_parser": table_parser_component,
"text2sql_prompt": text2sql_prompt,
"text2sql_llm": llm,
"sql_output_parser": sql_parser_component,
"sql_retriever": sql_retriever,
"response_synthesis_prompt": response_synthesis_prompt,
"response_synthesis_llm": llm,
"sql_pandas_component": sql_pandas_component,

"vector_tool": vector_tool_component

},
verbose=True,
input_keys = 'query'

)
qp.add_link("input", "table_retriever")
qp.add_link("input", "table_output_parser", dest_key="query_str")
qp.add_link(
"table_retriever", "table_output_parser", dest_key="table_schema_objs"
)
qp.add_link("input", "text2sql_prompt", dest_key="query_str")
qp.add_link("table_output_parser", "text2sql_prompt", dest_key="schema")
qp.add_chain(
["text2sql_prompt", "text2sql_llm", "sql_output_parser", "sql_retriever"]
)
qp.add_link(
"sql_output_parser", "response_synthesis_prompt", dest_key="sql_query"
)
qp.add_link(
"sql_retriever", "response_synthesis_prompt", dest_key="context_str"
)
#additions
qp.add_link(
"sql_output_parser", "sql_pandas_component", dest_key="sql_query"
)
qp.add_link(
"sql_pandas_component", "response_synthesis_prompt", dest_key="markdown_str"
)
qp.add_link("input", "response_synthesis_prompt", dest_key="query_str")
qp.add_link("response_synthesis_prompt", "response_synthesis_llm")

The sql retriever component just takes one input, sql query generated from sql_output_parser and one output to the response sysnthesis promt.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

query pipelines with router #15852

{{title}}

Replies: 2 comments 12 replies

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

Select a reply

query pipelines with router #15852

debarun-ghosh17 Sep 5, 2024

Replies: 2 comments · 12 replies

dosubot[bot] bot Sep 5, 2024

debarun-ghosh17 Sep 5, 2024 Author

dosubot[bot] bot Sep 5, 2024

debarun-ghosh17 Sep 6, 2024 Author

dosubot[bot] bot Sep 6, 2024

debarun-ghosh17 Sep 6, 2024 Author

dosubot[bot] bot Sep 6, 2024

debarun-ghosh17 Sep 9, 2024 Author

dosubot[bot] bot Sep 9, 2024

debarun-ghosh17 Sep 12, 2024 Author

dosubot[bot] bot Sep 12, 2024

Explanation:

debarun-ghosh17 Sep 12, 2024 Author

"vector_tool": vector_tool_component

debarun-ghosh17
Sep 5, 2024

Replies: 2 comments 12 replies

dosubot[bot]
bot Sep 5, 2024

debarun-ghosh17 Sep 5, 2024
Author

debarun-ghosh17 Sep 6, 2024
Author

debarun-ghosh17 Sep 6, 2024
Author

debarun-ghosh17 Sep 9, 2024
Author

debarun-ghosh17
Sep 12, 2024
Author

debarun-ghosh17 Sep 12, 2024
Author