Skip to main content

jamjet_models/
ollama.rs

1//! Ollama adapter — local model inference via Ollama's HTTP API.
2//!
3//! Supports any model available via `ollama pull`: qwen3, llama3, gemma2, phi3, etc.
4//! Reads `OLLAMA_HOST` from the environment (defaults to http://localhost:11434).
5//! Uses Ollama's native /api/chat endpoint for accurate token counts.
6
7use crate::adapter::{
8    warn_tools_not_forwarded, ChatMessage, ChatRole, ModelAdapter, ModelConfig, ModelError,
9    ModelRequest, ModelResponse, StructuredRequest,
10};
11use async_trait::async_trait;
12use serde_json::{json, Value};
13use tracing::{debug, instrument};
14
15const OLLAMA_DEFAULT_HOST: &str = "http://localhost:11434";
16const DEFAULT_MODEL: &str = "llama3.2:3b";
17const DEFAULT_MAX_TOKENS: u32 = 4096;
18
19/// Ollama adapter for local model inference.
20///
21/// Connects to a running Ollama server and uses its native /api/chat endpoint.
22/// All inference is free (local GPU/CPU), making this ideal for development,
23/// testing, and cost-sensitive workloads.
24pub struct OllamaAdapter {
25    client: reqwest::Client,
26    host: String,
27    default_model: String,
28}
29
30impl OllamaAdapter {
31    pub fn new(host: impl Into<String>) -> Self {
32        Self {
33            client: reqwest::Client::new(),
34            host: host.into(),
35            default_model: DEFAULT_MODEL.into(),
36        }
37    }
38
39    /// Create adapter from `OLLAMA_HOST` env var (defaults to localhost:11434).
40    pub fn from_env() -> Result<Self, ModelError> {
41        let host = std::env::var("OLLAMA_HOST").unwrap_or_else(|_| OLLAMA_DEFAULT_HOST.to_string());
42
43        // Quick check: if Ollama is not reachable, fail fast.
44        // We skip the actual health check here to keep construction sync;
45        // errors will surface on first call_api() instead.
46        Ok(Self::new(host))
47    }
48
49    pub fn with_default_model(mut self, model: impl Into<String>) -> Self {
50        self.default_model = model.into();
51        self
52    }
53
54    async fn call_api(&self, body: Value) -> Result<Value, ModelError> {
55        let resp = self
56            .client
57            .post(format!("{}/api/chat", self.host))
58            .json(&body)
59            .send()
60            .await
61            .map_err(|e| ModelError::Network(format!("Ollama unreachable: {e}")))?;
62
63        let status = resp.status().as_u16();
64        let body_text = resp
65            .text()
66            .await
67            .map_err(|e| ModelError::Network(e.to_string()))?;
68
69        if status != 200 {
70            return Err(ModelError::Api {
71                status,
72                body: body_text,
73            });
74        }
75
76        serde_json::from_str(&body_text).map_err(|e| ModelError::Serialization(e.to_string()))
77    }
78
79    fn build_request_body(
80        &self,
81        messages: &[ChatMessage],
82        config: &ModelConfig,
83        format: Option<&str>,
84    ) -> Value {
85        let model = config.model.as_deref().unwrap_or(&self.default_model);
86        let max_tokens = config.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS);
87
88        let mut ollama_messages: Vec<Value> = Vec::new();
89
90        // Prepend system prompt if provided in config.
91        if let Some(sys) = &config.system_prompt {
92            ollama_messages.push(json!({ "role": "system", "content": sys }));
93        }
94
95        for m in messages {
96            let role = match m.role {
97                ChatRole::System => "system",
98                ChatRole::User => "user",
99                ChatRole::Assistant => "assistant",
100                ChatRole::Tool => "tool",
101            };
102            ollama_messages.push(json!({ "role": role, "content": m.content }));
103        }
104
105        let mut body = json!({
106            "model": model,
107            "messages": ollama_messages,
108            "stream": false,
109            "options": {
110                "num_predict": max_tokens,
111            },
112        });
113
114        if let Some(temp) = config.temperature {
115            body["options"]["temperature"] = json!(temp);
116        }
117        if let Some(stops) = &config.stop_sequences {
118            body["options"]["stop"] = json!(stops);
119        }
120        if let Some(fmt) = format {
121            body["format"] = json!(fmt);
122        }
123
124        body
125    }
126
127    fn parse_response(&self, resp: Value) -> Result<ModelResponse, ModelError> {
128        let model = resp["model"]
129            .as_str()
130            .unwrap_or(&self.default_model)
131            .to_string();
132
133        let content = resp["message"]["content"]
134            .as_str()
135            .unwrap_or("")
136            .to_string();
137
138        // Ollama provides token counts in prompt_eval_count / eval_count.
139        let input_tokens = resp["prompt_eval_count"].as_u64().unwrap_or(0);
140        let output_tokens = resp["eval_count"].as_u64().unwrap_or(0);
141
142        // Ollama uses "done_reason" (not "finish_reason").
143        let finish_reason = resp["done_reason"].as_str().unwrap_or("stop").to_string();
144
145        Ok(ModelResponse {
146            content,
147            model,
148            finish_reason,
149            input_tokens,
150            output_tokens,
151            structured: None,
152            tool_calls: vec![],
153        })
154    }
155}
156
157#[async_trait]
158impl ModelAdapter for OllamaAdapter {
159    fn system_name(&self) -> &'static str {
160        "ollama"
161    }
162
163    fn default_model(&self) -> &str {
164        &self.default_model
165    }
166
167    #[instrument(skip(self, request), fields(
168        gen_ai.system = "ollama",
169        gen_ai.request.model = tracing::field::Empty,
170        gen_ai.usage.input_tokens = tracing::field::Empty,
171        gen_ai.usage.output_tokens = tracing::field::Empty,
172    ))]
173    async fn chat(&self, request: ModelRequest) -> Result<ModelResponse, ModelError> {
174        // Native adapter: tools are not forwarded to the provider. Warn (once) so
175        // a tool-carrying call does not silently degenerate the agent loop.
176        if !request.tools.is_empty() {
177            warn_tools_not_forwarded(self.system_name());
178        }
179
180        let model = request
181            .config
182            .model
183            .as_deref()
184            .unwrap_or(&self.default_model)
185            .to_string();
186        tracing::Span::current().record("gen_ai.request.model", model.as_str());
187
188        debug!(model = %model, host = %self.host, "Calling Ollama /api/chat");
189
190        let body = self.build_request_body(&request.messages, &request.config, None);
191        let resp_json = self.call_api(body).await?;
192        let response = self.parse_response(resp_json)?;
193
194        tracing::Span::current()
195            .record("gen_ai.usage.input_tokens", response.input_tokens)
196            .record("gen_ai.usage.output_tokens", response.output_tokens);
197
198        Ok(response)
199    }
200
201    #[instrument(skip(self, request), fields(
202        gen_ai.system = "ollama",
203        gen_ai.request.model = tracing::field::Empty,
204    ))]
205    async fn structured_output(
206        &self,
207        request: StructuredRequest,
208    ) -> Result<ModelResponse, ModelError> {
209        let model = request
210            .config
211            .model
212            .as_deref()
213            .unwrap_or(&self.default_model)
214            .to_string();
215        tracing::Span::current().record("gen_ai.request.model", model.as_str());
216
217        // Ollama supports format: "json" for JSON mode.
218        // Append the schema to the system prompt so the model knows the structure.
219        let mut config = request.config.clone();
220        let schema_str = serde_json::to_string_pretty(&request.output_schema)
221            .map_err(|e| ModelError::Serialization(e.to_string()))?;
222        let system = config.system_prompt.get_or_insert_with(String::new);
223        system.push_str(&format!(
224            "\n\nRespond ONLY with a valid JSON object matching this schema:\n{schema_str}"
225        ));
226
227        let body = self.build_request_body(&request.messages, &config, Some("json"));
228        let resp_json = self.call_api(body).await?;
229        let mut response = self.parse_response(resp_json)?;
230
231        // Parse JSON from response content.
232        let structured =
233            serde_json::from_str::<serde_json::Value>(&response.content).map_err(|e| {
234                ModelError::Serialization(format!("structured output parse error: {e}"))
235            })?;
236        response.structured = Some(structured);
237
238        Ok(response)
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245
246    #[test]
247    fn test_build_request_body() {
248        let adapter = OllamaAdapter::new("http://localhost:11434");
249        let messages = vec![ChatMessage::user("Hello")];
250        let config = ModelConfig {
251            model: Some("qwen3:8b".into()),
252            max_tokens: Some(100),
253            temperature: Some(0.7),
254            ..Default::default()
255        };
256        let body = adapter.build_request_body(&messages, &config, None);
257
258        assert_eq!(body["model"], "qwen3:8b");
259        assert_eq!(body["stream"], false);
260        assert_eq!(body["options"]["num_predict"], 100);
261        let temp = body["options"]["temperature"].as_f64().unwrap();
262        assert!((temp - 0.7).abs() < 0.01);
263    }
264
265    #[test]
266    fn test_parse_response() {
267        let adapter = OllamaAdapter::new("http://localhost:11434");
268        let resp = json!({
269            "model": "qwen3:8b",
270            "message": {"role": "assistant", "content": "Hello!"},
271            "done": true,
272            "done_reason": "stop",
273            "prompt_eval_count": 42,
274            "eval_count": 5,
275        });
276
277        let parsed = adapter.parse_response(resp).unwrap();
278        assert_eq!(parsed.content, "Hello!");
279        assert_eq!(parsed.model, "qwen3:8b");
280        assert_eq!(parsed.input_tokens, 42);
281        assert_eq!(parsed.output_tokens, 5);
282        assert_eq!(parsed.finish_reason, "stop");
283    }
284}