class DRAgentNestedReasoningStepAdaptor(StepAdaptor):
"""
Convert native agent steps to DRAgent AG-UI events.
This adaptor interprets LLM events from the root level function as TEXT, and
downstream functions as REASONING.
"""
def __init__(self, config: StepAdaptorConfig) -> None:
super().__init__(config)
self.function_level = 0
self.seen_llm_new_token = False
def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig) -> bool: # noqa: PLR0911
"""Returns True if this intermediate step should be included (based on the config.mode).""" # noqa: D401
if config.mode == StepAdaptorMode.OFF:
return False
if config.mode == StepAdaptorMode.DEFAULT:
# Process all steps
return True
if config.mode == StepAdaptorMode.CUSTOM:
# pass only what the user explicitly listed
return step.event_type in config.custom_event_types
return False
def process(self, step: IntermediateStep) -> ResponseSerializable | None:
result = super().process(step)
if not self._step_matches_filter(step, self.config):
return None
try:
payload = step.payload
ancestry = step.function_ancestry
if step.event_category == IntermediateStepCategory.WORKFLOW:
result = self._handle_workflow(payload, ancestry)
# FUNCTION fall-through leaks the run as a CustomEvent.
elif result is None and step.event_category != IntermediateStepCategory.FUNCTION:
result = self._handle_custom(payload, ancestry)
except Exception as e:
logger.exception("Error processing intermediate step: %s", e)
return None
if result is not None:
result.usage_metrics = self._get_usage_metrics(step.usage_info)
return result
return result
def _get_usage_metrics(self, usage_info: UsageInfo | None) -> dict[str, int]:
if usage_info is None:
return {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0,
}
return usage_info.token_usage.model_dump()
@staticmethod
def _extract_payload_text(payload: IntermediateStepPayload) -> str:
"""Extract text from an LLM_END payload across different frameworks.
LangChain payloads expose ``.text`` (a ``TextAccessor``), while
LlamaIndex ``ChatResponse`` objects use ``.message.content``
Always returns a plain ``str``
"""
data = payload.data.payload
if hasattr(data, "text"):
return str(data.text)
if hasattr(data, "message"):
return str(getattr(data.message, "content", data.message))
return str(data)
def _unknown_step_type(self, payload: IntermediateStepPayload) -> Exception:
return ValueError(
f"Unsupported intermediate step type: {payload.event_type}, payload: {payload}"
)
@staticmethod
def _primary_content_events(content: str | list[Any] | None, message_id: str) -> list[Event]:
"""Map a root-level LLM chunk/text to AG-UI events.
Reasoning models stream list-form content
(``[{"type": "thinking", ...}, {"type": "text", ...}]``); passing that
verbatim into a string-typed ``delta`` is what crashes the adaptor. We
normalize it: text becomes ``TextMessageContentEvent`` and thinking
becomes a self-contained ``ReasoningMessageChunkEvent`` (the shape the
CLI/web renderers already understand).
"""
events: list[Event] = []
for kind, delta in iter_content_blocks(content):
if kind == "thinking":
events.append(ReasoningMessageChunkEvent(message_id=message_id, delta=delta))
else:
events.append(TextMessageContentEvent(message_id=message_id, delta=delta))
return events
@staticmethod
def _reasoning_content_events(content: str | list[Any] | None, message_id: str) -> list[Event]:
"""Map a nested (sub-agent) LLM chunk/text to reasoning events.
Downstream LLM output renders entirely as reasoning, so both text and
thinking blocks become ``ReasoningMessageContentEvent``s inside the
surrounding reasoning envelope. Normalizing also flattens list-form
content so it is never passed verbatim into a string ``delta``.
"""
return [
ReasoningMessageContentEvent(message_id=message_id, delta=delta)
for _kind, delta in iter_content_blocks(content)
]
def _handle_llm(
self, payload: IntermediateStepPayload, ancestry: InvocationNode
) -> ResponseSerializable | None:
# Find the start in the history with matching run_id
if self.function_level == 1:
events = self._handle_llm_primary_function(payload)
else:
events = self._handle_llm_nested_function(payload)
response = DRAgentEventResponse(
events=events,
# Only for llm events we actually know the model name
# And its is passed as name
model=payload.name,
)
return response
def _handle_llm_primary_function(self, payload: IntermediateStepPayload) -> list[Event]:
events = []
if payload.event_type == IntermediateStepType.LLM_START:
events.append(TextMessageStartEvent(message_id=payload.UUID))
self.seen_llm_new_token = False
# Text might be sent in both LLM_END and LLM_NEW_TOKEN steps
# we need to send content only once
elif payload.event_type == IntermediateStepType.LLM_END:
if not self.seen_llm_new_token:
events.extend(
self._primary_content_events(self._extract_payload_text(payload), payload.UUID)
)
events.append(TextMessageEndEvent(message_id=payload.UUID))
elif payload.event_type == IntermediateStepType.LLM_NEW_TOKEN:
self.seen_llm_new_token = True
events.extend(self._primary_content_events(payload.data.chunk, payload.UUID))
else:
raise self._unknown_step_type(payload)
return events
def _handle_llm_nested_function(self, payload: IntermediateStepPayload) -> list[Event]:
events: list[Event] = []
if payload.event_type == IntermediateStepType.LLM_START:
events.append(ReasoningStartEvent(message_id=payload.UUID))
events.append(ReasoningMessageStartEvent(message_id=payload.UUID, role="reasoning"))
self.seen_llm_new_token = False
elif payload.event_type == IntermediateStepType.LLM_END:
if not self.seen_llm_new_token:
events.extend(
self._reasoning_content_events(
self._extract_payload_text(payload), payload.UUID
)
)
events.append(ReasoningMessageEndEvent(message_id=payload.UUID))
events.append(ReasoningEndEvent(message_id=payload.UUID))
elif payload.event_type == IntermediateStepType.LLM_NEW_TOKEN:
self.seen_llm_new_token = True
events.extend(self._reasoning_content_events(payload.data.chunk, payload.UUID))
else:
raise self._unknown_step_type(payload)
return events
def _handle_workflow(
self, payload: IntermediateStepPayload, ancestry: InvocationNode
) -> ResponseSerializable | None:
events = []
# run id and thread id are set on the API level, should not be set here
if payload.event_type == IntermediateStepType.WORKFLOW_START:
events.append(RunStartedEvent(run_id="", thread_id=""))
events.append(StepStartedEvent(step_name=payload.name))
elif payload.event_type == IntermediateStepType.WORKFLOW_END:
events.append(StepFinishedEvent(step_name=payload.name))
events.append(RunFinishedEvent(run_id="", thread_id=""))
else:
raise self._unknown_step_type(payload)
response = DRAgentEventResponse(events=events)
return response
@staticmethod
def _serialize_tool_args(payload: IntermediateStepPayload) -> str:
"""Extract tool call arguments as a JSON string.
Tries metadata.tool_inputs first, then data.input. Returns "{}"
when no usable arguments are found. Non-serializable values
(e.g. CrewStructuredTool leaked via nvidia-nat-crewai) are logged
and skipped.
"""
tool_inputs = getattr(payload.metadata, "tool_inputs", None)
if isinstance(tool_inputs, dict):
raw = tool_inputs
else:
raw = payload.data.input
if raw is None:
return "{}"
if isinstance(raw, str):
try:
json.loads(raw)
return raw
except (json.JSONDecodeError, ValueError):
logger.warning(
"Tool args not valid JSON for %s, skipping",
payload.name,
)
return "{}"
try:
return json.dumps(raw)
except (TypeError, ValueError):
logger.warning(
"Tool args not serializable for %s, skipping",
payload.name,
)
return "{}"
def _handle_tool(
self, payload: IntermediateStepPayload, ancestry: InvocationNode
) -> ResponseSerializable | None:
events = []
# run id and thread id are set on the API level, should not be set here
if payload.event_type == IntermediateStepType.TOOL_START:
events.append(
ToolCallStartEvent(tool_call_name=payload.name, tool_call_id=payload.UUID)
)
args_delta = self._serialize_tool_args(payload)
events.append(ToolCallArgsEvent(tool_call_id=payload.UUID, delta=args_delta))
elif payload.event_type == IntermediateStepType.TOOL_END:
events.append(ToolCallEndEvent(tool_call_id=payload.UUID))
tool_outputs = GlobalTypeConverter.get().convert(payload.metadata.tool_outputs, str)
events.append(
ToolCallResultEvent(
message_id=payload.UUID,
tool_call_id=payload.UUID,
content=tool_outputs,
role="tool",
)
)
else:
raise self._unknown_step_type(payload)
response = DRAgentEventResponse(events=events)
return response
def _handle_function(
self, payload: IntermediateStepPayload, ancestry: InvocationNode
) -> ResponseSerializable | None:
# Just track the function level so we can handle nested functions correctly
if payload.event_type == IntermediateStepType.FUNCTION_START:
self.function_level += 1
bind_tool_call(payload.name, payload.UUID)
elif payload.event_type == IntermediateStepType.FUNCTION_END:
self.function_level -= 1
# Sub-agents dispatched as tools fire only FUNCTION_*, no TOOL_*.
tool_call_id = pop_tool_call(payload.UUID)
if tool_call_id is not None:
output = ""
if payload.data is not None and getattr(payload.data, "output", None) is not None:
output = json.dumps(payload.data.output, default=str)
# End before Result; reversed order strands the UI in args streaming.
end_events = [
ToolCallEndEvent(tool_call_id=tool_call_id),
ToolCallResultEvent(
message_id=tool_call_id,
tool_call_id=tool_call_id,
content=output,
role="tool",
),
]
# If the stream converter has already finished sending all
# ToolCallArgsEvent chunks, emit immediately. Otherwise
# defer until the converter calls mark_args_done.
if is_args_done(tool_call_id):
return DRAgentEventResponse(events=end_events)
defer_tool_end(tool_call_id, end_events)
else:
raise self._unknown_step_type(payload)
return None
def _handle_custom(
self, payload: IntermediateStepPayload, ancestry: InvocationNode
) -> ResponseSerializable | None:
event = CustomEvent(name=payload.event_type, value=payload)
response = DRAgentEventResponse(events=[event])
return response