Skip to main content

volition_core/providers/
mod.rs

1// volition-agent-core/src/providers/mod.rs
2use crate::models::chat::{ApiResponse, ChatMessage};
3use crate::models::tools::ToolDefinition; // Import ToolDefinition
4use anyhow::{Result, anyhow};
5use async_trait::async_trait;
6use std::collections::HashMap;
7
8#[async_trait]
9pub trait Provider: Send + Sync {
10    // Add tools argument
11    async fn get_completion(
12        &self,
13        messages: Vec<ChatMessage>,
14        tools: Option<&[ToolDefinition]>, // Add tools argument
15    ) -> Result<ApiResponse>;
16    fn name(&self) -> &str;
17}
18
19// Add ProviderRegistry back
20pub struct ProviderRegistry {
21    providers: HashMap<String, Box<dyn Provider>>,
22    default_provider: String,
23}
24
25impl ProviderRegistry {
26    pub fn new(default_provider: String) -> Self {
27        Self {
28            providers: HashMap::new(),
29            default_provider,
30        }
31    }
32
33    pub fn register(&mut self, id: String, provider: Box<dyn Provider>) {
34        self.providers.insert(id, provider);
35    }
36
37    pub fn get(&self, id: &str) -> Result<&dyn Provider> {
38        self.providers
39            .get(id)
40            .map(|p| p.as_ref())
41            .ok_or_else(|| anyhow!("Provider not found: {}", id))
42    }
43
44    pub fn default(&self) -> Result<&dyn Provider> {
45        self.get(&self.default_provider)
46    }
47
48    pub fn default_provider_id(&self) -> &str {
49        &self.default_provider
50    }
51}
52
53pub mod gemini;
54pub mod ollama;
55pub mod openai; // Added OpenAI module