feat: Add wait_until_ready() method for agent deployment polling (#56)

cipheraxat · bbatha · web-flow · commit dcef3d5ebb4e · 2025-10-20T17:49:09.000-04:00
Resolves #40 - Add wait_until_ready() method to AgentsResource and AsyncAgentsResource - Add AgentDeploymentError and AgentDeploymentTimeoutError exceptions - Add comprehensive test coverage for sync and async variants - Add example usage documentation Users can now call client.agents.wait_until_ready(agent_id) instead of writing custom polling loops to wait for agent deployment completion. The method polls the agent status every 5 seconds (configurable) until: - SUCCESS: Agent reaches STATUS_RUNNING (returns AgentRetrieveResponse) - FAILURE: Agent reaches STATUS_FAILED/STATUS_UNDEPLOYMENT_FAILED/STATUS_DELETED (raises AgentDeploymentError) - TIMEOUT: Agent doesn't reach STATUS_RUNNING within timeout period (raises AgentDeploymentTimeoutError) Default timeout is 300 seconds (5 minutes) and poll_interval is 5 seconds. Both sync and async implementations are provided. -------------------- Co-authored-by: Ben Batha <bhbatha@gmail.com>
diff --git a/examples/agent_wait_until_ready.py b/examples/agent_wait_until_ready.py
@@ -0,0 +1,100 @@
+"""
+Example: Wait for Agent Deployment to Complete
+
+This example demonstrates how to use the wait_until_ready() method to wait for
+an agent to finish deploying before using it.
+"""
+
+from gradient import Gradient
+from gradient._exceptions import AgentDeploymentError, AgentDeploymentTimeoutError
+
+# Initialize the Gradient client
+client = Gradient()
+
+# Create a new agent
+agent_response = client.agents.create(
+    name="My Agent",
+    instruction="You are a helpful assistant",
+    model_uuid="<your-model-uuid>",
+    region="nyc1",
+)
+
+agent_id = agent_response.agent.uuid if agent_response.agent else None
+
+if agent_id:
+    print(f"Agent created with ID: {agent_id}")
+    print("Waiting for agent to be ready...")
+    
+    try:
+        # Wait for the agent to be deployed and ready
+        # This will poll the agent status every 5 seconds (default)
+        # and wait up to 5 minutes (default timeout=300 seconds)
+        ready_agent = client.agents.wait_until_ready(
+            agent_id,
+            poll_interval=5.0,  # Check every 5 seconds
+            timeout=300.0,      # Wait up to 5 minutes
+        )
+        
+        if ready_agent.agent and ready_agent.agent.deployment:
+            print(f"Agent is ready! Status: {ready_agent.agent.deployment.status}")
+            print(f"Agent URL: {ready_agent.agent.url}")
+        
+        # Now you can use the agent
+        # ...
+        
+    except AgentDeploymentError as e:
+        print(f"Agent deployment failed: {e}")
+        print(f"Failed status: {e.status}")
+        
+    except AgentDeploymentTimeoutError as e:
+        print(f"Agent deployment timed out: {e}")
+        print(f"Agent ID: {e.agent_id}")
+        
+    except Exception as e:
+        print(f"Unexpected error: {e}")
+
+
+# Async example
+from gradient import AsyncGradient
+
+
+async def main() -> None:
+    async_client = AsyncGradient()
+    
+    # Create a new agent
+    agent_response = await async_client.agents.create(
+        name="My Async Agent",
+        instruction="You are a helpful assistant",
+        model_uuid="<your-model-uuid>",
+        region="nyc1",
+    )
+    
+    agent_id = agent_response.agent.uuid if agent_response.agent else None
+    
+    if agent_id:
+        print(f"Agent created with ID: {agent_id}")
+        print("Waiting for agent to be ready...")
+        
+        try:
+            # Wait for the agent to be deployed and ready (async)
+            ready_agent = await async_client.agents.wait_until_ready(
+                agent_id,
+                poll_interval=5.0,
+                timeout=300.0,
+            )
+            
+            if ready_agent.agent and ready_agent.agent.deployment:
+                print(f"Agent is ready! Status: {ready_agent.agent.deployment.status}")
+                print(f"Agent URL: {ready_agent.agent.url}")
+            
+        except AgentDeploymentError as e:
+            print(f"Agent deployment failed: {e}")
+            print(f"Failed status: {e.status}")
+            
+        except AgentDeploymentTimeoutError as e:
+            print(f"Agent deployment timed out: {e}")
+            print(f"Agent ID: {e.agent_id}")
+
+
+# Uncomment to run async example
+# asyncio.run(main())
diff --git a/src/gradient/_exceptions.py b/src/gradient/_exceptions.py
@@ -15,6 +15,8 @@
     "UnprocessableEntityError",
     "RateLimitError",
     "InternalServerError",
+    "AgentDeploymentError",
+    "AgentDeploymentTimeoutError",
 ]
 
 
@@ -106,3 +108,19 @@ class RateLimitError(APIStatusError):
 
 class InternalServerError(APIStatusError):
     pass
+
+
+class AgentDeploymentError(GradientError):
+    """Raised when an agent deployment fails."""
+
+    def __init__(self, message: str, status: str) -> None:
+        super().__init__(message)
+        self.status = status
+
+
+class AgentDeploymentTimeoutError(GradientError):
+    """Raised when waiting for an agent deployment times out."""
+
+    def __init__(self, message: str, agent_id: str) -> None:
+        super().__init__(message)
+        self.agent_id = agent_id
diff --git a/src/gradient/resources/agents/agents.py b/src/gradient/resources/agents/agents.py
@@ -1,7 +1,8 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
-
 from __future__ import annotations
 
+import time
+
 import httpx
 
 from .routes import (
@@ -612,6 +613,92 @@ def update_status(
             cast_to=AgentUpdateStatusResponse,
         )
 
+    def wait_until_ready(
+        self,
+        uuid: str,
+        *,
+        timeout: float = 300.0,
+        poll_interval: float = 5.0,
+        # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+        # The extra values given here take precedence over values defined on the client or passed to this method.
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+    ) -> AgentRetrieveResponse:
+        """Wait for an agent to be ready (deployment status is STATUS_RUNNING).
+
+        This method polls the agent status until it reaches STATUS_RUNNING or a terminal
+        error state. It handles timeout and deployment failures automatically.
+
+        Args:
+          uuid: The unique identifier of the agent to wait for
+
+          timeout: Maximum time to wait in seconds (default: 300 seconds / 5 minutes)
+
+          poll_interval: Time to wait between status checks in seconds (default: 5 seconds)
+
+          extra_headers: Send extra headers
+
+          extra_query: Add additional query parameters to the request
+
+          extra_body: Add additional JSON properties to the request
+
+        Returns:
+          AgentRetrieveResponse: The agent response when it reaches STATUS_RUNNING
+
+        Raises:
+          AgentDeploymentError: If the agent deployment fails (STATUS_FAILED,
+              STATUS_UNDEPLOYMENT_FAILED, or STATUS_DELETED)
+          AgentDeploymentTimeoutError: If the agent doesn't reach STATUS_RUNNING
+              within the timeout period
+          ValueError: If uuid is empty
+        """
+        from ..._exceptions import AgentDeploymentError, AgentDeploymentTimeoutError
+
+        if not uuid:
+            raise ValueError(f"Expected a non-empty value for `uuid` but received {uuid!r}")
+
+        start_time = time.time()
+
+        while True:
+            agent_response = self.retrieve(
+                uuid,
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+            )
+
+            # Check if agent and deployment exist
+            if agent_response.agent and agent_response.agent.deployment:
+                status = agent_response.agent.deployment.status
+
+                # Success case
+                if status == "STATUS_RUNNING":
+                    return agent_response
+
+                # Failure cases
+                if status in ("STATUS_FAILED", "STATUS_UNDEPLOYMENT_FAILED", "STATUS_DELETED"):
+                    raise AgentDeploymentError(
+                        f"Agent deployment failed with status: {status}",
+                        status=status,
+                    )
+
+            # Check timeout
+            elapsed_time = time.time() - start_time
+            if elapsed_time >= timeout:
+                current_status = (
+                    agent_response.agent.deployment.status
+                    if agent_response.agent and agent_response.agent.deployment
+                    else "UNKNOWN"
+                )
+                raise AgentDeploymentTimeoutError(
+                    f"Agent did not reach STATUS_RUNNING within {timeout} seconds. Current status: {current_status}",
+                    agent_id=uuid,
+                )
+
+            # Wait before polling again
+            time.sleep(poll_interval)
+
 
 class AsyncAgentsResource(AsyncAPIResource):
     @cached_property
@@ -1108,6 +1195,94 @@ async def update_status(
             cast_to=AgentUpdateStatusResponse,
         )
 
+    async def wait_until_ready(
+        self,
+        uuid: str,
+        *,
+        timeout: float = 300.0,
+        poll_interval: float = 5.0,
+        # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+        # The extra values given here take precedence over values defined on the client or passed to this method.
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+    ) -> AgentRetrieveResponse:
+        """Wait for an agent to be ready (deployment status is STATUS_RUNNING).
+
+        This method polls the agent status until it reaches STATUS_RUNNING or a terminal
+        error state. It handles timeout and deployment failures automatically.
+
+        Args:
+          uuid: The unique identifier of the agent to wait for
+
+          timeout: Maximum time to wait in seconds (default: 300 seconds / 5 minutes)
+
+          poll_interval: Time to wait between status checks in seconds (default: 5 seconds)
+
+          extra_headers: Send extra headers
+
+          extra_query: Add additional query parameters to the request
+
+          extra_body: Add additional JSON properties to the request
+
+        Returns:
+          AgentRetrieveResponse: The agent response when it reaches STATUS_RUNNING
+
+        Raises:
+          AgentDeploymentError: If the agent deployment fails (STATUS_FAILED,
+              STATUS_UNDEPLOYMENT_FAILED, or STATUS_DELETED)
+          AgentDeploymentTimeoutError: If the agent doesn't reach STATUS_RUNNING
+              within the timeout period
+          ValueError: If uuid is empty
+        """
+        import asyncio
+
+        from ..._exceptions import AgentDeploymentError, AgentDeploymentTimeoutError
+
+        if not uuid:
+            raise ValueError(f"Expected a non-empty value for `uuid` but received {uuid!r}")
+
+        start_time = time.time()
+
+        while True:
+            agent_response = await self.retrieve(
+                uuid,
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+            )
+
+            # Check if agent and deployment exist
+            if agent_response.agent and agent_response.agent.deployment:
+                status = agent_response.agent.deployment.status
+
+                # Success case
+                if status == "STATUS_RUNNING":
+                    return agent_response
+
+                # Failure cases
+                if status in ("STATUS_FAILED", "STATUS_UNDEPLOYMENT_FAILED", "STATUS_DELETED"):
+                    raise AgentDeploymentError(
+                        f"Agent deployment failed with status: {status}",
+                        status=status,
+                    )
+
+            # Check timeout
+            elapsed_time = time.time() - start_time
+            if elapsed_time >= timeout:
+                current_status = (
+                    agent_response.agent.deployment.status
+                    if agent_response.agent and agent_response.agent.deployment
+                    else "UNKNOWN"
+                )
+                raise AgentDeploymentTimeoutError(
+                    f"Agent did not reach STATUS_RUNNING within {timeout} seconds. Current status: {current_status}",
+                    agent_id=uuid,
+                )
+
+            # Wait before polling again
+            await asyncio.sleep(poll_interval)
+
 
 class AgentsResourceWithRawResponse:
     def __init__(self, agents: AgentsResource) -> None:
@@ -1134,6 +1309,9 @@ def __init__(self, agents: AgentsResource) -> None:
         self.update_status = to_raw_response_wrapper(
             agents.update_status,
         )
+        self.wait_until_ready = to_raw_response_wrapper(
+            agents.wait_until_ready,
+        )
 
     @cached_property
     def api_keys(self) -> APIKeysResourceWithRawResponse:
@@ -1201,6 +1379,9 @@ def __init__(self, agents: AsyncAgentsResource) -> None:
         self.update_status = async_to_raw_response_wrapper(
             agents.update_status,
         )
+        self.wait_until_ready = async_to_raw_response_wrapper(
+            agents.wait_until_ready,
+        )
 
     @cached_property
     def api_keys(self) -> AsyncAPIKeysResourceWithRawResponse:
@@ -1268,6 +1449,9 @@ def __init__(self, agents: AgentsResource) -> None:
         self.update_status = to_streamed_response_wrapper(
             agents.update_status,
         )
+        self.wait_until_ready = to_streamed_response_wrapper(
+            agents.wait_until_ready,
+        )
 
     @cached_property
     def api_keys(self) -> APIKeysResourceWithStreamingResponse:
@@ -1335,6 +1519,9 @@ def __init__(self, agents: AsyncAgentsResource) -> None:
         self.update_status = async_to_streamed_response_wrapper(
             agents.update_status,
         )
+        self.wait_until_ready = async_to_streamed_response_wrapper(
+            agents.wait_until_ready,
+        )
 
     @cached_property
     def api_keys(self) -> AsyncAPIKeysResourceWithStreamingResponse:
diff --git a/tests/api_resources/test_agents.py b/tests/api_resources/test_agents.py
