| |
| from __future__ import annotations |
|
|
| import ast |
| import importlib.util |
| import json |
| import logging |
| import mimetypes |
| import os |
| import re |
| import sys |
| import time |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import requests |
|
|
|
|
| ROOT_DIR = Path(__file__).resolve().parent.parent |
| DEFAULT_API_BASE_URL = "https://generativelanguage.googleapis.com" |
| DEFAULT_MODEL_NAME = "models/gemini-2.5-pro" |
| DEFAULT_SYSTEM_INSTRUCTION = ( |
| "You are a careful video annotation assistant. " |
| "Return strict JSON only and prefer high-confidence visual facts." |
| ) |
|
|
|
|
| class GeminiBlockedContentError(RuntimeError): |
| def __init__(self, block_reason: str, response_payload: dict[str, Any]): |
| self.block_reason = (block_reason or "").strip() or "UNKNOWN" |
| self.response_payload = response_payload |
| response_snippet = json.dumps(response_payload, ensure_ascii=False)[:2000] |
| super().__init__( |
| f"Gemini request blocked by safety filter. " |
| f"block_reason={self.block_reason} response={response_snippet}" |
| ) |
|
|
|
|
| def load_repo_module(module_name: str, relative_path: str): |
| module_path = ROOT_DIR / relative_path |
| spec = importlib.util.spec_from_file_location(module_name, module_path) |
| if spec is None or spec.loader is None: |
| raise ImportError(f"Failed to load module from {module_path}") |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[module_name] = module |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def normalize_model_name(value: str | None) -> str: |
| cleaned = (value or "").strip() |
| if not cleaned: |
| return DEFAULT_MODEL_NAME |
| if cleaned.startswith("models/"): |
| return cleaned |
| return f"models/{cleaned}" |
|
|
|
|
| def extract_json_text(raw_text: str) -> str: |
| text = raw_text.strip() |
| fenced_matches = re.findall(r"```(?:json)?\s*(.*?)```", text, flags=re.DOTALL | re.IGNORECASE) |
| candidates = [match.strip() for match in fenced_matches if match.strip()] |
| candidates.append(text) |
| for candidate in candidates: |
| if candidate.startswith("{") or candidate.startswith("["): |
| return candidate |
| for open_char, close_char in (("{", "}"), ("[", "]")): |
| start = text.find(open_char) |
| if start == -1: |
| continue |
| depth = 0 |
| in_string = False |
| escape = False |
| for index in range(start, len(text)): |
| char = text[index] |
| if in_string: |
| if escape: |
| escape = False |
| elif char == "\\": |
| escape = True |
| elif char == '"': |
| in_string = False |
| continue |
| if char == '"': |
| in_string = True |
| continue |
| if char == open_char: |
| depth += 1 |
| elif char == close_char: |
| depth -= 1 |
| if depth == 0: |
| return text[start : index + 1] |
| raise ValueError("Failed to extract a JSON object from the model response.") |
|
|
|
|
| def parse_json_payload(raw_text: str) -> dict[str, Any]: |
| json_text = extract_json_text(raw_text) |
| try: |
| payload = json.loads(json_text) |
| except json.JSONDecodeError: |
| try: |
| payload = ast.literal_eval(json_text) |
| except (SyntaxError, ValueError) as exc: |
| snippet = raw_text.strip().replace("\n", " ")[:500] |
| raise ValueError(f"Model did not return valid JSON. Snippet: {snippet}") from exc |
| if isinstance(payload, list): |
| if len(payload) == 1 and isinstance(payload[0], dict): |
| return payload[0] |
| raise ValueError(f"Expected a JSON object, received list with {len(payload)} items.") |
| if not isinstance(payload, dict): |
| raise ValueError(f"Expected a JSON object, received {type(payload).__name__}.") |
| return payload |
|
|
|
|
| def guess_mime_type(path: Path) -> str: |
| mime_type, _ = mimetypes.guess_type(str(path)) |
| if mime_type: |
| return mime_type |
| suffix = path.suffix.lower() |
| if suffix == ".mp4": |
| return "video/mp4" |
| if suffix == ".mov": |
| return "video/quicktime" |
| if suffix == ".mkv": |
| return "video/x-matroska" |
| if suffix == ".avi": |
| return "video/x-msvideo" |
| if suffix == ".webm": |
| return "video/webm" |
| return "application/octet-stream" |
|
|
|
|
| def coerce_file_payload(payload: dict[str, Any]) -> dict[str, Any]: |
| if "file" in payload and isinstance(payload["file"], dict): |
| return payload["file"] |
| return payload |
|
|
|
|
| def extract_response_text(payload: dict[str, Any]) -> str: |
| texts: list[str] = [] |
| for candidate in payload.get("candidates", []): |
| content = candidate.get("content", {}) |
| for part in content.get("parts", []): |
| text = part.get("text") |
| if text: |
| texts.append(str(text).strip()) |
| return "\n".join(part for part in texts if part) |
|
|
|
|
| def extract_finish_reason(payload: dict[str, Any]) -> str: |
| candidates = payload.get("candidates", []) |
| if not candidates: |
| return "" |
| return str(candidates[0].get("finishReason", "")).strip() |
|
|
|
|
| def extract_block_reason(payload: dict[str, Any]) -> str: |
| prompt_feedback = payload.get("promptFeedback", {}) |
| if not isinstance(prompt_feedback, dict): |
| return "" |
| return str(prompt_feedback.get("blockReason", "")).strip() |
|
|
|
|
| def looks_like_detection_output(text: str) -> bool: |
| lowered = text.lower() |
| detection_markers = ( |
| "box_2d", |
| "\"bbox\"", |
| "\"label\"", |
| "\"labels\"", |
| "\"tracking_id\"", |
| "\"polygon\"", |
| "\"mask\"", |
| "\"score\"", |
| ) |
| return any(marker in lowered for marker in detection_markers) |
|
|
|
|
| def _string_schema(*, nullable: bool = False) -> dict[str, Any]: |
| schema: dict[str, Any] = {"type": "STRING"} |
| if nullable: |
| schema["nullable"] = True |
| return schema |
|
|
|
|
| def _boolean_schema(*, nullable: bool = False) -> dict[str, Any]: |
| schema: dict[str, Any] = {"type": "BOOLEAN"} |
| if nullable: |
| schema["nullable"] = True |
| return schema |
|
|
|
|
| def _object_schema( |
| properties: dict[str, dict[str, Any]], |
| *, |
| required: list[str] | None = None, |
| ) -> dict[str, Any]: |
| schema: dict[str, Any] = { |
| "type": "OBJECT", |
| "properties": properties, |
| } |
| if required: |
| schema["required"] = required |
| return schema |
|
|
|
|
| def _array_schema(items: dict[str, Any]) -> dict[str, Any]: |
| return { |
| "type": "ARRAY", |
| "items": items, |
| } |
|
|
|
|
| def build_subject_registry_item_schema() -> dict[str, Any]: |
| properties = { |
| "subject_id": _string_schema(), |
| "category": _string_schema(), |
| "term": _string_schema(), |
| "profile": _string_schema(), |
| } |
| return _object_schema(properties, required=list(properties)) |
|
|
|
|
| def build_global_context_response_schema() -> dict[str, Any]: |
| properties = { |
| "is_sports": _boolean_schema(), |
| "has_animal": _boolean_schema(), |
| "is_cooking": _boolean_schema(), |
| "subject_registry": _array_schema(build_subject_registry_item_schema()), |
| "stable_scene_summary": _string_schema(), |
| "stable_camera_summary": _string_schema(), |
| } |
| return _object_schema(properties, required=list(properties)) |
|
|
|
|
| def build_event_response_schema() -> dict[str, Any]: |
| properties = { |
| "caption": _string_schema(), |
| "has_appearance": _boolean_schema(), |
| "has_disappearance": _boolean_schema(), |
| "has_camera_motion": _boolean_schema(), |
| "has_focus_shift": _boolean_schema(nullable=True), |
| "is_multi_person": _boolean_schema(), |
| "has_environmental_change": _boolean_schema(), |
| "has_occlusion": _boolean_schema(), |
| "has_lighting_change": _boolean_schema(), |
| "has_animal": _boolean_schema(), |
| "is_cooking": _boolean_schema(), |
| "focus_object": _string_schema(), |
| "camera_motion_label": _string_schema(), |
| "camera_motion_desc": _string_schema(), |
| } |
| return _object_schema(properties, required=list(properties)) |
|
|
|
|
| def build_caption_only_response_schema() -> dict[str, Any]: |
| properties = {"caption": _string_schema()} |
| return _object_schema(properties, required=["caption"]) |
|
|
|
|
| def build_event_camera_repair_response_schema() -> dict[str, Any]: |
| properties = { |
| "has_camera_motion": _boolean_schema(), |
| "camera_motion_label": _string_schema(), |
| "camera_motion_desc": _string_schema(), |
| } |
| return _object_schema(properties, required=list(properties)) |
|
|
|
|
| def build_camera_desc_only_response_schema() -> dict[str, Any]: |
| properties = {"camera_motion_desc": _string_schema()} |
| return _object_schema(properties, required=["camera_motion_desc"]) |
|
|
|
|
| def build_global_camera_repair_response_schema() -> dict[str, Any]: |
| properties = { |
| "camera_movement_prefix": _string_schema(), |
| "camera_movement_label": _string_schema(), |
| "camera_movement_description": _string_schema(), |
| } |
| return _object_schema(properties, required=list(properties)) |
|
|
|
|
| def build_subject_profile_repair_response_schema() -> dict[str, Any]: |
| properties = {"subject_profile": _string_schema()} |
| return _object_schema(properties, required=["subject_profile"]) |
|
|
|
|
| def build_global_refine_response_schema(*, include_detailed_caption: bool) -> dict[str, Any]: |
| properties = { |
| "background": _string_schema(), |
| "theme": _string_schema(), |
| "style": _string_schema(), |
| "shot_type": _string_schema(), |
| "camera_movement_prefix": _string_schema(), |
| "camera_movement_label": _string_schema(), |
| "camera_movement_description": _string_schema(), |
| "lighting": _string_schema(), |
| "atmosphere": _string_schema(), |
| "subject_profile": _string_schema(), |
| "short_caption": _string_schema(), |
| "middle_caption": _string_schema(), |
| "long_caption": _string_schema(), |
| } |
| required = list(properties) |
| if include_detailed_caption: |
| properties["detailed_caption"] = _string_schema() |
| required.append("detailed_caption") |
| return _object_schema(properties, required=required) |
|
|
|
|
| def infer_response_schema(prompt: str) -> dict[str, Any] | None: |
| lowered = prompt.lower() |
| if ( |
| "writing the global annotation" in lowered |
| and '"background"' in prompt |
| and '"short_caption"' in prompt |
| and '"middle_caption"' in prompt |
| and '"long_caption"' in prompt |
| ): |
| return build_global_refine_response_schema( |
| include_detailed_caption='"detailed_caption"' in prompt |
| ) |
| if ( |
| '"subject_registry"' in prompt |
| and '"stable_scene_summary"' in prompt |
| and '"stable_camera_summary"' in prompt |
| and '"is_sports"' in prompt |
| ): |
| return build_global_context_response_schema() |
| if ( |
| '"has_appearance"' in prompt |
| and '"has_disappearance"' in prompt |
| and '"focus_object"' in prompt |
| and '"camera_motion_desc"' in prompt |
| ): |
| return build_event_response_schema() |
| if "revising the caption for event" in lowered and '"caption"' in prompt: |
| return build_caption_only_response_schema() |
| if "repairing the camera motion fields for event" in lowered: |
| return build_event_camera_repair_response_schema() |
| if "rewriting only the event-level camera_motion_desc" in lowered: |
| return build_camera_desc_only_response_schema() |
| if "repairing the global camera motion fields" in lowered: |
| return build_global_camera_repair_response_schema() |
| if "repairing the subject_profile field" in lowered: |
| return build_subject_profile_repair_response_schema() |
| if "return json in exactly this schema" in lowered: |
| return {"type": "OBJECT"} |
| return None |
|
|
|
|
| def resolve_thinking_budget(model_name: str, raw_value: str | None) -> int | None: |
| if raw_value is not None and raw_value.strip() != "": |
| return int(raw_value) |
| normalized = normalize_model_name(model_name) |
| if normalized == "models/gemini-2.5-pro": |
| return 128 |
| return None |
|
|
|
|
| def parse_cleanup_flag(raw_value: str | None) -> bool: |
| if raw_value is None: |
| return True |
| return raw_value.strip().lower() not in {"0", "false", "no", "off"} |
|
|
|
|
| def extract_file_state_name(value: Any) -> str: |
| if isinstance(value, dict): |
| name = value.get("name") |
| return str(name).upper() if name else "" |
| if value is None: |
| return "" |
| return str(value).upper() |
|
|
|
|
| @dataclass |
| class UploadedFile: |
| name: str |
| uri: str |
| mime_type: str |
| state: str |
|
|
|
|
| class GeminiRunner: |
| def __init__( |
| self, |
| model_path: Path, |
| device_map: str, |
| dtype: str, |
| attn_implementation: str, |
| trust_remote_code: bool, |
| fps: float, |
| max_pixels: int | None, |
| temperature: float, |
| top_p: float, |
| top_k: int, |
| repetition_penalty: float, |
| greedy: bool, |
| ) -> None: |
| del model_path, device_map, dtype, attn_implementation, trust_remote_code, repetition_penalty |
| self.session = requests.Session() |
| self.api_key = ( |
| os.environ.get("GEMINI_API_KEY") |
| or os.environ.get("GOOGLE_API_KEY") |
| or os.environ.get("GEMINI_API_TOKEN") |
| or "" |
| ).strip() |
| if not self.api_key: |
| raise RuntimeError("Missing Gemini API key. Set GEMINI_API_KEY or GOOGLE_API_KEY.") |
|
|
| self.model_name = normalize_model_name(os.environ.get("GEMINI_MODEL")) |
| self.api_base_url = os.environ.get("GEMINI_API_BASE_URL", DEFAULT_API_BASE_URL).rstrip("/") |
| self.system_instruction = os.environ.get("GEMINI_SYSTEM_INSTRUCTION", DEFAULT_SYSTEM_INSTRUCTION).strip() |
| self.temperature = 0.0 if greedy else temperature |
| self.top_p = top_p |
| self.top_k = top_k |
| self.greedy = greedy |
| self.fps = fps if fps and fps > 0 else None |
| self.max_pixels = max_pixels if max_pixels and max_pixels > 0 else None |
| self.http_timeout = float(os.environ.get("GEMINI_HTTP_TIMEOUT", "600")) |
| self.file_poll_interval = float(os.environ.get("GEMINI_FILE_POLL_INTERVAL", "5")) |
| self.file_poll_timeout = float(os.environ.get("GEMINI_FILE_POLL_TIMEOUT", "1800")) |
| self.http_retry_count = max(1, int(os.environ.get("GEMINI_HTTP_RETRY_COUNT", "32"))) |
| self.http_retry_base_delay = max(0.5, float(os.environ.get("GEMINI_HTTP_RETRY_BASE_DELAY", "10"))) |
| self.http_retry_max_delay = max( |
| self.http_retry_base_delay, |
| float(os.environ.get("GEMINI_HTTP_RETRY_MAX_DELAY", "120")), |
| ) |
| self.cleanup_uploaded_files = parse_cleanup_flag(os.environ.get("GEMINI_CLEANUP_FILES")) |
| self.thinking_budget = resolve_thinking_budget(self.model_name, os.environ.get("GEMINI_THINKING_BUDGET")) |
| self.json_retry_count = max(1, int(os.environ.get("GEMINI_JSON_RETRY_COUNT", "3"))) |
| self.retry_extra_tokens = max(128, int(os.environ.get("GEMINI_RETRY_EXTRA_TOKENS", "512"))) |
| self.max_retry_tokens = max(1024, int(os.environ.get("GEMINI_MAX_RETRY_TOKENS", "8192"))) |
| self.use_response_schema = parse_cleanup_flag(os.environ.get("GEMINI_USE_RESPONSE_SCHEMA")) |
| self.upload_cache: dict[Path, UploadedFile] = {} |
|
|
| logging.info( |
| "Initialized Gemini runner with model=%s api_base=%s cleanup_files=%s json_retries=%d http_retries=%d response_schema=%s", |
| self.model_name, |
| self.api_base_url, |
| self.cleanup_uploaded_files, |
| self.json_retry_count, |
| self.http_retry_count, |
| self.use_response_schema, |
| ) |
|
|
| @staticmethod |
| def _is_retryable_status(status_code: int) -> bool: |
| return status_code == 429 or status_code in {500, 502, 503, 504} |
|
|
| def _retry_delay(self, attempt: int, response: requests.Response | None = None) -> float: |
| if response is not None: |
| retry_after = response.headers.get("Retry-After", "").strip() |
| if retry_after: |
| try: |
| return min(float(retry_after), self.http_retry_max_delay) |
| except ValueError: |
| pass |
| delay = self.http_retry_base_delay * (2 ** max(attempt - 1, 0)) |
| return min(delay, self.http_retry_max_delay) |
|
|
| def _request( |
| self, |
| method: str, |
| url: str, |
| *, |
| json_payload: dict[str, Any] | None = None, |
| data: Any = None, |
| headers: dict[str, str] | None = None, |
| timeout: float | None = None, |
| expected_status: int | None = None, |
| include_api_key: bool = True, |
| ) -> requests.Response: |
| merged_headers = dict(headers or {}) |
| last_error: Exception | None = None |
| for attempt in range(1, self.http_retry_count + 1): |
| response: requests.Response | None = None |
| try: |
| response = self.session.request( |
| method=method, |
| url=url, |
| params={"key": self.api_key} if include_api_key else None, |
| json=json_payload, |
| data=data, |
| headers=merged_headers, |
| timeout=timeout or self.http_timeout, |
| ) |
| if expected_status is not None and response.status_code != expected_status: |
| if self._is_retryable_status(response.status_code) and attempt < self.http_retry_count: |
| delay = self._retry_delay(attempt, response) |
| logging.warning( |
| "%s %s returned retryable status %d on attempt %d/%d. Retrying in %.1fs.", |
| method, |
| url, |
| response.status_code, |
| attempt, |
| self.http_retry_count, |
| delay, |
| ) |
| time.sleep(delay) |
| continue |
| raise RuntimeError( |
| f"{method} {url} failed with {response.status_code}: {response.text.strip()}" |
| ) |
| try: |
| response.raise_for_status() |
| except requests.HTTPError as exc: |
| if self._is_retryable_status(response.status_code) and attempt < self.http_retry_count: |
| delay = self._retry_delay(attempt, response) |
| logging.warning( |
| "%s %s returned retryable status %d on attempt %d/%d. Retrying in %.1fs.", |
| method, |
| url, |
| response.status_code, |
| attempt, |
| self.http_retry_count, |
| delay, |
| ) |
| time.sleep(delay) |
| continue |
| raise RuntimeError( |
| f"{method} {url} failed with {response.status_code}: {response.text.strip()}" |
| ) from exc |
| return response |
| except requests.RequestException as exc: |
| last_error = exc |
| if attempt < self.http_retry_count: |
| delay = self._retry_delay(attempt, response) |
| logging.warning( |
| "%s %s failed with %s on attempt %d/%d. Retrying in %.1fs.", |
| method, |
| url, |
| exc.__class__.__name__, |
| attempt, |
| self.http_retry_count, |
| delay, |
| ) |
| time.sleep(delay) |
| continue |
| raise RuntimeError(f"{method} {url} failed after {attempt} attempts: {exc}") from exc |
| if last_error is not None: |
| raise RuntimeError(f"{method} {url} failed after {self.http_retry_count} attempts: {last_error}") from last_error |
| raise RuntimeError(f"{method} {url} failed without a response.") |
|
|
| def _request_json( |
| self, |
| method: str, |
| url: str, |
| *, |
| json_payload: dict[str, Any] | None = None, |
| data: Any = None, |
| headers: dict[str, str] | None = None, |
| timeout: float | None = None, |
| expected_status: int | None = None, |
| include_api_key: bool = True, |
| ) -> dict[str, Any]: |
| response = self._request( |
| method, |
| url, |
| json_payload=json_payload, |
| data=data, |
| headers=headers, |
| timeout=timeout, |
| expected_status=expected_status, |
| include_api_key=include_api_key, |
| ) |
| if not response.content: |
| return {} |
| return response.json() |
|
|
| def _upload_file(self, path: Path) -> UploadedFile: |
| resolved = path.expanduser().resolve() |
| cached = self.upload_cache.get(resolved) |
| if cached is not None: |
| return cached |
|
|
| if not resolved.exists(): |
| raise FileNotFoundError(f"Video file does not exist: {resolved}") |
|
|
| mime_type = guess_mime_type(resolved) |
| file_size = resolved.stat().st_size |
| start_headers = { |
| "X-Goog-Upload-Protocol": "resumable", |
| "X-Goog-Upload-Command": "start", |
| "X-Goog-Upload-Header-Content-Length": str(file_size), |
| "X-Goog-Upload-Header-Content-Type": mime_type, |
| "Content-Type": "application/json", |
| } |
| start_payload = {"file": {"display_name": resolved.name}} |
| start_response = self._request( |
| "POST", |
| f"{self.api_base_url}/upload/v1beta/files", |
| json_payload=start_payload, |
| headers=start_headers, |
| ) |
| upload_url = ( |
| start_response.headers.get("X-Goog-Upload-URL") |
| or start_response.headers.get("x-goog-upload-url") |
| or "" |
| ).strip() |
| if not upload_url: |
| raise RuntimeError("Gemini Files API did not return an upload URL.") |
|
|
| upload_headers = { |
| "Content-Length": str(file_size), |
| "X-Goog-Upload-Offset": "0", |
| "X-Goog-Upload-Command": "upload, finalize", |
| } |
| with resolved.open("rb") as handle: |
| upload_payload = self._request_json( |
| "POST", |
| upload_url, |
| data=handle, |
| headers=upload_headers, |
| timeout=self.http_timeout, |
| include_api_key=False, |
| ) |
|
|
| uploaded = self._coerce_uploaded_file(upload_payload, fallback_mime_type=mime_type) |
| uploaded = self._wait_until_active(uploaded) |
| self.upload_cache[resolved] = uploaded |
| return uploaded |
|
|
| def _coerce_uploaded_file(self, payload: dict[str, Any], fallback_mime_type: str) -> UploadedFile: |
| file_payload = coerce_file_payload(payload) |
| name = str(file_payload.get("name", "")).strip() |
| uri = str(file_payload.get("uri", "")).strip() |
| mime_type = str(file_payload.get("mimeType", "") or file_payload.get("mime_type", "")).strip() |
| state = extract_file_state_name(file_payload.get("state")) |
| if not name: |
| raise RuntimeError(f"Gemini file upload/get response did not contain a file name: {payload}") |
| return UploadedFile( |
| name=name, |
| uri=uri, |
| mime_type=mime_type or fallback_mime_type, |
| state=state, |
| ) |
|
|
| def _get_file(self, file_name: str, fallback_mime_type: str) -> UploadedFile: |
| payload = self._request_json( |
| "GET", |
| f"{self.api_base_url}/v1beta/{file_name}", |
| ) |
| return self._coerce_uploaded_file(payload, fallback_mime_type=fallback_mime_type) |
|
|
| def _wait_until_active(self, uploaded: UploadedFile) -> UploadedFile: |
| current = uploaded |
| deadline = time.time() + self.file_poll_timeout |
| while True: |
| state = current.state.upper() |
| if current.uri and state in {"", "ACTIVE"}: |
| return current |
| if state not in {"PROCESSING", "UNSPECIFIED", "STATE_UNSPECIFIED", ""}: |
| raise RuntimeError( |
| f"Gemini file {current.name} entered unexpected state {current.state or '<empty>'}." |
| ) |
| if time.time() >= deadline: |
| raise TimeoutError(f"Timed out waiting for Gemini file {current.name} to become ACTIVE.") |
| time.sleep(self.file_poll_interval) |
| current = self._get_file(current.name, fallback_mime_type=current.mime_type) |
|
|
| def _build_contents(self, video_paths: list[Path], prompt: str) -> list[dict[str, Any]]: |
| parts: list[dict[str, Any]] = [] |
| for path in video_paths: |
| uploaded = self._upload_file(path) |
| parts.append( |
| { |
| "file_data": { |
| "mime_type": uploaded.mime_type, |
| "file_uri": uploaded.uri, |
| } |
| } |
| ) |
| parts.append({"text": prompt}) |
| return [{"role": "user", "parts": parts}] |
|
|
| def _build_generation_config( |
| self, |
| max_new_tokens: int, |
| *, |
| force_deterministic: bool = False, |
| response_schema: dict[str, Any] | None = None, |
| ) -> dict[str, Any]: |
| config: dict[str, Any] = { |
| "temperature": 0.0 if force_deterministic else self.temperature, |
| "maxOutputTokens": max_new_tokens, |
| "responseMimeType": "application/json", |
| } |
| if not force_deterministic and self.top_p and self.top_p > 0: |
| config["topP"] = self.top_p |
| if not force_deterministic and self.top_k and self.top_k > 0: |
| config["topK"] = self.top_k |
| if self.thinking_budget is not None: |
| config["thinkingConfig"] = {"thinkingBudget": self.thinking_budget} |
| if response_schema is not None: |
| config["responseSchema"] = response_schema |
| return config |
|
|
| def _build_retry_prompt(self, prompt: str, retry_reason: str, attempt: int) -> str: |
| return ( |
| f"{prompt}\n\n" |
| "IMPORTANT RETRY INSTRUCTIONS:\n" |
| f"- The previous response was invalid because: {retry_reason}.\n" |
| f"- This is retry attempt {attempt}.\n" |
| "- Return exactly one complete top-level JSON object matching the requested schema.\n" |
| "- Do not return a JSON array or a list of events.\n" |
| "- If you previously returned one item per clip or one item per event, merge them into one global object that matches the requested schema.\n" |
| "- Do not return object-detection output, boxes, coordinates, labels-only results, OCR blocks, or grounding results.\n" |
| "- Never emit keys such as box_2d, bbox, label, labels, score, mask, polygon, or tracking_id unless the requested schema explicitly asks for them.\n" |
| "- Do not include markdown fences, explanations, or any extra text.\n" |
| "- Make sure every string is closed and the JSON is syntactically complete.\n" |
| "- Keep field values concise enough to finish within the token budget.\n" |
| ) |
|
|
| def _generate_once( |
| self, |
| *, |
| video_paths: list[Path], |
| prompt: str, |
| max_new_tokens: int, |
| force_deterministic: bool, |
| ) -> dict[str, Any]: |
| response_schema = infer_response_schema(prompt) if self.use_response_schema else None |
| request_payload = { |
| "systemInstruction": { |
| "parts": [{"text": self.system_instruction}], |
| }, |
| "contents": self._build_contents(video_paths, prompt), |
| "generationConfig": self._build_generation_config( |
| max_new_tokens, |
| force_deterministic=force_deterministic, |
| response_schema=response_schema, |
| ), |
| "store": False, |
| } |
| return self._request_json( |
| "POST", |
| f"{self.api_base_url}/v1beta/{self.model_name}:generateContent", |
| json_payload=request_payload, |
| ) |
|
|
| def generate_json( |
| self, |
| video_paths: list[Path], |
| prompt: str, |
| max_new_tokens: int, |
| ) -> dict[str, Any]: |
| current_max_tokens = max_new_tokens |
| retry_reason = "initial invalid JSON response" |
| last_error: Exception | None = None |
| last_response_text = "" |
| last_finish_reason = "" |
|
|
| for attempt in range(1, self.json_retry_count + 1): |
| attempt_prompt = ( |
| prompt if attempt == 1 else self._build_retry_prompt(prompt, retry_reason, attempt) |
| ) |
| response_payload = self._generate_once( |
| video_paths=video_paths, |
| prompt=attempt_prompt, |
| max_new_tokens=current_max_tokens, |
| force_deterministic=attempt > 1, |
| ) |
| finish_reason = extract_finish_reason(response_payload) |
| block_reason = extract_block_reason(response_payload) |
| output_text = extract_response_text(response_payload) |
| last_response_text = output_text |
| last_finish_reason = finish_reason |
|
|
| if block_reason == "PROHIBITED_CONTENT" or finish_reason == "PROHIBITED_CONTENT": |
| raise GeminiBlockedContentError( |
| block_reason=block_reason or finish_reason, |
| response_payload=response_payload, |
| ) |
|
|
| if not output_text: |
| retry_reason = f"empty text response (finishReason={finish_reason or '<empty>'})" |
| last_error = RuntimeError( |
| "Gemini generateContent did not return a text payload. " |
| f"finish_reason={finish_reason or '<empty>'} response={json.dumps(response_payload, ensure_ascii=False)[:2000]}" |
| ) |
| else: |
| try: |
| return parse_json_payload(output_text) |
| except Exception as exc: |
| snippet = output_text.strip().replace("\n", " ")[:500] |
| logging.warning( |
| "Gemini returned invalid JSON on attempt %d/%d (finish_reason=%s, max_tokens=%d). Snippet: %s", |
| attempt, |
| self.json_retry_count, |
| finish_reason or "<empty>", |
| current_max_tokens, |
| snippet, |
| ) |
| if looks_like_detection_output(output_text): |
| retry_reason = ( |
| "You returned object-detection or grounding JSON with boxes/labels instead of the requested schema." |
| ) |
| else: |
| retry_reason = f"{type(exc).__name__}: {exc}" |
| last_error = exc |
|
|
| if finish_reason == "MAX_TOKENS": |
| current_max_tokens = min(self.max_retry_tokens, current_max_tokens + self.retry_extra_tokens) |
|
|
| snippet = last_response_text.strip().replace("\n", " ")[:500] |
| raise RuntimeError( |
| "Gemini failed to return valid JSON after retries. " |
| f"finish_reason={last_finish_reason or '<empty>'} " |
| f"last_error={last_error} " |
| f"last_snippet={snippet}" |
| ) |
|
|
| def close(self) -> None: |
| if not self.cleanup_uploaded_files: |
| return |
| seen_names = {uploaded.name for uploaded in self.upload_cache.values() if uploaded.name} |
| for file_name in seen_names: |
| try: |
| self._request( |
| "DELETE", |
| f"{self.api_base_url}/v1beta/{file_name}", |
| expected_status=200, |
| ) |
| except Exception as exc: |
| logging.warning("Failed to delete Gemini file %s: %s", file_name, exc) |
| self.upload_cache.clear() |
|
|
