server-26/drb-c2-core/app/internal/transcript_correction.py

"""
Transcript correction — the second opinion on what was said.

Whisper hears a P25 vocoder through a narrowband channel and guesses at proper
nouns it has no reason to know: street names, business names, unit call signs.
It guesses confidently, so the output reads like speech and is wrong in exactly
the places that matter downstream — "Cool Parts, Illinois" and "Shout out to
Optum" both became incident locations.

Correction used to be a line in intelligence.py's EXTRACTION_PROMPT, which put
it in the wrong place twice over (server-26#36): the same model call that
extracted units, location and severity emitted the correction *afterwards*, so
extraction reasoned over uncorrected text; and it sat behind
`correlation_enabled`, so during a cost-controlled STT-only window nothing was
ever corrected at all. It belongs here, between transcription and everything
that consumes a transcript.

WHY A SEPARATE PASS AND NOT A WHISPER PROMPT: Whisper treats its prompt as
preceding transcript text and will happily continue a pattern it finds there —
an enumerated ten-code prompt made it emit "10-4. 10-5. 10-6. …" over silence
(see transcription.py). Vocabulary can never be a transcription prior. A
corrector that receives an already-produced transcript plus a reference list has
no series to extend; it can only substitute what it was given.

SCOPE RESOLUTION: reference data is merged from the talkgroup and the system,
**talkgroup first**. The specific beats the general — a system spanning several
counties may have one talkgroup covering a single municipality, and that
municipality's streets must not be buried under a county-wide list. A
single-municipality system is the degenerate case: populate the system level and
every talkgroup inherits it.
"""

import asyncio
import json
from typing import Any, Optional

from app.config import settings
from app.internal import firestore as fstore
from app.internal.logger import logger

# A transcript this short has no proper nouns to get wrong — "10-4.", "6-2,
# stand by." — and 9 of 29 calls in the 2026-08-23 sample sat at or under this.
# Skipping them is most of the cost saving for none of the value.
MIN_WORDS_FOR_CORRECTION = 4

_PROMPT = """You are correcting a police/fire radio transcript produced by an automatic speech recogniser.

The recogniser hears a low-bitrate vocoded radio channel. It reliably mishears proper nouns — street names, business names, town names, unit call signs — and substitutes common words that sound similar. Your job is to put back what was almost certainly said.

{context_block}
Rules:
- Change ONLY what is likely a mishearing. If a phrase is already plausible radio traffic, leave it exactly as it is.
- Prefer a name from the reference lists above when the transcript contains something that sounds like it. That is the entire point of this pass.
- NEVER add information. No new sentences, no invented units, no addresses that are not implied by the audio's own words.
- Keep radio language as radio language. Do NOT expand ten-codes or signals into plain English: "10-4" stays "10-4".
- Keep the speaker's structure and order. This is not a rewrite or a summary.
- If the text is clearly not speech at all — a counting run like "10-11. 10-12. 10-13.", or one phrase repeating many times over static — set not_speech to true.

Return JSON:
  corrected:  the corrected transcript, or null if nothing needed changing
  segments:   REQUIRED when numbered transmissions are given below — the corrected
              text for each one, as an array of exactly the same length and order.
              Never merge, split, reorder or drop a transmission; an unchanged one
              is returned verbatim. Omit this field entirely when no transmissions
              are numbered.
  not_speech: true if this is recogniser noise rather than a transmission
  changed:    list of ["heard" -> "corrected"] pairs you applied, for audit

{transcript}"""


def _render_input(text: str, segments: Optional[list[dict]]) -> str:
    """Numbered transmissions when we have them, so corrections stay aligned."""
    if segments and len(segments) > 1:
        lines = [f"{i + 1}. {s.get('text', '')}" for i, s in enumerate(segments)]
        body = "\n".join(lines)
        return f"Transmissions ({len(segments)}):\n{body}"
    return f"Transcript:\n{text}"


def _dedupe(items: list[str]) -> list[str]:
    """Preserve order, drop case-insensitive duplicates."""
    seen: set[str] = set()
    out: list[str] = []
    for item in items:
        key = (item or "").strip().lower()
        if key and key not in seen:
            seen.add(key)
            out.append(item.strip())
    return out


def _talkgroup_entry(system_doc: dict, talkgroup_id: Optional[int]) -> dict:
    """The config.talkgroups[] entry for this talkgroup, or {}."""
    if talkgroup_id is None:
        return {}
    try:
        wanted = int(talkgroup_id)
    except (TypeError, ValueError):
        return {}
    for tg in (system_doc.get("config") or {}).get("talkgroups", []) or []:
        try:
            if int(tg.get("id", -1)) == wanted:
                return tg
        except (TypeError, ValueError):
            continue
    return {}


def _area_lines(area: dict, label: str) -> list[str]:
    """Render an area_context dict as prompt lines. Empty when nothing is set."""
    if not area:
        return []
    parts: list[str] = []
    for key, heading in (
        ("municipality", "Municipality"),
        ("county", "County"),
        ("roads", "Roads and highways"),
        ("landmarks", "Landmarks, businesses and facilities"),
    ):
        value = area.get(key)
        if not value:
            continue
        if isinstance(value, list):
            value = ", ".join(str(v) for v in value if v)
        if value:
            parts.append(f"  {heading}: {value}")
    return [f"{label}:"] + parts if parts else []


async def resolve_context(system_id: Optional[str], talkgroup_id: Optional[int]) -> dict:
    """
    Merge the reference data a corrector needs, talkgroup ahead of system.

    Returns {"vocabulary": [...], "ten_codes": {...}, "area_lines": [...]}.
    Empty everywhere is legitimate — a system nobody has configured yet.
    """
    empty: dict[str, Any] = {"vocabulary": [], "ten_codes": {}, "area_lines": []}
    if not system_id:
        return empty

    system_doc = await fstore.doc_get_cached("systems", system_id)
    if not system_doc:
        return empty

    tg = _talkgroup_entry(system_doc, talkgroup_id)

    # Talkgroup terms first so they survive any downstream truncation.
    vocabulary = _dedupe(
        list(tg.get("vocabulary") or []) + list(system_doc.get("vocabulary") or [])
    )

    # Ten-codes: system-wide reference, with talkgroup entries overriding a
    # code that means something different on this channel.
    ten_codes = dict(system_doc.get("ten_codes") or {})
    ten_codes.update(tg.get("ten_codes") or {})

    area_lines = (
        _area_lines(tg.get("area_context") or {}, "Area covered by this talkgroup")
        + _area_lines(system_doc.get("area_context") or {}, "Area covered by this system")
    )

    return {"vocabulary": vocabulary, "ten_codes": ten_codes, "area_lines": area_lines}


def build_context_block(context: dict, talkgroup_name: Optional[str]) -> str:
    """Render resolved context into the prompt's reference section."""
    lines: list[str] = []
    if talkgroup_name:
        lines.append(f"Channel: {talkgroup_name}")
    lines.extend(context.get("area_lines") or [])
    vocabulary = context.get("vocabulary") or []
    if vocabulary:
        lines.append("Known local names and terms: " + ", ".join(vocabulary))
    ten_codes = context.get("ten_codes") or {}
    if ten_codes:
        rendered = ", ".join(f"{code}={meaning}" for code, meaning in sorted(ten_codes.items()))
        lines.append(f"Ten-codes used on this system: {rendered}")
    return ("\n".join(lines) + "\n") if lines else ""


def _sync_gemini(model_name: str, prompt: str) -> dict:
    import google.generativeai as genai  # lazy import — only when needed

    genai.configure(api_key=settings.gemini_api_key)
    model = genai.GenerativeModel(
        model_name,
        generation_config={"response_mime_type": "application/json"},
    )
    return json.loads(model.generate_content(prompt).text)


async def correct(
    call_id: str,
    text: str,
    segments: Optional[list[dict]] = None,
    system_id: Optional[str] = None,
    talkgroup_id: Optional[int] = None,
    talkgroup_name: Optional[str] = None,
) -> tuple[Optional[str], Optional[list[dict]], bool]:
    """
    Second-opinion pass over a transcript.

    Returns (corrected_text, corrected_segments, not_speech). ``None`` for
    either correction means "no change" — the corrector found nothing to fix,
    could not run, or returned segments that did not line up. Callers keep the
    original in that case; correction is an improvement, never a dependency.

    Segments matter as much as the joined text: intelligence.py builds its
    extraction prompt from NUMBERED SEGMENTS whenever there is more than one,
    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:
        return None, None, False

    if len((text or "").split()) < MIN_WORDS_FOR_CORRECTION:
        return None, None, False

    context = await resolve_context(system_id, talkgroup_id)
    prompt = _PROMPT.format(
        context_block=build_context_block(context, talkgroup_name),
        transcript=_render_input(text, segments),
    )

    try:
        raw = await asyncio.to_thread(
            _sync_gemini, settings.transcript_correction_model, prompt
        )
    except Exception as e:
        # Never fail the transcript over a failed correction — the raw text is
        # still worth having. ai_health reporting is the caller's business.
        logger.warning(f"Transcript correction failed for call {call_id}: {e}")
        return None, None, False

    not_speech = bool(raw.get("not_speech"))

    corrected = raw.get("corrected")
    if not isinstance(corrected, str) or not corrected.strip():
        corrected = None
    elif corrected.strip() == (text or "").strip():
        corrected = None

    # Segment alignment is non-negotiable: scene extraction maps scenes back to
    # transmissions by INDEX (segment_indices), so a returned array of the wrong
    # length would silently attribute the wrong audio to a scene. Wrong length,
    # wrong type, or any non-string entry and the segments are discarded whole —
    # the joined correction still stands.
    corrected_segments: Optional[list[dict]] = None
    if segments and len(segments) > 1:
        returned = raw.get("segments")
        if (
            isinstance(returned, list)
            and len(returned) == len(segments)
            and all(isinstance(x, str) for x in returned)
        ):
            corrected_segments = [
                {**seg, "text": new.strip() or seg.get("text", "")}
                for seg, new in zip(segments, returned)
            ]
            if all(s["text"] == o.get("text") for s, o in zip(corrected_segments, segments)):
                corrected_segments = None
        elif returned is not None:
            logger.warning(
                f"Transcript correction for call {call_id} returned "
                f"{len(returned) if isinstance(returned, list) else type(returned).__name__} "
                f"segment(s) against {len(segments)} — discarding segment corrections"
            )

    if corrected or corrected_segments or not_speech:
        changed = raw.get("changed") or []
        logger.info(
            f"Transcript correction ({settings.transcript_correction_model}): call {call_id} "
            f"not_speech={not_speech} segments={'yes' if corrected_segments else 'no'} "
            f"changes={changed if isinstance(changed, list) else '?'}"
        )
    return corrected, corrected_segments, not_speech