glaip-sdk 0.0.5b1__py3-none-any.whl → 0.0.7__py3-none-any.whl

glaip_sdk/cli/validators.py

@@ -7,7 +7,9 @@ Authors:
     Raymond Christopher (raymond.christopher@gdplabs.id)
 """
 
+from collections.abc import Callable
 from pathlib import Path
+from typing import Any
 
 import click
 
@@ -189,7 +191,7 @@ def validate_api_key_cli(api_key: str) -> str:
         raise click.ClickException(str(e))
 
 
-def coerce_timeout_cli(value) -> int:
+def coerce_timeout_cli(value: int | float | str) -> int:
     """Coerce timeout value to integer with CLI-friendly error handling.
 
     Args:
@@ -208,7 +210,10 @@ def coerce_timeout_cli(value) -> int:
 
 
 def validate_name_uniqueness_cli(
-    _client, name: str, resource_type: str, finder_func
+    _client: Any,
+    name: str,
+    resource_type: str,
+    finder_func: Callable[..., list[Any]],
 ) -> None:
     """Validate that a resource name is unique.
 

glaip_sdk/client/agents.py

@@ -42,6 +42,9 @@ from glaip_sdk.utils.validation import validate_agent_instruction
 # API endpoints
 AGENTS_ENDPOINT = "/agents/"
 
+# SSE content type
+SSE_CONTENT_TYPE = "text/event-stream"
+
 # Set up module-level logger
 logger = logging.getLogger("glaip_sdk.agents")
 
@@ -49,7 +52,12 @@ logger = logging.getLogger("glaip_sdk.agents")
 class AgentClient(BaseClient):
     """Client for agent operations."""
 
-    def __init__(self, *, parent_client: BaseClient | None = None, **kwargs):
+    def __init__(
+        self,
+        *,
+        parent_client: BaseClient | None = None,
+        **kwargs: Any,
+    ) -> None:
         """Initialize the agent client.
 
         Args:
@@ -161,7 +169,7 @@ class AgentClient(BaseClient):
         tools: list[str | Any] | None = None,
         agents: list[str | Any] | None = None,
         timeout: int = DEFAULT_AGENT_RUN_TIMEOUT,
-        **kwargs,
+        **kwargs: Any,
     ) -> dict[str, Any]:
         """Build payload for agent creation with proper LM selection and metadata handling.
 
@@ -365,7 +373,10 @@ class AgentClient(BaseClient):
                 agent_config.pop(key, None)
 
     def _finalize_update_payload(
-        self, update_data: dict[str, Any], current_agent: "Agent", **kwargs
+        self,
+        update_data: dict[str, Any],
+        current_agent: "Agent",
+        **kwargs: Any,
     ) -> dict[str, Any]:
         """Finalize the update payload with metadata and additional kwargs."""
         # Handle metadata preservation
@@ -386,7 +397,7 @@ class AgentClient(BaseClient):
         name: str | None = None,
         instruction: str | None = None,
         model: str | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> dict[str, Any]:
         """Build payload for agent update with proper LM selection and current state preservation.
 
@@ -434,7 +445,7 @@ class AgentClient(BaseClient):
         tools: list[str | Any] | None = None,
         agents: list[str | Any] | None = None,
         timeout: int = DEFAULT_AGENT_RUN_TIMEOUT,
-        **kwargs,
+        **kwargs: Any,
     ) -> "Agent":
         """Create a new agent."""
         # Client-side validation
@@ -470,7 +481,7 @@ class AgentClient(BaseClient):
         name: str | None = None,
         instruction: str | None = None,
         model: str | None = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> "Agent":
         """Update an existing agent."""
         # First, get the current agent data
@@ -493,39 +504,67 @@ class AgentClient(BaseClient):
         """Delete an agent."""
         self._request("DELETE", f"/agents/{agent_id}")
 
-    def _prepare_payload_and_headers(
-        self, message: str, files: list[str | BinaryIO] | None, tty: bool, **kwargs
-    ):
-        """Prepare payload and headers for agent run request."""
-        multipart_data = None
-        headers = None
+    def _prepare_sync_request_data(
+        self,
+        message: str,
+        files: list[str | BinaryIO] | None,
+        tty: bool,
+        **kwargs: Any,
+    ) -> tuple[dict | None, dict | None, list | None, dict, Any | None]:
+        """Prepare request data for synchronous agent runs with renderer support.
+
+        Args:
+            message: Message to send
+            files: Optional files to include
+            tty: Whether to enable TTY mode
+            **kwargs: Additional request parameters
+
+        Returns:
+            Tuple of (payload, data_payload, files_payload, headers, multipart_data)
+        """
+        headers = {"Accept": SSE_CONTENT_TYPE}
 
         if files:
+            # Handle multipart data for file uploads
             multipart_data = prepare_multipart_data(message, files)
-            # Inject optional multipart extras expected by backend
             if "chat_history" in kwargs and kwargs["chat_history"] is not None:
                 multipart_data.data["chat_history"] = kwargs["chat_history"]
             if "pii_mapping" in kwargs and kwargs["pii_mapping"] is not None:
                 multipart_data.data["pii_mapping"] = kwargs["pii_mapping"]
 
-        # When streaming, explicitly prefer SSE
-        headers = {**(headers or {}), "Accept": "text/event-stream"}
-
-        if files:
-            payload = None
-            data_payload = multipart_data.data
-            files_payload = multipart_data.files
+            return (
+                None,
+                multipart_data.data,
+                multipart_data.files,
+                headers,
+                multipart_data,
+            )
         else:
-            payload = {"input": message, **kwargs}
+            # Simple JSON payload for text-only requests
+            payload = {"input": message, "stream": True, **kwargs}
             if tty:
                 payload["tty"] = True
-            payload["stream"] = True
-            data_payload = None
-            files_payload = None
+            return payload, None, None, headers, None
+
+    def _get_timeout_values(
+        self, timeout: float | None, **kwargs: Any
+    ) -> tuple[float, float]:
+        """Get request timeout and execution timeout values.
+
+        Args:
+            timeout: Request timeout (overrides instance timeout)
+            **kwargs: Additional parameters including execution timeout
 
-        return payload, data_payload, files_payload, headers, multipart_data
+        Returns:
+            Tuple of (request_timeout, execution_timeout)
+        """
+        request_timeout = timeout or self.timeout
+        execution_timeout = kwargs.get("timeout", DEFAULT_AGENT_RUN_TIMEOUT)
+        return request_timeout, execution_timeout
 
-    def _create_renderer(self, renderer, **kwargs):
+    def _create_renderer(
+        self, renderer: RichStreamRenderer | None, **kwargs: Any
+    ) -> RichStreamRenderer:
         """Create appropriate renderer based on configuration."""
         if isinstance(renderer, RichStreamRenderer):
             return renderer
@@ -544,7 +583,7 @@ class AgentClient(BaseClient):
         else:
             return self._create_default_renderer(verbose)
 
-    def _create_silent_renderer(self):
+    def _create_silent_renderer(self) -> RichStreamRenderer:
         """Create a silent renderer that suppresses all output."""
         silent_config = RendererConfig(
             live=False,
@@ -558,7 +597,7 @@ class AgentClient(BaseClient):
             verbose=False,
         )
 
-    def _create_minimal_renderer(self):
+    def _create_minimal_renderer(self) -> RichStreamRenderer:
         """Create a minimal renderer with basic output."""
         minimal_config = RendererConfig(
             live=False,
@@ -572,7 +611,7 @@ class AgentClient(BaseClient):
             verbose=False,
         )
 
-    def _create_verbose_renderer(self):
+    def _create_verbose_renderer(self) -> RichStreamRenderer:
         """Create a verbose renderer for detailed output."""
         verbose_config = RendererConfig(
             theme="dark",
@@ -587,7 +626,7 @@ class AgentClient(BaseClient):
             verbose=True,
         )
 
-    def _create_default_renderer(self, verbose: bool):
+    def _create_default_renderer(self, verbose: bool) -> RichStreamRenderer:
         """Create the default renderer."""
         if verbose:
             return self._create_verbose_renderer()
@@ -595,22 +634,22 @@ class AgentClient(BaseClient):
             default_config = RendererConfig(show_delegate_tool_panels=True)
             return RichStreamRenderer(console=_Console(), cfg=default_config)
 
-    def _process_stream_events(
-        self, stream_response, renderer, timeout_seconds, agent_name, kwargs
-    ):
-        """Process streaming events and accumulate response."""
-        final_text = ""
-        stats_usage = {}
-        started_monotonic = None
-        finished_monotonic = None
-        meta = {
+    def _initialize_stream_metadata(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Initialize stream metadata."""
+        return {
             "agent_name": kwargs.get("agent_name", ""),
             "model": kwargs.get("model"),
             "run_id": None,
             "input_message": "",  # Will be set from kwargs if available
         }
 
-        # Capture request id if provided
+    def _capture_request_id(
+        self,
+        stream_response: httpx.Response,
+        meta: dict[str, Any],
+        renderer: RichStreamRenderer,
+    ) -> None:
+        """Capture request ID from response headers."""
         req_id = stream_response.headers.get(
             "x-request-id"
         ) or stream_response.headers.get("x-run-id")
@@ -618,50 +657,98 @@ class AgentClient(BaseClient):
             meta["run_id"] = req_id
             renderer.on_start(meta)
 
-        for event in iter_sse_events(stream_response, timeout_seconds, agent_name):
-            try:
-                ev = json.loads(event["data"])
-            except json.JSONDecodeError:
-                logger.debug("Non-JSON SSE fragment skipped")
-                continue
+    def _should_start_timer(self, ev: dict[str, Any]) -> bool:
+        """Check if timer should be started for this event."""
+        return "content" in ev or "status" in ev or ev.get("metadata")
+
+    def _handle_content_event(self, ev: dict[str, Any], final_text: str) -> str:
+        """Handle content events."""
+        content = ev.get("content", "")
+        if not content.startswith("Artifact received:"):
+            return content
+        return final_text
+
+    def _handle_usage_event(
+        self, ev: dict[str, Any], stats_usage: dict[str, Any]
+    ) -> None:
+        """Handle usage events."""
+        stats_usage.update(ev.get("usage") or {})
+
+    def _handle_run_info_event(
+        self, ev: dict[str, Any], meta: dict[str, Any], renderer: RichStreamRenderer
+    ) -> None:
+        """Handle run info events."""
+        if ev.get("model"):
+            meta["model"] = ev["model"]
+            renderer.on_start(meta)
+        if ev.get("run_id"):
+            meta["run_id"] = ev["run_id"]
+            renderer.on_start(meta)
+
+    def _process_single_event(
+        self,
+        event: dict[str, Any],
+        renderer: RichStreamRenderer,
+        final_text: str,
+        stats_usage: dict[str, Any],
+        meta: dict[str, Any],
+    ) -> tuple[str, dict[str, Any]]:
+        """Process a single streaming event."""
+        try:
+            ev = json.loads(event["data"])
+        except json.JSONDecodeError:
+            logger.debug("Non-JSON SSE fragment skipped")
+            return final_text, stats_usage
+
+        kind = (ev.get("metadata") or {}).get("kind")
+        renderer.on_event(ev)
+
+        # Skip artifacts from content accumulation
+        if kind == "artifact":
+            return final_text, stats_usage
+
+        # Handle different event types
+        if kind == "final_response" and ev.get("content"):
+            final_text = ev.get("content", "")
+        elif ev.get("content"):
+            final_text = self._handle_content_event(ev, final_text)
+        elif kind == "usage":
+            self._handle_usage_event(ev, stats_usage)
+        elif kind == "run_info":
+            self._handle_run_info_event(ev, meta, renderer)
+
+        return final_text, stats_usage
 
+    def _process_stream_events(
+        self,
+        stream_response: httpx.Response,
+        renderer: RichStreamRenderer,
+        timeout_seconds: float,
+        agent_name: str | None,
+        kwargs: dict[str, Any],
+    ) -> tuple[str, dict[str, Any], float | None, float | None]:
+        """Process streaming events and accumulate response."""
+        final_text = ""
+        stats_usage = {}
+        started_monotonic = None
+        finished_monotonic = None
+
+        meta = self._initialize_stream_metadata(kwargs)
+        self._capture_request_id(stream_response, meta, renderer)
+
+        for event in iter_sse_events(stream_response, timeout_seconds, agent_name):
             # Start timer at first meaningful event
-            if started_monotonic is None and (
-                "content" in ev or "status" in ev or ev.get("metadata")
-            ):
-                started_monotonic = monotonic()
-
-            kind = (ev.get("metadata") or {}).get("kind")
-            renderer.on_event(ev)
-
-            # Skip artifacts from content accumulation
-            if kind == "artifact":
-                continue
-
-            # Accumulate assistant content
-            if ev.get("content"):
-                if not ev["content"].startswith("Artifact received:"):
-                    final_text = ev["content"]
-                continue
-
-            # Handle final response
-            if kind == "final_response" and ev.get("content"):
-                final_text = ev["content"]
-                continue
-
-            # Handle usage stats
-            if kind == "usage":
-                stats_usage.update(ev.get("usage") or {})
-                continue
-
-            # Handle run info updates
-            if kind == "run_info":
-                if ev.get("model"):
-                    meta["model"] = ev["model"]
-                    renderer.on_start(meta)
-                if ev.get("run_id"):
-                    meta["run_id"] = ev["run_id"]
-                    renderer.on_start(meta)
+            if started_monotonic is None:
+                try:
+                    ev = json.loads(event["data"])
+                    if self._should_start_timer(ev):
+                        started_monotonic = monotonic()
+                except json.JSONDecodeError:
+                    pass
+
+            final_text, stats_usage = self._process_single_event(
+                event, renderer, final_text, stats_usage, meta
+            )
 
         finished_monotonic = monotonic()
         return final_text, stats_usage, started_monotonic, finished_monotonic
@@ -678,9 +765,13 @@ class AgentClient(BaseClient):
     ) -> str:
         """Run an agent with a message, streaming via a renderer."""
         # Prepare request payload and headers
-        payload, data_payload, files_payload, headers, multipart_data = (
-            self._prepare_payload_and_headers(message, files, tty, **kwargs)
-        )
+        (
+            payload,
+            data_payload,
+            files_payload,
+            headers,
+            multipart_data,
+        ) = self._prepare_sync_request_data(message, files, tty, **kwargs)
 
         # Create renderer
         r = self._create_renderer(renderer, **kwargs)
@@ -712,10 +803,13 @@ class AgentClient(BaseClient):
                 timeout_seconds = kwargs.get("timeout", DEFAULT_AGENT_RUN_TIMEOUT)
                 agent_name = kwargs.get("agent_name")
 
-                final_text, stats_usage, started_monotonic, finished_monotonic = (
-                    self._process_stream_events(
-                        stream_response, r, timeout_seconds, agent_name, kwargs
-                    )
+                (
+                    final_text,
+                    stats_usage,
+                    started_monotonic,
+                    finished_monotonic,
+                ) = self._process_stream_events(
+                    stream_response, r, timeout_seconds, agent_name, kwargs
                 )
 
         except KeyboardInterrupt:
@@ -749,6 +843,75 @@ class AgentClient(BaseClient):
         r.on_complete(st)
         return final_payload
 
+    def _prepare_request_data(
+        self,
+        message: str,
+        files: list[str | BinaryIO] | None,
+        **kwargs,
+    ) -> tuple[dict | None, dict | None, dict | None, dict | None]:
+        """Prepare request data for async agent runs.
+
+        Returns:
+            Tuple of (payload, data_payload, files_payload, headers)
+        """
+        if files:
+            # Handle multipart data for file uploads
+            multipart_data = prepare_multipart_data(message, files)
+            # Inject optional multipart extras expected by backend
+            if "chat_history" in kwargs and kwargs["chat_history"] is not None:
+                multipart_data.data["chat_history"] = kwargs["chat_history"]
+            if "pii_mapping" in kwargs and kwargs["pii_mapping"] is not None:
+                multipart_data.data["pii_mapping"] = kwargs["pii_mapping"]
+
+            headers = {"Accept": SSE_CONTENT_TYPE}
+            return None, multipart_data.data, multipart_data.files, headers
+        else:
+            # Simple JSON payload for text-only requests
+            payload = {"input": message, "stream": True, **kwargs}
+            headers = {"Accept": SSE_CONTENT_TYPE}
+            return payload, None, None, headers
+
+    def _create_async_client_config(
+        self, timeout: float | None, headers: dict | None
+    ) -> dict:
+        """Create async client configuration with proper headers and timeout."""
+        config = self._build_async_client(timeout or self.timeout)
+        if headers:
+            config["headers"] = {**config["headers"], **headers}
+        return config
+
+    async def _stream_agent_response(
+        self,
+        async_client: httpx.AsyncClient,
+        agent_id: str,
+        payload: dict | None,
+        data_payload: dict | None,
+        files_payload: dict | None,
+        headers: dict | None,
+        timeout_seconds: float,
+        agent_name: str | None,
+    ) -> AsyncGenerator[dict, None]:
+        """Stream the agent response and yield parsed JSON chunks."""
+        async with async_client.stream(
+            "POST",
+            f"/agents/{agent_id}/run",
+            json=payload,
+            data=data_payload,
+            files=files_payload,
+            headers=headers,
+        ) as stream_response:
+            stream_response.raise_for_status()
+
+            async for event in aiter_sse_events(
+                stream_response, timeout_seconds, agent_name
+            ):
+                try:
+                    chunk = json.loads(event["data"])
+                    yield chunk
+                except json.JSONDecodeError:
+                    logger.debug("Non-JSON SSE fragment skipped")
+                    continue
+
     async def arun_agent(
         self,
         agent_id: str,
@@ -775,74 +938,34 @@ class AgentClient(BaseClient):
             httpx.TimeoutException: When general timeout occurs
             Exception: For other unexpected errors
         """
-        # Prepare multipart data if files are provided
-        multipart_data = None
-        headers = None  # None means "don't override client defaults"
-
-        if files:
-            multipart_data = prepare_multipart_data(message, files)
-            # Inject optional multipart extras expected by backend
-            if "chat_history" in kwargs and kwargs["chat_history"] is not None:
-                multipart_data.data["chat_history"] = kwargs["chat_history"]
-            if "pii_mapping" in kwargs and kwargs["pii_mapping"] is not None:
-                multipart_data.data["pii_mapping"] = kwargs["pii_mapping"]
-            headers = None  # Let httpx set proper multipart boundaries
-
-        # When streaming, explicitly prefer SSE
-        headers = {**(headers or {}), "Accept": "text/event-stream"}
+        # Prepare request data
+        payload, data_payload, files_payload, headers = self._prepare_request_data(
+            message, files, **kwargs
+        )
 
-        if files:
-            payload = None
-            # Use multipart data
-            data_payload = multipart_data.data
-            files_payload = multipart_data.files
-        else:
-            payload = {"input": message, **kwargs}
-            # Explicitly send stream intent both ways
-            payload["stream"] = True
-            data_payload = None
-            files_payload = None
+        # Create async client configuration
+        async_client_config = self._create_async_client_config(timeout, headers)
 
-        # Use timeout from parameter or instance default
-        request_timeout = timeout or self.timeout
+        # Get execution timeout for streaming control
+        timeout_seconds = kwargs.get("timeout", DEFAULT_AGENT_RUN_TIMEOUT)
+        agent_name = kwargs.get("agent_name")
 
         try:
-            async_client_config = self._build_async_client(request_timeout)
-            if headers:
-                async_client_config["headers"] = {
-                    **async_client_config["headers"],
-                    **headers,
-                }
-
-            # Create async client for this request
+            # Create async client and stream response
             async with httpx.AsyncClient(**async_client_config) as async_client:
-                async with async_client.stream(
-                    "POST",
-                    f"/agents/{agent_id}/run",
-                    json=payload,
-                    data=data_payload,
-                    files=files_payload,
-                    headers=headers,
-                ) as stream_response:
-                    stream_response.raise_for_status()
-
-                    # Get agent run timeout for execution control
-                    # Prefer parameter timeout, otherwise use default
-                    timeout_seconds = kwargs.get("timeout", DEFAULT_AGENT_RUN_TIMEOUT)
-
-                    agent_name = kwargs.get("agent_name")
-
-                    async for event in aiter_sse_events(
-                        stream_response, timeout_seconds, agent_name
-                    ):
-                        try:
-                            chunk = json.loads(event["data"])
-                            yield chunk
-                        except json.JSONDecodeError:
-                            logger.debug("Non-JSON SSE fragment skipped")
-                            continue
+                async for chunk in self._stream_agent_response(
+                    async_client,
+                    agent_id,
+                    payload,
+                    data_payload,
+                    files_payload,
+                    headers,
+                    timeout_seconds,
+                    agent_name,
+                ):
+                    yield chunk
 
         finally:
-            # Ensure we close any opened file handles from multipart
-            if multipart_data:
-                multipart_data.close()
+            # Ensure cleanup - this is handled by the calling context
+            # but we keep this for safety in case of future changes
+            pass