Alice Hau
120 lines
3.6 KiB
Rust
120 lines
3.6 KiB
Rust
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
|
|
use super::errors::ProviderError;
|
|
use super::utils::{get_model, handle_response_openai_compat};
|
|
use crate::message::Message;
|
|
use crate::model::ModelConfig;
|
|
use crate::providers::formats::openai::{create_request, get_usage, response_to_message};
|
|
use anyhow::Result;
|
|
use async_trait::async_trait;
|
|
use mcp_core::tool::Tool;
|
|
use reqwest::Client;
|
|
use serde_json::Value;
|
|
use std::time::Duration;
|
|
|
|
pub const OLLAMA_HOST: &str = "http://localhost:11434";
|
|
pub const OLLAMA_DEFAULT_MODEL: &str = "qwen2.5";
|
|
// Ollama can run many models, we only provide the default
|
|
pub const OLLAMA_KNOWN_MODELS: &[&str] = &[OLLAMA_DEFAULT_MODEL];
|
|
pub const OLLAMA_DOC_URL: &str = "https://ollama.com/library";
|
|
|
|
#[derive(serde::Serialize)]
|
|
pub struct OllamaProvider {
|
|
#[serde(skip)]
|
|
client: Client,
|
|
host: String,
|
|
model: ModelConfig,
|
|
}
|
|
|
|
impl Default for OllamaProvider {
|
|
fn default() -> Self {
|
|
let model = ModelConfig::new(OllamaProvider::metadata().default_model);
|
|
OllamaProvider::from_env(model).expect("Failed to initialize Ollama provider")
|
|
}
|
|
}
|
|
|
|
impl OllamaProvider {
|
|
pub fn from_env(model: ModelConfig) -> Result<Self> {
|
|
let config = crate::config::Config::global();
|
|
let host: String = config
|
|
.get("OLLAMA_HOST")
|
|
.unwrap_or_else(|_| OLLAMA_HOST.to_string());
|
|
|
|
let client = Client::builder()
|
|
.timeout(Duration::from_secs(600))
|
|
.build()?;
|
|
|
|
Ok(Self {
|
|
client,
|
|
host,
|
|
model,
|
|
})
|
|
}
|
|
|
|
async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
|
|
let url = format!("{}/v1/chat/completions", self.host.trim_end_matches('/'));
|
|
|
|
let response = self.client.post(&url).json(&payload).send().await?;
|
|
|
|
handle_response_openai_compat(response).await
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for OllamaProvider {
|
|
fn metadata() -> ProviderMetadata {
|
|
ProviderMetadata::new(
|
|
"ollama",
|
|
"Ollama",
|
|
"Local open source models",
|
|
OLLAMA_DEFAULT_MODEL,
|
|
OLLAMA_KNOWN_MODELS.iter().map(|&s| s.to_string()).collect(),
|
|
OLLAMA_DOC_URL,
|
|
vec![ConfigKey::new(
|
|
"OLLAMA_HOST",
|
|
true,
|
|
false,
|
|
Some(OLLAMA_HOST),
|
|
)],
|
|
)
|
|
}
|
|
|
|
fn get_model_config(&self) -> ModelConfig {
|
|
self.model.clone()
|
|
}
|
|
|
|
#[tracing::instrument(
|
|
skip(self, system, messages, tools),
|
|
fields(model_config, input, output, input_tokens, output_tokens, total_tokens)
|
|
)]
|
|
async fn complete(
|
|
&self,
|
|
system: &str,
|
|
messages: &[Message],
|
|
tools: &[Tool],
|
|
) -> Result<(Message, ProviderUsage), ProviderError> {
|
|
let payload = create_request(
|
|
&self.model,
|
|
system,
|
|
messages,
|
|
tools,
|
|
&super::utils::ImageFormat::OpenAi,
|
|
)?;
|
|
|
|
let response = self.post(payload.clone()).await?;
|
|
|
|
// Parse response
|
|
let message = response_to_message(response.clone())?;
|
|
let usage = match get_usage(&response) {
|
|
Ok(usage) => usage,
|
|
Err(ProviderError::UsageError(e)) => {
|
|
tracing::warn!("Failed to get usage data: {}", e);
|
|
Usage::default()
|
|
}
|
|
Err(e) => return Err(e),
|
|
};
|
|
let model = get_model(&response);
|
|
super::utils::emit_debug_trace(self, &payload, &response, &usage);
|
|
Ok((message, ProviderUsage::new(model, usage)))
|
|
}
|
|
}
|