client.rs 4.8 KB

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