Refactor: Move Gemini SDK, centralize GCS utils, optimize downloads

src/a2e_avatar.py

@@ -9,7 +9,7 @@ from dataclasses import dataclass
 from enum import Enum
 import os
 from utils import logger
-import gemini_sdk
+from google_src import ai_studio_sdk
 import json_repair
 from data_holder import DataHolder
 from moviepy.editor import AudioFileClip
@@ -369,7 +369,7 @@ Image Prompt: {image_prompt}
 Available Voices: {available_voices}
 Available Avatars with Usage Count: {available_avatar}
 """
-        response = gemini_sdk.generate(model_input)
+        response = ai_studio_sdk.generate(model_input)
 
         response_text = response.strip()
 

src/api_clients.py

@@ -23,12 +23,13 @@ from google.cloud import storage, texttospeech
 import asyncio
 from utils import logger
 from data_holder import DataHolder
+from file_downloader import FileDownloader
 
 
 # --- NEW IMPORTS ---
 from google.oauth2 import service_account
 import vertexai
-from google_src.gcs_utils import get_gcs_client, get_gcs_credentials, upload_file_to_gcs, list_gcs_files
+from google_src.gcs_utils import get_gcs_client, get_gcs_credentials, upload_file_to_gcs, list_gcs_files, create_bucket_if_not_exists
 from google_src.setup_gcs_permissions import setup_bucket_permissions
 
 # --------------------
@@ -36,7 +37,7 @@ from google_src.setup_gcs_permissions import setup_bucket_permissions
 import base64
 from pathlib import Path
 
-import gemini_sdk
+from google_src import ai_studio_sdk
 import uuid
 
 class APIClients:
@@ -53,7 +54,7 @@ class APIClients:
         self.gcs_bucket = self.gcs_client.bucket(gcs_bucket_name)
         
         # Create main bucket if it doesn't exist
-        self.create_bucket_if_not_exists(self.gcs_client, gcs_bucket_name)
+        create_bucket_if_not_exists(self.gcs_client, gcs_bucket_name)
 
         # Apply permissions to the main bucket
         try:
@@ -100,36 +101,6 @@ class APIClients:
             await self.store_in_cache(file_path, f"{method_type}_{duration}", ".txt")
         except: pass
 
-    def _fallback_timed_transcript(self, audio_file_path: str) -> List[Dict]:
-        """Fallback transcript generation when Speech-to-Text fails"""
-        try:
-            from mutagen.mp3 import MP3
-
-            audio = MP3(audio_file_path)
-            duration = audio.info.length
-
-            # Simple word splitting with estimated timing
-            import tempfile
-
-            temp_script = "This is a fallback transcript with estimated timing."
-            words = temp_script.split()
-            word_duration = duration / len(words)
-
-            return [
-                {"word": word, "start_time": i * word_duration, "end_time": (i + 1) * word_duration, "confidence": 0.8}
-                for i, word in enumerate(words)
-            ]
-        except:
-            return []
-
-    def _format_srt_time(self, seconds: float) -> str:
-        """Convert seconds to SRT time format"""
-        hours = int(seconds // 3600)
-        minutes = int((seconds % 3600) // 60)
-        secs = seconds % 60
-        milliseconds = int((secs - int(secs)) * 1000)
-        return f"{hours:02d}:{minutes:02d}:{int(secs):02d},{milliseconds:03d}"
-
     async def generate_image(self, prompt: str) -> Optional[str]:
         """
         Generate image using Vertex AI Imagen 4 Ultra
@@ -149,7 +120,7 @@ class APIClients:
                 return url
 
             logger.info(f"🎨 Generating image with Imagen 4 Ultra: {prompt[:200]}...")
-            image_path = gemini_sdk.generate_image(prompt)
+            image_path = ai_studio_sdk.generate_image(prompt)
             if image_path:
                 await self.store_in_cache(image_path, "generate_image", ".jpg")
                 return image_path
@@ -200,34 +171,6 @@ class APIClients:
             logger.error(f"❌ Imagen 4 Ultra generation failed: {e}")
             raise
 
-    async def generate_captions(self) -> str:
-        """Auto-generate captions from TTS script"""
-        try:
-            logger.info("Auto-generating captions...")
-
-            instruction = f"""Create engaging social media captions for this TikTok/Instagram video script.
-
-    Requirements:
-    - Hook viewers in first 3 words
-    - Use line breaks for readability  
-    - Include relevant emojis naturally
-    - Keep total under 150 characters
-    - Make it scroll-stopping
-
-    Script: {self.data_holder.tts_script}
- 
-    Return ONLY the caption text, nothing else."""
-
-            response = gemini_sdk.generate(instruction)
-
-            caption = response.strip()
-            logger.info(f"✓ Generated caption: {caption[:100]}...")
-            return caption
-
-        except Exception as e:
-            logger.error(f"Caption generation failed: {e}")
-            return self.data_holder.tts_script.split(".")[0][:150]
-
     async def generate_video(self, prompt: str, duration: int, image_input: str = None) -> Dict:
         """
         Generate video using RunwayML gen4_turbo ($0.25 per video / 25 credits)
@@ -308,11 +251,11 @@ class APIClients:
             if os.getenv("USE_GEMIMI_VIDEO", "false").lower() == "true":
                 logger.info("Using Gemini SDK for video generation...")
 
-                output_path = await self.get_cache_url(f"gemini_sdk.generate_video_{model_name}", ".mp4")
+                output_path = await self.get_cache_url(f"ai_studio_sdk.generate_video_{model_name}", ".mp4")
                 if not output_path:
                     output_path = f'/tmp/video_{duration}_{model_name}_{uuid.uuid4().hex[:8]}.mp4'
-                    gemini_sdk.generate_video(prompt, output_path, image_input)
-                    await self.store_in_cache(output_path, f"gemini_sdk.generate_video_{model_name}", ".mp4")
+                    ai_studio_sdk.generate_video(prompt, output_path, image_input)
+                    await self.store_in_cache(output_path, f"ai_studio_sdk.generate_video_{model_name}", ".mp4")
 
                 video_result = {
                     "local_path": output_path,
@@ -394,30 +337,7 @@ class APIClients:
             logger.error(f"❌ Video generation error: {e}")
             raise
 
-    async def download_file(self, url: str, filename: str) -> str:
-        """Download file from URL to local temporary file"""
-        if os.getenv("TEST_AUTOMATION", "").lower() == "true":
-            if not url.startswith("http"):
-                return url
-        import aiohttp
-        import tempfile
-        from pathlib import Path
-
-        local_path = Path(tempfile.gettempdir()) / filename
 
-        try:
-            async with aiohttp.ClientSession() as session:
-                async with session.get(url) as response:
-                    if response.status == 200:
-                        with open(local_path, "wb") as f:
-                            f.write(await response.read())
-                        logger.info(f"✓ Downloaded {filename} from {url}")
-                        return str(local_path)
-                    else:
-                        raise Exception(f"Download failed: {response.status}")
-        except Exception as e:
-            logger.error(f"Failed to download {url}: {e}")
-            raise
 
 
     async def store_in_cache(self, file_path: str, method_type: str, file_ext: str = ".mp4") -> str:
@@ -515,197 +435,30 @@ class APIClients:
 
             logger.info(f"☁️ Found matching file: gs://{self.gcs_bucket.name}/{blob_name}")
 
-            blob = self.gcs_bucket.blob(blob_name)
-            if not blob.exists():
-                logger.error(f"❌ Blob not found in GCS: {blob_name}")
-                return None
-
-            os.makedirs(local_dir, exist_ok=True)
-            local_path = os.path.join(local_dir, os.path.basename(blob_name))
-
-            # Download the blob
-            blob.download_to_filename(local_path)
-            file_size = os.path.getsize(local_path)
-            logger.info(f"✅ Downloaded {blob_name} → {local_path} ({file_size/1024:.1f} KB)")
-
-            return local_path
-
-        except Exception as e:
-            logger.error(f"❌ GCS download failed: {e}")
-            return None
-
-    async def generate_srt_content(self, words: List[Dict]) -> str:
-        """Convert timed words to SRT subtitle format"""
-        srt_content = ""
-        subtitle_index = 1
-
-        # Group words into phrases (3-4 words per subtitle)
-        phrase_words = []
-        current_phrase = []
-
-        for word in words:
-            current_phrase.append(word)
-            if len(current_phrase) >= 3 or any(p in word["word"] for p in [".", ",", "!", "?"]):
-                phrase_words.append(current_phrase)
-                current_phrase = []
-
-        if current_phrase:
-            phrase_words.append(current_phrase)
-
-        # Create SRT entries
-        for phrase in phrase_words:
-            if not phrase:
-                continue
-
-            start_time = phrase[0]["start_time"]
-            end_time = phrase[-1]["end_time"]
-
-            # Format times for SRT (HH:MM:SS,mmm)
-            start_str = self._format_srt_time(start_time)
-            end_str = self._format_srt_time(end_time)
-
-            # Create phrase text
-            phrase_text = " ".join(word["word"] for word in phrase)
-
-            srt_content += f"{subtitle_index}\n"
-            srt_content += f"{start_str} --> {end_str}\n"
-            srt_content += f"{phrase_text}\n\n"
-
-            subtitle_index += 1
-
-        return srt_content
-
-    def _format_srt_time(self, seconds: float) -> str:
-        """Convert seconds to SRT time format"""
-        hours = int(seconds // 3600)
-        minutes = int((seconds % 3600) // 60)
-        secs = seconds % 60
-        milliseconds = int((secs - int(secs)) * 1000)
-        return f"{hours:02d}:{minutes:02d}:{int(secs):02d},{milliseconds:03d}"
-
-
-
-    # ===============================
-    # temp GCS UPLOAD SUPPORT
-    # ===============================
-
-    def init_temp_gcs(self):
-        """
-        One-time initialization of temp GCS client using VERTEX_TEMP_AI_CREDENTIALS_JSON.
-        This sets up a reusable client and bucket reference for later uploads.
-        """
-        try:
-            vertex_creds_json_str = os.getenv("VERTEX_TEMP_AI_CREDENTIALS_JSON")
-            if not vertex_creds_json_str:
-                logger.error("❌ Missing VERTEX_TEMP_AI_CREDENTIALS_JSON for temp GCS initialization")
-                return False
-
-            # Try as file path first, fall back to raw JSON if that fails
-            try:
-                if Path(vertex_creds_json_str).exists():
-                    logger.info(f"Loading GCP credentials from file: {vertex_creds_json_str}")
-                    with open(vertex_creds_json_str, "r") as f:
-                        creds_info = json.load(f)
-                else:
-                    raise FileNotFoundError("Not a valid file path")
-            except Exception:
-                # Fall back to parsing as raw JSON string
-                logger.info("Parsing GCP credentials as JSON string")
-                creds_info = json.loads(vertex_creds_json_str)
-
-            creds = service_account.Credentials.from_service_account_info(creds_info)
-            project_id = creds_info.get("project_id")
-
-            # Allow override or fallback naming
-            my_bucket_name = os.getenv("MY_TEMP_GCS_BUCKET") or f"{project_id}-temp-uploads"
-
-            # Create the storage client
-            self.my_temp_gcs_client = storage.Client(credentials=creds, project=project_id)
-            self.my_temp_gcs_bucket = self.my_temp_gcs_client.bucket(my_bucket_name)
-
-            # Create bucket if it doesn't exist
-            self.create_bucket_if_not_exists(self.my_temp_gcs_client, my_bucket_name)
-
-            logger.info(f"🔐 temp GCS authenticated successfully → gs://{my_bucket_name}")
-            return True
-
-        except Exception as e:
-            logger.error(f"❌ Failed to initialize temp GCS client: {e}")
-            self.my_temp_gcs_client = None
-            self.my_temp_gcs_bucket = None
-            return False
-
-    def create_bucket_if_not_exists(self, client, bucket_name: str, location: str = "us-central1") -> bool:
-        """
-        Create a GCS bucket if it doesn't already exist.
-        
-        Args:
-            client: GCS storage client to use
-            bucket_name: Name of the bucket to create
-            location: GCS location for the bucket (default: us-central1)
+            # Construct GCS URL
+            gs_url = f"gs://{self.gcs_bucket.name}/{blob_name}"
             
-        Returns:
-            True if bucket exists or was created successfully, False otherwise
-        """
-        try:
-            bucket = client.bucket(bucket_name)
+            # Use FileDownloader
+            downloader = FileDownloader()
+            # We can download mostly anywhere, but let's stick to the default behavior or local_dir if needed.
+            # safe_download uses temp dir by default. The original code used local_dir="/tmp/tts_downloads".
+            # FileDownloader uses its own temp dir logic.
+            # Let's see if we can pass output path. safe_download takes output_path.
             
-            if bucket.exists():
-                logger.info(f"✓ Bucket already exists: gs://{bucket_name}")
-                return True
+            local_path = os.path.join(local_dir, os.path.basename(blob_name))
             
-            # Bucket doesn't exist, create it
-            logger.info(f"📦 Creating new bucket: gs://{bucket_name} in {location}")
-            new_bucket = client.create_bucket(bucket_name, location=location)
-            logger.info(f"✅ Bucket created successfully: gs://{new_bucket.name}")
-            return True
+            downloaded = downloader.safe_download(gs_url, output_path=local_path)
             
-        except Exception as e:
-            logger.error(f"❌ Failed to create bucket {bucket_name}: {e}")
-            return False
-
-
-    async def upload_to_temp_gcs(self, file_path: str, folder: str = "video") -> Optional[str]:
-        """
-        Upload MP4 file to the temp GCS bucket initialized by init_temp_gcs().
-
-        Args:
-            file_path: Path to the .mp4 file
-            folder: Folder name in GCS (default: 'videos')
-
-        Returns:
-            Public URL of uploaded video, or None on failure.
-        """
-        try:
-            # Check initialization first
-            if not getattr(self, "my_temp_gcs_bucket", None):
-                logger.error("⚠️ temp GCS not initialized. Call init_temp_gcs() first.")
-                return None
-
-            if not os.path.exists(file_path):
-                logger.error(f"❌ File not found: {file_path}")
-                return None
-
-            filename = f'{self.data_holder.hash_tts_script}_{os.path.basename(file_path)}'
-            if not filename.endswith(".mp4"):
-                logger.warning("⚠️ Skipping non-MP4 file upload.")
-                return None
-
-            blob_name = f"{folder}/{filename}"
-            blob = self.my_temp_gcs_bucket.blob(blob_name)
-            blob.content_type = "video/mp4"
-
-            logger.info(f"☁️ Uploading {filename} → gs://{self.my_temp_gcs_bucket.name}/{blob_name}")
-            blob.upload_from_filename(file_path)
-
-            # Optional: make it public (comment out if you want private URLs)
-            # blob.make_public()
-            public_url = f"https://storage.googleapis.com/{self.my_temp_gcs_bucket.name}/{blob_name}"
-
-            logger.info(f"✅ temp GCS upload complete: {public_url}")
-            return public_url
+            if downloaded:
+                file_size = os.path.getsize(downloaded)
+                logger.info(f"✅ Downloaded {blob_name} → {downloaded} ({file_size/1024:.1f} KB)")
+                return str(downloaded)
+            else:
+                 logger.error(f"❌ FileDownloader failed for {gs_url}")
+                 return None
 
         except Exception as e:
-            logger.error(f"❌ temp GCS upload failed: {e}")
+            logger.error(f"❌ GCS download failed: {e}")
             return None
 
+

src/asset_manager/asset_processor.py

@@ -10,8 +10,7 @@ from typing import List, Dict, Optional, Tuple
 import pandas as pd
 import json_repair
 from moviepy.editor import VideoFileClip
-
-import gemini_sdk
+from google_src import ai_studio_sdk
 from utils import logger
 from .video_lib import get_video_lib
 
@@ -86,7 +85,7 @@ USER PROMPT:
 TTS Script: {tts_script}
 Video Options: {video_context}
 """
-            response = gemini_sdk.generate(model_input)
+            response = ai_studio_sdk.generate(model_input)
 
             response_text = response.strip()
 

src/automation.py

@@ -6,7 +6,7 @@ import asyncio
 import os
 import time
 import json
-import gemini_sdk
+from google_src import ai_studio_sdk
 from typing import Dict, List, Optional, Any
 from pathlib import Path
 from api_clients import APIClients
@@ -467,9 +467,6 @@ class ContentAutomation:
             if strategy.get("captions"):
                 captions = strategy["captions"]
                 logger.info(f"Using provided captions: {captions[:50]}...")
-            else:
-                captions = await self.api_clients.generate_captions()
-                logger.info(f"Auto-generated captions: {captions[:50]}...")
 
             # Step 1: Generate image using Imagen 4 Ultra
             image_path = await self.api_clients.generate_image(strategy["gemini_prompt"])
@@ -484,7 +481,7 @@ class ContentAutomation:
                 prompt=strategy["runway_prompt"], image_input=image_path, duration=strategy.get("duration", 3)
             )
 
-            video_data["captions"] = captions
+            video_data["captions"] = captions or "No captions provided"
             video_data["script"] = self.data_holder.tts_script
 
             if os.getenv("USE_VEO", "false").lower() == "true":
@@ -524,7 +521,7 @@ class ContentAutomation:
         "style": "casual"
     }}"""
 
-            response = gemini_sdk.generate(analysis_prompt)
+            response = ai_studio_sdk.generate(analysis_prompt)
 
             # Parse response
             response_text = response.strip()

src/generate_content.py

@@ -6,8 +6,7 @@ from datetime import datetime
 from pathlib import Path
 from dotenv import load_dotenv
 import argparse
-
-import gemini_sdk
+from google_src import ai_studio_sdk
 
 class GenerateContent:
     def __init__(self, commit: bool = False):
@@ -59,7 +58,7 @@ class GenerateContent:
         """Generate 10 unique scenarios using Gemini and save to CSV."""
         print("🤖 Generating scenarios...")
 
-        model = gemini_sdk.generate("gemini-2.5-pro")
+        model = ai_studio_sdk.generate("gemini-2.5-pro")
         model_input = f"""SYSTEM INSTRUCTION::
 {self._get_system_instruction()}
 

src/{gemini_sdk.py → google_src/ai_studio_sdk.py}

@@ -250,4 +250,4 @@ if __name__ == "__main__":
         "Morning violence against self visible, broken NPC movement, glitchy character lighting, camera shows malfunction, cinematic realistic style",
         "test.mp4",
         image="testData/image.png"
-    )
+    )

src/google_src/gcs_utils.py

@@ -203,3 +203,24 @@ def list_gcs_files(prefix: str = "video/", account_name: str = "final_data") ->
     except Exception as e:
         logger.error(f"❌ Failed to list files: {e}")
         return []
+
+def create_bucket_if_not_exists(client, bucket_name: str, location: str = "us-central1") -> bool:
+    """
+    Create a GCS bucket if it doesn't already exist.
+    """
+    try:
+        bucket = client.bucket(bucket_name)
+        
+        if bucket.exists():
+            logger.info(f"✓ Bucket already exists: gs://{bucket_name}")
+            return True
+        
+        # Bucket doesn't exist, create it
+        logger.info(f"📦 Creating new bucket: gs://{bucket_name} in {location}")
+        new_bucket = client.create_bucket(bucket_name, location=location)
+        logger.info(f"✅ Bucket created successfully: gs://{new_bucket.name}")
+        return True
+        
+    except Exception as e:
+        logger.error(f"❌ Failed to create bucket {bucket_name}: {e}")
+        return False