Skip to main content

codec_rs/
frame.rs

1// SPDX-License-Identifier: MIT
2//! `CodecFrame` and the pluggable map cache.
3
4use std::collections::HashMap;
5use std::sync::{Arc, Mutex};
6
7use crate::map::TokenizerMap;
8
9/// One streaming frame produced by a Codec-compliant server.
10///
11/// Identical shape across MessagePack and Protobuf wire modes — only
12/// serialization differs.
13#[derive(Debug, Clone, Default, PartialEq, Eq)]
14pub struct CodecFrame {
15    /// Token IDs emitted by the model in this chunk.
16    pub ids: Vec<u32>,
17    /// `true` on the final frame — no further frames follow.
18    pub done: bool,
19    /// Set on the final frame: e.g. `"length"`, `"stop"`, `"eos_token"`, `"error"`.
20    pub finish_reason: Option<String>,
21}
22
23impl CodecFrame {
24    pub fn new(ids: Vec<u32>, done: bool, finish_reason: Option<String>) -> Self {
25        Self { ids, done, finish_reason }
26    }
27}
28
29/// Pluggable cache for loaded maps. The default is [`MemoryMapCache`].
30///
31/// Implement this trait to plug in IndexedDB, Redis, on-disk storage, etc.
32pub trait MapCache: Send + Sync {
33    /// Returns the cached map for `key`, or `None`.
34    fn get(&self, key: &str) -> Option<Arc<TokenizerMap>>;
35    /// Stores `map` under `key`.
36    fn set(&self, key: &str, map: Arc<TokenizerMap>);
37}
38
39/// Alias matching the .NET `IMapCache` name for users familiar with the
40/// .NET surface.
41pub use MapCache as IMapCache;
42
43/// Default in-memory [`MapCache`] backed by a mutex-guarded `HashMap`.
44#[derive(Debug, Default)]
45pub struct MemoryMapCache {
46    inner: Mutex<HashMap<String, Arc<TokenizerMap>>>,
47}
48
49impl MemoryMapCache {
50    pub fn new() -> Self {
51        Self { inner: Mutex::new(HashMap::new()) }
52    }
53}
54
55impl MapCache for MemoryMapCache {
56    fn get(&self, key: &str) -> Option<Arc<TokenizerMap>> {
57        self.inner.lock().ok()?.get(key).cloned()
58    }
59
60    fn set(&self, key: &str, map: Arc<TokenizerMap>) {
61        if let Ok(mut g) = self.inner.lock() {
62            g.insert(key.to_string(), map);
63        }
64    }
65}