//! 🦙 Ollama & LM Studio Provider - Local LLM Auto-Detection
//!
//! Automatically detects and connects to local LLM servers:
//! - Ollama at localhost:11434
//! - LM Studio at localhost:1234
//!
//! Both use OpenAI-compatible APIs, so we handle them uniformly.
//!
//! "Why pay for clouds when you've got a llama at home?" - The Cheet 🦙
use anyhow::{Context, Result};
use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::time::Duration;
use super::{LlmProvider, LlmRequest, LlmResponse, LlmUsage};
/// Default ports for local LLM servers
pub const OLLAMA_PORT: u16 = 11434;
pub const LMSTUDIO_PORT: u16 = 1234;
/// Detected local LLM server type
#[derive(Debug, Clone, PartialEq)]
pub enum LocalLlmType {
Ollama,
LmStudio,
}
impl std::fmt::Display for LocalLlmType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
LocalLlmType::Ollama => write!(f, "Ollama"),
LocalLlmType::LmStudio => write!(f, "LM Studio"),
}
}
}
/// Information about a detected local LLM server
#[derive(Debug, Clone)]
pub struct LocalLlmInfo {
pub server_type: LocalLlmType,
pub base_url: String,
pub models: Vec<String>,
}
/// 🦙 Provider for local LLM servers (Ollama, LM Studio)
pub struct OllamaProvider {
client: Client,
base_url: String,
server_type: LocalLlmType,
default_model: String,
}
impl OllamaProvider {
/// Create a new Ollama provider with explicit URL
pub fn new(base_url: &str, server_type: LocalLlmType) -> Self {
Self {
client: Client::builder()
.timeout(Duration::from_secs(300)) // Local models can be slow
.build()
.expect("Failed to create HTTP client"),
base_url: base_url.trim_end_matches('/').to_string(),
server_type,
default_model: "llama3.2".to_string(),
}
}
/// Create provider for Ollama at default port
pub fn ollama() -> Self {
Self::new(
&format!("http://localhost:{}", OLLAMA_PORT),
LocalLlmType::Ollama,
)
}
/// Create provider for LM Studio at default port
pub fn lmstudio() -> Self {
Self::new(
&format!("http://localhost:{}", LMSTUDIO_PORT),
LocalLlmType::LmStudio,
)
}
/// Set the default model to use
pub fn with_model(mut self, model: &str) -> Self {
self.default_model = model.to_string();
self
}
/// List available models from the server
pub async fn list_models(&self) -> Result<Vec<String>> {
match self.server_type {
LocalLlmType::Ollama => self.list_ollama_models().await,
LocalLlmType::LmStudio => self.list_lmstudio_models().await,
}
}
async fn list_ollama_models(&self) -> Result<Vec<String>> {
let url = format!("{}/api/tags", self.base_url);
let response = self
.client
.get(&url)
.send()
.await
.context("Failed to connect to Ollama")?;
let tags: OllamaTagsResponse = response
.json()
.await
.context("Failed to parse Ollama models response")?;
Ok(tags.models.into_iter().map(|m| m.name).collect())
}
async fn list_lmstudio_models(&self) -> Result<Vec<String>> {
let url = format!("{}/v1/models", self.base_url);
let response = self
.client
.get(&url)
.send()
.await
.context("Failed to connect to LM Studio")?;
let models: OpenAiModelsResponse = response
.json()
.await
.context("Failed to parse LM Studio models response")?;
Ok(models.data.into_iter().map(|m| m.id).collect())
}
}
impl Default for OllamaProvider {
fn default() -> Self {
Self::ollama()
}
}
#[async_trait]
impl LlmProvider for OllamaProvider {
async fn complete(&self, request: LlmRequest) -> Result<LlmResponse> {
let url = format!("{}/v1/chat/completions", self.base_url);
let model = if request.model.is_empty() || request.model == "default" {
self.default_model.clone()
} else {
request.model.clone()
};
let openai_request = OpenAiChatRequest {
model: model.clone(),
messages: request
.messages
.iter()
.map(|m| OpenAiMessage {
role: match m.role {
super::LlmRole::System => "system".to_string(),
super::LlmRole::User => "user".to_string(),
super::LlmRole::Assistant => "assistant".to_string(),
},
content: m.content.clone(),
})
.collect(),
temperature: request.temperature,
max_tokens: request.max_tokens,
stream: false, // We don't handle streaming in this basic impl
};
let response = self
.client
.post(&url)
.json(&openai_request)
.send()
.await
.context(format!("Failed to send request to {}", self.server_type))?;
if !response.status().is_success() {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
return Err(anyhow::anyhow!(
"{} returned error {}: {}",
self.server_type,
status,
error_text
));
}
let openai_response: OpenAiChatResponse = response
.json()
.await
.context("Failed to parse response from local LLM")?;
let content = openai_response
.choices
.first()
.map(|c| c.message.content.clone())
.unwrap_or_default();
Ok(LlmResponse {
content,
model: openai_response.model,
usage: openai_response.usage.map(|u| LlmUsage {
prompt_tokens: u.prompt_tokens,
completion_tokens: u.completion_tokens,
total_tokens: u.total_tokens,
}),
})
}
fn name(&self) -> &'static str {
match self.server_type {
LocalLlmType::Ollama => "ollama",
LocalLlmType::LmStudio => "lmstudio",
}
}
}
// ============================================================================
// Auto-Detection
// ============================================================================
/// Check if a local LLM server is running at the given port
pub async fn check_server(host: &str, port: u16, timeout_ms: u64) -> bool {
let client = match Client::builder()
.timeout(Duration::from_millis(timeout_ms))
.build()
{
Ok(c) => c,
Err(_) => return false,
};
// Try the health/version endpoint first (fast)
let health_url = format!("http://{}:{}/", host, port);
if client.get(&health_url).send().await.is_ok() {
return true;
}
// Fallback: try the models endpoint
let models_url = format!("http://{}:{}/v1/models", host, port);
client.get(&models_url).send().await.is_ok()
}
/// Detect all available local LLM servers
pub async fn detect_local_llms() -> Vec<LocalLlmInfo> {
let mut detected = Vec::new();
// Check Ollama
if check_server("localhost", OLLAMA_PORT, 500).await {
let provider = OllamaProvider::ollama();
let models = provider.list_models().await.unwrap_or_default();
detected.push(LocalLlmInfo {
server_type: LocalLlmType::Ollama,
base_url: format!("http://localhost:{}", OLLAMA_PORT),
models,
});
}
// Check LM Studio
if check_server("localhost", LMSTUDIO_PORT, 500).await {
let provider = OllamaProvider::lmstudio();
let models = provider.list_models().await.unwrap_or_default();
detected.push(LocalLlmInfo {
server_type: LocalLlmType::LmStudio,
base_url: format!("http://localhost:{}", LMSTUDIO_PORT),
models,
});
}
detected
}
/// Quick check if any local LLM is available (non-blocking, fast timeout)
pub async fn any_local_llm_available() -> bool {
tokio::select! {
ollama = check_server("localhost", OLLAMA_PORT, 200) => {
if ollama { return true; }
}
lmstudio = check_server("localhost", LMSTUDIO_PORT, 200) => {
if lmstudio { return true; }
}
}
// Check remaining
check_server("localhost", OLLAMA_PORT, 200).await
|| check_server("localhost", LMSTUDIO_PORT, 200).await
}
// ============================================================================
// API Types
// ============================================================================
#[derive(Debug, Deserialize)]
struct OllamaTagsResponse {
models: Vec<OllamaModel>,
}
#[derive(Debug, Deserialize)]
struct OllamaModel {
name: String,
#[allow(dead_code)]
modified_at: Option<String>,
#[allow(dead_code)]
size: Option<u64>,
}
#[derive(Debug, Deserialize)]
struct OpenAiModelsResponse {
data: Vec<OpenAiModelInfo>,
}
#[derive(Debug, Deserialize)]
struct OpenAiModelInfo {
id: String,
}
#[derive(Debug, Serialize)]
struct OpenAiChatRequest {
model: String,
messages: Vec<OpenAiMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
max_tokens: Option<usize>,
stream: bool,
}
#[derive(Debug, Serialize, Deserialize)]
struct OpenAiMessage {
role: String,
content: String,
}
#[derive(Debug, Deserialize)]
struct OpenAiChatResponse {
model: String,
choices: Vec<OpenAiChoice>,
usage: Option<OpenAiUsageInfo>,
}
#[derive(Debug, Deserialize)]
struct OpenAiChoice {
message: OpenAiMessage,
}
#[derive(Debug, Deserialize)]
struct OpenAiUsageInfo {
prompt_tokens: usize,
completion_tokens: usize,
total_tokens: usize,
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_detect_local_llms() {
// This test will pass whether or not local LLMs are running
let detected = detect_local_llms().await;
println!("Detected {} local LLM server(s)", detected.len());
for info in &detected {
println!(
" - {} at {} with {} models",
info.server_type,
info.base_url,
info.models.len()
);
for model in &info.models {
println!(" • {}", model);
}
}
}
#[tokio::test]
async fn test_check_server_timeout() {
// Should timeout quickly on non-existent server
let start = std::time::Instant::now();
let result = check_server("localhost", 59999, 100).await;
let elapsed = start.elapsed();
assert!(!result);
assert!(
elapsed.as_millis() < 500,
"Timeout took too long: {:?}",
elapsed
);
}
}
