Creating Custom SQL Query Engine from Query Pipeline

[deybelman]

I want to create a SQL query engine using my existing query pipeline (see code below). My query pipeline is similar to NLSQLQueryEngine but more advanced. How do I connect it to an engine? Once I do that, I will wrap it in a tool and have an agent select that tool. Please help explain exactly how to get from a query pipeline to a query engine like I'm a newbie.

from llama_index.core.query_pipeline import (
    QueryPipeline as QP,
    Link,
    InputComponent,
    CustomQueryComponent,
)

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,
    },
    verbose=True,
)
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"
)
qp.add_link("input", "response_synthesis_prompt", dest_key="query_str")
qp.add_link("response_synthesis_prompt", "response_synthesis_llm")

[dosubot[bot]]

Hello, @deybelman! I'm here to help you with any bugs, questions, or contributions you might have. If you need assistance, feel free to ask!

To connect your existing query pipeline to a SQL query engine and wrap it in a tool for an agent to select, follow these steps:

1. Setup the SQL Database:

   from llama_index.core import SQLDatabase
   from sqlalchemy import create_engine
   
   engine = create_engine("sqlite:///chinook.db")
   sql_database = SQLDatabase(engine)

2. Setup the Text-to-SQL Query Engine and Tool:

   from llama_index.core.query_engine import NLSQLTableQueryEngine
   from llama_index.core.tools import QueryEngineTool, ToolMetadata
   
   sql_query_engine = NLSQLTableQueryEngine(
       sql_database=sql_database,
       tables=["albums", "tracks", "artists"],
       verbose=True,
   )
   sql_tool = QueryEngineTool(
       query_engine=sql_query_engine,
       metadata=ToolMetadata(
           name="sql_tool",
           description="Useful for translating a natural language query into a SQL query"
       ),
   )

3. Setup the ReAct Agent Pipeline:

   from llama_index.core.query_pipeline import QueryPipeline as QP
   from llama_index.llms.openai import OpenAI
   
   qp = QP(verbose=True)

4. Define Agent Input Component:

   from llama_index.core.agent.react.types import ObservationReasoningStep
   from llama_index.core.agent import Task
   from llama_index.core.query_pipeline import AgentInputComponent
   
   def agent_input_fn(task: Task, state: Dict[str, Any]) -> Dict[str, Any]:
       if "current_reasoning" not in state:
           state["current_reasoning"] = []
       reasoning_step = ObservationReasoningStep(observation=task.input)
       state["current_reasoning"].append(reasoning_step)
       return {"input": task.input}
   
   agent_input_component = AgentInputComponent(fn=agent_input_fn)

5. Define Agent Prompt:

   from llama_index.core.agent import ReActChatFormatter
   from llama_index.core.query_pipeline import AgentFnComponent
   from llama_index.core.llms import ChatMessage
   from llama_index.core.tools import BaseTool
   
   def react_prompt_fn(task: Task, state: Dict[str, Any], input: str, tools: List[BaseTool]) -> List[ChatMessage]:
       chat_formatter = ReActChatFormatter()
       return chat_formatter.format(
           tools,
           chat_history=task.memory.get() + state["memory"].get_all(),
           current_reasoning=state["current_reasoning"],
       )
   
   react_prompt_component = AgentFnComponent(
       fn=react_prompt_fn, partial_dict={"tools": [sql_tool]}
   )

6. Define Agent Output Parser and Tool Pipeline:

   from llama_index.core.agent.react.output_parser import ReActOutputParser
   from llama_index.core.llms import ChatResponse
   from llama_index.core.query_pipeline import ToolRunnerComponent
   
   def parse_react_output_fn(task: Task, state: Dict[str, Any], chat_response: ChatResponse):
       output_parser = ReActOutputParser()
       reasoning_step = output_parser.parse(chat_response.message.content)
       return {"done": reasoning_step.is_done, "reasoning_step": reasoning_step}
   
   parse_react_output = AgentFnComponent(fn=parse_react_output_fn)
   
   def run_tool_fn(task: Task, state: Dict[str, Any], reasoning_step: ActionReasoningStep):
       tool_runner_component = ToolRunnerComponent([sql_tool], callback_manager=task.callback_manager)
       tool_output = tool_runner_component.run_component(
           tool_name=reasoning_step.action,
           tool_input=reasoning_step.action_input,
       )
       observation_step = ObservationReasoningStep(observation=str(tool_output))
       state["current_reasoning"].append(observation_step)
       return {"response_str": observation_step.get_content(), "is_done": False}
   
   run_tool = AgentFnComponent(fn=run_tool_fn)

7. Stitch Together Agent Query Pipeline:

   qp.add_modules(
       {
           "agent_input": agent_input_component,
           "react_prompt": react_prompt_component,
           "llm": OpenAI(model="gpt-4-1106-preview"),
           "react_output_parser": parse_react_output,
           "run_tool": run_tool,
       }
   )
   
   qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])
   qp.add_link(
       "react_output_parser",
       "run_tool",
       condition_fn=lambda x: not x["done"],
       input_fn=lambda x: x["reasoning_step"],
   )

8. Setup Agent Worker:

   from llama_index.core.agent import QueryPipelineAgentWorker
   from llama_index.core.callbacks import CallbackManager
   
   agent_worker = QueryPipelineAgentWorker(qp)
   agent = agent_worker.as_agent(
       callback_manager=CallbackManager([]), verbose=True
   )

9. Run the Agent:

   task = agent.create_task("What are some tracks from the artist AC/DC? Limit it to 3")
   step_output = agent.run_step(task.task_id)
   response = agent.finalize_response(task.task_id)
   print(str(response))

This setup connects your query pipeline to a SQL query engine and wraps it in a tool for an agent to select ^[1][2].

To continue talking to Dosu, mention @dosu.

[dosubot[bot]]

To use your existing text-2-SQL query pipeline instead of NLSQLTableQueryEngine, you need to create a custom query engine that wraps your pipeline. Here’s a step-by-step guide to achieve this:

1. Define Your Custom Query Engine

First, you need to define a custom query engine class that will use your existing query pipeline. This class should inherit from CustomQueryEngine and implement the custom_query method.

from llama_index.core.query_engine.custom import CustomQueryEngine
from llama_index.core.query_pipeline import QueryPipeline

class YourCustomText2SQLQueryEngine(CustomQueryEngine):
    def __init__(self, query_pipeline: QueryPipeline, **kwargs):
        super().__init__(**kwargs)
        self.query_pipeline = query_pipeline

    def custom_query(self, query_str: str):
        # Run the query through your existing pipeline
        result = self.query_pipeline.run(query_str)
        return result

2. Instantiate Your Custom Query Engine

Next, instantiate your custom query engine using your existing query pipeline.

# Assuming `qp` is your existing query pipeline
sql_query_engine = YourCustomText2SQLQueryEngine(query_pipeline=qp)

3. Wrap Your Custom Query Engine in a Tool

Now, wrap your custom query engine in a QueryEngineTool so it can be used by the agent.

from llama_index.core.tools import QueryEngineTool, ToolMetadata

sql_tool = QueryEngineTool(
    query_engine=sql_query_engine,
    metadata=ToolMetadata(
        name="sql_tool",
        description="Useful for translating a natural language query into a SQL query"
    ),
)

4. Integrate the Tool into the ReAct Agent Pipeline

Finally, integrate the sql_tool into your ReAct Agent Pipeline. This involves setting up the agent input, prompt, output parser, and tool runner components.

from llama_index.core.agent.react.types import ObservationReasoningStep
from llama_index.core.agent import Task
from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent, ToolRunnerComponent
from llama_index.core.llms import OpenAI, ChatMessage, ChatResponse
from llama_index.core.agent.react.output_parser import ReActOutputParser

# Define Agent Input Component
def agent_input_fn(task: Task, state: Dict[str, Any]) -> Dict[str, Any]:
    if "current_reasoning" not in state:
        state["current_reasoning"] = []
    reasoning_step = ObservationReasoningStep(observation=task.input)
    state["current_reasoning"].append(reasoning_step)
    return {"input": task.input}

agent_input_component = AgentInputComponent(fn=agent_input_fn)

# Define Agent Prompt
def react_prompt_fn(task: Task, state: Dict[str, Any], input: str, tools: List[BaseTool]) -> List[ChatMessage]:
    chat_formatter = ReActChatFormatter()
    return chat_formatter.format(
        tools,
        chat_history=task.memory.get() + state["memory"].get_all(),
        current_reasoning=state["current_reasoning"],
    )

react_prompt_component = AgentFnComponent(
    fn=react_prompt_fn, partial_dict={"tools": [sql_tool]}
)

# Define Agent Output Parser
def parse_react_output_fn(task: Task, state: Dict[str, Any], chat_response: ChatResponse):
    output_parser = ReActOutputParser()
    reasoning_step = output_parser.parse(chat_response.message.content)
    return {"done": reasoning_step.is_done, "reasoning_step": reasoning_step}

parse_react_output = AgentFnComponent(fn=parse_react_output_fn)

# Define Tool Runner
def run_tool_fn(task: Task, state: Dict[str, Any], reasoning_step: ActionReasoningStep):
    tool_runner_component = ToolRunnerComponent([sql_tool], callback_manager=task.callback_manager)
    tool_output = tool_runner_component.run_component(
        tool_name=reasoning_step.action,
        tool_input=reasoning_step.action_input,
    )
    observation_step = ObservationReasoningStep(observation=str(tool_output))
    state["current_reasoning"].append(observation_step)
    return {"response_str": observation_step.get_content(), "is_done": False}

run_tool = AgentFnComponent(fn=run_tool_fn)

# Stitch Together Agent Query Pipeline
qp.add_modules(
    {
        "agent_input": agent_input_component,
        "react_prompt": react_prompt_component,
        "llm": OpenAI(model="gpt-4-1106-preview"),
        "react_output_parser": parse_react_output,
        "run_tool": run_tool,
    }
)

qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])
qp.add_link(
    "react_output_parser",
    "run_tool",
    condition_fn=lambda x: not x["done"],
    input_fn=lambda x: x["reasoning_step"],
)

5. Setup and Run the Agent

Finally, set up the agent worker and run the agent.

from llama_index.core.agent import QueryPipelineAgentWorker
from llama_index.core.callbacks import CallbackManager

agent_worker = QueryPipelineAgentWorker(qp)
agent = agent_worker.as_agent(
    callback_manager=CallbackManager([]), verbose=True
)

task = agent.create_task("What are some tracks from the artist AC/DC? Limit it to 3")
step_output = agent.run_step(task.task_id)
response = agent.finalize_response(task.task_id)
print(str(response))

By following these steps, you can replace the NLSQLTableQueryEngine with your existing text-2-SQL query pipeline and integrate it into the ReAct Agent Pipeline ^[1][2].

To continue talking to Dosu, mention @dosu.

[deybelman]

@dosu I ran steps 1 and 2 exactly as you said and got the error below. What should I do?

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[42], line 1
----> 1 sql_query_engine = YourCustomText2SQLQueryEngine(query_pipeline=qp)

Cell In[41], line 7, in YourCustomText2SQLQueryEngine.__init__(self, query_pipeline, **kwargs)
      5 def __init__(self, query_pipeline: QueryPipeline, **kwargs):
      6     super().__init__(**kwargs)
----> 7     self.query_pipeline = query_pipeline

File /Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/v1/main.py:357, in BaseModel.__setattr__(self, name, value)
    354     return object_setattr(self, name, value)
    356 if self.__config__.extra is not Extra.allow and name not in self.__fields__:
--> 357     raise ValueError(f'"{self.__class__.__name__}" object has no field "{name}"')
    358 elif not self.__config__.allow_mutation or self.__config__.frozen:
    359     raise TypeError(f'"{self.__class__.__name__}" is immutable and does not support item assignment')

ValueError: "YourCustomText2SQLQueryEngine" object has no field "query_pipeline"

[dosubot[bot]]

To resolve the error "YourCustomText2SQLQueryEngine object has no field 'query_pipeline'" and correctly use your existing query pipeline in the custom query engine, you need to ensure that the query_pipeline field is properly defined and initialized in your custom query engine class.

Here’s how you can do it:

1. Define Your Custom Query Engine

First, define the query_pipeline field using Pydantic's Field and ensure it is properly initialized when creating an instance of your custom query engine.

from llama_index.core.query_engine.custom import CustomQueryEngine
from llama_index.core.query_pipeline import QueryPipeline
from llama_index.legacy.bridge.pydantic import Field
from llama_index.legacy.core.response.schema import RESPONSE_TYPE

class YourCustomText2SQLQueryEngine(CustomQueryEngine):
    query_pipeline: QueryPipeline = Field(
        description="Query Pipeline to use for Q&A.",
    )

    def __init__(self, query_pipeline: QueryPipeline, **kwargs):
        super().__init__(**kwargs)
        self.query_pipeline = query_pipeline

    def custom_query(self, query_str: str) -> RESPONSE_TYPE:
        return self.query_pipeline.run(query_str=query_str)

    async def acustom_query(self, query_str: str) -> RESPONSE_TYPE:
        return await self.query_pipeline.arun(query_str=query_str)

2. Instantiate Your Custom Query Engine

Next, instantiate your custom query engine using your existing query pipeline.

# Assuming `qp` is your existing query pipeline
sql_query_engine = YourCustomText2SQLQueryEngine(query_pipeline=qp)

3. Wrap Your Custom Query Engine in a Tool

Now, wrap your custom query engine in a QueryEngineTool so it can be used by the agent.

from llama_index.core.tools import QueryEngineTool, ToolMetadata

sql_tool = QueryEngineTool(
    query_engine=sql_query_engine,
    metadata=ToolMetadata(
        name="sql_tool",
        description="Useful for translating a natural language query into a SQL query"
    ),
)

4. Integrate the Tool into the ReAct Agent Pipeline

Finally, integrate the sql_tool into your ReAct Agent Pipeline. This involves setting up the agent input, prompt, output parser, and tool runner components.

from llama_index.core.agent.react.types import ObservationReasoningStep
from llama_index.core.agent import Task
from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent, ToolRunnerComponent
from llama_index.core.llms import OpenAI, ChatMessage, ChatResponse
from llama_index.core.agent.react.output_parser import ReActOutputParser

# Define Agent Input Component
def agent_input_fn(task: Task, state: Dict[str, Any]) -> Dict[str, Any]:
    if "current_reasoning" not in state:
        state["current_reasoning"] = []
    reasoning_step = ObservationReasoningStep(observation=task.input)
    state["current_reasoning"].append(reasoning_step)
    return {"input": task.input}

agent_input_component = AgentInputComponent(fn=agent_input_fn)

# Define Agent Prompt
def react_prompt_fn(task: Task, state: Dict[str, Any], input: str, tools: List[BaseTool]) -> List[ChatMessage]:
    chat_formatter = ReActChatFormatter()
    return chat_formatter.format(
        tools,
        chat_history=task.memory.get() + state["memory"].get_all(),
        current_reasoning=state["current_reasoning"],
    )

react_prompt_component = AgentFnComponent(
    fn=react_prompt_fn, partial_dict={"tools": [sql_tool]}
)

# Define Agent Output Parser
def parse_react_output_fn(task: Task, state: Dict[str, Any], chat_response: ChatResponse):
    output_parser = ReActOutputParser()
    reasoning_step = output_parser.parse(chat_response.message.content)
    return {"done": reasoning_step.is_done, "reasoning_step": reasoning_step}

parse_react_output = AgentFnComponent(fn=parse_react_output_fn)

# Define Tool Runner
def run_tool_fn(task: Task, state: Dict[str, Any], reasoning_step: ActionReasoningStep):
    tool_runner_component = ToolRunnerComponent([sql_tool], callback_manager=task.callback_manager)
    tool_output = tool_runner_component.run_component(
        tool_name=reasoning_step.action,
        tool_input=reasoning_step.action_input,
    )
    observation_step = ObservationReasoningStep(observation=str(tool_output))
    state["current_reasoning"].append(observation_step)
    return {"response_str": observation_step.get_content(), "is_done": False}

run_tool = AgentFnComponent(fn=run_tool_fn)

# Stitch Together Agent Query Pipeline
qp.add_modules(
    {
        "agent_input": agent_input_component,
        "react_prompt": react_prompt_component,
        "llm": OpenAI(model="gpt-4-1106-preview"),
        "react_output_parser": parse_react_output,
        "run_tool": run_tool,
    }
)

qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])
qp.add_link(
    "react_output_parser",
    "run_tool",
    condition_fn=lambda x: not x["done"],
    input_fn=lambda x: x["reasoning_step"],
)

5. Setup and Run the Agent

Finally, set up the agent worker and run the agent.

from llama_index.core.agent import QueryPipelineAgentWorker
from llama_index.core.callbacks import CallbackManager

agent_worker = QueryPipelineAgentWorker(qp)
agent = agent_worker.as_agent(
    callback_manager=CallbackManager([]), verbose=True
)

task = agent.create_task("What are some tracks from the artist AC/DC? Limit it to 3")
step_output = agent.run_step(task.task_id)
response = agent.finalize_response(task.task_id)
print(str(response))

By following these steps, you can replace the NLSQLTableQueryEngine with your existing text-2-SQL query pipeline and integrate it into the ReAct Agent Pipeline ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[deybelman]

@dosu I just tried this code:

from llama_index.core.query_engine.custom import CustomQueryEngine
from llama_index.core.query_pipeline import QueryPipeline
from llama_index.legacy.bridge.pydantic import Field
from llama_index.legacy.core.response.schema import RESPONSE_TYPE

class YourCustomText2SQLQueryEngine(CustomQueryEngine):
    query_pipeline: QueryPipeline = Field(
        description="Query Pipeline to use for Q&A.",
    )

    def __init__(self, query_pipeline: QueryPipeline, **kwargs):
        super().__init__(**kwargs)
        self.query_pipeline = query_pipeline

    def custom_query(self, query_str: str) -> RESPONSE_TYPE:
        return self.query_pipeline.run(query_str=query_str)

    async def acustom_query(self, query_str: str) -> RESPONSE_TYPE:
        return await self.query_pipeline.arun(query_str=query_str)

# Assuming `qp` is your existing query pipeline
sql_query_engine = YourCustomText2SQLQueryEngine(query_pipeline=qp)

It still gives me the following error:

---------------------------------------------------------------------------
ValidationError                           Traceback (most recent call last)
Cell In[62], line 22
     19         return await self.query_pipeline.arun(query_str=query_str)
     21 # Assuming `qp` is your existing query pipeline
---> 22 sql_query_engine = YourCustomText2SQLQueryEngine(query_pipeline=qp)

Cell In[62], line 12, in YourCustomText2SQLQueryEngine.__init__(self, query_pipeline, **kwargs)
     11 def __init__(self, query_pipeline: QueryPipeline, **kwargs):
---> 12     super().__init__(**kwargs)
     13     self.query_pipeline = query_pipeline

File /Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/pydantic/v1/main.py:341, in BaseModel.__init__(__pydantic_self__, **data)
    339 values, fields_set, validation_error = validate_model(__pydantic_self__.__class__, data)
    340 if validation_error:
--> 341     raise validation_error
    342 try:
    343     object_setattr(__pydantic_self__, '__dict__', values)

ValidationError: 1 validation error for YourCustomText2SQLQueryEngine
query_pipeline
  field required (type=value_error.missing)

How do I fix this? I don't know any other way to wrap a query pipeline as a query engine.

[dosubot[bot]]

To fix the ValidationError and correctly wrap your query pipeline as a query engine, you need to ensure that the query_pipeline field in the YourCustomText2SQLQueryEngine class is properly initialized and validated. The error occurs because Pydantic expects the query_pipeline field to be passed during initialization, but it is not being set correctly.

Here’s how you can fix it:

1. Define Your Custom Query Engine

Ensure that the query_pipeline field is properly defined and initialized in your custom query engine class. Use Pydantic's Field to define the field and ensure it is properly validated.

from llama_index.core.query_engine.custom import CustomQueryEngine
from llama_index.core.query_pipeline import QueryPipeline
from llama_index.legacy.bridge.pydantic import Field
from llama_index.legacy.core.response.schema import RESPONSE_TYPE

class YourCustomText2SQLQueryEngine(CustomQueryEngine):
    query_pipeline: QueryPipeline = Field(
        description="Query Pipeline to use for Q&A.",
    )

    def __init__(self, query_pipeline: QueryPipeline, **kwargs):
        super().__init__(**kwargs)
        self.query_pipeline = query_pipeline

    def custom_query(self, query_str: str) -> RESPONSE_TYPE:
        return self.query_pipeline.run(query_str=query_str)

    async def acustom_query(self, query_str: str) -> RESPONSE_TYPE:
        return await self.query_pipeline.arun(query_str=query_str)

2. Instantiate Your Custom Query Engine

When instantiating your custom query engine, ensure that the query_pipeline is passed correctly.

# Assuming `qp` is your existing query pipeline
sql_query_engine = YourCustomText2SQLQueryEngine(query_pipeline=qp)

3. Wrap Your Custom Query Engine in a Tool

Wrap your custom query engine in a QueryEngineTool so it can be used by the agent.

from llama_index.core.tools import QueryEngineTool, ToolMetadata

sql_tool = QueryEngineTool(
    query_engine=sql_query_engine,
    metadata=ToolMetadata(
        name="sql_tool",
        description="Useful for translating a natural language query into a SQL query"
    ),
)

4. Integrate the Tool into the ReAct Agent Pipeline

Integrate the sql_tool into your ReAct Agent Pipeline. This involves setting up the agent input, prompt, output parser, and tool runner components.

from llama_index.core.agent.react.types import ObservationReasoningStep
from llama_index.core.agent import Task
from llama_index.core.query_pipeline import AgentInputComponent, AgentFnComponent, ToolRunnerComponent
from llama_index.core.llms import OpenAI, ChatMessage, ChatResponse
from llama_index.core.agent.react.output_parser import ReActOutputParser

# Define Agent Input Component
def agent_input_fn(task: Task, state: Dict[str, Any]) -> Dict[str, Any]:
    if "current_reasoning" not in state:
        state["current_reasoning"] = []
    reasoning_step = ObservationReasoningStep(observation=task.input)
    state["current_reasoning"].append(reasoning_step)
    return {"input": task.input}

agent_input_component = AgentInputComponent(fn=agent_input_fn)

# Define Agent Prompt
def react_prompt_fn(task: Task, state: Dict[str, Any], input: str, tools: List[BaseTool]) -> List[ChatMessage]:
    chat_formatter = ReActChatFormatter()
    return chat_formatter.format(
        tools,
        chat_history=task.memory.get() + state["memory"].get_all(),
        current_reasoning=state["current_reasoning"],
    )

react_prompt_component = AgentFnComponent(
    fn=react_prompt_fn, partial_dict={"tools": [sql_tool]}
)

# Define Agent Output Parser
def parse_react_output_fn(task: Task, state: Dict[str, Any], chat_response: ChatResponse):
    output_parser = ReActOutputParser()
    reasoning_step = output_parser.parse(chat_response.message.content)
    return {"done": reasoning_step.is_done, "reasoning_step": reasoning_step}

parse_react_output = AgentFnComponent(fn=parse_react_output_fn)

# Define Tool Runner
def run_tool_fn(task: Task, state: Dict[str, Any], reasoning_step: ActionReasoningStep):
    tool_runner_component = ToolRunnerComponent([sql_tool], callback_manager=task.callback_manager)
    tool_output = tool_runner_component.run_component(
        tool_name=reasoning_step.action,
        tool_input=reasoning_step.action_input,
    )
    observation_step = ObservationReasoningStep(observation=str(tool_output))
    state["current_reasoning"].append(observation_step)
    return {"response_str": observation_step.get_content(), "is_done": False}

run_tool = AgentFnComponent(fn=run_tool_fn)

# Stitch Together Agent Query Pipeline
qp.add_modules(
    {
        "agent_input": agent_input_component,
        "react_prompt": react_prompt_component,
        "llm": OpenAI(model="gpt-4-1106-preview"),
        "react_output_parser": parse_react_output,
        "run_tool": run_tool,
    }
)

qp.add_chain(["agent_input", "react_prompt", "llm", "react_output_parser"])
qp.add_link(
    "react_output_parser",
    "run_tool",
    condition_fn=lambda x: not x["done"],
    input_fn=lambda x: x["reasoning_step"],
)

5. Setup and Run the Agent

Set up the agent worker and run the agent.

from llama_index.core.agent import QueryPipelineAgentWorker
from llama_index.core.callbacks import CallbackManager

agent_worker = QueryPipelineAgentWorker(qp)
agent = agent_worker.as_agent(
    callback_manager=CallbackManager([]), verbose=True
)

task = agent.create_task("What are some tracks from the artist AC/DC? Limit it to 3")
step_output = agent.run_step(task.task_id)
response = agent.finalize_response(task.task_id)
print(str(response))

By following these steps, you can replace the NLSQLTableQueryEngine with your existing text-2-SQL query pipeline and integrate it into the ReAct Agent Pipeline ^[1][2].

To continue talking to Dosu, mention @dosu.

[Ahmed-Elbaqary]

@deybelman Hope you have solved it by now;
I have the same Idea as you, looked into the responses and llamaindex docs

Here is how you can do it:

from llama_index.core.query_engine import CustomQueryEngine

class CustomQueryEngine(CustomQueryEngine):
    """Custom Query Engine."""
    query_pipeline: QueryPipeline
    
    def custom_query(self, query_str: str):
        return self.query_pipeline.run(query_str=query_str)

    async def acustom_query(self, query_str: str) -> RESPONSE_TYPE:
        return await self.query_pipeline.arun(query_str=query_str)

Then build your query engine obj

sql_query_engine = CustomQueryEngine(query_pipeline=qp)

Finally the sub question query engine wrapper

from llama_index.core.tools import QueryEngineTool, ToolMetadata

# # setup base query engine as tool
query_engine_tools = [
    QueryEngineTool(
        query_engine=sql_query_engine,
        metadata=ToolMetadata(
        name="sql_tool",
        description="Useful for translating a natural language query into a SQL query"
        ),
    ),
]


query_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=query_engine_tools,
    use_async=True,
)

To use it, you may have to use asyncio like this:

import asyncio
import nest_asyncio

nest_asyncio.apply()

coroutine = query_engine.aquery(
    "your question here"
)
response = asyncio.run(coroutine)

I faced some issues to build this and had to add a final tool at the end of my query pipeline just to give the custom query engine the correct response as string, just like this;

from llama_index.core.query_pipeline import FnComponent
from llama_index.core.llms import ChatResponse

def parse_response_to_text(response: ChatResponse) -> str:
    """Parse response as text to use in subquestion query engine."""
    response = response.message.content
    return response
    
response_parser = FnComponent(fn=parse_response_to_text)

add this at the end of your query pipeline before giving it to the custom query engine