wraith_api/
client.rs

1use crate::error::ApiError;
2use crate::providers::anthropic::{self, AuthSource, AnthropicClient};
3use crate::providers::openai_compat::{self, OpenAiCompatClient, OpenAiCompatConfig};
4use crate::providers::{self, Provider, ProviderKind};
5use crate::types::{MessageRequest, MessageResponse, StreamEvent};
6
7async fn send_via_provider<P: Provider>(
8    provider: &P,
9    request: &MessageRequest,
10) -> Result<MessageResponse, ApiError> {
11    provider.send_message(request).await
12}
13
14async fn stream_via_provider<P: Provider>(
15    provider: &P,
16    request: &MessageRequest,
17) -> Result<P::Stream, ApiError> {
18    provider.stream_message(request).await
19}
20
21#[derive(Debug, Clone)]
22pub enum ProviderClient {
23    Anthropic(AnthropicClient),
24    Xai(OpenAiCompatClient),
25    OpenAi(OpenAiCompatClient),
26    Gemini(OpenAiCompatClient),
27    OpenRouter(OpenAiCompatClient),
28}
29
30impl ProviderClient {
31    pub fn from_model(model: &str) -> Result<Self, ApiError> {
32        Self::from_model_with_default_auth(model, None)
33    }
34
35    pub fn from_model_with_default_auth(
36        model: &str,
37        default_auth: Option<AuthSource>,
38    ) -> Result<Self, ApiError> {
39        let resolved_model = providers::resolve_model_alias(model);
40        match providers::detect_provider_kind(&resolved_model) {
41            ProviderKind::Anthropic => Ok(Self::Anthropic(match default_auth {
42                Some(auth) => AnthropicClient::from_auth(auth),
43                None => AnthropicClient::from_env()?,
44            })),
45            ProviderKind::Xai => Ok(Self::Xai(OpenAiCompatClient::from_env(
46                OpenAiCompatConfig::xai(),
47            )?)),
48            ProviderKind::OpenAi => Ok(Self::OpenAi(OpenAiCompatClient::from_env(
49                OpenAiCompatConfig::openai(),
50            )?)),
51            ProviderKind::Gemini => Ok(Self::Gemini(OpenAiCompatClient::from_env(
52                OpenAiCompatConfig::gemini(),
53            )?)),
54            ProviderKind::OpenRouter => Ok(Self::OpenRouter(OpenAiCompatClient::from_env(
55                OpenAiCompatConfig::openrouter(),
56            )?)),
57        }
58    }
59
60    #[must_use]
61    pub const fn provider_kind(&self) -> ProviderKind {
62        match self {
63            Self::Anthropic(_) => ProviderKind::Anthropic,
64            Self::Xai(_) => ProviderKind::Xai,
65            Self::OpenAi(_) => ProviderKind::OpenAi,
66            Self::Gemini(_) => ProviderKind::Gemini,
67            Self::OpenRouter(_) => ProviderKind::OpenRouter,
68        }
69    }
70
71    pub async fn send_message(
72        &self,
73        request: &MessageRequest,
74    ) -> Result<MessageResponse, ApiError> {
75        match self {
76            Self::Anthropic(client) => send_via_provider(client, request).await,
77            Self::Xai(client) | Self::OpenAi(client) | Self::Gemini(client) | Self::OpenRouter(client) => {
78                send_via_provider(client, request).await
79            }
80        }
81    }
82
83    pub async fn stream_message(
84        &self,
85        request: &MessageRequest,
86    ) -> Result<MessageStream, ApiError> {
87        match self {
88            Self::Anthropic(client) => stream_via_provider(client, request)
89                .await
90                .map(MessageStream::Anthropic),
91            Self::Xai(client) | Self::OpenAi(client) | Self::Gemini(client) | Self::OpenRouter(client) => {
92                stream_via_provider(client, request)
93                    .await
94                    .map(MessageStream::OpenAiCompat)
95            }
96        }
97    }
98}
99
100#[derive(Debug)]
101pub enum MessageStream {
102    Anthropic(anthropic::MessageStream),
103    OpenAiCompat(openai_compat::MessageStream),
104}
105
106impl MessageStream {
107    #[must_use]
108    pub fn request_id(&self) -> Option<&str> {
109        match self {
110            Self::Anthropic(stream) => stream.request_id(),
111            Self::OpenAiCompat(stream) => stream.request_id(),
112        }
113    }
114
115    pub async fn next_event(&mut self) -> Result<Option<StreamEvent>, ApiError> {
116        match self {
117            Self::Anthropic(stream) => stream.next_event().await,
118            Self::OpenAiCompat(stream) => stream.next_event().await,
119        }
120    }
121}
122
123pub use anthropic::{
124    oauth_token_is_expired, resolve_saved_oauth_token, resolve_startup_auth_source, OAuthTokenSet,
125};
126#[must_use]
127pub fn read_base_url() -> String {
128    anthropic::read_base_url()
129}
130
131#[must_use]
132pub fn read_xai_base_url() -> String {
133    openai_compat::read_base_url(OpenAiCompatConfig::xai())
134}
135
136#[cfg(test)]
137mod tests {
138    use crate::providers::{detect_provider_kind, resolve_model_alias, ProviderKind};
139
140    #[test]
141    fn resolves_existing_and_grok_aliases() {
142        assert_eq!(resolve_model_alias("opus"), "claude-opus-4-6");
143        assert_eq!(resolve_model_alias("grok"), "grok-3");
144        assert_eq!(resolve_model_alias("grok-mini"), "grok-3-mini");
145    }
146
147    #[test]
148    fn provider_detection_prefers_model_family() {
149        assert_eq!(detect_provider_kind("grok-3"), ProviderKind::Xai);
150        assert_eq!(
151            detect_provider_kind("claude-sonnet-4-6"),
152            ProviderKind::Anthropic
153        );
154    }
155}
wraith_api/client.rs

wraith_api/
client.rs
