tokenmiser_providers/
ollama.rs

1//! Ollama client — uses the OpenAI-compatible endpoint Ollama exposes on
2//! `/v1/chat/completions`, so this is essentially the OpenAI client with no
3//! auth header and a localhost base URL.
4
5use async_trait::async_trait;
6use bytes::Bytes;
7use futures::stream::{BoxStream, StreamExt};
8use reqwest::header::{HeaderMap, HeaderValue, CONTENT_TYPE};
9use serde::Deserialize;
10use tokenmiser_config::ProviderConfig;
11
12use crate::{ChatRequest, ChatResponse, Provider, ProviderError, StreamChunk};
13
14pub struct OllamaProvider {
15    cfg: ProviderConfig,
16    client: reqwest::Client,
17}
18
19#[derive(Debug, Deserialize)]
20pub struct OllamaTag {
21    pub name: String,
22}
23
24#[derive(Debug, Deserialize)]
25struct OllamaTagsResponse {
26    models: Vec<OllamaTag>,
27}
28
29impl OllamaProvider {
30    pub fn new(cfg: ProviderConfig) -> Self {
31        let client = reqwest::Client::builder()
32            .pool_max_idle_per_host(32)
33            .build()
34            .expect("reqwest client construction");
35        Self { cfg, client }
36    }
37
38    /// Probe localhost:11434 for a running Ollama. Returns the loaded model
39    /// names if reachable; used by the daemon at startup for zero-config
40    /// local routing (architecture §7).
41    pub async fn detect(base_url: &str) -> Result<Vec<String>, ProviderError> {
42        let client = reqwest::Client::builder()
43            .timeout(std::time::Duration::from_millis(500))
44            .build()?;
45        let url = format!("{}/api/tags", base_url.trim_end_matches('/'));
46        let res = client.get(&url).send().await?;
47        if !res.status().is_success() {
48            return Err(ProviderError::Upstream {
49                status: res.status().as_u16(),
50                body: res.text().await.unwrap_or_default(),
51            });
52        }
53        let body: OllamaTagsResponse = res.json().await?;
54        Ok(body.models.into_iter().map(|m| m.name).collect())
55    }
56}
57
58#[async_trait]
59impl Provider for OllamaProvider {
60    fn name(&self) -> &str {
61        &self.cfg.name
62    }
63
64    fn config(&self) -> &ProviderConfig {
65        &self.cfg
66    }
67
68    async fn complete(&self, req: &ChatRequest) -> Result<ChatResponse, ProviderError> {
69        let url = format!(
70            "{}/v1/chat/completions",
71            self.cfg.base_url.trim_end_matches('/')
72        );
73        let mut headers = HeaderMap::new();
74        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
75
76        let mut body = req.clone();
77        body.stream = Some(false);
78
79        // Ollama prefixes-tolerant: strip `ollama:` if present.
80        if let Some(rest) = body.model.strip_prefix("ollama:") {
81            body.model = rest.to_string();
82        }
83
84        let res = self
85            .client
86            .post(&url)
87            .headers(headers)
88            .json(&body)
89            .send()
90            .await?;
91        let status = res.status();
92        let text = res.text().await?;
93
94        if !status.is_success() {
95            return Err(ProviderError::Upstream {
96                status: status.as_u16(),
97                body: text,
98            });
99        }
100
101        let parsed: ChatResponse = serde_json::from_str(&text)?;
102        Ok(parsed)
103    }
104
105    async fn stream(
106        &self,
107        req: &ChatRequest,
108    ) -> Result<BoxStream<'static, Result<StreamChunk, ProviderError>>, ProviderError> {
109        let url = format!(
110            "{}/v1/chat/completions",
111            self.cfg.base_url.trim_end_matches('/')
112        );
113        let mut headers = HeaderMap::new();
114        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
115
116        let mut body = req.clone();
117        body.stream = Some(true);
118        if let Some(rest) = body.model.strip_prefix("ollama:") {
119            body.model = rest.to_string();
120        }
121
122        let res = self
123            .client
124            .post(&url)
125            .headers(headers)
126            .json(&body)
127            .send()
128            .await?;
129        let status = res.status();
130        if !status.is_success() {
131            let text = res.text().await?;
132            return Err(ProviderError::Upstream {
133                status: status.as_u16(),
134                body: text,
135            });
136        }
137
138        let stream = res
139            .bytes_stream()
140            .map(|chunk_res| match chunk_res {
141                Ok(b) => Ok(StreamChunk::Sse(Bytes::from(b.to_vec()))),
142                Err(e) => Err(ProviderError::Http(e)),
143            })
144            .chain(futures::stream::once(async { Ok(StreamChunk::Done) }));
145
146        Ok(stream.boxed())
147    }
148}
tokenmiser_providers/ollama.rs

tokenmiser_providers/
ollama.rs
