glaip-sdk 0.6.5b3__py3-none-any.whl → 0.7.17__py3-none-any.whl

glaip_sdk/client/run_rendering.py

@@ -9,9 +9,12 @@ from __future__ import annotations
 
 import json
 import logging
-from collections.abc import Callable
+from collections.abc import AsyncIterable, Callable
 from time import monotonic
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from glaip_sdk.hitl.remote import RemoteHITLHandler
 
 import httpx
 from rich.console import Console as _Console
@@ -30,6 +33,7 @@ from glaip_sdk.utils.rendering.renderer import (
 from glaip_sdk.utils.rendering.state import TranscriptBuffer
 
 NO_AGENT_RESPONSE_FALLBACK = "No agent response received."
+_FINAL_EVENT_TYPES = {"final_response", "error", "step_limit_exceeded"}
 
 
 def _coerce_to_string(value: Any) -> str:
@@ -129,6 +133,7 @@ class AgentRunRenderingManager:
         timeout_seconds: float,
         agent_name: str | None,
         meta: dict[str, Any],
+        hitl_handler: RemoteHITLHandler | None = None,
     ) -> tuple[str, dict[str, Any], float | None, float | None]:
         """Process streaming events and accumulate response."""
         final_text = ""
@@ -152,10 +157,132 @@ class AgentRunRenderingManager:
                     final_text,
                     stats_usage,
                     meta,
+                    hitl_handler=hitl_handler,
                 )
 
                 if controller and getattr(controller, "enabled", False):
                     controller.poll(renderer)
+                parsed_event = self._parse_event(event)
+                if parsed_event and self._is_final_event(parsed_event):
+                    break
+        finally:
+            if controller and getattr(controller, "enabled", False):
+                controller.on_stream_complete()
+
+        finished_monotonic = monotonic()
+        return final_text, stats_usage, started_monotonic, finished_monotonic
+
+    async def _consume_event_stream(
+        self,
+        event_stream: AsyncIterable[dict[str, Any]],
+        renderer: RichStreamRenderer,
+        final_text: str,
+        stats_usage: dict[str, Any],
+        meta: dict[str, Any],
+        skip_final_render: bool,
+        last_rendered_content: str | None,
+        controller: Any | None,
+    ) -> tuple[str, dict[str, Any], float | None]:
+        """Consume event stream and update state.
+
+        Args:
+            event_stream: Async iterable yielding SSE-like event dicts.
+            renderer: Renderer to use for displaying events.
+            final_text: Current accumulated final text.
+            stats_usage: Usage statistics dictionary.
+            meta: Metadata dictionary.
+            skip_final_render: If True, skip rendering final_response events.
+            last_rendered_content: Last rendered content to avoid duplicates.
+            controller: Controller instance.
+
+        Returns:
+            Tuple of (final_text, stats_usage, started_monotonic).
+        """
+        started_monotonic: float | None = None
+
+        async for event in event_stream:
+            if started_monotonic is None:
+                started_monotonic = monotonic()
+
+            parsed_event = self._parse_event(event)
+            if parsed_event is None:
+                continue
+
+            final_text, stats_usage = self._handle_parsed_event(
+                parsed_event,
+                renderer,
+                final_text,
+                stats_usage,
+                meta,
+                skip_final_render=skip_final_render,
+                last_rendered_content=last_rendered_content,
+            )
+
+            content_str = self._extract_content_string(parsed_event)
+            if content_str:
+                last_rendered_content = content_str
+
+            if controller and getattr(controller, "enabled", False):
+                controller.poll(renderer)
+            if parsed_event and self._is_final_event(parsed_event):
+                break
+
+        return final_text, stats_usage, started_monotonic
+
+    async def async_process_stream_events(
+        self,
+        event_stream: AsyncIterable[dict[str, Any]],
+        renderer: RichStreamRenderer,
+        meta: dict[str, Any],
+        *,
+        skip_final_render: bool = True,
+    ) -> tuple[str, dict[str, Any], float | None, float | None]:
+        """Process streaming events from an async event source.
+
+        This method provides unified stream processing for both remote (HTTP)
+        and local (LangGraph) agent execution, ensuring consistent behavior.
+
+        Args:
+            event_stream: Async iterable yielding SSE-like event dicts.
+                Each event should have a "data" key with JSON string, or be
+                a pre-parsed dict with "content", "metadata", etc.
+            renderer: Renderer to use for displaying events.
+            meta: Metadata dictionary for renderer context.
+            skip_final_render: If True, skip rendering final_response events
+                (they are rendered separately via finalize_renderer).
+
+        Returns:
+            Tuple of (final_text, stats_usage, started_monotonic, finished_monotonic).
+        """
+        final_text = ""
+        stats_usage: dict[str, Any] = {}
+        started_monotonic: float | None = None
+        last_rendered_content: str | None = None
+
+        controller = getattr(renderer, "transcript_controller", None)
+        if controller and getattr(controller, "enabled", False):
+            controller.on_stream_start(renderer)
+
+        try:
+            final_text, stats_usage, started_monotonic = await self._consume_event_stream(
+                event_stream,
+                renderer,
+                final_text,
+                stats_usage,
+                meta,
+                skip_final_render,
+                last_rendered_content,
+                controller,
+            )
+        except Exception as e:
+            err_msg = str(e)
+            reason = getattr(getattr(e, "result", None), "reason", None)
+            if reason:
+                final_text = f"⚠️ Guardrail violation: {reason}"
+            elif "⚠️ Guardrail violation" in err_msg or "Content blocked by guardrails" in err_msg:
+                final_text = err_msg
+            else:
+                raise e
         finally:
             if controller and getattr(controller, "enabled", False):
                 controller.on_stream_complete()
@@ -163,6 +290,229 @@ class AgentRunRenderingManager:
         finished_monotonic = monotonic()
         return final_text, stats_usage, started_monotonic, finished_monotonic
 
+    def _parse_event(self, event: dict[str, Any]) -> dict[str, Any] | None:
+        """Parse an SSE event dict into a usable format.
+
+        Args:
+            event: Raw event dict, either with "data" key (SSE format) or
+                pre-parsed with "content", "metadata", etc.
+
+        Returns:
+            Parsed event dict, or None if parsing fails.
+        """
+        if "data" in event:
+            try:
+                return json.loads(event["data"])
+            except json.JSONDecodeError:
+                self._logger.debug("Non-JSON SSE fragment skipped")
+                return None
+        # Already parsed (e.g., from local runner)
+        return event if event else None
+
+    def _handle_parsed_event(
+        self,
+        ev: dict[str, Any],
+        renderer: RichStreamRenderer,
+        final_text: str,
+        stats_usage: dict[str, Any],
+        meta: dict[str, Any],
+        *,
+        skip_final_render: bool = True,
+        last_rendered_content: str | None = None,
+    ) -> tuple[str, dict[str, Any]]:
+        """Handle a parsed event and update accumulators.
+
+        Args:
+            ev: Parsed event dictionary.
+            renderer: Renderer instance.
+            final_text: Current accumulated final text.
+            stats_usage: Usage statistics dictionary.
+            meta: Metadata dictionary.
+            skip_final_render: If True, skip rendering final_response events.
+            last_rendered_content: Last rendered content to avoid duplicates.
+
+        Returns:
+            Tuple of (updated_final_text, updated_stats_usage).
+        """
+        kind = self._get_event_kind(ev)
+
+        # Dispatch to specialized handlers based on event kind
+        handler = self._get_event_handler(kind, ev)
+        if handler:
+            return handler(ev, renderer, final_text, stats_usage, meta, skip_final_render)
+
+        # Default: handle content events
+        return self._handle_content_event_async(ev, renderer, final_text, stats_usage, last_rendered_content)
+
+    def _get_event_handler(
+        self,
+        kind: str | None,
+        ev: dict[str, Any],
+    ) -> Callable[..., tuple[str, dict[str, Any]]] | None:
+        """Get the appropriate handler for an event kind.
+
+        Args:
+            kind: Event kind string.
+            ev: Event dictionary (for checking is_final flag).
+
+        Returns:
+            Handler function or None for default content handling.
+        """
+        if kind == "usage":
+            return self._handle_usage_event
+        if kind == "final_response" or ev.get("is_final"):
+            return self._handle_final_response_event
+        if kind == "run_info":
+            return self._handle_run_info_event_wrapper
+        if kind in ("artifact", "status_update"):
+            return self._handle_render_only_event
+        return None
+
+    def _handle_usage_event(
+        self,
+        ev: dict[str, Any],
+        _renderer: RichStreamRenderer,
+        final_text: str,
+        stats_usage: dict[str, Any],
+        _meta: dict[str, Any],
+        _skip_final_render: bool,
+    ) -> tuple[str, dict[str, Any]]:
+        """Handle usage events."""
+        stats_usage.update(ev.get("usage") or {})
+        return final_text, stats_usage
+
+    def _handle_final_response_event(
+        self,
+        ev: dict[str, Any],
+        renderer: RichStreamRenderer,
+        final_text: str,
+        stats_usage: dict[str, Any],
+        _meta: dict[str, Any],
+        skip_final_render: bool,
+    ) -> tuple[str, dict[str, Any]]:
+        """Handle final_response events."""
+        content = ev.get("content")
+        if content:
+            final_text = str(content)
+        if not skip_final_render:
+            renderer.on_event(ev)
+        return final_text, stats_usage
+
+    def _handle_run_info_event_wrapper(
+        self,
+        ev: dict[str, Any],
+        renderer: RichStreamRenderer,
+        final_text: str,
+        stats_usage: dict[str, Any],
+        meta: dict[str, Any],
+        _skip_final_render: bool,
+    ) -> tuple[str, dict[str, Any]]:
+        """Handle run_info events."""
+        self._handle_run_info_event(ev, meta, renderer)
+        return final_text, stats_usage
+
+    def _handle_render_only_event(
+        self,
+        ev: dict[str, Any],
+        renderer: RichStreamRenderer,
+        final_text: str,
+        stats_usage: dict[str, Any],
+        _meta: dict[str, Any],
+        _skip_final_render: bool,
+    ) -> tuple[str, dict[str, Any]]:
+        """Handle events that only need rendering (artifact, status_update)."""
+        renderer.on_event(ev)
+        return final_text, stats_usage
+
+    def _handle_content_event_async(
+        self,
+        ev: dict[str, Any],
+        renderer: RichStreamRenderer,
+        final_text: str,
+        stats_usage: dict[str, Any],
+        last_rendered_content: str | None,
+    ) -> tuple[str, dict[str, Any]]:
+        """Handle content events with deduplication."""
+        content = ev.get("content")
+        if content:
+            content_str = str(content)
+            if not content_str.startswith("Artifact received:"):
+                kind = self._get_event_kind(ev)
+                # Skip accumulating content for status updates and agent steps
+                if kind in ("agent_step", "status_update"):
+                    renderer.on_event(ev)
+                    return final_text, stats_usage
+
+                if self._is_token_event(ev):
+                    renderer.on_event(ev)
+                    final_text = f"{final_text}{content_str}"
+                else:
+                    if content_str != last_rendered_content:
+                        renderer.on_event(ev)
+                    final_text = content_str
+        else:
+            renderer.on_event(ev)
+        return final_text, stats_usage
+
+    def _get_event_kind(self, ev: dict[str, Any]) -> str | None:
+        """Extract normalized event kind from parsed event.
+
+        Args:
+            ev: Parsed event dictionary.
+
+        Returns:
+            Event kind string or None.
+        """
+        metadata = ev.get("metadata") or {}
+        kind = metadata.get("kind")
+        if kind:
+            return str(kind)
+        event_type = ev.get("event_type")
+        return str(event_type) if event_type else None
+
+    def _is_token_event(self, ev: dict[str, Any]) -> bool:
+        """Return True when the event represents token streaming output.
+
+        Args:
+            ev: Parsed event dictionary.
+
+        Returns:
+            True when the event is a token chunk, otherwise False.
+        """
+        metadata = ev.get("metadata") or {}
+        kind = metadata.get("kind")
+        return str(kind).lower() == "token"
+
+    def _is_final_event(self, ev: dict[str, Any]) -> bool:
+        """Return True when the event marks stream termination.
+
+        Args:
+            ev: Parsed event dictionary.
+
+        Returns:
+            True when the event is terminal, otherwise False.
+        """
+        if ev.get("is_final") is True or ev.get("final") is True:
+            return True
+        kind = self._get_event_kind(ev)
+        return kind in _FINAL_EVENT_TYPES
+
+    def _extract_content_string(self, event: dict[str, Any]) -> str | None:
+        """Extract textual content from a parsed event.
+
+        Args:
+            event: Parsed event dictionary.
+
+        Returns:
+            Content string or None.
+        """
+        if not event:
+            return None
+        content = event.get("content")
+        if content:
+            return str(content)
+        return None
+
     def _capture_request_id(
         self,
         stream_response: httpx.Response,
@@ -206,6 +556,7 @@ class AgentRunRenderingManager:
         final_text: str,
         stats_usage: dict[str, Any],
         meta: dict[str, Any],
+        hitl_handler: RemoteHITLHandler | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """Process a single streaming event.
 
@@ -215,6 +566,7 @@ class AgentRunRenderingManager:
             final_text: Accumulated text so far.
             stats_usage: Usage statistics dictionary.
             meta: Metadata dictionary.
+            hitl_handler: Optional HITL handler for approval callbacks.
 
         Returns:
             Tuple of (updated_final_text, updated_stats_usage).
@@ -225,6 +577,17 @@ class AgentRunRenderingManager:
             self._logger.debug("Non-JSON SSE fragment skipped")
             return final_text, stats_usage
 
+        # Handle HITL event (non-blocking via thread)
+        if hitl_handler and self._is_hitl_pending_event(ev):
+            try:
+                hitl_handler.handle_hitl_event(ev)
+            except Exception as e:
+                # Log but don't crash stream
+                self._logger.error(
+                    f"HITL handler error: {e}",
+                    exc_info=True,
+                )
+
         kind = (ev.get("metadata") or {}).get("kind")
         renderer.on_event(ev)
 
@@ -239,7 +602,9 @@ class AgentRunRenderingManager:
         if handled is not None:
             return handled
 
-        if ev.get("content"):
+        # Only accumulate content for actual content events, not status updates or agent steps
+        # Status updates (agent_step) should be rendered but not accumulated in final_text
+        if ev.get("content") and kind not in ("agent_step", "status_update"):
             final_text = self._handle_content_event(ev, final_text)
 
         return final_text, stats_usage
@@ -285,9 +650,25 @@ class AgentRunRenderingManager:
         """
         content = ev.get("content", "")
         if not content.startswith("Artifact received:"):
+            if self._is_token_event(ev):
+                return f"{final_text}{content}"
             return content
         return final_text
 
+    @staticmethod
+    def _is_hitl_pending_event(event: dict[str, Any]) -> bool:
+        """Check if event is a pending HITL approval request.
+
+        Args:
+            event: Parsed event dictionary.
+
+        Returns:
+            True if event is a pending HITL request.
+        """
+        metadata = event.get("metadata", {})
+        hitl_meta = metadata.get("hitl", {})
+        return hitl_meta.get("required") is True and hitl_meta.get("decision") == "pending"
+
     def _handle_run_info_event(
         self,
         ev: dict[str, Any],
@@ -413,3 +794,33 @@ def compute_timeout_seconds(kwargs: dict[str, Any]) -> float:
         if not specified in kwargs.
     """
     return kwargs.get("timeout", DEFAULT_AGENT_RUN_TIMEOUT)
+
+
+def finalize_render_manager(
+    manager: AgentRunRenderingManager,
+    renderer: RichStreamRenderer,
+    final_text: str,
+    stats_usage: dict[str, Any],
+    started_monotonic: float | None,
+    finished_monotonic: float | None,
+) -> str:
+    """Helper to finalize renderer via manager and return final text.
+
+    Args:
+        manager: The rendering manager instance.
+        renderer: Renderer to finalize.
+        final_text: Final text content.
+        stats_usage: Usage statistics dictionary.
+        started_monotonic: Start time (monotonic).
+        finished_monotonic: Finish time (monotonic).
+
+    Returns:
+        Final text string.
+    """
+    return manager.finalize_renderer(
+        renderer,
+        final_text,
+        stats_usage,
+        started_monotonic,
+        finished_monotonic,
+    )