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

glaip_sdk/client/agents.py

@@ -5,11 +5,14 @@ Authors:
     Raymond Christopher (raymond.christopher@gdplabs.id)
 """
 
+import io
 import json
 import logging
+from collections.abc import AsyncGenerator
 from time import monotonic
 from typing import Any, BinaryIO
 
+import httpx
 from rich.console import Console as _Console
 
 from glaip_sdk.client.base import BaseClient
@@ -21,8 +24,10 @@ from glaip_sdk.config.constants import (
     DEFAULT_AGENT_VERSION,
     DEFAULT_MODEL,
 )
+from glaip_sdk.exceptions import NotFoundError
 from glaip_sdk.models import Agent
 from glaip_sdk.utils.client_utils import (
+    aiter_sse_events,
     create_model_instances,
     extract_ids,
     find_by_name,
@@ -32,6 +37,10 @@ from glaip_sdk.utils.client_utils import (
 from glaip_sdk.utils.rendering.models import RunStats
 from glaip_sdk.utils.rendering.renderer import RichStreamRenderer
 from glaip_sdk.utils.rendering.renderer.config import RendererConfig
+from glaip_sdk.utils.validation import validate_agent_instruction
+
+# API endpoints
+AGENTS_ENDPOINT = "/agents/"
 
 # Set up module-level logger
 logger = logging.getLogger("glaip_sdk.agents")
@@ -49,14 +58,87 @@ class AgentClient(BaseClient):
         """
         super().__init__(parent_client=parent_client, **kwargs)
 
-    def list_agents(self) -> list[Agent]:
-        """List all agents."""
-        data = self._request("GET", "/agents/")
+    def list_agents(
+        self,
+        agent_type: str | None = None,
+        framework: str | None = None,
+        name: str | None = None,
+        version: str | None = None,
+        sync_langflow_agents: bool = False,
+    ) -> list[Agent]:
+        """List agents with optional filtering.
+
+        Args:
+            agent_type: Filter by agent type (config, code, a2a)
+            framework: Filter by framework (langchain, langgraph, google_adk)
+            name: Filter by partial name match (case-insensitive)
+            version: Filter by exact version match
+            sync_langflow_agents: Sync with LangFlow server before listing (only applies when agent_type=langflow)
+
+        Returns:
+            List of agents matching the filters
+        """
+        params = {}
+        if agent_type is not None:
+            params["agent_type"] = agent_type
+        if framework is not None:
+            params["framework"] = framework
+        if name is not None:
+            params["name"] = name
+        if version is not None:
+            params["version"] = version
+        if sync_langflow_agents:
+            params["sync_langflow_agents"] = "true"
+
+        if params:
+            data = self._request("GET", AGENTS_ENDPOINT, params=params)
+        else:
+            data = self._request("GET", AGENTS_ENDPOINT)
         return create_model_instances(data, Agent, self)
 
+    def sync_langflow_agents(
+        self,
+        base_url: str | None = None,
+        api_key: str | None = None,
+    ) -> dict[str, Any]:
+        """Sync LangFlow agents by fetching flows from the LangFlow server.
+
+        This method synchronizes agents with LangFlow flows. It fetches all flows
+        from the configured LangFlow server and creates/updates corresponding agents.
+
+        Args:
+            base_url: Custom LangFlow server base URL. If not provided, uses LANGFLOW_BASE_URL env var.
+            api_key: Custom LangFlow API key. If not provided, uses LANGFLOW_API_KEY env var.
+
+        Returns:
+            Response containing sync results and statistics
+
+        Raises:
+            ValueError: If LangFlow server configuration is missing
+        """
+        payload = {}
+        if base_url is not None:
+            payload["base_url"] = base_url
+        if api_key is not None:
+            payload["api_key"] = api_key
+
+        return self._request("POST", "/agents/langflow/sync", json=payload)
+
     def get_agent_by_id(self, agent_id: str) -> Agent:
         """Get agent by ID."""
         data = self._request("GET", f"/agents/{agent_id}")
+
+        if isinstance(data, str):
+            # Some backends may respond with plain text for missing agents.
+            message = data.strip() or f"Agent '{agent_id}' not found"
+            raise NotFoundError(message, status_code=404)
+
+        if not isinstance(data, dict):
+            raise NotFoundError(
+                f"Agent '{agent_id}' not found (unexpected response type)",
+                status_code=404,
+            )
+
         return Agent(**data)._set_client(self)
 
     def find_agents(self, name: str | None = None) -> list[Agent]:
@@ -65,13 +147,13 @@ class AgentClient(BaseClient):
         if name:
             params["name"] = name
 
-        data = self._request("GET", "/agents/", params=params)
+        data = self._request("GET", AGENTS_ENDPOINT, params=params)
         agents = create_model_instances(data, Agent, self)
         if name is None:
             return agents
         return find_by_name(agents, name, case_sensitive=False)
 
-    def create_agent(
+    def _build_create_payload(
         self,
         name: str,
         instruction: str,
@@ -80,38 +162,53 @@ class AgentClient(BaseClient):
         agents: list[str | Any] | None = None,
         timeout: int = DEFAULT_AGENT_RUN_TIMEOUT,
         **kwargs,
-    ) -> "Agent":
-        """Create a new agent."""
-        # Client-side validation
-        if not name or not name.strip():
-            raise ValueError("Agent name cannot be empty or whitespace")
-
-        if not instruction or not instruction.strip():
-            raise ValueError("Agent instruction cannot be empty or whitespace")
+    ) -> dict[str, Any]:
+        """Build payload for agent creation with proper LM selection and metadata handling.
 
-        if len(instruction.strip()) < 10:
-            raise ValueError("Agent instruction must be at least 10 characters long")
+        CENTRALIZED PAYLOAD BUILDING LOGIC:
+        - LM exclusivity: Uses language_model_id if provided, otherwise provider/model_name
+        - Always includes required backend metadata
+        - Preserves mem0 keys in agent_config
+        - Handles tool/agent ID extraction from objects
 
-        # Prepare the creation payload
+        Args:
+            name: Agent name
+            instruction: Agent instruction
+            model: Language model name (used when language_model_id not provided)
+            tools: List of tools to attach
+            agents: List of sub-agents to attach
+            timeout: Agent execution timeout
+            **kwargs: Additional parameters (language_model_id, agent_config, etc.)
+
+        Returns:
+            Complete payload dictionary for agent creation
+        """
+        # Prepare the creation payload with required fields
         payload: dict[str, Any] = {
             "name": name.strip(),
             "instruction": instruction.strip(),
             "type": DEFAULT_AGENT_TYPE,
             "framework": DEFAULT_AGENT_FRAMEWORK,
             "version": DEFAULT_AGENT_VERSION,
-            "provider": DEFAULT_AGENT_PROVIDER,
-            "model_name": model or DEFAULT_MODEL,  # Ensure model_name is never None
         }
 
-        # Include default execution timeout if provided
+        # Language model selection with exclusivity:
+        # Priority: language_model_id (if provided) > provider/model_name (fallback)
+        if kwargs.get("language_model_id"):
+            # Use language_model_id - defer to kwargs update below
+            pass
+        else:
+            # Use provider/model_name fallback
+            payload["provider"] = DEFAULT_AGENT_PROVIDER
+            payload["model_name"] = model or DEFAULT_MODEL
+
+        # Include execution timeout if provided
         if timeout is not None:
             payload["timeout"] = str(timeout)
 
         # Ensure minimum required metadata for visibility
         if "metadata" not in kwargs:
             kwargs["metadata"] = {}
-
-        # Always include the minimum required metadata for visibility
         if "type" not in kwargs["metadata"]:
             kwargs["metadata"]["type"] = "custom"
 
@@ -125,32 +222,16 @@ class AgentClient(BaseClient):
         if agent_ids:
             payload["agents"] = agent_ids
 
-        # Add any additional kwargs
+        # Add any additional kwargs (including language_model_id, agent_config, etc.)
         payload.update(kwargs)
 
-        # Create the agent and fetch full details
-        full_agent_data = self._post_then_fetch(
-            id_key="id",
-            post_endpoint="/agents/",
-            get_endpoint_fmt="/agents/{id}",
-            json=payload,
-        )
-        return Agent(**full_agent_data)._set_client(self)
-
-    def update_agent(
-        self,
-        agent_id: str,
-        name: str | None = None,
-        instruction: str | None = None,
-        model: str | None = None,
-        **kwargs,
-    ) -> "Agent":
-        """Update an existing agent."""
-        # First, get the current agent data
-        current_agent = self.get_agent_by_id(agent_id)
+        return payload
 
-        # Prepare the update payload with current values as defaults
-        update_data = {
+    def _build_basic_update_payload(
+        self, current_agent: "Agent", name: str | None, instruction: str | None
+    ) -> dict[str, Any]:
+        """Build the basic update payload with required fields."""
+        return {
             "name": name if name is not None else current_agent.name,
             "instruction": instruction
             if instruction is not None
@@ -160,48 +241,250 @@ class AgentClient(BaseClient):
             "version": DEFAULT_AGENT_VERSION,  # Required by backend
         }
 
-        # Handle model specification
-        if model is not None:
-            update_data["provider"] = DEFAULT_AGENT_PROVIDER  # Default provider
+    def _handle_language_model_selection(
+        self,
+        update_data: dict[str, Any],
+        current_agent: "Agent",
+        model: str | None,
+        language_model_id: str | None,
+    ) -> None:
+        """Handle language model selection with proper priority and fallbacks."""
+        if language_model_id:
+            # Use language_model_id if provided
+            update_data["language_model_id"] = language_model_id
+        elif model is not None:
+            # Use explicit model parameter
+            update_data["provider"] = DEFAULT_AGENT_PROVIDER
             update_data["model_name"] = model
         else:
-            # Use current model if available
-            if hasattr(current_agent, "agent_config") and current_agent.agent_config:
-                if "lm_provider" in current_agent.agent_config:
-                    update_data["provider"] = current_agent.agent_config["lm_provider"]
-                if "lm_name" in current_agent.agent_config:
-                    update_data["model_name"] = current_agent.agent_config["lm_name"]
-            else:
-                # Default values
-                update_data["provider"] = DEFAULT_AGENT_PROVIDER
-                update_data["model_name"] = DEFAULT_MODEL
+            # Use current agent config or fallbacks
+            self._set_language_model_from_current_agent(update_data, current_agent)
+
+    def _set_language_model_from_current_agent(
+        self, update_data: dict[str, Any], current_agent: "Agent"
+    ) -> None:
+        """Set language model from current agent config or use defaults."""
+        if hasattr(current_agent, "agent_config") and current_agent.agent_config:
+            agent_config = current_agent.agent_config
+            if "lm_provider" in agent_config:
+                update_data["provider"] = agent_config["lm_provider"]
+            if "lm_name" in agent_config:
+                update_data["model_name"] = agent_config["lm_name"]
+        else:
+            # Default fallback values
+            update_data["provider"] = DEFAULT_AGENT_PROVIDER
+            update_data["model_name"] = DEFAULT_MODEL
 
-        # Handle tools and agents
-        if "tools" in kwargs:
-            tool_ids = extract_ids(kwargs["tools"])
-            if tool_ids:
-                update_data["tools"] = tool_ids
-        elif current_agent.tools:
-            update_data["tools"] = [
+    def _handle_tools_and_agents(
+        self,
+        update_data: dict[str, Any],
+        current_agent: "Agent",
+        tools: list | None,
+        agents: list | None,
+    ) -> None:
+        """Handle tools and agents with proper ID extraction."""
+        # Handle tools
+        if tools is not None:
+            tool_ids = extract_ids(tools)
+            update_data["tools"] = tool_ids if tool_ids else []
+        else:
+            update_data["tools"] = self._extract_current_tool_ids(current_agent)
+
+        # Handle agents
+        if agents is not None:
+            agent_ids = extract_ids(agents)
+            update_data["agents"] = agent_ids if agent_ids else []
+        else:
+            update_data["agents"] = self._extract_current_agent_ids(current_agent)
+
+    def _extract_current_tool_ids(self, current_agent: "Agent") -> list[str]:
+        """Extract tool IDs from current agent."""
+        if current_agent.tools:
+            return [
                 tool["id"] if isinstance(tool, dict) else tool
                 for tool in current_agent.tools
             ]
+        return []
 
-        if "agents" in kwargs:
-            agent_ids = extract_ids(kwargs["agents"])
-            if agent_ids:
-                update_data["agents"] = agent_ids
-        elif current_agent.agents:
-            update_data["agents"] = [
+    def _extract_current_agent_ids(self, current_agent: "Agent") -> list[str]:
+        """Extract agent IDs from current agent."""
+        if current_agent.agents:
+            return [
                 agent["id"] if isinstance(agent, dict) else agent
                 for agent in current_agent.agents
             ]
+        return []
 
-        # Add any other kwargs
+    def _handle_agent_config(
+        self,
+        update_data: dict[str, Any],
+        current_agent: "Agent",
+        agent_config: dict | None,
+    ) -> None:
+        """Handle agent_config with proper merging and cleanup."""
+        if agent_config is not None:
+            # Use provided agent_config, merging with current if needed
+            update_data["agent_config"] = self._merge_agent_configs(
+                current_agent, agent_config
+            )
+        elif hasattr(current_agent, "agent_config") and current_agent.agent_config:
+            # Preserve existing agent_config
+            update_data["agent_config"] = current_agent.agent_config.copy()
+        else:
+            # Default agent_config
+            update_data["agent_config"] = {
+                "lm_provider": DEFAULT_AGENT_PROVIDER,
+                "lm_name": DEFAULT_MODEL,
+                "lm_hyperparameters": {"temperature": 0.0},
+            }
+
+        # Clean LM keys from agent_config to prevent conflicts
+        self._clean_agent_config_lm_keys(update_data)
+
+    def _merge_agent_configs(self, current_agent: "Agent", new_config: dict) -> dict:
+        """Merge current agent config with new config."""
+        if hasattr(current_agent, "agent_config") and current_agent.agent_config:
+            merged_config = current_agent.agent_config.copy()
+            merged_config.update(new_config)
+            return merged_config
+        return new_config
+
+    def _clean_agent_config_lm_keys(self, update_data: dict[str, Any]) -> None:
+        """Remove LM keys from agent_config to prevent conflicts."""
+        if "agent_config" in update_data and isinstance(
+            update_data["agent_config"], dict
+        ):
+            agent_config = update_data["agent_config"]
+            lm_keys_to_remove = {
+                "lm_provider",
+                "lm_name",
+                "lm_base_url",
+                "lm_hyperparameters",
+            }
+            for key in lm_keys_to_remove:
+                agent_config.pop(key, None)
+
+    def _finalize_update_payload(
+        self, update_data: dict[str, Any], current_agent: "Agent", **kwargs
+    ) -> dict[str, Any]:
+        """Finalize the update payload with metadata and additional kwargs."""
+        # Handle metadata preservation
+        if hasattr(current_agent, "metadata") and current_agent.metadata:
+            update_data["metadata"] = current_agent.metadata.copy()
+
+        # Add any other kwargs (excluding already handled ones)
+        excluded_keys = {"tools", "agents", "agent_config", "language_model_id"}
         for key, value in kwargs.items():
-            if key not in ["tools", "agents"]:
+            if key not in excluded_keys:
                 update_data[key] = value
 
+        return update_data
+
+    def _build_update_payload(
+        self,
+        current_agent: "Agent",
+        name: str | None = None,
+        instruction: str | None = None,
+        model: str | None = None,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """Build payload for agent update with proper LM selection and current state preservation.
+
+        Args:
+            current_agent: Current agent object to update
+            name: New agent name (None to keep current)
+            instruction: New instruction (None to keep current)
+            model: New language model name (None to use current or fallback)
+            **kwargs: Additional parameters including language_model_id, agent_config, etc.
+
+        Returns:
+            Complete payload dictionary for agent update
+
+        Notes:
+            - LM exclusivity: Uses language_model_id if provided, otherwise provider/model_name
+            - Preserves current values as defaults when new values not provided
+            - Handles tools/agents updates with proper ID extraction
+        """
+        # Build basic payload
+        update_data = self._build_basic_update_payload(current_agent, name, instruction)
+
+        # Handle language model selection
+        language_model_id = kwargs.get("language_model_id")
+        self._handle_language_model_selection(
+            update_data, current_agent, model, language_model_id
+        )
+
+        # Handle tools and agents
+        tools = kwargs.get("tools")
+        agents = kwargs.get("agents")
+        self._handle_tools_and_agents(update_data, current_agent, tools, agents)
+
+        # Handle agent config
+        agent_config = kwargs.get("agent_config")
+        self._handle_agent_config(update_data, current_agent, agent_config)
+
+        # Finalize payload
+        return self._finalize_update_payload(update_data, current_agent, **kwargs)
+
+    def create_agent(
+        self,
+        name: str,
+        instruction: str,
+        model: str = DEFAULT_MODEL,
+        tools: list[str | Any] | None = None,
+        agents: list[str | Any] | None = None,
+        timeout: int = DEFAULT_AGENT_RUN_TIMEOUT,
+        **kwargs,
+    ) -> "Agent":
+        """Create a new agent."""
+        # Client-side validation
+        if not name or not name.strip():
+            raise ValueError("Agent name cannot be empty or whitespace")
+
+        # Validate instruction using centralized validation
+        instruction = validate_agent_instruction(instruction)
+
+        # Build payload using centralized builder
+        payload = self._build_create_payload(
+            name=name,
+            instruction=instruction,
+            model=model,
+            tools=tools,
+            agents=agents,
+            timeout=timeout,
+            **kwargs,
+        )
+
+        # Create the agent and fetch full details
+        full_agent_data = self._post_then_fetch(
+            id_key="id",
+            post_endpoint=AGENTS_ENDPOINT,
+            get_endpoint_fmt=f"{AGENTS_ENDPOINT}{{id}}",
+            json=payload,
+        )
+        return Agent(**full_agent_data)._set_client(self)
+
+    def update_agent(
+        self,
+        agent_id: str,
+        name: str | None = None,
+        instruction: str | None = None,
+        model: str | None = None,
+        **kwargs,
+    ) -> "Agent":
+        """Update an existing agent."""
+        # First, get the current agent data
+        current_agent = self.get_agent_by_id(agent_id)
+
+        # Build payload using centralized builder
+        update_data = self._build_update_payload(
+            current_agent=current_agent,
+            name=name,
+            instruction=instruction,
+            model=model,
+            **kwargs,
+        )
+
         # Send the complete payload
         data = self._request("PUT", f"/agents/{agent_id}", json=update_data)
         return Agent(**data)._set_client(self)
@@ -210,20 +493,12 @@ class AgentClient(BaseClient):
         """Delete an agent."""
         self._request("DELETE", f"/agents/{agent_id}")
 
-    def run_agent(
-        self,
-        agent_id: str,
-        message: str,
-        files: list[str | BinaryIO] | None = None,
-        tty: bool = False,
-        *,
-        renderer: RichStreamRenderer | str | None = "auto",
-        **kwargs,
-    ) -> str:
-        """Run an agent with a message, streaming via a renderer."""
-        # Prepare multipart data if files are provided
+    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  # None means "don't override client defaults"
+        headers = None
 
         if files:
             multipart_data = prepare_multipart_data(message, files)
@@ -232,69 +507,195 @@ class AgentClient(BaseClient):
                 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"}
 
         if files:
             payload = None
-            # Use multipart data
             data_payload = multipart_data.data
             files_payload = multipart_data.files
         else:
             payload = {"input": message, **kwargs}
             if tty:
                 payload["tty"] = True
-            # Explicitly send stream intent both ways
             payload["stream"] = True
             data_payload = None
             files_payload = None
 
-        # Choose renderer: use provided instance or create a default
+        return payload, data_payload, files_payload, headers, multipart_data
+
+    def _create_renderer(self, renderer, **kwargs):
+        """Create appropriate renderer based on configuration."""
         if isinstance(renderer, RichStreamRenderer):
-            r = renderer
-        else:
-            # Check if verbose mode is requested
-            verbose = kwargs.get("verbose", False)
-            if verbose:
-                # Create a verbose renderer similar to CLI --verbose
-                verbose_config = RendererConfig(
-                    theme="dark",
-                    style="debug",  # CLI uses "debug" style for verbose
-                    live=False,  # CLI disables live updates for verbose
-                    show_delegate_tool_panels=True,  # CLI always shows tool panels
-                    append_finished_snapshots=False,
-                )
-                r = RichStreamRenderer(
-                    console=_Console(),
-                    cfg=verbose_config,
-                    verbose=True,
-                )
+            return renderer
+
+        verbose = kwargs.get("verbose", False)
+
+        if isinstance(renderer, str):
+            if renderer == "silent":
+                return self._create_silent_renderer()
+            elif renderer == "minimal":
+                return self._create_minimal_renderer()
             else:
-                # Default to a standard rich renderer with tool panels enabled
-                default_config = RendererConfig(
-                    show_delegate_tool_panels=True,  # Enable tool panels by default
-                )
-                r = RichStreamRenderer(console=_Console(), cfg=default_config)
+                return self._create_default_renderer(verbose)
+        elif verbose:
+            return self._create_verbose_renderer()
+        else:
+            return self._create_default_renderer(verbose)
+
+    def _create_silent_renderer(self):
+        """Create a silent renderer that suppresses all output."""
+        silent_config = RendererConfig(
+            live=False,
+            persist_live=False,
+            show_delegate_tool_panels=False,
+            render_thinking=False,
+        )
+        return RichStreamRenderer(
+            console=_Console(file=io.StringIO(), force_terminal=False),
+            cfg=silent_config,
+            verbose=False,
+        )
 
-        # Try to set some meta early; refine as we receive events
-        meta = {
-            "agent_name": kwargs.get("agent_name", agent_id),
-            "model": kwargs.get("model"),
-            "run_id": None,
-            "input_message": message,  # Add the original query for context
-        }
-        r.on_start(meta)
+    def _create_minimal_renderer(self):
+        """Create a minimal renderer with basic output."""
+        minimal_config = RendererConfig(
+            live=False,
+            persist_live=False,
+            show_delegate_tool_panels=False,
+            render_thinking=False,
+        )
+        return RichStreamRenderer(
+            console=_Console(),
+            cfg=minimal_config,
+            verbose=False,
+        )
+
+    def _create_verbose_renderer(self):
+        """Create a verbose renderer for detailed output."""
+        verbose_config = RendererConfig(
+            theme="dark",
+            style="debug",
+            live=False,
+            show_delegate_tool_panels=True,
+            append_finished_snapshots=False,
+        )
+        return RichStreamRenderer(
+            console=_Console(),
+            cfg=verbose_config,
+            verbose=True,
+        )
+
+    def _create_default_renderer(self, verbose: bool):
+        """Create the default renderer."""
+        if verbose:
+            return self._create_verbose_renderer()
+        else:
+            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 = {
+            "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
+        req_id = stream_response.headers.get(
+            "x-request-id"
+        ) or stream_response.headers.get("x-run-id")
+        if req_id:
+            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
+
+            # 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)
+
+        finished_monotonic = monotonic()
+        return final_text, stats_usage, started_monotonic, finished_monotonic
 
-        # MultipartData handles file cleanup automatically
+    def run_agent(
+        self,
+        agent_id: str,
+        message: str,
+        files: list[str | BinaryIO] | None = None,
+        tty: bool = False,
+        *,
+        renderer: RichStreamRenderer | str | None = "auto",
+        **kwargs,
+    ) -> 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)
+        )
+
+        # Create renderer
+        r = self._create_renderer(renderer, **kwargs)
+
+        # Initialize renderer
+        meta = {
+            "agent_name": kwargs.get("agent_name", agent_id),
+            "model": kwargs.get("model"),
+            "run_id": None,
+            "input_message": message,
+        }
+        r.on_start(meta)
 
         try:
+            # Make streaming request
             response = self.http_client.stream(
                 "POST",
                 f"/agents/{agent_id}/run",
@@ -307,71 +708,16 @@ class AgentClient(BaseClient):
             with response as stream_response:
                 stream_response.raise_for_status()
 
-                # capture request id if provided
-                req_id = stream_response.headers.get(
-                    "x-request-id"
-                ) or stream_response.headers.get("x-run-id")
-                if req_id:
-                    meta["run_id"] = req_id
-                    r.on_start(meta)  # refresh header with run_id
-
-                # Get agent run timeout for execution control
-                # Prefer CLI-provided timeout, otherwise use default
+                # Process streaming events
                 timeout_seconds = kwargs.get("timeout", DEFAULT_AGENT_RUN_TIMEOUT)
-
                 agent_name = kwargs.get("agent_name")
 
-                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
-
-                    # 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")
-
-                    # Pass event to the renderer (always, don't filter)
-                    r.on_event(ev)
-
-                    # Hide "artifact" chatter from content accumulation only
-                    if kind == "artifact":
-                        continue
-
-                    # Accumulate assistant content, but do not print here
-                    if "content" in ev and ev["content"]:
-                        # Filter weird backend text like "Artifact received: ..."
-                        if not ev["content"].startswith("Artifact received:"):
-                            final_text = ev["content"]  # replace with latest
-                        continue
-
-                    # Also treat final_response like content for CLI return value
-                    if kind == "final_response" and ev.get("content"):
-                        final_text = ev["content"]  # ensure CLI non-empty
-                        continue
-
-                    # Usage/cost event (if your backend emits it)
-                    if kind == "usage":
-                        stats_usage.update(ev.get("usage") or {})
-                        continue
-
-                    # Model/run info (if emitted mid-stream)
-                    if kind == "run_info":
-                        if ev.get("model"):
-                            meta["model"] = ev["model"]
-                            r.on_start(meta)
-                        if ev.get("run_id"):
-                            meta["run_id"] = ev["run_id"]
-                            r.on_start(meta)
-
-            finished_monotonic = monotonic()
+                final_text, stats_usage, started_monotonic, finished_monotonic = (
+                    self._process_stream_events(
+                        stream_response, r, timeout_seconds, agent_name, kwargs
+                    )
+                )
+
         except KeyboardInterrupt:
             try:
                 r.close()
@@ -383,26 +729,120 @@ class AgentClient(BaseClient):
             finally:
                 raise
         finally:
-            # Ensure we close any opened file handles from multipart
+            # Ensure cleanup
             if multipart_data:
                 multipart_data.close()
 
-        # Finalize stats
+        # Finalize and return result
         st = RunStats()
-        # Ensure monotonic order (avoid negative -0.0s)
-        if started_monotonic is None:
-            started_monotonic = finished_monotonic
-
         st.started_at = started_monotonic or st.started_at
         st.finished_at = finished_monotonic or st.started_at
         st.usage = stats_usage
 
-        # Prefer explicit content, otherwise fall back to what the renderer saw
+        # Get final content
         if hasattr(r, "state") and hasattr(r.state, "buffer"):
             rendered_text = "".join(r.state.buffer)
         else:
             rendered_text = ""
-        final_payload = final_text or rendered_text or "No response content received."
 
+        final_payload = final_text or rendered_text or "No response content received."
         r.on_complete(st)
         return final_payload
+
+    async def arun_agent(
+        self,
+        agent_id: str,
+        message: str,
+        files: list[str | BinaryIO] | None = None,
+        *,
+        timeout: float | None = None,
+        **kwargs,
+    ) -> AsyncGenerator[dict, None]:
+        """Async run an agent with a message, yielding streaming JSON chunks.
+
+        Args:
+            agent_id: ID of the agent to run
+            message: Message to send to the agent
+            files: Optional list of files to include
+            timeout: Request timeout in seconds
+            **kwargs: Additional arguments (chat_history, pii_mapping, etc.)
+
+        Yields:
+            Dictionary containing parsed JSON chunks from the streaming response
+
+        Raises:
+            AgentTimeoutError: When agent execution times out
+            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"}
+
+        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
+
+        # Use timeout from parameter or instance default
+        request_timeout = timeout or self.timeout
+
+        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
+            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
+
+        finally:
+            # Ensure we close any opened file handles from multipart
+            if multipart_data:
+                multipart_data.close()