agent-starter-pack 0.17.5__py3-none-any.whl → 0.18.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of agent-starter-pack might be problematic. Click here for more details.

@@ -127,6 +127,8 @@ playground-remote: build-frontend-if-needed
127
127
  @echo "==============================================================================="
128
128
  uv run python -m {{cookiecutter.agent_directory}}.utils.expose_app --mode remote
129
129
  {%- endif %}
130
+ {%- endif %}
131
+ {%- if cookiecutter.is_adk_live and cookiecutter.deployment_target == 'agent_engine' %}
130
132
 
131
133
  # Start the frontend UI separately for development (requires backend running separately)
132
134
  ui:
@@ -154,9 +156,9 @@ playground-dev:
154
156
 
155
157
  # Deploy the agent remotely
156
158
  {%- if cookiecutter.deployment_target == 'cloud_run' %}
157
- # Usage: make backend [IAP=true] [PORT=8080] - Set IAP=true to enable Identity-Aware Proxy, PORT to specify container port
159
+ # Usage: make deploy [IAP=true] [PORT=8080] - Set IAP=true to enable Identity-Aware Proxy, PORT to specify container port
158
160
  {%- endif %}
159
- backend:
161
+ deploy:
160
162
  {%- if cookiecutter.deployment_target == 'cloud_run' %}
161
163
  PROJECT_ID=$$(gcloud config get-value project) && \
162
164
  gcloud beta run deploy {{cookiecutter.project_name}} \
@@ -177,6 +179,9 @@ backend:
177
179
  uv export --no-hashes --no-header --no-dev --no-emit-project > .requirements.txt && uv run {{cookiecutter.agent_directory}}/agent_engine_app.py
178
180
  {%- endif %}
179
181
 
182
+ # Alias for 'make deploy' for backward compatibility
183
+ backend: deploy
184
+
180
185
 
181
186
  # ==============================================================================
182
187
  # Infrastructure Setup
@@ -226,3 +231,20 @@ lint:
226
231
  uv run ruff check . --diff
227
232
  uv run ruff format . --check --diff
228
233
  uv run mypy .
234
+ {%- if cookiecutter.is_adk and cookiecutter.deployment_target == 'agent_engine' %}
235
+
236
+ # ==============================================================================
237
+ # Gemini Enterprise Integration
238
+ # ==============================================================================
239
+
240
+ # Register the deployed agent to Gemini Enterprise
241
+ # Usage: make register-gemini-enterprise GEMINI_ENTERPRISE_APP_ID=projects/{project_number}/locations/{location}/collections/{collection}/engines/{engine_id} [AGENT_ENGINE_ID=<id>] # Defaults to deployment_metadata.json
242
+ register-gemini-enterprise:
243
+ uvx --from agent-starter-pack agent-starter-pack-register-gemini-enterprise \
244
+ $(if $(GEMINI_ENTERPRISE_APP_ID),--gemini-enterprise-app-id="$(GEMINI_ENTERPRISE_APP_ID)",) \
245
+ $(if $(AGENT_ENGINE_ID),--agent-engine-id="$(AGENT_ENGINE_ID)",) \
246
+ $(if $(GEMINI_DISPLAY_NAME),--display-name="$(GEMINI_DISPLAY_NAME)",) \
247
+ $(if $(GEMINI_DESCRIPTION),--description="$(GEMINI_DESCRIPTION)",) \
248
+ $(if $(GEMINI_TOOL_DESCRIPTION),--tool-description="$(GEMINI_TOOL_DESCRIPTION)",) \
249
+ $(if $(GEMINI_AUTHORIZATION_ID),--authorization-id="$(GEMINI_AUTHORIZATION_ID)",)
250
+ {%- endif %}
@@ -59,16 +59,21 @@ make install && make playground
59
59
  {%- endif %}
60
60
  {%- if cookiecutter.deployment_target == 'cloud_run' %}
61
61
  | `make playground` | Launch local development environment with backend and frontend{%- if cookiecutter.is_adk %} - leveraging `adk web` command. {%- endif %}|
62
- | `make backend` | Deploy agent to Cloud Run (use `IAP=true` to enable Identity-Aware Proxy) |
63
- | `make local-backend` | Launch local development server |
64
- {%- if cookiecutter.deployment_target == 'cloud_run' %}
62
+ | `make deploy` | Deploy agent to Cloud Run (use `IAP=true` to enable Identity-Aware Proxy, `PORT=8080` to specify container port) |
63
+ | `make local-backend` | Launch local development server with hot-reload |
64
+ {%- elif cookiecutter.deployment_target == 'agent_engine' %}
65
+ | `make playground` | Launch Streamlit interface for testing agent locally and remotely |
66
+ | `make deploy` | Deploy agent to Agent Engine |
65
67
  {%- if cookiecutter.is_adk_live %}
66
- | `make ui` | Launch Agent Playground front-end only |
68
+ | `make local-backend` | Launch local development server with hot-reload |
69
+ | `make ui` | Start the frontend UI separately for development (requires backend running separately) |
70
+ | `make playground-dev` | Launch dev playground with both frontend and backend hot-reload |
71
+ | `make playground-remote` | Connect to remote deployed agent with local frontend |
72
+ | `make build-frontend` | Build the frontend for production |
67
73
  {%- endif %}
74
+ {%- if cookiecutter.is_adk %}
75
+ | `make register-gemini-enterprise` | Register deployed agent to Gemini Enterprise (see Makefile for parameters) |
68
76
  {%- endif %}
69
- {%- elif cookiecutter.deployment_target == 'agent_engine' %}
70
- | `make playground` | Launch Streamlit interface for testing agent locally and remotely |
71
- | `make backend` | Deploy agent to Agent Engine |
72
77
  {%- endif %}
73
78
  | `make test` | Run unit and integration tests |
74
79
  | `make lint` | Run code quality checks (codespell, ruff, mypy) |
@@ -76,7 +81,6 @@ make install && make playground
76
81
  {%- if cookiecutter.data_ingestion %}
77
82
  | `make data-ingestion`| Run data ingestion pipeline in the Dev environment |
78
83
  {%- endif %}
79
- | `uv run jupyter lab` | Launch Jupyter notebook |
80
84
 
81
85
  For full command options and usage, refer to the [Makefile](Makefile).
82
86
 
@@ -156,10 +160,10 @@ You can test deployment towards a Dev Environment using the following command:
156
160
 
157
161
  ```bash
158
162
  gcloud config set project <your-dev-project-id>
159
- make backend
163
+ make deploy
160
164
  ```
161
165
  {% if cookiecutter.is_adk_live %}
162
- **Note:** For secure access to your deployed backend, consider using Identity-Aware Proxy (IAP) by running `make backend IAP=true`.
166
+ **Note:** For secure access to your deployed backend, consider using Identity-Aware Proxy (IAP) by running `make deploy IAP=true`.
163
167
  {%- endif %}
164
168
 
165
169
  The repository includes a Terraform configuration for the setup of the Dev Google Cloud project.
@@ -276,6 +276,20 @@ data "google_secret_manager_secret" "github_pat" {
276
276
  secret_id = var.github_pat_secret_id
277
277
  }
278
278
 
279
+ # Get CICD project data for Cloud Build service account
280
+ data "google_project" "cicd_project" {
281
+ project_id = var.cicd_runner_project_id
282
+ }
283
+
284
+ # Grant Cloud Build service account access to GitHub PAT secret
285
+ resource "google_secret_manager_secret_iam_member" "cloudbuild_secret_accessor" {
286
+ project = var.cicd_runner_project_id
287
+ secret_id = data.google_secret_manager_secret.github_pat.secret_id
288
+ role = "roles/secretmanager.secretAccessor"
289
+ member = "serviceAccount:service-${data.google_project.cicd_project.number}@gcp-sa-cloudbuild.iam.gserviceaccount.com"
290
+ depends_on = [resource.google_project_service.cicd_services]
291
+ }
292
+
279
293
  # Create the GitHub connection (fallback for manual Terraform usage)
280
294
  resource "google_cloudbuildv2_connection" "github_connection" {
281
295
  count = var.create_cb_connection ? 0 : 1
@@ -289,7 +303,11 @@ resource "google_cloudbuildv2_connection" "github_connection" {
289
303
  oauth_token_secret_version = "${data.google_secret_manager_secret.github_pat.id}/versions/latest"
290
304
  }
291
305
  }
292
- depends_on = [resource.google_project_service.cicd_services, resource.google_project_service.deploy_project_services]
306
+ depends_on = [
307
+ resource.google_project_service.cicd_services,
308
+ resource.google_project_service.deploy_project_services,
309
+ resource.google_secret_manager_secret_iam_member.cloudbuild_secret_accessor
310
+ ]
293
311
  }
294
312
 
295
313
 
@@ -0,0 +1,406 @@
1
+ #!/usr/bin/env python3
2
+ # Copyright 2025 Google LLC
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ """Utility to register an Agent Engine to Gemini Enterprise."""
17
+
18
+ import argparse
19
+ import json
20
+ import os
21
+ import sys
22
+ from pathlib import Path
23
+
24
+ import requests
25
+ import vertexai
26
+ from google.auth import default
27
+ from google.auth.transport.requests import Request as GoogleAuthRequest
28
+
29
+
30
+ def get_agent_engine_id(
31
+ agent_engine_id: str | None, metadata_file: str = "deployment_metadata.json"
32
+ ) -> str:
33
+ """Get the agent engine ID from parameter or deployment metadata.
34
+
35
+ Args:
36
+ agent_engine_id: Optional agent engine resource name
37
+ metadata_file: Path to deployment metadata JSON file
38
+
39
+ Returns:
40
+ The agent engine resource name
41
+
42
+ Raises:
43
+ ValueError: If agent_engine_id is not provided and metadata file doesn't exist
44
+ """
45
+ if agent_engine_id:
46
+ return agent_engine_id
47
+
48
+ # Try to read from deployment_metadata.json
49
+ metadata_path = Path(metadata_file)
50
+ if not metadata_path.exists():
51
+ raise ValueError(
52
+ f"No agent engine ID provided and {metadata_file} not found. "
53
+ "Please provide --agent-engine-id or deploy your agent first."
54
+ )
55
+
56
+ with open(metadata_path) as f:
57
+ metadata = json.load(f)
58
+ return metadata["remote_agent_engine_id"]
59
+
60
+
61
+ def get_access_token() -> str:
62
+ """Get Google Cloud access token.
63
+
64
+ Returns:
65
+ Access token string
66
+
67
+ Raises:
68
+ SystemExit: If authentication fails
69
+ """
70
+ try:
71
+ credentials, _ = default()
72
+ auth_req = GoogleAuthRequest()
73
+ credentials.refresh(auth_req)
74
+ return credentials.token
75
+ except Exception as e:
76
+ print(f"Error getting access token: {e}", file=sys.stderr)
77
+ print(
78
+ "Please ensure you are authenticated with 'gcloud auth application-default login'",
79
+ file=sys.stderr,
80
+ )
81
+ raise RuntimeError("Failed to get access token") from e
82
+
83
+
84
+ def get_agent_engine_metadata(agent_engine_id: str) -> tuple[str | None, str | None]:
85
+ """Fetch display_name and description from deployed Agent Engine.
86
+
87
+ Args:
88
+ agent_engine_id: Agent Engine resource name
89
+
90
+ Returns:
91
+ Tuple of (display_name, description) - either can be None if not found
92
+ """
93
+ parts = agent_engine_id.split("/")
94
+ if len(parts) < 6:
95
+ return None, None
96
+
97
+ project_id = parts[1]
98
+ location = parts[3]
99
+
100
+ try:
101
+ client = vertexai.Client(project=project_id, location=location)
102
+ agent_engine = client.agent_engines.get(name=agent_engine_id)
103
+
104
+ display_name = getattr(agent_engine.api_resource, "display_name", None)
105
+ description = getattr(agent_engine.api_resource, "description", None)
106
+
107
+ return display_name, description
108
+ except Exception as e:
109
+ print(
110
+ f"Warning: Could not fetch metadata from Agent Engine: {e}", file=sys.stderr
111
+ )
112
+ return None, None
113
+
114
+
115
+ def register_agent(
116
+ agent_engine_id: str,
117
+ gemini_enterprise_app_id: str,
118
+ display_name: str,
119
+ description: str,
120
+ tool_description: str,
121
+ project_id: str | None = None,
122
+ authorization_id: str | None = None,
123
+ ) -> dict:
124
+ """Register an agent engine to Gemini Enterprise.
125
+
126
+ This function attempts to create a new agent registration. If the agent is already
127
+ registered (same reasoning engine), it will automatically update the existing
128
+ registration instead.
129
+
130
+ Args:
131
+ agent_engine_id: Agent engine resource name (e.g., projects/.../reasoningEngines/...)
132
+ gemini_enterprise_app_id: Full Gemini Enterprise app resource name
133
+ (e.g., projects/{project_number}/locations/{location}/collections/{collection}/engines/{engine_id})
134
+ display_name: Display name for the agent in Gemini Enterprise
135
+ description: Description of the agent
136
+ tool_description: Description of what the tool does
137
+ project_id: Optional GCP project ID for billing (extracted from agent_engine_id if not provided)
138
+ authorization_id: Optional OAuth authorization ID
139
+ (e.g., projects/{project_number}/locations/global/authorizations/{auth_id})
140
+
141
+ Returns:
142
+ API response as dictionary
143
+
144
+ Raises:
145
+ requests.HTTPError: If the API request fails
146
+ ValueError: If gemini_enterprise_app_id format is invalid
147
+ """
148
+ # Parse Gemini Enterprise app resource name
149
+ # Format: projects/{project_number}/locations/{location}/collections/{collection}/engines/{engine_id}
150
+ parts = gemini_enterprise_app_id.split("/")
151
+ if (
152
+ len(parts) != 8
153
+ or parts[0] != "projects"
154
+ or parts[2] != "locations"
155
+ or parts[4] != "collections"
156
+ or parts[6] != "engines"
157
+ ):
158
+ raise ValueError(
159
+ f"Invalid GEMINI_ENTERPRISE_APP_ID format. Expected: "
160
+ f"projects/{{project_number}}/locations/{{location}}/collections/{{collection}}/engines/{{engine_id}}, "
161
+ f"got: {gemini_enterprise_app_id}"
162
+ )
163
+
164
+ project_number = parts[1]
165
+ as_location = parts[3]
166
+ collection = parts[5]
167
+ engine_id = parts[7]
168
+
169
+ # Use project from agent engine if not explicitly provided (for billing header)
170
+ if not project_id:
171
+ # Extract from agent_engine_id: projects/{project}/locations/{location}/reasoningEngines/{id}
172
+ agent_parts = agent_engine_id.split("/")
173
+ if len(agent_parts) > 1 and agent_parts[0] == "projects":
174
+ project_id = agent_parts[1]
175
+ else:
176
+ # Fallback to the project number from the Gemini Enterprise App ID.
177
+ project_id = project_number
178
+
179
+ # Get access token
180
+ access_token = get_access_token()
181
+
182
+ # Build API endpoint
183
+ url = (
184
+ f"https://discoveryengine.googleapis.com/v1alpha/projects/{project_number}/"
185
+ f"locations/{as_location}/collections/{collection}/engines/{engine_id}/"
186
+ "assistants/default_assistant/agents"
187
+ )
188
+
189
+ # Request headers
190
+ headers = {
191
+ "Authorization": f"Bearer {access_token}",
192
+ "Content-Type": "application/json",
193
+ "x-goog-user-project": project_id,
194
+ }
195
+
196
+ # Request body
197
+ adk_agent_definition: dict = {
198
+ "tool_settings": {"tool_description": tool_description},
199
+ "provisioned_reasoning_engine": {"reasoningEngine": agent_engine_id},
200
+ }
201
+
202
+ # Add OAuth authorization if provided
203
+ if authorization_id:
204
+ adk_agent_definition["authorizations"] = [authorization_id]
205
+
206
+ payload = {
207
+ "displayName": display_name,
208
+ "description": description,
209
+ "icon": {
210
+ "uri": "https://fonts.gstatic.com/s/i/short-term/release/googlesymbols/smart_toy/default/24px.svg"
211
+ },
212
+ "adk_agent_definition": adk_agent_definition,
213
+ }
214
+
215
+ print("Registering agent to Gemini Enterprise...")
216
+ print(f" Agent Engine: {agent_engine_id}")
217
+ print(f" Gemini Enterprise App: {gemini_enterprise_app_id}")
218
+ print(f" Display Name: {display_name}")
219
+ print(f" API Endpoint: {url}")
220
+
221
+ try:
222
+ # Try to create a new registration first
223
+ response = requests.post(url, headers=headers, json=payload, timeout=30)
224
+ response.raise_for_status()
225
+
226
+ result = response.json()
227
+ print("\n✅ Successfully registered agent to Gemini Enterprise!")
228
+ print(f" Agent Name: {result.get('name', 'N/A')}")
229
+ return result
230
+
231
+ except requests.exceptions.HTTPError as http_err:
232
+ # Check if the error is because the agent already exists
233
+ if response.status_code in (400, 409):
234
+ try:
235
+ error_data = response.json()
236
+ error_message = error_data.get("error", {}).get("message", "")
237
+
238
+ # Check if error indicates the agent already exists
239
+ if (
240
+ "already exists" in error_message.lower()
241
+ or "duplicate" in error_message.lower()
242
+ ):
243
+ print(
244
+ "\n⚠️ Agent already registered. Updating existing registration..."
245
+ )
246
+
247
+ # For update, we need to use the specific agent resource name
248
+ # The agent name should be in the error or we can construct it
249
+ # Format: {url}/{agent_id} but we need to find existing agent first
250
+
251
+ # List existing agents to find the one for this reasoning engine
252
+ list_response = requests.get(url, headers=headers, timeout=30)
253
+ list_response.raise_for_status()
254
+ agents_list = list_response.json().get("agents", [])
255
+
256
+ # Find the agent that matches our reasoning engine
257
+ existing_agent = None
258
+ for agent in agents_list:
259
+ re_name = (
260
+ agent.get("adk_agent_definition", {})
261
+ .get("provisioned_reasoning_engine", {})
262
+ .get("reasoningEngine", "")
263
+ )
264
+ if re_name == agent_engine_id:
265
+ existing_agent = agent
266
+ break
267
+
268
+ if existing_agent:
269
+ agent_name = existing_agent.get("name")
270
+ update_url = f"https://discoveryengine.googleapis.com/v1alpha/{agent_name}"
271
+
272
+ print(f" Updating agent: {agent_name}")
273
+
274
+ # PATCH request to update
275
+ update_response = requests.patch(
276
+ update_url, headers=headers, json=payload, timeout=30
277
+ )
278
+ update_response.raise_for_status()
279
+
280
+ result = update_response.json()
281
+ print(
282
+ "\n✅ Successfully updated agent registration in Gemini Enterprise!"
283
+ )
284
+ print(f" Agent Name: {result.get('name', 'N/A')}")
285
+ return result
286
+ else:
287
+ print(
288
+ "\n❌ Could not find existing agent to update",
289
+ file=sys.stderr,
290
+ )
291
+ raise
292
+ except (ValueError, KeyError):
293
+ # Failed to parse error response, raise original error
294
+ pass
295
+
296
+ # If not an "already exists" error, or update failed, raise the original error
297
+ print(f"\n❌ HTTP error occurred: {http_err}", file=sys.stderr)
298
+ print(f" Response: {response.text}", file=sys.stderr)
299
+ raise
300
+ except requests.exceptions.RequestException as req_err:
301
+ print(f"\n❌ Request error occurred: {req_err}", file=sys.stderr)
302
+ raise
303
+
304
+
305
+ def main() -> None:
306
+ """Main entry point for CLI."""
307
+ parser = argparse.ArgumentParser(
308
+ description="Register an Agent Engine to Gemini Enterprise"
309
+ )
310
+ parser.add_argument(
311
+ "--agent-engine-id",
312
+ help="Agent Engine resource name (e.g., projects/.../reasoningEngines/...). "
313
+ "If not provided, reads from deployment_metadata.json",
314
+ )
315
+ parser.add_argument(
316
+ "--metadata-file",
317
+ default="deployment_metadata.json",
318
+ help="Path to deployment metadata file (default: deployment_metadata.json)",
319
+ )
320
+ parser.add_argument(
321
+ "--gemini-enterprise-app-id",
322
+ help="Gemini Enterprise app full resource name "
323
+ "(e.g., projects/{project_number}/locations/{location}/collections/{collection}/engines/{engine_id}). "
324
+ "Can also be set via GEMINI_ENTERPRISE_APP_ID env var",
325
+ )
326
+ parser.add_argument(
327
+ "--display-name",
328
+ help="Display name for the agent. Can also be set via GEMINI_DISPLAY_NAME env var",
329
+ )
330
+ parser.add_argument(
331
+ "--description",
332
+ help="Description of the agent. Can also be set via GEMINI_DESCRIPTION env var",
333
+ )
334
+ parser.add_argument(
335
+ "--tool-description",
336
+ help="Description of what the tool does. Can also be set via GEMINI_TOOL_DESCRIPTION env var",
337
+ )
338
+ parser.add_argument(
339
+ "--project-id",
340
+ help="GCP project ID (extracted from agent-engine-id if not provided). "
341
+ "Can also be set via GOOGLE_CLOUD_PROJECT env var",
342
+ )
343
+ parser.add_argument(
344
+ "--authorization-id",
345
+ help="OAuth authorization resource name "
346
+ "(e.g., projects/{project_number}/locations/global/authorizations/{auth_id}). "
347
+ "Can also be set via GEMINI_AUTHORIZATION_ID env var",
348
+ )
349
+
350
+ args = parser.parse_args()
351
+
352
+ # Get agent engine ID
353
+ try:
354
+ agent_engine_id = get_agent_engine_id(args.agent_engine_id, args.metadata_file)
355
+ except ValueError as e:
356
+ print(f"Error: {e}", file=sys.stderr)
357
+ sys.exit(1)
358
+
359
+ # Auto-detect display_name and description from Agent Engine
360
+ auto_display_name, auto_description = get_agent_engine_metadata(agent_engine_id)
361
+
362
+ gemini_enterprise_app_id = args.gemini_enterprise_app_id or os.getenv(
363
+ "GEMINI_ENTERPRISE_APP_ID"
364
+ )
365
+ if not gemini_enterprise_app_id:
366
+ print(
367
+ "Error: --gemini-enterprise-app-id or GEMINI_ENTERPRISE_APP_ID env var required",
368
+ file=sys.stderr,
369
+ )
370
+ sys.exit(1)
371
+
372
+ display_name = (
373
+ args.display_name
374
+ or os.getenv("GEMINI_DISPLAY_NAME")
375
+ or auto_display_name
376
+ or "My Agent"
377
+ )
378
+ description = (
379
+ args.description
380
+ or os.getenv("GEMINI_DESCRIPTION")
381
+ or auto_description
382
+ or "AI Agent"
383
+ )
384
+ tool_description = (
385
+ args.tool_description or os.getenv("GEMINI_TOOL_DESCRIPTION") or description
386
+ )
387
+ project_id = args.project_id or os.getenv("GOOGLE_CLOUD_PROJECT")
388
+ authorization_id = args.authorization_id or os.getenv("GEMINI_AUTHORIZATION_ID")
389
+
390
+ try:
391
+ register_agent(
392
+ agent_engine_id=agent_engine_id,
393
+ gemini_enterprise_app_id=gemini_enterprise_app_id,
394
+ display_name=display_name,
395
+ description=description,
396
+ tool_description=tool_description,
397
+ project_id=project_id,
398
+ authorization_id=authorization_id,
399
+ )
400
+ except Exception as e:
401
+ print(f"Error during registration: {e}", file=sys.stderr)
402
+ sys.exit(1)
403
+
404
+
405
+ if __name__ == "__main__":
406
+ main()
@@ -17,6 +17,7 @@ import logging
17
17
  import os
18
18
  import sys
19
19
 
20
+ import backoff
20
21
  from data_ingestion_pipeline.pipeline import pipeline
21
22
  from google.cloud import aiplatform
22
23
  from kfp import compiler
@@ -136,6 +137,22 @@ def parse_args() -> argparse.Namespace:
136
137
  return parsed_args
137
138
 
138
139
 
140
+ @backoff.on_exception(
141
+ backoff.expo,
142
+ Exception,
143
+ max_tries=3,
144
+ max_time=3600,
145
+ on_backoff=lambda details: logging.warning(
146
+ f"Pipeline attempt {details['tries']} failed, retrying in {details['wait']:.1f}s..."
147
+ ),
148
+ )
149
+ def submit_and_wait_pipeline(pipeline_job_params: dict, service_account: str) -> None:
150
+ """Submit pipeline job and wait for completion with retry logic."""
151
+ job = aiplatform.PipelineJob(**pipeline_job_params)
152
+ job.submit(service_account=service_account)
153
+ job.wait()
154
+
155
+
139
156
  if __name__ == "__main__":
140
157
  args = parse_args()
141
158
 
@@ -182,17 +199,14 @@ if __name__ == "__main__":
182
199
  )
183
200
  {%- endif %}
184
201
 
185
- # Create pipeline job instance
186
- job = aiplatform.PipelineJob(**pipeline_job_params)
187
-
188
202
  if not args.schedule_only:
189
203
  logging.info("Running pipeline and waiting for completion...")
190
- job.submit(service_account=args.service_account)
191
- job.wait()
204
+ submit_and_wait_pipeline(pipeline_job_params, args.service_account)
192
205
  logging.info("Pipeline completed!")
193
206
 
194
207
  if args.cron_schedule and args.schedule_only:
195
- # No need to create new job instance since we already have one with the same params
208
+ # Create pipeline job instance for scheduling
209
+ job = aiplatform.PipelineJob(**pipeline_job_params)
196
210
  pipeline_job_schedule = aiplatform.PipelineJobSchedule(
197
211
  pipeline_job=job,
198
212
  display_name=f"{args.pipeline_name} Weekly Ingestion Job",
@@ -5,6 +5,7 @@ description = "Data ingestion pipeline for RAG retriever"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.9, <=3.13"
7
7
  dependencies = [
8
+ "backoff>=2.2.0",
8
9
  "google-cloud-aiplatform>=1.80.0",
9
10
  "google-cloud-pipeline-components>=2.19.0",
10
11
  "kfp>=1.4.0",