datarobot_genai.llama_index.agent

agent

DataRobotLiteLLM

Bases: LiteLLM

LiteLLM wrapper providing chat/function capability metadata for LlamaIndex.

Source code in datarobot_genai/llama_index/agent.py

class DataRobotLiteLLM(LiteLLM):
    """LiteLLM wrapper providing chat/function capability metadata for LlamaIndex."""

    @property
    def metadata(self) -> LLMMetadata:
        """Return LLM metadata."""
        return LLMMetadata(
            context_window=128000,
            num_output=self.max_tokens or -1,
            is_chat_model=True,
            is_function_calling_model=True,
            model_name=self.model,
        )

metadata property

metadata: LLMMetadata

Return LLM metadata.

LlamaIndexAgent

Bases: BaseAgent[BaseTool], ABC

Abstract base agent for LlamaIndex workflows.

Framework-specific parameters:

• allow_parallel_tool_calls: when true (default), LlamaIndex workflow agents that support it may issue parallel tool calls

Source code in datarobot_genai/llama_index/agent.py

class LlamaIndexAgent(BaseAgent[BaseTool], abc.ABC):
    """Abstract base agent for LlamaIndex workflows.

    Framework-specific parameters:

    - ``allow_parallel_tool_calls``: when true (default), LlamaIndex workflow
      agents that support it may issue parallel tool calls
    """

    def __init__(
        self,
        api_key: str | None = None,
        api_base: str | None = None,
        llm: Any | None = None,
        tools: list[BaseTool] | None = None,
        verbose: bool = True,
        timeout: int = 90,
        forwarded_headers: dict[str, str] | None = None,
        max_history_messages: int | None = None,
        model: str | None = None,
        structured_history: bool = True,
        allow_parallel_tool_calls: bool = True,
    ) -> None:
        super().__init__(
            api_key=api_key,
            api_base=api_base,
            llm=llm,
            tools=tools,
            verbose=verbose,
            timeout=timeout,
            forwarded_headers=forwarded_headers,
            max_history_messages=max_history_messages,
            model=model,
        )
        self._structured_history = structured_history
        self.allow_parallel_tool_calls = allow_parallel_tool_calls

    @property
    def structured_history(self) -> bool:
        """When true, prior turns are fed to the model as structured native
        ``ChatMessage`` history (tool calls preserved) instead of the text
        ``{chat_history}`` summary. Only applies when the prompt has no
        ``{chat_history}`` placeholder.

        Defaults to ``True`` so multi-turn history is replayed by default; pass
        ``structured_history=False`` to opt out.
        """
        return self._structured_history

    @abc.abstractmethod
    async def build_workflow(self) -> Any:
        """Return an AgentWorkflow instance ready to run."""
        raise NotImplementedError

    @abc.abstractmethod
    def extract_response_text(self, result_state: Any, events: list[Any]) -> str:
        """Extract final response text from workflow state and/or events."""
        raise NotImplementedError

    async def invoke(self, run_agent_input: RunAgentInput) -> InvokeReturn:
        """Run the LlamaIndex workflow with the provided completion parameters."""
        user_prompt_content = extract_user_prompt_content(run_agent_input)
        input_message = str(user_prompt_content)

        # Prior turns reach the model as a text {chat_history} summary when the
        # prompt uses the placeholder, or as structured native ChatMessage history
        # (tool calls preserved) when structured_history is opt-in.
        structured_chat_history = None
        if "{chat_history}" in input_message:
            history_summary = self.build_history_summary(run_agent_input)
            formatted_history = (
                f"\n\nPrior conversation:\n{history_summary}" if history_summary else ""
            )
            input_message = input_message.replace("{chat_history}", formatted_history)
        elif self.structured_history:
            structured_chat_history = (
                ag_ui_history_to_chat_messages(self.history_messages(run_agent_input)) or None
            )

        input_message = prepend_streaming_memory_to_prompt(input_message, run_agent_input)

        logger.info(f"Running agent with user prompt: {input_message}")

        thread_id = run_agent_input.thread_id
        run_id = run_agent_input.run_id
        usage_metrics: UsageMetrics = default_usage_metrics()

        # Partial AG-UI: lifecycle + text + tool calls + steps
        yield (
            RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=run_id),
            None,
            usage_metrics,
        )

        # Subclasses may implement build_workflow as async or sync; support both.
        built: Any = self.build_workflow()
        workflow = await built if inspect.isawaitable(built) else built
        run_kwargs: dict[str, Any] = {"user_msg": input_message}
        if structured_chat_history:
            run_kwargs["chat_history"] = structured_chat_history
        handler = workflow.run(**run_kwargs)

        events: list[Any] = []
        current_agent_name: str | None = None
        message_id = str(uuid.uuid4())
        text_started = False
        any_text_emitted = False
        # Id of the most recently closed assistant text bubble. A following tool call
        # names it as its parent_message_id so the call attaches to the message that
        # introduced it (AG-UI grouping) instead of an orphan id; reset on a tool
        # result so a later text-less step can't reuse a stale bubble id.
        last_text_message_id: str | None = None

        async for event in handler.stream_events():
            events.append(event)

            # Detect agent transition first so the first delta of a new agent
            # is attributed to the new agent's text bubble, not the previous one.
            new_agent = getattr(event, "current_agent_name", None)
            if new_agent is not None and new_agent != current_agent_name:
                # Close any open text message from the outgoing agent so each
                # step's text commits as its own AG-UI message.
                if text_started:
                    yield (
                        TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=message_id),
                        None,
                        usage_metrics,
                    )
                    text_started = False
                if current_agent_name is not None:
                    yield (
                        StepFinishedEvent(
                            type=EventType.STEP_FINISHED, step_name=current_agent_name
                        ),
                        None,
                        usage_metrics,
                    )
                yield (
                    StepStartedEvent(type=EventType.STEP_STARTED, step_name=new_agent),
                    None,
                    usage_metrics,
                )
                current_agent_name = new_agent
                # Fresh id for the new step's first text bubble.
                message_id = str(uuid.uuid4())
                logger.info(f"Agent: {current_agent_name}")

            # Reasoning models surface incremental thinking on AgentStream.thinking_delta
            # when the LLM integration populates it; the stock litellm wrapper does not, so
            # fall back to the raw LiteLLM chunk (event.raw). Emit as a self-contained AG-UI
            # reasoning chunk before any text for this event.
            thinking_delta = getattr(event, "thinking_delta", None) or _thinking_delta_from_raw(
                getattr(event, "raw", None)
            )
            if thinking_delta:
                yield (
                    ReasoningMessageChunkEvent(
                        type=EventType.REASONING_MESSAGE_CHUNK,
                        # Its own message id, distinct from the assistant text so a frontend
                        # grouping by id renders reasoning as its own block instead of folding
                        # it into the text bubble. Derived (uuid5) from the text id: a valid
                        # UUID, stable across this message's chunks (so they group), no state.
                        message_id=str(uuid.uuid5(uuid.NAMESPACE_OID, f"{message_id}-reasoning")),
                        delta=thinking_delta,
                    ),
                    None,
                    usage_metrics,
                )

            # Best-effort extraction of incremental text from LlamaIndex events
            delta: str | None = None
            try:
                if hasattr(event, "delta") and isinstance(getattr(event, "delta"), str):
                    delta = getattr(event, "delta")
                # Some event types may carry incremental text under "text" or similar
                elif hasattr(event, "text") and isinstance(getattr(event, "text"), str):
                    delta = getattr(event, "text")
            except Exception:
                # Ignore malformed events and continue
                delta = None

            if delta:
                if not text_started:
                    yield (
                        TextMessageStartEvent(
                            type=EventType.TEXT_MESSAGE_START, message_id=message_id
                        ),
                        None,
                        usage_metrics,
                    )
                    text_started = True
                    any_text_emitted = True
                yield (
                    TextMessageContentEvent(
                        type=EventType.TEXT_MESSAGE_CONTENT, message_id=message_id, delta=delta
                    ),
                    None,
                    usage_metrics,
                )

            event_type = type(event).__name__
            if event_type == "AgentInput" and hasattr(event, "input"):
                logger.info(f"Input: {getattr(event, 'input')}")
            elif event_type == "AgentOutput":
                # When no deltas streamed for this turn, attach the full response
                # content to the current step's open text message. Step transition
                # or end-of-run closes it.
                resp = getattr(event, "response", None)
                if resp is not None and hasattr(resp, "content") and getattr(resp, "content"):
                    content_str = str(getattr(resp, "content"))
                    logger.info(f"Output: {content_str}")
                    if not text_started:
                        yield (
                            TextMessageStartEvent(
                                type=EventType.TEXT_MESSAGE_START, message_id=message_id
                            ),
                            None,
                            usage_metrics,
                        )
                        text_started = True
                        any_text_emitted = True
                        yield (
                            TextMessageContentEvent(
                                type=EventType.TEXT_MESSAGE_CONTENT,
                                message_id=message_id,
                                delta=content_str,
                            ),
                            None,
                            usage_metrics,
                        )
                # Planned tool calls
                tcalls = getattr(event, "tool_calls", None)
                if isinstance(tcalls, list) and tcalls:
                    names = []
                    for c in tcalls:
                        try:
                            nm = getattr(c, "tool_name", None) or (
                                c.get("tool_name") if isinstance(c, dict) else None
                            )
                            if nm:
                                names.append(str(nm))
                        except Exception:
                            pass
                    if names:
                        logger.info(f"Planning to use tools: {names}")
            elif event_type == "ToolCallResult":
                # The tool-calling step is complete; clear the captured bubble id so a
                # later text-less step doesn't attach its tool calls to a stale bubble.
                last_text_message_id = None
                tname = getattr(event, "tool_name", None)
                tid = getattr(event, "tool_id", None)
                tkwargs = getattr(event, "tool_kwargs", None)
                tout = getattr(event, "tool_output", None)
                logger.info(f"Tool Result: {tname}")
                logger.debug(f"Arguments: {tkwargs}")
                logger.debug(f"Output: {tout}")
                yield (
                    ToolCallEndEvent(
                        type=EventType.TOOL_CALL_END,
                        tool_call_id=tid,
                    ),
                    None,
                    usage_metrics,
                )
                yield (
                    ToolCallResultEvent(
                        type=EventType.TOOL_CALL_RESULT,
                        message_id=tid,
                        tool_call_id=tid,
                        content=json.dumps(tout, default=str),
                        role="tool",
                    ),
                    None,
                    usage_metrics,
                )
            elif event_type == "ToolCall":
                # Close any open text message so the tool widget renders below
                # the preceding text and the next text run starts a fresh bubble.
                if text_started:
                    yield (
                        TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=message_id),
                        None,
                        usage_metrics,
                    )
                    text_started = False
                    # Capture the closed bubble's id (before minting a fresh one) so the
                    # tool call below can name it as its parent_message_id.
                    last_text_message_id = message_id
                    message_id = str(uuid.uuid4())
                tname = getattr(event, "tool_name", None)
                tkwargs = getattr(event, "tool_kwargs", None)
                tid = getattr(event, "tool_id", None)
                logger.info(f"Calling Tool: {tname}")
                logger.debug(f"With arguments: {tkwargs}")
                yield (
                    ToolCallStartEvent(
                        type=EventType.TOOL_CALL_START,
                        tool_call_id=tid,
                        tool_call_name=tname,
                        parent_message_id=last_text_message_id or "",
                    ),
                    None,
                    usage_metrics,
                )
                yield (
                    ToolCallArgsEvent(
                        type=EventType.TOOL_CALL_ARGS,
                        tool_call_id=tid,
                        delta=json.dumps(tkwargs, default=str),
                    ),
                    None,
                    usage_metrics,
                )

        # Close any text message still open from the final agent before the
        # step is finished so the message commits inside its step boundaries.
        if text_started:
            yield (
                TextMessageEndEvent(
                    type=EventType.TEXT_MESSAGE_END,
                    message_id=message_id,
                ),
                None,
                usage_metrics,
            )
            text_started = False
            message_id = str(uuid.uuid4())

        # After streaming completes, build final interactions and finish chunk
        # Extract state from workflow context (supports sync/async get or attribute)
        state = None
        ctx = getattr(handler, "ctx", None)
        try:
            if ctx is not None:
                get = getattr(ctx, "get", None)
                if callable(get):
                    result = get("state")
                    state = await result if inspect.isawaitable(result) else result
                elif hasattr(ctx, "state"):
                    state = getattr(ctx, "state")
        except (AttributeError, TypeError):
            state = None

        # Use subclass-defined response extraction as final fallback when no text was streamed
        fallback_text = self.extract_response_text(state, events)
        if not any_text_emitted and fallback_text:
            yield (
                TextMessageStartEvent(type=EventType.TEXT_MESSAGE_START, message_id=message_id),
                None,
                usage_metrics,
            )
            yield (
                TextMessageContentEvent(
                    type=EventType.TEXT_MESSAGE_CONTENT,
                    message_id=message_id,
                    delta=fallback_text,
                ),
                None,
                usage_metrics,
            )
            yield (
                TextMessageEndEvent(
                    type=EventType.TEXT_MESSAGE_END,
                    message_id=message_id,
                ),
                None,
                usage_metrics,
            )
            message_id = str(uuid.uuid4())

        # Close the final agent's step. Each STEP_STARTED needs a matching
        # STEP_FINISHED; without this, the UI shows the last step as still running.
        if current_agent_name is not None:
            yield (
                StepFinishedEvent(type=EventType.STEP_FINISHED, step_name=current_agent_name),
                None,
                usage_metrics,
            )

        pipeline_interactions = self.create_pipeline_interactions_from_events(events)
        # TODO: find a way to count usage (LlamaIndex does not report it)
        yield (
            RunFinishedEvent(type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=run_id),
            pipeline_interactions,
            usage_metrics,
        )

    @classmethod
    def create_pipeline_interactions_from_events(
        cls,
        events: list[Event] | None,
    ) -> MultiTurnSample | None:
        if not events:
            return None
        # Lazy import to reduce memory overhead when ragas is not used
        from ragas import MultiTurnSample
        from ragas.integrations.llama_index import convert_to_ragas_messages
        from ragas.messages import AIMessage
        from ragas.messages import HumanMessage
        from ragas.messages import ToolMessage

        # convert_to_ragas_messages expects a list[Event]
        ragas_trace = convert_to_ragas_messages(list(events))
        ragas_messages = cast(list[HumanMessage | AIMessage | ToolMessage], ragas_trace)
        return MultiTurnSample(user_input=ragas_messages)

structured_history property

structured_history: bool

When true, prior turns are fed to the model as structured native ChatMessage history (tool calls preserved) instead of the text {chat_history} summary. Only applies when the prompt has no {chat_history} placeholder.

Defaults to True so multi-turn history is replayed by default; pass structured_history=False to opt out.

build_workflow abstractmethod async

build_workflow() -> Any

Return an AgentWorkflow instance ready to run.

Source code in datarobot_genai/llama_index/agent.py

@abc.abstractmethod
async def build_workflow(self) -> Any:
    """Return an AgentWorkflow instance ready to run."""
    raise NotImplementedError

extract_response_text abstractmethod

extract_response_text(result_state: Any, events: list[Any]) -> str

Extract final response text from workflow state and/or events.

Source code in datarobot_genai/llama_index/agent.py

@abc.abstractmethod
def extract_response_text(self, result_state: Any, events: list[Any]) -> str:
    """Extract final response text from workflow state and/or events."""
    raise NotImplementedError

invoke async

invoke(run_agent_input: RunAgentInput) -> InvokeReturn

Run the LlamaIndex workflow with the provided completion parameters.

Source code in datarobot_genai/llama_index/agent.py

async def invoke(self, run_agent_input: RunAgentInput) -> InvokeReturn:
    """Run the LlamaIndex workflow with the provided completion parameters."""
    user_prompt_content = extract_user_prompt_content(run_agent_input)
    input_message = str(user_prompt_content)

    # Prior turns reach the model as a text {chat_history} summary when the
    # prompt uses the placeholder, or as structured native ChatMessage history
    # (tool calls preserved) when structured_history is opt-in.
    structured_chat_history = None
    if "{chat_history}" in input_message:
        history_summary = self.build_history_summary(run_agent_input)
        formatted_history = (
            f"\n\nPrior conversation:\n{history_summary}" if history_summary else ""
        )
        input_message = input_message.replace("{chat_history}", formatted_history)
    elif self.structured_history:
        structured_chat_history = (
            ag_ui_history_to_chat_messages(self.history_messages(run_agent_input)) or None
        )

    input_message = prepend_streaming_memory_to_prompt(input_message, run_agent_input)

    logger.info(f"Running agent with user prompt: {input_message}")

    thread_id = run_agent_input.thread_id
    run_id = run_agent_input.run_id
    usage_metrics: UsageMetrics = default_usage_metrics()

    # Partial AG-UI: lifecycle + text + tool calls + steps
    yield (
        RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=run_id),
        None,
        usage_metrics,
    )

    # Subclasses may implement build_workflow as async or sync; support both.
    built: Any = self.build_workflow()
    workflow = await built if inspect.isawaitable(built) else built
    run_kwargs: dict[str, Any] = {"user_msg": input_message}
    if structured_chat_history:
        run_kwargs["chat_history"] = structured_chat_history
    handler = workflow.run(**run_kwargs)

    events: list[Any] = []
    current_agent_name: str | None = None
    message_id = str(uuid.uuid4())
    text_started = False
    any_text_emitted = False
    # Id of the most recently closed assistant text bubble. A following tool call
    # names it as its parent_message_id so the call attaches to the message that
    # introduced it (AG-UI grouping) instead of an orphan id; reset on a tool
    # result so a later text-less step can't reuse a stale bubble id.
    last_text_message_id: str | None = None

    async for event in handler.stream_events():
        events.append(event)

        # Detect agent transition first so the first delta of a new agent
        # is attributed to the new agent's text bubble, not the previous one.
        new_agent = getattr(event, "current_agent_name", None)
        if new_agent is not None and new_agent != current_agent_name:
            # Close any open text message from the outgoing agent so each
            # step's text commits as its own AG-UI message.
            if text_started:
                yield (
                    TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=message_id),
                    None,
                    usage_metrics,
                )
                text_started = False
            if current_agent_name is not None:
                yield (
                    StepFinishedEvent(
                        type=EventType.STEP_FINISHED, step_name=current_agent_name
                    ),
                    None,
                    usage_metrics,
                )
            yield (
                StepStartedEvent(type=EventType.STEP_STARTED, step_name=new_agent),
                None,
                usage_metrics,
            )
            current_agent_name = new_agent
            # Fresh id for the new step's first text bubble.
            message_id = str(uuid.uuid4())
            logger.info(f"Agent: {current_agent_name}")

        # Reasoning models surface incremental thinking on AgentStream.thinking_delta
        # when the LLM integration populates it; the stock litellm wrapper does not, so
        # fall back to the raw LiteLLM chunk (event.raw). Emit as a self-contained AG-UI
        # reasoning chunk before any text for this event.
        thinking_delta = getattr(event, "thinking_delta", None) or _thinking_delta_from_raw(
            getattr(event, "raw", None)
        )
        if thinking_delta:
            yield (
                ReasoningMessageChunkEvent(
                    type=EventType.REASONING_MESSAGE_CHUNK,
                    # Its own message id, distinct from the assistant text so a frontend
                    # grouping by id renders reasoning as its own block instead of folding
                    # it into the text bubble. Derived (uuid5) from the text id: a valid
                    # UUID, stable across this message's chunks (so they group), no state.
                    message_id=str(uuid.uuid5(uuid.NAMESPACE_OID, f"{message_id}-reasoning")),
                    delta=thinking_delta,
                ),
                None,
                usage_metrics,
            )

        # Best-effort extraction of incremental text from LlamaIndex events
        delta: str | None = None
        try:
            if hasattr(event, "delta") and isinstance(getattr(event, "delta"), str):
                delta = getattr(event, "delta")
            # Some event types may carry incremental text under "text" or similar
            elif hasattr(event, "text") and isinstance(getattr(event, "text"), str):
                delta = getattr(event, "text")
        except Exception:
            # Ignore malformed events and continue
            delta = None

        if delta:
            if not text_started:
                yield (
                    TextMessageStartEvent(
                        type=EventType.TEXT_MESSAGE_START, message_id=message_id
                    ),
                    None,
                    usage_metrics,
                )
                text_started = True
                any_text_emitted = True
            yield (
                TextMessageContentEvent(
                    type=EventType.TEXT_MESSAGE_CONTENT, message_id=message_id, delta=delta
                ),
                None,
                usage_metrics,
            )

        event_type = type(event).__name__
        if event_type == "AgentInput" and hasattr(event, "input"):
            logger.info(f"Input: {getattr(event, 'input')}")
        elif event_type == "AgentOutput":
            # When no deltas streamed for this turn, attach the full response
            # content to the current step's open text message. Step transition
            # or end-of-run closes it.
            resp = getattr(event, "response", None)
            if resp is not None and hasattr(resp, "content") and getattr(resp, "content"):
                content_str = str(getattr(resp, "content"))
                logger.info(f"Output: {content_str}")
                if not text_started:
                    yield (
                        TextMessageStartEvent(
                            type=EventType.TEXT_MESSAGE_START, message_id=message_id
                        ),
                        None,
                        usage_metrics,
                    )
                    text_started = True
                    any_text_emitted = True
                    yield (
                        TextMessageContentEvent(
                            type=EventType.TEXT_MESSAGE_CONTENT,
                            message_id=message_id,
                            delta=content_str,
                        ),
                        None,
                        usage_metrics,
                    )
            # Planned tool calls
            tcalls = getattr(event, "tool_calls", None)
            if isinstance(tcalls, list) and tcalls:
                names = []
                for c in tcalls:
                    try:
                        nm = getattr(c, "tool_name", None) or (
                            c.get("tool_name") if isinstance(c, dict) else None
                        )
                        if nm:
                            names.append(str(nm))
                    except Exception:
                        pass
                if names:
                    logger.info(f"Planning to use tools: {names}")
        elif event_type == "ToolCallResult":
            # The tool-calling step is complete; clear the captured bubble id so a
            # later text-less step doesn't attach its tool calls to a stale bubble.
            last_text_message_id = None
            tname = getattr(event, "tool_name", None)
            tid = getattr(event, "tool_id", None)
            tkwargs = getattr(event, "tool_kwargs", None)
            tout = getattr(event, "tool_output", None)
            logger.info(f"Tool Result: {tname}")
            logger.debug(f"Arguments: {tkwargs}")
            logger.debug(f"Output: {tout}")
            yield (
                ToolCallEndEvent(
                    type=EventType.TOOL_CALL_END,
                    tool_call_id=tid,
                ),
                None,
                usage_metrics,
            )
            yield (
                ToolCallResultEvent(
                    type=EventType.TOOL_CALL_RESULT,
                    message_id=tid,
                    tool_call_id=tid,
                    content=json.dumps(tout, default=str),
                    role="tool",
                ),
                None,
                usage_metrics,
            )
        elif event_type == "ToolCall":
            # Close any open text message so the tool widget renders below
            # the preceding text and the next text run starts a fresh bubble.
            if text_started:
                yield (
                    TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=message_id),
                    None,
                    usage_metrics,
                )
                text_started = False
                # Capture the closed bubble's id (before minting a fresh one) so the
                # tool call below can name it as its parent_message_id.
                last_text_message_id = message_id
                message_id = str(uuid.uuid4())
            tname = getattr(event, "tool_name", None)
            tkwargs = getattr(event, "tool_kwargs", None)
            tid = getattr(event, "tool_id", None)
            logger.info(f"Calling Tool: {tname}")
            logger.debug(f"With arguments: {tkwargs}")
            yield (
                ToolCallStartEvent(
                    type=EventType.TOOL_CALL_START,
                    tool_call_id=tid,
                    tool_call_name=tname,
                    parent_message_id=last_text_message_id or "",
                ),
                None,
                usage_metrics,
            )
            yield (
                ToolCallArgsEvent(
                    type=EventType.TOOL_CALL_ARGS,
                    tool_call_id=tid,
                    delta=json.dumps(tkwargs, default=str),
                ),
                None,
                usage_metrics,
            )

    # Close any text message still open from the final agent before the
    # step is finished so the message commits inside its step boundaries.
    if text_started:
        yield (
            TextMessageEndEvent(
                type=EventType.TEXT_MESSAGE_END,
                message_id=message_id,
            ),
            None,
            usage_metrics,
        )
        text_started = False
        message_id = str(uuid.uuid4())

    # After streaming completes, build final interactions and finish chunk
    # Extract state from workflow context (supports sync/async get or attribute)
    state = None
    ctx = getattr(handler, "ctx", None)
    try:
        if ctx is not None:
            get = getattr(ctx, "get", None)
            if callable(get):
                result = get("state")
                state = await result if inspect.isawaitable(result) else result
            elif hasattr(ctx, "state"):
                state = getattr(ctx, "state")
    except (AttributeError, TypeError):
        state = None

    # Use subclass-defined response extraction as final fallback when no text was streamed
    fallback_text = self.extract_response_text(state, events)
    if not any_text_emitted and fallback_text:
        yield (
            TextMessageStartEvent(type=EventType.TEXT_MESSAGE_START, message_id=message_id),
            None,
            usage_metrics,
        )
        yield (
            TextMessageContentEvent(
                type=EventType.TEXT_MESSAGE_CONTENT,
                message_id=message_id,
                delta=fallback_text,
            ),
            None,
            usage_metrics,
        )
        yield (
            TextMessageEndEvent(
                type=EventType.TEXT_MESSAGE_END,
                message_id=message_id,
            ),
            None,
            usage_metrics,
        )
        message_id = str(uuid.uuid4())

    # Close the final agent's step. Each STEP_STARTED needs a matching
    # STEP_FINISHED; without this, the UI shows the last step as still running.
    if current_agent_name is not None:
        yield (
            StepFinishedEvent(type=EventType.STEP_FINISHED, step_name=current_agent_name),
            None,
            usage_metrics,
        )

    pipeline_interactions = self.create_pipeline_interactions_from_events(events)
    # TODO: find a way to count usage (LlamaIndex does not report it)
    yield (
        RunFinishedEvent(type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=run_id),
        pipeline_interactions,
        usage_metrics,
    )

datarobot_agent_class_from_llamaindex

datarobot_agent_class_from_llamaindex(workflow: AgentWorkflow, agents: list[BaseWorkflowAgent], extract_response_text: Callable[[Any, list[Any]], str]) -> type[LlamaIndexAgent]

Create a LlamaIndex agent class from a pre-built workflow and agents.

This is a convenience helper that dynamically builds a concrete :class:LlamaIndexAgent subclass so that callers can define an agent entirely from an existing :class:AgentWorkflow and its constituent agents without writing a class by hand.

When the returned class is instantiated, calling set_llm or set_tools propagates the LLM / tools to every agent in agents while preserving each agent's originally configured tools.

Parameters

workflow : AgentWorkflow A fully configured LlamaIndex :class:AgentWorkflow instance that orchestrates the provided agents. agents : list[BaseWorkflowAgent] The list of LlamaIndex workflow agents participating in the workflow. Their llm and tools attributes are updated at runtime when the DataRobot platform injects the LLM and MCP tools. extract_response_text : Callable[[Any, list[Any]], str] A callback that extracts the final human-readable response from the workflow result state and the list of streamed events. Receives (result_state, events) and must return a str.

Returns

type[LlamaIndexAgent] A new :class:LlamaIndexAgent subclass wired to the provided workflow, agents, and response extractor.

Source code in datarobot_genai/llama_index/agent.py

def datarobot_agent_class_from_llamaindex(
    workflow: AgentWorkflow,
    agents: list[BaseWorkflowAgent],
    extract_response_text: Callable[[Any, list[Any]], str],
) -> type[LlamaIndexAgent]:
    """Create a LlamaIndex agent class from a pre-built workflow and agents.

    This is a convenience helper that dynamically builds a concrete
    :class:`LlamaIndexAgent` subclass so that callers can define an agent
    entirely from an existing :class:`AgentWorkflow` and its constituent
    agents without writing a class by hand.

    When the returned class is instantiated, calling ``set_llm`` or
    ``set_tools`` propagates the LLM / tools to every agent in *agents*
    while preserving each agent's originally configured tools.

    Parameters
    ----------
    workflow : AgentWorkflow
        A fully configured LlamaIndex :class:`AgentWorkflow` instance that
        orchestrates the provided agents.
    agents : list[BaseWorkflowAgent]
        The list of LlamaIndex workflow agents participating in the workflow.
        Their ``llm`` and ``tools`` attributes are updated at runtime when the
        DataRobot platform injects the LLM and MCP tools.
    extract_response_text : Callable[[Any, list[Any]], str]
        A callback that extracts the final human-readable response from the
        workflow result state and the list of streamed events.  Receives
        ``(result_state, events)`` and must return a ``str``.

    Returns
    -------
    type[LlamaIndexAgent]
        A new :class:`LlamaIndexAgent` subclass wired to the provided
        workflow, agents, and response extractor.
    """
    original_agent_tools = {agent.name: agent.tools for agent in agents}

    class DataRobotLlamaIndexAgent(LlamaIndexAgent):
        def __init__(
            self,
            api_key: str | None = None,
            api_base: str | None = None,
            llm: Any | None = None,
            tools: list[BaseTool] | None = None,
            verbose: bool = True,
            timeout: int = 90,
            forwarded_headers: dict[str, str] | None = None,
            max_history_messages: int | None = None,
            model: str | None = None,
            structured_history: bool = True,
            allow_parallel_tool_calls: bool = True,
        ) -> None:
            super().__init__(
                api_key=api_key,
                api_base=api_base,
                llm=llm,
                tools=tools,
                verbose=verbose,
                timeout=timeout,
                forwarded_headers=forwarded_headers,
                max_history_messages=max_history_messages,
                model=model,
                structured_history=structured_history,
                allow_parallel_tool_calls=allow_parallel_tool_calls,
            )
            for agent in agents:
                if isinstance(agent, FunctionAgent):
                    agent.allow_parallel_tool_calls = allow_parallel_tool_calls

        def set_llm(self, llm: Any) -> None:
            super().set_llm(llm)
            for agent in agents:
                agent.llm = llm

        def set_tools(self, tools: list[BaseTool]) -> None:
            super().set_tools(tools)
            for agent in agents:
                agent.tools = original_agent_tools[agent.name] + tools

        def set_allow_parallel_tool_calls(self, allow_parallel_tool_calls: bool) -> None:
            self.allow_parallel_tool_calls = allow_parallel_tool_calls
            for agent in agents:
                if isinstance(agent, FunctionAgent):
                    agent.allow_parallel_tool_calls = allow_parallel_tool_calls

        async def build_workflow(self) -> AgentWorkflow:
            return workflow

        def extract_response_text(self, result_state: Any, events: list[Any]) -> str:
            return extract_response_text(result_state, events)

    return DataRobotLlamaIndexAgent