Skip to content

Commit 9edc2a6

Browse files
areeb1501areeb1501
authored
feat: add wait_for_completion method to IndexingJobs resource with sy… (#49)
1 parent 911c72f commit 9edc2a6

File tree

7 files changed

+741
-37
lines changed

7 files changed

+741
-37
lines changed

‎examples/agent_wait_until_ready.py‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -24,32 +24,32 @@
2424
if agent_id:
2525
print(f"Agent created with ID: {agent_id}")
2626
print("Waiting for agent to be ready...")
27-
27+
2828
try:
2929
# Wait for the agent to be deployed and ready
3030
# This will poll the agent status every 5 seconds (default)
3131
# and wait up to 5 minutes (default timeout=300 seconds)
3232
ready_agent = client.agents.wait_until_ready(
3333
agent_id,
3434
poll_interval=5.0, # Check every 5 seconds
35-
timeout=300.0, # Wait up to 5 minutes
35+
timeout=300.0, # Wait up to 5 minutes
3636
)
37-
37+
3838
if ready_agent.agent and ready_agent.agent.deployment:
3939
print(f"Agent is ready! Status: {ready_agent.agent.deployment.status}")
4040
print(f"Agent URL: {ready_agent.agent.url}")
41-
41+
4242
# Now you can use the agent
4343
# ...
44-
44+
4545
except AgentDeploymentError as e:
4646
print(f"Agent deployment failed: {e}")
4747
print(f"Failed status: {e.status}")
48-
48+
4949
except AgentDeploymentTimeoutError as e:
5050
print(f"Agent deployment timed out: {e}")
5151
print(f"Agent ID: {e.agent_id}")
52-
52+
5353
except Exception as e:
5454
print(f"Unexpected error: {e}")
5555

@@ -60,37 +60,37 @@
6060

6161
async def main() -> None:
6262
async_client = AsyncGradient()
63-
63+
6464
# Create a new agent
6565
agent_response = await async_client.agents.create(
6666
name="My Async Agent",
6767
instruction="You are a helpful assistant",
6868
model_uuid="<your-model-uuid>",
6969
region="nyc1",
7070
)
71-
71+
7272
agent_id = agent_response.agent.uuid if agent_response.agent else None
73-
73+
7474
if agent_id:
7575
print(f"Agent created with ID: {agent_id}")
7676
print("Waiting for agent to be ready...")
77-
77+
7878
try:
7979
# Wait for the agent to be deployed and ready (async)
8080
ready_agent = await async_client.agents.wait_until_ready(
8181
agent_id,
8282
poll_interval=5.0,
8383
timeout=300.0,
8484
)
85-
85+
8686
if ready_agent.agent and ready_agent.agent.deployment:
8787
print(f"Agent is ready! Status: {ready_agent.agent.deployment.status}")
8888
print(f"Agent URL: {ready_agent.agent.url}")
89-
89+
9090
except AgentDeploymentError as e:
9191
print(f"Agent deployment failed: {e}")
9292
print(f"Failed status: {e.status}")
93-
93+
9494
except AgentDeploymentTimeoutError as e:
9595
print(f"Agent deployment timed out: {e}")
9696
print(f"Agent ID: {e.agent_id}")

‎examples/knowledge_base_indexing_wait.py‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Example: Waiting for Knowledge Base Indexing Job Completion
4+
5+
This example demonstrates how to use the wait_for_completion() method
6+
to automatically wait for a knowledge base indexing job to finish,
7+
without needing to write manual polling loops.
8+
"""
9+
10+
import os
11+
12+
from gradient import Gradient, IndexingJobError, IndexingJobTimeoutError
13+
14+
15+
def main() -> None:
16+
# Initialize the Gradient client
17+
client = Gradient()
18+
19+
# Example 1: Basic usage - wait for indexing job to complete
20+
print("Example 1: Basic usage")
21+
print("-" * 50)
22+
23+
# Create an indexing job (replace with your actual knowledge base UUID)
24+
knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here")
25+
26+
print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}")
27+
indexing_job = client.knowledge_bases.indexing_jobs.create(
28+
knowledge_base_uuid=knowledge_base_uuid,
29+
)
30+
31+
job_uuid = indexing_job.job.uuid if indexing_job.job else None
32+
if not job_uuid:
33+
print("Error: Could not create indexing job")
34+
return
35+
36+
print(f"Indexing job created with UUID: {job_uuid}")
37+
print("Waiting for indexing job to complete...")
38+
39+
try:
40+
# Wait for the job to complete (polls every 5 seconds by default)
41+
completed_job = client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
42+
43+
print("\n✅ Indexing job completed successfully!")
44+
if completed_job.job:
45+
print(f"Phase: {completed_job.job.phase}")
46+
print(f"Total items indexed: {completed_job.job.total_items_indexed}")
47+
print(f"Total items failed: {completed_job.job.total_items_failed}")
48+
print(f"Total datasources: {completed_job.job.total_datasources}")
49+
print(f"Completed datasources: {completed_job.job.completed_datasources}")
50+
51+
except IndexingJobTimeoutError as e:
52+
print(f"\n⏱️ Timeout: {e}")
53+
except IndexingJobError as e:
54+
print(f"\n❌ Error: {e}")
55+
except Exception as e:
56+
print(f"\n❌ Unexpected error: {e}")
57+
58+
59+
def example_with_custom_polling() -> None:
60+
"""Example with custom polling interval and timeout"""
61+
print("\n\nExample 2: Custom polling interval and timeout")
62+
print("-" * 50)
63+
64+
client = Gradient()
65+
knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here")
66+
67+
print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}")
68+
indexing_job = client.knowledge_bases.indexing_jobs.create(
69+
knowledge_base_uuid=knowledge_base_uuid,
70+
)
71+
72+
job_uuid = indexing_job.job.uuid if indexing_job.job else None
73+
if not job_uuid:
74+
print("Error: Could not create indexing job")
75+
return
76+
77+
print(f"Indexing job created with UUID: {job_uuid}")
78+
print("Waiting for indexing job to complete (polling every 10 seconds, 5 minute timeout)...")
79+
80+
try:
81+
# Wait with custom poll interval (10 seconds) and timeout (5 minutes = 300 seconds)
82+
completed_job = client.knowledge_bases.indexing_jobs.wait_for_completion(
83+
job_uuid,
84+
poll_interval=10, # Poll every 10 seconds
85+
timeout=300, # Timeout after 5 minutes
86+
)
87+
88+
print("\n✅ Indexing job completed successfully!")
89+
if completed_job.job:
90+
print(f"Phase: {completed_job.job.phase}")
91+
92+
except IndexingJobTimeoutError:
93+
print("\n⏱️ Job did not complete within 5 minutes")
94+
# You can still check the current status
95+
current_status = client.knowledge_bases.indexing_jobs.retrieve(job_uuid)
96+
if current_status.job:
97+
print(f"Current phase: {current_status.job.phase}")
98+
print(
99+
f"Completed datasources: {current_status.job.completed_datasources}/{current_status.job.total_datasources}"
100+
)
101+
except IndexingJobError as e:
102+
print(f"\n❌ Job failed: {e}")
103+
104+
105+
def example_manual_polling() -> None:
106+
"""Example of the old manual polling approach (for comparison)"""
107+
print("\n\nExample 3: Manual polling (old approach)")
108+
print("-" * 50)
109+
110+
client = Gradient()
111+
knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here")
112+
113+
indexing_job = client.knowledge_bases.indexing_jobs.create(
114+
knowledge_base_uuid=knowledge_base_uuid,
115+
)
116+
117+
job_uuid = indexing_job.job.uuid if indexing_job.job else None
118+
if not job_uuid:
119+
print("Error: Could not create indexing job")
120+
return
121+
122+
print(f"Indexing job created with UUID: {job_uuid}")
123+
print("Manual polling (old approach)...")
124+
125+
import time
126+
127+
while True:
128+
indexing_job_status = client.knowledge_bases.indexing_jobs.retrieve(job_uuid)
129+
130+
if indexing_job_status.job and indexing_job_status.job.phase:
131+
phase = indexing_job_status.job.phase
132+
print(f"Current phase: {phase}")
133+
134+
if phase in ["BATCH_JOB_PHASE_UNKNOWN", "BATCH_JOB_PHASE_PENDING", "BATCH_JOB_PHASE_RUNNING"]:
135+
time.sleep(5)
136+
continue
137+
elif phase == "BATCH_JOB_PHASE_SUCCEEDED":
138+
print("✅ Job completed successfully!")
139+
break
140+
else:
141+
print(f"❌ Job ended with phase: {phase}")
142+
break
143+
144+
145+
async def example_async() -> None:
146+
"""Example using async/await"""
147+
print("\n\nExample 4: Async usage")
148+
print("-" * 50)
149+
150+
from gradient import AsyncGradient
151+
152+
client = AsyncGradient()
153+
knowledge_base_uuid = os.getenv("KNOWLEDGE_BASE_UUID", "your-kb-uuid-here")
154+
155+
print(f"Creating indexing job for knowledge base: {knowledge_base_uuid}")
156+
indexing_job = await client.knowledge_bases.indexing_jobs.create(
157+
knowledge_base_uuid=knowledge_base_uuid,
158+
)
159+
160+
job_uuid = indexing_job.job.uuid if indexing_job.job else None
161+
if not job_uuid:
162+
print("Error: Could not create indexing job")
163+
return
164+
165+
print(f"Indexing job created with UUID: {job_uuid}")
166+
print("Waiting for indexing job to complete (async)...")
167+
168+
try:
169+
completed_job = await client.knowledge_bases.indexing_jobs.wait_for_completion(
170+
job_uuid,
171+
poll_interval=5,
172+
timeout=600, # 10 minute timeout
173+
)
174+
175+
print("\n✅ Indexing job completed successfully!")
176+
if completed_job.job:
177+
print(f"Phase: {completed_job.job.phase}")
178+
179+
except IndexingJobTimeoutError as e:
180+
print(f"\n⏱️ Timeout: {e}")
181+
except IndexingJobError as e:
182+
print(f"\n❌ Error: {e}")
183+
finally:
184+
await client.close()
185+
186+
187+
if __name__ == "__main__":
188+
# Run the basic example
189+
main()
190+
191+
# Uncomment to run other examples:
192+
# example_with_custom_polling()
193+
# example_manual_polling()
194+
195+
# For async example, you would need to run:
196+
# import asyncio
197+
# asyncio.run(example_async())

‎src/gradient/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,16 @@
2929
RateLimitError,
3030
APITimeoutError,
3131
BadRequestError,
32+
IndexingJobError,
3233
APIConnectionError,
3334
AuthenticationError,
3435
InternalServerError,
36+
AgentDeploymentError,
3537
PermissionDeniedError,
38+
IndexingJobTimeoutError,
3639
UnprocessableEntityError,
3740
APIResponseValidationError,
41+
AgentDeploymentTimeoutError,
3842
)
3943
from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
4044
from ._utils._logs import setup_logging as _setup_logging
@@ -65,6 +69,10 @@
6569
"UnprocessableEntityError",
6670
"RateLimitError",
6771
"InternalServerError",
72+
"IndexingJobError",
73+
"IndexingJobTimeoutError",
74+
"AgentDeploymentError",
75+
"AgentDeploymentTimeoutError",
6876
"Timeout",
6977
"RequestOptions",
7078
"Client",

‎src/gradient/_exceptions.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
"UnprocessableEntityError",
1616
"RateLimitError",
1717
"InternalServerError",
18+
"IndexingJobError",
19+
"IndexingJobTimeoutError",
1820
"AgentDeploymentError",
1921
"AgentDeploymentTimeoutError",
2022
]
@@ -110,6 +112,32 @@ class InternalServerError(APIStatusError):
110112
pass
111113

112114

115+
class IndexingJobError(GradientError):
116+
"""Raised when an indexing job fails, encounters an error, or is cancelled."""
117+
118+
uuid: str
119+
phase: str
120+
121+
def __init__(self, message: str, *, uuid: str, phase: str) -> None:
122+
super().__init__(message)
123+
self.uuid = uuid
124+
self.phase = phase
125+
126+
127+
class IndexingJobTimeoutError(GradientError):
128+
"""Raised when polling for an indexing job times out."""
129+
130+
uuid: str
131+
phase: str
132+
timeout: float
133+
134+
def __init__(self, message: str, *, uuid: str, phase: str, timeout: float) -> None:
135+
super().__init__(message)
136+
self.uuid = uuid
137+
self.phase = phase
138+
self.timeout = timeout
139+
140+
113141
class AgentDeploymentError(GradientError):
114142
"""Raised when an agent deployment fails."""
115143

0 commit comments

Comments
 (0)