use std::io::Write;
use std::process::{Command, Stdio};
use crate::{AiBackend, AiError, GenerateRequest, GenerateResponse};
#[derive(Debug, Clone)]
pub struct Ollama {
model: String,
bin: String,
}
impl Ollama {
pub fn new(model: impl Into<String>) -> Self {
Self {
model: model.into(),
bin: "ollama".into(),
}
}
pub fn with_model(mut self, model: impl Into<String>) -> Self {
self.model = model.into();
self
}
pub fn with_bin(mut self, bin: impl Into<String>) -> Self {
self.bin = bin.into();
self
}
pub fn model(&self) -> &str {
&self.model
}
pub fn bin(&self) -> &str {
&self.bin
}
}
impl AiBackend for Ollama {
fn name(&self) -> &'static str {
"ollama"
}
fn generate(&self, request: &GenerateRequest) -> Result<GenerateResponse, AiError> {
let full_prompt = match &request.instructions {
Some(instructions) if !instructions.trim().is_empty() => {
format!(
"Instructions:\n{}\n\nPrompt:\n{}",
instructions, request.prompt
)
}
_ => request.prompt.clone(),
};
let mut child = Command::new(&self.bin)
.args(["run", &self.model])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.spawn()?;
if let Some(stdin) = child.stdin.as_mut() {
stdin.write_all(full_prompt.as_bytes())?;
}
let output = child.wait_with_output()?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
return Err(AiError::CommandFailed(format!(
"{} exited with status {}{}",
self.bin,
output.status,
if stderr.is_empty() {
String::new()
} else {
format!(": {}", stderr)
}
)));
}
Ok(GenerateResponse {
text: String::from_utf8(output.stdout)?.trim().to_string(),
})
}
}