playbook/antigravity-awesome-skills/skills/gemini-omni-flash-api/scripts/upload_file.py

#!/usr/bin/env python3
"""
Uploads a file to the Gemini Files API and waits for it to become ACTIVE.
Uses the official google-genai SDK.
"""

import argparse
import json
import mimetypes
import os
import sys
import time
from google import genai
from google.genai import types

def get_api_key(args):
    """Retrieves API key from command args or environment."""
    if args.api_key:
        return args.api_key
    return os.environ.get("GEMINI_API_KEY")

def detect_mime_type(file_path):
    """Determines MIME type based on file extension, falling back to standard mimetypes module."""
    ext = os.path.splitext(file_path)[1].lower()
    mime_map = {
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".webp": "image/webp",
        ".mp4": "video/mp4",
        ".mp3": "audio/mpeg",
        ".wav": "audio/wav",
        ".pdf": "application/pdf",
        ".txt": "text/plain",
    }

    if ext in mime_map:
        return mime_map[ext]

    mime_type, _ = mimetypes.guess_type(file_path)
    if mime_type:
        return mime_type

    return "application/octet-stream"

def upload_file(file_path, display_name=None, api_key=None):
    """Performs an upload using google-genai SDK, with automatic pre-processing for large videos."""
    file_size = os.path.getsize(file_path)
    mime_type = detect_mime_type(file_path)

    # Large video file size check (>25MB)
    is_video = mime_type.startswith("video/")
    if is_video and file_size > 25 * 1024 * 1024:
        size_mb = file_size / (1024 * 1024)
        print(f"\nWARNING: Video file '{file_path}' is very large ({size_mb:.2f} MB)!")
        print("Note: Gemini Omni Flash is optimized for 10s videos at 720p and 24fps. Uploading very large or")
        print("high-resolution videos will significantly increase upload times and may cause Out-Of-Memory (OOM) errors.")

        # Determine if terminal is interactive
        if sys.stdin.isatty():
            print("\nWould you like to automatically pre-process this video first using prep_video.py?")
            print("This will trim, scale, and optimize the video to ensure a fast, OOM-safe upload.")
            try:
                choice = input("Pre-process video? [Y/n]: ").strip().lower()
                if choice in ("", "y", "yes"):
                    prepped_output_path = os.path.join("media", f"prepped_{os.path.basename(file_path)}")
                    os.makedirs("media", exist_ok=True)

                    # Resolve prep_video.py script path
                    import subprocess
                    prep_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "video", "prep_video.py")
                    if not os.path.exists(prep_script):
                        prep_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "prep_video.py")

                    cmd = [sys.executable, prep_script, file_path, "--output", prepped_output_path]
                    print(f"Running: {' '.join(cmd)}")

                    try:
                        result = subprocess.run(cmd)
                        if result.returncode == 0 and os.path.exists(prepped_output_path):
                            file_path = prepped_output_path
                            file_size = os.path.getsize(file_path)
                            print(f"\nPre-processing completed successfully! Proceeding with upload of prepped video ({file_size / (1024*1024):.2f} MB)...")
                        else:
                            raise RuntimeError("Error: Video pre-processing failed. Proceeding with original file upload is not recommended.")
                    except Exception as e:
                        raise RuntimeError(f"Error executing prep_video.py: {e}")
                else:
                    proceed_choice = input("Do you want to proceed with uploading the original large video anyway? [y/N]: ").strip().lower()
                    if proceed_choice not in ("y", "yes"):
                        raise RuntimeError("Upload cancelled by user. Please pre-process the video manually first.")
            except (KeyboardInterrupt, EOFError):
                raise RuntimeError("\nNo input received. Upload cancelled to prevent OOM.")
        else:
            # Non-interactive mode
            if file_size > 100 * 1024 * 1024: # Block files larger than 100MB in non-interactive mode
                err_msg = (
                    f"Error: Video file is extremely large ({size_mb:.2f} MB) and script is running in non-interactive mode.\n"
                    "To prevent Out-Of-Memory (OOM) errors, upload has been blocked.\n"
                    "Please pre-process the video first using prep_video.py."
                )
                raise RuntimeError(err_msg)
            else:
                print("Proceeding with upload in non-interactive mode...", file=sys.stderr)

    if not display_name:
        display_name = os.path.basename(file_path)

    print(f"Preparing upload of '{file_path}' ({file_size} bytes, type: {mime_type})...")

    # Step 1: Initialize Client
    client = genai.Client(api_key=api_key)

    # Step 2: Upload file using SDK
    print("Uploading file bytes using google-genai SDK...")
    try:
        config = types.UploadFileConfig(
            display_name=display_name,
            mime_type=mime_type,
        )
        file_obj = client.files.upload(file=file_path, config=config)
        # Convert Pydantic File model to dictionary with both camelCase and snake_case keys for compatibility
        file_dict = json.loads(file_obj.model_dump_json())
        # Add camelCase field for mimeType
        if "mime_type" in file_dict:
            file_dict["mimeType"] = file_dict["mime_type"]
        return file_dict
    except Exception as e:
        raise RuntimeError(f"Error uploading file via SDK: {e}")

def wait_for_active(file_name, api_key, poll_interval=3, max_attempts=30, backoff_factor=1.5, max_interval=30):
    """Polls the file status until state is ACTIVE or FAILED using exponential backoff via the SDK."""
    print(f"Waiting for file {file_name} to finish processing...")

    client = genai.Client(api_key=api_key)
    attempt = 0
    current_interval = poll_interval
    consecutive_errors = 0
    max_consecutive_errors = 5

    while attempt < max_attempts:
        try:
            file_obj = client.files.get(name=file_name)
            state = file_obj.state

            # Extract state name as a string (handling both Enum and String)
            state_str = state.name if hasattr(state, "name") else str(state)

            # Reset consecutive errors on successful API response
            consecutive_errors = 0

            if state_str == "ACTIVE":
                print("File is ACTIVE and ready for generations!")
                file_dict = json.loads(file_obj.model_dump_json())
                if "mime_type" in file_dict:
                    file_dict["mimeType"] = file_dict["mime_type"]
                return file_dict
            elif state_str == "FAILED":
                raise RuntimeError("Error: File processing failed on the backend.")

            print(f"Current state: {state_str}. Retrying in {current_interval:.1f}s...")
            time.sleep(current_interval)

            # Increase interval for the next poll (backoff)
            current_interval = min(current_interval * backoff_factor, max_interval)
            attempt += 1

        except Exception as e:
            consecutive_errors += 1
            if consecutive_errors >= max_consecutive_errors:
                raise RuntimeError(f"Error: Too many consecutive errors checking status ({e}). Exiting.")

            print(f"Warning: Error checking status ({e}). Retrying in {current_interval:.1f}s...")
            time.sleep(current_interval)
            current_interval = min(current_interval * backoff_factor, max_interval)
            attempt += 1

    raise RuntimeError(f"Error: Maximum polling attempts ({max_attempts}) reached. File is still not ACTIVE.")

def main():
    parser = argparse.ArgumentParser(description="Upload files to Gemini Files API using google-genai SDK.")
    parser.add_argument("file", help="Path to the file to upload")
    parser.add_argument("--name", help="Custom display name for the file")
    parser.add_argument("--api-key", help="Gemini API Key (overrides env)")
    parser.add_argument("--no-wait", action="store_true", help="Don't wait for ACTIVE status")

    args = parser.parse_args()

    api_key = get_api_key(args)
    if not api_key:
        print("Error: API key is not set. Use --api-key or set GEMINI_API_KEY environment variable.", file=sys.stderr)
        sys.exit(1)

    if not os.path.exists(args.file):
        print(f"Error: File '{args.file}' not found.", file=sys.stderr)
        sys.exit(1)

    try:
        file_meta = upload_file(args.file, args.name, api_key)
        file_name = file_meta.get("name")

        print(f"File metadata created:")
        print(f"  Name: {file_name}")
        print(f"  URI:  {file_meta.get('uri')}")
        print(f"  Type: {file_meta.get('mimeType')}")

        if not args.no_wait:
            file_meta = wait_for_active(file_name, api_key)

        print("\nFile upload successfully completed! JSON Output:")
        print(json.dumps(file_meta, indent=2))
        sys.exit(0)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()