zag_agent/providers/
ollama.rs

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

zag_agent/providers/
ollama.rs
