Skip to main content

rig/providers/anthropic/
client.rs

1//! Anthropic client api implementation
2use http::{HeaderName, HeaderValue};
3
4use super::completion::{ANTHROPIC_VERSION_LATEST, CompletionModel};
5use crate::{
6    client::{
7        self, ApiKey, Capabilities, Capable, DebugExt, Nothing, Provider, ProviderBuilder,
8        ProviderClient,
9    },
10    http_client::{self, HttpClientExt},
11    providers::anthropic::model_listing::AnthropicModelLister,
12};
13
14// ================================================================
15// Main Anthropic Client
16// ================================================================
17#[derive(Debug, Default, Clone)]
18pub struct AnthropicExt;
19
20impl Provider for AnthropicExt {
21    type Builder = AnthropicBuilder;
22    const VERIFY_PATH: &'static str = "/v1/models";
23}
24
25impl<H> Capabilities<H> for AnthropicExt {
26    type Completion = Capable<CompletionModel<H>>;
27
28    type Embeddings = Nothing;
29    type Transcription = Nothing;
30    type ModelListing = Capable<AnthropicModelLister<H>>;
31    #[cfg(feature = "image")]
32    type ImageGeneration = Nothing;
33    #[cfg(feature = "audio")]
34    type AudioGeneration = Nothing;
35}
36
37#[derive(Debug, Clone)]
38pub struct AnthropicBuilder {
39    anthropic_version: String,
40    anthropic_betas: Vec<String>,
41}
42
43#[derive(Debug, Clone)]
44pub struct AnthropicKey(String);
45
46impl<S> From<S> for AnthropicKey
47where
48    S: Into<String>,
49{
50    fn from(value: S) -> Self {
51        Self(value.into())
52    }
53}
54
55impl ApiKey for AnthropicKey {
56    fn into_header(self) -> Option<http_client::Result<(http::HeaderName, HeaderValue)>> {
57        Some(
58            HeaderValue::from_str(&self.0)
59                .map(|val| (HeaderName::from_static("x-api-key"), val))
60                .map_err(Into::into),
61        )
62    }
63}
64
65pub type Client<H = reqwest::Client> = client::Client<AnthropicExt, H>;
66pub type ClientBuilder<H = reqwest::Client> =
67    client::ClientBuilder<AnthropicBuilder, AnthropicKey, H>;
68
69impl Default for AnthropicBuilder {
70    fn default() -> Self {
71        Self {
72            anthropic_version: ANTHROPIC_VERSION_LATEST.into(),
73            anthropic_betas: Vec::new(),
74        }
75    }
76}
77
78impl ProviderBuilder for AnthropicBuilder {
79    type Extension<H>
80        = AnthropicExt
81    where
82        H: HttpClientExt;
83    type ApiKey = AnthropicKey;
84
85    const BASE_URL: &'static str = "https://api.anthropic.com";
86
87    fn build<H>(
88        _builder: &client::ClientBuilder<Self, Self::ApiKey, H>,
89    ) -> http_client::Result<Self::Extension<H>>
90    where
91        H: HttpClientExt,
92    {
93        Ok(AnthropicExt)
94    }
95
96    fn finish<H>(
97        &self,
98        mut builder: client::ClientBuilder<Self, AnthropicKey, H>,
99    ) -> http_client::Result<client::ClientBuilder<Self, AnthropicKey, H>> {
100        builder.headers_mut().insert(
101            "anthropic-version",
102            HeaderValue::from_str(&self.anthropic_version)?,
103        );
104
105        if !self.anthropic_betas.is_empty() {
106            builder.headers_mut().insert(
107                "anthropic-beta",
108                HeaderValue::from_str(&self.anthropic_betas.join(","))?,
109            );
110        }
111
112        Ok(builder)
113    }
114}
115
116impl DebugExt for AnthropicExt {}
117
118impl ProviderClient for Client {
119    type Input = String;
120
121    fn from_env() -> Self
122    where
123        Self: Sized,
124    {
125        let key = std::env::var("ANTHROPIC_API_KEY").expect("ANTHROPIC_API_KEY not set");
126
127        Self::builder().api_key(key).build().unwrap()
128    }
129
130    fn from_val(input: Self::Input) -> Self
131    where
132        Self: Sized,
133    {
134        Self::builder().api_key(input).build().unwrap()
135    }
136}
137
138/// Create a new anthropic client using the builder
139///
140/// # Example
141/// ```
142/// use rig::providers::anthropic::{ClientBuilder, self};
143///
144/// // Initialize the Anthropic client
145/// let anthropic_client = ClientBuilder::new("your-claude-api-key")
146///    .anthropic_version(ANTHROPIC_VERSION_LATEST)
147///    .anthropic_beta("prompt-caching-2024-07-31")
148///    .build()
149/// ```
150impl<H> ClientBuilder<H> {
151    pub fn anthropic_version(self, anthropic_version: &str) -> Self {
152        self.over_ext(|ext| AnthropicBuilder {
153            anthropic_version: anthropic_version.into(),
154            ..ext
155        })
156    }
157
158    pub fn anthropic_betas(self, anthropic_betas: &[&str]) -> Self {
159        self.over_ext(|mut ext| {
160            ext.anthropic_betas
161                .extend(anthropic_betas.iter().copied().map(String::from));
162
163            ext
164        })
165    }
166
167    pub fn anthropic_beta(self, anthropic_beta: &str) -> Self {
168        self.over_ext(|mut ext| {
169            ext.anthropic_betas.push(anthropic_beta.into());
170
171            ext
172        })
173    }
174}
175#[cfg(test)]
176mod tests {
177    #[test]
178    fn test_client_initialization() {
179        let _client =
180            crate::providers::anthropic::Client::new("dummy-key").expect("Client::new() failed");
181        let _client_from_builder = crate::providers::anthropic::Client::builder()
182            .api_key("dummy-key")
183            .build()
184            .expect("Client::builder() failed");
185    }
186}