pub const DEFAULT_ENDPOINT: &str = "http://localhost:11434";
pub const DEFAULT_MODEL: &str = "qwen2.5";
#[derive(Debug, Clone)]
pub struct Ollama {
endpoint: String,
model: String,
client: reqwest::Client,
}
impl Ollama {
pub fn new(endpoint: impl Into<String>, model: impl Into<String>) -> Self {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(120))
.connect_timeout(std::time::Duration::from_secs(5))
.build()
.expect("failed to build HTTP client");
Self {
endpoint: endpoint.into(),
model: model.into(),
client,
}
}
pub async fn generate(&self, prompt: &str) -> crate::Result<String> {
let url = format!("{}/api/generate", self.endpoint);
let body = serde_json::json!({
"model": self.model,
"prompt": prompt,
"stream": false,
"options": {
"temperature": 0.0,
"num_predict": 512,
},
});
let response = self
.client
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| crate::Error::OllamaUnreachable(format!("{}: {e}", self.endpoint)))?;
let status = response.status();
let body = response
.text()
.await
.map_err(|e| crate::Error::Parse(e.to_string()))?;
if !status.is_success() {
let reason = serde_json::from_str::<serde_json::Value>(&body)
.ok()
.and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from))
.unwrap_or_else(|| format!("Ollama returned HTTP {}", status.as_u16()));
return Err(crate::Error::OllamaModel(format!(
"{reason} (model `{}`) — run `ollama pull {}`",
self.model, self.model
)));
}
let json: serde_json::Value =
serde_json::from_str(&body).map_err(|e| crate::Error::Parse(e.to_string()))?;
if let Some(err) = json.get("error").and_then(|e| e.as_str()) {
return Err(crate::Error::OllamaModel(format!(
"{err} (model `{}`) — run `ollama pull {}`",
self.model, self.model
)));
}
json["response"]
.as_str()
.map(String::from)
.ok_or_else(|| crate::Error::Parse("missing 'response' field".into()))
}
}
impl Default for Ollama {
fn default() -> Self {
Self::new(DEFAULT_ENDPOINT, DEFAULT_MODEL)
}
}