Compare commits
7 Commits
d4cebe21c7
...
dadc178fd6
| Author | SHA1 | Date | |
|---|---|---|---|
| dadc178fd6 | |||
| 9f1b5f908d | |||
| b52d82dec5 | |||
| aa778579ee | |||
| d874194d95 | |||
| bd8350b52a | |||
| b66d474cd3 |
@@ -76,9 +76,24 @@ impl Agent {
|
||||
|
||||
/// Get the model for a specific task type
|
||||
fn get_model(&self, task_type: TaskType) -> Result<String> {
|
||||
self.config
|
||||
.get_model_for_task(task_type)
|
||||
.ok_or_else(|| AgentError::ConfigError(format!("No model configured for {:?}", task_type)))
|
||||
// First try to get model from agent config
|
||||
if let Some(model) = self.config.get_model_for_task(task_type) {
|
||||
return Ok(model);
|
||||
}
|
||||
|
||||
// Fallback: use the default model from any available provider
|
||||
// Prefer ollama for local-first usage
|
||||
if self.provider_registry.get("ollama").is_some() {
|
||||
return Ok("ollama".to_string());
|
||||
}
|
||||
if self.provider_registry.get("openai").is_some() {
|
||||
return Ok("openai".to_string());
|
||||
}
|
||||
|
||||
// No providers available
|
||||
Err(AgentError::ConfigError(
|
||||
"No model configured and no providers available".to_string()
|
||||
))
|
||||
}
|
||||
|
||||
/// Get available tools as provider Tool definitions
|
||||
|
||||
@@ -147,15 +147,8 @@ pub struct AgentRuntimeConfig {
|
||||
|
||||
impl AgentRuntimeConfig {
|
||||
pub fn get_model_for_task(&self, task_type: TaskType) -> Option<String> {
|
||||
let model_str = match task_type {
|
||||
TaskType::Chat => self.config.models.chat.as_ref(),
|
||||
TaskType::Code => self.config.models.code.as_ref(),
|
||||
TaskType::Planning => self.config.models.planning.as_ref(),
|
||||
TaskType::ToolUse => self.config.models.tool_use.as_ref(),
|
||||
TaskType::Default => self.config.models.default.as_ref(),
|
||||
};
|
||||
|
||||
model_str.cloned()
|
||||
// Use the AgentModels helper which handles fallback to default
|
||||
self.config.models.for_task(task_type)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,21 +1,26 @@
|
||||
use crate::agents::{AgentRequest, AgentTrait};
|
||||
use crate::providers::TaskType;
|
||||
use crate::providers::{ProviderRegistry, TaskType};
|
||||
use std::io::{self, Write};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Interactive REPL for svibe
|
||||
pub struct Repl {
|
||||
agent: std::sync::Arc<dyn AgentTrait>,
|
||||
agent: Arc<dyn AgentTrait>,
|
||||
provider_registry: Arc<ProviderRegistry>,
|
||||
}
|
||||
|
||||
impl Repl {
|
||||
pub fn new(agent: std::sync::Arc<dyn AgentTrait>) -> Self {
|
||||
Self { agent }
|
||||
pub fn new(agent: Arc<dyn AgentTrait>, provider_registry: Arc<ProviderRegistry>) -> Self {
|
||||
Self {
|
||||
agent,
|
||||
provider_registry,
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the REPL
|
||||
pub fn run(&self) -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// Run the REPL (async version)
|
||||
pub async fn run(&self) -> Result<(), Box<dyn std::error::Error>> {
|
||||
println!("svibe REPL - Type 'exit' or 'quit' to leave");
|
||||
println!("Commands: /help, /clear, /status, /agents, /tools");
|
||||
println!("Commands: /help, /clear, /status, /agents, /tools, /connect");
|
||||
println!();
|
||||
|
||||
let stdin = io::stdin();
|
||||
@@ -44,25 +49,13 @@ impl Repl {
|
||||
}
|
||||
|
||||
if input.starts_with('/') {
|
||||
self.handle_command(input)?;
|
||||
self.handle_command(input).await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Process as a task - need to use tokio runtime
|
||||
let input = input.to_string();
|
||||
let agent = self.agent.clone();
|
||||
|
||||
let rt = tokio::runtime::Handle::current();
|
||||
let result = rt.block_on(async move {
|
||||
let request = AgentRequest {
|
||||
task: input,
|
||||
task_type: TaskType::Default,
|
||||
context: vec![],
|
||||
history: vec![],
|
||||
};
|
||||
agent.process(request).await
|
||||
});
|
||||
|
||||
// Process as a task
|
||||
let result = self.process_task(input).await;
|
||||
|
||||
match result {
|
||||
Ok(response) => {
|
||||
println!("\n{}", response.content);
|
||||
@@ -81,7 +74,17 @@ impl Repl {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_command(&self, input: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
async fn process_task(&self, input: &str) -> Result<crate::agents::AgentResponse, crate::agents::AgentError> {
|
||||
let request = AgentRequest {
|
||||
task: input.to_string(),
|
||||
task_type: TaskType::Default,
|
||||
context: vec![],
|
||||
history: vec![],
|
||||
};
|
||||
self.agent.process(request).await
|
||||
}
|
||||
|
||||
async fn handle_command(&self, input: &str) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let parts: Vec<&str> = input.split_whitespace().collect();
|
||||
|
||||
match parts.get(0).copied() {
|
||||
@@ -94,6 +97,7 @@ Available commands:
|
||||
/status Show agent status and metrics
|
||||
/agents List available agents
|
||||
/tools List available tools
|
||||
/connect Check connectivity to all providers
|
||||
/exit, /quit, :q Exit the REPL
|
||||
|
||||
Just type your task to get started!
|
||||
@@ -128,6 +132,9 @@ Metrics:
|
||||
Some("tools") | Some("/tools") => {
|
||||
println!("Tool listing not available in this context");
|
||||
}
|
||||
Some("connect") | Some("/connect") => {
|
||||
self.check_connectivity().await?;
|
||||
}
|
||||
_ => {
|
||||
println!("Unknown command. Type /help for available commands.");
|
||||
}
|
||||
@@ -135,4 +142,35 @@ Metrics:
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn check_connectivity(&self) -> Result<(), Box<dyn std::error::Error>> {
|
||||
println!("\nChecking provider connectivity...\n");
|
||||
|
||||
let health = self.provider_registry.health_check_all().await;
|
||||
let providers = self.provider_registry.list();
|
||||
|
||||
if providers.is_empty() {
|
||||
println!("No providers configured.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut all_healthy = true;
|
||||
for provider in providers {
|
||||
let healthy = health.get(provider).copied().unwrap_or(false);
|
||||
let status = if healthy { "✓" } else { "✗" };
|
||||
println!(" {} {}: {}", provider, status, if healthy { "healthy" } else { "unhealthy" });
|
||||
if !healthy {
|
||||
all_healthy = false;
|
||||
}
|
||||
}
|
||||
|
||||
println!();
|
||||
if all_healthy {
|
||||
println!("All providers are healthy!");
|
||||
} else {
|
||||
println!("Some providers are unhealthy. Check your configuration.");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
use super::{Config, ConfigError, Result, OpenAiConfig, OllamaConfig, AgentConfig, AgentModels};
|
||||
use super::types::{default_openai_base_url, default_openai_model, default_ollama_base_url, default_ollama_model};
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Configuration loader
|
||||
@@ -138,17 +137,20 @@ fn find_local_config() -> Option<PathBuf> {
|
||||
/// Merge environment variables into config
|
||||
impl Config {
|
||||
fn merge_env(&mut self) {
|
||||
// OpenAI API key from environment
|
||||
if let Some(ref mut openai) = self.providers.openai {
|
||||
// OpenAI API key from environment (applies to first instance without key)
|
||||
for openai in &mut self.providers.openai {
|
||||
if openai.api_key.is_none() {
|
||||
openai.api_key = std::env::var("OPENAI_API_KEY").ok();
|
||||
if let Ok(key) = std::env::var("OPENAI_API_KEY") {
|
||||
openai.api_key = Some(key);
|
||||
break; // Only apply to first one
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Ollama base URL from environment
|
||||
if let Some(ref mut ollama) = self.providers.ollama {
|
||||
// Ollama base URL from environment (applies to all instances)
|
||||
for ollama in &mut self.providers.ollama {
|
||||
if let Ok(url) = std::env::var("OLLAMA_HOST") {
|
||||
if !url.is_empty() {
|
||||
if !url.is_empty() && ollama.base_url == default_ollama_base_url() {
|
||||
ollama.base_url = url;
|
||||
}
|
||||
}
|
||||
@@ -157,10 +159,16 @@ impl Config {
|
||||
|
||||
/// Merge another config into this one (other overrides self)
|
||||
fn merge(&mut self, other: Config) {
|
||||
// Merge providers
|
||||
if let Some(other_openai) = other.providers.openai {
|
||||
if let Some(ref mut self_openai) = self.providers.openai {
|
||||
// Merge OpenAI configs
|
||||
// Merge OpenAI providers (extend list)
|
||||
for other_openai in other.providers.openai {
|
||||
// Check if we should merge with existing or add new
|
||||
let existing = self.providers.openai.iter_mut().find(|p| {
|
||||
p.name == other_openai.name ||
|
||||
(p.name.is_none() && other_openai.name.is_none())
|
||||
});
|
||||
|
||||
if let Some(self_openai) = existing {
|
||||
// Merge with existing provider
|
||||
if other_openai.api_key.is_some() {
|
||||
self_openai.api_key = other_openai.api_key;
|
||||
}
|
||||
@@ -171,13 +179,20 @@ impl Config {
|
||||
self_openai.default_model = other_openai.default_model;
|
||||
}
|
||||
} else {
|
||||
self.providers.openai = Some(other_openai);
|
||||
// Add as new provider
|
||||
self.providers.openai.push(other_openai);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(other_ollama) = other.providers.ollama {
|
||||
if let Some(ref mut self_ollama) = self.providers.ollama {
|
||||
// Merge Ollama configs
|
||||
// Merge Ollama providers (extend list)
|
||||
for other_ollama in other.providers.ollama {
|
||||
let existing = self.providers.ollama.iter_mut().find(|p| {
|
||||
p.name == other_ollama.name ||
|
||||
(p.name.is_none() && other_ollama.name.is_none())
|
||||
});
|
||||
|
||||
if let Some(self_ollama) = existing {
|
||||
// Merge with existing provider
|
||||
if other_ollama.base_url != default_ollama_base_url() {
|
||||
self_ollama.base_url = other_ollama.base_url;
|
||||
}
|
||||
@@ -194,7 +209,8 @@ impl Config {
|
||||
self_ollama.password = other_ollama.password;
|
||||
}
|
||||
} else {
|
||||
self.providers.ollama = Some(other_ollama);
|
||||
// Add as new provider
|
||||
self.providers.ollama.push(other_ollama);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -232,11 +248,9 @@ fn default_agent_config() -> AgentConfig {
|
||||
AgentConfig {
|
||||
description: "Default general-purpose agent".to_string(),
|
||||
models: AgentModels {
|
||||
chat: Some("ollama/llama3".to_string()),
|
||||
code: Some("openai/gpt-4o".to_string()),
|
||||
planning: Some("openai/gpt-4o".to_string()),
|
||||
tool_use: Some("openai/gpt-4o".to_string()),
|
||||
// Only set default model - will be used for all task types
|
||||
default: Some("ollama/llama3".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
system_prompt: "You are a helpful AI assistant. Help the user with their tasks.".to_string(),
|
||||
tools: vec!["file".to_string(), "shell".to_string(), "git".to_string()],
|
||||
@@ -249,13 +263,17 @@ fn default_agent_config() -> AgentConfig {
|
||||
pub fn create_default_config() -> Result<Config> {
|
||||
let mut config = Config::default();
|
||||
|
||||
config.providers.openai = Some(OpenAiConfig {
|
||||
api_key: None, // User must set this
|
||||
// Add default OpenAI provider
|
||||
config.providers.openai.push(OpenAiConfig {
|
||||
name: Some("default".to_string()),
|
||||
api_key: None, // User must set this or use env var
|
||||
base_url: default_openai_base_url(),
|
||||
default_model: default_openai_model(),
|
||||
});
|
||||
|
||||
config.providers.ollama = Some(OllamaConfig {
|
||||
// Add default Ollama provider
|
||||
config.providers.ollama.push(OllamaConfig {
|
||||
name: Some("default".to_string()),
|
||||
base_url: default_ollama_base_url(),
|
||||
default_model: default_ollama_model(),
|
||||
bearer_token: None,
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::providers::TaskType;
|
||||
|
||||
/// Main configuration structure
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Config {
|
||||
@@ -28,18 +30,22 @@ fn default_agent() -> String {
|
||||
/// Provider configurations
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct ProvidersConfig {
|
||||
/// OpenAI provider config
|
||||
/// OpenAI provider instances (can have multiple)
|
||||
#[serde(default)]
|
||||
pub openai: Option<OpenAiConfig>,
|
||||
pub openai: Vec<OpenAiConfig>,
|
||||
|
||||
/// Ollama provider config
|
||||
/// Ollama provider instances (can have multiple)
|
||||
#[serde(default)]
|
||||
pub ollama: Option<OllamaConfig>,
|
||||
pub ollama: Vec<OllamaConfig>,
|
||||
}
|
||||
|
||||
/// OpenAI configuration
|
||||
/// Named OpenAI configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OpenAiConfig {
|
||||
/// Name for this provider instance (optional, for reference)
|
||||
#[serde(default)]
|
||||
pub name: Option<String>,
|
||||
|
||||
/// API key (can also be set via OPENAI_API_KEY env var)
|
||||
pub api_key: Option<String>,
|
||||
|
||||
@@ -63,6 +69,7 @@ pub fn default_openai_model() -> String {
|
||||
impl Default for OpenAiConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
name: None,
|
||||
api_key: None,
|
||||
base_url: default_openai_base_url(),
|
||||
default_model: default_openai_model(),
|
||||
@@ -70,9 +77,13 @@ impl Default for OpenAiConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// Ollama configuration
|
||||
/// Named Ollama configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OllamaConfig {
|
||||
/// Name for this provider instance (optional, for reference)
|
||||
#[serde(default)]
|
||||
pub name: Option<String>,
|
||||
|
||||
/// Base URL
|
||||
#[serde(default = "default_ollama_base_url")]
|
||||
pub base_url: String,
|
||||
@@ -105,6 +116,7 @@ pub fn default_ollama_model() -> String {
|
||||
impl Default for OllamaConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
name: None,
|
||||
base_url: default_ollama_base_url(),
|
||||
default_model: default_ollama_model(),
|
||||
bearer_token: None,
|
||||
@@ -121,7 +133,8 @@ pub struct AgentConfig {
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
|
||||
/// Model routing configuration
|
||||
/// Model routing configuration (optional, uses default model for all tasks if not specified)
|
||||
#[serde(default)]
|
||||
pub models: AgentModels,
|
||||
|
||||
/// System prompt for this agent
|
||||
@@ -142,6 +155,7 @@ pub struct AgentConfig {
|
||||
}
|
||||
|
||||
/// Model routing for different task types
|
||||
/// If no models are specified, the default model is used for all tasks
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct AgentModels {
|
||||
/// Model for general chat/conversation
|
||||
@@ -160,21 +174,34 @@ pub struct AgentModels {
|
||||
#[serde(default)]
|
||||
pub tool_use: Option<String>,
|
||||
|
||||
/// Fallback model when task type is unclear
|
||||
/// Fallback model for all tasks (used if task-specific models not set)
|
||||
#[serde(default)]
|
||||
pub default: Option<String>,
|
||||
}
|
||||
|
||||
impl AgentModels {
|
||||
/// Get the model for a specific task type
|
||||
pub fn for_task(&self, task_type: &str) -> Option<&str> {
|
||||
match task_type {
|
||||
"chat" => self.chat.as_deref(),
|
||||
"code" => self.code.as_deref(),
|
||||
"planning" => self.planning.as_deref(),
|
||||
"tool_use" => self.tool_use.as_deref(),
|
||||
_ => self.default.as_deref(),
|
||||
}
|
||||
/// Falls back to default model if task-specific model not set
|
||||
pub fn for_task(&self, task_type: TaskType) -> Option<String> {
|
||||
let task_model = match task_type {
|
||||
TaskType::Chat => self.chat.as_ref(),
|
||||
TaskType::Code => self.code.as_ref(),
|
||||
TaskType::Planning => self.planning.as_ref(),
|
||||
TaskType::ToolUse => self.tool_use.as_ref(),
|
||||
TaskType::Default => self.default.as_ref(),
|
||||
};
|
||||
|
||||
// Return task-specific model if set, otherwise use default
|
||||
task_model.or(self.default.as_ref()).cloned()
|
||||
}
|
||||
|
||||
/// Check if any models are configured
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.chat.is_none() &&
|
||||
self.code.is_none() &&
|
||||
self.planning.is_none() &&
|
||||
self.tool_use.is_none() &&
|
||||
self.default.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -74,8 +74,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
match cli.command {
|
||||
None | Some(Commands::Repl) => {
|
||||
// Start REPL
|
||||
let repl = Repl::new(agent);
|
||||
if let Err(e) = repl.run() {
|
||||
let repl = Repl::new(agent, provider_registry.clone());
|
||||
if let Err(e) = repl.run().await {
|
||||
eprintln!("REPL error: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,7 +46,11 @@ impl OllamaProvider {
|
||||
// Handle "ollama/model-name" format
|
||||
if let Some(stripped) = model_spec.strip_prefix("ollama/") {
|
||||
stripped.to_string()
|
||||
} else if model_spec == "ollama" {
|
||||
// Just provider name - use default model
|
||||
self.default_model.clone()
|
||||
} else {
|
||||
// Use as-is (could be a model name)
|
||||
model_spec.to_string()
|
||||
}
|
||||
}
|
||||
@@ -94,9 +98,20 @@ impl Provider for OllamaProvider {
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> Result<bool> {
|
||||
// First check if the API is reachable
|
||||
let url = format!("{}/api/tags", self.base_url);
|
||||
let response = self.apply_auth(self.client.get(&url)).send().await?;
|
||||
Ok(response.status().is_success())
|
||||
let response = self.apply_auth(self.client.get(&url))
|
||||
.timeout(std::time::Duration::from_secs(5))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// Also verify we can parse the response (deeper health check)
|
||||
let body: serde_json::Value = response.json().await?;
|
||||
Ok(body.get("models").is_some())
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>> {
|
||||
@@ -155,6 +170,7 @@ impl Provider for OllamaProvider {
|
||||
self.client
|
||||
.post(&url)
|
||||
.header("Content-Type", "application/json")
|
||||
.timeout(std::time::Duration::from_secs(120)) // 2 minute timeout for chat
|
||||
)
|
||||
.json(&ollama_request)
|
||||
.send()
|
||||
|
||||
@@ -27,7 +27,11 @@ impl OpenAiProvider {
|
||||
// Handle "openai/model-name" format
|
||||
if let Some(stripped) = model_spec.strip_prefix("openai/") {
|
||||
stripped.to_string()
|
||||
} else if model_spec == "openai" {
|
||||
// Just provider name - use default model
|
||||
self.default_model.clone()
|
||||
} else {
|
||||
// Use as-is (could be a model name)
|
||||
model_spec.to_string()
|
||||
}
|
||||
}
|
||||
@@ -45,10 +49,17 @@ impl Provider for OpenAiProvider {
|
||||
.client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.timeout(std::time::Duration::from_secs(5))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
Ok(response.status().is_success())
|
||||
if !response.status().is_success() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// Verify we can parse the response
|
||||
let body: serde_json::Value = response.json().await?;
|
||||
Ok(body.get("data").is_some())
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>> {
|
||||
@@ -148,6 +159,7 @@ impl Provider for OpenAiProvider {
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.timeout(std::time::Duration::from_secs(120)) // 2 minute timeout for chat
|
||||
.json(&req_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
@@ -16,27 +16,50 @@ impl ProviderRegistry {
|
||||
}
|
||||
|
||||
/// Create a registry from configuration
|
||||
/// Does not require API key if at least one provider is configured
|
||||
pub fn from_config(config: &Config) -> Result<Self> {
|
||||
let mut registry = Self::new();
|
||||
let mut has_any_provider = false;
|
||||
|
||||
// Register OpenAI provider if configured
|
||||
if let Some(ref openai_config) = config.providers.openai {
|
||||
// Register all OpenAI provider instances
|
||||
for (i, openai_config) in config.providers.openai.iter().enumerate() {
|
||||
let name = openai_config
|
||||
.name
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("openai-{}", i + 1));
|
||||
let registry_name = if i == 0 { "openai".to_string() } else { name.clone() };
|
||||
|
||||
// Get API key from config or environment
|
||||
let api_key = openai_config
|
||||
.api_key
|
||||
.clone()
|
||||
.ok_or_else(|| ProviderError::AuthError("OpenAI API key not configured".to_string()))?;
|
||||
.or_else(|| std::env::var("OPENAI_API_KEY").ok());
|
||||
|
||||
let provider = OpenAiProvider::new(
|
||||
&openai_config.base_url,
|
||||
&api_key,
|
||||
&openai_config.default_model,
|
||||
);
|
||||
|
||||
registry.register("openai", Arc::new(provider));
|
||||
if let Some(key) = api_key {
|
||||
let provider = OpenAiProvider::new(
|
||||
&openai_config.base_url,
|
||||
&key,
|
||||
&openai_config.default_model,
|
||||
);
|
||||
registry.register(®istry_name, Arc::new(provider));
|
||||
has_any_provider = true;
|
||||
} else if config.providers.ollama.is_empty() {
|
||||
// Only error if no other providers are available
|
||||
return Err(ProviderError::AuthError(
|
||||
"OpenAI API key not configured. Set api_key in config or OPENAI_API_KEY env var, or configure Ollama provider.".to_string()
|
||||
));
|
||||
}
|
||||
// If no API key but Ollama is available, just skip this OpenAI instance
|
||||
}
|
||||
|
||||
// Register Ollama provider if configured
|
||||
if let Some(ref ollama_config) = config.providers.ollama {
|
||||
// Register all Ollama provider instances
|
||||
for (i, ollama_config) in config.providers.ollama.iter().enumerate() {
|
||||
let name = ollama_config
|
||||
.name
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("ollama-{}", i + 1));
|
||||
let registry_name = if i == 0 { "ollama".to_string() } else { name.clone() };
|
||||
|
||||
let mut provider = OllamaProvider::new(
|
||||
&ollama_config.base_url,
|
||||
&ollama_config.default_model,
|
||||
@@ -45,12 +68,19 @@ impl ProviderRegistry {
|
||||
// Apply authentication
|
||||
if let Some(ref token) = ollama_config.bearer_token {
|
||||
provider = provider.with_bearer_token(token);
|
||||
} else if let (Some(ref username), Some(ref password)) =
|
||||
} else if let (Some(ref username), Some(ref password)) =
|
||||
(&ollama_config.username, &ollama_config.password) {
|
||||
provider = provider.with_basic_auth(username, password);
|
||||
}
|
||||
|
||||
registry.register("ollama", Arc::new(provider));
|
||||
registry.register(®istry_name, Arc::new(provider));
|
||||
has_any_provider = true;
|
||||
}
|
||||
|
||||
if !has_any_provider && config.providers.openai.is_empty() && config.providers.ollama.is_empty() {
|
||||
return Err(ProviderError::AuthError(
|
||||
"No providers configured. Add at least one OpenAI or Ollama provider.".to_string()
|
||||
));
|
||||
}
|
||||
|
||||
Ok(registry)
|
||||
@@ -73,20 +103,29 @@ impl ProviderRegistry {
|
||||
}
|
||||
|
||||
/// Resolve a model specification to a provider
|
||||
/// Format: "provider/model-name" or just "model-name" (uses default)
|
||||
/// Format:
|
||||
/// - "provider/model-name" - use specific model
|
||||
/// - "provider" - use provider's default model
|
||||
/// - "model-name" - try to find provider with this model as default
|
||||
pub fn resolve_provider(&self, model_spec: &str) -> Result<(Arc<dyn Provider>, String)> {
|
||||
if let Some((provider_name, model_name)) = model_spec.split_once('/') {
|
||||
let provider = self.get_or_error(provider_name)?;
|
||||
Ok((provider, model_name.to_string()))
|
||||
} else {
|
||||
// Try to find a provider that has this model as default
|
||||
// For now, default to openai if available
|
||||
if let Some(provider) = self.get("openai") {
|
||||
// Check if it's just a provider name (no model specified)
|
||||
// In this case, use the provider's default model
|
||||
if let Some(provider) = self.get(model_spec) {
|
||||
return Ok((provider, model_spec.to_string()));
|
||||
}
|
||||
|
||||
// Try to find a provider that has this model as default
|
||||
// Prefer Ollama for local-first
|
||||
if let Some(provider) = self.get("ollama") {
|
||||
return Ok((provider, model_spec.to_string()));
|
||||
}
|
||||
if let Some(provider) = self.get("openai") {
|
||||
return Ok((provider, model_spec.to_string()));
|
||||
}
|
||||
Err(ProviderError::UnknownProvider(
|
||||
"No providers available".to_string(),
|
||||
))
|
||||
@@ -109,6 +148,11 @@ impl ProviderRegistry {
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
/// Check if any provider is available
|
||||
pub fn has_providers(&self) -> bool {
|
||||
!self.providers.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ProviderRegistry {
|
||||
@@ -124,7 +168,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_resolve_provider_format() {
|
||||
let registry = ProviderRegistry::new();
|
||||
|
||||
|
||||
// Test parsing format
|
||||
assert_eq!(
|
||||
"openai/gpt-4".split_once('/'),
|
||||
|
||||
Reference in New Issue
Block a user