Replay usage key gemini_usage -> llm_usage (backend + ReplayTab), re-price the correction/correlation estimates for gpt-4o-mini, no temperature for o-series or gpt-5 models, and refresh comments still naming Gemini. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
"""
|
|
One entry point for the two small JSON-mode LLM jobs: the correlation decision
|
|
on ambiguous scenes and transcript correction.
|
|
|
|
The provider follows the model name: a `gemini-*` model goes through
|
|
app/internal/gemini.py, anything else is an OpenAI chat model. Since the
|
|
2026-11-08 provider consolidation both defaults are OpenAI (gpt-4o-mini), so
|
|
one funded account (OpenAI) covers extraction, embeddings, summaries and these
|
|
two jobs; Groq stays optional for transcription. Pointing either model setting
|
|
back at a gemini-* ID needs no code change.
|
|
|
|
Token counts from both providers land in the same usage sink, so a replay run
|
|
(app/internal/replay.py) reports what it spent whichever provider served it.
|
|
"""
|
|
import json
|
|
from contextvars import ContextVar
|
|
from typing import Optional
|
|
|
|
from app.config import settings
|
|
from app.internal.logger import logger
|
|
|
|
_usage_sink: ContextVar[Optional[dict]] = ContextVar("drb_llm_usage", default=None)
|
|
|
|
|
|
def collect_usage(sink: Optional[dict]):
|
|
"""Route token counts for the current context into `sink` (a replay run). Returns a reset token."""
|
|
return _usage_sink.set(sink)
|
|
|
|
|
|
def reset_usage(token) -> None:
|
|
_usage_sink.reset(token)
|
|
|
|
|
|
def record_usage(purpose: str, model: str, prompt: int, output: int, thinking: int = 0) -> None:
|
|
logger.info(f"llm usage {purpose} {model}: in={prompt} out={output} thinking={thinking}")
|
|
sink = _usage_sink.get()
|
|
if sink is not None:
|
|
row = sink.setdefault(f"{purpose}:{model}", {"calls": 0, "in": 0, "out": 0, "thinking": 0})
|
|
row["calls"] += 1
|
|
row["in"] += prompt
|
|
row["out"] += output
|
|
row["thinking"] += thinking
|
|
|
|
|
|
def provider_for(model: str) -> str:
|
|
return "gemini" if (model or "").startswith("gemini") else "openai"
|
|
|
|
|
|
def available(model: str) -> bool:
|
|
"""Whether the provider serving `model` has a key configured."""
|
|
key = settings.gemini_api_key if provider_for(model) == "gemini" else settings.openai_api_key
|
|
return bool(key)
|
|
|
|
|
|
def timeout_seconds(model: str) -> float:
|
|
if provider_for(model) == "gemini":
|
|
return settings.gemini_timeout_seconds
|
|
return (1 + max(0, settings.openai_max_retries)) * settings.openai_timeout_seconds
|
|
|
|
|
|
def billing_fix(model: str) -> str:
|
|
if provider_for(model) == "gemini":
|
|
return "Top up billing at https://ai.studio/projects"
|
|
return "Top up billing at https://platform.openai.com/settings/organization/billing"
|
|
|
|
|
|
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 so callers' ai_health classification (billing /
|
|
dead model / transient) is unchanged. `thinking_level` only applies to
|
|
Gemini 3.x; OpenAI's gpt-4o family has no such setting.
|
|
"""
|
|
if provider_for(model) == "gemini":
|
|
from app.internal import gemini
|
|
return gemini.generate_json(model, prompt, purpose=purpose, thinking_level=thinking_level)
|
|
|
|
from app.internal.ai_clients import openai_client
|
|
# JSON mode requires the word "JSON" in the prompt; both prompts say
|
|
# "Return JSON" / "Respond with JSON only".
|
|
kwargs = {}
|
|
# Reasoning models (o-series, gpt-5) reject a temperature; chat models get 0.
|
|
if not model.startswith(("o1", "o3", "o4", "gpt-5")):
|
|
kwargs["temperature"] = 0
|
|
resp = openai_client().chat.completions.create(
|
|
model=model,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
response_format={"type": "json_object"},
|
|
**kwargs,
|
|
)
|
|
usage = getattr(resp, "usage", None)
|
|
record_usage(purpose, model,
|
|
getattr(usage, "prompt_tokens", 0) or 0,
|
|
getattr(usage, "completion_tokens", 0) or 0)
|
|
return json.loads(resp.choices[0].message.content or "{}")
|