Expose agent deployment exceptions via public API and fix example import

examples/agent_wait_until_ready.py

@@ -5,8 +5,7 @@
 an agent to finish deploying before using it.
 """
 
-from gradient import Gradient
-from gradient._exceptions import AgentDeploymentError, AgentDeploymentTimeoutError
+from gradient import Gradient, AgentDeploymentError, AgentDeploymentTimeoutError
 
 # Initialize the Gradient client
 client = Gradient()
@@ -24,32 +23,32 @@
 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
+            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}")
 
@@ -60,37 +59,37 @@
 
 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}")

src/gradient/__init__.py

@@ -32,9 +32,11 @@
     APIConnectionError,
     AuthenticationError,
     InternalServerError,
+    AgentDeploymentError,
     PermissionDeniedError,
     UnprocessableEntityError,
     APIResponseValidationError,
+    AgentDeploymentTimeoutError,
 )
 from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
 from ._utils._logs import setup_logging as _setup_logging
@@ -57,6 +59,8 @@
     "APITimeoutError",
     "APIConnectionError",
     "APIResponseValidationError",
+    "AgentDeploymentError",
+    "AgentDeploymentTimeoutError",
     "BadRequestError",
     "AuthenticationError",
     "PermissionDeniedError",

tests/api_resources/test_agents.py

@@ -368,9 +368,10 @@ def test_path_params_update_status(self, client: Gradient) -> None:
     def test_method_wait_until_ready(self, client: Gradient, respx_mock: Any) -> None:
         """Test successful wait_until_ready when agent becomes ready."""
         agent_uuid = "test-agent-id"
-        
+
         # Create side effect that returns different responses
         call_count = [0]
+
         def get_response(_: httpx.Request) -> httpx.Response:
             call_count[0] += 1
             if call_count[0] == 1:
@@ -395,9 +396,9 @@ def get_response(_: httpx.Request) -> httpx.Response:
                         }
                     },
                 )
-        
+
         respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock(side_effect=get_response)
-        
+
         agent = client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=10.0)
         assert_matches_type(AgentRetrieveResponse, agent, path=["response"])
         assert agent.agent is not None
@@ -408,9 +409,9 @@ def get_response(_: httpx.Request) -> httpx.Response:
     def test_wait_until_ready_timeout(self, client: Gradient, respx_mock: Any) -> None:
         """Test that wait_until_ready raises timeout error."""
         from gradient._exceptions import AgentDeploymentTimeoutError
-        
+
         agent_uuid = "test-agent-id"
-        
+
         # Mock always returns deploying
         respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock(
             return_value=httpx.Response(
@@ -423,20 +424,20 @@ def test_wait_until_ready_timeout(self, client: Gradient, respx_mock: Any) -> No
                 },
             )
         )
-        
+
         with pytest.raises(AgentDeploymentTimeoutError) as exc_info:
             client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=0.5)
-        
+
         assert "did not reach STATUS_RUNNING within" in str(exc_info.value)
         assert exc_info.value.agent_id == agent_uuid
 
     @parametrize
     def test_wait_until_ready_deployment_failed(self, client: Gradient, respx_mock: Any) -> None:
         """Test that wait_until_ready raises error on deployment failure."""
         from gradient._exceptions import AgentDeploymentError
-        
+
         agent_uuid = "test-agent-id"
-        
+
         # Mock returns failed status
         respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock(
             return_value=httpx.Response(
@@ -449,10 +450,10 @@ def test_wait_until_ready_deployment_failed(self, client: Gradient, respx_mock:
                 },
             )
         )
-        
+
         with pytest.raises(AgentDeploymentError) as exc_info:
             client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=10.0)
-        
+
         assert "deployment failed with status: STATUS_FAILED" in str(exc_info.value)
         assert exc_info.value.status == "STATUS_FAILED"
 
@@ -810,9 +811,10 @@ async def test_path_params_update_status(self, async_client: AsyncGradient) -> N
     async def test_method_wait_until_ready(self, async_client: AsyncGradient, respx_mock: Any) -> None:
         """Test successful async wait_until_ready when agent becomes ready."""
         agent_uuid = "test-agent-id"
-        
+
         # Create side effect that returns different responses
         call_count = [0]
+
         def get_response(_: httpx.Request) -> httpx.Response:
             call_count[0] += 1
             if call_count[0] == 1:
@@ -837,9 +839,9 @@ def get_response(_: httpx.Request) -> httpx.Response:
                         }
                     },
                 )
-        
+
         respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock(side_effect=get_response)
-        
+
         agent = await async_client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=10.0)
         assert_matches_type(AgentRetrieveResponse, agent, path=["response"])
         assert agent.agent is not None
@@ -850,9 +852,9 @@ def get_response(_: httpx.Request) -> httpx.Response:
     async def test_wait_until_ready_timeout(self, async_client: AsyncGradient, respx_mock: Any) -> None:
         """Test that async wait_until_ready raises timeout error."""
         from gradient._exceptions import AgentDeploymentTimeoutError
-        
+
         agent_uuid = "test-agent-id"
-        
+
         # Mock always returns deploying
         respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock(
             return_value=httpx.Response(
@@ -865,20 +867,20 @@ async def test_wait_until_ready_timeout(self, async_client: AsyncGradient, respx
                 },
             )
         )
-        
+
         with pytest.raises(AgentDeploymentTimeoutError) as exc_info:
             await async_client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=0.5)
-        
+
         assert "did not reach STATUS_RUNNING within" in str(exc_info.value)
         assert exc_info.value.agent_id == agent_uuid
 
     @parametrize
     async def test_wait_until_ready_deployment_failed(self, async_client: AsyncGradient, respx_mock: Any) -> None:
         """Test that async wait_until_ready raises error on deployment failure."""
         from gradient._exceptions import AgentDeploymentError
-        
+
         agent_uuid = "test-agent-id"
-        
+
         # Mock returns failed status
         respx_mock.get(f"/v2/gen-ai/agents/{agent_uuid}").mock(
             return_value=httpx.Response(
@@ -891,9 +893,9 @@ async def test_wait_until_ready_deployment_failed(self, async_client: AsyncGradi
                 },
             )
         )
-        
+
         with pytest.raises(AgentDeploymentError) as exc_info:
             await async_client.agents.wait_until_ready(agent_uuid, poll_interval=0.1, timeout=10.0)
-        
+
         assert "deployment failed with status: STATUS_FAILED" in str(exc_info.value)
         assert exc_info.value.status == "STATUS_FAILED"