77 lines
2.4 KiB
Rust
77 lines
2.4 KiB
Rust
use async_trait::async_trait;
|
|
use serde_json::{Value, json};
|
|
|
|
use crate::{ChatOptions, ChatResponse, ChatbotClient, Message, Role};
|
|
|
|
/// Ollama client using the native `/api/chat` endpoint.
|
|
///
|
|
/// Defaults to `http://localhost:11434`. No API key required.
|
|
pub struct OllamaClient {
|
|
base_url: String,
|
|
http: reqwest::Client,
|
|
}
|
|
|
|
impl OllamaClient {
|
|
/// `base_url` defaults to `http://localhost:11434` if `None`.
|
|
pub fn new(base_url: Option<impl Into<String>>) -> Self {
|
|
let url = base_url
|
|
.map(|u| u.into())
|
|
.unwrap_or_else(|| "http://localhost:11434".to_string());
|
|
Self { base_url: url, http: reqwest::Client::new() }
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl ChatbotClient for OllamaClient {
|
|
async fn chat(
|
|
&self,
|
|
messages: &[Message],
|
|
options: &ChatOptions,
|
|
) -> anyhow::Result<ChatResponse> {
|
|
let msgs: Vec<Value> = messages
|
|
.iter()
|
|
.map(|m| {
|
|
let role = match m.role {
|
|
Role::System => "system",
|
|
Role::User => "user",
|
|
Role::Assistant => "assistant",
|
|
};
|
|
json!({ "role": role, "content": m.content })
|
|
})
|
|
.collect();
|
|
|
|
let mut options_obj = json!({});
|
|
if let Some(t) = options.temperature { options_obj["temperature"] = t.into(); }
|
|
if let Some(n) = options.max_tokens { options_obj["num_predict"] = n.into(); }
|
|
|
|
let body = json!({
|
|
"model": options.model,
|
|
"messages": msgs,
|
|
"stream": false,
|
|
"options": options_obj,
|
|
});
|
|
|
|
let url = format!("{}/api/chat", self.base_url.trim_end_matches('/'));
|
|
|
|
let resp: Value = self
|
|
.http
|
|
.post(&url)
|
|
.json(&body)
|
|
.send()
|
|
.await?
|
|
.error_for_status()?
|
|
.json()
|
|
.await?;
|
|
|
|
let content = resp["message"]["content"]
|
|
.as_str()
|
|
.ok_or_else(|| anyhow::anyhow!("Missing content in Ollama response"))?
|
|
.to_string();
|
|
|
|
let input_tokens = resp["prompt_eval_count"].as_u64().map(|n| n as u32);
|
|
let output_tokens = resp["eval_count"].as_u64().map(|n| n as u32);
|
|
|
|
Ok(ChatResponse { content, input_tokens, output_tokens, truncated: false, reasoning_content: None, cache_read_tokens: None, cache_creation_tokens: None, cost: None })
|
|
}
|
|
}
|