doctate/server/src/gazetteer/mod.rs

//! Gazetteer: domain-vocabulary lookup that injects correction hints
//! before the analysis LLM runs. See
//! `~/.claude/plans/happy-mapping-snail.md` for the design document.
//!
//! The LLM cannot know invented proper names (drug brands, anatomical
//! Latin/Greek terms). Whisper misrenders them phonetically
//! (`Vomex` → `Womax`). This module exposes a static lookup that, given
//! a transcript, finds tokens that are likely typos of a known entry
//! and rewrites them as `token [?Canonical]`. The LLM decides in context
//! whether to adopt the hint.

pub mod koelner;

use std::collections::{HashMap, HashSet};
use std::io;
use std::path::Path;

use strsim::damerau_levenshtein;

/// Minimum char count for *both* gazetteer entries and transcript tokens.
/// Calibrated against curated medication lists: 5 keeps prototypical
/// 5-char brand names (Vomex, Tavor, Lasix, Nasic) while still blocking
/// the 4-char deutsche Alltagswort ↔ eponym collisions that unkurated
/// crawls produce (Bein↔Behn, nach↔Nuck).
const MIN_TOKEN_LEN: usize = 5;

/// Maximum Damerau-Levenshtein distance for a candidate to qualify as
/// a typo of a gazetteer entry. Distance 0 is excluded — the `exact`
/// set short-circuits that case.
const MAX_EDIT_DISTANCE: usize = 2;

/// Domain-specific vocabulary, populated once at server start from
/// `*.txt` files under a vocab directory. Read-only thereafter; share
/// via `Arc<Gazetteer>`.
#[derive(Debug, Default)]
pub struct Gazetteer {
    /// Kölner-Phonetik code → canonical entries sharing that code. The
    /// canonicals keep their original casing for display.
    by_phonetic: HashMap<String, Vec<Box<str>>>,
    /// Lowercased canonicals; short-circuits annotation when the token
    /// is already spelled correctly.
    exact: HashSet<Box<str>>,
}

impl Gazetteer {
    /// Empty gazetteer. `annotate` is a pure pass-through on this instance.
    pub fn empty() -> Self {
        Self::default()
    }

    /// Load every `*.txt` file (non-recursive) under `dir`. Blank lines
    /// and lines beginning with `#` are skipped. Case-insensitive
    /// duplicates are silently collapsed.
    pub async fn load_dir(dir: &Path) -> io::Result<Self> {
        let mut g = Self::default();
        let mut entries = tokio::fs::read_dir(dir).await?;
        while let Some(entry) = entries.next_entry().await? {
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) != Some("txt") {
                continue;
            }
            let content = tokio::fs::read_to_string(&path).await?;
            for line in content.lines() {
                let trimmed = line.trim();
                if trimmed.is_empty() || trimmed.starts_with('#') {
                    continue;
                }
                g.insert(trimmed);
            }
        }
        Ok(g)
    }

    /// Number of unique entries (case-insensitive).
    pub fn len(&self) -> usize {
        self.exact.len()
    }

    pub fn is_empty(&self) -> bool {
        self.exact.is_empty()
    }

    fn insert(&mut self, entry: &str) {
        // Skip short entries — they cause noisy phonetic collisions with
        // ordinary German words. Same threshold as token matching.
        if entry.chars().count() < MIN_TOKEN_LEN {
            return;
        }
        let lower = entry.to_lowercase();
        if !self.exact.insert(lower.into_boxed_str()) {
            // Duplicate (case-insensitive) — don't add to phonetic index either.
            return;
        }
        let code = koelner::encode(entry);
        if code.is_empty() {
            return;
        }
        self.by_phonetic
            .entry(code)
            .or_default()
            .push(entry.into());
    }

    /// Annotate `text` by appending ` [?Canonical]` after any word-token
    /// that looks like a misspelling of a gazetteer entry. Non-matching
    /// tokens and all non-word segments (whitespace, punctuation, digits)
    /// pass through byte-identical.
    pub fn annotate(&self, text: &str) -> String {
        if self.is_empty() {
            return text.to_string();
        }
        let mut out = String::with_capacity(text.len());
        for seg in tokenize(text) {
            match seg {
                Segment::Word(w) => match self.find_candidate(w) {
                    Some(c) => {
                        out.push_str(w);
                        out.push_str(" [?");
                        out.push_str(c);
                        out.push(']');
                    }
                    None => out.push_str(w),
                },
                Segment::Other(s) => out.push_str(s),
            }
        }
        out
    }

    fn find_candidate(&self, token: &str) -> Option<&str> {
        if token.chars().count() < MIN_TOKEN_LEN {
            return None;
        }
        let lower = token.to_lowercase();
        if self.exact.contains(lower.as_str()) {
            return None;
        }

        // Stage 1: phonetic-indexed lookup. Fast, filters out the vast
        // majority of non-matches in O(1). Catches same-code neighbours
        // like Klexane↔Clexane, Diacepam↔Diazepam.
        let code = koelner::encode(token);
        if !code.is_empty()
            && let Some(candidates) = self.by_phonetic.get(&code)
            && let Some(hit) = best_candidate(&lower, candidates.iter().map(|c| c.as_ref()))
        {
            return Some(hit);
        }

        // Stage 2: brute-force edit-distance over every entry. Catches
        // ASR errors where Whisper substitutes phonetically-distinct
        // consonants (Carvenilol↔Carvedilol: N↔D) or drops a letter
        // that shifts the Kölner code (Acoxia↔Arcoxia: missing R).
        // O(n) per unknown token; at n≈200 entries and ~50 tokens per
        // transcript the absolute cost stays well under a millisecond.
        best_candidate(
            &lower,
            self.by_phonetic.values().flatten().map(|c| c.as_ref()),
        )
    }
}

/// Select the canonical entry from `candidates` that is closest to
/// `lower_token` by Damerau-Levenshtein distance, within the allowed
/// threshold. Distance 0 is rejected because that's already handled by
/// the `exact` short-circuit. Ties break alphabetically for determinism.
fn best_candidate<'a>(
    lower_token: &str,
    candidates: impl Iterator<Item = &'a str>,
) -> Option<&'a str> {
    candidates
        .map(|c| (damerau_levenshtein(lower_token, &c.to_lowercase()), c))
        .filter(|(d, _)| *d >= 1 && *d <= MAX_EDIT_DISTANCE)
        .min_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(b.1)))
        .map(|(_, c)| c)
}

/// A slice of the input text — either a word or a non-word run.
enum Segment<'a> {
    Word(&'a str),
    Other(&'a str),
}

/// Split `text` into alternating runs of alphabetic and non-alphabetic
/// chars. Round-trip: concatenating every segment yields the original
/// string byte-for-byte.
fn tokenize(text: &str) -> Vec<Segment<'_>> {
    let mut out = Vec::new();
    let mut chars = text.char_indices().peekable();
    while let Some((start, c)) = chars.next() {
        let is_alpha = c.is_alphabetic();
        let mut end = start + c.len_utf8();
        while let Some(&(next_start, next_c)) = chars.peek() {
            if next_c.is_alphabetic() != is_alpha {
                break;
            }
            end = next_start + next_c.len_utf8();
            chars.next();
        }
        let slice = &text[start..end];
        out.push(if is_alpha {
            Segment::Word(slice)
        } else {
            Segment::Other(slice)
        });
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    /// Build a gazetteer from an in-memory list of canonical entries —
    /// avoids filesystem I/O in the matcher tests.
    fn from_entries(entries: &[&str]) -> Gazetteer {
        let mut g = Gazetteer::empty();
        for e in entries {
            g.insert(e);
        }
        g
    }

    #[test]
    fn empty_gazetteer_is_pass_through() {
        let g = Gazetteer::empty();
        let input = "Patient nahm Zerebrum-Medikation.";
        assert_eq!(g.annotate(input), input);
    }

    /// Core regression: a token that's a phonetic neighbor at edit
    /// distance ≤ 2 gets a hint appended. Cerebrum and Zerebrum both
    /// encode to Kölner `87176` and differ by one char — a plausible
    /// Whisper-style transcription slip.
    #[test]
    fn annotates_typo_for_phonetic_neighbor() {
        let g = from_entries(&["Cerebrum"]);
        let out = g.annotate("Patient mit Blutung im Zerebrum.");
        assert!(
            out.contains("Zerebrum [?Cerebrum]"),
            "expected annotation, got: {out:?}"
        );
    }

    #[test]
    fn noop_on_exact_match() {
        let g = from_entries(&["Cerebrum"]);
        assert_eq!(
            g.annotate("Patient mit Blutung im Cerebrum."),
            "Patient mit Blutung im Cerebrum.",
        );
    }

    /// Lowercase variant of an entry is still "exact" — we never nag
    /// the LLM with a hint for a token that's only case-different.
    #[test]
    fn noop_on_case_only_difference() {
        let g = from_entries(&["Cerebrum"]);
        assert_eq!(
            g.annotate("Patient mit Blutung im cerebrum."),
            "Patient mit Blutung im cerebrum.",
        );
    }

    #[test]
    fn noop_on_unrelated_token() {
        let g = from_entries(&["Cerebrum"]);
        let input = "Patient hat Husten und Fieber.";
        assert_eq!(g.annotate(input), input);
    }

    /// Gazetteer entries shorter than MIN_TOKEN_LEN are rejected at load
    /// time. This prevents 4-char eponyms (Nuck, Behn, Bell, …) from
    /// generating noisy false-positive matches against everyday German
    /// words — the empirical reason MIN_TOKEN_LEN exists.
    #[test]
    fn rejects_gazetteer_entries_below_min_length() {
        let g = from_entries(&["Nuck", "Bein", "Behn"]);
        assert_eq!(g.len(), 0, "short entries must be dropped");
        // Input with a matching short word passes through untouched.
        let input = "Patient klagt nach schwerem Heben.";
        assert_eq!(g.annotate(input), input);
    }

    /// Edit-distance fallback: Whisper substitutes phonetically-distinct
    /// consonants (here N↔D in Carvenilol↔Carvedilol). The Kölner codes
    /// differ (N=6 vs D=2), so stage-1 misses — the brute-force stage
    /// must catch the distance-1 match.
    #[test]
    fn edit_distance_fallback_catches_consonant_substitution() {
        let g = from_entries(&["Carvedilol"]);
        let out = g.annotate("Patient nimmt Carvenilol.");
        assert!(
            out.contains("Carvenilol [?Carvedilol]"),
            "expected fallback annotation, got: {out:?}"
        );
    }

    /// Edit-distance fallback: Whisper drops a consonant (leading B in
    /// Berodual). Kölner codes are completely different — only the
    /// linear edit-distance scan saves the match.
    #[test]
    fn edit_distance_fallback_catches_dropped_consonant() {
        let g = from_entries(&["Berodual"]);
        let out = g.annotate("Patient bekommt Erodual.");
        assert!(
            out.contains("Erodual [?Berodual]"),
            "expected fallback annotation, got: {out:?}"
        );
    }

    /// With two phonetically-colliding candidates, the one with smaller
    /// Damerau-Levenshtein distance wins. Zerebrum↔Cerebrum = 1,
    /// Zerebrum↔Cerebrom = 2 → Cerebrum is chosen.
    #[test]
    fn picks_closest_candidate_by_edit_distance() {
        let g = from_entries(&["Cerebrum", "Cerebrom"]);
        let out = g.annotate("Patient mit Blutung im Zerebrum.");
        assert!(
            out.contains("Zerebrum [?Cerebrum]"),
            "expected Cerebrum as closer candidate, got: {out:?}"
        );
    }

    #[test]
    fn preserves_whitespace_and_punctuation() {
        let g = from_entries(&["Cerebrum"]);
        let input = "Line 1.\n\nLine 2: foo-bar; baz!";
        assert_eq!(g.annotate(input), input);
    }

    #[tokio::test]
    async fn load_dir_reads_all_txt_files() {
        let dir = tempdir().unwrap();
        tokio::fs::write(dir.path().join("meds.txt"), "Aspirin\nIbuprofen\n")
            .await
            .unwrap();
        tokio::fs::write(dir.path().join("anatomy.txt"), "Coracoideus\n")
            .await
            .unwrap();
        let g = Gazetteer::load_dir(dir.path()).await.unwrap();
        assert_eq!(g.len(), 3);
    }

    #[tokio::test]
    async fn load_dir_skips_blank_lines_and_comments() {
        let dir = tempdir().unwrap();
        tokio::fs::write(
            dir.path().join("v.txt"),
            "# medications\nAspirin\n\n   \n# another\nIbuprofen\n",
        )
        .await
        .unwrap();
        let g = Gazetteer::load_dir(dir.path()).await.unwrap();
        assert_eq!(g.len(), 2);
    }

    #[tokio::test]
    async fn load_dir_ignores_non_txt_files() {
        let dir = tempdir().unwrap();
        tokio::fs::write(dir.path().join("v.txt"), "Aspirin\n")
            .await
            .unwrap();
        tokio::fs::write(dir.path().join("readme.md"), "Ibuprofen\n")
            .await
            .unwrap();
        let g = Gazetteer::load_dir(dir.path()).await.unwrap();
        assert_eq!(g.len(), 1);
    }

    #[tokio::test]
    async fn load_dir_missing_dir_is_err() {
        let dir = tempdir().unwrap();
        let missing = dir.path().join("nonexistent");
        assert!(Gazetteer::load_dir(&missing).await.is_err());
    }
}