""" DocuLens - Google Gemini Vision Extraction Provider. Uses the Gemini REST API (Google AI Studio free tier) for document extraction. Designed as a drop-in alternative to the HF Inference pipeline — accepts the same prompt and returns raw text that _parse_vlm_json / _vlm_data_to_result can process. Environment variables: GEMINI_API_KEY — Google AI Studio API key (required for Gemini models) GEMINI_MODEL — Model name override (default: gemini-3.6-flash) GEMINI_TIMEOUT — Request timeout in seconds (default: 120) GEMINI_MAX_RETRIES — Retry count on transient errors (default: 2) Usage: from gemini_client import extract_with_gemini result_text = extract_with_gemini(image_b64, prompt, model="gemini-3.6-flash") """ import os import time import json import logging from typing import Optional import requests logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "") GEMINI_DEFAULT_MODEL = os.environ.get("GEMINI_MODEL", "gemini-3.5-flash-lite") GEMINI_TIMEOUT = int(os.environ.get("GEMINI_TIMEOUT", "120")) GEMINI_MAX_RETRIES = int(os.environ.get("GEMINI_MAX_RETRIES", "2")) GEMINI_RETRY_BASE_DELAY = float(os.environ.get("GEMINI_RETRY_BASE_DELAY", "1.0")) # Google AI Studio REST endpoint _BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models" # Transient HTTP status codes worth retrying _TRANSIENT_CODES = {429, 500, 502, 503} # --------------------------------------------------------------------------- # Low-level API call with retry # --------------------------------------------------------------------------- def _call_gemini( model: str, prompt: str, image_b64: str, *, max_tokens: int = 4096, temperature: float = 0.1, timeout: Optional[int] = None, max_retries: Optional[int] = None, ) -> str: """Call the Gemini generateContent API and return the text response. Uses header-based auth (x-goog-api-key) per Google's recommended approach. For Gemini 3.x thinking models, sets thinkingConfig to minimal to reduce latency for structured extraction tasks. Raises on non-transient errors or after exhausting retries. """ api_key = GEMINI_API_KEY if not api_key: raise RuntimeError("GEMINI_API_KEY environment variable is not set") url = f"{_BASE_URL}/{model}:generateContent" headers = { "Content-Type": "application/json", "x-goog-api-key": api_key, } # Build generationConfig — adapt thinkingConfig to model generation: # Gemini 3.x uses thinkingLevel: "minimal" | "low" | "medium" | "high" # Gemini 2.5 uses thinkingBudget: 0..24576 # Omitting thinkingConfig entirely also works (uses model default). generation_config: dict = { "temperature": temperature, "maxOutputTokens": max_tokens, } # Use minimal thinking for extraction — we want deterministic JSON, not # chain-of-thought reasoning, and it reduces latency + token usage. if model.startswith("gemini-3"): generation_config["thinkingConfig"] = {"thinkingLevel": "minimal"} elif model.startswith("gemini-2.5"): generation_config["thinkingConfig"] = {"thinkingBudget": 0} payload = { "contents": [ { "parts": [ { "inline_data": { "mime_type": "image/jpeg", "data": image_b64, } }, {"text": prompt}, ], } ], "generationConfig": generation_config, } retries = max_retries if max_retries is not None else GEMINI_MAX_RETRIES delay = GEMINI_RETRY_BASE_DELAY t = timeout or GEMINI_TIMEOUT # Log payload size for latency diagnostics payload_json = json.dumps(payload) payload_kb = len(payload_json) / 1024 logger.info( "Gemini %s payload: %.0f KB (image b64: %.0f KB, prompt: %d chars)", model, payload_kb, len(image_b64) / 1024, len(prompt), ) last_exc: Optional[Exception] = None for attempt in range(1 + retries): try: logger.info( "Gemini %s attempt %d/%d", model, attempt + 1, 1 + retries, ) t0 = time.monotonic() resp = requests.post(url, json=payload, headers=headers, timeout=t) elapsed_s = time.monotonic() - t0 logger.info( "Gemini %s HTTP %d in %.1fs (%.0f KB response)", model, resp.status_code, elapsed_s, len(resp.content) / 1024, ) if resp.status_code in _TRANSIENT_CODES and attempt < retries: wait = delay * (2 ** attempt) logger.warning( "Gemini %s transient %d (attempt %d/%d). Retrying in %.1fs... Body: %s", model, resp.status_code, attempt + 1, 1 + retries, wait, resp.text[:300], ) time.sleep(wait) continue if not resp.ok: body = resp.text[:500] logger.error( "Gemini %s HTTP %d: %s", model, resp.status_code, body, ) resp.raise_for_status() data = resp.json() # Extract text from response — thinking models may include # "thought" parts alongside the actual text output. candidates = data.get("candidates", []) if not candidates: raise RuntimeError(f"Gemini returned no candidates: {json.dumps(data)[:300]}") parts = candidates[0].get("content", {}).get("parts", []) # Filter out "thought" parts — only keep regular text output text_parts = [ p["text"] for p in parts if "text" in p and not p.get("thought") ] if not text_parts: # Fallback: include thought parts if no regular text found text_parts = [p["text"] for p in parts if "text" in p] if not text_parts: raise RuntimeError(f"Gemini returned no text parts: {json.dumps(data)[:300]}") return "\n".join(text_parts) except requests.exceptions.Timeout as e: last_exc = e if attempt < retries: wait = delay * (2 ** attempt) logger.warning( "Gemini %s timeout (attempt %d/%d). Retrying in %.1fs...", model, attempt + 1, 1 + retries, wait, ) time.sleep(wait) else: raise except requests.exceptions.HTTPError as e: # Non-transient HTTP error — don't retry raise except Exception as e: last_exc = e if attempt < retries: wait = delay * (2 ** attempt) logger.warning( "Gemini %s error (attempt %d/%d): %s. Retrying in %.1fs...", model, attempt + 1, 1 + retries, str(e)[:120], wait, ) time.sleep(wait) else: raise if last_exc: raise last_exc raise RuntimeError("Gemini call failed silently") # --------------------------------------------------------------------------- # Public extraction function # --------------------------------------------------------------------------- def extract_with_gemini( image_b64: str, prompt: str, model: str = "", ) -> str: """Extract document data using Google Gemini Vision API. Args: image_b64: Base64-encoded JPEG image data (no data-URI prefix). prompt: The extraction prompt text. model: Gemini model id (e.g. "gemini-3.6-flash"). Falls back to GEMINI_DEFAULT_MODEL env var. Returns: Raw text response from Gemini (JSON string to be parsed by caller). Raises: RuntimeError: If GEMINI_API_KEY is not configured. requests.exceptions.HTTPError: On non-transient API errors. """ model_id = model or GEMINI_DEFAULT_MODEL logger.info("Gemini extraction using model=%s", model_id) return _call_gemini(model_id, prompt, image_b64)