cerno_host/
router.rs

1//! Ollama serves two endpoints cerno can ask, and [`OllamaRouter`] picks one per model.
2//!
3//! A model Ollama lists with the `decision` capability (`nimble`, `tev1`) is asked through
4//! `/v1/systemone`, where Ollama writes the prompt and answers with a probability per offered
5//! option. Every other model gets cerno's letter prompt over `/api/chat`. Detection is one
6//! `/api/show` per model, cached for the life of the process.
7
8use crate::{
9    FirstTokenDistribution, FirstTokenRequest, HostCapabilities, HostError, ModelHost, OllamaHost,
10    SystemOneHost,
11};
12use async_trait::async_trait;
13use serde::Deserialize;
14use std::collections::HashMap;
15use std::sync::Mutex;
16use std::time::Duration;
17
18#[derive(Deserialize)]
19struct Show {
20    /// Absent on Ollama versions that predate the field; those have no decision models either.
21    #[serde(default)]
22    capabilities: Vec<String>,
23}
24
25pub struct OllamaRouter {
26    client: reqwest::Client,
27    base_url: String,
28    timeout: Duration,
29    chat: OllamaHost,
30    systemone: SystemOneHost,
31    /// Per model: whether Ollama lists it with the `decision` capability. Asked once per model.
32    decision: Mutex<HashMap<String, bool>>,
33}
34
35impl OllamaRouter {
36    pub fn new(base_url: impl Into<String>, timeout: Duration) -> Result<Self, HostError> {
37        let url = crate::base_url(&base_url.into())?;
38        let client = reqwest::Client::builder()
39            .timeout(timeout)
40            .build()
41            .map_err(|e| HostError::Unavailable(e.to_string()))?;
42        Ok(Self {
43            client,
44            chat: OllamaHost::new(url.clone(), timeout)?,
45            systemone: SystemOneHost::new(url.clone(), timeout)?,
46            base_url: url,
47            timeout,
48            decision: Mutex::default(),
49        })
50    }
51
52    /// Whether `model` is a decision model.
53    ///
54    /// A refused `/api/show` is answered "no" and not remembered: `/api/chat` then fails the
55    /// same way it always has, so an unknown model is still reported as one, and a model pulled
56    /// later is recognised on its next question.
57    async fn is_decision(&self, model: &str) -> Result<bool, HostError> {
58        let known = self
59            .decision
60            .lock()
61            .expect("lock is never poisoned")
62            .get(model)
63            .copied();
64        if let Some(known) = known {
65            return Ok(known);
66        }
67
68        let response = self
69            .client
70            .post(format!("{}/api/show", self.base_url))
71            .json(&serde_json::json!({ "model": model }))
72            .send()
73            .await
74            .map_err(|e| HostError::transport(&e, self.timeout))?;
75
76        let status = response.status().as_u16();
77        if status >= 400 {
78            tracing::debug!(
79                model,
80                status,
81                "could not read the model's capabilities; asking /api/chat"
82            );
83            return Ok(false);
84        }
85
86        let text = response
87            .text()
88            .await
89            .map_err(|e| HostError::transport(&e, self.timeout))?;
90        let show: Show = serde_json::from_str(&text)
91            .map_err(|e| HostError::Protocol(format!("could not read /api/show: {e}")))?;
92        let decision = show.capabilities.iter().any(|c| c == "decision");
93
94        // Concurrent first questions may each have asked; only the first to finish says so.
95        let first = self
96            .decision
97            .lock()
98            .expect("lock is never poisoned")
99            .insert(model.to_string(), decision)
100            .is_none();
101        if decision && first {
102            tracing::info!(model, "decision model; answering through /v1/systemone");
103        }
104        Ok(decision)
105    }
106}
107
108#[async_trait]
109impl ModelHost for OllamaRouter {
110    /// The `/api/chat` ceiling. `/v1/systemone` takes more options, but `Engine` caps a choice at
111    /// `MAX_OPTIONS` either way, so the smaller one costs nothing.
112    fn capabilities(&self) -> HostCapabilities {
113        self.chat.capabilities()
114    }
115
116    fn name(&self) -> &str {
117        "ollama"
118    }
119
120    async fn first_token(
121        &self,
122        req: FirstTokenRequest,
123    ) -> Result<FirstTokenDistribution, HostError> {
124        if self.is_decision(&req.model).await? {
125            self.systemone.first_token(req).await
126        } else {
127            self.chat.first_token(req).await
128        }
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use mockito::Matcher;
136    use serde_json::json;
137
138    fn request(model: &str) -> FirstTokenRequest {
139        FirstTokenRequest {
140            model: model.into(),
141            system: None,
142            user: "Which?".into(),
143            top_logprobs: 20,
144            keep_alive: None,
145            offer: crate::test_offer(),
146        }
147    }
148
149    async fn show(
150        server: &mut mockito::Server,
151        model: &str,
152        capabilities: &[&str],
153    ) -> mockito::Mock {
154        server
155            .mock("POST", "/api/show")
156            .match_body(Matcher::PartialJson(json!({ "model": model })))
157            .with_body(json!({ "capabilities": capabilities }).to_string())
158            .expect(1)
159            .create_async()
160            .await
161    }
162
163    #[tokio::test]
164    async fn a_decision_model_is_answered_through_systemone_and_asked_once() {
165        let mut server = mockito::Server::new_async().await;
166        let show = show(&mut server, "nimble", &["decision"]).await;
167        let systemone = server
168            .mock("POST", "/v1/systemone")
169            .with_body(
170                json!({"answers": {"q": {"type": "choice", "probabilities": {"x": 0.8, "y": 0.2}}}})
171                    .to_string(),
172            )
173            .expect(2)
174            .create_async()
175            .await;
176        let chat = server
177            .mock("POST", "/api/chat")
178            .expect(0)
179            .create_async()
180            .await;
181        let host = OllamaRouter::new(server.url(), Duration::from_secs(5)).unwrap();
182
183        for _ in 0..2 {
184            let dist = host.first_token(request("nimble")).await.unwrap();
185            assert_eq!(dist.tokens[0].0, "A");
186            assert!((dist.logprob("A").unwrap().exp() - 0.8).abs() < 1e-9);
187        }
188
189        show.assert_async().await;
190        systemone.assert_async().await;
191        chat.assert_async().await;
192    }
193
194    #[tokio::test]
195    async fn any_other_model_gets_the_letter_prompt() {
196        let mut server = mockito::Server::new_async().await;
197        let show = show(&mut server, "gemma4", &["completion", "tools"]).await;
198        let chat = server
199            .mock("POST", "/api/chat")
200            .with_body(
201                json!({
202                    "message": {"role": "assistant", "content": "B"},
203                    "logprobs": [{"token": "B", "logprob": -0.1, "top_logprobs": [
204                        {"token": "B", "logprob": -0.1},
205                        {"token": "A", "logprob": -2.4}
206                    ]}]
207                })
208                .to_string(),
209            )
210            .expect(1)
211            .create_async()
212            .await;
213        let systemone = server
214            .mock("POST", "/v1/systemone")
215            .expect(0)
216            .create_async()
217            .await;
218        let host = OllamaRouter::new(server.url(), Duration::from_secs(5)).unwrap();
219
220        let dist = host.first_token(request("gemma4")).await.unwrap();
221
222        assert_eq!(dist.tokens[0], ("B".to_string(), -0.1));
223        show.assert_async().await;
224        chat.assert_async().await;
225        systemone.assert_async().await;
226    }
227
228    #[tokio::test]
229    async fn an_unreadable_capability_list_falls_back_to_chat_and_is_asked_again() {
230        let mut server = mockito::Server::new_async().await;
231        let show = server
232            .mock("POST", "/api/show")
233            .with_status(404)
234            .with_body(r#"{"error":"model 'x:latest' not found"}"#)
235            .expect(2)
236            .create_async()
237            .await;
238        let chat = server
239            .mock("POST", "/api/chat")
240            .with_status(404)
241            .with_body(r#"{"error":"model \"x\" not found"}"#)
242            .expect(2)
243            .create_async()
244            .await;
245        let host = OllamaRouter::new(server.url(), Duration::from_secs(5)).unwrap();
246
247        for _ in 0..2 {
248            let err = host.first_token(request("x")).await.unwrap_err();
249            assert!(
250                matches!(err, HostError::Status { status: 404, ref body } if body.contains("not found")),
251                "{err}"
252            );
253        }
254
255        show.assert_async().await;
256        chat.assert_async().await;
257    }
258}
cerno_host/router.rs

cerno_host/
router.rs
