Refomatted the sample with black and flake8

Author: happyhuman

python/agents/product-catalog-ad-generation/agent_engine_deploy.py

@@ -35,7 +35,9 @@
 flags.DEFINE_string("project_id", None, "GCP project ID.")
 flags.DEFINE_string("location", None, "GCP location.")
 flags.DEFINE_string("bucket", None, "GCP Cloud Storage bucket for staging.")
-flags.DEFINE_string("resource_id", None, "ReasoningEngine resource ID for deletion.")
+flags.DEFINE_string(
+    "resource_id", None, "ReasoningEngine resource ID for deletion."
+)
 
 flags.DEFINE_bool("list", False, "List all deployed agent engines.")
 flags.DEFINE_bool("create", False, "Creates a new agent engine.")
@@ -263,14 +265,22 @@ def main(argv: list[str]) -> None:
     del argv  # unused
 
     project_id = (
-        FLAGS.project_id if FLAGS.project_id else os.getenv("GOOGLE_CLOUD_PROJECT_ID")
+        FLAGS.project_id
+        if FLAGS.project_id
+        else os.getenv("GOOGLE_CLOUD_PROJECT_ID")
     )
     # --- THIS LINE IS THE FIX ---
     # Corrected the typo from GOOGLE_CLOUD_LOCATION_REGION_REGION to GOOGLE_CLOUD_LOCATION_REGION
     location = (
-        FLAGS.location if FLAGS.location else os.getenv("GOOGLE_CLOUD_LOCATION_REGION")
+        FLAGS.location
+        if FLAGS.location
+        else os.getenv("GOOGLE_CLOUD_LOCATION_REGION")
+    )
+    bucket = (
+        FLAGS.bucket
+        if FLAGS.bucket
+        else os.getenv("GOOGLE_CLOUD_STORAGE_BUCKET")
     )
-    bucket = FLAGS.bucket if FLAGS.bucket else os.getenv("GOOGLE_CLOUD_STORAGE_BUCKET")
 
     logging.info("Using Project ID: %s", project_id)
     logging.info("Using Location:   %s", location)

python/agents/product-catalog-ad-generation/content_gen_agent/agent.py

@@ -22,7 +22,9 @@
 
 from google.adk.agents import Agent
 from google.adk.apps import App
-from google.adk.plugins.save_files_as_artifacts_plugin import SaveFilesAsArtifactsPlugin
+from google.adk.plugins.save_files_as_artifacts_plugin import (
+    SaveFilesAsArtifactsPlugin,
+)
 from google.adk.tools import FunctionTool, load_artifacts
 
 from .func_tools.combine_video import combine

python/agents/product-catalog-ad-generation/content_gen_agent/func_tools/combine_video.py

@@ -55,7 +55,9 @@ def _get_storyline_schema(num_images: int) -> List[Dict]:
 
     schema = []
     if num_images > 1:
-        schema.append({"name": "before", "generate": True, "step": 0, "duration": 3})
+        schema.append(
+            {"name": "before", "generate": True, "step": 0, "duration": 3}
+        )
 
     for i in range(num_images - 2):
         schema.append(
@@ -106,7 +108,9 @@ def _upload_to_gcs(video_bytes: bytes, filename: str) -> Optional[str]:
             logging.error("GCP_PROJECT environment variable not set.")
             return None
 
-        bucket_name = os.getenv("GCS_BUCKET") or f"{project_id}-contentgen-static"
+        bucket_name = (
+            os.getenv("GCS_BUCKET") or f"{project_id}-contentgen-static"
+        )
         folder_path = _get_datetime_folder_path()
         blob_name = f"{folder_path}{filename}"
 
@@ -132,7 +136,9 @@ async def _load_single_clip(
     """Loads and processes a single video clip artifact."""
     try:
         artifact = await tool_context.load_artifact(filename)
-        if not (artifact and artifact.inline_data and artifact.inline_data.data):
+        if not (
+            artifact and artifact.inline_data and artifact.inline_data.data
+        ):
             logging.warning("Could not load artifact data for %s.", filename)
             return None
 
@@ -184,7 +190,9 @@ async def _load_and_process_video_clips(
             logging.warning("Skipping video file with missing filename.")
             continue
 
-        result = await _load_single_clip(filename, tool_context, temp_dir, storyline)
+        result = await _load_single_clip(
+            filename, tool_context, temp_dir, storyline
+        )
         if result:
             clip, temp_path = result
             video_clips.append(clip)
@@ -225,8 +233,14 @@ async def _load_and_process_audio_clips(
         # Voiceover
         if voiceover_file:
             vo_artifact = await tool_context.load_artifact(voiceover_file)
-            if vo_artifact and vo_artifact.inline_data and vo_artifact.inline_data.data:
-                temp_path = os.path.join(temp_dir, os.path.basename(voiceover_file))
+            if (
+                vo_artifact
+                and vo_artifact.inline_data
+                and vo_artifact.inline_data.data
+            ):
+                temp_path = os.path.join(
+                    temp_dir, os.path.basename(voiceover_file)
+                )
                 with open(temp_path, "wb") as f:
                     f.write(vo_artifact.inline_data.data)
                 vo_clip = AudioFileClip(temp_path)
@@ -268,7 +282,9 @@ async def _combine_and_upload_video(
         gcs_uri = _upload_to_gcs(video_bytes, filename)
         await tool_context.save_artifact(
             filename,
-            genai.types.Part.from_bytes(data=video_bytes, mime_type="video/mp4"),
+            genai.types.Part.from_bytes(
+                data=video_bytes, mime_type="video/mp4"
+            ),
         )
 
         result = {"name": filename}
@@ -321,7 +337,11 @@ async def combine(
         try:
             final_clip = concatenate_videoclips(video_clips, method="compose")
             final_clip.audio = await _load_and_process_audio_clips(
-                audio_file, voiceover_file, final_clip.duration, tool_context, temp_dir
+                audio_file,
+                voiceover_file,
+                final_clip.duration,
+                tool_context,
+                temp_dir,
             )
 
             filename = f"combined_video_{int(time.time())}.mp4"
@@ -334,7 +354,9 @@ async def combine(
             gcs_uri = _upload_to_gcs(video_bytes, filename)
             await tool_context.save_artifact(
                 filename,
-                genai.types.Part.from_bytes(data=video_bytes, mime_type="video/mp4"),
+                genai.types.Part.from_bytes(
+                    data=video_bytes, mime_type="video/mp4"
+                ),
             )
 
             result = {"name": filename}
@@ -343,7 +365,8 @@ async def combine(
             return result
         except Exception as e:
             logging.error(
-                f"An error occurred during video combination: {e}", exc_info=True
+                f"An error occurred during video combination: {e}",
+                exc_info=True,
             )
             return None
         finally:

python/agents/product-catalog-ad-generation/content_gen_agent/func_tools/generate_audio.py

@@ -19,13 +19,13 @@
 import os
 import re  # noqa: F401
 import time
-from typing import Any, Dict, List, Optional  # noqa: F401
+import random
+from typing import Dict, List, Optional, Union
 
 import aiohttp
 import google.auth
 import google.auth.transport.requests
 from google.adk.tools import ToolContext
-from google.api_core.exceptions import GoogleAPICallError
 from google.cloud import texttospeech
 from google.genai import types
 
@@ -39,11 +39,17 @@
 TTS_MODEL_NAME = "gemini-2.5-flash-preview-tts"
 TTS_VOICE_NAME = "Schedar"
 LYRIA_MODEL_NAME = "lyria-002"
+LYRIA_LOCATION = os.environ.get(
+    "LYRIA_LOCATION", "us-central1"
+)  # Default to us-central1
+LYRIA_MODEL_ID = os.environ.get("LYRIA_MODEL_ID", LYRIA_MODEL_NAME)
 
 
 async def _send_google_api_request(
     api_endpoint: str,
-    data: Optional[Dict[str, Union[List[Dict[str, str]], Dict[str, None]]]] = None,
+    data: Optional[
+        Dict[str, Union[List[Dict[str, str]], Dict[str, None]]]
+    ] = None,
 ) -> Optional[Dict[str, List[Dict[str, str]]]]:
     """Sends an authenticated HTTP request to a Google API endpoint.
 
@@ -74,7 +80,10 @@ async def _send_google_api_request(
                         response.raise_for_status()
                         return await response.json()
                 except aiohttp.ClientResponseError as e:
-                    if e.status in [400, 429, 500, 503] and attempt < MAX_RETRIES - 1:
+                    if (
+                        e.status in [400, 429, 500, 503]
+                        and attempt < MAX_RETRIES - 1
+                    ):
                         wait_time = (2**attempt) + (random.uniform(0, 1))
                         logging.warning(
                             "Attempt %s/%s failed with status %s. Retrying in %.2f seconds...",
@@ -152,15 +161,19 @@ async def generate_audio(
                 MAX_RETRIES,
             )
             async with aiohttp.ClientSession() as session:
-                async with session.post(url, headers=headers, json=payload) as resp:
+                async with session.post(
+                    url, headers=headers, json=payload
+                ) as resp:
                     if resp.status == 200:
                         data = await resp.json()
                         predictions = data.get("predictions")
                         if (
                             not predictions
                             or "bytesBase64Encoded" not in predictions[0]
                         ):
-                            logging.warning("No audioContent in Lyria response.")
+                            logging.warning(
+                                "No audioContent in Lyria response."
+                            )
                             raise ValueError("Invalid response format")
 
                         audio_b64 = predictions[0]["bytesBase64Encoded"]
@@ -222,7 +235,9 @@ async def generate_audio(
     return {"name": STATIC_AUDIO_FALLBACK}
 
 
-async def _generate_voiceover_content(prompt: str, text: str) -> Optional[bytes]:
+async def _generate_voiceover_content(
+    prompt: str, text: str
+) -> Optional[bytes]:
     """Synthesizes speech using Gemini-TTS.
 
     Args:
@@ -248,7 +263,9 @@ async def _generate_voiceover_content(prompt: str, text: str) -> Optional[bytes]
         )
         return response.audio_content
     except Exception as e:
-        logging.error(f"Failed to generate voiceover content: {e}", exc_info=True)
+        logging.error(
+            f"Failed to generate voiceover content: {e}", exc_info=True
+        )
         return None
 
 
@@ -339,7 +356,9 @@ async def generate_audio_and_voiceover(
     if generation_mode in ["audio", "both"]:
         audio_res = results[result_index]
         if isinstance(audio_res, Exception) or not audio_res:
-            response["failures"].append(f"audio: {audio_res or 'Unknown error'}")
+            response["failures"].append(
+                f"audio: {audio_res or 'Unknown error'}"
+            )
             response["audio_name"] = STATIC_AUDIO_FALLBACK
         else:
             response["audio_name"] = audio_res["name"]
@@ -348,7 +367,9 @@ async def generate_audio_and_voiceover(
     if generation_mode in ["voiceover", "both"]:
         vo_res = results[result_index]
         if isinstance(vo_res, Exception) or not vo_res:
-            response["failures"].append(f"voiceover: {vo_res or 'Unknown error'}")
+            response["failures"].append(
+                f"voiceover: {vo_res or 'Unknown error'}"
+            )
         else:
             response["voiceover_name"] = vo_res["name"]
 

python/agents/product-catalog-ad-generation/content_gen_agent/func_tools/generate_image.py

@@ -17,9 +17,9 @@
 import json
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Awaitable
+
 
-import vertexai
 from content_gen_agent.utils.evaluate_media import (
     calculate_evaluation_score,
     evaluate_media,
@@ -30,7 +30,6 @@
 from google.adk.tools import ToolContext
 from google.cloud import storage
 from google.genai import types
-from google.genai.types import HarmBlockThreshold, HarmCategory
 from vertexai.preview.vision_models import ImageGenerationModel
 
 # --- Configuration ---
@@ -43,6 +42,17 @@
 GCP_PROJECT = os.getenv("GCP_PROJECT")
 
 IMAGE_GEN_MODEL_GEMINI = "gemini-3-pro-image-preview"
+IMAGE_GEN_MODEL_IMAGEN = "imagen-v4"
+
+# Configuration for Gemini image generation
+GENERATE_CONTENT_CONFIG = genai.types.GenerateContentConfig(
+    temperature=0.4,
+    top_p=0.8,
+    top_k=40,
+)
+
+# Type alias for image generation results
+ImageGenerationResult = Dict[str, Any]
 
 MAX_RETRIES = 3
 ASSET_SHEET_FILENAME = "asset_sheet.png"
@@ -189,14 +199,18 @@ async def generate_one_image(
         }
 
     contents = [prompt, *input_images]
-    tasks = [_call_gemini_image_api(contents, prompt) for _ in range(MAX_RETRIES)]
+    tasks = [
+        _call_gemini_image_api(contents, prompt) for _ in range(MAX_RETRIES)
+    ]
     results = await asyncio.gather(*tasks)
     successful_attempts = [res for res in results if res]
 
     if not successful_attempts:
         return {
             "status": "failed",
-            "detail": (f"All image generation attempts failed for prompt: '{prompt}'."),
+            "detail": (
+                f"All image generation attempts failed for prompt: '{prompt}'."
+            ),
         }
 
     best_attempt = max(
@@ -233,7 +247,9 @@ async def _save_generated_images(
             save_tasks.append(
                 tool_context.save_artifact(
                     filename,
-                    types.Part.from_bytes(data=image_bytes, mime_type=IMAGE_MIME_TYPE),
+                    types.Part.from_bytes(
+                        data=image_bytes, mime_type=IMAGE_MIME_TYPE
+                    ),
                 )
             )
             result["detail"] = f"Image stored as {filename}."
@@ -253,9 +269,7 @@ def _create_image_generation_task(
     """Creates a task for generating a single image."""
     filename_prefix = f"{scene_num}_"
     if is_logo_scene:
-        logo_prompt = (
-            f"Place the company logo centered on the following background: {prompt}"
-        )
+        logo_prompt = f"Place the company logo centered on the following background: {prompt}"
         return generate_one_image(logo_prompt, [logo_image], filename_prefix)
 
     return generate_one_image(prompt, [asset_sheet_image], filename_prefix)
@@ -312,7 +326,9 @@ async def generate_images_from_storyline(
     """
     if not client:
         return [
-            json.dumps({"status": "failed", "detail": "Gemini client not initialized."})
+            json.dumps(
+                {"status": "failed", "detail": "Gemini client not initialized."}
+            )
         ]
 
     storage_client = storage.Client()
@@ -328,7 +344,9 @@ async def generate_images_from_storyline(
         ]
 
     try:
-        asset_sheet_image = await tool_context.load_artifact(ASSET_SHEET_FILENAME)
+        asset_sheet_image = await tool_context.load_artifact(
+            ASSET_SHEET_FILENAME
+        )
         if not asset_sheet_image:
             raise ValueError("Asset sheet artifact is empty.")
     except Exception as e:
@@ -355,10 +373,10 @@ async def generate_images_from_storyline(
         filename_prefix = f"{scene_num}_"
 
         if is_logo_scene:
-            logo_prompt = (
-                f"Place the company logo centered on the following background: {prompt}"
+            logo_prompt = f"Place the company logo centered on the following background: {prompt}"
+            tasks.append(
+                generate_one_image(logo_prompt, [logo_image], filename_prefix)
             )
-            tasks.append(generate_one_image(logo_prompt, [logo_image], filename_prefix))
         else:
             tasks.append(
                 generate_one_image(prompt, [asset_sheet_image], filename_prefix)
@@ -375,7 +393,9 @@ async def generate_images_from_storyline(
             save_tasks.append(
                 tool_context.save_artifact(
                     filename,
-                    types.Part.from_bytes(data=image_bytes, mime_type=IMAGE_MIME_TYPE),
+                    types.Part.from_bytes(
+                        data=image_bytes, mime_type=IMAGE_MIME_TYPE
+                    ),
                 )
             )
             result["detail"] = f"Image stored as {filename}."