Skip to main content

context_harness/agent_model/
ollama.rs

1//! Native, non-streaming Ollama chat adapter with a loopback-only transport.
2
3use 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            // Some local models return the canonical name from the tool
147            // description instead of the hashed name supplied in the
148            // function schema. Accept it only when it exactly matches a tool
149            // declared for this request.
150            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;