Skip to main content

superfluid_engine/
engine.rs

1//! The [`Engine`] trait.
2
3use superfluid_abi::{
4    MatchResult, OpStatus, StateSpaceDesc, StrategyGrant,
5    StrategyRegistration, TickEvents, TickPlan, TokenRange, Status,
6};
7
8use crate::rings::Rings;
9
10pub type SeqHandle = u64;
11
12pub type OpHandle = u64;
13
14#[derive(Debug, Clone)]
15pub struct SpaceConfig {
16    pub space_id: u32,
17    pub kind: u32,
18    pub name: &'static str,
19    pub version_tag: u64,
20    pub bytes_per_token: u64,
21    pub blob_bytes: u64,
22    pub page_size_tokens: u32,
23    pub snapshot_cadence: u32,
24    pub snapshot_interval_tokens: u32,
25    pub flags: u32,
26}
27
28#[derive(Debug, Clone)]
29pub struct EngineConfig {
30    pub max_batch_size: u32,
31    pub pool_bytes: u64,
32    pub spaces: Vec<SpaceConfig>,
33    pub seed_ttl_ticks: u64,
34    pub vocab: u32,
35    pub kernel_caps: u32,
36    pub image_token_id: u32,
37    pub media_tokens_per_image: u32,
38    pub scripted: Vec<u32>,
39    pub tick_delay: std::time::Duration,
40}
41
42impl Default for EngineConfig {
43    fn default() -> Self {
44        use superfluid_abi::space_kind;
45        EngineConfig {
46            max_batch_size: 8,
47            pool_bytes: 1 << 20,
48            spaces: vec![
49                SpaceConfig {
50                    space_id: 1,
51                    kind: space_kind::PAGED_TOKEN_KV,
52                    name: "kv.full",
53                    version_tag: 1,
54                    bytes_per_token: 64,
55                    blob_bytes: 0,
56                    page_size_tokens: 16,
57                    snapshot_cadence: 0,
58                    snapshot_interval_tokens: 0,
59                    flags: 0,
60                },
61                SpaceConfig {
62                    space_id: 2,
63                    kind: space_kind::RECURRENT_BLOB,
64                    name: "gdn.blob",
65                    version_tag: 1,
66                    bytes_per_token: 0,
67                    blob_bytes: 4096,
68                    page_size_tokens: 0,
69                    snapshot_cadence: 2,
70                    snapshot_interval_tokens: 32,
71                    flags: 0,
72                },
73            ],
74            seed_ttl_ticks: 8,
75            image_token_id: 0,
76            media_tokens_per_image: 4,
77            vocab: 32000,
78            kernel_caps: 0b111,
79            scripted: Vec::new(),
80            tick_delay: std::time::Duration::ZERO,
81        }
82    }
83}
84
85#[derive(Debug, Clone, PartialEq, Eq)]
86pub enum StageOutput {
87    Hidden(Vec<u8>),
88    Token(u32),
89}
90
91#[derive(Debug, Clone, PartialEq)]
92pub struct TranscriptSegment {
93    pub start_ms: i32,
94    pub end_ms: i32,
95    pub text: String,
96    pub avg_logprob: f32,
97    pub no_speech_prob: f32,
98    pub compression_ratio: f32,
99    pub temperature: f32,
100}
101
102#[derive(Debug, Clone, Default, PartialEq, Eq)]
103pub struct TranscribeParams {
104    pub language: Option<String>,
105    pub translate: bool,
106    pub timestamps: bool,
107    pub prompt: Option<String>,
108}
109
110pub type SegmentSink<'a> = &'a mut (dyn FnMut(i32, i32, &str) -> bool + 'a);
111
112#[derive(Debug, Clone, PartialEq, Default)]
113pub struct Transcription {
114    pub text: String,
115    pub language: String,
116    pub duration_ms: i32,
117    pub segments: Vec<TranscriptSegment>,
118}
119
120pub trait Engine {
121    fn tick<'a>(
122        &'a mut self,
123        plan: &TickPlan,
124        arena: &superfluid_abi::RecordArena,
125        rings: &mut dyn Rings,
126    ) -> Result<&'a TickEvents, Status>;
127
128    fn pump(&mut self);
129
130    fn attach_ring(&mut self, _ring: &crate::rings::RingAttachment) -> Result<(), Status> {
131        Ok(())
132    }
133
134    fn state_spaces(&self) -> &[StateSpaceDesc];
135
136    fn kv_bits(&self) -> u32 {
137        0
138    }
139
140    fn strategy_register(
141        &mut self,
142        reg: &StrategyRegistration,
143    ) -> Result<&StrategyGrant, Status>;
144
145    fn space_match(
146        &mut self,
147        space_id: u32,
148        span_tokens: &[u32],
149        media_deps: &[[u8; 32]],
150    ) -> Result<&MatchResult, Status>;
151
152    fn space_match_ref(
153        &mut self,
154        space_id: u32,
155        _span: &superfluid_abi::TokenRef,
156        span_tokens: &[u32],
157        media_deps: &[[u8; 32]],
158    ) -> Result<&MatchResult, Status> {
159        self.space_match(space_id, span_tokens, media_deps)
160    }
161
162    fn seed_acquire(
163        &mut self,
164        span_tokens: &[u32],
165        prefix_len: u64,
166        determinism_class: u8,
167    ) -> Result<u64, Status>;
168
169    fn seed_acquire_ref(
170        &mut self,
171        _span: &superfluid_abi::TokenRef,
172        span_tokens: &[u32],
173        prefix_len: u64,
174        determinism_class: u8,
175    ) -> Result<u64, Status> {
176        self.seed_acquire(span_tokens, prefix_len, determinism_class)
177    }
178
179    fn seed_release(&mut self, seed_handle: u64) -> Result<(), Status>;
180
181    fn seed_lease_ticks(&self) -> Option<u64> {
182        None
183    }
184
185    fn seed_adopt(&mut self, _seq: SeqHandle, _span_tokens: &[u32]) -> Result<u64, Status> {
186        Err(Status::Unsupported)
187    }
188
189    fn media_probe(&mut self, _image_path: &str) -> Result<superfluid_abi::MediaInfo, Status> {
190        Err(Status::Unsupported)
191    }
192
193    fn media_encode(
194        &mut self,
195        _image_path: &str,
196    ) -> Result<(u64, superfluid_abi::MediaInfo), Status> {
197        Err(Status::Unsupported)
198    }
199
200    fn media_release(&mut self, _media_handle: u64) -> Result<(), Status> {
201        Err(Status::Unsupported)
202    }
203
204    fn media_bind(&mut self, _lane_tag: u64, _media_handle: u64, _token_offset: u32) -> Result<(), Status> {
205        Err(Status::Unsupported)
206    }
207
208    fn seed_adopt_ref(
209        &mut self,
210        seq: SeqHandle,
211        _span: &superfluid_abi::TokenRef,
212        span_tokens: &[u32],
213    ) -> Result<u64, Status> {
214        self.seed_adopt(seq, span_tokens)
215    }
216
217    fn seq_fork(&mut self, parent: SeqHandle, flags: u32) -> Result<SeqHandle, Status>;
218
219    fn lane_sequence(&self, _lane_tag: u64) -> Option<SeqHandle> {
220        None
221    }
222
223    fn grammar_create(&mut self, _json_schema: &str) -> Result<u32, Status> {
224        Err(Status::Unsupported)
225    }
226
227    fn grammar_create_structural(&mut self, _tag_json: &str) -> Result<u32, Status> {
228        Err(Status::Unsupported)
229    }
230
231    fn grammar_free(&mut self, _handle: u32) -> Result<(), Status> {
232        Err(Status::Unsupported)
233    }
234
235    fn logit_bias_create(&mut self, _tokens: &[i32], _values: &[f32]) -> Result<u32, Status> {
236        Err(Status::Unsupported)
237    }
238
239    fn logit_bias_free(&mut self, _handle: u32) -> Result<(), Status> {
240        Err(Status::Unsupported)
241    }
242
243    fn embed(&mut self, _tokens: &[u32]) -> Result<Vec<f32>, Status> {
244        Err(Status::Unsupported)
245    }
246
247    fn transcribe(
248        &mut self,
249        _audio_path: &str,
250        _params: &TranscribeParams,
251        _on_segment: Option<SegmentSink<'_>>,
252    ) -> Result<Transcription, Status> {
253        Err(Status::Unsupported)
254    }
255
256    fn lora_load(&mut self, _adapter_path: &str) -> Result<(), Status> {
257        Err(Status::Unsupported)
258    }
259
260    fn lora_unload(&mut self) -> Result<(), Status> {
261        Err(Status::Unsupported)
262    }
263
264    fn lora_id(&self) -> Option<String> {
265        None
266    }
267
268    fn capability_descriptor(&self) -> Option<String> {
269        None
270    }
271
272    fn create_sequence(&mut self) -> Result<SeqHandle, Status> {
273        Err(Status::Unsupported)
274    }
275
276    fn free_sequence(&mut self, _seq: SeqHandle) -> Result<(), Status> {
277        Err(Status::Unsupported)
278    }
279
280    fn publish_sequence(&mut self, _seq: SeqHandle, _tokens: &[u32]) -> Result<(), Status> {
281        Err(Status::Unsupported)
282    }
283
284    /// Drops the prefix-cache entry published under exactly `tokens`, unless
285    /// a seed holds it; none there is not an error.
286    fn unpublish(&mut self, _tokens: &[u32]) -> Result<(), Status> {
287        Err(Status::Unsupported)
288    }
289
290    fn space_export_size(
291        &mut self,
292        seq: SeqHandle,
293        space_id: u32,
294        range: TokenRange,
295        encoding: u8,
296    ) -> Result<(u64, u64), Status>;
297
298    fn space_snapshot(
299        &mut self,
300        seq: SeqHandle,
301        space_id: u32,
302        boundary_pos: u64,
303        dst_len: u64,
304        sizing_gen: u64,
305    ) -> Result<OpHandle, Status>;
306
307    fn space_snapshot_boundary(
308        &mut self,
309        _seq: SeqHandle,
310        _space_id: u32,
311        _cap: u64,
312    ) -> Result<u64, Status> {
313        Err(Status::Unsupported)
314    }
315
316    fn cache_evict_entries(&mut self, _keys: &[([u8; 32], [u8; 32])]) -> Result<(), Status> {
317        Err(Status::Unsupported)
318    }
319
320    fn space_restore(
321        &mut self,
322        seq: SeqHandle,
323        space_id: u32,
324        src: &[u8],
325    ) -> Result<OpHandle, Status>;
326
327    fn space_trim(&mut self, seq: SeqHandle, space_id: u32, new_len: u64) -> Result<(), Status>;
328
329    fn space_demote(
330        &mut self,
331        seq: SeqHandle,
332        space_id: u32,
333        range: TokenRange,
334        encoding: u8,
335        dst_len: u64,
336        sizing_gen: u64,
337    ) -> Result<OpHandle, Status>;
338
339    fn space_promote(
340        &mut self,
341        seq: SeqHandle,
342        space_id: u32,
343        range: TokenRange,
344        src: &[u8],
345    ) -> Result<OpHandle, Status>;
346
347    fn op_poll(&mut self, op: OpHandle) -> Result<OpStatus, Status>;
348
349    fn take_op_output(&mut self, _op: OpHandle) -> Option<Vec<u8>> {
350        None
351    }
352
353    fn n_layers(&self) -> u32 {
354        0
355    }
356
357    fn forward_stage(
358        &mut self,
359        _tokens: &[u32],
360        _start_layer: u32,
361        _end_layer: u32,
362        _hidden_in: Option<Vec<u8>>,
363    ) -> Result<StageOutput, Status> {
364        Err(Status::Unsupported)
365    }
366
367    fn op_cancel(&mut self, op: OpHandle) -> Result<(), Status>;
368
369    fn cache_evict(
370        &mut self,
371        evictable_cache_classes: u64,
372        protected_quota_bytes: u64,
373        max_evict_bytes: u64,
374        bytes_target: u64,
375    ) -> Result<u64, Status>;
376}