graphify_extract/semantic/
provider.rs

1use anyhow::{Context, Result};
2
3#[derive(Debug, Clone, PartialEq, Eq)]
4pub enum LLMProvider {
5    Anthropic,
6    OpenAI,
7    Ollama,
8    OpenAICompatible,
9}
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum AuthType {
13    ApiKey,
14    Bearer,
15}
16
17/// Raw LLM configuration from `graphify.toml`'s `[llm]` section.
18/// Used to avoid a 9-parameter resolve() signature.
19#[derive(Debug, Default)]
20pub struct LLMConfigRaw {
21    pub provider: String,
22    pub model: String,
23    pub anthropic_api_key: Option<String>,
24    pub anthropic_base_url: Option<String>,
25    pub openai_api_key: Option<String>,
26    pub openai_base_url: Option<String>,
27    pub ollama_base_url: Option<String>,
28    pub openai_compatible_api_key: Option<String>,
29    pub openai_compatible_base_url: Option<String>,
30}
31
32#[derive(Debug, Clone)]
33pub struct LLMProviderConfig {
34    pub provider: LLMProvider,
35    pub model: String,
36    pub api_key: Option<String>,
37    pub base_url: String,
38    pub auth_type: AuthType,
39}
40
41impl LLMProviderConfig {
42    const ANTHROPIC_DEFAULT_URL: &str = "https://api.anthropic.com";
43    const OPENAI_DEFAULT_URL: &str = "https://api.openai.com/v1";
44    const OLLAMA_DEFAULT_URL: &str = "http://localhost:11434/v1";
45
46    pub fn resolve(raw: &LLMConfigRaw) -> Result<Self> {
47        let provider = match raw.provider.as_str() {
48            "anthropic" => LLMProvider::Anthropic,
49            "openai" => LLMProvider::OpenAI,
50            "ollama" => LLMProvider::Ollama,
51            "openai_compatible" => LLMProvider::OpenAICompatible,
52            other => anyhow::bail!(
53                "Unknown LLM provider: '{other}'. Supported: anthropic, openai, ollama, openai_compatible"
54            ),
55        };
56
57        if raw.model.is_empty() {
58            anyhow::bail!("LLM model is required in [llm] config");
59        }
60
61        let (api_key, base_url, auth_type) = match provider {
62            LLMProvider::Anthropic => {
63                let (key, at) = if let Some(ref k) = raw.anthropic_api_key {
64                    (Some(k.clone()), AuthType::ApiKey)
65                } else if let Ok(k) = std::env::var("ANTHROPIC_API_KEY") {
66                    (Some(k), AuthType::ApiKey)
67                } else if let Some(token) = super::anthropic_oauth::read_claude_code_oauth_token() {
68                    (Some(token), AuthType::Bearer)
69                } else {
70                    (None, AuthType::ApiKey)
71                };
72                let url = raw
73                    .anthropic_base_url
74                    .clone()
75                    .unwrap_or_else(|| Self::ANTHROPIC_DEFAULT_URL.to_string());
76                (key, url, at)
77            }
78            LLMProvider::OpenAI => {
79                let key = raw
80                    .openai_api_key
81                    .clone()
82                    .or_else(|| std::env::var("OPENAI_API_KEY").ok());
83                let url = raw
84                    .openai_base_url
85                    .clone()
86                    .unwrap_or_else(|| Self::OPENAI_DEFAULT_URL.to_string());
87                (key, url, AuthType::Bearer)
88            }
89            LLMProvider::Ollama => {
90                let url = raw
91                    .ollama_base_url
92                    .clone()
93                    .unwrap_or_else(|| Self::OLLAMA_DEFAULT_URL.to_string());
94                (None, url, AuthType::Bearer)
95            }
96            LLMProvider::OpenAICompatible => {
97                let key = raw.openai_compatible_api_key.clone();
98                let url = raw.openai_compatible_base_url.clone().context(
99                    "openai_compatible_base_url is required for openai_compatible provider",
100                )?;
101                (key, url, AuthType::Bearer)
102            }
103        };
104
105        Ok(LLMProviderConfig {
106            provider,
107            model: raw.model.clone(),
108            api_key,
109            base_url,
110            auth_type,
111        })
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    fn raw(provider: &str, model: &str) -> LLMConfigRaw {
120        LLMConfigRaw {
121            provider: provider.to_string(),
122            model: model.to_string(),
123            ..Default::default()
124        }
125    }
126
127    #[test]
128    fn resolve_anthropic_with_api_key() {
129        let r = LLMConfigRaw {
130            provider: "anthropic".into(),
131            model: "claude-sonnet-4.6".into(),
132            anthropic_api_key: Some("sk-test-key".into()),
133            ..Default::default()
134        };
135        let config = LLMProviderConfig::resolve(&r).unwrap();
136        assert_eq!(config.provider, LLMProvider::Anthropic);
137        assert_eq!(config.model, "claude-sonnet-4.6");
138        assert_eq!(config.api_key.as_deref(), Some("sk-test-key"));
139        assert_eq!(config.base_url, "https://api.anthropic.com");
140        assert_eq!(config.auth_type, AuthType::ApiKey);
141    }
142
143    #[test]
144    fn resolve_openai_with_base_url_override() {
145        let r = LLMConfigRaw {
146            provider: "openai".into(),
147            model: "gpt-4o".into(),
148            openai_api_key: Some("sk-openai-key".into()),
149            openai_base_url: Some("https://custom.api.com/v1".into()),
150            ..Default::default()
151        };
152        let config = LLMProviderConfig::resolve(&r).unwrap();
153        assert_eq!(config.provider, LLMProvider::OpenAI);
154        assert_eq!(config.base_url, "https://custom.api.com/v1");
155        assert_eq!(config.auth_type, AuthType::Bearer);
156    }
157
158    #[test]
159    fn resolve_ollama_defaults() {
160        let config = LLMProviderConfig::resolve(&raw("ollama", "llama3")).unwrap();
161        assert_eq!(config.provider, LLMProvider::Ollama);
162        assert_eq!(config.base_url, "http://localhost:11434/v1");
163        assert!(config.api_key.is_none());
164    }
165
166    #[test]
167    fn resolve_openai_compatible_requires_base_url() {
168        let result = LLMProviderConfig::resolve(&raw("openai_compatible", "my-model"));
169        assert!(result.is_err());
170        assert!(
171            result
172                .unwrap_err()
173                .to_string()
174                .contains("openai_compatible_base_url")
175        );
176    }
177
178    #[test]
179    fn resolve_openai_compatible_with_base_url() {
180        let r = LLMConfigRaw {
181            provider: "openai_compatible".into(),
182            model: "my-model".into(),
183            openai_compatible_api_key: Some("optional-key".into()),
184            openai_compatible_base_url: Some("http://localhost:8000/v1".into()),
185            ..Default::default()
186        };
187        let config = LLMProviderConfig::resolve(&r).unwrap();
188        assert_eq!(config.provider, LLMProvider::OpenAICompatible);
189        assert_eq!(config.base_url, "http://localhost:8000/v1");
190        assert_eq!(config.api_key.as_deref(), Some("optional-key"));
191    }
192
193    #[test]
194    fn reject_unknown_provider() {
195        let result = LLMProviderConfig::resolve(&raw("unknown", "model"));
196        assert!(result.is_err());
197        assert!(
198            result
199                .unwrap_err()
200                .to_string()
201                .contains("Unknown LLM provider")
202        );
203    }
204
205    #[test]
206    fn reject_empty_model() {
207        let result = LLMProviderConfig::resolve(&raw("anthropic", ""));
208        assert!(result.is_err());
209        assert!(
210            result
211                .unwrap_err()
212                .to_string()
213                .contains("model is required")
214        );
215    }
216}
graphify_extract/semantic/provider.rs

graphify_extract/semantic/
provider.rs
