LlamaIndex SubQuestionQueryEngine — over-decomposition and vague sub-queries

from llama_index.core.query_engine import SubQuestionQueryEngine
from llama_index.core.question_gen import LLMQuestionGenerator
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.llms.openai import OpenAI

llm = OpenAI(model="gpt-4o-mini")

# --- Tools with rich descriptions ---
policy_tool = QueryEngineTool(
    query_engine=policy_index.as_query_engine(similarity_top_k=6),
    metadata=ToolMetadata(
        name="policies",
        description="Company policies: refund, return, warranty, shipping, ToS.",
    ),
)

product_tool = QueryEngineTool(
    query_engine=product_index.as_query_engine(similarity_top_k=6),
    metadata=ToolMetadata(
        name="products",
        description="Product catalog: features, specs, pricing.",
    ),
)

# --- Custom question generator prompt (capped) ---
CUSTOM_QUESTION_GEN_PROMPT = """
Given a user question and a list of tools, generate the MINIMUM number of
sub-questions (at most 4) needed to answer the user's question. Each sub-question
should be targeted at exactly one tool and be specific enough to retrieve
relevant information.

Do NOT generate sub-questions that are vague or redundant.
Do NOT generate more than 4 sub-questions total.

Tools:
{tools_str}

User question: {query_str}

Sub-questions (JSON array):
"""

# --- SubQuestionQueryEngine with async and bounded generator ---
question_gen = LLMQuestionGenerator.from_defaults(
    llm=llm,
    prompt_template_str=CUSTOM_QUESTION_GEN_PROMPT,
)

subq_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=[policy_tool, product_tool],
    question_gen=question_gen,
    llm=llm,
    use_async=True,           # parallel sub-question execution
    verbose=True,             # prints sub-questions
)

response = subq_engine.query(
    "Compare the refund policy with product warranty terms for the ACME-5000."
)
print(response)
print(f"\nSource nodes: {len(response.source_nodes)}")

from llama_index.core.query_engine import SubQuestionQueryEngine
from llama_index.core.callbacks import CallbackManager, LlamaDebugHandler, CBEventType

# Enable the debug handler to see all events
debug = LlamaDebugHandler(print_trace_on_end=False)
cb_manager = CallbackManager([debug])

subq_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=[policy_tool, product_tool],
    llm=llm,
    use_async=True,
    verbose=True,
    callback_manager=cb_manager,
)

response = subq_engine.query("...")
print(f"Final: {response.response}")

# --- Extract per-sub-question intermediate steps ---
sub_events = debug.get_events(CBEventType.SUB_QUESTION)

for event in sub_events:
    payload = event.payload or {}
    sub_q_ans = payload.get("sub_question")   # type: SubQuestionAnswerPair
    if sub_q_ans:
        print(f"\nSub-question: {sub_q_ans.sub_q.sub_question}")
        print(f"Tool:         {sub_q_ans.sub_q.tool_name}")
        print(f"Answer:       {sub_q_ans.answer[:200]}")
        print(f"Sources:      {[n.metadata.get('source') for n in sub_q_ans.sources or []]}")

# --- For UI: attach citations per sub-answer ---
def render_response_with_citations(response, debug):
    output = [response.response, "\n\n**Sources:**"]
    for event in debug.get_events(CBEventType.SUB_QUESTION):
        sq = (event.payload or {}).get("sub_question")
        if sq:
            output.append(
                f"- **{sq.sub_q.sub_question}** ({sq.sub_q.tool_name}): "
                f"{[n.metadata.get('source', '?') for n in sq.sources or []]}"
            )
    return "\n".join(output)

print(render_response_with_citations(response, debug))

# --- Bad match for SubQuestion — single-domain query ---
# query = "What is the refund policy?"
# Just use policy_engine directly

# --- Good match for SubQuestion — multi-domain query ---
# query = "Compare the refund policy against the ACME-5000 warranty"
# Needs policies AND products; sub-questions naturally decompose

# --- Bad match for SubQuestion — iterative reasoning ---
# query = "Investigate why sales dropped last quarter and recommend a fix"
# Needs iterative reasoning + planning; use an agent

# --- Simple heuristic router ---
from llama_index.core.query_engine import RouterQueryEngine, SubQuestionQueryEngine

# Set up both engines
router_engine = RouterQueryEngine(
    selector=PydanticSingleSelector.from_defaults(llm=llm),
    query_engine_tools=[policy_tool, product_tool],
)
subq_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=[policy_tool, product_tool],
    llm=llm,
    use_async=True,
)

# --- Meta-router that picks between them ---
def smart_query(query: str) -> str:
    # Heuristic: multi-domain markers
    multi_domain_markers = [
        "compare", "and", "also", "both", "in addition",
        "how does X relate to Y", "along with",
    ]
    if any(marker in query.lower() for marker in multi_domain_markers):
        return subq_engine.query(query).response
    return router_engine.query(query).response

# --- Or use an LLM upfront to route ---
def llm_smart_query(query: str) -> str:
    """Use a cheap LLM call to decide."""
    decision = llm.complete(
        f"Should this query be answered by decomposing into sub-questions "
        f"across multiple domains, or by a single specialized tool?\n\n"
        f"Query: {query}\n\n"
        f"Answer 'single' or 'decompose':"
    ).text.strip().lower()

    if "decompose" in decision:
        return subq_engine.query(query).response
    return router_engine.query(query).response