Spaces:
Runtime error
Runtime error
| """ | |
| 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) | |