cerno_host/
lib.rs

1//! The seam between cerno and whatever runs the model.
2//!
3//! A host knows nothing about Noul, Choice or Score beyond the [`Shape`] of what was offered,
4//! which only a host that takes questions rather than prompts reads. It answers exactly one
5//! question: given a prompt, what is the probability distribution over the *first* token the
6//! model would generate? All primitive logic lives in `cerno-core` and therefore holds for every
7//! host equally.
8//!
9//! Three adapters cover the runtimes: [`OllamaHost`] for Ollama's native API,
10//! [`OpenAiCompatHost`] for everything speaking `/v1/chat/completions`, and [`SystemOneHost`]
11//! for Ollama's `/v1/systemone`, which writes its own prompt for its decision models. [`connect`]
12//! picks one from a [`HostKind`]; for [`HostKind::Ollama`] it builds an [`OllamaRouter`] over
13//! the two Ollama adapters, which sends each model to the endpoint made for it.
14
15mod ollama;
16mod openai;
17mod router;
18mod systemone;
19
20pub use ollama::OllamaHost;
21pub use openai::{Flavour, OpenAiCompatHost};
22pub use router::OllamaRouter;
23pub use systemone::SystemOneHost;
24
25use async_trait::async_trait;
26use std::sync::Arc;
27use std::time::Duration;
28
29/// Which runtime a server talks to. One per process.
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum HostKind {
32    Ollama,
33    /// OpenAI itself, or any server that speaks its API and nothing more.
34    OpenAi,
35    Vllm,
36    LlamaCpp,
37    LmStudio,
38}
39
40impl HostKind {
41    /// Where the runtime listens out of the box.
42    pub fn default_url(self) -> &'static str {
43        match self {
44            HostKind::Ollama => "http://localhost:11434",
45            HostKind::OpenAi => "https://api.openai.com/v1",
46            HostKind::Vllm => "http://localhost:8000/v1",
47            HostKind::LlamaCpp => "http://localhost:8080/v1",
48            HostKind::LmStudio => "http://localhost:1234/v1",
49        }
50    }
51}
52
53#[derive(Debug, thiserror::Error)]
54#[error("unknown host {0:?}; expected one of ollama, openai, vllm, llamacpp, lmstudio")]
55pub struct UnknownHostKind(String);
56
57impl std::str::FromStr for HostKind {
58    type Err = UnknownHostKind;
59
60    fn from_str(s: &str) -> Result<Self, Self::Err> {
61        match s.trim().to_ascii_lowercase().as_str() {
62            "ollama" => Ok(HostKind::Ollama),
63            "openai" => Ok(HostKind::OpenAi),
64            "vllm" => Ok(HostKind::Vllm),
65            "llamacpp" | "llama.cpp" => Ok(HostKind::LlamaCpp),
66            "lmstudio" => Ok(HostKind::LmStudio),
67            _ => Err(UnknownHostKind(s.to_string())),
68        }
69    }
70}
71
72impl std::fmt::Display for HostKind {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.write_str(match self {
75            HostKind::Ollama => "ollama",
76            HostKind::OpenAi => "openai",
77            HostKind::Vllm => "vllm",
78            HostKind::LlamaCpp => "llamacpp",
79            HostKind::LmStudio => "lmstudio",
80        })
81    }
82}
83
84/// Build the adapter for `kind`. Ollama takes no API key; the others send it as a bearer token.
85pub fn connect(
86    kind: HostKind,
87    url: &str,
88    api_key: Option<String>,
89    timeout: Duration,
90) -> Result<Arc<dyn ModelHost>, HostError> {
91    let flavour = match kind {
92        HostKind::Ollama => return Ok(Arc::new(OllamaRouter::new(url, timeout)?)),
93        HostKind::OpenAi => Flavour::Generic,
94        HostKind::Vllm => Flavour::Vllm,
95        HostKind::LlamaCpp => Flavour::LlamaCpp,
96        HostKind::LmStudio => Flavour::LmStudio,
97    };
98    Ok(Arc::new(OpenAiCompatHost::new(
99        url, api_key, flavour, timeout,
100    )?))
101}
102
103/// What a host can and cannot do. `cerno-core` reads this to size its label alphabet.
104#[derive(Debug, Clone, Copy, PartialEq, Eq)]
105pub struct HostCapabilities {
106    /// How many ranked tokens the host will report for one position.
107    ///
108    /// Ollama and OpenAI cap this at 20. A label beyond that rank is invisible, which is exactly
109    /// the bound that limits a Choice to 20 options.
110    pub max_top_logprobs: usize,
111}
112
113/// One prompt, one token, full distribution.
114#[derive(Debug, Clone)]
115pub struct FirstTokenRequest {
116    pub model: String,
117    pub system: Option<String>,
118    pub user: String,
119    pub top_logprobs: usize,
120    /// How long the host should keep the model resident after answering. Only Ollama has this;
121    /// the OpenAI-compatible runtimes keep their model loaded for the life of the process.
122    pub keep_alive: Option<String>,
123    /// The question behind `user`; see [`Offer`].
124    pub offer: Offer,
125}
126
127/// What the model was asked, before it was rendered into a prompt.
128///
129/// Token hosts ignore it and read `system`/`user`. A host that is asked questions rather than
130/// prompts — [`SystemOneHost`] — reads nothing else.
131#[derive(Debug, Clone, PartialEq)]
132pub struct Offer {
133    /// The question's id, which the structured request keys its answer by.
134    pub question_id: String,
135    pub state: String,
136    /// The question text; `None` when a choice or score has none worth asking.
137    pub question: Option<String>,
138    pub shape: Shape,
139    /// `(label, text)` per option, in request order. A yes/no offer is `[("A","Yes"),("B","No")]`.
140    pub options: Vec<(String, String)>,
141}
142
143/// How the offered options relate to each other.
144#[derive(Debug, Clone, Copy, PartialEq, Eq)]
145pub enum Shape {
146    /// Two options, "Yes" first.
147    YesNo,
148    /// Any one of the options.
149    Pick,
150    /// Ordered levels, lowest first.
151    Scale,
152}
153
154/// The distribution over the first generated token.
155#[derive(Debug, Clone)]
156pub struct FirstTokenDistribution {
157    /// Ranked tokens, highest logprob first.
158    pub tokens: Vec<(String, f64)>,
159    /// The lowest logprob the host reported.
160    ///
161    /// A label missing from `tokens` ranked below every entry, so this value is a strict upper
162    /// bound on its logprob. `cerno-core` substitutes it and flags the answer as truncated.
163    pub floor: f64,
164    pub input_tokens: u32,
165    pub latency: Duration,
166}
167
168impl FirstTokenDistribution {
169    /// The logprob of `token`, or `None` when it fell outside the reported window.
170    pub fn logprob(&self, token: &str) -> Option<f64> {
171        self.tokens
172            .iter()
173            .find(|(t, _)| t == token)
174            .map(|(_, lp)| *lp)
175    }
176}
177
178#[derive(Debug, thiserror::Error)]
179pub enum HostError {
180    #[error("host unreachable: {0}")]
181    Unavailable(String),
182
183    #[error("host timed out after {0:?}")]
184    Timeout(Duration),
185
186    #[error("host returned {status}: {body}")]
187    Status { status: u16, body: String },
188
189    /// The host answered, but not in the shape we need — most often because the model or the
190    /// runtime does not report logprobs at all.
191    #[error(
192        "host response carried no logprobs (model {model:?}); \
193             the runtime must support top_logprobs for cerno to work"
194    )]
195    NoLogprobs { model: String },
196
197    #[error("could not parse host response: {0}")]
198    Protocol(String),
199
200    /// A base URL no request could be sent to. Caught when the host is built, because otherwise
201    /// the server starts and every request fails with reqwest's "builder error".
202    #[error("host URL {url:?} is not usable: {reason}")]
203    InvalidUrl { url: String, reason: String },
204}
205
206/// The base URL with any trailing slash removed, or why it cannot be one.
207fn base_url(raw: &str) -> Result<String, HostError> {
208    let trimmed = raw.trim().trim_end_matches('/');
209    let invalid = |reason: String| HostError::InvalidUrl {
210        url: raw.to_string(),
211        reason,
212    };
213
214    let parsed = reqwest::Url::parse(trimmed).map_err(|e| invalid(e.to_string()))?;
215    if !matches!(parsed.scheme(), "http" | "https") {
216        // `localhost:11434` parses, with `localhost` as its scheme, so naming the scheme found
217        // would only confuse.
218        return Err(invalid("it must start with http:// or https://".into()));
219    }
220    if parsed.host_str().is_none() {
221        return Err(invalid("it names no host".into()));
222    }
223    Ok(trimmed.to_string())
224}
225
226/// How much of a failed response's body travels on in the error. The whole body is logged; the
227/// error reaches the caller of the service, and an upstream body can carry account or request
228/// details that are the operator's business, not theirs.
229const MAX_ERROR_BODY: usize = 300;
230
231impl HostError {
232    /// A non-success status from the host, with its body logged in full and kept short.
233    ///
234    /// An authentication failure keeps none of it. OpenAI answers a wrong key with "Incorrect
235    /// API key provided: sk-…abcd", and part of the operator's key is not the caller's business
236    /// however short the excerpt.
237    fn status(status: u16, body: String) -> Self {
238        tracing::warn!(status, body = %body, "host answered with an error");
239
240        if matches!(status, 401 | 403) {
241            return Self::Status {
242                status,
243                body: "the host refused cerno's credentials; the service log has its answer".into(),
244            };
245        }
246
247        let body = match body.char_indices().nth(MAX_ERROR_BODY) {
248            Some((cut, _)) => format!("{}…", &body[..cut]),
249            None => body,
250        };
251        Self::Status { status, body }
252    }
253
254    /// Map a transport failure, keeping a timeout a timeout wherever in the exchange it struck —
255    /// sending the request or reading the body.
256    fn transport(error: &reqwest::Error, timeout: Duration) -> Self {
257        if error.is_timeout() {
258            Self::Timeout(timeout)
259        } else {
260            Self::Unavailable(error.to_string())
261        }
262    }
263}
264
265#[async_trait]
266pub trait ModelHost: Send + Sync {
267    fn capabilities(&self) -> HostCapabilities;
268
269    /// A human-readable name for the host, used in logs and diagnostics.
270    fn name(&self) -> &str;
271
272    async fn first_token(
273        &self,
274        req: FirstTokenRequest,
275    ) -> Result<FirstTokenDistribution, HostError>;
276}
277
278/// A small `Pick` offer for tests that only need one to exist.
279#[cfg(test)]
280pub(crate) fn test_offer() -> Offer {
281    Offer {
282        question_id: "q".into(),
283        state: "s".into(),
284        question: Some("Which?".into()),
285        shape: Shape::Pick,
286        options: vec![("A".into(), "x".into()), ("B".into(), "y".into())],
287    }
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293
294    #[test]
295    fn every_host_kind_parses_from_its_own_name() {
296        for kind in [
297            HostKind::Ollama,
298            HostKind::OpenAi,
299            HostKind::Vllm,
300            HostKind::LlamaCpp,
301            HostKind::LmStudio,
302        ] {
303            assert_eq!(kind.to_string().parse::<HostKind>().unwrap(), kind);
304        }
305        assert_eq!(
306            " LLAMA.CPP ".parse::<HostKind>().unwrap(),
307            HostKind::LlamaCpp
308        );
309    }
310
311    #[test]
312    fn an_unknown_host_kind_names_the_choices() {
313        let err = "tgi".parse::<HostKind>().unwrap_err().to_string();
314
315        assert!(err.contains("\"tgi\""), "{err}");
316        assert!(err.contains("lmstudio"), "{err}");
317    }
318
319    /// The OpenAI-compatible defaults carry the `/v1` prefix; Ollama's native API has none.
320    #[test]
321    fn default_urls_match_where_each_runtime_listens() {
322        assert_eq!(HostKind::Ollama.default_url(), "http://localhost:11434");
323        assert_eq!(HostKind::Vllm.default_url(), "http://localhost:8000/v1");
324        assert_eq!(HostKind::LlamaCpp.default_url(), "http://localhost:8080/v1");
325        assert_eq!(HostKind::LmStudio.default_url(), "http://localhost:1234/v1");
326    }
327
328    #[test]
329    fn a_long_error_body_is_cut_short() {
330        let HostError::Status { body, .. } = HostError::status(400, "x".repeat(5000)) else {
331            panic!()
332        };
333
334        assert_eq!(
335            body.chars().count(),
336            MAX_ERROR_BODY + 1,
337            "300 characters and an ellipsis"
338        );
339        assert!(body.ends_with('…'));
340    }
341
342    #[test]
343    fn an_authentication_failure_passes_on_none_of_the_body() {
344        for status in [401, 403] {
345            let HostError::Status { body, .. } = HostError::status(
346                status,
347                r#"{"error":{"message":"Incorrect API key provided: sk-proj-********abcd."}}"#
348                    .into(),
349            ) else {
350                panic!()
351            };
352
353            assert!(!body.contains("sk-"), "{status}: {body}");
354            assert!(!body.contains("abcd"), "{status}: {body}");
355        }
356    }
357
358    #[test]
359    fn a_short_error_body_is_kept_whole() {
360        let HostError::Status { body, .. } = HostError::status(404, "model not found".into())
361        else {
362            panic!()
363        };
364
365        assert_eq!(body, "model not found");
366    }
367
368    #[test]
369    fn a_base_url_must_be_http_with_a_host() {
370        assert_eq!(
371            base_url(" http://localhost:11434/ ").unwrap(),
372            "http://localhost:11434"
373        );
374        assert_eq!(
375            base_url("https://api.openai.com/v1").unwrap(),
376            "https://api.openai.com/v1"
377        );
378
379        for bad in ["localhost:11434", "ftp://x", "http://", "not a url", ""] {
380            assert!(
381                matches!(base_url(bad), Err(HostError::InvalidUrl { .. })),
382                "{bad:?} was accepted"
383            );
384        }
385    }
386
387    #[test]
388    fn connect_refuses_a_url_without_a_scheme() {
389        let err = connect(
390            HostKind::Ollama,
391            "localhost:11434",
392            None,
393            Duration::from_secs(1),
394        )
395        .err()
396        .expect("refused");
397
398        assert!(err.to_string().contains("http://"), "{err}");
399    }
400
401    #[test]
402    fn connect_names_the_host_it_built() {
403        let timeout = Duration::from_secs(1);
404
405        assert_eq!(
406            connect(HostKind::Ollama, "http://x", None, timeout)
407                .unwrap()
408                .name(),
409            "ollama"
410        );
411        assert_eq!(
412            connect(HostKind::Vllm, "http://x/v1", None, timeout)
413                .unwrap()
414                .name(),
415            "vllm"
416        );
417    }
418}
cerno_host/lib.rs

cerno_host/
lib.rs
