Skip to main content

cortexai_llm_client/
provider.rs

1//! LLM Provider definitions.
2
3use serde::{Deserialize, Serialize};
4use std::fmt;
5
6/// Supported LLM providers.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
8#[serde(rename_all = "lowercase")]
9pub enum Provider {
10    /// OpenAI (GPT-4, GPT-3.5, etc.)
11    OpenAI,
12    /// Anthropic (Claude 3, etc.)
13    Anthropic,
14    /// OpenRouter (100+ models)
15    OpenRouter,
16}
17
18impl Provider {
19    /// Get the API endpoint URL for this provider.
20    pub fn endpoint(&self) -> &'static str {
21        match self {
22            Self::OpenAI => "https://api.openai.com/v1/chat/completions",
23            Self::Anthropic => "https://api.anthropic.com/v1/messages",
24            Self::OpenRouter => "https://openrouter.ai/api/v1/chat/completions",
25        }
26    }
27
28    /// Get the header name for the API key.
29    pub fn auth_header(&self) -> &'static str {
30        match self {
31            Self::OpenAI | Self::OpenRouter => "Authorization",
32            Self::Anthropic => "x-api-key",
33        }
34    }
35
36    /// Format the API key for the authorization header.
37    pub fn format_auth(&self, api_key: &str) -> String {
38        match self {
39            Self::OpenAI | Self::OpenRouter => format!("Bearer {}", api_key),
40            Self::Anthropic => api_key.to_string(),
41        }
42    }
43
44    /// Get additional required headers for this provider.
45    pub fn extra_headers(&self) -> Vec<(&'static str, &'static str)> {
46        match self {
47            Self::OpenAI => vec![],
48            Self::Anthropic => vec![
49                ("anthropic-version", "2023-06-01"),
50                ("anthropic-dangerous-direct-browser-access", "true"),
51            ],
52            Self::OpenRouter => vec![("HTTP-Referer", "https://cortex.dev")],
53        }
54    }
55
56    /// Check if this provider uses OpenAI-compatible format.
57    pub fn is_openai_compatible(&self) -> bool {
58        matches!(self, Self::OpenAI | Self::OpenRouter)
59    }
60}
61
62impl fmt::Display for Provider {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        match self {
65            Self::OpenAI => write!(f, "openai"),
66            Self::Anthropic => write!(f, "anthropic"),
67            Self::OpenRouter => write!(f, "openrouter"),
68        }
69    }
70}
71
72impl std::str::FromStr for Provider {
73    type Err = crate::LlmClientError;
74
75    fn from_str(s: &str) -> Result<Self, Self::Err> {
76        match s.to_lowercase().as_str() {
77            "openai" => Ok(Self::OpenAI),
78            "anthropic" => Ok(Self::Anthropic),
79            "openrouter" => Ok(Self::OpenRouter),
80            _ => Err(crate::LlmClientError::UnknownProvider(s.to_string())),
81        }
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[test]
90    fn test_provider_endpoints() {
91        assert!(Provider::OpenAI.endpoint().contains("openai.com"));
92        assert!(Provider::Anthropic.endpoint().contains("anthropic.com"));
93        assert!(Provider::OpenRouter.endpoint().contains("openrouter.ai"));
94    }
95
96    #[test]
97    fn test_provider_from_str() {
98        assert_eq!("openai".parse::<Provider>().unwrap(), Provider::OpenAI);
99        assert_eq!(
100            "ANTHROPIC".parse::<Provider>().unwrap(),
101            Provider::Anthropic
102        );
103        assert!("invalid".parse::<Provider>().is_err());
104    }
105
106    #[test]
107    fn test_auth_format() {
108        assert!(Provider::OpenAI.format_auth("key").starts_with("Bearer "));
109        assert_eq!(Provider::Anthropic.format_auth("key"), "key");
110    }
111}