multiple inputs in Query Pipeline (advanced text2sql with run,multi)

[sumitbindra]

Hello All,

I am trying to figure out how to use multiple inputs in advanced text2sql with a query pipeline. I think using run.multi here is the answer but not sure how it will be implemented if the inputs are used in several modules in query pipeline.

https://docs.llamaindex.ai/en/stable/examples/pipeline/query_pipeline_sql/?h=sql#2-advanced-capability-2-text-to-sql-with-query-time-row-retrieval-along-with-table-retrieval

Run.multi example: https://docs.llamaindex.ai/en/stable/module_guides/querying/pipeline/usage_pattern/#multi-inputmulti-output

The eg above is what I want to implement for multiple inputs in the advanced text to sql query pipeline:

My functions takes in two variables : query_str (string) and tablename (string). The example has only one input component. I tried using run.multi with a dictionary but am confused on how it will be setup if the multiple inputs go in a few different modules.

Would appreciate any help on this.

[dosubot[bot]]

Hey @sumitbindra! I'm here to help you with any questions or issues you have. Let's work together to solve this problem.

To set up run.multi for multiple inputs in an advanced text-to-SQL query pipeline, especially when the inputs are used in different modules, you can use the run_multi method provided in the QueryPipeline class. Here is a detailed example to help you implement this:

1. Define Your Modules: Create the necessary modules for your pipeline. For instance, you might have modules for handling the query string and the table name.

2. Set Up the Query Pipeline: Add the modules to the pipeline and define the links between them.

3. Run the Pipeline with Multiple Inputs: Use the run_multi method to execute the pipeline with the required inputs.

Here is a practical example:

from llama_index.core.query_pipeline.query import QueryPipeline
from llama_index.core.query_pipeline.components import QueryComponent

# Define your custom components
class QueryComponent1(QueryComponent):
    def run_component(self, query_str, tablename):
        # Process the query_str and tablename
        return {"processed_query": f"Processed {query_str} for {tablename}"}

class QueryComponent2(QueryComponent):
    def run_component(self, processed_query):
        # Further process the query
        return {"final_output": f"Final output: {processed_query}"}

# Initialize the components
qc1 = QueryComponent1()
qc2 = QueryComponent2()

# Define the query pipeline
pipeline = QueryPipeline()
pipeline.add_modules({"qc1": qc1, "qc2": qc2})
pipeline.add_link("qc1", "qc2", dest_key="processed_query")

# Run the pipeline with multiple inputs
inputs = {
    "qc1": {"query_str": "SELECT * FROM users", "tablename": "users"}
}
output = pipeline.run_multi(inputs)

print(output)

In this example:

• QueryComponent1 processes the query_str and tablename.
• QueryComponent2 takes the processed query from QueryComponent1 and produces the final output.
• The run_multi method is used to run the pipeline with the specified inputs.

This setup ensures that the inputs are correctly distributed to the respective modules and processed accordingly. You can customize the components and their processing logic as needed for your specific use case.

For more detailed information and examples, you can refer to the LlamaIndex documentation on multi-input/multi-output pipelines ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

@dosu thanks.

What if there are more than one module in my pipeline that take both or part of the inputs?

For example

# Define your custom components
class QueryComponent1(QueryComponent):
    def run_component(self, query_str, tablename):
        # Process the query_str and tablename
        return {"processed_query": f"Processed {query_str} for {tablename}"}

class QueryComponent2(QueryComponent):
    def run_component(self, processed_query):
        # Further process the query
        return {"final_output": f"Final output: {processed_query}"}

# Define your custom components
class QueryComponent3(QueryComponent):
    def run_component(self, query_str, tablename):
        # Process the query_str and tablename
        return {"processed_query": f"Processed {query_str} for {tablename} yet again"}

how will I specify run.multi for this?

[dosubot[bot]]

To specify run.multi for a setup where multiple modules need to take both or part of the inputs, you can use the run_multi method provided in the QueryPipeline class. This method is designed to handle multiple root modules and their respective inputs.

Here is an example of how you can specify run.multi for your scenario:

from llama_index.core.query_pipeline.query import QueryPipeline
from llama_index.core.query_pipeline.components import QueryComponent

# Define your custom components
class QueryComponent1(QueryComponent):
    def run_component(self, query_str, tablename):
        # Process the query_str and tablename
        return {"processed_query1": f"Processed {query_str} for {tablename}"}

class QueryComponent2(QueryComponent):
    def run_component(self, processed_query1):
        # Further process the query
        return {"final_output": f"Final output: {processed_query1}"}

class QueryComponent3(QueryComponent):
    def run_component(self, query_str, tablename):
        # Process the query_str and tablename again
        return {"processed_query3": f"Processed {query_str} for {tablename} yet again"}

# Initialize the components
qc1 = QueryComponent1()
qc2 = QueryComponent2()
qc3 = QueryComponent3()

# Define the query pipeline
pipeline = QueryPipeline()
pipeline.add_modules({"qc1": qc1, "qc2": qc2, "qc3": qc3})
pipeline.add_link("qc1", "qc2", dest_key="processed_query1")
pipeline.add_link("qc3", "qc2", dest_key="processed_query3")

# Run the pipeline with multiple inputs
inputs = {
    "qc1": {"query_str": "SELECT * FROM users", "tablename": "users"},
    "qc3": {"query_str": "SELECT * FROM users", "tablename": "users"}
}
output = pipeline.run_multi(inputs)

print(output)

In this example:

• QueryComponent1 processes the query_str and tablename and produces processed_query1.
• QueryComponent2 takes processed_query1 from QueryComponent1 and processed_query3 from QueryComponent3 to produce the final output.
• QueryComponent3 also processes the query_str and tablename and produces processed_query3.
• The run_multi method is used to run the pipeline with the specified inputs for both qc1 and qc3.

This setup ensures that the inputs are correctly distributed to the respective modules and processed accordingly. You can customize the components and their processing logic as needed for your specific use case ^[1][2].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

when i run this code:

from llama_index.core.query_pipeline.query import QueryPipeline
from llama_index.core.query_pipeline import CustomQueryComponent

class QueryComponent1(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        # Process the query_str and tablename
        return {"processed_query1": f"Processed {query_str} for {tablename}"}

class QueryComponent2(CustomQueryComponent):
    def _run_component(self, **kwargs):
        processed_query1 = kwargs.get('processed_query1')
        processed_query3 = kwargs.get('processed_query3')
        # Further process the query
        return {"final_output": f"Final output: {processed_query1} and {processed_query3}"}

class QueryComponent3(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        # Process the query_str and tablename again
        return {"processed_query3": f"Processed {query_str} for {tablename} yet again"}

# Initialize the components
qc1 = QueryComponent1()
qc2 = QueryComponent2()
qc3 = QueryComponent3()

# Define the query pipeline
pipeline = QueryPipeline()
pipeline.add_modules({"qc1": qc1, "qc2": qc2, "qc3": qc3})
pipeline.add_link("qc1", "qc2", dest_key="processed_query1")
pipeline.add_link("qc3", "qc2", dest_key="processed_query3")

# Run the pipeline with multiple inputs
inputs = {
    "qc1": {"query_str": "SELECT * FROM users", "tablename": "users"},
    "qc3": {"query_str": "SELECT * FROM users", "tablename": "users"}
}
output = pipeline.run_multi(inputs)

print(output)

I get the following error:

Traceback (most recent call last):
File "/home/runner/randmo-python/main.py", line 41, in
output = pipeline.run_multi(inputs)
File "/home/runner/randmo-python/.pythonlibs/lib/python3.10/site-packages/llama_index/core/query_pipeline/query.py", line 425, in run_multi
outputs, _ = self._run_multi(module_input_dict)
File "/home/runner/randmo-python/.pythonlibs/lib/python3.10/site-packages/llama_index/core/query_pipeline/query.py", line 876, in _run_multi
output_dict = module.run_component(**module_input)
File "/home/runner/randmo-python/.pythonlibs/lib/python3.10/site-packages/llama_index/core/base/query_pipeline/query.py", line 198, in run_component
kwargs = self.validate_component_inputs(kwargs)
File "/home/runner/randmo-python/.pythonlibs/lib/python3.10/site-packages/llama_index/core/base/query_pipeline/query.py", line 186, in validate_component_inputs
self.input_keys.validate(set(input.keys()))
File "/home/runner/randmo-python/.pythonlibs/lib/python3.10/site-packages/llama_index/core/base/query_pipeline/query.py", line 285, in input_keys
required_keys=self._input_keys, optional_keys=self._optional_input_keys
File "/home/runner/randmo-python/.pythonlibs/lib/python3.10/site-packages/llama_index/core/base/query_pipeline/query.py", line 266, in _input_keys
raise NotImplementedError("Not implemented yet. Please override this method.")
NotImplementedError: Not implemented yet. Please override this method.

How do I fix it?

[dosubot[bot]]

To fix the NotImplementedError in your code, you need to implement the _input_keys, _optional_input_keys, and _output_keys properties in your CustomQueryComponent subclasses. These properties are required to specify the input and output keys for each component.

Here is an updated version of your code with the necessary implementations:

from llama_index.core.query_pipeline.query import QueryPipeline
from llama_index.core.query_pipeline.components import CustomQueryComponent

class QueryComponent1(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        # Process the query_str and tablename
        return {"processed_query1": f"Processed {query_str} for {tablename}"}

    @property
    def _input_keys(self):
        return {"query_str", "tablename"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"processed_query1"}

class QueryComponent2(CustomQueryComponent):
    def _run_component(self, **kwargs):
        processed_query1 = kwargs.get('processed_query1')
        processed_query3 = kwargs.get('processed_query3')
        # Further process the query
        return {"final_output": f"Final output: {processed_query1} and {processed_query3}"}

    @property
    def _input_keys(self):
        return {"processed_query1", "processed_query3"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"final_output"}

class QueryComponent3(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        # Process the query_str and tablename again
        return {"processed_query3": f"Processed {query_str} for {tablename} yet again"}

    @property
    def _input_keys(self):
        return {"query_str", "tablename"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"processed_query3"}

# Initialize the components
qc1 = QueryComponent1()
qc2 = QueryComponent2()
qc3 = QueryComponent3()

# Define the query pipeline
pipeline = QueryPipeline()
pipeline.add_modules({"qc1": qc1, "qc2": qc2, "qc3": qc3})
pipeline.add_link("qc1", "qc2", dest_key="processed_query1")
pipeline.add_link("qc3", "qc2", dest_key="processed_query3")

# Run the pipeline with multiple inputs
inputs = {
    "qc1": {"query_str": "SELECT * FROM users", "tablename": "users"},
    "qc3": {"query_str": "SELECT * FROM users", "tablename": "users"}
}
output = pipeline.run_multi(inputs)

print(output)

This should resolve the NotImplementedError by providing the required input and output keys for each component ^[1].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

Hey @dosu,

How would I set it up if one of my query components has a input sent as function call parameter (like tablename) and another is an output from a different query component?

lets call it querycomponent4 and imagine it needs two parameters ... tablename and output of qc3 i.e. processed_query3

[dosubot[bot]]

To set up QueryComponent4 with tablename as a function call parameter and processed_query3 as an output from QueryComponent3, you need to define QueryComponent4 and create links to connect it with the previous components. Here’s how you can modify your code:

1. Define QueryComponent4 with the required input keys.
2. Create links to connect QueryComponent4 with the previous components.

Here’s the updated code:

from llama_index.core.query_pipeline.query import QueryPipeline
from llama_index.core.query_pipeline.components import CustomQueryComponent

class QueryComponent1(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        return {"processed_query1": f"Processed {query_str} for {tablename}"}

    @property
    def _input_keys(self):
        return {"query_str", "tablename"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"processed_query1"}

class QueryComponent2(CustomQueryComponent):
    def _run_component(self, **kwargs):
        processed_query1 = kwargs.get('processed_query1')
        processed_query3 = kwargs.get('processed_query3')
        return {"final_output": f"Final output: {processed_query1} and {processed_query3}"}

    @property
    def _input_keys(self):
        return {"processed_query1", "processed_query3"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"final_output"}

class QueryComponent3(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        return {"processed_query3": f"Processed {query_str} for {tablename} yet again"}

    @property
    def _input_keys(self):
        return {"query_str", "tablename"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"processed_query3"}

class QueryComponent4(CustomQueryComponent):
    def _run_component(self, **kwargs):
        tablename = kwargs.get('tablename')
        processed_query3 = kwargs.get('processed_query3')
        # Implement the logic for QueryComponent4 here
        return {"output": f"Processed {tablename} and {processed_query3}"}

    @property
    def _input_keys(self):
        return {"tablename", "processed_query3"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"output"}

# Initialize the components
qc1 = QueryComponent1()
qc2 = QueryComponent2()
qc3 = QueryComponent3()
qc4 = QueryComponent4()

# Define the query pipeline
pipeline = QueryPipeline()
pipeline.add_modules({"qc1": qc1, "qc2": qc2, "qc3": qc3, "qc4": qc4})
pipeline.add_link("qc1", "qc2", dest_key="processed_query1")
pipeline.add_link("qc3", "qc2", dest_key="processed_query3")
pipeline.add_link("qc3", "qc4", dest_key="processed_query3")

# Assuming you have another link or a way to provide `tablename` to QueryComponent4
pipeline.add_link("qc1", "qc4", dest_key="tablename")

# Run the pipeline with multiple inputs
inputs = {
    "qc1": {"query_str": "SELECT * FROM users", "tablename": "users"},
    "qc3": {"query_str": "SELECT * FROM users", "tablename": "users"}
}
output = pipeline.run_multi(inputs)

print(output)

In this setup:

• QueryComponent4 is defined to take tablename and processed_query3 as inputs.
• Links are created to connect QueryComponent4 with QueryComponent3 for processed_query3 and with QueryComponent1 for tablename.

This ensures that QueryComponent4 receives the necessary inputs from the appropriate components ^[1][2].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

@dosu what I need is something like:

response_synthesis_prompt_str = (
"Given an input question, synthesize a response from the query results.\n"
"table: {tablename}\n"
"q3 Response: {processed_query3}\n"
"Response: ")

Run the pipeline with multiple inputs

inputs = {
"qc1": {"query_str": "SELECT * FROM users", "tablename": "users"},
"qc3": {"query_str": "SELECT * FROM users", "tablename": "users"},
"qc4": {"tablename": "users"}
}

but when I set it up like this I get an error:

ValueError: Expected root keys do not match up with input keys.
Expected root keys: ['qc1', 'qc3']
Input keys: dict_keys(['qc1', 'qc3', 'qc4'])

[dosubot[bot]]

To resolve the ValueError regarding the mismatch between expected root keys and input keys, you need to ensure that qc4 is not treated as a root module. Instead, it should receive its inputs from other components in the pipeline.

Here’s how you can set it up:

1. Define QueryComponent4 to take tablename and processed_query3 as inputs.
2. Create links to connect QueryComponent4 with the previous components.
3. Run the pipeline without specifying qc4 as a root input.

Here’s the updated code:

from llama_index.core.query_pipeline.query import QueryPipeline
from llama_index.core.query_pipeline.components import CustomQueryComponent

class QueryComponent1(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        return {"processed_query1": f"Processed {query_str} for {tablename}"}

    @property
    def _input_keys(self):
        return {"query_str", "tablename"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"processed_query1"}

class QueryComponent2(CustomQueryComponent):
    def _run_component(self, **kwargs):
        processed_query1 = kwargs.get('processed_query1')
        processed_query3 = kwargs.get('processed_query3')
        return {"final_output": f"Final output: {processed_query1} and {processed_query3}"}

    @property
    def _input_keys(self):
        return {"processed_query1", "processed_query3"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"final_output"}

class QueryComponent3(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        return {"processed_query3": f"Processed {query_str} for {tablename} yet again"}

    @property
    def _input_keys(self):
        return {"query_str", "tablename"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"processed_query3"}

class QueryComponent4(CustomQueryComponent):
    def _run_component(self, **kwargs):
        tablename = kwargs.get('tablename')
        processed_query3 = kwargs.get('processed_query3')
        response_synthesis_prompt_str = (
            f"Given an input question, synthesize a response from the query results.\n"
            f"table: {tablename}\n"
            f"q3 Response: {processed_query3}\n"
            f"Response: "
        )
        return {"output": response_synthesis_prompt_str}

    @property
    def _input_keys(self):
        return {"tablename", "processed_query3"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"output"}

# Initialize the components
qc1 = QueryComponent1()
qc2 = QueryComponent2()
qc3 = QueryComponent3()
qc4 = QueryComponent4()

# Define the query pipeline
pipeline = QueryPipeline()
pipeline.add_modules({"qc1": qc1, "qc2": qc2, "qc3": qc3, "qc4": qc4})
pipeline.add_link("qc1", "qc2", dest_key="processed_query1")
pipeline.add_link("qc3", "qc2", dest_key="processed_query3")
pipeline.add_link("qc3", "qc4", dest_key="processed_query3")
pipeline.add_link("qc1", "qc4", dest_key="tablename")

# Run the pipeline with multiple inputs
inputs = {
    "qc1": {"query_str": "SELECT * FROM users", "tablename": "users"},
    "qc3": {"query_str": "SELECT * FROM users", "tablename": "users"}
}
output = pipeline.run_multi(inputs)

print(output)

In this setup:

• QueryComponent4 is defined to take tablename and processed_query3 as inputs.
• Links are created to connect QueryComponent4 with QueryComponent3 for processed_query3 and with QueryComponent1 for tablename.
• qc4 is not included in the root inputs, thus avoiding the ValueError.

This ensures that QueryComponent4 receives the necessary inputs from the appropriate components and synthesizes the response correctly ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

@dosu

do I have to add "tablename" to output keys of Querycomponent1?

[dosubot[bot]]

Yes, you need to add "tablename" to the output keys of QueryComponent1 to ensure it is available for QueryComponent4. This is because QueryComponent4 requires "tablename" as an input, and it needs to be passed through the pipeline from QueryComponent1.

Here's how you can modify QueryComponent1 to include "tablename" in its output keys:

class QueryComponent1(CustomQueryComponent):
    def _run_component(self, **kwargs):
        query_str = kwargs.get('query_str')
        tablename = kwargs.get('tablename')
        return {
            "processed_query1": f"Processed {query_str} for {tablename}",
            "tablename": tablename  # Include tablename in the output
        }

    @property
    def _input_keys(self):
        return {"query_str", "tablename"}

    @property
    def _optional_input_keys(self):
        return set()

    @property
    def _output_keys(self):
        return {"processed_query1", "tablename"}  # Add tablename to output keys

With this change, QueryComponent1 will output "tablename", making it available for QueryComponent4 through the pipeline.

To continue talking to Dosu, mention @dosu.