1use 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
19pub 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 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 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 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 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 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 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 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 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}