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(
31        provider: LlmProvider,
32        api_key: impl Into<String>,
33        model: impl Into<String>,
34    ) -> Self {
35        Self {
36            provider,
37            api_key: api_key.into(),
38            model: model.into(),
39            base_url: None,
40        }
41    }
42
43    pub fn from_env() -> Option<Self> {
44        let api_key = std::env::var("LLM_API_KEY").ok()?;
45        let model = std::env::var("LLM_MODEL").unwrap_or_else(|_| "gpt-4o".to_string());
46        let base_url = std::env::var("LLM_BASE_URL").ok();
47        let provider_str = std::env::var("LLM_PROVIDER").unwrap_or_else(|_| "openai".to_string());
48
49        Some(Self {
50            provider: LlmProvider::from_str(&provider_str),
51            api_key,
52            model,
53            base_url,
54        })
55    }
56
57    pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
58        self.base_url = Some(base_url.into());
59        self
60    }
61
62    pub fn build(self) -> Arc<dyn LlmClient> {
63        let base_url = self.base_url;
64        match self.provider {
65            LlmProvider::OpenAi => {
66                let url = base_url.unwrap_or_else(|| "https://api.openai.com/v1".to_string());
67                Arc::new(OpenAiClient::new(self.api_key, self.model, Some(url)))
68            }
69            LlmProvider::Anthropic => {
70                let url = base_url.unwrap_or_else(|| "https://api.anthropic.com".to_string());
71                Arc::new(AnthropicClient::new(self.api_key, self.model, Some(url)))
72            }
73            LlmProvider::Custom(_) => {
74                let url = base_url.unwrap_or_else(|| "https://api.openai.com/v1".to_string());
75                Arc::new(OpenAiClient::new(self.api_key, self.model, Some(url)))
76            }
77        }
78    }
79}
agent_base/llm/registry.rs

agent_base/llm/
registry.rs
