context_harness/agent_model/
ollama.rs1use super::*;
4use anyhow::{ensure, Context};
5use reqwest::{redirect::Policy, Client, Url};
6use serde_json::{json, Map, Value};
7use std::{collections::BTreeMap, net::IpAddr, time::Duration};
8
9const DEFAULT_BASE_URL: &str = "http://127.0.0.1:11434";
10const DEFAULT_TIMEOUT_SECONDS: u64 = 120;
11const MAX_BODY_BYTES: usize = 4 * 1024 * 1024;
12
13pub struct OllamaProvider {
14 model: String,
15 endpoint: Url,
16 client: Client,
17}
18
19impl OllamaProvider {
20 pub fn from_definition(definition: &ModelDefinition) -> anyhow::Result<Self> {
21 let (base_url, timeout_seconds) = effective_config(definition)?;
22 let mut endpoint = Url::parse(&base_url)?;
23 endpoint.set_path("/api/chat");
24 Ok(Self {
25 model: definition.model.clone(),
26 endpoint,
27 client: Client::builder()
28 .timeout(Duration::from_secs(timeout_seconds))
29 .connect_timeout(Duration::from_secs(10))
30 .redirect(Policy::none())
31 .no_proxy()
32 .build()?,
33 })
34 }
35
36 fn body(&self, request: &ModelRequest) -> ModelResult<Vec<u8>> {
37 request.validate()?;
38 let mut messages = Vec::new();
39 let mut calls = BTreeMap::new();
40 for message in &request.messages {
41 match message {
42 ModelMessage::System { content } => {
43 messages.push(json!({"role":"system", "content":content}));
44 }
45 ModelMessage::User { content } => {
46 messages.push(json!({"role":"user", "content":content}));
47 }
48 ModelMessage::Assistant {
49 text, tool_calls, ..
50 } => {
51 let tool_calls: Vec<_> = tool_calls
52 .iter()
53 .map(|call| {
54 calls.insert(call.id.as_str(), wire_name(&call.name));
55 json!({"type":"function", "function":{
56 "name":wire_name(&call.name), "arguments":call.arguments
57 }})
58 })
59 .collect();
60 messages.push(json!({
61 "role":"assistant", "content":text, "tool_calls":tool_calls
62 }));
63 }
64 ModelMessage::Tool { call_id, content } => {
65 let name = calls
66 .get(call_id.as_str())
67 .ok_or_else(|| ModelError::new(ModelErrorKind::InvalidRequest))?;
68 messages.push(json!({
69 "role":"tool", "tool_name":name, "content":content
70 }));
71 }
72 }
73 }
74
75 let tools: Vec<_> = request
76 .tools
77 .iter()
78 .map(|tool| {
79 json!({"type":"function", "function":{
80 "name":wire_name(&tool.name),
81 "description":format!("{}: {}", tool.name, tool.description),
82 "parameters":tool.parameters
83 }})
84 })
85 .collect();
86 let mut body = Map::from_iter([
87 ("model".into(), json!(self.model)),
88 ("messages".into(), Value::Array(messages)),
89 ("stream".into(), Value::Bool(false)),
90 ("think".into(), Value::Bool(false)),
91 ]);
92 if !tools.is_empty() {
93 body.insert("tools".into(), Value::Array(tools));
94 }
95 if let Some(output) = &request.output_schema {
96 body.insert("format".into(), output.schema.clone());
97 }
98 if let Some(limit) = request.max_output_tokens {
99 body.insert("options".into(), json!({"num_predict":limit}));
100 }
101 let bytes = serde_json::to_vec(&body)
102 .map_err(|_| ModelError::new(ModelErrorKind::InvalidRequest))?;
103 if bytes.len() > MAX_BODY_BYTES {
104 return Err(ModelError::new(ModelErrorKind::InvalidRequest));
105 }
106 Ok(bytes)
107 }
108
109 fn decode(&self, body: Value, request: &ModelRequest) -> ModelResult<ModelResponse> {
110 let invalid = || ModelError::new(ModelErrorKind::InvalidResponse);
111 if body["done"].as_bool() != Some(true)
112 || body["model"].as_str().is_none_or(str::is_empty)
113 || body["message"]["role"].as_str() != Some("assistant")
114 {
115 return Err(invalid());
116 }
117 let text = body["message"]["content"]
118 .as_str()
119 .ok_or_else(invalid)?
120 .to_owned();
121 let names: BTreeMap<_, _> = request
122 .tools
123 .iter()
124 .map(|tool| (wire_name(&tool.name), tool.name.clone()))
125 .collect();
126 let declared_names: BTreeMap<_, _> = request
127 .tools
128 .iter()
129 .map(|tool| (tool.name.as_str(), tool.name.clone()))
130 .collect();
131 let assistant_ordinal = request
132 .messages
133 .iter()
134 .filter(|message| matches!(message, ModelMessage::Assistant { .. }))
135 .count();
136 let raw_calls = match body["message"].get("tool_calls") {
137 None | Some(Value::Null) => &[][..],
138 Some(value) => value.as_array().ok_or_else(invalid)?.as_slice(),
139 };
140 if raw_calls.is_empty() && looks_like_unstructured_tool_call(&text) {
141 return Err(invalid());
142 }
143 let mut tool_calls = Vec::with_capacity(raw_calls.len());
144 for (index, call) in raw_calls.iter().enumerate() {
145 let wire = call["function"]["name"].as_str().ok_or_else(invalid)?;
146 let name = names
151 .get(wire)
152 .or_else(|| declared_names.get(wire))
153 .ok_or_else(invalid)?
154 .clone();
155 let arguments = call["function"]["arguments"].clone();
156 if !arguments.is_object() {
157 return Err(invalid());
158 }
159 tool_calls.push(ToolCall {
160 id: format!("ollama-call-{assistant_ordinal}-{index}"),
161 name,
162 arguments,
163 });
164 }
165 let finish_reason = if !tool_calls.is_empty() {
166 FinishReason::ToolCalls
167 } else {
168 match body.get("done_reason").and_then(Value::as_str) {
169 None | Some("") | Some("stop") => FinishReason::Completed,
170 Some("length") => FinishReason::Length,
171 _ => return Err(invalid()),
172 }
173 };
174 let usage = match (
175 body.get("prompt_eval_count").and_then(Value::as_u64),
176 body.get("eval_count").and_then(Value::as_u64),
177 ) {
178 (Some(input_tokens), Some(output_tokens)) => Some(Usage {
179 input_tokens,
180 output_tokens,
181 total_tokens: input_tokens
182 .checked_add(output_tokens)
183 .ok_or_else(invalid)?,
184 }),
185 _ => None,
186 };
187 let structured_output =
188 if request.output_schema.is_some() && finish_reason == FinishReason::Completed {
189 Some(serde_json::from_str(&text).map_err(|_| invalid())?)
190 } else {
191 None
192 };
193 let response = ModelResponse {
194 text,
195 tool_calls,
196 usage,
197 finish_reason,
198 structured_output,
199 continuation: None,
200 };
201 response.validate(request)?;
202 Ok(response)
203 }
204}
205
206fn looks_like_unstructured_tool_call(text: &str) -> bool {
207 text.contains("<function=") && text.contains("<parameter=") && text.contains("</tool_call>")
208}
209
210pub(crate) fn validate_definition(definition: &ModelDefinition) -> anyhow::Result<()> {
211 effective_config(definition).map(|_| ())
212}
213
214pub(crate) fn effective_config(definition: &ModelDefinition) -> anyhow::Result<(String, u64)> {
215 definition.validate()?;
216 ensure!(
217 definition.provider == "ollama",
218 "Ollama provider requires provider = 'ollama'"
219 );
220 ensure!(
221 definition.api_key_env.is_none(),
222 "api_key_env is not supported by the ollama provider"
223 );
224 let raw = definition.base_url.as_deref().unwrap_or(DEFAULT_BASE_URL);
225 let url = Url::parse(raw).context("invalid Ollama base_url")?;
226 ensure!(url.scheme() == "http", "Ollama base_url must use http");
227 ensure!(
228 url.username().is_empty() && url.password().is_none(),
229 "Ollama base_url must not contain credentials"
230 );
231 ensure!(
232 url.path() == "/" && url.query().is_none() && url.fragment().is_none(),
233 "Ollama base_url must be an origin without a path, query, or fragment"
234 );
235 let host = url.host_str().context("Ollama base_url host is required")?;
236 let address: IpAddr = host
237 .parse()
238 .context("Ollama base_url host must be a literal loopback address")?;
239 ensure!(address.is_loopback(), "Ollama base_url must be loopback");
240 ensure!(
241 url.port_or_known_default().is_some_and(|port| port > 0),
242 "Ollama base_url port must be nonzero"
243 );
244 let timeout = definition
245 .timeout_seconds
246 .unwrap_or(DEFAULT_TIMEOUT_SECONDS);
247 ensure!(
248 (1..=3600).contains(&timeout),
249 "Ollama timeout_seconds must be in 1..=3600"
250 );
251 Ok((url.as_str().trim_end_matches('/').into(), timeout))
252}
253
254fn transport_error(error: reqwest::Error) -> ModelError {
255 ModelError::new(if error.is_timeout() {
256 ModelErrorKind::Timeout
257 } else if error.is_connect() {
258 ModelErrorKind::Unavailable
259 } else {
260 ModelErrorKind::Transport
261 })
262}
263
264async fn read_limited(mut response: reqwest::Response) -> ModelResult<Vec<u8>> {
265 let mut bytes = Vec::new();
266 while let Some(chunk) = response.chunk().await.map_err(transport_error)? {
267 if bytes.len().saturating_add(chunk.len()) > MAX_BODY_BYTES {
268 return Err(ModelError::new(ModelErrorKind::InvalidResponse));
269 }
270 bytes.extend_from_slice(&chunk);
271 }
272 Ok(bytes)
273}
274
275#[async_trait]
276impl ModelProvider for OllamaProvider {
277 async fn generate(&self, request: &ModelRequest) -> ModelResult<ModelResponse> {
278 let body = self.body(request)?;
279 let response = self
280 .client
281 .post(self.endpoint.clone())
282 .header(reqwest::header::CONTENT_TYPE, "application/json")
283 .body(body)
284 .send()
285 .await
286 .map_err(transport_error)?;
287 let status = response.status();
288 if status.is_redirection() {
289 return Err(ModelError {
290 kind: ModelErrorKind::ProviderFailure,
291 http_status: Some(status.as_u16()),
292 });
293 }
294 let bytes = read_limited(response).await?;
295 if !status.is_success() {
296 let error = serde_json::from_slice::<Value>(&bytes)
297 .ok()
298 .and_then(|body| body["error"].as_str().map(str::to_owned))
299 .ok_or_else(|| ModelError::new(ModelErrorKind::InvalidResponse))?;
300 let lower = error.to_ascii_lowercase();
301 let unsupported = lower.contains("does not support tools")
302 || (lower.contains("does not support")
303 && (lower.contains("format") || lower.contains("structured output")));
304 let kind = match status.as_u16() {
305 404 => ModelErrorKind::ModelUnavailable,
306 429 => ModelErrorKind::RateLimited,
307 500..=599 => ModelErrorKind::Unavailable,
308 400 | 422 if unsupported => ModelErrorKind::UnsupportedCapability,
309 400 | 422 => ModelErrorKind::InvalidRequest,
310 _ => ModelErrorKind::ProviderFailure,
311 };
312 return Err(ModelError {
313 kind,
314 http_status: Some(status.as_u16()),
315 });
316 }
317 let body = serde_json::from_slice(&bytes)
318 .map_err(|_| ModelError::new(ModelErrorKind::InvalidResponse))?;
319 self.decode(body, request)
320 }
321}
322
323#[cfg(test)]
324mod tests;