kcode-speaker-system 0.2.0

#![forbid(unsafe_code)]

mod ktool;

pub use kcode_speaker_types::{FEATURE_COUNT, FeatureVector};
pub use ktool::{
    DELETE_TOOL, IDENTIFY_TOOL, KTOOLS, KtoolCall, KtoolError, TRAIN_TOOL, decode_ktool,
};

use kcode_speaker_model::{
    FitRowsInput, ModelConfig, ModelError, ModelSnapshot, TrainingRow, fit_rows,
    identify as identify_model,
};
use kcode_speaker_types::{FeatureMask, Key};
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, HashMap};
use std::fmt;
use std::path::Path;
use std::sync::{Mutex, MutexGuard};
use std::time::Duration;

/// The frozen 24-value feature row used by the classifier.
pub type FeatureRow = FeatureVector;

/// One caller-owned observation identity.
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct ObservationKey {
    pub object_id: String,
    pub piece_index: u32,
}

impl ObservationKey {
    fn validate(&self) -> Result<(), Error> {
        validate_nonempty("key.object_id", &self.object_id)
    }
}

/// The provider and feature-contract boundary within which rows may be compared.
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct Cohort {
    pub provider: String,
    pub model: String,
    pub prompt_version: String,
    pub schema_version: String,
    pub primary_language: String,
}

impl Cohort {
    fn validate(&self) -> Result<(), Error> {
        validate_nonempty("cohort.provider", &self.provider)?;
        validate_nonempty("cohort.model", &self.model)?;
        validate_nonempty("cohort.prompt_version", &self.prompt_version)?;
        validate_nonempty("cohort.schema_version", &self.schema_version)?;
        if self.primary_language.len() != 3
            || !self
                .primary_language
                .bytes()
                .all(|byte| byte.is_ascii_lowercase())
        {
            return Err(Error::validation(
                "cohort.primary_language",
                "must be a lowercase ISO 639-3 code",
            ));
        }
        Ok(())
    }
}

/// One ranked classifier candidate. Lower cost is better.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct CandidateEvidence {
    pub speaker_id: String,
    pub cost: f64,
}

/// Evidence supporting an identify result.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct IdentifyEvidence {
    pub best: CandidateEvidence,
    pub runner_up: Option<CandidateEvidence>,
    pub background_population_cost: f64,
    pub absolute_gap: f64,
    pub runner_up_gap: Option<f64>,
    pub confidence_score: f64,
}

/// The accepted speaker, if any, and the available ranking evidence.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct IdentifyOutcome {
    pub speaker_id: Option<String>,
    pub evidence: Option<IdentifyEvidence>,
}

/// Result of adding or correcting labelled training data.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum TrainOutcome {
    Added,
    Unchanged,
    Corrected,
}

/// Result of idempotently deleting an observation.
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum DeleteOutcome {
    Deleted,
    NotFound,
}

/// A classifier operation failure.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Error {
    Validation {
        field: String,
        message: String,
    },
    Conflict {
        key: ObservationKey,
        message: String,
    },
    Storage(String),
    CorruptStorage(String),
    Model(String),
}

impl Error {
    fn validation(field: impl Into<String>, message: impl Into<String>) -> Self {
        Self::Validation {
            field: field.into(),
            message: message.into(),
        }
    }

    fn conflict(key: &ObservationKey) -> Self {
        Self::Conflict {
            key: key.clone(),
            message:
                "observation key already has conflicting data or assignment; use train or delete"
                    .to_owned(),
        }
    }

    fn corrupt(message: impl Into<String>) -> Self {
        Self::CorruptStorage(message.into())
    }

    /// Stable error category for callers that need machine-readable handling.
    pub const fn code(&self) -> &'static str {
        match self {
            Self::Validation { .. } => "validation",
            Self::Conflict { .. } => "conflict",
            Self::Storage(_) => "storage",
            Self::CorruptStorage(_) => "corrupt_storage",
            Self::Model(_) => "model",
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Validation { field, message } => write!(formatter, "{field}: {message}"),
            Self::Conflict { key, message } => {
                write!(
                    formatter,
                    "{}:{}: {message}",
                    key.object_id, key.piece_index
                )
            }
            Self::Storage(message) => write!(formatter, "storage error: {message}"),
            Self::CorruptStorage(message) => write!(formatter, "corrupt storage: {message}"),
            Self::Model(message) => write!(formatter, "speaker model error: {message}"),
        }
    }
}

impl std::error::Error for Error {}

impl From<rusqlite::Error> for Error {
    fn from(error: rusqlite::Error) -> Self {
        Self::Storage(error.to_string())
    }
}

impl From<ModelError> for Error {
    fn from(error: ModelError) -> Self {
        Self::Model(error.to_string())
    }
}

/// A thread-safe classifier backed by one SQLite database.
pub struct SpeechClassifier {
    connection: Mutex<Connection>,
    models: Mutex<HashMap<String, CachedCohort>>,
}

/// Descriptive alias for the classifier owned by this package.
pub type SpeakerSystem = SpeechClassifier;

impl SpeechClassifier {
    /// Opens or creates the classifier's current-row store.
    pub fn open(path: impl AsRef<Path>) -> Result<Self, Error> {
        let connection = Connection::open(path)?;
        connection.busy_timeout(Duration::from_secs(30))?;
        connection.execute_batch(
            "PRAGMA journal_mode = WAL;
             PRAGMA synchronous = FULL;
             CREATE TABLE IF NOT EXISTS speaker_observations_v2 (
                 object_id TEXT NOT NULL,
                 piece_index INTEGER NOT NULL,
                 cohort_json TEXT NOT NULL,
                 features BLOB NOT NULL,
                 speaker_id TEXT NOT NULL,
                 PRIMARY KEY (object_id, piece_index)
             );
             CREATE INDEX IF NOT EXISTS speaker_observations_v2_cohort
                 ON speaker_observations_v2 (cohort_json);",
        )?;
        Ok(Self {
            connection: Mutex::new(connection),
            models: Mutex::new(HashMap::new()),
        })
    }

    /// Scores an observation and atomically retains it when the threshold is met.
    pub fn identify(
        &self,
        key: ObservationKey,
        cohort: Cohort,
        row: FeatureRow,
        threshold: f64,
    ) -> Result<IdentifyOutcome, Error> {
        key.validate()?;
        cohort.validate()?;
        if !threshold.is_finite() {
            return Err(Error::validation("threshold", "must be finite"));
        }

        let cohort_json = encode_cohort(&cohort)?;
        let mut connection = self.lock()?;
        let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;

        if let Some(existing) = load_observation(&transaction, &key)? {
            if existing.cohort_json != cohort_json || existing.features != row {
                return Err(Error::conflict(&key));
            }
            transaction.commit()?;
            return Ok(IdentifyOutcome {
                speaker_id: Some(existing.speaker_id),
                evidence: None,
            });
        }

        let mut models = self
            .models
            .lock()
            .map_err(|_| Error::Model("classifier model cache lock was poisoned".to_owned()))?;
        if !models.contains_key(&cohort_json) {
            let training = load_cohort(&transaction, &cohort_json)?;
            models.insert(cohort_json.clone(), fit_cohort(&training, &cohort_json)?);
        }
        let outcome = classify(
            models
                .get(&cohort_json)
                .ok_or_else(|| Error::Model("fitted cohort was not cached".to_owned()))?,
            &row,
            threshold,
        )?;
        if let Some(speaker_id) = &outcome.speaker_id {
            insert_observation(&transaction, &key, &cohort_json, &row, speaker_id)?;
            models.clear();
        }
        transaction.commit()?;
        Ok(outcome)
    }

    /// Adds known training data or atomically corrects the current assignment.
    pub fn train(
        &self,
        key: ObservationKey,
        cohort: Cohort,
        row: FeatureRow,
        speaker_id: String,
    ) -> Result<TrainOutcome, Error> {
        key.validate()?;
        cohort.validate()?;
        validate_nonempty("speaker_id", &speaker_id)?;

        let cohort_json = encode_cohort(&cohort)?;
        let mut connection = self.lock()?;
        let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?;
        let existing = load_observation(&transaction, &key)?;
        let outcome = match existing {
            None => TrainOutcome::Added,
            Some(existing)
                if existing.cohort_json == cohort_json
                    && existing.features == row
                    && existing.speaker_id == speaker_id =>
            {
                TrainOutcome::Unchanged
            }
            Some(_) => TrainOutcome::Corrected,
        };
        if outcome != TrainOutcome::Unchanged {
            transaction.execute(
                "INSERT INTO speaker_observations_v2
                     (object_id, piece_index, cohort_json, features, speaker_id)
                 VALUES (?1, ?2, ?3, ?4, ?5)
                 ON CONFLICT(object_id, piece_index) DO UPDATE SET
                     cohort_json = excluded.cohort_json,
                     features = excluded.features,
                     speaker_id = excluded.speaker_id",
                params![
                    key.object_id,
                    i64::from(key.piece_index),
                    cohort_json,
                    row.as_ref().as_slice(),
                    speaker_id,
                ],
            )?;
        }
        transaction.commit()?;
        if outcome != TrainOutcome::Unchanged {
            self.clear_models()?;
        }
        Ok(outcome)
    }

    /// Idempotently deletes the current observation for a key.
    pub fn delete(&self, key: ObservationKey) -> Result<DeleteOutcome, Error> {
        key.validate()?;
        let connection = self.lock()?;
        let changed = connection.execute(
            "DELETE FROM speaker_observations_v2
             WHERE object_id = ?1 AND piece_index = ?2",
            params![key.object_id, i64::from(key.piece_index)],
        )?;
        if changed != 0 {
            self.clear_models()?;
        }
        Ok(if changed == 0 {
            DeleteOutcome::NotFound
        } else {
            DeleteOutcome::Deleted
        })
    }

    fn lock(&self) -> Result<MutexGuard<'_, Connection>, Error> {
        self.connection
            .lock()
            .map_err(|_| Error::Storage("classifier database lock was poisoned".to_owned()))
    }

    fn clear_models(&self) -> Result<(), Error> {
        self.models
            .lock()
            .map_err(|_| Error::Model("classifier model cache lock was poisoned".to_owned()))?
            .clear();
        Ok(())
    }
}

#[derive(Debug)]
struct StoredObservation {
    key: ObservationKey,
    cohort_json: String,
    features: FeatureRow,
    speaker_id: String,
}

struct CachedCohort {
    cohort_id: Key,
    names: BTreeMap<Key, String>,
    model: Option<ModelSnapshot>,
}

fn load_observation(
    connection: &Connection,
    key: &ObservationKey,
) -> Result<Option<StoredObservation>, Error> {
    connection
        .query_row(
            "SELECT cohort_json, features, speaker_id
             FROM speaker_observations_v2
             WHERE object_id = ?1 AND piece_index = ?2",
            params![key.object_id, i64::from(key.piece_index)],
            |record| {
                Ok((
                    record.get::<_, String>(0)?,
                    record.get::<_, Vec<u8>>(1)?,
                    record.get::<_, String>(2)?,
                ))
            },
        )
        .optional()?
        .map(|(cohort_json, features, speaker_id)| {
            Ok(StoredObservation {
                key: key.clone(),
                cohort_json,
                features: decode_features(&features)?,
                speaker_id,
            })
        })
        .transpose()
}

fn load_cohort(
    connection: &Connection,
    cohort_json: &str,
) -> Result<Vec<StoredObservation>, Error> {
    let mut statement = connection.prepare(
        "SELECT object_id, piece_index, features, speaker_id
         FROM speaker_observations_v2
         WHERE cohort_json = ?1
         ORDER BY object_id, piece_index",
    )?;
    let records = statement.query_map([cohort_json], |record| {
        Ok((
            record.get::<_, String>(0)?,
            record.get::<_, u32>(1)?,
            record.get::<_, Vec<u8>>(2)?,
            record.get::<_, String>(3)?,
        ))
    })?;

    let mut observations = Vec::new();
    for record in records {
        let (object_id, piece_index, features, speaker_id) = record?;
        observations.push(StoredObservation {
            key: ObservationKey {
                object_id,
                piece_index,
            },
            cohort_json: cohort_json.to_owned(),
            features: decode_features(&features)?,
            speaker_id,
        });
    }
    Ok(observations)
}

fn insert_observation(
    connection: &Connection,
    key: &ObservationKey,
    cohort_json: &str,
    features: &FeatureRow,
    speaker_id: &str,
) -> Result<(), Error> {
    connection.execute(
        "INSERT INTO speaker_observations_v2
             (object_id, piece_index, cohort_json, features, speaker_id)
         VALUES (?1, ?2, ?3, ?4, ?5)",
        params![
            key.object_id,
            i64::from(key.piece_index),
            cohort_json,
            features.as_ref().as_slice(),
            speaker_id,
        ],
    )?;
    Ok(())
}

fn fit_cohort(
    observations: &[StoredObservation],
    cohort_json: &str,
) -> Result<CachedCohort, Error> {
    let cohort_id = digest_key("cohort", &[cohort_json.as_bytes()])?;
    let mut names = BTreeMap::new();
    let mut rows = Vec::with_capacity(observations.len());
    for observation in observations {
        let speaker_key = digest_key("speaker", &[observation.speaker_id.as_bytes()])?;
        names.insert(speaker_key.clone(), observation.speaker_id.clone());
        rows.push(TrainingRow {
            sample_id: digest_key(
                "sample",
                &[
                    observation.key.object_id.as_bytes(),
                    &observation.key.piece_index.to_be_bytes(),
                ],
            )?,
            speaker_id: speaker_key,
            cohort_id: cohort_id.clone(),
            features: observation.features,
        });
    }

    let model = varying_features(observations)
        .map(|mask| {
            fit_rows(FitRowsInput {
                cohort_id: &cohort_id,
                rows: &rows,
                config: ModelConfig {
                    mask,
                    components: 1,
                    relevance: 2.0,
                    variance_floor: 0.01,
                    absolute_threshold: 0.0,
                    margin_threshold: 0.0,
                },
            })
        })
        .transpose()?;
    Ok(CachedCohort {
        cohort_id,
        names,
        model,
    })
}

fn classify(
    cached: &CachedCohort,
    probe: &FeatureRow,
    threshold: f64,
) -> Result<IdentifyOutcome, Error> {
    let Some(model) = cached.model.as_ref() else {
        return Ok(IdentifyOutcome {
            speaker_id: None,
            evidence: None,
        });
    };
    let identification = identify_model(model, &cached.cohort_id, probe)?;
    let accepted = identification.best.llr >= threshold
        && identification
            .runner_up
            .as_ref()
            .is_none_or(|candidate| identification.best.llr - candidate.llr >= threshold);
    let speaker_id = if accepted {
        Some(
            cached
                .names
                .get(&identification.best.speaker_id)
                .cloned()
                .ok_or_else(|| Error::corrupt("model returned an unknown speaker key"))?,
        )
    } else {
        None
    };
    let best_name = cached
        .names
        .get(&identification.best.speaker_id)
        .cloned()
        .ok_or_else(|| Error::corrupt("model returned an unknown best speaker key"))?;
    let runner_up = identification
        .runner_up
        .as_ref()
        .map(|candidate| {
            let speaker_id = cached
                .names
                .get(&candidate.speaker_id)
                .cloned()
                .ok_or_else(|| Error::corrupt("model returned an unknown runner-up speaker key"))?;
            Ok::<_, Error>(CandidateEvidence {
                speaker_id,
                cost: -candidate.llr,
            })
        })
        .transpose()?;
    let absolute_gap = identification.best.llr;
    let runner_up_gap = identification
        .runner_up
        .as_ref()
        .map(|candidate| identification.best.llr - candidate.llr);
    let confidence_score = runner_up_gap.map_or(absolute_gap, |gap| absolute_gap.min(gap));

    Ok(IdentifyOutcome {
        speaker_id,
        evidence: Some(IdentifyEvidence {
            best: CandidateEvidence {
                speaker_id: best_name,
                cost: -identification.best.llr,
            },
            runner_up,
            background_population_cost: 0.0,
            absolute_gap,
            runner_up_gap,
            confidence_score,
        }),
    })
}

fn varying_features(observations: &[StoredObservation]) -> Option<FeatureMask> {
    let first = observations.first()?.features.as_ref();
    let mut bits = 0_u64;
    for (index, first_value) in first.iter().enumerate() {
        if observations
            .iter()
            .skip(1)
            .any(|observation| observation.features.as_ref()[index] != *first_value)
        {
            bits |= 1_u64 << index;
        }
    }
    FeatureMask::from_bits(bits).ok()
}

fn digest_key(prefix: &str, fields: &[&[u8]]) -> Result<Key, Error> {
    let mut digest = Sha256::new();
    for field in fields {
        digest.update((field.len() as u64).to_be_bytes());
        digest.update(field);
    }
    let hex = format!("{prefix}:{:x}", digest.finalize());
    Key::parse(&hex).map_err(|error| Error::Model(error.to_string()))
}

fn encode_cohort(cohort: &Cohort) -> Result<String, Error> {
    serde_json::to_string(cohort).map_err(|error| Error::Storage(error.to_string()))
}

fn decode_features(bytes: &[u8]) -> Result<FeatureRow, Error> {
    let values: [u8; FEATURE_COUNT] = bytes
        .try_into()
        .map_err(|_| Error::corrupt("stored feature row does not contain exactly 24 values"))?;
    FeatureRow::new(values).map_err(|error| Error::corrupt(error.to_string()))
}

fn validate_nonempty(field: &str, value: &str) -> Result<(), Error> {
    if value.trim().is_empty() {
        Err(Error::validation(field, "must not be empty"))
    } else {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::path::{Path, PathBuf};
    use std::sync::atomic::{AtomicU64, Ordering};

    static NEXT_PATH: AtomicU64 = AtomicU64::new(0);

    fn database_path(label: &str) -> PathBuf {
        std::env::temp_dir().join(format!(
            "kcode-speaker-system-{label}-{}-{}.sqlite3",
            std::process::id(),
            NEXT_PATH.fetch_add(1, Ordering::Relaxed)
        ))
    }

    fn cleanup(path: &Path) {
        for suffix in ["", "-wal", "-shm"] {
            let _ = fs::remove_file(format!("{}{}", path.display(), suffix));
        }
    }

    fn key(value: &str) -> ObservationKey {
        ObservationKey {
            object_id: value.to_owned(),
            piece_index: 0,
        }
    }

    fn cohort() -> Cohort {
        Cohort {
            provider: "google".to_owned(),
            model: "gemini-example".to_owned(),
            prompt_version: "speaker-24-1".to_owned(),
            schema_version: "speaker-24-1".to_owned(),
            primary_language: "eng".to_owned(),
        }
    }

    fn row(seed: u8) -> FeatureRow {
        let mut values = [0_u8; FEATURE_COUNT];
        for (index, value) in values.iter_mut().enumerate() {
            *value = (seed + u8::try_from(index).unwrap()) % 100;
        }
        FeatureRow::new(values).unwrap()
    }

    #[test]
    fn train_correct_delete_and_reopen_use_only_current_rows() {
        let path = database_path("mutations");
        let classifier = SpeechClassifier::open(&path).unwrap();
        assert_eq!(
            classifier
                .train(key("one"), cohort(), row(10), "Alice Example".to_owned())
                .unwrap(),
            TrainOutcome::Added
        );
        assert_eq!(
            classifier
                .train(key("one"), cohort(), row(10), "Alice Example".to_owned())
                .unwrap(),
            TrainOutcome::Unchanged
        );
        assert_eq!(
            classifier
                .train(key("one"), cohort(), row(11), "Alice Corrected".to_owned())
                .unwrap(),
            TrainOutcome::Corrected
        );
        drop(classifier);

        let classifier = SpeechClassifier::open(&path).unwrap();
        assert_eq!(
            classifier.delete(key("one")).unwrap(),
            DeleteOutcome::Deleted
        );
        assert_eq!(
            classifier.delete(key("one")).unwrap(),
            DeleteOutcome::NotFound
        );
        drop(classifier);
        cleanup(&path);
    }

    #[test]
    fn identify_is_cohort_scoped_thresholded_and_retained() {
        let path = database_path("identify");
        let classifier = SpeechClassifier::open(&path).unwrap();
        for (id, seed, speaker) in [
            ("alice-1", 10, "Alice Example"),
            ("alice-2", 12, "Alice Example"),
            ("bob-1", 70, "Bob Example"),
            ("bob-2", 72, "Bob Example"),
        ] {
            classifier
                .train(key(id), cohort(), row(seed), speaker.to_owned())
                .unwrap();
        }

        let rejected = classifier
            .identify(key("rejected"), cohort(), row(11), f64::MAX)
            .unwrap();
        assert_eq!(rejected.speaker_id, None);
        assert!(rejected.evidence.is_some());
        assert_eq!(classifier.models.lock().unwrap().len(), 1);

        let accepted = classifier
            .identify(key("accepted"), cohort(), row(11), -1_000_000.0)
            .unwrap();
        assert_eq!(accepted.speaker_id.as_deref(), Some("Alice Example"));
        assert!(classifier.models.lock().unwrap().is_empty());
        assert_eq!(
            classifier
                .identify(key("accepted"), cohort(), row(11), f64::MAX)
                .unwrap()
                .speaker_id
                .as_deref(),
            Some("Alice Example")
        );

        let mut other = cohort();
        other.prompt_version = "other".to_owned();
        assert_eq!(
            classifier
                .identify(key("isolated"), other, row(11), -1_000_000.0)
                .unwrap(),
            IdentifyOutcome {
                speaker_id: None,
                evidence: None,
            }
        );
        drop(classifier);
        cleanup(&path);
    }

    #[test]
    fn public_boundary_rejects_bad_keys_labels_and_thresholds() {
        let path = database_path("validation");
        let classifier = SpeechClassifier::open(&path).unwrap();
        assert!(matches!(
            classifier.delete(key(" ")),
            Err(Error::Validation { .. })
        ));
        assert!(matches!(
            classifier.train(key("one"), cohort(), row(1), " ".to_owned()),
            Err(Error::Validation { .. })
        ));
        assert!(matches!(
            classifier.identify(key("one"), cohort(), row(1), f64::NAN),
            Err(Error::Validation { .. })
        ));
        drop(classifier);
        cleanup(&path);
    }
}