Sub-questions not being triggered in ChatCallbackHandler using lastest version LlamaIndex OpenAIAgent with SubQuestionQueryEngines

[byamasu-patrick]

Description:

I'm encountering an issue after updating the version of LlamaIndex in our project (SEC Insights). The problem specifically relates to sub-questions not being displayed on the frontend.

Steps to Reproduce:

Update the LlamaIndex to the latest version.
Use OpenAIAgent with SubQuestionQueryEngine.
Trigger sub-questions during a query.
Expected Behavior: The CBEventType.SUB_QUESTION should be triggered inside ChatCallbackHandler, allowing sub-questions and their responses to be streamed and displayed on the frontend.

Observed Behavior: After updating the version, the CBEventType.SUB_QUESTION is not triggered, preventing subquery answers from streaming to the frontend.

Additional Context: This issue is impacting the ability to handle sub-questions in a live environment. Has anyone else encountered this problem, or are there any workarounds?

Environment:

LlamaIndex Version: "^0.11.13"
Llama Index Agent Openai: "^0.3.4"

[tannonk]

Hey @byamasu-patrick,

I'm running into the same issue with llama-index-agent-openai==0.4.6 and llama-index-core==0.12.23. Do you know of any work arounds?

Thanks in advance!

[byamasu-patrick]

Hey @tannonk ,

Yes, I solved it!

I created a custom class that inherits from SubQuestionQueryEngine and used that to restore the expected behavior. I shared the full solution in the sec-insights Discord server, but here's a quick summary for you:

I defined a CustomSubQuestionQueryEngine and used it like this:

CustomSubQuestionQueryEngine.from_defaults(
    query_engine_tools=[vector_sql_query_engine_tool],
    llm=Settings.llm,
    response_synthesizer=response_synth,
    verbose=settings.VERBOSE,
    question_gen=question_gen,
    use_async=True,
    callback_manager=callback_manager,
)

This helped ensure CBEventType.SUB_QUESTION is correctly triggered and streamed to the frontend via ChatCallbackHandler.

import logging
from typing import Optional, Sequence

from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.llms.llm import LLM
from llama_index.core.query_engine.sub_question_query_engine import (
  SubQuestionQueryEngine,
)
from llama_index.core.question_gen.llm_generators import LLMQuestionGenerator
from llama_index.core.question_gen.types import BaseQuestionGenerator
from llama_index.core.response_synthesizers import (
  BaseSynthesizer,
  get_response_synthesizer,
)
from llama_index.core.settings import Settings
from llama_index.core.tools.query_engine import QueryEngineTool

logger = logging.getLogger(__name__)


class CustomSubQuestionQueryEngine(SubQuestionQueryEngine):
    """Custom Sub question query engine.

    A query engine that breaks down a complex query (e.g. compare and contrast) into
    many sub questions and their target query engine for execution.
    After executing all sub questions, all responses are gathered and sent to
    response synthesizer to produce the final response.

    Args:
        question_gen (BaseQuestionGenerator): A module for generating sub questions
            given a complex question and tools.
        response_synthesizer (BaseSynthesizer): A response synthesizer for
            generating the final response
        query_engine_tools (Sequence[QueryEngineTool]): Tools to answer the
            sub questions.
        verbose (bool): whether to print intermediate questions and answers.
            Defaults to True
        use_async (bool): whether to execute the sub questions with asyncio.
            Defaults to True
    """

    def __init__(
        self,
        question_gen: BaseQuestionGenerator,
        response_synthesizer: BaseSynthesizer,
        query_engine_tools: Sequence[QueryEngineTool],
        callback_manager: Optional[CallbackManager] = None,
        verbose: bool = True,
        use_async: bool = False,
    ) -> None:
        self._question_gen = question_gen
        self._response_synthesizer = response_synthesizer
        self._metadatas = [x.metadata for x in query_engine_tools]
        self._query_engines = {
            tool.metadata.name: tool.query_engine for tool in query_engine_tools
        }
        self._verbose = verbose
        self._use_async = use_async
        super().__init__(
            question_gen=question_gen,
            response_synthesizer=response_synthesizer,
            callback_manager=callback_manager,
            query_engine_tools=query_engine_tools,
            use_async=use_async,
            verbose=verbose,
        )

    @classmethod
    def from_defaults(
        cls,
        query_engine_tools: Sequence[QueryEngineTool],
        llm: Optional[LLM] = None,
        question_gen: Optional[BaseQuestionGenerator] = None,
        response_synthesizer: Optional[BaseSynthesizer] = None,
        verbose: bool = True,
        use_async: bool = True,
        callback_manager: Optional[CallbackManager] = None,
    ) -> "CustomSubQuestionQueryEngine":
        llm = llm or Settings.llm
        if question_gen is None:
            try:
                from llama_index.question_gen.openai import (  # pants: no-infer-dep
                  OpenAIQuestionGenerator,
                )

                # try to use OpenAI function calling based question generator.
                # if incompatible, use general LLM question generator
                question_gen = OpenAIQuestionGenerator.from_defaults(llm=llm)

            except ImportError as e:
                raise ImportError(
                    "`llama-index-question-gen-openai` package cannot be found. "
                    "Please install it by using `pip install `llama-index-question-gen-openai`"
                )
            except ValueError:
                question_gen = LLMQuestionGenerator.from_defaults(llm=llm)

        synth = response_synthesizer or get_response_synthesizer(
            llm=llm,
            callback_manager=callback_manager,
            use_async=use_async,
        )

        return cls(
            question_gen,
            synth,
            query_engine_tools,
            callback_manager=callback_manager,
            verbose=verbose,
            use_async=use_async,
        )

Let me know if you’d like the full class implementation or help integrating it!

Thanks,

[byamasu-patrick]

Note: vector_sql_query_engine_tool can be replaced based on how you're implementing your tools.

The key part is using the custom implementation, which ensures SEC Insights continues to work as expected, especially since it's no longer actively maintained. This workaround helps preserve sub-question handling and streaming without depending on future updates.

[byamasu-patrick]

@tannok The link of where I explain the main issue can be found here https://discord.com/channels/1059199217496772688/1150942525968879636/1290681061994528939

[tannonk]

Hey @byamasu-patrick, thanks for the detailed and blazingly fast response! Looks like a nice solution you came up with!

I'm still quite new to the app so it might take me a while to integrate the necessary changes. Would you happen to have your working version somewhere on github? That might help me to compare between the versions and clarify a few questions...

Thanks!

[byamasu-patrick]

Hey @tannonk,

Yeah sure! I have a working customized solution and I’ll be happy to share it with you.

The version I have had access to DB tables (so some parts are tailored to SEC Insights), but the overall structure and working behavior, especially how it handles sub-questions, is the same and should be easy to adapt to your setup.

I’ll send over the GitHub link shortly so you can compare and adjust as needed!

Cheers,
@byamasu-patrick

[byamasu-patrick]

Hey @tannonk,

I’ve shared the repo with you, feel free to explore it!

You might want to check out these key parts related to the SubQuestionQueryEngine implementation:

• Custom implementation:
  [custom_sub_question_query_engine.py](https://github.com/byamasu-patrick/modellmind.info4.org/blob/master/chat/custom_sub_question_query_engine.py)

• Where it’s used:
  [engine.py#L452](https://github.com/byamasu-patrick/modellmind.info4.org/blob/master/chat/engine.py#L452)

• Relevant endpoint:
  [conversation.py#L130](https://github.com/byamasu-patrick/modellmind.info4.org/blob/master/api/endpoints/conversation.py#L130)

Let me know if you have any questions or need help integrating it into your setup.

Thanks,
@byamasu-patrick