Files
server-26/drb-c2-core/app/internal/gemini.py
T
Logan CusanoandClaude Opus 5.5 0060ccde75 AI providers: correlation decision and transcript correction move to OpenAI gpt-4o-mini
One funded account (OpenAI) now covers extraction, embeddings, summaries,
the ambiguous-scene correlation call and transcript correction; Groq stays
optional for STT. New app/internal/llm.py picks the provider from the model
name, so setting either model back to a gemini-* ID needs no code change.
Usage accounting for replay runs now covers both providers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-11-08 01:25:21 -05:00

87 lines
3.8 KiB
Python

"""
One place every Gemini call goes through: JSON-mode generation, an explicit
thinking level, and token accounting.
Why it exists: a day of replay runs (server-26#170) cost ~$5 of Gemini for
~7 two-hour windows — roughly $0.70 per 290 calls, which projects to several
dollars a day per live deployment for correlation alone — and nothing in DRB
could say where it went (server-26#45). Gemini 3.x models "think" by default
and bill that as output; the old google-generativeai SDK these calls used
cannot even set a thinking level. A link/new/orphan choice or a transcript
cleanup does not need extended reasoning.
Every call logs its token counts, and inside a replay run they are also added
to the run's own usage sink (see app/internal/replay.py), so a run reports
what it actually spent instead of an estimate.
"""
import json
import threading
from typing import Optional
from app.config import settings
from app.internal.logger import logger
_client = None
_client_lock = threading.Lock()
# Models that rejected a thinking level: retried without one from then on.
_no_thinking_level: set[str] = set()
def _get_client():
global _client
with _client_lock:
if _client is None:
from google import genai # lazy — only when a Gemini call is made
from google.genai import types
# server-26#153: google-genai's default is no timeout at all, so a
# hung call held its pipeline (and the sweep hold-back sized
# against it) open indefinitely. HttpOptions.timeout is in ms.
# The client does not retry unless retry_options is set, so this
# is the whole bound per request. A timeout raises from
# generate_json like any API error, so every caller's fallback
# (rules-only correlation, uncorrected transcript) is unchanged.
_client = genai.Client(
api_key=settings.gemini_api_key,
http_options=types.HttpOptions(timeout=int(settings.gemini_timeout_seconds * 1000)),
)
return _client
def _config(thinking_level: Optional[str]):
from google.genai import types
kwargs = {"response_mime_type": "application/json"}
if thinking_level:
kwargs["thinking_config"] = types.ThinkingConfig(thinking_level=thinking_level)
return types.GenerateContentConfig(**kwargs)
def _record(purpose: str, model: str, usage) -> None:
prompt = getattr(usage, "prompt_token_count", None) or 0
output = getattr(usage, "candidates_token_count", None) or 0
thoughts = getattr(usage, "thoughts_token_count", None) or 0
from app.internal.llm import record_usage
record_usage(purpose, model, prompt, output, thoughts)
def generate_json(model: str, prompt: str, *, purpose: str,
thinking_level: Optional[str] = "minimal") -> dict:
"""
Synchronous (run it via asyncio.to_thread). Returns the parsed JSON body.
Raises on API failure, exactly like the old per-module helpers, so callers'
ai_health classification (billing / dead model / transient) is unchanged.
"""
client = _get_client()
level = None if model in _no_thinking_level else thinking_level
try:
resp = client.models.generate_content(model=model, contents=prompt, config=_config(level))
except Exception as e:
# A model that doesn't accept this thinking level answers 400 for
# every call; drop the setting for that model rather than lose the tier.
if level and "thinking" in str(e).lower():
logger.warning(f"gemini: {model} rejected thinking_level={level!r} ({e}); retrying without it")
_no_thinking_level.add(model)
resp = client.models.generate_content(model=model, contents=prompt, config=_config(None))
else:
raise
_record(purpose, model, getattr(resp, "usage_metadata", None))
return json.loads(resp.text)