diff --git a/drb-c2-core/.env.example b/drb-c2-core/.env.example index 352dd28..310a2b6 100644 --- a/drb-c2-core/.env.example +++ b/drb-c2-core/.env.example @@ -38,6 +38,9 @@ NWS_AREA=NY # OpenAI — transcription (whisper-1, and the STT fallback), extraction, embeddings OPENAI_API_KEY= +# Since 2026-11-08 OpenAI also serves the correlation decision and transcript +# correction (CORR_CHEAP_MODEL / TRANSCRIPT_CORRECTION_MODEL, default gpt-4o-mini). +# GEMINI_API_KEY is only needed if you point either of those at a gemini-* model. # OpenAI transcription model: whisper-1 | gpt-4o-mini-transcribe | gpt-4o-transcribe STT_MODEL=whisper-1 diff --git a/drb-c2-core/app/config.py b/drb-c2-core/app/config.py index 354c4f8..0525a32 100644 --- a/drb-c2-core/app/config.py +++ b/drb-c2-core/app/config.py @@ -133,12 +133,12 @@ class Settings(BaseSettings): # a failed LLM call falls back to the rules decision, so nothing broke # loudly. Verify against https://ai.google.dev/gemini-api/docs/models # before changing. - corr_cheap_model: str = "gemini-3.6-flash" # was gemini-2.0-flash (shut down) + corr_cheap_model: str = "gpt-4o-mini" # OpenAI since 2026-11-08 (was gemini-3.6-flash); a gemini-* ID still works # Transcript correction (server-26#36). Runs inside transcription, once per # transcribed call above MIN_WORDS_FOR_CORRECTION, so it is priced like STT # rather than like the correlation tier — cheap model on purpose. transcript_correction_enabled: bool = True - transcript_correction_model: str = "gemini-3.6-flash" + transcript_correction_model: str = "gpt-4o-mini" # OpenAI since 2026-11-08; a gemini-* ID still works # Retry Whisper once when its output is degenerate. The same clip produced a # 56-word ten-code counting run on one attempt and real speech on the next # (2026-08-23, call e49ea32c), so a hallucination is a coin-flip rather than diff --git a/drb-c2-core/app/internal/gemini.py b/drb-c2-core/app/internal/gemini.py index 5966d27..642e510 100644 --- a/drb-c2-core/app/internal/gemini.py +++ b/drb-c2-core/app/internal/gemini.py @@ -16,7 +16,6 @@ what it actually spent instead of an estimate. """ import json import threading -from contextvars import ContextVar from typing import Optional from app.config import settings @@ -27,18 +26,6 @@ _client_lock = threading.Lock() # Models that rejected a thinking level: retried without one from then on. _no_thinking_level: set[str] = set() -_usage_sink: ContextVar[Optional[dict]] = ContextVar("drb_gemini_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 _get_client(): global _client with _client_lock: @@ -71,14 +58,8 @@ 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 - logger.info(f"gemini usage {purpose} {model}: in={prompt} out={output} thinking={thoughts}") - 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"] += thoughts + from app.internal.llm import record_usage + record_usage(purpose, model, prompt, output, thoughts) def generate_json(model: str, prompt: str, *, purpose: str, diff --git a/drb-c2-core/app/internal/llm.py b/drb-c2-core/app/internal/llm.py new file mode 100644 index 0000000..c08e7b6 --- /dev/null +++ b/drb-c2-core/app/internal/llm.py @@ -0,0 +1,92 @@ +""" +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". + resp = openai_client().chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + response_format={"type": "json_object"}, + temperature=0, + ) + 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 "{}") diff --git a/drb-c2-core/app/internal/llm_correlator.py b/drb-c2-core/app/internal/llm_correlator.py index ce5be1f..1f0aa37 100644 --- a/drb-c2-core/app/internal/llm_correlator.py +++ b/drb-c2-core/app/internal/llm_correlator.py @@ -1,5 +1,5 @@ """ -LLM incident correlator (Gemini, settings.corr_cheap_model). +LLM incident correlator (settings.corr_cheap_model, via app/internal/llm.py). One function is exposed: decide(call_id, ctx) — called by routers/upload.py ONLY when the rules @@ -182,12 +182,12 @@ def _build_decide_prompt(ctx: dict, candidates: list[dict]) -> str: # ───────────────────────────────────────────────────────────────────────────── -# Gemini API call (sync, runs in thread pool) +# LLM call (sync, runs in thread pool) # ───────────────────────────────────────────────────────────────────────────── def _sync_gemini(model_name: str, prompt: str) -> dict: - from app.internal import gemini - return gemini.generate_json(model_name, prompt, purpose="correlation") + from app.internal import llm + return llm.generate_json(model_name, prompt, purpose="correlation") # ───────────────────────────────────────────────────────────────────────────── @@ -253,12 +253,13 @@ def _parse_response(raw: dict, ctx: dict, candidates: list[dict]) -> Optional[di async def decide(call_id: str, ctx: dict) -> Optional[dict]: """ One cheap-model correlation decision for an ambiguous scene. Returns None - (keep the rules decision) when GEMINI_API_KEY is unset, there is no - candidate to show, the answer fails validation, or Gemini fails. + (keep the rules decision) when the model's provider has no API key, there + is no candidate to show, the answer fails validation, or the call fails. """ from app.internal.incident_correlator import llm_candidates - if not settings.gemini_api_key: + from app.internal import llm + if not llm.available(settings.corr_cheap_model): return None candidates = llm_candidates(ctx) if not candidates: @@ -311,20 +312,23 @@ async def _log_llm_failure(where: str, tier: str, call_id: str, model: str, exc: # dead account never does. ai_health.classify() keeps a burst of rate limits # at WARNING while an empty account escalates like a bad model ID. if kind == "billing": - await _log_tier_down(where, tier, model, "the Gemini account is out of credit", - "Top up billing at https://ai.studio/projects", text) + from app.internal import llm + await _log_tier_down(where, tier, model, f"the {llm.provider_for(model)} account is out of credit", + llm.billing_fix(model), text) return logger.warning(f"{where} failed for call {call_id}: {text}") + from app.internal import llm await ai_health.report_degraded( - tier, "gemini", model, "transient API error", + tier, llm.provider_for(model), model, "transient API error", "no action needed unless this persists", permanent=False, ) async def _log_tier_down(where: str, tier: str, model: str, problem: str, fix: str, text: str) -> None: """ERROR once per model, not once per call — this runs at radio-traffic volume.""" - await ai_health.report_degraded(tier, "gemini", model, problem, fix, permanent=True) + from app.internal import llm + await ai_health.report_degraded(tier, llm.provider_for(model), model, problem, fix, permanent=True) if model in _dead_models: return _dead_models.add(model) diff --git a/drb-c2-core/app/internal/recorrelation_sweep.py b/drb-c2-core/app/internal/recorrelation_sweep.py index 9f57b88..94b9912 100644 --- a/drb-c2-core/app/internal/recorrelation_sweep.py +++ b/drb-c2-core/app/internal/recorrelation_sweep.py @@ -70,13 +70,13 @@ def pipeline_worst_case_seconds(n_scenes: int = _BOUND_SCENES) -> float: from app.internal.transcription import stt_targets stt_providers = max(1, len(stt_targets())) stt = stt_providers * stt_rounds * openai_attempts * settings.stt_timeout_seconds - # Transcript correction (inside transcribe_call). gemini.generate_json only - # re-issues a request when the model rejects a thinking level — a fast - # 400, not a timeout — and google-genai itself does not retry. - correction = settings.gemini_timeout_seconds + # Transcript correction and the correlation decision go through + # llm.generate_json; the bound depends on which provider serves the model. + from app.internal import llm + correction = llm.timeout_seconds(settings.transcript_correction_model) extraction = openai_attempts * settings.openai_timeout_seconds embeddings = n_scenes * openai_attempts * settings.openai_timeout_seconds - correlation = settings.gemini_timeout_seconds # at most one LLM call, scene 0 (#299) + correlation = llm.timeout_seconds(settings.corr_cheap_model) # at most one LLM call, scene 0 (#299) return stt + correction + extraction + embeddings + correlation + _NON_AI_SLACK_SECONDS diff --git a/drb-c2-core/app/internal/replay.py b/drb-c2-core/app/internal/replay.py index f74de6d..6acde7a 100644 --- a/drb-c2-core/app/internal/replay.py +++ b/drb-c2-core/app/internal/replay.py @@ -474,9 +474,9 @@ async def _run(run_id: str, org_id: str, calls: list[dict], mode: str, fl_token = force_flags(_flags_for(mode)) ai_failures: list = [] ai_token = ai_health.collect_sandbox_failures(ai_failures) - from app.internal import gemini + from app.internal import llm usage: dict = {} - usage_token = gemini.collect_usage(usage) + usage_token = llm.collect_usage(usage) try: sem = asyncio.Semaphore(PREFETCH) @@ -577,7 +577,7 @@ async def _run(run_id: str, org_id: str, calls: list[dict], mode: str, logger.error(f"Replay {run_id} failed: {e}") finally: ai_health._sandbox_failures.reset(ai_token) - gemini.reset_usage(usage_token) + llm.reset_usage(usage_token) unforce_flags(fl_token) fstore.exit_sandbox(sb_token) _cancel.discard(run_id) diff --git a/drb-c2-core/app/internal/transcript_correction.py b/drb-c2-core/app/internal/transcript_correction.py index 2f09fd4..b546f5b 100644 --- a/drb-c2-core/app/internal/transcript_correction.py +++ b/drb-c2-core/app/internal/transcript_correction.py @@ -253,11 +253,11 @@ def build_context_block(context: dict, talkgroup_name: Optional[str]) -> str: def _sync_gemini(model_name: str, prompt: str) -> dict: - from app.internal import gemini + from app.internal import llm # Correction rewrites text against vocabulary; keep a little reasoning # ("low") rather than the correlator's "minimal" until a replay shows # minimal doesn't hurt it. - return gemini.generate_json(model_name, prompt, purpose="correction", thinking_level="low") + return llm.generate_json(model_name, prompt, purpose="correction", thinking_level="low") async def correct( @@ -282,7 +282,8 @@ async def correct( so a correction that only fixed the joined transcript would never reach the model on exactly the multi-transmission calls that carry the most content. """ - if not settings.gemini_api_key or not settings.transcript_correction_enabled: + from app.internal import llm + if not llm.available(settings.transcript_correction_model) or not settings.transcript_correction_enabled: return None, None, False if len((text or "").split()) < MIN_WORDS_FOR_CORRECTION: diff --git a/drb-c2-core/tests/test_ai_timeouts.py b/drb-c2-core/tests/test_ai_timeouts.py index 140ddac..c7f11f7 100644 --- a/drb-c2-core/tests/test_ai_timeouts.py +++ b/drb-c2-core/tests/test_ai_timeouts.py @@ -164,7 +164,7 @@ async def test_gemini_timeout_in_decide_falls_back_to_rules(): def hang(model, prompt): raise httpx.ReadTimeout("The read operation timed out") - with patch.object(settings, "gemini_api_key", "test-key"), \ + with patch.object(settings, "openai_api_key", "test-key"), \ patch.object(llm_correlator, "_build_decide_prompt", lambda ctx, cands: "p"), \ patch.object(llm_correlator, "_sync_gemini", hang): assert await llm_correlator.decide("c1", _corr_ctx()) is None # None == keep the rules decision @@ -179,7 +179,7 @@ async def test_gemini_timeout_in_transcript_correction_keeps_raw_text(): raise httpx.ReadTimeout("The read operation timed out") text = "engine 5 respond to a structure fire at 12 main street smoke showing" - with patch.object(settings, "gemini_api_key", "test-key"), \ + with patch.object(settings, "openai_api_key", "test-key"), \ patch.object(transcript_correction, "resolve_context", AsyncMock(return_value={})), \ patch.object(transcript_correction, "_sync_gemini", hang): assert await transcript_correction.correct("c1", text) == (None, None, False) diff --git a/drb-c2-core/tests/test_gemini_helper.py b/drb-c2-core/tests/test_gemini_helper.py index 327f89e..707cde3 100644 --- a/drb-c2-core/tests/test_gemini_helper.py +++ b/drb-c2-core/tests/test_gemini_helper.py @@ -34,12 +34,13 @@ def test_minimal_thinking_by_default_and_usage_lands_in_the_sink(): models = _FakeModels() a, b = _patched(models) sink = {} - tok = gemini.collect_usage(sink) + from app.internal import llm + tok = llm.collect_usage(sink) try: with a, b: assert gemini.generate_json("m1", "p", purpose="correlation") == {"action": "link"} finally: - gemini.reset_usage(tok) + llm.reset_usage(tok) assert models.configs == [{"thinking_level": "minimal"}] assert sink == {"correlation:m1": {"calls": 1, "in": 1200, "out": 30, "thinking": 0}} diff --git a/drb-c2-core/tests/test_llm_provider.py b/drb-c2-core/tests/test_llm_provider.py new file mode 100644 index 0000000..b4f540f --- /dev/null +++ b/drb-c2-core/tests/test_llm_provider.py @@ -0,0 +1,48 @@ +"""app/internal/llm.py: provider follows the model name (2026-11-08 consolidation).""" +from types import SimpleNamespace +from unittest.mock import patch + +from app.config import settings +from app.internal import llm + + +def test_provider_follows_model_name(): + assert llm.provider_for("gpt-4o-mini") == "openai" + assert llm.provider_for("gemini-3.6-flash") == "gemini" + + +def test_available_checks_the_right_key(): + with patch.object(settings, "openai_api_key", "k"), patch.object(settings, "gemini_api_key", None): + assert llm.available("gpt-4o-mini") and not llm.available("gemini-3.6-flash") + + +def test_defaults_are_openai(): + assert llm.provider_for(settings.corr_cheap_model) == "openai" + assert llm.provider_for(settings.transcript_correction_model) == "openai" + + +def test_openai_path_parses_json_and_records_usage(): + resp = SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content='{"action": "new"}'))], + usage=SimpleNamespace(prompt_tokens=120, completion_tokens=9), + ) + calls = {} + + class FakeClient: + def __init__(self): + self.chat = SimpleNamespace(completions=SimpleNamespace(create=self.create)) + + def create(self, **kw): + calls.update(kw) + return resp + + sink = {} + tok = llm.collect_usage(sink) + try: + with patch("app.internal.ai_clients.openai_client", FakeClient): + out = llm.generate_json("gpt-4o-mini", "Return JSON", purpose="correlation") + finally: + llm.reset_usage(tok) + assert out == {"action": "new"} + assert calls["response_format"] == {"type": "json_object"} + assert sink["correlation:gpt-4o-mini"] == {"calls": 1, "in": 120, "out": 9, "thinking": 0} diff --git a/drb-c2-core/tests/test_transcript_correction.py b/drb-c2-core/tests/test_transcript_correction.py index dabd279..adbc94c 100644 --- a/drb-c2-core/tests/test_transcript_correction.py +++ b/drb-c2-core/tests/test_transcript_correction.py @@ -57,7 +57,7 @@ def _api_key(): The dev venv has no GEMINI_API_KEY, and correct() returns early without one — which would make every assertion below pass for the wrong reason. """ - with patch.object(tc.settings, "gemini_api_key", "test-key"): + with patch.object(tc.settings, "openai_api_key", "test-key"): yield