leocrm/app/ai/agent_stream.py

"""SSE streaming for the ReAct agent loop.

Wraps ``run_react_loop`` from ``app.ai.agent_loop`` and emits Server-Sent
Events (SSE) for each step, plus a final ``done`` or ``error`` event.

Events emitted:
- ``event: step``   — JSON {step_number, thought, action, action_input, observation, cost_usd}
- ``event: status`` — JSON {status: "running", step: N}
- ``event: done``   — JSON {status, total_cost, steps_taken, final_content}
- ``event: error``  — JSON {error, trace_id}

Trace modes:
- ``standard`` — step events include action + result only (no thought)
- ``extended`` — step events also include the thought/reasoning
"""

from __future__ import annotations

import asyncio
import json
import logging
import uuid
from typing import TYPE_CHECKING, Any, AsyncGenerator

from app.ai.agent_loop import ReActStep, run_react_loop

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)


# ──────────────────────────────────────────────────────────────────────────
# SSE helpers
# ──────────────────────────────────────────────────────────────────────────


def _sse(event: str, data: dict[str, Any]) -> str:
    """Format a single SSE event as ``event: <name>\ndata: <json>\n\n``."""
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"


def _step_event(step: ReActStep, trace_mode: str) -> str:
    """Build the SSE ``step`` event for a ReAct step."""
    data: dict[str, Any] = {
        "step_number": step.step_number,
        "action": step.action,
        "action_input": step.action_input,
        "observation": step.observation,
        "cost_usd": step.cost_usd,
    }
    if trace_mode == "extended":
        data["thought"] = step.thought
    return _sse("step", data)


# ──────────────────────────────────────────────────────────────────────────
# Streaming loop
# ──────────────────────────────────────────────────────────────────────────


async def stream_react_loop(
    agent_definition: Any,
    user_message: str,
    tools: list[dict],
    tool_registry: Any,
    db: AsyncSession | None,
    tenant_id: uuid.UUID,
    user_id: uuid.UUID,
    agent_run_id: uuid.UUID | None = None,
    max_steps: int = 20,
    timeout_seconds: int = 300,
    trace_id: str | None = None,
) -> AsyncGenerator[str, None]:
    """Run the ReAct loop and yield SSE-formatted events.

    Args:
        agent_definition: AgentDefinition with llm_model, system_prompt, etc.
        user_message: The user's message to the agent.
        tools: OpenAI-format tool schemas for function calling.
        tool_registry: ToolRegistry instance for tool execution.
        db: Async DB session.
        tenant_id: Tenant ID for multi-tenancy.
        user_id: User ID for permission context.
        agent_run_id: Optional AgentRun ID for step persistence.
        max_steps: Maximum loop iterations (default 20).
        timeout_seconds: Overall timeout (default 300).
        trace_id: Optional trace ID for correlation.

    Yields:
        SSE-formatted event strings.
    """
    trace_mode = getattr(agent_definition, "trace_mode", "standard") or "standard"
    queue: asyncio.Queue[str | None] = asyncio.Queue()

    async def on_step(step: ReActStep) -> None:
        """Push step + status events into the queue."""
        await queue.put(_step_event(step, trace_mode))
        await queue.put(
            _sse("status", {"status": "running", "step": step.step_number})
        )

    async def _producer() -> None:
        """Run the loop and push the final done/error event."""
        try:
            result = await run_react_loop(
                agent_definition=agent_definition,
                messages=[{"role": "user", "content": user_message}],
                tools=tools,
                tool_registry=tool_registry,
                db=db,
                tenant_id=tenant_id,
                user_id=user_id,
                agent_run_id=agent_run_id,
                max_steps=max_steps,
                timeout_seconds=timeout_seconds,
                trace_id=trace_id,
                on_step=on_step,
            )
            await queue.put(
                _sse(
                    "done",
                    {
                        "status": result.status,
                        "total_cost": result.total_cost_usd,
                        "steps_taken": result.steps_taken,
                        "final_content": result.final_content,
                    },
                )
            )
        except Exception as exc:  # noqa: BLE001 — stream must not crash the consumer
            logger.exception("ReAct streaming loop failed")
            await queue.put(
                _sse("error", {"error": str(exc), "trace_id": trace_id})
            )
        finally:
            await queue.put(None)  # sentinel

    producer_task = asyncio.create_task(_producer())
    try:
        while True:
            event = await queue.get()
            if event is None:
                break
            yield event
    finally:
        if not producer_task.done():
            producer_task.cancel()
            try:
                await producer_task
            except asyncio.CancelledError:
                pass


__all__ = ["stream_react_loop"]