agent_base/llm/
registry.rs

1use std::sync::Arc;
2
3use super::{AnthropicClient, LlmClient, OpenAiClient};
4
5#[derive(Clone, Debug)]
6pub enum LlmProvider {
7    OpenAi,
8    Anthropic,
9    Custom(String),
10}
11
12impl LlmProvider {
13    pub fn from_str(s: &str) -> Self {
14        match s.to_lowercase().as_str() {
15            "openai" => Self::OpenAi,
16            "anthropic" => Self::Anthropic,
17            other => Self::Custom(other.to_string()),
18        }
19    }
20}
21
22pub struct LlmClientBuilder {
23    provider: LlmProvider,
24    api_key: String,
25    model: String,
26    base_url: Option<String>,
27}
28
29impl LlmClientBuilder {
30    pub fn new(provider: LlmProvider, api_key: impl Into<String>, model: impl Into<String>) -> Self {
31        Self {
32            provider,
33            api_key: api_key.into(),
34            model: model.into(),
35            base_url: None,
36        }
37    }
38
39    pub fn from_env() -> Option<Self> {
40        let api_key = std::env::var("LLM_API_KEY").ok()?;
41        let model = std::env::var("LLM_MODEL")
42            .unwrap_or_else(|_| "gpt-4o".to_string());
43        let base_url = std::env::var("LLM_BASE_URL").ok();
44        let provider_str = std::env::var("LLM_PROVIDER")
45            .unwrap_or_else(|_| "openai".to_string());
46
47        Some(Self {
48            provider: LlmProvider::from_str(&provider_str),
49            api_key,
50            model,
51            base_url,
52        })
53    }
54
55    pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
56        self.base_url = Some(base_url.into());
57        self
58    }
59
60    pub fn build(self) -> Arc<dyn LlmClient> {
61        let base_url = self.base_url;
62        match self.provider {
63            LlmProvider::OpenAi => {
64                let url = base_url.unwrap_or_else(|| "https://api.openai.com/v1".to_string());
65                Arc::new(OpenAiClient::new(self.api_key, self.model, Some(url)))
66            }
67            LlmProvider::Anthropic => {
68                let url = base_url.unwrap_or_else(|| "https://api.anthropic.com".to_string());
69                Arc::new(AnthropicClient::new(self.api_key, self.model, Some(url)))
70            }
71            LlmProvider::Custom(_) => {
72                let url = base_url.unwrap_or_else(|| "https://api.openai.com/v1".to_string());
73                Arc::new(OpenAiClient::new(self.api_key, self.model, Some(url)))
74            }
75        }
76    }
77}
agent_base/llm/registry.rs

agent_base/llm/
registry.rs
