cuttlefish_host/
infer.rs

1//! Where inference comes from.
2//!
3//! The runner talks to models only through [`InferBackend`], so the whole
4//! reactor loop can be tested without a model. llama.cpp arrives behind this
5//! same trait later; nothing above it should need to change.
6
7use async_trait::async_trait;
8
9/// What one completed generation produced.
10pub struct InferResult {
11    /// The generated text.
12    pub text: String,
13    /// Tokens consumed by the prompt.
14    pub tokens_in: u32,
15    /// Tokens generated.
16    pub tokens_out: u32,
17}
18
19/// Anything that can serve an inference request.
20#[async_trait]
21pub trait InferBackend: Send + Sync {
22    /// Generate from `prompt`, invoking `on_token` once per token.
23    ///
24    /// `on_token` returns whether to keep going; returning `false` ends
25    /// generation early, which is how a guest's `Stop` verdict is honoured.
26    ///
27    /// The `for<'t>` is load-bearing. `#[async_trait]` rewrites elided lifetimes
28    /// into named ones, which would make this closure non-generic over the
29    /// token's lifetime and leave implementations unable to hand it a local
30    /// `&str` — an E0597 that appears only in the implementor, with a message
31    /// that does not obviously point back here.
32    async fn infer(
33        &self,
34        prompt: &str,
35        max_tokens: u32,
36        on_token: &mut (dyn for<'t> FnMut(&'t str) -> bool + Send),
37    ) -> anyhow::Result<InferResult>;
38
39    /// Identifier recorded in the job's usage accounting.
40    fn model_name(&self) -> String;
41}
42
43/// A deterministic fake that streams a fixed reply word by word.
44///
45/// Enough to exercise streaming, early-stop, and token accounting with no model
46/// present.
47pub struct StubBackend {
48    /// What every request generates.
49    pub reply: String,
50}
51
52impl Default for StubBackend {
53    fn default() -> Self {
54        Self {
55            reply: "a stub summary".into(),
56        }
57    }
58}
59
60#[async_trait]
61impl InferBackend for StubBackend {
62    async fn infer(
63        &self,
64        prompt: &str,
65        max_tokens: u32,
66        on_token: &mut (dyn for<'t> FnMut(&'t str) -> bool + Send),
67    ) -> anyhow::Result<InferResult> {
68        let mut out = String::new();
69        let mut tokens_out = 0u32;
70
71        for word in self.reply.split_whitespace().take(max_tokens as usize) {
72            let piece = if out.is_empty() {
73                word.to_string()
74            } else {
75                format!(" {word}")
76            };
77            tokens_out += 1;
78
79            let keep_going = on_token(&piece);
80            out.push_str(&piece);
81            if !keep_going {
82                break;
83            }
84
85            // Yielding between tokens is not cosmetic. Without an await point
86            // the entire loop runs inside a single poll, every token lands in
87            // the channel at once, and the host can never interleave a guest's
88            // Stop verdict — the early-stop path would look implemented while
89            // being unreachable. A real backend awaits naturally; this one has
90            // to do it deliberately.
91            tokio::task::yield_now().await;
92        }
93
94        Ok(InferResult {
95            text: out,
96            tokens_in: prompt.split_whitespace().count() as u32,
97            tokens_out,
98        })
99    }
100
101    fn model_name(&self) -> String {
102        "stub".into()
103    }
104}
cuttlefish_host/infer.rs

cuttlefish_host/
infer.rs
