zag_agent/providers/
ollama.rs

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

zag_agent/providers/
ollama.rs
