#![forbid(unsafe_code)]
use kcode_speaker_types::{FeatureVector, Key, LabeledSample, ObjectId, RecordingKind, SegmentRef};
use rusqlite::{Connection, params};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::{
collections::{BTreeMap, BTreeSet},
error::Error,
fmt,
path::Path,
sync::Mutex,
};
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct EventEnvelope {
pub event_id: Key,
pub event: Event,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum Event {
RegisterSegment(SegmentRegistration),
RecordAttempt(AttemptRecord),
SelectAttempt(AttemptSelection),
SetSampleState(SampleStateChange),
PutArtifact(ArtifactRecord),
SetActiveModel(ActiveModelChange),
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct SegmentRegistration {
pub segment: SegmentRef,
pub group_id: Key,
pub recording_kind: RecordingKind,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct AttemptRecord {
pub attempt_id: Key,
pub clip_object: ObjectId,
pub extraction_key: Key,
pub provider_result_object: Option<ObjectId>,
pub outcome: StoredAttemptOutcome,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum StoredAttemptOutcome {
Scored {
recording_quality: u8,
speakers: Vec<StoredSpeaker>,
additional_speakers: Vec<StoredAdditionalSpeaker>,
},
Unscorable {
reason: String,
additional_speakers: Vec<StoredAdditionalSpeaker>,
},
InvalidResponse {
error: String,
},
ProviderFailure {
code: Key,
detail: String,
},
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct StoredSpeaker {
pub speaker_ordinal: u16,
pub primary_language: Key,
pub closest_dialect: String,
pub usable_speech_ms: u32,
pub features: FeatureVector,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct StoredAdditionalSpeaker {
pub speaker_ordinal: u16,
pub description: String,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct AttemptSelection {
pub clip_object: ObjectId,
pub extraction_key: Key,
pub attempt_id: Key,
pub reason: String,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct SampleStateChange {
pub sample_id: Key,
pub state: SampleState,
pub reason: String,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum SampleState {
Unlabeled,
Confirmed { speaker_id: Key },
Retracted,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ArtifactRecord {
pub artifact_id: Key,
pub cohort_id: Key,
pub kind: ArtifactKind,
pub sha256: [u8; 32],
pub bytes: Vec<u8>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum ArtifactKind {
Model,
Evaluation,
ReplayManifest,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ActiveModelChange {
pub cohort_id: Key,
pub model_artifact_id: Option<Key>,
pub reason: String,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum Request {
Commit(Vec<EventEnvelope>),
Read(Query),
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum Query {
Attempt {
attempt_id: Key,
},
Attempts {
extraction_key: Option<Key>,
source_object: Option<ObjectId>,
},
TrainingRows {
cohort_id: Key,
},
RepeatabilityRows {
cohort_id: Key,
},
Artifact {
artifact_id: Key,
},
ActiveModel {
cohort_id: Key,
},
Inventory {
extraction_key: Option<Key>,
},
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct Response {
pub revision: u64,
pub result: ResponseKind,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub enum ResponseKind {
Commit { applied: u64 },
Attempt(Option<AttemptSnapshot>),
Attempts(Vec<AttemptSnapshot>),
TrainingRows(Vec<LabeledSample>),
RepeatabilityRows(Vec<LabeledSample>),
Artifact(Option<ArtifactRecord>),
ActiveModel(Option<Key>),
Inventory(InventorySnapshot),
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct AttemptSnapshot {
pub attempt: AttemptRecord,
pub sample_ids: Vec<Key>,
pub selected: bool,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct SelectionSnapshot {
pub selection: AttemptSelection,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct SampleStateSnapshot {
pub change: SampleStateChange,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ArtifactMetadata {
pub artifact_id: Key,
pub cohort_id: Key,
pub kind: ArtifactKind,
pub sha256: [u8; 32],
pub byte_len: u64,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct ActiveModelSnapshot {
pub change: ActiveModelChange,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct MissingSegmentOrdinals {
pub source_object: ObjectId,
pub ordinals: Vec<u16>,
}
#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
pub struct OutcomeCounts {
pub scored: u64,
pub unscorable: u64,
pub invalid_response: u64,
pub provider_failure: u64,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct InventorySnapshot {
pub registrations: Vec<SegmentRegistration>,
pub attempts: Vec<AttemptSnapshot>,
pub selections: Vec<SelectionSnapshot>,
pub sample_states: Vec<SampleStateSnapshot>,
pub artifacts: Vec<ArtifactMetadata>,
pub active_models: Vec<ActiveModelSnapshot>,
pub missing_segment_ordinals: Vec<MissingSegmentOrdinals>,
pub outcome_counts: OutcomeCounts,
}
#[derive(Debug)]
pub enum StoreError {
Validation(String),
Conflict(String),
SchemaVersion(i64),
Storage(String),
}
impl fmt::Display for StoreError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Validation(value) => write!(formatter, "validation: {value}"),
Self::Conflict(value) => write!(formatter, "conflict: {value}"),
Self::SchemaVersion(value) => {
write!(formatter, "unsupported schema version {value}")
}
Self::Storage(value) => write!(formatter, "storage: {value}"),
}
}
}
impl Error for StoreError {}
impl From<rusqlite::Error> for StoreError {
fn from(value: rusqlite::Error) -> Self {
Self::Storage(value.to_string())
}
}
impl From<serde_json::Error> for StoreError {
fn from(value: serde_json::Error) -> Self {
Self::Storage(value.to_string())
}
}
pub struct Store {
connection: Mutex<Connection>,
}
pub fn open(path: impl AsRef<Path>) -> Result<Store, StoreError> {
let connection = Connection::open(path)?;
let version: i64 = connection.query_row("PRAGMA user_version", [], |row| row.get(0))?;
if version == 0 {
let objects: i64 = connection.query_row(
"SELECT count(*) FROM sqlite_schema WHERE name NOT LIKE 'sqlite_%'",
[],
|row| row.get(0),
)?;
if objects != 0 {
return Err(StoreError::SchemaVersion(0));
}
connection.execute_batch(
"CREATE TABLE events(
event_id TEXT PRIMARY KEY NOT NULL,
event_json BLOB NOT NULL
);
PRAGMA user_version = 2;",
)?;
} else if version != 2 {
return Err(StoreError::SchemaVersion(version));
}
connection.prepare("SELECT event_id,event_json FROM events LIMIT 0")?;
Ok(Store {
connection: Mutex::new(connection),
})
}
impl Store {
pub fn execute(&self, request: Request) -> Result<Response, StoreError> {
let mut connection = self
.connection
.lock()
.map_err(|_| StoreError::Storage("connection lock poisoned".into()))?;
let (mut event_ids, mut projection) = load(&connection)?;
match request {
Request::Read(query) => Ok(Response {
revision: event_ids.len() as u64,
result: projection.read(query)?,
}),
Request::Commit(events) => {
let mut additions = Vec::new();
for envelope in events {
let id = envelope.event_id.as_ref().to_owned();
let json = serde_json::to_vec(&envelope.event)?;
if let Some(old) = event_ids.get(&id) {
if old != &json {
return Err(StoreError::Conflict(format!(
"event ID {id} has different content"
)));
}
continue;
}
projection.apply(&envelope.event)?;
event_ids.insert(id.clone(), json.clone());
additions.push((id, json));
}
let transaction = connection.transaction()?;
for (id, json) in &additions {
transaction.execute(
"INSERT INTO events(event_id,event_json) VALUES(?1,?2)",
params![id, json],
)?;
}
transaction.commit()?;
Ok(Response {
revision: event_ids.len() as u64,
result: ResponseKind::Commit {
applied: additions.len() as u64,
},
})
}
}
}
}
#[derive(Default)]
struct Projection {
registrations: BTreeMap<ObjectId, SegmentRegistration>,
attempts: BTreeMap<Key, AttemptRecord>,
selections: BTreeMap<(ObjectId, Key), AttemptSelection>,
states: BTreeMap<Key, SampleStateChange>,
artifacts: BTreeMap<Key, ArtifactRecord>,
active_models: BTreeMap<Key, ActiveModelChange>,
}
fn load(connection: &Connection) -> Result<(BTreeMap<String, Vec<u8>>, Projection), StoreError> {
let mut statement =
connection.prepare("SELECT event_id,event_json FROM events ORDER BY rowid")?;
let rows = statement.query_map([], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, Vec<u8>>(1)?))
})?;
let mut ids = BTreeMap::new();
let mut projection = Projection::default();
for row in rows {
let (id, json) = row?;
Key::parse(&id)
.map_err(|error| StoreError::Storage(format!("invalid stored event ID: {error}")))?;
let event: Event = serde_json::from_slice(&json)?;
projection.apply(&event)?;
ids.insert(id, json);
}
Ok((ids, projection))
}
impl Projection {
fn apply(&mut self, event: &Event) -> Result<(), StoreError> {
match event {
Event::RegisterSegment(registration) => self.register(registration),
Event::RecordAttempt(attempt) => self.record_attempt(attempt),
Event::SelectAttempt(selection) => self.select(selection),
Event::SetSampleState(change) => self.set_sample_state(change),
Event::PutArtifact(artifact) => self.put_artifact(artifact),
Event::SetActiveModel(change) => self.set_active_model(change),
}
}
fn register(&mut self, registration: &SegmentRegistration) -> Result<(), StoreError> {
registration
.segment
.validate()
.map_err(|error| StoreError::Validation(error.to_string()))?;
let segment = ®istration.segment;
let duration = segment.end_ms - segment.start_ms;
if duration == 0 || duration >= 240_000 {
return validation("segment duration must be 1..239999 ms");
}
if segment.ordinal == 0 && segment.start_ms != 0 {
return validation("source ordinal zero must begin at zero");
}
if self.registrations.contains_key(&segment.clip_object) {
return conflict("clip object is already registered");
}
for old in self
.registrations
.values()
.filter(|old| old.segment.source_object == segment.source_object)
{
if old.segment.policy != segment.policy
|| old.segment.segment_count != segment.segment_count
|| old.group_id != registration.group_id
|| old.recording_kind != registration.recording_kind
{
return conflict("source plan, group, or recording kind differs");
}
if old.segment.ordinal == segment.ordinal {
return conflict("source ordinal is already registered");
}
let old_duration = old.segment.end_ms - old.segment.start_ms;
if old_duration.abs_diff(duration) > 1 {
return validation("source segment lengths differ by more than one millisecond");
}
let ordered = if old.segment.ordinal < segment.ordinal {
old.segment.start_ms < segment.start_ms && old.segment.end_ms < segment.end_ms
} else {
segment.start_ms < old.segment.start_ms && segment.end_ms < old.segment.end_ms
};
if !ordered {
return validation("source segment ranges do not follow ordinal order");
}
if old.segment.ordinal.checked_add(1) == Some(segment.ordinal)
&& old.segment.end_ms.checked_sub(segment.start_ms) != Some(5_000)
{
return validation("adjacent source segments must overlap by 5000 ms");
}
if segment.ordinal.checked_add(1) == Some(old.segment.ordinal)
&& segment.end_ms.checked_sub(old.segment.start_ms) != Some(5_000)
{
return validation("adjacent source segments must overlap by 5000 ms");
}
}
self.registrations
.insert(segment.clip_object.clone(), registration.clone());
Ok(())
}
fn record_attempt(&mut self, attempt: &AttemptRecord) -> Result<(), StoreError> {
if self.attempts.contains_key(&attempt.attempt_id) {
return conflict("attempt ID is already recorded");
}
let registration = self
.registrations
.get(&attempt.clip_object)
.ok_or_else(|| StoreError::Conflict("clip must be registered first".into()))?;
if !matches!(
attempt.outcome,
StoredAttemptOutcome::ProviderFailure { .. }
) && attempt.provider_result_object.is_none()
{
return validation("response-bearing attempt needs a result object");
}
match &attempt.outcome {
StoredAttemptOutcome::Scored {
recording_quality,
speakers,
additional_speakers,
} => {
if *recording_quality > 100 || speakers.is_empty() {
return validation(
"scored attempt needs recording quality <=100 and a complete speaker",
);
}
validate_speaker_ordinals(speakers, additional_speakers)?;
let duration = registration.segment.end_ms - registration.segment.start_ms;
for speaker in speakers {
bounded(&speaker.closest_dialect, 128, "dialect")?;
if speaker.usable_speech_ms == 0
|| u64::from(speaker.usable_speech_ms) > duration
{
return validation("usable speech is outside segment duration");
}
sample_id(&attempt.attempt_id, speaker.speaker_ordinal)?;
}
}
StoredAttemptOutcome::Unscorable {
reason,
additional_speakers,
} => {
bounded(reason, 240, "reason")?;
validate_speaker_ordinals(&[], additional_speakers)?;
}
StoredAttemptOutcome::InvalidResponse { error } => bounded(error, 240, "error")?,
StoredAttemptOutcome::ProviderFailure { detail, .. } => {
bounded(detail, 240, "detail")?;
}
}
self.attempts
.insert(attempt.attempt_id.clone(), attempt.clone());
Ok(())
}
fn select(&mut self, selection: &AttemptSelection) -> Result<(), StoreError> {
bounded(&selection.reason, 240, "selection reason")?;
let attempt = self
.attempts
.get(&selection.attempt_id)
.ok_or_else(|| StoreError::Conflict("selected attempt does not exist".into()))?;
if attempt.clip_object != selection.clip_object
|| attempt.extraction_key != selection.extraction_key
|| !matches!(attempt.outcome, StoredAttemptOutcome::Scored { .. })
{
return conflict("selection does not match a scored clip attempt");
}
self.selections.insert(
(
selection.clip_object.clone(),
selection.extraction_key.clone(),
),
selection.clone(),
);
Ok(())
}
fn set_sample_state(&mut self, change: &SampleStateChange) -> Result<(), StoreError> {
bounded(&change.reason, 240, "sample-state reason")?;
if self.sample_owner(&change.sample_id).is_none() {
return conflict("sample ID does not exist");
}
self.states.insert(change.sample_id.clone(), change.clone());
Ok(())
}
fn put_artifact(&mut self, artifact: &ArtifactRecord) -> Result<(), StoreError> {
let actual: [u8; 32] = Sha256::digest(&artifact.bytes).into();
if actual != artifact.sha256 {
return validation("artifact SHA-256 does not match bytes");
}
if let Some(old) = self.artifacts.get(&artifact.artifact_id) {
if old != artifact {
return conflict("artifact ID is immutable");
}
return Ok(());
}
self.artifacts
.insert(artifact.artifact_id.clone(), artifact.clone());
Ok(())
}
fn set_active_model(&mut self, change: &ActiveModelChange) -> Result<(), StoreError> {
bounded(&change.reason, 240, "activation reason")?;
if let Some(id) = &change.model_artifact_id {
let artifact = self
.artifacts
.get(id)
.ok_or_else(|| StoreError::Conflict("model artifact does not exist".into()))?;
if artifact.kind != ArtifactKind::Model || artifact.cohort_id != change.cohort_id {
return conflict("active artifact must be a Model in the same cohort");
}
}
self.active_models
.insert(change.cohort_id.clone(), change.clone());
Ok(())
}
fn sample_owner(&self, wanted: &Key) -> Option<(&AttemptRecord, &StoredSpeaker, u8)> {
for attempt in self.attempts.values() {
if let StoredAttemptOutcome::Scored {
recording_quality,
speakers,
..
} = &attempt.outcome
{
for speaker in speakers {
if sample_id(&attempt.attempt_id, speaker.speaker_ordinal)
.ok()
.as_ref()
== Some(wanted)
{
return Some((attempt, speaker, *recording_quality));
}
}
}
}
None
}
fn snapshot(&self, attempt: &AttemptRecord) -> Result<AttemptSnapshot, StoreError> {
let sample_ids = match &attempt.outcome {
StoredAttemptOutcome::Scored { speakers, .. } => speakers
.iter()
.map(|speaker| sample_id(&attempt.attempt_id, speaker.speaker_ordinal))
.collect::<Result<_, _>>()?,
_ => Vec::new(),
};
let selected = self
.selections
.get(&(attempt.clip_object.clone(), attempt.extraction_key.clone()))
.is_some_and(|selection| selection.attempt_id == attempt.attempt_id);
Ok(AttemptSnapshot {
attempt: attempt.clone(),
sample_ids,
selected,
})
}
fn rows(&self, cohort_id: &Key, selected_only: bool) -> Result<Vec<LabeledSample>, StoreError> {
let mut rows = Vec::new();
for attempt in self
.attempts
.values()
.filter(|attempt| &attempt.extraction_key == cohort_id)
{
if selected_only
&& !self
.selections
.get(&(attempt.clip_object.clone(), attempt.extraction_key.clone()))
.is_some_and(|selection| selection.attempt_id == attempt.attempt_id)
{
continue;
}
let StoredAttemptOutcome::Scored {
recording_quality,
speakers,
..
} = &attempt.outcome
else {
continue;
};
let registration = self
.registrations
.get(&attempt.clip_object)
.ok_or_else(|| StoreError::Storage("attempt lost its registration".into()))?;
for speaker in speakers {
let id = sample_id(&attempt.attempt_id, speaker.speaker_ordinal)?;
let Some(SampleStateChange {
state: SampleState::Confirmed { speaker_id },
..
}) = self.states.get(&id)
else {
continue;
};
rows.push(LabeledSample {
sample_id: id,
attempt_id: attempt.attempt_id.clone(),
speaker_id: speaker_id.clone(),
cohort_id: attempt.extraction_key.clone(),
group_id: registration.group_id.clone(),
clip_object: attempt.clip_object.clone(),
primary_language: speaker.primary_language.clone(),
recording_kind: registration.recording_kind,
usable_speech_ms: speaker.usable_speech_ms,
recording_quality: *recording_quality,
features: speaker.features,
});
}
}
Ok(rows)
}
fn read(&self, query: Query) -> Result<ResponseKind, StoreError> {
match query {
Query::Attempt { attempt_id } => Ok(ResponseKind::Attempt(
self.attempts
.get(&attempt_id)
.map(|attempt| self.snapshot(attempt))
.transpose()?,
)),
Query::Attempts {
extraction_key,
source_object,
} => Ok(ResponseKind::Attempts(
self.attempts
.values()
.filter(|attempt| {
extraction_key
.as_ref()
.is_none_or(|key| key == &attempt.extraction_key)
&& source_object.as_ref().is_none_or(|source| {
self.registrations.get(&attempt.clip_object).is_some_and(
|registration| ®istration.segment.source_object == source,
)
})
})
.map(|attempt| self.snapshot(attempt))
.collect::<Result<_, _>>()?,
)),
Query::TrainingRows { cohort_id } => {
Ok(ResponseKind::TrainingRows(self.rows(&cohort_id, true)?))
}
Query::RepeatabilityRows { cohort_id } => Ok(ResponseKind::RepeatabilityRows(
self.rows(&cohort_id, false)?,
)),
Query::Artifact { artifact_id } => Ok(ResponseKind::Artifact(
self.artifacts.get(&artifact_id).cloned(),
)),
Query::ActiveModel { cohort_id } => Ok(ResponseKind::ActiveModel(
self.active_models
.get(&cohort_id)
.and_then(|change| change.model_artifact_id.clone()),
)),
Query::Inventory { extraction_key } => {
Ok(ResponseKind::Inventory(self.inventory(extraction_key)?))
}
}
}
fn inventory(&self, filter: Option<Key>) -> Result<InventorySnapshot, StoreError> {
let included = |attempt: &&AttemptRecord| {
filter
.as_ref()
.is_none_or(|key| key == &attempt.extraction_key)
};
let attempts = self
.attempts
.values()
.filter(included)
.map(|attempt| self.snapshot(attempt))
.collect::<Result<Vec<_>, _>>()?;
let included_samples: BTreeSet<Key> = attempts
.iter()
.flat_map(|attempt| attempt.sample_ids.iter().cloned())
.collect();
let selections = self
.selections
.values()
.filter(|selection| {
filter
.as_ref()
.is_none_or(|key| key == &selection.extraction_key)
})
.cloned()
.map(|selection| SelectionSnapshot { selection })
.collect();
let sample_states = self
.states
.values()
.filter(|change| filter.is_none() || included_samples.contains(&change.sample_id))
.cloned()
.map(|change| SampleStateSnapshot { change })
.collect();
let mut registrations: Vec<_> = self.registrations.values().cloned().collect();
registrations.sort_by(|left, right| {
(&left.segment.source_object, left.segment.ordinal)
.cmp(&(&right.segment.source_object, right.segment.ordinal))
});
let mut sources: BTreeMap<ObjectId, (u16, BTreeSet<u16>)> = BTreeMap::new();
for registration in ®istrations {
let entry = sources
.entry(registration.segment.source_object.clone())
.or_insert((registration.segment.segment_count, BTreeSet::new()));
entry.1.insert(registration.segment.ordinal);
}
let missing_segment_ordinals = sources
.into_iter()
.filter_map(|(source_object, (count, present))| {
let ordinals: Vec<_> = (0..count)
.filter(|ordinal| !present.contains(ordinal))
.collect();
(!ordinals.is_empty()).then_some(MissingSegmentOrdinals {
source_object,
ordinals,
})
})
.collect();
let mut outcome_counts = OutcomeCounts::default();
for attempt in self.attempts.values().filter(included) {
match attempt.outcome {
StoredAttemptOutcome::Scored { .. } => outcome_counts.scored += 1,
StoredAttemptOutcome::Unscorable { .. } => outcome_counts.unscorable += 1,
StoredAttemptOutcome::InvalidResponse { .. } => {
outcome_counts.invalid_response += 1;
}
StoredAttemptOutcome::ProviderFailure { .. } => {
outcome_counts.provider_failure += 1;
}
}
}
Ok(InventorySnapshot {
registrations,
attempts,
selections,
sample_states,
artifacts: self
.artifacts
.values()
.map(|artifact| ArtifactMetadata {
artifact_id: artifact.artifact_id.clone(),
cohort_id: artifact.cohort_id.clone(),
kind: artifact.kind,
sha256: artifact.sha256,
byte_len: artifact.bytes.len() as u64,
})
.collect(),
active_models: self
.active_models
.values()
.cloned()
.map(|change| ActiveModelSnapshot { change })
.collect(),
missing_segment_ordinals,
outcome_counts,
})
}
}
fn validate_speaker_ordinals(
speakers: &[StoredSpeaker],
additional_speakers: &[StoredAdditionalSpeaker],
) -> Result<(), StoreError> {
let mut ordinals = BTreeSet::new();
for ordinal in speakers
.iter()
.map(|speaker| speaker.speaker_ordinal)
.chain(
additional_speakers
.iter()
.map(|speaker| speaker.speaker_ordinal),
)
{
if !ordinals.insert(ordinal) {
return validation("complete and additional speaker ordinals must be unique");
}
}
if ordinals
.iter()
.enumerate()
.any(|(expected, actual)| expected != usize::from(*actual))
{
return validation("complete and additional speaker ordinals must be contiguous from zero");
}
for speaker in additional_speakers {
bounded(&speaker.description, 240, "additional-speaker description")?;
}
Ok(())
}
fn sample_id(attempt_id: &Key, ordinal: u16) -> Result<Key, StoreError> {
let mut digest = Sha256::new();
digest.update(attempt_id.as_ref().as_bytes());
digest.update([0]);
digest.update(ordinal.to_be_bytes());
let hex = digest
.finalize()
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
Key::parse(&format!("sample:sha256:{hex}"))
.map_err(|error| StoreError::Storage(format!("generated invalid sample ID: {error}")))
}
fn bounded(value: &str, maximum: usize, name: &str) -> Result<(), StoreError> {
if value.trim().is_empty() || value.len() > maximum {
return validation(format!(
"{name} must be nonblank and at most {maximum} bytes"
));
}
Ok(())
}
fn validation<T>(message: impl Into<String>) -> Result<T, StoreError> {
Err(StoreError::Validation(message.into()))
}
fn conflict<T>(message: impl Into<String>) -> Result<T, StoreError> {
Err(StoreError::Conflict(message.into()))
}
#[cfg(test)]
mod tests {
use super::*;
use kcode_speaker_types::FEATURE_COUNT;
use std::{
fs,
path::PathBuf,
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
thread,
};
static NEXT_PATH: AtomicU64 = AtomicU64::new(0);
fn path() -> PathBuf {
std::env::temp_dir().join(format!(
"kcode-speaker-store-{}-{}.db",
std::process::id(),
NEXT_PATH.fetch_add(1, Ordering::Relaxed)
))
}
fn key(value: &str) -> Key {
Key::parse(value).unwrap()
}
fn object(value: &str) -> ObjectId {
ObjectId::parse(value).unwrap()
}
fn features(seed: u8) -> FeatureVector {
FeatureVector::new([seed; FEATURE_COUNT]).unwrap()
}
fn additional(speaker_ordinal: u16, description: &str) -> StoredAdditionalSpeaker {
StoredAdditionalSpeaker {
speaker_ordinal,
description: description.into(),
}
}
fn registration(source: &str, clip: &str, ordinal: u16, count: u16) -> SegmentRegistration {
SegmentRegistration {
segment: SegmentRef {
source_object: object(source),
clip_object: object(clip),
ordinal,
segment_count: count,
start_ms: u64::from(ordinal) * 95_000,
end_ms: u64::from(ordinal) * 95_000 + 100_000,
policy: key("speaker-segments/1"),
},
group_id: key("group/1"),
recording_kind: RecordingKind::VoiceNote,
}
}
fn scored(id: &str, clip: &str, extraction: &str, speakers: u16) -> AttemptRecord {
AttemptRecord {
attempt_id: key(id),
clip_object: object(clip),
extraction_key: key(extraction),
provider_result_object: Some(object("RESULT01")),
outcome: StoredAttemptOutcome::Scored {
recording_quality: 80,
speakers: (0..speakers)
.map(|speaker_ordinal| StoredSpeaker {
speaker_ordinal,
primary_language: key("en"),
closest_dialect: "General American".into(),
usable_speech_ms: 20_000,
features: features(50),
})
.collect(),
additional_speakers: Vec::new(),
},
}
}
fn failure(id: &str, clip: &str, extraction: &str) -> AttemptRecord {
AttemptRecord {
attempt_id: key(id),
clip_object: object(clip),
extraction_key: key(extraction),
provider_result_object: None,
outcome: StoredAttemptOutcome::ProviderFailure {
code: key("timeout"),
detail: "provider timed out".into(),
},
}
}
fn event(id: &str, event: Event) -> EventEnvelope {
EventEnvelope {
event_id: key(id),
event,
}
}
fn commit(store: &Store, events: Vec<EventEnvelope>) -> Response {
store.execute(Request::Commit(events)).unwrap()
}
fn attempt(store: &Store, id: &str) -> AttemptSnapshot {
let response = store
.execute(Request::Read(Query::Attempt {
attempt_id: key(id),
}))
.unwrap();
let ResponseKind::Attempt(Some(attempt)) = response.result else {
panic!("missing attempt");
};
attempt
}
fn selection(id: &str, clip: &str, extraction: &str) -> AttemptSelection {
AttemptSelection {
clip_object: object(clip),
extraction_key: key(extraction),
attempt_id: key(id),
reason: "reviewed".into(),
}
}
fn state(sample_id: Key, state: SampleState) -> SampleStateChange {
SampleStateChange {
sample_id,
state,
reason: "human review".into(),
}
}
#[test]
fn fresh_schema_two_database_restarts_with_additional_evidence() {
let path = path();
{
let store = open(&path).unwrap();
let mut scored = scored("attempt/1", "CLIP0001", "extract/1", 1);
let StoredAttemptOutcome::Scored {
additional_speakers,
..
} = &mut scored.outcome
else {
unreachable!();
};
additional_speakers.push(additional(1, "Brief background interjection"));
let response = commit(
&store,
vec![
event(
"event/register",
Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
),
event("event/attempt", Event::RecordAttempt(scored)),
],
);
assert_eq!(response.revision, 2);
let version: i64 = store
.connection
.lock()
.unwrap()
.query_row("PRAGMA user_version", [], |row| row.get(0))
.unwrap();
assert_eq!(version, 2);
}
let store = open(&path).unwrap();
let stored = attempt(&store, "attempt/1");
assert_eq!(stored.sample_ids.len(), 1);
let StoredAttemptOutcome::Scored {
speakers,
additional_speakers,
..
} = stored.attempt.outcome
else {
panic!();
};
assert_eq!(speakers.len(), 1);
assert_eq!(
additional_speakers,
vec![additional(1, "Brief background interjection")]
);
drop(store);
fs::remove_file(path).unwrap();
}
#[test]
fn commit_is_atomic_and_event_ids_are_idempotent() {
let path = path();
let store = open(&path).unwrap();
let registration_event = event(
"event/register",
Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
);
let bad = event(
"event/bad",
Event::RecordAttempt(scored("attempt/1", "UNKNOWN1", "extract/1", 1)),
);
assert!(
store
.execute(Request::Commit(vec![registration_event.clone(), bad]))
.is_err()
);
let response = store
.execute(Request::Read(Query::Inventory {
extraction_key: None,
}))
.unwrap();
assert_eq!(response.revision, 0);
assert!(matches!(
commit(&store, vec![registration_event.clone()]).result,
ResponseKind::Commit { applied: 1 }
));
assert!(matches!(
commit(
&store,
vec![registration_event.clone(), registration_event.clone()]
)
.result,
ResponseKind::Commit { applied: 0 }
));
let mut changed = registration_event;
changed.event = Event::RegisterSegment(registration("SOURCE02", "CLIP0002", 0, 1));
assert!(matches!(
store.execute(Request::Commit(vec![changed])),
Err(StoreError::Conflict(_))
));
assert_eq!(
store
.execute(Request::Read(Query::Inventory {
extraction_key: None,
}))
.unwrap()
.revision,
1
);
drop(store);
fs::remove_file(path).unwrap();
}
#[test]
fn registration_order_and_source_plan_conflicts_fail() {
let path = path();
let store = open(&path).unwrap();
assert!(matches!(
store.execute(Request::Commit(vec![event(
"attempt/early",
Event::RecordAttempt(scored("attempt/1", "CLIP0001", "extract/1", 1)),
)])),
Err(StoreError::Conflict(_))
));
commit(
&store,
vec![event(
"register/0",
Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 2)),
)],
);
let mut wrong_count = registration("SOURCE01", "CLIP0002", 1, 3);
assert!(
store
.execute(Request::Commit(vec![event(
"register/wrong-count",
Event::RegisterSegment(wrong_count.clone()),
)]))
.is_err()
);
wrong_count.segment.segment_count = 2;
wrong_count.segment.policy = key("different-plan");
assert!(
store
.execute(Request::Commit(vec![event(
"register/wrong-plan",
Event::RegisterSegment(wrong_count),
)]))
.is_err()
);
let duplicate_ordinal = registration("SOURCE01", "CLIP0003", 0, 2);
assert!(
store
.execute(Request::Commit(vec![event(
"register/duplicate",
Event::RegisterSegment(duplicate_ordinal),
)]))
.is_err()
);
drop(store);
fs::remove_file(path).unwrap();
}
#[test]
fn complete_and_additional_speaker_contract_is_enforced() {
let path = path();
let store = open(&path).unwrap();
commit(
&store,
vec![event(
"register",
Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
)],
);
let empty_scored = scored("empty", "CLIP0001", "extract/1", 0);
assert!(matches!(
store.execute(Request::Commit(vec![event(
"attempt/empty",
Event::RecordAttempt(empty_scored),
)])),
Err(StoreError::Validation(_))
));
let mut duplicate = scored("duplicate", "CLIP0001", "extract/1", 1);
let StoredAttemptOutcome::Scored {
additional_speakers,
..
} = &mut duplicate.outcome
else {
unreachable!();
};
additional_speakers.push(additional(0, "Same ordinal"));
assert!(matches!(
store.execute(Request::Commit(vec![event(
"attempt/duplicate",
Event::RecordAttempt(duplicate),
)])),
Err(StoreError::Validation(_))
));
let mut gap = scored("gap", "CLIP0001", "extract/1", 1);
let StoredAttemptOutcome::Scored {
additional_speakers,
..
} = &mut gap.outcome
else {
unreachable!();
};
additional_speakers.push(additional(2, "Ordinal one is missing"));
assert!(matches!(
store.execute(Request::Commit(vec![event(
"attempt/gap",
Event::RecordAttempt(gap),
)])),
Err(StoreError::Validation(_))
));
let unscorable = AttemptRecord {
attempt_id: key("unscorable"),
clip_object: object("CLIP0001"),
extraction_key: key("extract/1"),
provider_result_object: Some(object("RESULT02")),
outcome: StoredAttemptOutcome::Unscorable {
reason: "No complete profile".into(),
additional_speakers: vec![
additional(0, "A few clear words"),
additional(1, "Overlapped background speech"),
],
},
};
commit(
&store,
vec![event(
"attempt/unscorable",
Event::RecordAttempt(unscorable),
)],
);
let snapshot = attempt(&store, "unscorable");
assert!(snapshot.sample_ids.is_empty());
let StoredAttemptOutcome::Unscorable {
additional_speakers,
..
} = snapshot.attempt.outcome
else {
panic!();
};
assert_eq!(additional_speakers.len(), 2);
let bad_description = AttemptRecord {
attempt_id: key("bad-description"),
clip_object: object("CLIP0001"),
extraction_key: key("extract/1"),
provider_result_object: Some(object("RESULT03")),
outcome: StoredAttemptOutcome::Unscorable {
reason: "No complete profile".into(),
additional_speakers: vec![additional(0, " ")],
},
};
assert!(matches!(
store.execute(Request::Commit(vec![event(
"attempt/bad-description",
Event::RecordAttempt(bad_description),
)])),
Err(StoreError::Validation(_))
));
let bad_reason = AttemptRecord {
attempt_id: key("bad-reason"),
clip_object: object("CLIP0001"),
extraction_key: key("extract/1"),
provider_result_object: Some(object("RESULT04")),
outcome: StoredAttemptOutcome::Unscorable {
reason: "x".repeat(241),
additional_speakers: Vec::new(),
},
};
assert!(matches!(
store.execute(Request::Commit(vec![event(
"attempt/bad-reason",
Event::RecordAttempt(bad_reason),
)])),
Err(StoreError::Validation(_))
));
drop(store);
fs::remove_file(path).unwrap();
}
#[test]
fn retries_replacement_ids_and_complete_label_projections() {
let path = path();
let store = open(&path).unwrap();
let mut first_attempt = scored("try/one", "CLIP0001", "extract/1", 2);
let StoredAttemptOutcome::Scored {
additional_speakers,
..
} = &mut first_attempt.outcome
else {
unreachable!();
};
additional_speakers.push(additional(2, "Too brief for a complete profile"));
commit(
&store,
vec![
event(
"register/1",
Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
),
event("attempt/1", Event::RecordAttempt(first_attempt)),
event(
"attempt/2",
Event::RecordAttempt(scored("try/two", "CLIP0001", "extract/1", 1)),
),
event(
"select/1",
Event::SelectAttempt(selection("try/one", "CLIP0001", "extract/1")),
),
],
);
let first = attempt(&store, "try/one");
assert_eq!(first.sample_ids.len(), 2);
assert_ne!(first.sample_ids[0], first.sample_ids[1]);
let mut digest = Sha256::new();
digest.update(b"try/one");
digest.update([0]);
digest.update(0u16.to_be_bytes());
let expected = key(&format!("sample:sha256:{:x}", digest.finalize()));
assert_eq!(first.sample_ids[0], expected);
let additional_id = sample_id(&key("try/one"), 2).unwrap();
assert!(matches!(
store.execute(Request::Commit(vec![event(
"state/additional",
Event::SetSampleState(state(
additional_id,
SampleState::Confirmed {
speaker_id: key("speaker/background"),
},
)),
)])),
Err(StoreError::Conflict(_))
));
let second = attempt(&store, "try/two");
commit(
&store,
vec![
event(
"state/1",
Event::SetSampleState(state(
first.sample_ids[0].clone(),
SampleState::Confirmed {
speaker_id: key("speaker/alice"),
},
)),
),
event(
"state/2",
Event::SetSampleState(state(
second.sample_ids[0].clone(),
SampleState::Confirmed {
speaker_id: key("speaker/alice"),
},
)),
),
],
);
let ResponseKind::TrainingRows(rows) = store
.execute(Request::Read(Query::TrainingRows {
cohort_id: key("extract/1"),
}))
.unwrap()
.result
else {
panic!();
};
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].attempt_id, key("try/one"));
assert_eq!(rows[0].recording_quality, 80);
assert_eq!(rows[0].features.as_ref().len(), FEATURE_COUNT);
let ResponseKind::RepeatabilityRows(rows) = store
.execute(Request::Read(Query::RepeatabilityRows {
cohort_id: key("extract/1"),
}))
.unwrap()
.result
else {
panic!();
};
assert_eq!(rows.len(), 2);
commit(
&store,
vec![event(
"select/2",
Event::SelectAttempt(selection("try/two", "CLIP0001", "extract/1")),
)],
);
let ResponseKind::TrainingRows(rows) = store
.execute(Request::Read(Query::TrainingRows {
cohort_id: key("extract/1"),
}))
.unwrap()
.result
else {
panic!();
};
assert_eq!(rows[0].attempt_id, key("try/two"));
commit(
&store,
vec![event(
"state/correct",
Event::SetSampleState(state(
second.sample_ids[0].clone(),
SampleState::Confirmed {
speaker_id: key("speaker/bob"),
},
)),
)],
);
let ResponseKind::TrainingRows(rows) = store
.execute(Request::Read(Query::TrainingRows {
cohort_id: key("extract/1"),
}))
.unwrap()
.result
else {
panic!();
};
assert_eq!(rows[0].speaker_id, key("speaker/bob"));
commit(
&store,
vec![event(
"state/retract",
Event::SetSampleState(state(second.sample_ids[0].clone(), SampleState::Retracted)),
)],
);
let ResponseKind::TrainingRows(rows) = store
.execute(Request::Read(Query::TrainingRows {
cohort_id: key("extract/1"),
}))
.unwrap()
.result
else {
panic!();
};
assert!(rows.is_empty());
drop(store);
fs::remove_file(path).unwrap();
}
#[test]
fn artifacts_are_hashed_immutable_and_cohort_checked() {
let path = path();
let store = open(&path).unwrap();
let bytes = b"model bytes".to_vec();
let hash: [u8; 32] = Sha256::digest(&bytes).into();
let model = ArtifactRecord {
artifact_id: key("artifact/model"),
cohort_id: key("cohort/1"),
kind: ArtifactKind::Model,
sha256: hash,
bytes,
};
let evaluation = ArtifactRecord {
artifact_id: key("artifact/eval"),
cohort_id: key("cohort/1"),
kind: ArtifactKind::Evaluation,
sha256: Sha256::digest(b"evaluation").into(),
bytes: b"evaluation".to_vec(),
};
commit(
&store,
vec![
event("artifact/model", Event::PutArtifact(model.clone())),
event("artifact/eval", Event::PutArtifact(evaluation)),
],
);
let activate = |id: &str, cohort: &str, artifact: &str| {
event(
id,
Event::SetActiveModel(ActiveModelChange {
cohort_id: key(cohort),
model_artifact_id: Some(key(artifact)),
reason: "approved".into(),
}),
)
};
assert!(
store
.execute(Request::Commit(vec![activate(
"activate/wrong-kind",
"cohort/1",
"artifact/eval",
)]))
.is_err()
);
assert!(
store
.execute(Request::Commit(vec![activate(
"activate/wrong-cohort",
"cohort/2",
"artifact/model",
)]))
.is_err()
);
let mut changed = model.clone();
changed.bytes.push(1);
changed.sha256 = Sha256::digest(&changed.bytes).into();
assert!(matches!(
store.execute(Request::Commit(vec![event(
"artifact/change",
Event::PutArtifact(changed),
)])),
Err(StoreError::Conflict(_))
));
let mut bad_hash = model;
bad_hash.artifact_id = key("artifact/bad");
bad_hash.sha256 = [0; 32];
assert!(matches!(
store.execute(Request::Commit(vec![event(
"artifact/bad",
Event::PutArtifact(bad_hash),
)])),
Err(StoreError::Validation(_))
));
commit(
&store,
vec![activate("activate/good", "cohort/1", "artifact/model")],
);
let ResponseKind::ActiveModel(Some(active)) = store
.execute(Request::Read(Query::ActiveModel {
cohort_id: key("cohort/1"),
}))
.unwrap()
.result
else {
panic!();
};
assert_eq!(active, key("artifact/model"));
commit(
&store,
vec![event(
"activate/none",
Event::SetActiveModel(ActiveModelChange {
cohort_id: key("cohort/1"),
model_artifact_id: None,
reason: "retired".into(),
}),
)],
);
assert!(matches!(
store
.execute(Request::Read(Query::ActiveModel {
cohort_id: key("cohort/1"),
}))
.unwrap()
.result,
ResponseKind::ActiveModel(None)
));
drop(store);
fs::remove_file(path).unwrap();
}
#[test]
fn filtering_inventory_missing_and_counts_are_typed() {
let path = path();
let store = open(&path).unwrap();
let unscorable = AttemptRecord {
attempt_id: key("unscorable"),
clip_object: object("CLIP0001"),
extraction_key: key("extract/1"),
provider_result_object: Some(object("RESULT02")),
outcome: StoredAttemptOutcome::Unscorable {
reason: "too faint".into(),
additional_speakers: vec![additional(0, "Faint substantive voice")],
},
};
let invalid = AttemptRecord {
attempt_id: key("invalid"),
clip_object: object("CLIP0001"),
extraction_key: key("extract/1"),
provider_result_object: Some(object("RESULT03")),
outcome: StoredAttemptOutcome::InvalidResponse {
error: "extra field".into(),
},
};
commit(
&store,
vec![
event(
"register",
Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 2)),
),
event(
"attempt/scored",
Event::RecordAttempt(scored("scored", "CLIP0001", "extract/1", 1)),
),
event("attempt/unscorable", Event::RecordAttempt(unscorable)),
event("attempt/invalid", Event::RecordAttempt(invalid)),
event(
"attempt/failure",
Event::RecordAttempt(failure("failure", "CLIP0001", "extract/1")),
),
event(
"attempt/other",
Event::RecordAttempt(failure("other", "CLIP0001", "extract/2")),
),
],
);
let ResponseKind::Attempts(filtered) = store
.execute(Request::Read(Query::Attempts {
extraction_key: Some(key("extract/1")),
source_object: Some(object("SOURCE01")),
}))
.unwrap()
.result
else {
panic!();
};
assert_eq!(filtered.len(), 4);
let ResponseKind::Inventory(inventory) = store
.execute(Request::Read(Query::Inventory {
extraction_key: Some(key("extract/1")),
}))
.unwrap()
.result
else {
panic!();
};
assert_eq!(inventory.attempts.len(), 4);
assert_eq!(
inventory.outcome_counts,
OutcomeCounts {
scored: 1,
unscorable: 1,
invalid_response: 1,
provider_failure: 1,
}
);
assert_eq!(inventory.missing_segment_ordinals.len(), 1);
assert_eq!(inventory.missing_segment_ordinals[0].ordinals, vec![1]);
drop(store);
fs::remove_file(path).unwrap();
}
#[test]
fn schema_one_unversioned_and_future_databases_are_rejected() {
for version in [1, 3] {
let path = path();
let connection = Connection::open(&path).unwrap();
connection
.execute_batch(&format!("PRAGMA user_version = {version};"))
.unwrap();
drop(connection);
assert!(matches!(
open(&path),
Err(StoreError::SchemaVersion(found)) if found == version
));
fs::remove_file(path).unwrap();
}
let path = path();
let connection = Connection::open(&path).unwrap();
connection
.execute_batch("CREATE TABLE legacy(value INTEGER);")
.unwrap();
drop(connection);
assert!(matches!(open(&path), Err(StoreError::SchemaVersion(0))));
fs::remove_file(path).unwrap();
}
#[test]
fn one_store_serializes_concurrent_calls() {
let path = path();
let store = Arc::new(open(&path).unwrap());
commit(
&store,
vec![event(
"register",
Event::RegisterSegment(registration("SOURCE01", "CLIP0001", 0, 1)),
)],
);
let threads: Vec<_> = (0..8)
.map(|index| {
let store = Arc::clone(&store);
thread::spawn(move || {
commit(
&store,
vec![event(
&format!("event/{index}"),
Event::RecordAttempt(failure(
&format!("attempt/{index}"),
"CLIP0001",
"extract/1",
)),
)],
);
})
})
.collect();
for thread in threads {
thread.join().unwrap();
}
let response = store
.execute(Request::Read(Query::Attempts {
extraction_key: None,
source_object: None,
}))
.unwrap();
assert_eq!(response.revision, 9);
let ResponseKind::Attempts(attempts) = response.result else {
panic!();
};
assert_eq!(attempts.len(), 8);
drop(store);
fs::remove_file(path).unwrap();
}
}