"""Gemini API integration for video segment extraction.""" import json import tempfile from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from google import genai from google.genai import types from .config import settings from .ffmpeg_utils import downsample_video, probe_video from .models import PerformanceWindow, SegmentsResponse, VideoSegment, VideoSource class IntelligenceEngine: """Handles video analysis using Gemini API.""" def __init__(self) -> None: """Initialize Gemini client.""" self.client = genai.Client(api_key=settings.gemini_api_key) self.model_name = settings.gemini_model def extract_segments( self, video_path: Path, num_segments: int = 3, min_duration: float | None = None, max_duration: float | None = None, extra_videos: list[Path] | None = None, video_source: VideoSource = VideoSource.ALL, ) -> list[VideoSegment]: """Extract high-energy segments from one or more videos using Gemini API. When extra_videos are provided and video_source includes them, all videos are uploaded in parallel and each gets its own Gemini query. Extra-video segments use a UGC-aware prompt (audio-agnostic) and are tagged with their source_video path. The top num_segments by score are returned. Returns: List of VideoSegment objects tagged with source_video (None = main video). """ extra = extra_videos or [] if video_source == VideoSource.ALL: source_videos = [video_path] + extra elif video_source == VideoSource.MAIN_ONLY: source_videos = [video_path] else: # EXTRAS_ONLY if not extra: raise ValueError("video_source='extras' requires at least one extra_video") source_videos = extra # All videos to upload (main always needed for performance-window query) all_to_upload = list(dict.fromkeys([video_path] + extra)) uploaded: dict[Path, object] = {} metadata_map: dict[Path, object] = {} temp_paths: list[Path] = [] try: # Phase A: downsample + upload all videos in parallel with ThreadPoolExecutor(max_workers=len(all_to_upload)) as executor: upload_futures = { executor.submit(self._prepare_and_upload, vp): vp for vp in all_to_upload } for future in as_completed(upload_futures): vp = upload_futures[future] exc = future.exception() if exc is not None: raise RuntimeError( f"Failed to prepare/upload {vp.name}: {exc}" ) from exc video_file, meta, temp_path = future.result() uploaded[vp] = video_file metadata_map[vp] = meta temp_paths.append(temp_path) main_file = uploaded[video_path] main_meta = metadata_map[video_path] generation_config = types.GenerateContentConfig( response_mime_type="application/json", response_schema=SegmentsResponse, media_resolution=types.MediaResolution.MEDIA_RESOLUTION_LOW, ) # Phase B: run all Gemini queries concurrently # Main video: segments query + performance window query # Each extra video (in source_videos): UGC segments query def _run_main_segments() -> SegmentsResponse: prompt = self._build_main_prompt(num_segments, min_duration, max_duration, main_meta.duration) resp = self.client.models.generate_content( model=self.model_name, contents=[main_file, prompt], config=generation_config, ) return SegmentsResponse(**json.loads(resp.text)) def _run_ugc_segments(vp: Path) -> SegmentsResponse: meta = metadata_map[vp] prompt = self._build_ugc_prompt(num_segments, min_duration, max_duration, meta.duration) resp = self.client.models.generate_content( model=self.model_name, contents=[uploaded[vp], prompt], config=generation_config, ) return SegmentsResponse(**json.loads(resp.text)) n_workers = 2 + len([v for v in source_videos if v != video_path]) with ThreadPoolExecutor(max_workers=max(n_workers, 2)) as executor: futures: dict = {} if video_path in source_videos: futures[executor.submit(_run_main_segments)] = ("main_segments", video_path) futures[executor.submit( self._query_performance_window, main_file, main_meta.duration )] = ("window", video_path) for vp in source_videos: if vp != video_path: futures[executor.submit(_run_ugc_segments, vp)] = ("ugc_segments", vp) results: dict = {"window": None, "segments_by_video": {}} for future in as_completed(futures): tag, vp = futures[future] exc = future.exception() if tag == "window": results["window"] = ( PerformanceWindow( performance_start=0.0, performance_end=main_meta.duration, ) if exc is not None else future.result() ) elif tag == "main_segments": if exc is not None: raise RuntimeError( f"Gemini segments query failed: {exc}" ) from exc results["segments_by_video"][vp] = future.result() else: # ugc_segments if exc is not None: raise RuntimeError( f"Gemini UGC segments query failed for {vp.name}: {exc}" ) from exc results["segments_by_video"][vp] = future.result() window: PerformanceWindow = results["window"] # Clean up uploaded files for video_file in uploaded.values(): try: self.client.files.delete(name=video_file.name) except Exception: pass # Build tagged segment list all_segments: list[VideoSegment] = [] for vp, segs_response in results["segments_by_video"].items(): meta = metadata_map[vp] validated = self._validate_segments( segs_response.segments, meta.duration, min_duration, max_duration ) # Performance window filtering only applies to main video if vp == video_path: validated = [ s for s in validated if s.start >= window.performance_start and s.end <= window.performance_end ] # Tag with source video (None means main) for seg in validated: all_segments.append(seg.model_copy(update={ "source_video": None if vp == video_path else vp })) # Return top num_segments by score all_segments.sort(key=lambda s: s.score, reverse=True) return all_segments[:num_segments] finally: for temp_path in temp_paths: if temp_path.exists(): temp_path.unlink() def _prepare_and_upload(self, video_path: Path) -> tuple: """Downsample a video, upload to Gemini Files API, return (file, metadata, temp_path).""" import time with tempfile.NamedTemporaryFile(suffix="_downsampled.mp4", delete=False) as tmp: temp_path = Path(tmp.name) downsampled_path = downsample_video(video_path, temp_path) metadata = probe_video(video_path) video_file = self.client.files.upload(file=str(downsampled_path)) while video_file.state.name == "PROCESSING": time.sleep(1) video_file = self.client.files.get(name=video_file.name) if video_file.state.name != "ACTIVE": raise RuntimeError(f"Video processing failed: {video_file.state.name}") return video_file, metadata, temp_path def _build_main_prompt( self, num_segments: int, min_duration: float | None, max_duration: float | None, video_duration: float, ) -> str: if min_duration is not None and max_duration is not None: duration_line = f"- Clip duration: {min_duration}–{max_duration} seconds each" else: duration_line = ( "- Clip duration: choose appropriate lengths at your discretion, " "maximum 30 seconds each" ) return f"""You are selecting clips from a music video to use as paid marketing creatives on TikTok and Instagram Reels (vertical 9:16, 10–15 seconds each). GOAL: Return EXACTLY {num_segments} non-overlapping segments most likely to make a viewer stop scrolling, recognize the track, and feel compelled to engage. AUDIO PREFERENCES (ranked, but all {num_segments} slots must be filled): - Best: chorus or hook — the most recognizable, singable section that defines the song - Good: pre-chorus build into chorus payoff, or a distinct recurring riff/vocal phrase - Acceptable when needed to reach the required count: any verse, bridge, or section that has a clear melodic identity, strong artist presence, and compelling visuals — prefer whichever remaining sections score highest on energy + visual quality AUDIO RULES: - Every segment must start and end at a natural phrase boundary (downbeat, start of a lyric line, top of a chorus) — never mid-phrase or mid-word - The clip must feel musically complete, not cut off mid-thought - Avoid the very opening intro (before the song establishes itself), fade-outs, and dead-air outros — but use anything else if needed to fill all {num_segments} slots VISUAL CRITERIA: - The artist must be clearly visible and prominent - Prefer high-energy performance: intense dancing, expressive delivery, striking imagery - Best creative = compelling audio + strong visual at the same time VISUAL EXCLUSIONS (never select these): - Segments where a full-screen title card or video title occupies most of the frame (e.g. the song/artist name displayed large at the start of the video) - Segments where closing credits, "follow us", or social-handle screens occupy most of the frame - Any shot where static text is the primary visual element rather than the performance NOTE: incidental or small text overlays (e.g. a subtle "out now" corner badge) are fine — only exclude when text dominates the majority of the visible frame. SCORING: Rate each segment 0.0–1.0 by trend/viral potential — how recognizable is this phrase as THE song, and would a TikTok user pick it as their sound? CONTEXT: - Total video duration: {video_duration:.1f} seconds {duration_line} - Segments must not overlap - You MUST return exactly {num_segments} segments even if the song has fewer than {num_segments} choruses — use the best remaining sections to fill the count Return exactly {num_segments} segments, ordered highest to lowest trend potential.""" def _build_ugc_prompt( self, num_segments: int, min_duration: float | None, max_duration: float | None, video_duration: float, ) -> str: """Prompt for UGC / artist-social-media content. Audio will be replaced by the official track, so Gemini should evaluate visuals only — energy, expressiveness, and visual impact. """ if min_duration is not None and max_duration is not None: duration_line = f"- Clip duration: {min_duration}–{max_duration} seconds each" else: duration_line = ( "- Clip duration: choose appropriate lengths at your discretion, " "maximum 30 seconds each" ) return f"""You are selecting clips from a UGC video or artist social-media content (e.g. a fan recording, concert clip, behind-the-scenes footage, or a social media post) to use as marketing creatives on TikTok and Instagram Reels (vertical 9:16). IMPORTANT: This is NOT the official music video. The audio from this clip will be REPLACED by the official song track in post-production. You must evaluate VISUALS ONLY. Completely IGNORE the audio — sound quality, lyrics, crowd noise, and audio sync are irrelevant. Base your entire selection on visual energy and impact. GOAL: Return EXACTLY {num_segments} non-overlapping segments with the highest visual impact — moments most likely to captivate a viewer scrolling through their feed. VISUAL SELECTION CRITERIA (in order of priority): 1. Peak visual energy: dynamic movement, expressive performance, crowd reaction, striking lighting, camera movement, or anything visually arresting 2. Artist visibility: segments where the artist is clearly visible and dominant in frame 3. Visual storytelling: moments that convey emotion, excitement, or authenticity 4. Technical quality: prefer well-lit, in-focus, stable shots over shaky/dark footage VISUAL EXCLUSIONS: - Static shots with little movement or visual interest - Heavily blurred or out-of-focus segments - Segments dominated by crowd backs or obstructions (unless crowd energy is the focus) SCORING: Rate each segment 0.0–1.0 by visual impact alone — how compelling is this footage visually, independent of any audio? CONTEXT: - Total video duration: {video_duration:.1f} seconds {duration_line} - Segments must not overlap - You MUST return exactly {num_segments} segments to fill all slots Return exactly {num_segments} segments, ordered highest to lowest visual impact.""" def _query_performance_window( self, video_file: object, video_duration: float, ) -> PerformanceWindow: """Ask Gemini where actual performance footage begins and ends. Used to hard-exclude title cards and closing credits before returning segments. """ generation_config = types.GenerateContentConfig( response_mime_type="application/json", response_schema=PerformanceWindow, media_resolution=types.MediaResolution.MEDIA_RESOLUTION_LOW, ) prompt = f"""Identify the usable performance window in this music video. TASK: Return two timestamps marking where actual performance footage begins and ends. performance_start: The timestamp (in seconds) when actual performance footage begins. - Return 0.0 if the video starts immediately with the performance - Return a positive value only if there is a title card, "artist – song name" screen, or any introductory graphic that occupies most of the frame before the performance performance_end: The timestamp (in seconds) of the LAST frame of actual performance footage, immediately before any non-performance sequence begins. - Return {video_duration:.1f} if the video ends on performance footage - Return the timestamp BEFORE any closing credits, crew/cast lists, "follow us" screens, social media handles, "out now" full-frame cards, or any outro sequence begin - Be conservative: when uncertain, return an EARLIER (lower) timestamp IMPORTANT: Credits and social screens are common in official music videos even when brief. Err on the side of a lower performance_end rather than a higher one. Total video duration: {video_duration:.1f} seconds""" response = self.client.models.generate_content( model=self.model_name, contents=[video_file, prompt], config=generation_config, ) data = json.loads(response.text) pw = PerformanceWindow(**data) return PerformanceWindow( performance_start=max(0.0, min(pw.performance_start, video_duration)), performance_end=max(pw.performance_start, min(pw.performance_end, video_duration)), ) def _validate_segments( self, segments: list[VideoSegment], video_duration: float, min_duration: float | None = None, max_duration: float | None = None, ) -> list[VideoSegment]: """ Validate and filter segments. Args: segments: List of segments from API video_duration: Total video duration min_duration: Minimum allowed duration (None = no lower bound) max_duration: Maximum allowed duration (None = 30s safety cap) Returns: List of valid segments """ max_cap = max_duration if max_duration is not None else 30.0 valid_segments = [] for segment in segments: # Check bounds if segment.start < 0 or segment.end > video_duration: continue # Check ordering if segment.start >= segment.end: continue # Check duration duration = segment.end - segment.start if min_duration is not None and duration < min_duration: continue if duration > max_cap: continue valid_segments.append(segment) # Check for overlaps valid_segments.sort(key=lambda s: s.start) non_overlapping = [] for segment in valid_segments: if not non_overlapping or segment.start >= non_overlapping[-1].end: non_overlapping.append(segment) return non_overlapping