diff --git a/crates/owlen-providers/src/ollama/local.rs b/crates/owlen-providers/src/ollama/local.rs new file mode 100644 index 0000000..6975ab9 --- /dev/null +++ b/crates/owlen-providers/src/ollama/local.rs @@ -0,0 +1,79 @@ +use std::time::Duration; + +use async_trait::async_trait; +use owlen_core::provider::{ + GenerateRequest, GenerateStream, ModelInfo, ModelProvider, ProviderMetadata, ProviderStatus, + ProviderType, +}; +use owlen_core::{Error as CoreError, Result as CoreResult}; +use serde_json::Value; +use tokio::time::timeout; + +use super::OllamaClient; + +const DEFAULT_BASE_URL: &str = "http://localhost:11434"; +const DEFAULT_HEALTH_TIMEOUT_SECS: u64 = 5; + +/// ModelProvider implementation for a local Ollama daemon. +pub struct OllamaLocalProvider { + client: OllamaClient, + health_timeout: Duration, +} + +impl OllamaLocalProvider { + /// Construct a new local provider using the shared [`OllamaClient`]. + pub fn new( + base_url: Option, + request_timeout: Option, + health_timeout: Option, + ) -> CoreResult { + let base_url = base_url.unwrap_or_else(|| DEFAULT_BASE_URL.to_string()); + let health_timeout = + health_timeout.unwrap_or_else(|| Duration::from_secs(DEFAULT_HEALTH_TIMEOUT_SECS)); + + let mut metadata = + ProviderMetadata::new("ollama_local", "Ollama (Local)", ProviderType::Local, false); + metadata + .metadata + .insert("base_url".into(), Value::String(base_url.clone())); + if let Some(timeout) = request_timeout { + metadata.metadata.insert( + "request_timeout_ms".into(), + Value::String(timeout.as_millis().to_string()), + ); + } + + let client = OllamaClient::new(&base_url, None, metadata, request_timeout)?; + + Ok(Self { + client, + health_timeout, + }) + } +} + +#[async_trait] +impl ModelProvider for OllamaLocalProvider { + fn metadata(&self) -> &ProviderMetadata { + self.client.metadata() + } + + async fn health_check(&self) -> CoreResult { + match timeout(self.health_timeout, self.client.health_check()).await { + Ok(Ok(status)) => Ok(status), + Ok(Err(CoreError::Network(_))) | Ok(Err(CoreError::Timeout(_))) => { + Ok(ProviderStatus::Unavailable) + } + Ok(Err(err)) => Err(err), + Err(_) => Ok(ProviderStatus::Unavailable), + } + } + + async fn list_models(&self) -> CoreResult> { + self.client.list_models().await + } + + async fn generate_stream(&self, request: GenerateRequest) -> CoreResult { + self.client.generate_stream(request).await + } +} diff --git a/crates/owlen-providers/src/ollama/mod.rs b/crates/owlen-providers/src/ollama/mod.rs index 663f6f4..fd89c07 100644 --- a/crates/owlen-providers/src/ollama/mod.rs +++ b/crates/owlen-providers/src/ollama/mod.rs @@ -1,3 +1,5 @@ +pub mod local; pub mod shared; +pub use local::OllamaLocalProvider; pub use shared::OllamaClient; diff --git a/crates/owlen-providers/src/ollama/shared.rs b/crates/owlen-providers/src/ollama/shared.rs index caec128..25bb846 100644 --- a/crates/owlen-providers/src/ollama/shared.rs +++ b/crates/owlen-providers/src/ollama/shared.rs @@ -29,12 +29,14 @@ impl OllamaClient { base_url: impl AsRef, api_key: Option, provider_metadata: ProviderMetadata, + request_timeout: Option, ) -> CoreResult { let base_url = Url::parse(base_url.as_ref()) .map_err(|err| CoreError::Config(format!("invalid base url: {}", err)))?; + let timeout = request_timeout.unwrap_or_else(|| Duration::from_secs(DEFAULT_TIMEOUT_SECS)); let http = Client::builder() - .timeout(Duration::from_secs(DEFAULT_TIMEOUT_SECS)) + .timeout(timeout) .build() .map_err(map_reqwest_error)?;