zkr/
store.rs

1use crate::{
2    Claim, ClaimEvidence, ClaimId, ClaimKind, ClaimStatus, DailyReview, DailyReviewId, Evidence,
3    EvidenceId, EvidenceRelation, MemoryProcessingState, MemoryRef, MemoryTier, PersonId,
4    ProfileEntry, ProfileEntryId, ProfileStability, RetrievalItem, RetrievalPack, Source, SourceId,
5    SourceKind, TenantId, Timestamp, assert_legal_state,
6};
7use rusqlite::{Connection, OptionalExtension, Transaction, TransactionBehavior, params};
8use serde::{Deserialize, Serialize};
9use std::path::Path;
10
11mod apply;
12mod embeddings;
13mod export;
14mod lifecycle;
15mod repair;
16mod retrieval;
17mod schema;
18
19use embeddings::*;
20#[cfg(test)]
21use retrieval::MAX_EXCERPT_BYTES;
22
23pub trait Embedder {
24    type Error: std::error::Error + Send + Sync + 'static;
25
26    fn embed(&self, input: &str) -> std::result::Result<Embedding, Self::Error>;
27}
28
29#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
30pub struct Embedding {
31    pub vector: Vec<f32>,
32    pub model: String,
33    pub version: String,
34    pub input_hash: String,
35    pub normalization: VectorNormalization,
36    pub distance: VectorDistance,
37}
38
39#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
40#[serde(rename_all = "snake_case")]
41pub enum VectorNormalization {
42    None,
43    L2,
44}
45
46#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
47#[serde(rename_all = "snake_case")]
48pub enum VectorDistance {
49    Cosine,
50    Dot,
51    Euclidean,
52}
53
54#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
55#[serde(tag = "kind", content = "id", rename_all = "snake_case")]
56pub enum EmbeddingTarget {
57    Source(SourceId),
58    Evidence(EvidenceId),
59    Claim(ClaimId),
60}
61
62#[derive(Debug, Deserialize)]
63pub struct EmbeddingInput {
64    pub tenant_id: TenantId,
65    pub person_id: PersonId,
66    pub target: EmbeddingTarget,
67    pub embedding: Embedding,
68}
69
70#[derive(Debug, Serialize)]
71pub struct StoredEmbedding {
72    pub target: EmbeddingTarget,
73    pub dimension: usize,
74    pub target_revision: i64,
75    pub input_hash: String,
76    pub created_at: Timestamp,
77}
78
79#[derive(Debug, Deserialize)]
80pub struct ProjectionAuditInput {
81    pub tenant_id: TenantId,
82    pub person_id: PersonId,
83    pub model: String,
84    pub version: String,
85    #[serde(default = "default_limit")]
86    pub limit: u32,
87}
88
89#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
90#[serde(rename_all = "snake_case")]
91pub enum ProjectionState {
92    Missing,
93    Stale,
94}
95
96#[derive(Debug, Serialize)]
97pub struct ProjectionInput {
98    pub target: EmbeddingTarget,
99    pub text: String,
100    pub target_revision: i64,
101    pub input_hash: String,
102}
103
104#[derive(Debug, Serialize)]
105pub struct ProjectionIssue {
106    pub state: ProjectionState,
107    pub input: ProjectionInput,
108    pub stored_target_revision: Option<i64>,
109    pub stored_input_hash: Option<String>,
110    pub stored_created_at: Option<Timestamp>,
111}
112
113#[derive(Debug, thiserror::Error)]
114pub enum Error {
115    #[error(transparent)]
116    Sql(#[from] rusqlite::Error),
117    #[error(transparent)]
118    Json(#[from] serde_json::Error),
119    #[error("{0}")]
120    Invalid(String),
121    #[error("record not found")]
122    NotFound,
123}
124
125pub type Result<T> = std::result::Result<T, Error>;
126
127#[derive(Debug, Deserialize)]
128pub struct RememberInput {
129    pub tenant_id: TenantId,
130    pub person_id: PersonId,
131    #[serde(default)]
132    pub ingestion_key: Option<String>,
133    pub kind: SourceKind,
134    pub text: String,
135    pub captured_at: Timestamp,
136    pub recorded_at: Timestamp,
137    pub claim: Option<ClaimInput>,
138    #[serde(default)]
139    pub feature_flag: Option<String>,
140}
141
142#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
143pub struct TranscriptLocator {
144    pub device_id: String,
145    pub provider: String,
146    pub stream_id: String,
147    pub segment_id: String,
148    pub start_ms: u64,
149    pub end_ms: u64,
150}
151
152#[derive(Debug, Deserialize)]
153pub struct EvidenceLocatorInput {
154    pub tenant_id: TenantId,
155    pub person_id: PersonId,
156    pub evidence_id: EvidenceId,
157}
158
159#[derive(Debug, Deserialize)]
160pub struct RememberRequest {
161    #[serde(flatten)]
162    pub memory: RememberInput,
163    #[serde(default)]
164    pub locator: Option<TranscriptLocator>,
165}
166
167#[derive(Debug, Deserialize)]
168pub struct ClaimInput {
169    pub subject: String,
170    pub predicate: String,
171    pub value: String,
172    pub kind: ClaimKind,
173    pub valid_from: Timestamp,
174    #[serde(default)]
175    pub tier: MemoryTier,
176    #[serde(default)]
177    pub processing_state: MemoryProcessingState,
178}
179
180#[derive(Debug, Serialize)]
181pub struct Remembered {
182    pub source_id: SourceId,
183    pub evidence_id: EvidenceId,
184    pub claim_id: Option<ClaimId>,
185}
186
187#[derive(Debug, Deserialize)]
188pub struct SearchInput {
189    pub tenant_id: TenantId,
190    pub person_id: PersonId,
191    pub query: String,
192    #[serde(default = "default_limit")]
193    pub limit: u32,
194    #[serde(default)]
195    pub query_embedding: Option<DenseQuery>,
196    #[serde(default)]
197    pub as_of: Option<TemporalQuery>,
198    #[serde(default)]
199    pub enabled_features: Vec<String>,
200}
201
202#[derive(Clone, Debug, Deserialize)]
203pub struct TemporalQuery {
204    pub valid_at: Timestamp,
205    pub recorded_at: Timestamp,
206}
207
208#[derive(Debug, Deserialize)]
209pub struct GetInput {
210    pub tenant_id: TenantId,
211    pub person_id: PersonId,
212    pub target: EmbeddingTarget,
213}
214
215#[derive(Debug, Deserialize)]
216pub struct DenseQuery {
217    pub vector: Vec<f32>,
218    pub model: String,
219    pub version: String,
220}
221
222#[derive(Debug, Deserialize)]
223pub struct CorrectInput {
224    pub tenant_id: TenantId,
225    pub person_id: PersonId,
226    pub claim_id: ClaimId,
227    pub text: String,
228    pub value: String,
229    pub valid_at: Timestamp,
230    pub recorded_at: Timestamp,
231}
232
233#[derive(Debug, Serialize)]
234pub struct Corrected {
235    pub source_id: SourceId,
236    pub evidence_id: EvidenceId,
237    pub claim_id: ClaimId,
238    pub superseded_claim_id: ClaimId,
239}
240
241#[derive(Debug, Deserialize)]
242pub struct DeleteInput {
243    pub tenant_id: TenantId,
244    pub person_id: PersonId,
245    pub source_id: SourceId,
246    pub deleted_at: Timestamp,
247}
248
249#[derive(Debug, Deserialize)]
250#[serde(deny_unknown_fields)]
251pub struct ProfileInput {
252    pub tenant_id: TenantId,
253    pub person_id: PersonId,
254    pub stability: ProfileStability,
255    pub claim_id: ClaimId,
256    pub recorded_at: Timestamp,
257}
258
259#[derive(Debug, Deserialize)]
260pub struct ProfilesInput {
261    pub tenant_id: TenantId,
262    pub person_id: PersonId,
263    #[serde(default = "default_limit")]
264    pub limit: u32,
265}
266
267#[derive(Debug, Serialize)]
268pub struct Deleted {
269    pub source_id: SourceId,
270    pub evidence_count: u64,
271    pub claim_count: u64,
272}
273
274#[derive(Debug, Deserialize)]
275pub struct PromoteInput {
276    pub tenant_id: TenantId,
277    pub person_id: PersonId,
278    pub claim_id: ClaimId,
279    pub recorded_at: Timestamp,
280}
281
282#[derive(Debug, Serialize)]
283pub struct Promoted {
284    pub claim_id: ClaimId,
285}
286
287#[derive(Debug, Deserialize)]
288pub struct ArchiveInput {
289    pub tenant_id: TenantId,
290    pub person_id: PersonId,
291    pub claim_id: ClaimId,
292    pub recorded_at: Timestamp,
293}
294
295#[derive(Debug, Serialize)]
296pub struct Archived {
297    pub claim_id: ClaimId,
298}
299
300#[derive(Debug, Deserialize)]
301pub struct RepairInput {
302    pub tenant_id: TenantId,
303    pub person_id: PersonId,
304    #[serde(default = "default_limit")]
305    pub limit: u32,
306}
307
308#[derive(Debug, Serialize)]
309pub struct RepairResult {
310    pub processed: u32,
311}
312
313#[derive(Debug, Deserialize)]
314pub struct ReviewInput {
315    pub tenant_id: TenantId,
316    pub person_id: PersonId,
317    pub day: String,
318    pub summary: String,
319    pub evidence_ids: Vec<EvidenceId>,
320    pub recorded_at: Timestamp,
321}
322
323#[derive(Debug, Deserialize)]
324pub struct ReviewsInput {
325    pub tenant_id: TenantId,
326    pub person_id: PersonId,
327    #[serde(default = "default_limit")]
328    pub limit: u32,
329}
330
331#[derive(Debug, Serialize)]
332pub struct StoredReview {
333    pub id: DailyReviewId,
334}
335
336#[derive(Debug, Serialize)]
337pub struct ReviewRecord {
338    pub id: DailyReviewId,
339    pub day: String,
340    pub summary: String,
341    pub evidence_ids: Vec<EvidenceId>,
342    pub recorded_at: Timestamp,
343}
344
345#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
346pub struct SourceRecord {
347    pub source: Source,
348    pub ingestion_key: Option<String>,
349    pub origin_evidence_id: Option<EvidenceId>,
350    pub origin_claim_id: Option<ClaimId>,
351}
352
353#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
354pub struct EvidenceRecord {
355    pub evidence: Evidence,
356    pub locator: Option<TranscriptLocator>,
357    pub deleted_at: Option<Timestamp>,
358}
359
360#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
361pub struct CorrectionRecord {
362    pub tenant_id: TenantId,
363    pub person_id: PersonId,
364    pub superseded_claim_id: ClaimId,
365    pub claim_id: ClaimId,
366    pub source_id: SourceId,
367    pub evidence_id: EvidenceId,
368    pub valid_at: Timestamp,
369    pub recorded_at: Timestamp,
370}
371
372#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
373pub struct DeletionRecord {
374    pub tenant_id: TenantId,
375    pub person_id: PersonId,
376    pub target: MemoryRef,
377    pub deleted_at: Timestamp,
378}
379
380#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
381#[serde(tag = "kind", content = "record", rename_all = "snake_case")]
382pub enum ExportRecord {
383    Source(SourceRecord),
384    Evidence(EvidenceRecord),
385    Claim(Claim),
386    ClaimEvidence(ClaimEvidence),
387    Correction(CorrectionRecord),
388    Deletion(DeletionRecord),
389    Profile(ProfileEntry),
390    DailyReview(DailyReview),
391}
392
393#[derive(Debug, Deserialize)]
394pub struct ApplyInput {
395    pub export_format: u32,
396    #[serde(default)]
397    pub database_schema_version: Option<i64>,
398    pub tenant_id: TenantId,
399    pub person_id: PersonId,
400    #[serde(default)]
401    pub commits: Vec<ExportCommit>,
402}
403
404#[derive(Debug, Serialize)]
405pub struct Applied {
406    pub commits_applied: u64,
407    pub commits_skipped: u64,
408    pub records_applied: u64,
409    pub records_skipped: u64,
410}
411
412#[derive(Debug, Deserialize)]
413pub struct ExportInput {
414    pub export_format: u32,
415    pub tenant_id: TenantId,
416    pub person_id: PersonId,
417    #[serde(default)]
418    pub after_commit: i64,
419    #[serde(default = "default_after_event_index")]
420    pub after_event_index: i64,
421    #[serde(default)]
422    pub high_water_mark: Option<i64>,
423    #[serde(default = "default_limit")]
424    pub limit: u32,
425}
426
427#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
428pub struct ExportCommit {
429    pub sequence: i64,
430    pub recorded_at: Timestamp,
431    pub event_count: i64,
432    pub first_event_index: i64,
433    pub records: Vec<ExportRecord>,
434}
435
436#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
437pub struct ExportPage {
438    pub export_format: u32,
439    pub database_schema_version: i64,
440    pub high_water_mark: i64,
441    pub next_after_commit: i64,
442    pub next_after_event_index: i64,
443    pub complete: bool,
444    pub commits: Vec<ExportCommit>,
445}
446
447pub const EXPORT_FORMAT_VERSION: u32 = 1;
448pub const DATABASE_SCHEMA_VERSION: i64 = schema::SCHEMA_VERSION;
449pub const MAX_EXPORT_RECORD_BYTES: usize = 1024 * 1024;
450pub const MAX_EXPORT_PAGE_BYTES: usize = 1024 * 1024;
451pub const MAX_PROJECTION_PAGE_BYTES: usize = 1_000_000;
452pub const MAX_READ_PAGE_BYTES: usize = 1_000_000;
453
454pub struct MemoryDb {
455    connection: Connection,
456}
457
458impl MemoryDb {
459    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
460        let connection = Connection::open(path)?;
461        let mut database = Self { connection };
462        database.migrate()?;
463        Ok(database)
464    }
465
466    fn migrate(&mut self) -> Result<()> {
467        schema::migrate(&mut self.connection)
468    }
469}
470
471fn require_scope(tenant_id: &TenantId, person_id: &PersonId) -> Result<()> {
472    require_text("tenant_id", &tenant_id.0)?;
473    require_text("person_id", &person_id.0)
474}
475
476fn require_text(field: &str, value: &str) -> Result<()> {
477    if value.trim().is_empty() {
478        return Err(Error::Invalid(format!("{field} must not be empty")));
479    }
480    Ok(())
481}
482
483const fn default_limit() -> u32 {
484    10
485}
486const fn default_after_event_index() -> i64 {
487    -1
488}
489const fn bounded_limit(limit: u32) -> u32 {
490    if limit == 0 {
491        10
492    } else if limit > 100 {
493        100
494    } else {
495        limit
496    }
497}
498
499fn collect_json_page<T: Serialize>(
500    rows: impl Iterator<Item = rusqlite::Result<T>>,
501) -> Result<Vec<T>> {
502    let mut items = Vec::new();
503    let mut page_bytes = 2;
504    for item in rows {
505        let item = item?;
506        let item_bytes = serde_json::to_vec(&item)?.len();
507        if item_bytes + 2 > MAX_READ_PAGE_BYTES {
508            return Err(Error::Invalid(format!(
509                "read record exceeds the {MAX_READ_PAGE_BYTES}-byte page limit"
510            )));
511        }
512        if !items.is_empty() && page_bytes + item_bytes + 1 > MAX_READ_PAGE_BYTES {
513            return Err(Error::Invalid(format!(
514                "read page exceeds the {MAX_READ_PAGE_BYTES}-byte page limit"
515            )));
516        }
517        page_bytes += item_bytes + usize::from(!items.is_empty());
518        items.push(item);
519    }
520    Ok(items)
521}
522
523#[cfg(test)]
524mod tests;
zkr/store.rs

zkr/
store.rs
