Spaces:
Runtime error
Runtime error
File size: 8,527 Bytes
04c4194 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 | """
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)
|