claude_agent/client/adapter/
foundry.rs1use std::sync::Arc;
4use std::time::Duration;
5
6use async_trait::async_trait;
7use azure_core::credentials::TokenCredential;
8use azure_identity::DeveloperToolsCredential;
9
10use super::base::RequestExecutor;
11use super::config::ProviderConfig;
12use super::request::build_messages_body;
13use super::token_cache::{CachedToken, TokenCache, new_token_cache};
14use super::traits::ProviderAdapter;
15use crate::client::messages::CreateMessageRequest;
16use crate::config::FoundryConfig;
17use crate::types::ApiResponse;
18use crate::{Error, Result};
19
20const ANTHROPIC_VERSION: &str = "2023-06-01";
21const COGNITIVE_SERVICES_SCOPE: &str = "https://cognitiveservices.azure.com/.default";
22
23#[derive(Debug)]
24pub struct FoundryAdapter {
25 config: ProviderConfig,
26 resource_name: Option<String>,
27 base_url: Option<String>,
28 credential: Arc<DeveloperToolsCredential>,
29 api_key: Option<String>,
30 token_cache: TokenCache,
31}
32
33impl FoundryAdapter {
34 pub async fn from_env(config: ProviderConfig) -> Result<Self> {
35 let foundry_config = FoundryConfig::from_env();
36 Self::from_config(config, foundry_config).await
37 }
38
39 pub async fn from_config(config: ProviderConfig, foundry: FoundryConfig) -> Result<Self> {
40 if foundry.resource.is_none() && foundry.base_url.is_none() {
41 return Err(Error::auth(
42 "Either ANTHROPIC_FOUNDRY_RESOURCE or ANTHROPIC_FOUNDRY_BASE_URL must be set",
43 ));
44 }
45
46 let credential = DeveloperToolsCredential::new(None)
47 .map_err(|e| Error::auth(format!("Failed to create Azure credential: {}", e)))?;
48
49 Ok(Self {
50 config,
51 resource_name: foundry.resource,
52 base_url: foundry.base_url,
53 credential,
54 api_key: foundry.api_key,
55 token_cache: new_token_cache(),
56 })
57 }
58
59 pub fn with_resource(mut self, resource: impl Into<String>) -> Self {
61 self.resource_name = Some(resource.into());
62 self
63 }
64
65 pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
67 self.base_url = Some(base_url.into());
68 self
69 }
70
71 pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
73 self.api_key = Some(key.into());
74 self
75 }
76
77 fn build_messages_url(&self) -> String {
78 if let Some(ref base_url) = self.base_url {
79 let base = base_url.trim_end_matches('/');
80 format!("{}/v1/messages", base)
81 } else if let Some(ref resource) = self.resource_name {
82 format!(
83 "https://{}.services.ai.azure.com/anthropic/v1/messages",
84 resource
85 )
86 } else {
87 unreachable!("validated in from_config")
88 }
89 }
90
91 fn build_request_body(&self, request: &CreateMessageRequest) -> serde_json::Value {
92 build_messages_body(request, None, self.config.thinking_budget)
93 }
94
95 async fn get_auth_header(&self) -> Result<(String, String)> {
96 if let Some(ref api_key) = self.api_key {
97 return Ok(("api-key".into(), api_key.clone()));
98 }
99 let token = self.get_token().await?;
100 Ok(("Authorization".into(), format!("Bearer {}", token)))
101 }
102
103 async fn get_token(&self) -> Result<String> {
104 {
105 let cache = self.token_cache.read().await;
106 if let Some(ref cached) = *cache
107 && !cached.is_expired()
108 {
109 return Ok(cached.token().to_string());
110 }
111 }
112
113 let token_response = self
114 .credential
115 .get_token(&[COGNITIVE_SERVICES_SCOPE], None)
116 .await
117 .map_err(|e| Error::auth(format!("Failed to get Azure token: {}", e)))?;
118
119 let token_str = token_response.token.secret().to_string();
120 let cached = CachedToken::new(token_str.clone(), Duration::from_secs(3600));
121 *self.token_cache.write().await = Some(cached);
122
123 Ok(token_str)
124 }
125
126 async fn execute_request(
127 &self,
128 http: &reqwest::Client,
129 url: &str,
130 body: &serde_json::Value,
131 ) -> Result<reqwest::Response> {
132 let (header_name, header_value) = self.get_auth_header().await?;
133 let headers = vec![
134 (header_name, header_value),
135 ("anthropic-version".into(), ANTHROPIC_VERSION.into()),
136 ];
137 RequestExecutor::post(http, url, body, headers).await
138 }
139}
140
141#[async_trait]
142impl ProviderAdapter for FoundryAdapter {
143 fn config(&self) -> &ProviderConfig {
144 &self.config
145 }
146
147 fn name(&self) -> &'static str {
148 "foundry"
149 }
150
151 async fn build_url(&self, _model: &str, _stream: bool) -> String {
152 self.build_messages_url()
153 }
154
155 async fn transform_request(&self, request: CreateMessageRequest) -> Result<serde_json::Value> {
156 Ok(self.build_request_body(&request))
157 }
158
159 fn transform_response(&self, response: serde_json::Value) -> Result<ApiResponse> {
160 serde_json::from_value(response).map_err(|e| Error::Parse(e.to_string()))
161 }
162
163 async fn send(
164 &self,
165 http: &reqwest::Client,
166 request: CreateMessageRequest,
167 ) -> Result<ApiResponse> {
168 let url = self.build_messages_url();
169 let body = self.build_request_body(&request);
170 let response = self.execute_request(http, &url, &body).await?;
171 let json: serde_json::Value = response.json().await?;
172 self.transform_response(json)
173 }
174
175 async fn send_stream(
176 &self,
177 http: &reqwest::Client,
178 mut request: CreateMessageRequest,
179 ) -> Result<reqwest::Response> {
180 request.stream = Some(true);
181 let url = self.build_messages_url();
182 let body = self.build_request_body(&request);
183 self.execute_request(http, &url, &body).await
184 }
185
186 async fn refresh_credentials(&self) -> Result<()> {
187 *self.token_cache.write().await = None;
188 if self.api_key.is_none() {
189 self.get_token().await?;
190 }
191 Ok(())
192 }
193}
194
195#[cfg(test)]
196mod tests {
197 use crate::client::adapter::ModelConfig;
198
199 #[test]
200 fn test_build_url_with_resource() {
201 let url = format!(
202 "https://{}.services.ai.azure.com/anthropic/v1/messages",
203 "my-resource"
204 );
205 assert!(url.contains("services.ai.azure.com"));
206 assert!(url.contains("/anthropic/v1/messages"));
207 }
208
209 #[test]
210 fn test_build_url_with_base_url() {
211 let base_url = "https://custom-endpoint.azure.com/anthropic";
212 let url = format!("{}/v1/messages", base_url.trim_end_matches('/'));
213 assert!(url.contains("custom-endpoint.azure.com"));
214 assert!(url.contains("/v1/messages"));
215 }
216
217 #[test]
218 fn test_model_config() {
219 let config = ModelConfig::foundry();
220 assert!(config.primary.contains("sonnet"));
221 }
222
223 #[test]
224 fn test_anthropic_version() {
225 assert_eq!(super::ANTHROPIC_VERSION, "2023-06-01");
226 }
227}