cadSet/CADDesigner-Code/tools/subagent_utils.py

from __future__ import annotations

from typing import Any, Awaitable, Callable, Optional

from SimpleLLMFunc.hooks import (
    CustomEvent,
    ReactEndEvent,
    ToolCallEndEvent,
    ToolCallErrorEvent,
    ToolCallStartEvent,
    is_event_yield,
    is_response_yield,
)
from react_stream import extract_output_text, extract_response_reasoning


async def emit_subagent_progress(
    event_emitter: Any,
    event_name: str,
    data: dict[str, Any],
) -> None:
    if event_emitter is None:
        return

    emit = getattr(event_emitter, "emit", None)
    if not callable(emit):
        return

    awaitable_result: Any = emit(event_name, data)
    await awaitable_result


async def run_subagent_with_events(
    *,
    specialist_callable: Callable[..., Any],
    specialist_kwargs: dict[str, Any],
    subagent_label: str,
    event_emitter: Any,
    status_payload: Optional[dict[str, Any]] = None,
    response_transform: Optional[Callable[[str], str]] = None,
    captured_tool_results: Optional[list[tuple[str, str]]] = None,
) -> str:
    """Run a streaming llm_function specialist and bridge its events outward."""

    report = ""
    start_payload = {"phase": "started", "message": f"{subagent_label} started."}
    if status_payload:
        start_payload.update(status_payload)
    start_payload["subagent_label"] = subagent_label

    await emit_subagent_progress(event_emitter, "subagent_status", start_payload)

    async for output in specialist_callable(**specialist_kwargs):
        if is_response_yield(output):
            response_text = extract_output_text(output, f"subagent_{subagent_label}")
            if response_text:
                report += response_text
                await emit_subagent_progress(
                    event_emitter,
                    "subagent_response",
                    {
                        "subagent_label": subagent_label,
                        "delta_text": response_text,
                    },
                )

            reasoning_text = extract_response_reasoning(output.response)
            if reasoning_text:
                await emit_subagent_progress(
                    event_emitter,
                    "subagent_reasoning",
                    {
                        "subagent_label": subagent_label,
                        "delta_reasoning": reasoning_text,
                    },
                )
            continue

        if not is_event_yield(output):
            continue

        event = output.event

        if isinstance(event, ToolCallStartEvent):
            await emit_subagent_progress(
                event_emitter,
                "subagent_tool_start",
                {
                    "subagent_label": subagent_label,
                    "nested_tool_name": event.tool_name,
                    "nested_tool_call_id": event.tool_call_id,
                    "arguments": event.arguments,
                },
            )
            continue

        if isinstance(event, ToolCallEndEvent):
            if captured_tool_results is not None and event.result is not None:
                captured_tool_results.append((event.tool_name, str(event.result)))
            await emit_subagent_progress(
                event_emitter,
                "subagent_tool_end",
                {
                    "subagent_label": subagent_label,
                    "nested_tool_name": event.tool_name,
                    "nested_tool_call_id": event.tool_call_id,
                    "arguments": event.arguments,
                    "result": event.result,
                    "execution_time": event.execution_time,
                    "success": event.success,
                },
            )
            continue

        if isinstance(event, ToolCallErrorEvent):
            await emit_subagent_progress(
                event_emitter,
                "subagent_tool_error",
                {
                    "subagent_label": subagent_label,
                    "nested_tool_name": event.tool_name,
                    "nested_tool_call_id": event.tool_call_id,
                    "arguments": event.arguments,
                    "error_message": event.error_message,
                    "execution_time": event.execution_time,
                },
            )
            continue

        if isinstance(event, CustomEvent):
            if event.event_name.startswith("subagent_") and isinstance(
                event.data, dict
            ):
                forwarded_data = dict(event.data)
                forwarded_data.setdefault("subagent_label", subagent_label)
                forwarded_data.setdefault("source_tool_name", event.tool_name)
                forwarded_data.setdefault("source_tool_call_id", event.tool_call_id)
                await emit_subagent_progress(
                    event_emitter,
                    event.event_name,
                    forwarded_data,
                )
                continue

            await emit_subagent_progress(
                event_emitter,
                "subagent_custom",
                {
                    "subagent_label": subagent_label,
                    "nested_tool_name": event.tool_name,
                    "nested_tool_call_id": event.tool_call_id,
                    "custom_event_name": event.event_name,
                    "data": event.data,
                },
            )
            continue

        if isinstance(event, ReactEndEvent):
            if not report and isinstance(event.final_response, str):
                report = event.final_response

    finish_payload = {"phase": "finished", "message": f"{subagent_label} finished."}
    if status_payload:
        finish_payload.update(status_payload)
    finish_payload["subagent_label"] = subagent_label

    await emit_subagent_progress(event_emitter, "subagent_status", finish_payload)

    final_report = response_transform(report) if response_transform else report
    return final_report


__all__ = ["emit_subagent_progress", "run_subagent_with_events"]