agent_framework_anthropic/
vertex.rs

1//! [`AnthropicVertexClient`]: Anthropic (Claude) models hosted on
2//! [Google Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude),
3//! spoken over Vertex's `rawPredict` / `streamRawPredict` publisher-model
4//! routes
5//! (`POST https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/publishers/anthropic/models/{model}:rawPredict`).
6//!
7//! Like [`crate::bedrock::AnthropicBedrockClient`], this is a *transport*, not
8//! a new wire format: the request/response body is the same Anthropic
9//! Messages API shape [`AnthropicClient`](crate::AnthropicClient) speaks
10//! directly, built via [`crate::convert::build_cloud_request`] (which carries
11//! `anthropic_version: "vertex-2023-10-16"` instead of a top-level `model`
12//! field, since the model is already selected by the URL).
13//!
14//! ## Authentication
15//!
16//! Vertex AI authenticates with a Google OAuth2 access token
17//! (`Authorization: Bearer <token>`) minted via Google's
18//! [Application Default Credentials](https://cloud.google.com/docs/authentication/application-default-credentials)
19//! (ADC) flow. This workspace has no Google Cloud SDK dependency to perform
20//! that flow, so instead of reimplementing it, token acquisition is factored
21//! out behind the small, synchronous [`VertexTokenProvider`] trait: the
22//! caller supplies a token however they like (most simply, the output of
23//! `gcloud auth print-access-token`, wrapped in [`StaticVertexToken`]).
24//! Wiring up full ADC auto-discovery (metadata-server tokens on GCE/GKE,
25//! service-account JSON key files, `gcloud`'s cached ADC token, …) behind a
26//! richer provider implementation is a documented extension point, not
27//! attempted here.
28//!
29//! ```no_run
30//! use std::sync::Arc;
31//! use agent_framework_anthropic::vertex::{AnthropicVertexClient, StaticVertexToken};
32//! use agent_framework_core::prelude::*;
33//!
34//! # async fn demo() -> Result<()> {
35//! // `$ gcloud auth print-access-token`
36//! let token = std::env::var("VERTEX_ACCESS_TOKEN").unwrap_or_default();
37//! let client = AnthropicVertexClient::new(
38//!     "my-gcp-project",
39//!     "us-east5",
40//!     "claude-sonnet-4-5@20250929",
41//!     Arc::new(StaticVertexToken::new(token)),
42//! );
43//! let agent = Agent::builder(client)
44//!     .instructions("You are concise.")
45//!     .build();
46//! let reply = agent.run_once("Say hi").await?;
47//! println!("{}", reply.text());
48//! # Ok(())
49//! # }
50//! ```
51
52use std::sync::Arc;
53
54use agent_framework_core::client::{ChatClient, ChatStream};
55use agent_framework_core::error::{Error, Result};
56use agent_framework_core::types::{
57    ChatOptions, ChatResponse, ChatResponseUpdate, Content, Message, Role, UsageContent,
58};
59use futures::stream::{self, StreamExt};
60use serde_json::Value;
61
62use crate::convert;
63
64/// The `anthropic_version` Vertex AI's `rawPredict`/`streamRawPredict` routes
65/// expect in the body of every Claude request (per Google's
66/// [Claude-on-Vertex documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude)).
67pub const ANTHROPIC_VERTEX_VERSION: &str = "vertex-2023-10-16";
68
69/// `max_tokens` is required by the Anthropic Messages API; used whenever
70/// neither `ChatOptions::max_tokens` nor a client-level override is set.
71/// Matches [`crate::AnthropicClient`]'s default.
72const DEFAULT_MAX_TOKENS: u32 = 1024;
73
74/// The region [`AnthropicVertexClient::from_env`] falls back to when neither
75/// `CLOUD_ML_REGION` nor `ANTHROPIC_VERTEX_REGION` is set — Anthropic's
76/// models on Vertex AI are only available in a handful of regions, and
77/// `us-east5` is the one Google's own quickstart documentation defaults to.
78const DEFAULT_REGION: &str = "us-east5";
79
80/// Supplies a Google OAuth2 access token to authenticate Vertex AI requests
81/// (`Authorization: Bearer <token>`).
82///
83/// Deliberately synchronous and minimal — see the [module docs](self) for why
84/// this crate doesn't perform Google's Application Default Credentials flow
85/// itself. Implement this to plug in real token acquisition/refresh (a
86/// metadata-server client, a cached `gcloud` token, a service-account JWT
87/// exchange, …); [`StaticVertexToken`] covers the common case of a
88/// caller-supplied, pre-fetched token.
89pub trait VertexTokenProvider: Send + Sync {
90    /// Return the current access token to send as `Authorization: Bearer
91    /// <token>`. Implementations that cache/refresh a token should do so
92    /// internally (e.g. behind a `Mutex`); this is called once per request.
93    fn access_token(&self) -> Result<String>;
94}
95
96/// A [`VertexTokenProvider`] that always returns the same, pre-fetched token
97/// — e.g. the output of `gcloud auth print-access-token`.
98#[derive(Debug, Clone)]
99pub struct StaticVertexToken(String);
100
101impl StaticVertexToken {
102    /// Wrap a fixed access token.
103    pub fn new(token: impl Into<String>) -> Self {
104        Self(token.into())
105    }
106}
107
108impl VertexTokenProvider for StaticVertexToken {
109    fn access_token(&self) -> Result<String> {
110        Ok(self.0.clone())
111    }
112}
113
114/// Classify a non-success Vertex AI HTTP response into a granular [`Error`].
115/// Google's Vertex AI error bodies use a `{"error": {"code": ..., "status":
116/// "PERMISSION_DENIED" | "INVALID_ARGUMENT" | ...}}` shape, but the HTTP
117/// status alone is sufficient and unambiguous for the classification this
118/// crate's other clients perform: `401`/`403` -> invalid auth, `400` ->
119/// invalid request, everything else (notably `429` and `5xx`) stays a
120/// generic, retry-layer-visible [`Error::ServiceStatus`].
121fn classify_vertex_error(status: u16, message: impl Into<String>) -> Error {
122    let message = message.into();
123    match status {
124        401 | 403 => Error::service_invalid_auth(message),
125        400 => Error::service_invalid_request(message),
126        _ => Error::service_status(status, message, None),
127    }
128}
129
130/// An Anthropic Messages API transport for Claude models on Google Vertex AI
131/// (`POST https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/publishers/anthropic/models/{model}:rawPredict`).
132///
133/// See the [module docs](self) for how this relates to
134/// [`AnthropicClient`](crate::AnthropicClient) and for the authentication
135/// model.
136#[derive(Clone)]
137pub struct AnthropicVertexClient {
138    inner: Arc<Inner>,
139}
140
141#[derive(Clone)]
142struct Inner {
143    http: reqwest::Client,
144    project_id: String,
145    region: String,
146    model: String,
147    token_provider: Arc<dyn VertexTokenProvider>,
148    max_tokens: u32,
149}
150
151impl std::fmt::Debug for AnthropicVertexClient {
152    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153        f.debug_struct("AnthropicVertexClient")
154            .field("project_id", &self.inner.project_id)
155            .field("region", &self.inner.region)
156            .field("model", &self.inner.model)
157            .field("max_tokens", &self.inner.max_tokens)
158            .finish_non_exhaustive()
159    }
160}
161
162impl AnthropicVertexClient {
163    /// Create a client for the given GCP project, region, default model id,
164    /// and token provider.
165    pub fn new(
166        project_id: impl Into<String>,
167        region: impl Into<String>,
168        model: impl Into<String>,
169        token_provider: Arc<dyn VertexTokenProvider>,
170    ) -> Self {
171        Self {
172            inner: Arc::new(Inner {
173                http: reqwest::Client::new(),
174                project_id: project_id.into(),
175                region: region.into(),
176                model: model.into(),
177                token_provider,
178                max_tokens: DEFAULT_MAX_TOKENS,
179            }),
180        }
181    }
182
183    /// Build a client from `GOOGLE_CLOUD_PROJECT` (falling back to
184    /// `ANTHROPIC_VERTEX_PROJECT_ID`) and `CLOUD_ML_REGION` (falling back to
185    /// `ANTHROPIC_VERTEX_REGION`, defaulting to `us-east5` when neither is
186    /// set).
187    pub fn from_env(
188        model: impl Into<String>,
189        token_provider: Arc<dyn VertexTokenProvider>,
190    ) -> Result<Self> {
191        let project_id = std::env::var("GOOGLE_CLOUD_PROJECT")
192            .or_else(|_| std::env::var("ANTHROPIC_VERTEX_PROJECT_ID"))
193            .map_err(|_| {
194                Error::Configuration(
195                    "neither GOOGLE_CLOUD_PROJECT nor ANTHROPIC_VERTEX_PROJECT_ID is set".into(),
196                )
197            })?;
198        let region = std::env::var("CLOUD_ML_REGION")
199            .or_else(|_| std::env::var("ANTHROPIC_VERTEX_REGION"))
200            .unwrap_or_else(|_| DEFAULT_REGION.to_string());
201        Ok(Self::new(project_id, region, model, token_provider))
202    }
203
204    /// Override the default `max_tokens` sent when `ChatOptions::max_tokens`
205    /// is unset (the Anthropic Messages API requires this field).
206    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
207        Arc::make_mut(&mut self.inner).max_tokens = max_tokens;
208        self
209    }
210
211    /// The default model id.
212    pub fn model(&self) -> &str {
213        &self.inner.model
214    }
215
216    /// The configured GCP project id.
217    pub fn project_id(&self) -> &str {
218        &self.inner.project_id
219    }
220
221    /// The configured Vertex AI region.
222    pub fn region(&self) -> &str {
223        &self.inner.region
224    }
225
226    fn effective_model(&self, options: &ChatOptions) -> String {
227        options
228            .model
229            .clone()
230            .unwrap_or_else(|| self.inner.model.clone())
231    }
232
233    /// The `rawPredict` (or `streamRawPredict`) URL for a given model id.
234    fn url(&self, model: &str, streaming: bool) -> String {
235        let action = if streaming {
236            "streamRawPredict"
237        } else {
238            "rawPredict"
239        };
240        format!(
241            "https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/publishers/anthropic/models/{model}:{action}",
242            region = self.inner.region,
243            project = self.inner.project_id,
244        )
245    }
246
247    async fn send(&self, model: &str, body: &Value) -> Result<reqwest::Response> {
248        let token = self.inner.token_provider.access_token()?;
249        let resp = self
250            .inner
251            .http
252            .post(self.url(model, false))
253            .bearer_auth(token)
254            .json(body)
255            .send()
256            .await
257            .map_err(|e| Error::service(format!("request failed: {e}")))?;
258        if !resp.status().is_success() {
259            let status = resp.status();
260            let text = resp.text().await.unwrap_or_default();
261            return Err(classify_vertex_error(
262                status.as_u16(),
263                format!("Vertex AI rawPredict error {status}: {text}"),
264            ));
265        }
266        Ok(resp)
267    }
268}
269
270#[async_trait::async_trait]
271impl ChatClient for AnthropicVertexClient {
272    async fn get_response(
273        &self,
274        messages: Vec<Message>,
275        options: ChatOptions,
276    ) -> Result<ChatResponse> {
277        let model = self.effective_model(&options);
278        let max_tokens = options.max_tokens.unwrap_or(self.inner.max_tokens);
279        let body = convert::build_cloud_request(
280            &messages,
281            &options,
282            max_tokens,
283            false,
284            ANTHROPIC_VERTEX_VERSION,
285        );
286        let resp = self.send(&model, &body).await?;
287        let value: Value = resp
288            .json()
289            .await
290            .map_err(|e| Error::service(format!("invalid response json: {e}")))?;
291        Ok(convert::parse_response(&value))
292    }
293
294    /// Get a streaming response.
295    ///
296    /// Vertex AI's `streamRawPredict` route does stream Anthropic's usual SSE
297    /// framing, but wiring up a second SSE consumer here (duplicating
298    /// [`crate::AnthropicClient`]'s `parse_sse_stream`) plus a
299    /// synchronous-token-provider-vs-long-lived-stream story is deferred as a
300    /// documented extension point for a first cut of this transport. This
301    /// method calls the non-streaming `rawPredict` route (the same request
302    /// [`ChatClient::get_response`] sends) and adapts the complete
303    /// [`ChatResponse`] into a single [`ChatResponseUpdate`] — the same
304    /// tactic [`crate::bedrock::AnthropicBedrockClient::get_streaming_response`]
305    /// and [`agent_framework_bedrock::BedrockChatClient::get_streaming_response`]
306    /// use. Callers driving this client through
307    /// [`ChatResponse::from_updates`](agent_framework_core::types::ChatResponse::from_updates)
308    /// still get a correct aggregated result; they just don't see partial
309    /// text arrive incrementally.
310    async fn get_streaming_response(
311        &self,
312        messages: Vec<Message>,
313        options: ChatOptions,
314    ) -> Result<ChatStream> {
315        let response = self.get_response(messages, options).await?;
316
317        let mut contents: Vec<Content> = response
318            .messages
319            .iter()
320            .flat_map(|m| m.contents.iter().cloned())
321            .collect();
322        if let Some(usage) = response.usage_details.clone() {
323            contents.push(Content::Usage(UsageContent { details: usage }));
324        }
325
326        let update = ChatResponseUpdate {
327            contents,
328            role: Some(Role::assistant()),
329            response_id: response.response_id.clone(),
330            model: response.model.clone(),
331            finish_reason: response.finish_reason.clone(),
332            ..Default::default()
333        };
334        Ok(stream::once(async move { Ok(update) }).boxed())
335    }
336
337    fn model(&self) -> Option<&str> {
338        Some(&self.inner.model)
339    }
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    fn client() -> AnthropicVertexClient {
347        AnthropicVertexClient::new(
348            "my-project",
349            "us-east5",
350            "claude-sonnet-4-5@20250929",
351            Arc::new(StaticVertexToken::new("test-token")),
352        )
353    }
354
355    #[test]
356    fn static_vertex_token_returns_configured_token() {
357        let provider = StaticVertexToken::new("abc123");
358        assert_eq!(provider.access_token().unwrap(), "abc123");
359    }
360
361    #[test]
362    fn url_contains_publishers_anthropic_models_and_raw_predict() {
363        let c = client();
364        let url = c.url("claude-sonnet-4-5@20250929", false);
365        assert!(url.contains("/projects/my-project/"));
366        assert!(url.contains("/locations/us-east5/"));
367        assert!(url.contains("/publishers/anthropic/models/claude-sonnet-4-5@20250929:rawPredict"));
368        assert_eq!(
369            url,
370            "https://us-east5-aiplatform.googleapis.com/v1/projects/my-project/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-5@20250929:rawPredict"
371        );
372    }
373
374    #[test]
375    fn url_streaming_uses_stream_raw_predict() {
376        let c = client();
377        let url = c.url("claude-sonnet-4-5@20250929", true);
378        assert!(url.ends_with(":streamRawPredict"));
379    }
380
381    #[test]
382    fn accessors_expose_project_region_and_model() {
383        let c = client();
384        assert_eq!(c.project_id(), "my-project");
385        assert_eq!(c.region(), "us-east5");
386        assert_eq!(c.model(), "claude-sonnet-4-5@20250929");
387        assert_eq!(ChatClient::model(&c), Some("claude-sonnet-4-5@20250929"));
388    }
389
390    #[test]
391    fn request_body_has_vertex_anthropic_version_and_no_model_key() {
392        let body = convert::build_cloud_request(
393            &[Message::user("hi")],
394            &ChatOptions::new(),
395            1024,
396            false,
397            ANTHROPIC_VERTEX_VERSION,
398        );
399        assert_eq!(
400            body["anthropic_version"],
401            serde_json::json!("vertex-2023-10-16")
402        );
403        assert!(body.get("model").is_none());
404    }
405
406    // region: env-var constructor
407
408    static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());
409
410    fn clear_vertex_env() {
411        // SAFETY: serialized by ENV_MUTEX against the other env-var tests in
412        // this module.
413        unsafe {
414            std::env::remove_var("GOOGLE_CLOUD_PROJECT");
415            std::env::remove_var("ANTHROPIC_VERTEX_PROJECT_ID");
416            std::env::remove_var("CLOUD_ML_REGION");
417            std::env::remove_var("ANTHROPIC_VERTEX_REGION");
418        }
419    }
420
421    #[test]
422    fn from_env_errors_without_project() {
423        let _guard = ENV_MUTEX.lock().unwrap();
424        clear_vertex_env();
425        let result =
426            AnthropicVertexClient::from_env("claude-x", Arc::new(StaticVertexToken::new("t")));
427        assert!(matches!(result, Err(Error::Configuration(_))));
428        clear_vertex_env();
429    }
430
431    #[test]
432    fn from_env_reads_project_and_region() {
433        let _guard = ENV_MUTEX.lock().unwrap();
434        clear_vertex_env();
435        // SAFETY: serialized by ENV_MUTEX; see clear_vertex_env.
436        unsafe {
437            std::env::set_var("GOOGLE_CLOUD_PROJECT", "proj-1");
438            std::env::set_var("CLOUD_ML_REGION", "europe-west1");
439        }
440        let client =
441            AnthropicVertexClient::from_env("claude-x", Arc::new(StaticVertexToken::new("t")))
442                .unwrap();
443        assert_eq!(client.project_id(), "proj-1");
444        assert_eq!(client.region(), "europe-west1");
445        clear_vertex_env();
446    }
447
448    #[test]
449    fn from_env_falls_back_to_anthropic_specific_vars_and_default_region() {
450        let _guard = ENV_MUTEX.lock().unwrap();
451        clear_vertex_env();
452        // SAFETY: serialized by ENV_MUTEX; see clear_vertex_env.
453        unsafe {
454            std::env::set_var("ANTHROPIC_VERTEX_PROJECT_ID", "proj-2");
455        }
456        let client =
457            AnthropicVertexClient::from_env("claude-x", Arc::new(StaticVertexToken::new("t")))
458                .unwrap();
459        assert_eq!(client.project_id(), "proj-2");
460        assert_eq!(client.region(), DEFAULT_REGION);
461        clear_vertex_env();
462    }
463
464    // endregion
465}
agent_framework_anthropic/vertex.rs

agent_framework_anthropic/
vertex.rs
