whiskers.git / crates / whiskers-conformance / src / memory.rs
memory.rsannotatedmemory.rssource555 lines · 18.3 KB · raw

The in-memory reference adapter: every port, held in a Mutex and nothing else. It is what the conformance suite is first proven against (the suite must pass on the simplest honest implementation before it is allowed to judge another), what the service's tests run on, and a readable statement of what each port means.

The outbound ports (Judge, Embedder, Voice, Thinker) are scripted: a test says what they answer and reads how often they were asked.

Handles are cheap clones of shared state, so "the process restarted" is a clone of the handle: the state survives, which is exactly what a restart of a durable adapter must look like from outside.

12use std::collections::BTreeMap;
13use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
14use std::sync::{Arc, Mutex, MutexGuard, PoisonError};

16use ::log::{debug, trace};
17use whiskers_core::{Age, ChatState, Direction, Household, IconId, MemoryDoc, MemorySnapshot, TokenLimit, Verdict};
18use whiskers_ports::{
19    Allowance, Appended, Audio, CreditsError, DeviceCursor, DeviceName, Diagnostic, EmbedBatch, EmbedError, Embedder, EntryLine,
20    Hub, Journal, Judge, JudgeError, Lookup, LogCursor, ManualClock, Merged, ModelName, Page, PictureShelf, Put, ReserveError, Reply,
21    SafePictureId, SecretName, Secrets, Settlement, SpeakError, SpeechLine, StoreError, StoredLine, ThinkError, ThinkRequest, Thinker,
22    ThinkingLedger, ThinkingRoom, ThinkingUsage, TimedAudio, TimingError, AudioMs, LineTiming, Tokens, Voice, VoiceCredits, VoiceDay, VoiceReservation, VoiceSpend,
23    Clock, HasAllowance, HasChat, HasClock, HasEmbedder, HasHousehold, HasJournal, HasJudge, HasMemory, HasPictures, HasThinker, HasVoice, Millis,
24};
25
26fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
27    m.lock().unwrap_or_else(PoisonError::into_inner)
28}

The memory hub.

31#[derive(Clone, Default)]
32pub struct MemMemoryHub(Arc<Mutex<MemoryDoc>>);

34impl Hub for MemMemoryHub {
35    type Document = MemorySnapshot;
36
37    async fn merge(&self, theirs: MemorySnapshot) -> Result<Merged<MemorySnapshot>, StoreError> {
38        let mut doc = lock(&self.0);
39        let changed = doc.merge(theirs);
40        Ok(Merged { document: doc.snapshot(), changed })
41    }
42
43    async fn current(&self) -> Result<MemorySnapshot, StoreError> {
44        Ok(lock(&self.0).snapshot())
45    }
46}

The household hub.

49#[derive(Clone, Default)]
50pub struct MemHouseholdHub(Arc<Mutex<Household>>);

52impl Hub for MemHouseholdHub {
53    type Document = Household;
54
55    async fn merge(&self, theirs: Household) -> Result<Merged<Household>, StoreError> {
56        let mut doc = lock(&self.0);
57        let changed = doc.merge(&theirs);
58        Ok(Merged { document: doc.clone(), changed: usize::from(changed) })
59    }
60
61    async fn current(&self) -> Result<Household, StoreError> {
62        Ok(lock(&self.0).clone())
63    }
64}

The chat hub. Its storage can be told to refuse writes, to see what a failed save does.

67#[derive(Clone, Default)]
68pub struct MemChatHub {
69    doc: Arc<Mutex<ChatState>>,
70    refuse_writes: Arc<AtomicBool>,
71}

73impl MemChatHub {

From now on (until told otherwise) a copy that would have to be saved is refused.

75    pub fn refuse_writes(&self, refuse: bool) {
76        self.refuse_writes.store(refuse, Ordering::SeqCst);
77    }
78}

80impl Hub for MemChatHub {
81    type Document = ChatState;
82
83    async fn merge(&self, theirs: ChatState) -> Result<Merged<ChatState>, StoreError> {
84        let mut doc = lock(&self.doc);
85        if !theirs.is_newer_than(&doc) {
86            return Ok(Merged { document: doc.clone(), changed: 0 });
87        }
88        if self.refuse_writes.load(Ordering::SeqCst) {
89            debug!("mem chat hub: refusing to keep v{}", theirs.version);
90            return Err(StoreError::Unavailable);
91        }
92        let adopted = doc.merge(theirs);
93        Ok(Merged { document: doc.clone(), changed: usize::from(adopted) })
94    }
95
96    async fn current(&self) -> Result<ChatState, StoreError> {
97        Ok(lock(&self.doc).clone())
98    }
99}

The journal.

102#[derive(Clone, Default)]
103pub struct MemJournal(Arc<Mutex<Vec<(DeviceName, StoredLine)>>>);

105impl MemJournal {
106    fn held_in(log: &[(DeviceName, StoredLine)], device: &DeviceName) -> DeviceCursor {
107        DeviceCursor::new(log.iter().filter(|(d, _)| d == device).count() as u64)
108    }
109}
110
111impl Journal for MemJournal {
112    async fn held(&self, device: &DeviceName) -> Result<DeviceCursor, StoreError> {
113        Ok(Self::held_in(&lock(&self.0), device))
114    }
115
116    async fn append(&self, device: &DeviceName, at: DeviceCursor, lines: &[EntryLine]) -> Result<Appended, StoreError> {
117        let mut log = lock(&self.0);
118        let held = Self::held_in(&log, device);
119        if lines.is_empty() {
120            return Ok(Appended::Accepted { held });
121        }
122        if at != held {
123            debug!("mem journal: {} asked to append at {} but holds {}", device.as_str(), at.get(), held.get());
124            return Ok(Appended::Misaligned { held });
125        }
126        for l in lines {
127            log.push((device.clone(), StoredLine::of(device, l)));
128        }
129        Ok(Appended::Accepted { held: held.after(lines.len()) })
130    }
131
132    async fn pull(&self, since: LogCursor) -> Result<Page, StoreError> {
133        let all: Vec<StoredLine> = lock(&self.0).iter().map(|(_, l)| l.clone()).collect();
134        Ok(Page::of(&all, since))
135    }
136}

The shelf of pictures.

139#[derive(Clone, Default)]
140pub struct MemPictures(Arc<Mutex<BTreeMap<String, Vec<u8>>>>);

142impl PictureShelf for MemPictures {
143    async fn has(&self, id: &SafePictureId) -> Result<bool, StoreError> {
144        Ok(lock(&self.0).contains_key(id.as_str()))
145    }
146
147    async fn put(&self, id: &SafePictureId, bytes: &[u8]) -> Result<Put, StoreError> {
148        let mut shelf = lock(&self.0);
149        if shelf.contains_key(id.as_str()) {
150            trace!("mem pictures: {} already kept", id.as_str());
151            return Ok(Put::AlreadyKept);
152        }
153        shelf.insert(id.as_str().to_owned(), bytes.to_vec());
154        Ok(Put::Stored)
155    }
156
157    async fn get(&self, id: &SafePictureId) -> Result<Option<Vec<u8>>, StoreError> {
158        Ok(lock(&self.0).get(id.as_str()).cloned())
159    }
160}
161
162#[derive(Default)]
163struct AllowanceState {
164    thinking: ThinkingLedger,
165    voice: VoiceDay,
166    last_failure: Option<SpeakError>,
167}

The allowance.

170#[derive(Clone, Default)]
171pub struct MemAllowance(Arc<Mutex<AllowanceState>>);

173impl Allowance for MemAllowance {
174    async fn thinking_room(&self, now: Millis, limit: &TokenLimit) -> Result<ThinkingRoom, StoreError> {
175        Ok(lock(&self.0).thinking.room(now, limit))
176    }
177
178    async fn charge_thinking(&self, now: Millis, cost: Tokens) -> Result<(), StoreError> {
179        lock(&self.0).thinking.charge(now, cost);
180        Ok(())
181    }
182
183    async fn thinking_usage(&self, now: Millis, limit: &TokenLimit) -> Result<ThinkingUsage, StoreError> {
184        Ok(lock(&self.0).thinking.usage(now, limit))
185    }
186
187    async fn reserve_voice(&self, now: Millis, chars: u32, cap: u32) -> Result<VoiceReservation, ReserveError> {
188        lock(&self.0).voice.reserve(now, chars, cap)
189    }
190
191    async fn settle_voice(&self, reservation: VoiceReservation, how: Settlement) -> Result<(), StoreError> {
192        let mut state = lock(&self.0);
193        match how {
194            Settlement::Spoken => state.last_failure = None,
195            Settlement::NotSpoken(why) => {
196                state.voice.release(reservation);
197                state.last_failure = Some(why);
198            }
199        }
200        Ok(())
201    }
202
203    async fn voice_spend(&self, now: Millis) -> Result<VoiceSpend, StoreError> {
204        let state = lock(&self.0);
205        Ok(VoiceSpend { spent_today: state.voice.spent_on(whiskers_ports::Day::containing(now)), last_failure: state.last_failure.clone() })
206    }
207}

Secrets, as raw text exactly as a store would hold it.

210#[derive(Clone, Default)]
211pub struct MemSecrets(Arc<Mutex<BTreeMap<String, String>>>);

213impl MemSecrets {
214    pub fn with(entries: &[(SecretName, &str)]) -> Self {
215        let s = Self::default();
216        for (n, v) in entries {
217            s.set(*n, v);
218        }
219        s
220    }
221
222    pub fn set(&self, name: SecretName, raw: &str) {
223        lock(&self.0).insert(name.operator_name().to_owned(), raw.to_owned());
224    }
225}
226
227impl Secrets for MemSecrets {
228    async fn get(&self, name: SecretName) -> Result<Lookup, StoreError> {
229        Ok(Lookup::of(lock(&self.0).get(name.operator_name()).map(String::as_str)))
230    }
231}

A Jev that answers what it is told to, and counts how often it was asked.

234pub struct ScriptedJudge {
235    pub verdict: Mutex<Result<Verdict, JudgeError>>,
236    pub ranking: Mutex<Result<Vec<f32>, JudgeError>>,

What choose_icon answers (the options it was last given are in icon_options).

238    pub icon: Mutex<Result<Option<IconId>, JudgeError>>,
239    pub icon_options: Mutex<Vec<IconId>>,
240    pub asked: AtomicUsize,
241}

243impl Default for ScriptedJudge {
244    fn default() -> Self {
245        Self {
246            verdict: Mutex::new(Ok(Verdict::Allow)),
247            ranking: Mutex::new(Ok(Vec::new())),
248            icon: Mutex::new(Ok(None)),
249            icon_options: Mutex::new(Vec::new()),
250            asked: AtomicUsize::new(0),
251        }
252    }
253}
254
255impl Judge for ScriptedJudge {
256    async fn check(&self, _direction: Direction, _age: Age, _text: &str) -> Result<Verdict, JudgeError> {
257        self.asked.fetch_add(1, Ordering::SeqCst);
258        lock(&self.verdict).clone()
259    }
260
261    async fn rerank(&self, _query: &str, candidates: &[String]) -> Result<Vec<f32>, JudgeError> {
262        self.asked.fetch_add(1, Ordering::SeqCst);
263        match &*lock(&self.ranking) {
264            Ok(p) if p.is_empty() => Ok(vec![0.5; candidates.len()]),
265            other => other.clone(),
266        }
267    }
268
269    async fn choose_icon(&self, _memory: &str, options: &[IconId]) -> Result<Option<IconId>, JudgeError> {
270        self.asked.fetch_add(1, Ordering::SeqCst);
271        *lock(&self.icon_options) = options.to_vec();
272        lock(&self.icon).clone()
273    }
274}

How a scripted embedder turns the texts of a batch into vectors, when a test wants it to answer by what it was asked rather than with one fixed answer.

278pub type EmbedWith = std::sync::Arc<dyn Fn(&[String]) -> Vec<Vec<f32>> + Send + Sync>;

An embedding model that answers what it is told to.

281pub struct ScriptedEmbedder {
282    pub answer: Mutex<Result<Vec<Vec<f32>>, EmbedError>>,

If set, used instead of answer.

284    pub with: Mutex<Option<EmbedWith>>,
285    pub asked: AtomicUsize,
286}

288impl Default for ScriptedEmbedder {
289    fn default() -> Self {
290        Self {
291            answer: Mutex::new(Err(EmbedError::Unreachable(Diagnostic::new("embedding server: not scripted")))),
292            with: Mutex::new(None),
293            asked: AtomicUsize::new(0),
294        }
295    }
296}
297
298impl Embedder for ScriptedEmbedder {
299    async fn embed(&self, batch: &EmbedBatch) -> Result<Vec<Vec<f32>>, EmbedError> {
300        self.asked.fetch_add(1, Ordering::SeqCst);
301        if let Some(f) = lock(&self.with).clone() {
302            return Ok(f(batch.texts()));
303        }
304        lock(&self.answer).clone()
305    }
306}

A voice that answers what it is told to.

309pub struct ScriptedVoice {
310    pub configured: Mutex<bool>,
311    pub speech: Mutex<Result<Vec<u8>, SpeakError>>,
312    pub credits: Mutex<Result<VoiceCredits, CreditsError>>,
313    pub asked: AtomicUsize,
314}

316impl Default for ScriptedVoice {
317    fn default() -> Self {
318        Self {
319            configured: Mutex::new(true),
320            speech: Mutex::new(Ok(vec![0xFF, 0xFB, 1, 2])),
321            credits: Mutex::new(Ok(VoiceCredits { tier: Some("free".into()), used: Some(10), limit: Some(10_000), resets_at_unix: Some(1_800_000_000) })),
322            asked: AtomicUsize::new(0),
323        }
324    }
325}
326
327impl Voice for ScriptedVoice {
328    async fn configured(&self) -> bool {
329        *lock(&self.configured)
330    }
331
332    async fn speak(&self, _line: &SpeechLine) -> Result<Audio, SpeakError> {
333        self.asked.fetch_add(1, Ordering::SeqCst);
334        if !*lock(&self.configured) {
335            return Err(SpeakError::NotConfigured);
336        }
337        lock(&self.speech).clone().map(Audio)
338    }

The scripted audio, and every character said for 100 ms in turn.

341    async fn speak_timed(&self, line: &SpeechLine) -> Result<TimedAudio, SpeakError> {
342        self.asked.fetch_add(1, Ordering::SeqCst);
343        if !*lock(&self.configured) {
344            return Err(SpeakError::NotConfigured);
345        }
346        let audio = Audio(lock(&self.speech).clone()?);
347        let spans = (0..line.chars()).map(|i| (AudioMs::new(i * 100), AudioMs::new(i * 100 + 100))).collect();
348        let unreadable = |e: TimingError| SpeakError::Unreadable(Diagnostic::new(e.to_string()));
349        let timing = LineTiming::new(line.as_str().to_owned(), spans).map_err(unreadable)?;
350        TimedAudio::new(audio, timing, line).map_err(unreadable)
351    }

353    async fn credits(&self) -> Result<VoiceCredits, CreditsError> {
354        lock(&self.credits).clone()
355    }
356}

A model that answers what it is told to.

359pub struct ScriptedThinker {
360    pub model: ModelName,
361    pub answer: Mutex<Result<Reply, ThinkError>>,
362    pub asked: AtomicUsize,
363}

365impl Default for ScriptedThinker {
366    fn default() -> Self {
367        Self {
368            model: ModelName::new("m").expect("a model name"),
369            answer: Mutex::new(Ok(Reply { status: 200, body: r#"{"content":[],"usage":{"input_tokens":70,"output_tokens":30}}"#.into() })),
370            asked: AtomicUsize::new(0),
371        }
372    }
373}
374
375impl Thinker for ScriptedThinker {
376    fn serves(&self) -> &ModelName {
377        &self.model
378    }
379
380    async fn think(&self, _request: &ThinkRequest) -> Result<Reply, ThinkError> {
381        self.asked.fetch_add(1, Ordering::SeqCst);
382        lock(&self.answer).clone()
383    }
384}

Every port, in memory. The fields are public so a test can script the outbound ports and read the state the service left behind.

388#[derive(Default)]
389pub struct MemBackend {
390    pub memory: MemMemoryHub,
391    pub household: MemHouseholdHub,
392    pub chat: MemChatHub,
393    pub journal: MemJournal,
394    pub pictures: MemPictures,
395    pub allowance: MemAllowance,
396    pub judge: ScriptedJudge,
397    pub embedder: ScriptedEmbedder,
398    pub icons: whiskers_ports::IconVectors,
399    pub voice: ScriptedVoice,
400    pub thinker: ScriptedThinker,
401    pub clock: ManualClock,
402}

Every capability, so a MemBackend is a complete Backend.

405impl HasMemory for MemBackend {
406    type Memory = MemMemoryHub;
407    fn memory(&self) -> &MemMemoryHub {
408        &self.memory
409    }
410}
411impl HasHousehold for MemBackend {
412    type Household = MemHouseholdHub;
413    fn household(&self) -> &MemHouseholdHub {
414        &self.household
415    }
416}
417impl HasChat for MemBackend {
418    type Chat = MemChatHub;
419    fn chat(&self) -> &MemChatHub {
420        &self.chat
421    }
422}
423impl HasJournal for MemBackend {
424    type Journal = MemJournal;
425    fn journal(&self) -> &MemJournal {
426        &self.journal
427    }
428}
429impl HasPictures for MemBackend {
430    type Pictures = MemPictures;
431    fn pictures(&self) -> &MemPictures {
432        &self.pictures
433    }
434}
435impl HasAllowance for MemBackend {
436    type Allowance = MemAllowance;
437    fn allowance(&self) -> &MemAllowance {
438        &self.allowance
439    }
440}
441impl HasJudge for MemBackend {
442    type Judge = ScriptedJudge;
443    fn judge(&self) -> &ScriptedJudge {
444        &self.judge
445    }
446}
447impl whiskers_ports::HasIcons for MemBackend {
448    fn icon_vectors(&self) -> &whiskers_ports::IconVectors {
449        &self.icons
450    }
451}
452impl HasEmbedder for MemBackend {
453    type Embedder = ScriptedEmbedder;
454    fn embedder(&self) -> &ScriptedEmbedder {
455        &self.embedder
456    }
457}
458impl HasVoice for MemBackend {
459    type Voice = ScriptedVoice;
460    fn voice(&self) -> &ScriptedVoice {
461        &self.voice
462    }
463}
464impl HasThinker for MemBackend {
465    type Thinker = ScriptedThinker;
466    fn thinker(&self) -> &ScriptedThinker {
467        &self.thinker
468    }
469}
470impl HasClock for MemBackend {
471    type Clock = ManualClock;
472    fn clock(&self) -> &ManualClock {
473        &self.clock
474    }
475}

477impl MemBackend {

A backend whose clock reads ms.

479    pub fn at(ms: u64) -> Self {
480        let b = Self::default();
481        b.clock.set(Millis::new(ms));
482        b
483    }

485    pub fn now(&self) -> Millis {
486        self.clock.now()
487    }
488}

The harness of the in-memory adapter: a fixture for every capability, so the whole suite runs on it (Suite::complete(MemFixtures)). A restart is the identity: what is in memory stays in memory, which is the most a process that keeps nothing on disk can promise.

493#[derive(Clone, Copy, Default)]
494pub struct MemFixtures;

496macro_rules! mem_stateful {
497    ($fixture:ident, $port:ident) => {
498        impl crate::suite::fixture::$fixture for MemFixtures {
499            type Port = $port;
500            fn fresh(&self) -> $port {
501                $port::default()
502            }
503            fn restart(&self, port: $port) -> $port {
504                port
505            }
506        }
507    };
508}
509
510mem_stateful!(MemoryFixture, MemMemoryHub);
511mem_stateful!(HouseholdFixture, MemHouseholdHub);
512mem_stateful!(ChatFixture, MemChatHub);
513mem_stateful!(JournalFixture, MemJournal);
514mem_stateful!(PicturesFixture, MemPictures);
515mem_stateful!(AllowanceFixture, MemAllowance);
516
517impl crate::suite::fixture::ChatFaultFixture for MemFixtures {
518    type Port = MemChatHub;
519    fn fresh_refusing(&self) -> (MemChatHub, crate::suite::hubs::WriteSwitch) {
520        let hub = MemChatHub::default();
521        let switch = hub.clone();
522        (hub, Box::new(move |refuse| switch.refuse_writes(refuse)))
523    }
524    fn restart_refusing(&self, port: MemChatHub) -> MemChatHub {
525        port
526    }
527}
528
529impl crate::suite::fixture::JudgeFixture for MemFixtures {
530    type Port = ScriptedJudge;
531    fn without_credentials(&self) -> ScriptedJudge {
532        let judge = ScriptedJudge::default();
533        *lock(&judge.verdict) = Err(JudgeError::NotConfigured(SecretName::JevKey));
534        *lock(&judge.ranking) = Err(JudgeError::NotConfigured(SecretName::JevKey));
535        *lock(&judge.icon) = Err(JudgeError::NotConfigured(SecretName::JevKey));
536        judge
537    }
538}
539
540impl crate::suite::fixture::VoiceFixture for MemFixtures {
541    type Port = ScriptedVoice;
542    fn without_credentials(&self) -> ScriptedVoice {
543        let voice = ScriptedVoice::default();
544        *lock(&voice.configured) = false;
545        voice
546    }
547}
548
549impl crate::suite::fixture::ClockFixture for MemFixtures {
550    type Port = ManualClock;
551    fn clock(&self) -> ManualClock {
552        // A manual clock that reads a plausible time (the law asks for one after 2020).
553        ManualClock::at(1_800_000_000_000)
554    }
555}