zag_agent/providers/
ollama.rs

1// provider-updated: 2026-04-05
2use crate::agent::{Agent, ModelSize};
3use crate::output::AgentOutput;
4use crate::sandbox::SandboxConfig;
5use crate::session_log::HistoricalLogAdapter;
6use anyhow::{Context, Result};
7use async_trait::async_trait;
8use std::process::Stdio;
9use tokio::process::Command;
10
11pub const DEFAULT_MODEL: &str = "qwen3.5";
12pub const DEFAULT_SIZE: &str = "9b";
13
14pub const AVAILABLE_SIZES: &[&str] = &["0.8b", "2b", "4b", "9b", "27b", "35b", "122b"];
15
16pub struct Ollama {
17    system_prompt: String,
18    model: String,
19    size: String,
20    root: Option<String>,
21    skip_permissions: bool,
22    output_format: Option<String>,
23    add_dirs: Vec<String>,
24    capture_output: bool,
25    max_turns: Option<u32>,
26    sandbox: Option<SandboxConfig>,
27    env_vars: Vec<(String, String)>,
28}
29
30pub struct OllamaHistoricalLogAdapter;
31
32impl Ollama {
33    pub fn new() -> Self {
34        Self {
35            system_prompt: String::new(),
36            model: DEFAULT_MODEL.to_string(),
37            size: DEFAULT_SIZE.to_string(),
38            root: None,
39            skip_permissions: false,
40            output_format: None,
41            add_dirs: Vec::new(),
42            capture_output: false,
43            max_turns: None,
44            sandbox: None,
45            env_vars: Vec::new(),
46        }
47    }
48
49    pub fn set_size(&mut self, size: String) {
50        self.size = size;
51    }
52
53    /// Get the display string for the model (e.g., "qwen3.5:9b").
54    pub fn display_model(&self) -> String {
55        self.model_tag()
56    }
57
58    /// Get the full model tag (e.g., "qwen3.5:9b").
59    fn model_tag(&self) -> String {
60        format!("{}:{}", self.model, self.size)
61    }
62
63    /// Build the argument list for a run invocation.
64    fn build_run_args(&self, interactive: bool, prompt: Option<&str>) -> Vec<String> {
65        let mut args = vec!["run".to_string()];
66
67        if let Some(ref format) = self.output_format
68            && format == "json"
69        {
70            args.extend(["--format".to_string(), "json".to_string()]);
71        }
72
73        if !interactive {
74            // --nowordwrap for clean piped output
75            args.push("--nowordwrap".to_string());
76        }
77
78        args.push("--hidethinking".to_string());
79
80        args.push(self.model_tag());
81
82        // ollama run has no --system flag; prepend system prompt to user prompt
83        let effective_prompt = match (self.system_prompt.is_empty(), prompt) {
84            (false, Some(p)) => Some(format!("{}\n\n{}", self.system_prompt, p)),
85            (false, None) => Some(self.system_prompt.clone()),
86            (true, p) => p.map(String::from),
87        };
88
89        if let Some(p) = effective_prompt {
90            args.push(p);
91        }
92
93        args
94    }
95
96    /// Create a `Command` either directly or wrapped in sandbox.
97    fn make_command(&self, agent_args: Vec<String>) -> Command {
98        if let Some(ref sb) = self.sandbox {
99            // For ollama in sandbox, we use the shell template:
100            // docker sandbox run shell <workspace> -- -c "ollama run ..."
101            let shell_cmd = format!(
102                "ollama {}",
103                agent_args
104                    .iter()
105                    .map(|a| shell_escape(a))
106                    .collect::<Vec<_>>()
107                    .join(" ")
108            );
109            let mut std_cmd = std::process::Command::new("docker");
110            std_cmd.args([
111                "sandbox",
112                "run",
113                "--name",
114                &sb.name,
115                &sb.template,
116                &sb.workspace,
117                "--",
118                "-c",
119                &shell_cmd,
120            ]);
121            log::debug!(
122                "Sandbox command: docker sandbox run --name {} {} {} -- -c {:?}",
123                sb.name,
124                sb.template,
125                sb.workspace,
126                shell_cmd
127            );
128            Command::from(std_cmd)
129        } else {
130            let mut cmd = Command::new("ollama");
131            if let Some(ref root) = self.root {
132                cmd.current_dir(root);
133            }
134            cmd.args(&agent_args);
135            for (key, value) in &self.env_vars {
136                cmd.env(key, value);
137            }
138            cmd
139        }
140    }
141
142    async fn execute(
143        &self,
144        interactive: bool,
145        prompt: Option<&str>,
146    ) -> Result<Option<AgentOutput>> {
147        let agent_args = self.build_run_args(interactive, prompt);
148        log::debug!("Ollama command: ollama {}", agent_args.join(" "));
149        if !self.system_prompt.is_empty() {
150            log::debug!("Ollama system prompt: {}", self.system_prompt);
151        }
152        if let Some(p) = prompt {
153            log::debug!("Ollama user prompt: {}", p);
154        }
155        let mut cmd = self.make_command(agent_args);
156
157        if interactive {
158            cmd.stdin(Stdio::inherit())
159                .stdout(Stdio::inherit())
160                .stderr(Stdio::inherit());
161            let status = cmd
162                .status()
163                .await
164                .context("Failed to execute 'ollama' CLI. Is it installed and in PATH?")?;
165            if !status.success() {
166                return Err(crate::process::ProcessError {
167                    exit_code: status.code(),
168                    stderr: String::new(),
169                    agent_name: "Ollama".to_string(),
170                }
171                .into());
172            }
173            Ok(None)
174        } else if self.capture_output {
175            let text = crate::process::run_captured(&mut cmd, "Ollama").await?;
176            log::debug!("Ollama raw response ({} bytes): {}", text.len(), text);
177            Ok(Some(AgentOutput::from_text("ollama", &text)))
178        } else {
179            cmd.stdin(Stdio::inherit()).stdout(Stdio::inherit());
180            crate::process::run_with_captured_stderr(&mut cmd).await?;
181            Ok(None)
182        }
183    }
184
185    /// Resolve a size alias to the appropriate parameter size.
186    pub fn size_for_model_size(size: ModelSize) -> &'static str {
187        match size {
188            ModelSize::Small => "2b",
189            ModelSize::Medium => "9b",
190            ModelSize::Large => "35b",
191        }
192    }
193}
194
195/// Escape a string for shell use. Wraps in single quotes if it contains special chars.
196fn shell_escape(s: &str) -> String {
197    if s.contains(' ')
198        || s.contains('\'')
199        || s.contains('"')
200        || s.contains('\\')
201        || s.contains('$')
202        || s.contains('`')
203        || s.contains('!')
204    {
205        format!("'{}'", s.replace('\'', "'\\''"))
206    } else {
207        s.to_string()
208    }
209}
210
211#[cfg(test)]
212#[path = "ollama_tests.rs"]
213mod tests;
214
215impl Default for Ollama {
216    fn default() -> Self {
217        Self::new()
218    }
219}
220
221impl HistoricalLogAdapter for OllamaHistoricalLogAdapter {
222    fn backfill(&self, _root: Option<&str>) -> Result<Vec<crate::session_log::BackfilledSession>> {
223        Ok(Vec::new())
224    }
225}
226
227#[async_trait]
228impl Agent for Ollama {
229    fn name(&self) -> &str {
230        "ollama"
231    }
232
233    fn default_model() -> &'static str
234    where
235        Self: Sized,
236    {
237        DEFAULT_MODEL
238    }
239
240    fn model_for_size(size: ModelSize) -> &'static str
241    where
242        Self: Sized,
243    {
244        // For ollama, model_for_size returns the size parameter, not the model name
245        Self::size_for_model_size(size)
246    }
247
248    fn available_models() -> &'static [&'static str]
249    where
250        Self: Sized,
251    {
252        // Ollama accepts any model — return common sizes for validation/help
253        AVAILABLE_SIZES
254    }
255
256    /// Ollama uses open model names — skip strict validation.
257    fn validate_model(_model: &str, _agent_name: &str) -> Result<()>
258    where
259        Self: Sized,
260    {
261        Ok(())
262    }
263
264    fn system_prompt(&self) -> &str {
265        &self.system_prompt
266    }
267
268    fn set_system_prompt(&mut self, prompt: String) {
269        self.system_prompt = prompt;
270    }
271
272    fn get_model(&self) -> &str {
273        &self.model
274    }
275
276    fn set_model(&mut self, model: String) {
277        self.model = model;
278    }
279
280    fn set_root(&mut self, root: String) {
281        self.root = Some(root);
282    }
283
284    fn set_skip_permissions(&mut self, _skip: bool) {
285        // Ollama runs locally — no permission concept
286        self.skip_permissions = true;
287    }
288
289    fn set_output_format(&mut self, format: Option<String>) {
290        self.output_format = format;
291    }
292
293    fn set_capture_output(&mut self, capture: bool) {
294        self.capture_output = capture;
295    }
296
297    fn set_max_turns(&mut self, turns: u32) {
298        self.max_turns = Some(turns);
299    }
300
301    fn set_sandbox(&mut self, config: SandboxConfig) {
302        self.sandbox = Some(config);
303    }
304
305    fn set_add_dirs(&mut self, dirs: Vec<String>) {
306        self.add_dirs = dirs;
307    }
308
309    fn set_env_vars(&mut self, vars: Vec<(String, String)>) {
310        self.env_vars = vars;
311    }
312
313    fn as_any_ref(&self) -> &dyn std::any::Any {
314        self
315    }
316
317    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
318        self
319    }
320
321    async fn run(&self, prompt: Option<&str>) -> Result<Option<AgentOutput>> {
322        self.execute(false, prompt).await
323    }
324
325    async fn run_interactive(&self, prompt: Option<&str>) -> Result<()> {
326        self.execute(true, prompt).await?;
327        Ok(())
328    }
329
330    async fn run_resume(&self, _session_id: Option<&str>, _last: bool) -> Result<()> {
331        anyhow::bail!("Ollama does not support session resume")
332    }
333
334    async fn cleanup(&self) -> Result<()> {
335        Ok(())
336    }
337}
zag_agent/providers/ollama.rs

zag_agent/providers/
ollama.rs
