Query Transform Cookbook

A user query can be transformed and decomposed in many ways before being executed as part of a RAG query engine, agent, or any other pipeline.

In this guide we show you different ways to transform, decompose queries, and find the set of relevant tools. Each technique might be applicable for different use cases!

For naming purposes, we define the underlying pipeline as a "tool". Here are the different query transformations:

Routing: Keep the query, but identify the relevant subset of tools that the query applies to. Output those tools as the relevant choices.
Query-Rewriting: Keep the tools, but rewrite the query in a variety of different ways to execute against the same tools.
Sub-Questions: Decompose queries into multiple sub-questions over different tools (identified by their metadata).
ReAct Agent Tool Picking: Given the initial query, identify 1) the tool to pick, and 2) the query to execute on the tool.

The goal of this guide is to show you how to use these query transforms as modular components. Of course, each of these components plug into a bigger system (e.g. the sub-question generator is a part of our SubQuestionQueryEngine) - and the guides for each of these are linked below.

Take a look and let us know your thoughts!

python

%pip install llama-index-question-gen-openai
%pip install llama-index-llms-openai

python

from IPython.display import Markdown, display


# define prompt viewing function
def display_prompt_dict(prompts_dict):
    for k, p in prompts_dict.items():
        text_md = f"**Prompt Key**: {k}
" f"**Text:** 
"
        display(Markdown(text_md))
        print(p.get_template())
        display(Markdown("

"))

Routing

In this example, we show how a query can be used to select the set of relevant tool choices.

We use our selector abstraction to pick the relevant tool(s) - it can be a single tool, or a multiple tool depending on the abstraction.

We have four selectors: combination of (LLM or function calling) x (single selection or multi-selection)

python

from llama_index.core.selectors import LLMSingleSelector, LLMMultiSelector
from llama_index.core.selectors import (
    PydanticMultiSelector,
    PydanticSingleSelector,
)

python

# pydantic selectors feed in pydantic objects to a function calling API
# single selector (pydantic, function calling)
# selector = PydanticSingleSelector.from_defaults()

# multi selector (pydantic, function calling)
# selector = PydanticMultiSelector.from_defaults()

# LLM selectors use text completion endpoints
# single selector (LLM)
# selector = LLMSingleSelector.from_defaults()
# multi selector (LLM)
selector = LLMMultiSelector.from_defaults()

python

from llama_index.core.tools import ToolMetadata

tool_choices = [
    ToolMetadata(
        name="covid_nyt",
        description=("This tool contains a NYT news article about COVID-19"),
    ),
    ToolMetadata(
        name="covid_wiki",
        description=("This tool contains the Wikipedia page about COVID-19"),
    ),
    ToolMetadata(
        name="covid_tesla",
        description=("This tool contains the Wikipedia page about apples"),
    ),
]

python

display_prompt_dict(selector.get_prompts())

python

selector_result = selector.select(
    tool_choices, query="Tell me more about COVID-19"
)

python

selector_result.selections

Learn more about our routing abstractions in our dedicated Router page.

Query Rewriting

In this section, we show you how to rewrite queries into multiple queries. You can then execute all these queries against a retriever.

This is a key step in advanced retrieval techniques. By doing query rewriting, you can generate multiple queries for [ensemble retrieval] and [fusion], leading to higher-quality retrieved results.

Unlike the sub-question generator, this is just a prompt call, and exists independently of tools.

Query Rewriting (Custom)

Here we show you how to use a prompt to generate multiple queries, using our LLM and prompt abstractions.

python

from llama_index.core import PromptTemplate
from llama_index.llms.openai import OpenAI

query_gen_str = """\
You are a helpful assistant that generates multiple search queries based on a \
single input query. Generate {num_queries} search queries, one on each line, \
related to the following input query:
Query: {query}
Queries:
"""
query_gen_prompt = PromptTemplate(query_gen_str)

llm = OpenAI(model="gpt-3.5-turbo")


def generate_queries(query: str, llm, num_queries: int = 4):
    response = llm.predict(
        query_gen_prompt, num_queries=num_queries, query=query
    )
    # assume LLM proper put each query on a newline
    queries = response.split("\n")
    queries_str = "\n".join(queries)
    print(f"Generated queries:\n{queries_str}")
    return queries

python

queries = generate_queries("What happened at Interleaf and Viaweb?", llm)

python

queries

For more details about an e2e implementation with a retriever, check out our guides on our fusion retriever:

Query Rewriting (using QueryTransform)

In this section we show you how to do query transformations using our QueryTransform class.

python

from llama_index.core.indices.query.query_transform import HyDEQueryTransform
from llama_index.llms.openai import OpenAI

python

hyde = HyDEQueryTransform(include_original=True)
llm = OpenAI(model="gpt-3.5-turbo")

query_bundle = hyde.run("What is Bel?")

This generates a query bundle that contains the original query, but also custom_embedding_strs representing the queries that should be embedded.

python

new_query.custom_embedding_strs

Sub-Questions

Given a set of tools and a user query, decide both the 1) set of sub-questions to generate, and 2) the tools that each sub-question should run over.

We run through an example using the OpenAIQuestionGenerator, which depends on function calling, and also the LLMQuestionGenerator, which depends on prompting.

python

from llama_index.core.question_gen import LLMQuestionGenerator
from llama_index.question_gen.openai import OpenAIQuestionGenerator
from llama_index.llms.openai import OpenAI

python

llm = OpenAI()
question_gen = OpenAIQuestionGenerator.from_defaults(llm=llm)

python

display_prompt_dict(question_gen.get_prompts())

python

from llama_index.core.tools import ToolMetadata

tool_choices = [
    ToolMetadata(
        name="uber_2021_10k",
        description=(
            "Provides information about Uber financials for year 2021"
        ),
    ),
    ToolMetadata(
        name="lyft_2021_10k",
        description=(
            "Provides information about Lyft financials for year 2021"
        ),
    ),
]

python

from llama_index.core import QueryBundle

query_str = "Compare and contrast Uber and Lyft"
choices = question_gen.generate(tool_choices, QueryBundle(query_str=query_str))

The outputs are SubQuestion Pydantic objects.

python

choices

For details on how to plug this into your RAG pipeline in a more packaged fashion, check out our SubQuestionQueryEngine.

Query Transformation with ReAct Prompt

ReAct is a popular framework for agents, and here we show how the core ReAct prompt can be used to transform queries.

We use the ReActChatFormatter to get the set of input messages for the LLM.

python

from llama_index.core.agent import ReActChatFormatter
from llama_index.core.agent.react.output_parser import ReActOutputParser
from llama_index.core.tools import FunctionTool
from llama_index.core.llms import ChatMessage

python

def execute_sql(sql: str) -> str:
    """Given a SQL input string, execute it."""
    # NOTE: This is a mock function
    return f"Executed {sql}"


def add(a: int, b: int) -> int:
    """Add two numbers."""
    return a + b


tool1 = FunctionTool.from_defaults(fn=execute_sql)
tool2 = FunctionTool.from_defaults(fn=add)
tools = [tool1, tool2]

Here we get the input prompt messages to pass to the LLM. Take a look!

python

chat_formatter = ReActChatFormatter()
output_parser = ReActOutputParser()
input_msgs = chat_formatter.format(
    tools,
    [
        ChatMessage(
            content="Can you find the top three rows from the table named `revenue_years`",
            role="user",
        )
    ],
)
input_msgs

Next we get the output from the model.

python

llm = OpenAI(model="gpt-4-1106-preview")

python

response = llm.chat(input_msgs)

Finally we use our ReActOutputParser to parse the content into a structured output, and analyze the action inputs.

python

reasoning_step = output_parser.parse(response.message.content)

python

reasoning_step.action_input