mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
修复 Vertex AI 服务账号访问
This commit is contained in:
2
Cargo.lock
generated
2
Cargo.lock
generated
@@ -280,9 +280,11 @@ dependencies = [
|
|||||||
"aether-video-tasks-core",
|
"aether-video-tasks-core",
|
||||||
"async-trait",
|
"async-trait",
|
||||||
"axum",
|
"axum",
|
||||||
|
"base64 0.22.1",
|
||||||
"http",
|
"http",
|
||||||
"regex",
|
"regex",
|
||||||
"reqwest",
|
"reqwest",
|
||||||
|
"rsa",
|
||||||
"serde",
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"sha2",
|
"sha2",
|
||||||
|
|||||||
@@ -6,7 +6,9 @@ use crate::ai_serving::planner::candidate_preparation::{
|
|||||||
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
use crate::ai_serving::planner::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||||
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
|
||||||
use crate::ai_serving::transport::kiro::KiroRequestAuth;
|
use crate::ai_serving::transport::kiro::KiroRequestAuth;
|
||||||
use crate::ai_serving::transport::vertex::resolve_local_vertex_api_key_query_auth;
|
use crate::ai_serving::transport::vertex::{
|
||||||
|
is_vertex_api_key_transport_context, resolve_local_vertex_api_key_query_auth,
|
||||||
|
};
|
||||||
use crate::ai_serving::transport::SameFormatProviderRequestBehavior;
|
use crate::ai_serving::transport::SameFormatProviderRequestBehavior;
|
||||||
use crate::ai_serving::{
|
use crate::ai_serving::{
|
||||||
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||||
@@ -134,7 +136,10 @@ pub(super) async fn prepare_local_same_format_provider_candidate(
|
|||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
if behavior.is_vertex && vertex_query_auth.is_none() {
|
if behavior.is_vertex
|
||||||
|
&& is_vertex_api_key_transport_context(&transport)
|
||||||
|
&& vertex_query_auth.is_none()
|
||||||
|
{
|
||||||
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
super::super::payload::mark_skipped_local_same_format_provider_candidate(
|
||||||
state,
|
state,
|
||||||
input,
|
input,
|
||||||
|
|||||||
@@ -362,8 +362,8 @@ fn provider_query_transport_supports_standard_test_execution(
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
"gemini:generate_content" => {
|
"gemini:generate_content" => {
|
||||||
if crate::provider_transport::is_vertex_api_key_transport_context(transport) {
|
if crate::provider_transport::is_vertex_transport_context(transport) {
|
||||||
aether_provider_transport::vertex::supports_local_vertex_api_key_gemini_transport_with_network(transport)
|
aether_provider_transport::vertex::supports_local_vertex_gemini_transport_with_network(transport)
|
||||||
} else {
|
} else {
|
||||||
state.supports_local_gemini_transport_with_network(transport, api_format)
|
state.supports_local_gemini_transport_with_network(transport, api_format)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -194,7 +194,7 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let oauth_auth = match format_value.as_str() {
|
let oauth_auth = match format_value.as_str() {
|
||||||
"openai:chat" | "claude:messages" => {
|
"openai:chat" | "claude:messages" | "gemini:generate_content" => {
|
||||||
match state.resolve_local_oauth_request_auth(&transport).await {
|
match state.resolve_local_oauth_request_auth(&transport).await {
|
||||||
Ok(Some(crate::provider_transport::LocalResolvedOAuthRequestAuth::Header {
|
Ok(Some(crate::provider_transport::LocalResolvedOAuthRequestAuth::Header {
|
||||||
name,
|
name,
|
||||||
@@ -217,16 +217,26 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
|
|||||||
}
|
}
|
||||||
"gemini:generate_content" => {
|
"gemini:generate_content" => {
|
||||||
crate::provider_transport::auth::resolve_local_gemini_auth(&transport)
|
crate::provider_transport::auth::resolve_local_gemini_auth(&transport)
|
||||||
|
.or(oauth_auth.clone())
|
||||||
}
|
}
|
||||||
_ => None,
|
_ => None,
|
||||||
};
|
};
|
||||||
let Some((auth_header, auth_value)) = auth else {
|
|
||||||
return None;
|
|
||||||
};
|
|
||||||
let uses_vertex_query_auth = crate::provider_transport::uses_vertex_api_key_query_auth(
|
let uses_vertex_query_auth = crate::provider_transport::uses_vertex_api_key_query_auth(
|
||||||
&transport,
|
&transport,
|
||||||
format_value.as_str(),
|
format_value.as_str(),
|
||||||
);
|
);
|
||||||
|
let vertex_query_auth = if uses_vertex_query_auth {
|
||||||
|
crate::provider_transport::vertex::resolve_local_vertex_api_key_query_auth(&transport)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let (auth_header, auth_value) = match auth {
|
||||||
|
Some((auth_header, auth_value)) => (auth_header, auth_value),
|
||||||
|
None if uses_vertex_query_auth && vertex_query_auth.is_some() => {
|
||||||
|
(String::new(), String::new())
|
||||||
|
}
|
||||||
|
None => return None,
|
||||||
|
};
|
||||||
|
|
||||||
let upstream_url = crate::provider_transport::build_transport_request_url(
|
let upstream_url = crate::provider_transport::build_transport_request_url(
|
||||||
&transport,
|
&transport,
|
||||||
|
|||||||
@@ -57,6 +57,17 @@ fn oauth_access_token_expired(expires_at_unix_secs: Option<u64>, now_unix_secs:
|
|||||||
expires_at_unix_secs.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs)
|
expires_at_unix_secs.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn local_oauth_refresh_entry_should_stay_memory_only(
|
||||||
|
transport: &provider_transport::GatewayProviderTransportSnapshot,
|
||||||
|
entry: &provider_transport::CachedOAuthEntry,
|
||||||
|
) -> bool {
|
||||||
|
entry
|
||||||
|
.provider_type
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case(provider_transport::vertex::VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE)
|
||||||
|
&& provider_transport::is_vertex_service_account_transport_context(transport)
|
||||||
|
}
|
||||||
|
|
||||||
fn oauth_auth_config_refresh_token_fingerprint(auth_config: Option<&str>) -> Option<String> {
|
fn oauth_auth_config_refresh_token_fingerprint(auth_config: Option<&str>) -> Option<String> {
|
||||||
let parsed = auth_config
|
let parsed = auth_config
|
||||||
.map(str::trim)
|
.map(str::trim)
|
||||||
@@ -1094,6 +1105,17 @@ impl AppState {
|
|||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if local_oauth_refresh_entry_should_stay_memory_only(transport, entry) {
|
||||||
|
tracing::info!(
|
||||||
|
key_id = %key_id,
|
||||||
|
provider_id = %transport.provider.id,
|
||||||
|
provider_type = %transport.provider.provider_type,
|
||||||
|
expires_at_unix_secs = ?entry.expires_at_unix_secs,
|
||||||
|
"gateway local oauth refresh entry kept in memory only"
|
||||||
|
);
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
let Some(encryption_key) = self.data.encryption_key() else {
|
let Some(encryption_key) = self.data.encryption_key() else {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
};
|
};
|
||||||
@@ -1703,4 +1725,71 @@ mod tests {
|
|||||||
Some("[OAUTH_EXPIRED] access token invalid".to_string()),
|
Some("[OAUTH_EXPIRED] access token invalid".to_string()),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn vertex_service_account_refresh_entry_stays_memory_only() {
|
||||||
|
let transport = crate::provider_transport::GatewayProviderTransportSnapshot {
|
||||||
|
provider: crate::provider_transport::snapshot::GatewayProviderTransportProvider {
|
||||||
|
id: "provider-1".to_string(),
|
||||||
|
name: "Vertex".to_string(),
|
||||||
|
provider_type: "vertex_ai".to_string(),
|
||||||
|
website: None,
|
||||||
|
is_active: true,
|
||||||
|
keep_priority_on_conversion: false,
|
||||||
|
enable_format_conversion: false,
|
||||||
|
concurrent_limit: None,
|
||||||
|
max_retries: None,
|
||||||
|
proxy: None,
|
||||||
|
request_timeout_secs: None,
|
||||||
|
stream_first_byte_timeout_secs: None,
|
||||||
|
config: None,
|
||||||
|
},
|
||||||
|
endpoint: crate::provider_transport::snapshot::GatewayProviderTransportEndpoint {
|
||||||
|
id: "endpoint-1".to_string(),
|
||||||
|
provider_id: "provider-1".to_string(),
|
||||||
|
api_format: "gemini:generate_content".to_string(),
|
||||||
|
api_family: Some("gemini".to_string()),
|
||||||
|
endpoint_kind: Some("chat".to_string()),
|
||||||
|
is_active: true,
|
||||||
|
base_url: "https://aiplatform.googleapis.com".to_string(),
|
||||||
|
header_rules: None,
|
||||||
|
body_rules: None,
|
||||||
|
max_retries: None,
|
||||||
|
custom_path: None,
|
||||||
|
config: None,
|
||||||
|
format_acceptance_config: None,
|
||||||
|
proxy: None,
|
||||||
|
},
|
||||||
|
key: crate::provider_transport::snapshot::GatewayProviderTransportKey {
|
||||||
|
id: "key-1".to_string(),
|
||||||
|
provider_id: "provider-1".to_string(),
|
||||||
|
name: "Gemini".to_string(),
|
||||||
|
auth_type: "service_account".to_string(),
|
||||||
|
is_active: true,
|
||||||
|
api_formats: Some(vec!["gemini:generate_content".to_string()]),
|
||||||
|
auth_type_by_format: None,
|
||||||
|
allow_auth_channel_mismatch_formats: None,
|
||||||
|
allowed_models: None,
|
||||||
|
capabilities: None,
|
||||||
|
rate_multipliers: None,
|
||||||
|
global_priority_by_format: None,
|
||||||
|
expires_at_unix_secs: None,
|
||||||
|
proxy: None,
|
||||||
|
fingerprint: None,
|
||||||
|
decrypted_api_key: "__placeholder__".to_string(),
|
||||||
|
decrypted_auth_config: Some("{\"project_id\":\"demo\"}".to_string()),
|
||||||
|
},
|
||||||
|
};
|
||||||
|
let entry = crate::provider_transport::CachedOAuthEntry {
|
||||||
|
provider_type: "vertex_ai".to_string(),
|
||||||
|
auth_header_name: "authorization".to_string(),
|
||||||
|
auth_header_value: "Bearer access-token".to_string(),
|
||||||
|
expires_at_unix_secs: Some(4_102_444_800),
|
||||||
|
metadata: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
assert!(super::local_oauth_refresh_entry_should_stay_memory_only(
|
||||||
|
&transport, &entry
|
||||||
|
));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,2 @@
|
|||||||
|
-- Historical compatibility no-op.
|
||||||
|
-- The auth module tables are already present in the MySQL baseline.
|
||||||
@@ -0,0 +1,3 @@
|
|||||||
|
-- Historical compatibility no-op.
|
||||||
|
-- The auth module tables are already present in the squashed baseline and
|
||||||
|
-- 20260505000000_sync_core_export_columns for existing Postgres deployments.
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
-- Historical compatibility no-op.
|
||||||
|
-- The auth module tables are already present in the SQLite baseline.
|
||||||
@@ -8,6 +8,7 @@ use aether_provider_transport::{
|
|||||||
};
|
};
|
||||||
use base64::engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD};
|
use base64::engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD};
|
||||||
use base64::Engine as _;
|
use base64::Engine as _;
|
||||||
|
use rsa::pkcs1::DecodeRsaPrivateKey;
|
||||||
use rsa::pkcs1v15::SigningKey;
|
use rsa::pkcs1v15::SigningKey;
|
||||||
use rsa::pkcs8::DecodePrivateKey;
|
use rsa::pkcs8::DecodePrivateKey;
|
||||||
use rsa::signature::{SignatureEncoding, Signer};
|
use rsa::signature::{SignatureEncoding, Signer};
|
||||||
@@ -675,8 +676,7 @@ fn build_vertex_service_account_assertion(
|
|||||||
.map_err(|err| format!("vertex_ai(service_account): jwt payload encode failed: {err}"))?,
|
.map_err(|err| format!("vertex_ai(service_account): jwt payload encode failed: {err}"))?,
|
||||||
);
|
);
|
||||||
let message = format!("{header}.{payload}");
|
let message = format!("{header}.{payload}");
|
||||||
let private_key = RsaPrivateKey::from_pkcs8_pem(private_key_pem)
|
let private_key = decode_vertex_service_account_private_key(private_key_pem)?;
|
||||||
.map_err(|err| format!("vertex_ai(service_account): private_key parse failed: {err}"))?;
|
|
||||||
let signing_key = SigningKey::<Sha256>::new(private_key);
|
let signing_key = SigningKey::<Sha256>::new(private_key);
|
||||||
let signature = signing_key.sign(message.as_bytes());
|
let signature = signing_key.sign(message.as_bytes());
|
||||||
Ok(format!(
|
Ok(format!(
|
||||||
@@ -685,6 +685,19 @@ fn build_vertex_service_account_assertion(
|
|||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn decode_vertex_service_account_private_key(
|
||||||
|
private_key_pem: &str,
|
||||||
|
) -> Result<RsaPrivateKey, String> {
|
||||||
|
match RsaPrivateKey::from_pkcs8_pem(private_key_pem) {
|
||||||
|
Ok(private_key) => Ok(private_key),
|
||||||
|
Err(pkcs8_err) => RsaPrivateKey::from_pkcs1_pem(private_key_pem).map_err(|pkcs1_err| {
|
||||||
|
format!(
|
||||||
|
"vertex_ai(service_account): private_key parse failed: pkcs8: {pkcs8_err}; pkcs1: {pkcs1_err}"
|
||||||
|
)
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn execution_result_json_body(result: &ExecutionResult) -> Result<Value, String> {
|
fn execution_result_json_body(result: &ExecutionResult) -> Result<Value, String> {
|
||||||
if result.status_code != 200 {
|
if result.status_code != 200 {
|
||||||
return Err(execution_result_error_message(result));
|
return Err(execution_result_error_message(result));
|
||||||
|
|||||||
@@ -15,12 +15,14 @@ aether-oauth.workspace = true
|
|||||||
aether-runtime-state.workspace = true
|
aether-runtime-state.workspace = true
|
||||||
aether-video-tasks-core.workspace = true
|
aether-video-tasks-core.workspace = true
|
||||||
async-trait.workspace = true
|
async-trait.workspace = true
|
||||||
|
base64.workspace = true
|
||||||
http.workspace = true
|
http.workspace = true
|
||||||
regex.workspace = true
|
regex.workspace = true
|
||||||
reqwest.workspace = true
|
reqwest.workspace = true
|
||||||
|
rsa = "0.9.10"
|
||||||
serde.workspace = true
|
serde.workspace = true
|
||||||
serde_json.workspace = true
|
serde_json.workspace = true
|
||||||
sha2.workspace = true
|
sha2 = { workspace = true, features = ["oid"] }
|
||||||
thiserror.workspace = true
|
thiserror.workspace = true
|
||||||
tokio.workspace = true
|
tokio.workspace = true
|
||||||
tracing.workspace = true
|
tracing.workspace = true
|
||||||
|
|||||||
@@ -15,8 +15,8 @@ use crate::policy::{
|
|||||||
local_standard_transport_unsupported_reason_with_network,
|
local_standard_transport_unsupported_reason_with_network,
|
||||||
};
|
};
|
||||||
use crate::vertex::{
|
use crate::vertex::{
|
||||||
is_vertex_api_key_transport_context,
|
is_vertex_api_key_transport_context, is_vertex_transport_context,
|
||||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
|
local_vertex_gemini_transport_unsupported_reason_with_network,
|
||||||
resolve_local_vertex_api_key_query_auth, VERTEX_API_KEY_QUERY_PARAM,
|
resolve_local_vertex_api_key_query_auth, VERTEX_API_KEY_QUERY_PARAM,
|
||||||
};
|
};
|
||||||
use crate::GatewayProviderTransportSnapshot;
|
use crate::GatewayProviderTransportSnapshot;
|
||||||
@@ -106,8 +106,8 @@ pub fn request_conversion_transport_unsupported_reason(
|
|||||||
"claude:messages" => {
|
"claude:messages" => {
|
||||||
local_standard_transport_unsupported_reason_with_network(transport, "claude:messages")
|
local_standard_transport_unsupported_reason_with_network(transport, "claude:messages")
|
||||||
}
|
}
|
||||||
"gemini:generate_content" if is_vertex_api_key_transport_context(transport) => {
|
"gemini:generate_content" if is_vertex_transport_context(transport) => {
|
||||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport)
|
local_vertex_gemini_transport_unsupported_reason_with_network(transport)
|
||||||
}
|
}
|
||||||
"gemini:generate_content" => local_gemini_transport_unsupported_reason_with_network(
|
"gemini:generate_content" => local_gemini_transport_unsupported_reason_with_network(
|
||||||
transport,
|
transport,
|
||||||
|
|||||||
@@ -103,7 +103,10 @@ pub use standard::{
|
|||||||
StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput,
|
StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput,
|
||||||
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
|
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
|
||||||
};
|
};
|
||||||
pub use vertex::{is_vertex_api_key_transport_context, uses_vertex_api_key_query_auth};
|
pub use vertex::{
|
||||||
|
is_vertex_api_key_transport_context, is_vertex_service_account_transport_context,
|
||||||
|
is_vertex_transport_context, uses_vertex_api_key_query_auth,
|
||||||
|
};
|
||||||
pub use video::{
|
pub use video::{
|
||||||
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
|
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
|
||||||
reconstruct_local_video_task_snapshot, resolve_local_video_task_transport,
|
reconstruct_local_video_task_snapshot, resolve_local_video_task_transport,
|
||||||
|
|||||||
@@ -19,6 +19,9 @@ use super::kiro::{
|
|||||||
supports_local_kiro_request_auth_resolution, KiroOAuthRefreshAdapter, KiroRequestAuth,
|
supports_local_kiro_request_auth_resolution, KiroOAuthRefreshAdapter, KiroRequestAuth,
|
||||||
};
|
};
|
||||||
use super::snapshot::GatewayProviderTransportSnapshot;
|
use super::snapshot::GatewayProviderTransportSnapshot;
|
||||||
|
use super::vertex::{
|
||||||
|
supports_local_vertex_service_account_auth_resolution, VertexServiceAccountRefreshAdapter,
|
||||||
|
};
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
#[allow(clippy::large_enum_variant)]
|
#[allow(clippy::large_enum_variant)]
|
||||||
@@ -339,6 +342,7 @@ impl LocalOAuthRefreshCoordinator {
|
|||||||
Self {
|
Self {
|
||||||
adapters: vec![
|
adapters: vec![
|
||||||
Arc::new(KiroOAuthRefreshAdapter::default()),
|
Arc::new(KiroOAuthRefreshAdapter::default()),
|
||||||
|
Arc::new(VertexServiceAccountRefreshAdapter),
|
||||||
Arc::new(GenericOAuthRefreshAdapter::default()),
|
Arc::new(GenericOAuthRefreshAdapter::default()),
|
||||||
],
|
],
|
||||||
cache: Mutex::new(BTreeMap::new()),
|
cache: Mutex::new(BTreeMap::new()),
|
||||||
@@ -547,6 +551,7 @@ pub fn supports_local_oauth_request_auth_resolution(
|
|||||||
transport: &GatewayProviderTransportSnapshot,
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
supports_local_kiro_request_auth_resolution(transport)
|
supports_local_kiro_request_auth_resolution(transport)
|
||||||
|
|| supports_local_vertex_service_account_auth_resolution(transport)
|
||||||
|| supports_local_generic_oauth_request_auth_resolution(transport)
|
|| supports_local_generic_oauth_request_auth_resolution(transport)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -15,7 +15,8 @@ use crate::url::{
|
|||||||
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
|
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
|
||||||
};
|
};
|
||||||
use crate::vertex::{
|
use crate::vertex::{
|
||||||
build_vertex_api_key_gemini_content_url, resolve_local_vertex_api_key_query_auth,
|
build_vertex_api_key_gemini_content_url, build_vertex_service_account_gemini_content_url,
|
||||||
|
resolve_local_vertex_api_key_query_auth, resolve_local_vertex_service_account_auth_config,
|
||||||
};
|
};
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy)]
|
#[derive(Debug, Clone, Copy)]
|
||||||
@@ -256,6 +257,14 @@ fn build_transport_hook_url(
|
|||||||
params.request_query,
|
params.request_query,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
if let Some(auth_config) = resolve_local_vertex_service_account_auth_config(transport) {
|
||||||
|
return build_vertex_service_account_gemini_content_url(
|
||||||
|
params.mapped_model?,
|
||||||
|
params.upstream_is_stream,
|
||||||
|
&auth_config,
|
||||||
|
params.request_query,
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if is_antigravity_provider_transport(transport) {
|
if is_antigravity_provider_transport(transport) {
|
||||||
@@ -526,6 +535,43 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn uses_vertex_service_account_hook_before_default_gemini_url() {
|
||||||
|
let mut transport = sample_transport(
|
||||||
|
"vertex_ai",
|
||||||
|
"gemini:generate_content",
|
||||||
|
"https://aiplatform.googleapis.com",
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
transport.key.auth_type = "service_account".to_string();
|
||||||
|
transport.key.decrypted_api_key = "__placeholder__".to_string();
|
||||||
|
transport.key.decrypted_auth_config = Some(
|
||||||
|
r#"{
|
||||||
|
"client_email":"svc@example.iam.gserviceaccount.com",
|
||||||
|
"private_key":"TEST-PRIVATE-KEY",
|
||||||
|
"project_id":"demo-project"
|
||||||
|
}"#
|
||||||
|
.to_string(),
|
||||||
|
);
|
||||||
|
|
||||||
|
let url = build_transport_request_url(
|
||||||
|
&transport,
|
||||||
|
TransportRequestUrlParams {
|
||||||
|
provider_api_format: "gemini:generate_content",
|
||||||
|
mapped_model: Some("gemini-3.1-pro-preview"),
|
||||||
|
upstream_is_stream: false,
|
||||||
|
request_query: Some("foo=bar&beta=1"),
|
||||||
|
kiro_api_region: None,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.expect("vertex service account hook url");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
url,
|
||||||
|
"https://aiplatform.googleapis.com/v1/projects/demo-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:generateContent?foo=bar"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn builds_openai_responses_url_for_formal_format_name() {
|
fn builds_openai_responses_url_for_formal_format_name() {
|
||||||
let transport = sample_transport(
|
let transport = sample_transport(
|
||||||
|
|||||||
@@ -23,8 +23,8 @@ use crate::rules::{
|
|||||||
};
|
};
|
||||||
use crate::snapshot::GatewayProviderTransportSnapshot;
|
use crate::snapshot::GatewayProviderTransportSnapshot;
|
||||||
use crate::vertex::{
|
use crate::vertex::{
|
||||||
is_vertex_api_key_transport_context,
|
is_vertex_service_account_transport_context, is_vertex_transport_context,
|
||||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
|
local_vertex_gemini_transport_unsupported_reason_with_network,
|
||||||
};
|
};
|
||||||
use crate::{build_transport_request_url, ensure_upstream_auth_header, TransportRequestUrlParams};
|
use crate::{build_transport_request_url, ensure_upstream_auth_header, TransportRequestUrlParams};
|
||||||
|
|
||||||
@@ -103,7 +103,7 @@ pub fn classify_same_format_provider_request_behavior(
|
|||||||
.provider_type
|
.provider_type
|
||||||
.trim()
|
.trim()
|
||||||
.eq_ignore_ascii_case("claude_code");
|
.eq_ignore_ascii_case("claude_code");
|
||||||
let is_vertex = is_vertex_api_key_transport_context(transport);
|
let is_vertex = is_vertex_transport_context(transport);
|
||||||
let is_kiro = is_kiro_provider_transport(transport);
|
let is_kiro = is_kiro_provider_transport(transport);
|
||||||
let upstream_is_stream = aether_ai_formats::resolve_upstream_is_stream_from_endpoint_config(
|
let upstream_is_stream = aether_ai_formats::resolve_upstream_is_stream_from_endpoint_config(
|
||||||
transport.endpoint.config.as_ref(),
|
transport.endpoint.config.as_ref(),
|
||||||
@@ -328,7 +328,7 @@ pub fn same_format_provider_transport_unsupported_reason(
|
|||||||
} else if behavior.is_claude_code {
|
} else if behavior.is_claude_code {
|
||||||
local_claude_code_transport_unsupported_reason_with_network(transport, api_format)
|
local_claude_code_transport_unsupported_reason_with_network(transport, api_format)
|
||||||
} else if behavior.is_vertex {
|
} else if behavior.is_vertex {
|
||||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport)
|
local_vertex_gemini_transport_unsupported_reason_with_network(transport)
|
||||||
} else {
|
} else {
|
||||||
match family {
|
match family {
|
||||||
SameFormatProviderFamily::Standard => {
|
SameFormatProviderFamily::Standard => {
|
||||||
@@ -391,6 +391,9 @@ pub fn should_try_same_format_provider_oauth_auth(
|
|||||||
behavior.is_kiro
|
behavior.is_kiro
|
||||||
|| matches!(family, SameFormatProviderFamily::Standard)
|
|| matches!(family, SameFormatProviderFamily::Standard)
|
||||||
&& resolve_local_standard_auth(transport).is_none()
|
&& resolve_local_standard_auth(transport).is_none()
|
||||||
|
|| matches!(family, SameFormatProviderFamily::Gemini)
|
||||||
|
&& behavior.is_vertex
|
||||||
|
&& is_vertex_service_account_transport_context(transport)
|
||||||
|| matches!(family, SameFormatProviderFamily::Gemini)
|
|| matches!(family, SameFormatProviderFamily::Gemini)
|
||||||
&& !behavior.is_vertex
|
&& !behavior.is_vertex
|
||||||
&& resolve_local_gemini_auth(transport).is_none()
|
&& resolve_local_gemini_auth(transport).is_none()
|
||||||
|
|||||||
@@ -1,6 +1,29 @@
|
|||||||
|
use std::collections::BTreeMap;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||||
|
use base64::Engine as _;
|
||||||
|
use rsa::pkcs1::DecodeRsaPrivateKey;
|
||||||
|
use rsa::pkcs1v15::SigningKey;
|
||||||
|
use rsa::pkcs8::DecodePrivateKey;
|
||||||
|
use rsa::signature::{SignatureEncoding, Signer};
|
||||||
|
use rsa::RsaPrivateKey;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use sha2::Sha256;
|
||||||
|
use url::form_urlencoded;
|
||||||
|
|
||||||
|
use super::super::oauth_refresh::{
|
||||||
|
CachedOAuthEntry, LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthRefreshAdapter,
|
||||||
|
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
|
||||||
|
};
|
||||||
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
||||||
|
|
||||||
pub const VERTEX_API_KEY_QUERY_PARAM: &str = "key";
|
pub const VERTEX_API_KEY_QUERY_PARAM: &str = "key";
|
||||||
|
pub const VERTEX_SERVICE_ACCOUNT_AUTH_HEADER: &str = "authorization";
|
||||||
|
pub const VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE: &str = "vertex_ai";
|
||||||
|
pub const GOOGLE_OAUTH_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
|
||||||
|
const GOOGLE_CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
|
||||||
|
const SERVICE_ACCOUNT_REFRESH_SKEW_SECS: u64 = 120;
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
pub struct VertexApiKeyQueryAuth {
|
pub struct VertexApiKeyQueryAuth {
|
||||||
@@ -8,6 +31,16 @@ pub struct VertexApiKeyQueryAuth {
|
|||||||
pub value: String,
|
pub value: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct VertexServiceAccountAuthConfig {
|
||||||
|
pub client_email: String,
|
||||||
|
pub private_key: String,
|
||||||
|
pub project_id: String,
|
||||||
|
pub token_uri: String,
|
||||||
|
pub region: Option<String>,
|
||||||
|
pub model_regions: BTreeMap<String, String>,
|
||||||
|
}
|
||||||
|
|
||||||
pub fn resolve_local_vertex_api_key_query_auth(
|
pub fn resolve_local_vertex_api_key_query_auth(
|
||||||
transport: &GatewayProviderTransportSnapshot,
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
) -> Option<VertexApiKeyQueryAuth> {
|
) -> Option<VertexApiKeyQueryAuth> {
|
||||||
@@ -39,14 +72,274 @@ pub fn resolve_local_vertex_api_key_query_auth(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn resolve_local_vertex_service_account_auth_config(
|
||||||
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
) -> Option<VertexServiceAccountAuthConfig> {
|
||||||
|
if !super::is_vertex_service_account_transport_context(transport) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
parse_vertex_service_account_auth_config(transport.key.decrypted_auth_config.as_deref())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn supports_local_vertex_service_account_auth_resolution(
|
||||||
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
) -> bool {
|
||||||
|
resolve_local_vertex_service_account_auth_config(transport).is_some()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn parse_vertex_service_account_auth_config(
|
||||||
|
raw: Option<&str>,
|
||||||
|
) -> Option<VertexServiceAccountAuthConfig> {
|
||||||
|
let raw = raw.map(str::trim).filter(|value| !value.is_empty())?;
|
||||||
|
let value: Value = serde_json::from_str(raw).ok()?;
|
||||||
|
parse_vertex_service_account_auth_config_value(&value)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_vertex_service_account_auth_config_value(
|
||||||
|
value: &Value,
|
||||||
|
) -> Option<VertexServiceAccountAuthConfig> {
|
||||||
|
let client_email = json_string(value.get("client_email"))?;
|
||||||
|
let private_key = json_string(value.get("private_key"))?;
|
||||||
|
let project_id = json_string(value.get("project_id"))?;
|
||||||
|
let token_uri =
|
||||||
|
json_string(value.get("token_uri")).unwrap_or_else(|| GOOGLE_OAUTH_TOKEN_URL.to_string());
|
||||||
|
let region = json_string(value.get("region"));
|
||||||
|
let model_regions = value
|
||||||
|
.get("model_regions")
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.map(|items| {
|
||||||
|
items
|
||||||
|
.iter()
|
||||||
|
.filter_map(|(model, region)| {
|
||||||
|
let model = model.trim();
|
||||||
|
let region = region.as_str()?.trim();
|
||||||
|
(!model.is_empty() && !region.is_empty())
|
||||||
|
.then(|| (model.to_string(), region.to_string()))
|
||||||
|
})
|
||||||
|
.collect::<BTreeMap<_, _>>()
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
Some(VertexServiceAccountAuthConfig {
|
||||||
|
client_email,
|
||||||
|
private_key,
|
||||||
|
project_id,
|
||||||
|
token_uri,
|
||||||
|
region,
|
||||||
|
model_regions,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn json_string(value: Option<&Value>) -> Option<String> {
|
||||||
|
value
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct VertexServiceAccountRefreshAdapter;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl LocalOAuthRefreshAdapter for VertexServiceAccountRefreshAdapter {
|
||||||
|
fn provider_type(&self) -> &'static str {
|
||||||
|
VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE
|
||||||
|
}
|
||||||
|
|
||||||
|
fn supports(&self, transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||||
|
supports_local_vertex_service_account_auth_resolution(transport)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve_cached(
|
||||||
|
&self,
|
||||||
|
_transport: &GatewayProviderTransportSnapshot,
|
||||||
|
entry: &CachedOAuthEntry,
|
||||||
|
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||||
|
if !entry
|
||||||
|
.provider_type
|
||||||
|
.eq_ignore_ascii_case(VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE)
|
||||||
|
{
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
if service_account_token_expires_soon(entry.expires_at_unix_secs) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let name = entry.auth_header_name.trim();
|
||||||
|
let value = entry.auth_header_value.trim();
|
||||||
|
if name.is_empty() || value.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some(LocalResolvedOAuthRequestAuth::Header {
|
||||||
|
name: name.to_ascii_lowercase(),
|
||||||
|
value: value.to_string(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve_without_refresh(
|
||||||
|
&self,
|
||||||
|
_transport: &GatewayProviderTransportSnapshot,
|
||||||
|
) -> Option<LocalResolvedOAuthRequestAuth> {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
|
fn should_refresh(
|
||||||
|
&self,
|
||||||
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
entry: Option<&CachedOAuthEntry>,
|
||||||
|
) -> bool {
|
||||||
|
supports_local_vertex_service_account_auth_resolution(transport)
|
||||||
|
&& entry
|
||||||
|
.and_then(|cached| self.resolve_cached(transport, cached))
|
||||||
|
.is_none()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn refresh(
|
||||||
|
&self,
|
||||||
|
executor: &dyn LocalOAuthHttpExecutor,
|
||||||
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
_entry: Option<&CachedOAuthEntry>,
|
||||||
|
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
|
||||||
|
let Some(auth_config) = resolve_local_vertex_service_account_auth_config(transport) else {
|
||||||
|
return Ok(None);
|
||||||
|
};
|
||||||
|
let now = aether_oauth::core::current_unix_secs();
|
||||||
|
let assertion = build_vertex_service_account_assertion(&auth_config, now)?;
|
||||||
|
let body = form_urlencoded::Serializer::new(String::new())
|
||||||
|
.append_pair(
|
||||||
|
"grant_type",
|
||||||
|
"urn:ietf:params:oauth:grant-type:jwt-bearer",
|
||||||
|
)
|
||||||
|
.append_pair("assertion", &assertion)
|
||||||
|
.finish();
|
||||||
|
let response = executor
|
||||||
|
.execute(
|
||||||
|
VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
|
||||||
|
transport,
|
||||||
|
&LocalOAuthHttpRequest {
|
||||||
|
request_id: "vertex_ai:service-account-token",
|
||||||
|
method: reqwest::Method::POST,
|
||||||
|
url: auth_config.token_uri.clone(),
|
||||||
|
headers: BTreeMap::from([(
|
||||||
|
"content-type".to_string(),
|
||||||
|
"application/x-www-form-urlencoded".to_string(),
|
||||||
|
)]),
|
||||||
|
json_body: None,
|
||||||
|
body_bytes: Some(body.into_bytes()),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.await?;
|
||||||
|
if response.status_code != 200 {
|
||||||
|
return Err(LocalOAuthRefreshError::HttpStatus {
|
||||||
|
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
|
||||||
|
status_code: response.status_code,
|
||||||
|
body_excerpt: body_excerpt(&response.body_text),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let body_json: Value =
|
||||||
|
serde_json::from_str(&response.body_text).map_err(|err| {
|
||||||
|
LocalOAuthRefreshError::InvalidResponse {
|
||||||
|
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
|
||||||
|
message: format!("vertex service account token response is not JSON: {err}"),
|
||||||
|
}
|
||||||
|
})?;
|
||||||
|
let access_token = json_string(body_json.get("access_token")).ok_or_else(|| {
|
||||||
|
LocalOAuthRefreshError::InvalidResponse {
|
||||||
|
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
|
||||||
|
message: "vertex service account token response missing access_token".to_string(),
|
||||||
|
}
|
||||||
|
})?;
|
||||||
|
let expires_in = body_json
|
||||||
|
.get("expires_in")
|
||||||
|
.and_then(Value::as_u64)
|
||||||
|
.unwrap_or(3600);
|
||||||
|
|
||||||
|
Ok(Some(CachedOAuthEntry {
|
||||||
|
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE.to_string(),
|
||||||
|
auth_header_name: VERTEX_SERVICE_ACCOUNT_AUTH_HEADER.to_string(),
|
||||||
|
auth_header_value: format!("Bearer {access_token}"),
|
||||||
|
expires_at_unix_secs: Some(now.saturating_add(expires_in)),
|
||||||
|
metadata: Some(json!({
|
||||||
|
"project_id": auth_config.project_id,
|
||||||
|
"client_email": auth_config.client_email,
|
||||||
|
})),
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn build_vertex_service_account_assertion(
|
||||||
|
auth_config: &VertexServiceAccountAuthConfig,
|
||||||
|
now_unix_secs: u64,
|
||||||
|
) -> Result<String, LocalOAuthRefreshError> {
|
||||||
|
let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"RS256","typ":"JWT"}"#);
|
||||||
|
let payload = URL_SAFE_NO_PAD.encode(
|
||||||
|
serde_json::to_string(&json!({
|
||||||
|
"iss": auth_config.client_email,
|
||||||
|
"sub": auth_config.client_email,
|
||||||
|
"scope": GOOGLE_CLOUD_PLATFORM_SCOPE,
|
||||||
|
"aud": auth_config.token_uri,
|
||||||
|
"iat": now_unix_secs,
|
||||||
|
"exp": now_unix_secs.saturating_add(3600),
|
||||||
|
}))
|
||||||
|
.map_err(|err| LocalOAuthRefreshError::InvalidResponse {
|
||||||
|
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
|
||||||
|
message: format!("vertex service account jwt payload encode failed: {err}"),
|
||||||
|
})?,
|
||||||
|
);
|
||||||
|
let message = format!("{header}.{payload}");
|
||||||
|
let private_key = decode_vertex_service_account_private_key(auth_config.private_key.as_str())?;
|
||||||
|
let signing_key = SigningKey::<Sha256>::new(private_key);
|
||||||
|
let signature = signing_key.sign(message.as_bytes());
|
||||||
|
Ok(format!(
|
||||||
|
"{message}.{}",
|
||||||
|
URL_SAFE_NO_PAD.encode(signature.to_bytes())
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decode_vertex_service_account_private_key(
|
||||||
|
private_key_pem: &str,
|
||||||
|
) -> Result<RsaPrivateKey, LocalOAuthRefreshError> {
|
||||||
|
match RsaPrivateKey::from_pkcs8_pem(private_key_pem) {
|
||||||
|
Ok(private_key) => return Ok(private_key),
|
||||||
|
Err(pkcs8_err) => RsaPrivateKey::from_pkcs1_pem(private_key_pem).map_err(|pkcs1_err| {
|
||||||
|
LocalOAuthRefreshError::InvalidResponse {
|
||||||
|
provider_type: VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
|
||||||
|
message: format!(
|
||||||
|
"vertex service account private_key parse failed: pkcs8: {pkcs8_err}; pkcs1: {pkcs1_err}"
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn service_account_token_expires_soon(expires_at_unix_secs: Option<u64>) -> bool {
|
||||||
|
expires_at_unix_secs
|
||||||
|
.map(|expires_at_unix_secs| {
|
||||||
|
aether_oauth::core::current_unix_secs()
|
||||||
|
>= expires_at_unix_secs.saturating_sub(SERVICE_ACCOUNT_REFRESH_SKEW_SECS)
|
||||||
|
})
|
||||||
|
.unwrap_or(true)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn body_excerpt(value: &str) -> String {
|
||||||
|
value.chars().take(500).collect()
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::super::super::snapshot::{
|
use super::super::super::snapshot::{
|
||||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||||
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
|
||||||
};
|
};
|
||||||
|
use rsa::pkcs1::{EncodeRsaPrivateKey, LineEnding};
|
||||||
|
use rsa::rand_core::OsRng;
|
||||||
|
use rsa::RsaPrivateKey;
|
||||||
|
|
||||||
use super::{resolve_local_vertex_api_key_query_auth, VERTEX_API_KEY_QUERY_PARAM};
|
use super::{
|
||||||
|
decode_vertex_service_account_private_key, parse_vertex_service_account_auth_config,
|
||||||
|
resolve_local_vertex_api_key_query_auth,
|
||||||
|
supports_local_vertex_service_account_auth_resolution, VERTEX_API_KEY_QUERY_PARAM,
|
||||||
|
};
|
||||||
|
|
||||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||||
GatewayProviderTransportSnapshot {
|
GatewayProviderTransportSnapshot {
|
||||||
@@ -136,4 +429,57 @@ mod tests {
|
|||||||
.expect("custom aiplatform transport should resolve");
|
.expect("custom aiplatform transport should resolve");
|
||||||
assert_eq!(auth.value, "vertex-secret");
|
assert_eq!(auth.value, "vertex-secret");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_service_account_auth_config() {
|
||||||
|
let config = parse_vertex_service_account_auth_config(Some(
|
||||||
|
r#"{
|
||||||
|
"client_email":"svc@example.iam.gserviceaccount.com",
|
||||||
|
"private_key":"TEST-PRIVATE-KEY",
|
||||||
|
"project_id":"demo-project",
|
||||||
|
"region":"global",
|
||||||
|
"model_regions":{"gemini-2.0-flash":"us-central1"}
|
||||||
|
}"#,
|
||||||
|
))
|
||||||
|
.expect("service account config should parse");
|
||||||
|
|
||||||
|
assert_eq!(config.client_email, "svc@example.iam.gserviceaccount.com");
|
||||||
|
assert_eq!(config.project_id, "demo-project");
|
||||||
|
assert_eq!(config.region.as_deref(), Some("global"));
|
||||||
|
assert_eq!(
|
||||||
|
config.model_regions.get("gemini-2.0-flash").map(String::as_str),
|
||||||
|
Some("us-central1")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn supports_vertex_service_account_auth_resolution() {
|
||||||
|
let mut transport = sample_transport();
|
||||||
|
transport.key.auth_type = "service_account".to_string();
|
||||||
|
transport.key.decrypted_api_key = "__placeholder__".to_string();
|
||||||
|
transport.key.decrypted_auth_config = Some(
|
||||||
|
r#"{
|
||||||
|
"client_email":"svc@example.iam.gserviceaccount.com",
|
||||||
|
"private_key":"TEST-PRIVATE-KEY",
|
||||||
|
"project_id":"demo-project"
|
||||||
|
}"#
|
||||||
|
.to_string(),
|
||||||
|
);
|
||||||
|
|
||||||
|
assert!(supports_local_vertex_service_account_auth_resolution(
|
||||||
|
&transport
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn decodes_pkcs1_service_account_private_key() {
|
||||||
|
let mut rng = OsRng;
|
||||||
|
let private_key = RsaPrivateKey::new(&mut rng, 1024)
|
||||||
|
.expect("test RSA private key should generate")
|
||||||
|
.to_pkcs1_pem(LineEnding::LF)
|
||||||
|
.expect("test RSA private key should encode as PKCS#1 PEM");
|
||||||
|
|
||||||
|
decode_vertex_service_account_private_key(private_key.as_str())
|
||||||
|
.expect("PKCS#1 private key should decode");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -27,28 +27,36 @@ pub fn looks_like_vertex_ai_host(base_url: &str) -> bool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn is_vertex_api_key_transport_context(transport: &GatewayProviderTransportSnapshot) -> bool {
|
pub fn is_vertex_api_key_transport_context(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||||
if transport
|
if is_vertex_provider_type(transport) {
|
||||||
.provider
|
|
||||||
.provider_type
|
|
||||||
.trim()
|
|
||||||
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
|
|
||||||
{
|
|
||||||
return resolve_local_auth_type_for_transport_format(transport)
|
return resolve_local_auth_type_for_transport_format(transport)
|
||||||
.eq_ignore_ascii_case("api_key");
|
.eq_ignore_ascii_case("api_key");
|
||||||
}
|
}
|
||||||
|
|
||||||
if !looks_like_vertex_ai_host(&transport.endpoint.base_url) {
|
if !is_vertex_host_format_context(transport) {
|
||||||
return false;
|
|
||||||
}
|
|
||||||
|
|
||||||
let endpoint_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
|
|
||||||
if !endpoint_api_format.starts_with("gemini:") && !endpoint_api_format.starts_with("claude:") {
|
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
resolve_local_auth_type_for_transport_format(transport).eq_ignore_ascii_case("api_key")
|
resolve_local_auth_type_for_transport_format(transport).eq_ignore_ascii_case("api_key")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn is_vertex_service_account_transport_context(
|
||||||
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
) -> bool {
|
||||||
|
if !is_vertex_provider_type(transport) && !is_vertex_host_format_context(transport) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
matches!(
|
||||||
|
resolve_local_auth_type_for_transport_format(transport).as_str(),
|
||||||
|
"service_account" | "vertex_ai"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_vertex_transport_context(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||||
|
is_vertex_api_key_transport_context(transport)
|
||||||
|
|| is_vertex_service_account_transport_context(transport)
|
||||||
|
}
|
||||||
|
|
||||||
pub fn uses_vertex_api_key_query_auth(
|
pub fn uses_vertex_api_key_query_auth(
|
||||||
transport: &GatewayProviderTransportSnapshot,
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
provider_api_format: &str,
|
provider_api_format: &str,
|
||||||
@@ -60,11 +68,28 @@ pub fn uses_vertex_api_key_query_auth(
|
|||||||
.starts_with("gemini:")
|
.starts_with("gemini:")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_vertex_provider_type(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||||
|
transport
|
||||||
|
.provider
|
||||||
|
.provider_type
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_vertex_host_format_context(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||||
|
if !looks_like_vertex_ai_host(&transport.endpoint.base_url) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
let endpoint_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
|
||||||
|
endpoint_api_format.starts_with("gemini:") || endpoint_api_format.starts_with("claude:")
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{
|
use super::{
|
||||||
is_vertex_api_key_transport_context, looks_like_vertex_ai_host,
|
is_vertex_api_key_transport_context, is_vertex_service_account_transport_context,
|
||||||
uses_vertex_api_key_query_auth,
|
is_vertex_transport_context, looks_like_vertex_ai_host, uses_vertex_api_key_query_auth,
|
||||||
};
|
};
|
||||||
use crate::snapshot::{
|
use crate::snapshot::{
|
||||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||||
@@ -150,6 +175,17 @@ mod tests {
|
|||||||
assert!(!is_vertex_api_key_transport_context(&transport));
|
assert!(!is_vertex_api_key_transport_context(&transport));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn infers_vertex_service_account_context_for_fixed_provider() {
|
||||||
|
let mut transport = sample_transport();
|
||||||
|
transport.provider.provider_type = "vertex_ai".to_string();
|
||||||
|
transport.key.auth_type = "service_account".to_string();
|
||||||
|
|
||||||
|
assert!(is_vertex_service_account_transport_context(&transport));
|
||||||
|
assert!(is_vertex_transport_context(&transport));
|
||||||
|
assert!(!is_vertex_api_key_transport_context(&transport));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn detects_vertex_query_auth_usage_for_gemini_formats() {
|
fn detects_vertex_query_auth_usage_for_gemini_formats() {
|
||||||
let transport = sample_transport();
|
let transport = sample_transport();
|
||||||
|
|||||||
@@ -4,20 +4,29 @@ mod policy;
|
|||||||
mod url;
|
mod url;
|
||||||
|
|
||||||
pub use auth::{
|
pub use auth::{
|
||||||
resolve_local_vertex_api_key_query_auth, VertexApiKeyQueryAuth, VERTEX_API_KEY_QUERY_PARAM,
|
parse_vertex_service_account_auth_config, resolve_local_vertex_api_key_query_auth,
|
||||||
|
resolve_local_vertex_service_account_auth_config,
|
||||||
|
supports_local_vertex_service_account_auth_resolution, VertexApiKeyQueryAuth,
|
||||||
|
VertexServiceAccountAuthConfig, VertexServiceAccountRefreshAdapter,
|
||||||
|
VERTEX_API_KEY_QUERY_PARAM, VERTEX_SERVICE_ACCOUNT_AUTH_HEADER,
|
||||||
|
VERTEX_SERVICE_ACCOUNT_PROVIDER_TYPE,
|
||||||
};
|
};
|
||||||
pub use context::{
|
pub use context::{
|
||||||
is_vertex_api_key_transport_context, looks_like_vertex_ai_host, uses_vertex_api_key_query_auth,
|
is_vertex_api_key_transport_context, is_vertex_service_account_transport_context,
|
||||||
|
is_vertex_transport_context, looks_like_vertex_ai_host, uses_vertex_api_key_query_auth,
|
||||||
};
|
};
|
||||||
pub use policy::{
|
pub use policy::{
|
||||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
|
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
|
||||||
|
local_vertex_gemini_transport_unsupported_reason_with_network,
|
||||||
supports_local_vertex_api_key_gemini_transport,
|
supports_local_vertex_api_key_gemini_transport,
|
||||||
supports_local_vertex_api_key_gemini_transport_with_network,
|
supports_local_vertex_api_key_gemini_transport_with_network,
|
||||||
supports_local_vertex_api_key_imagen_transport,
|
supports_local_vertex_api_key_imagen_transport,
|
||||||
supports_local_vertex_api_key_imagen_transport_with_network,
|
supports_local_vertex_api_key_imagen_transport_with_network,
|
||||||
|
supports_local_vertex_gemini_transport_with_network,
|
||||||
};
|
};
|
||||||
pub use url::{
|
pub use url::{
|
||||||
build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url,
|
build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url,
|
||||||
|
build_vertex_service_account_gemini_content_url, resolve_vertex_service_account_region,
|
||||||
VERTEX_API_KEY_BASE_URL,
|
VERTEX_API_KEY_BASE_URL,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@@ -1,10 +1,13 @@
|
|||||||
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
use super::super::snapshot::GatewayProviderTransportSnapshot;
|
||||||
use super::super::{
|
use super::super::{
|
||||||
body_rules_are_locally_supported, header_rules_are_locally_supported,
|
body_rules_are_locally_supported, header_rules_are_locally_supported,
|
||||||
|
supports_local_oauth_request_auth_resolution,
|
||||||
resolve_transport_profile, transport_profile_is_configured,
|
resolve_transport_profile, transport_profile_is_configured,
|
||||||
transport_proxy_is_locally_supported,
|
transport_proxy_is_locally_supported,
|
||||||
};
|
};
|
||||||
use super::auth::resolve_local_vertex_api_key_query_auth;
|
use super::auth::{
|
||||||
|
resolve_local_vertex_api_key_query_auth, supports_local_vertex_service_account_auth_resolution,
|
||||||
|
};
|
||||||
|
|
||||||
fn is_vertex_transport_family(transport: &GatewayProviderTransportSnapshot) -> bool {
|
fn is_vertex_transport_family(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||||
transport
|
transport
|
||||||
@@ -17,6 +20,19 @@ fn is_vertex_transport_family(transport: &GatewayProviderTransportSnapshot) -> b
|
|||||||
|
|
||||||
pub fn local_vertex_api_key_gemini_transport_unsupported_reason_with_network(
|
pub fn local_vertex_api_key_gemini_transport_unsupported_reason_with_network(
|
||||||
transport: &GatewayProviderTransportSnapshot,
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
) -> Option<&'static str> {
|
||||||
|
local_vertex_gemini_transport_unsupported_reason_with_network_impl(transport, true)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn local_vertex_gemini_transport_unsupported_reason_with_network(
|
||||||
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
) -> Option<&'static str> {
|
||||||
|
local_vertex_gemini_transport_unsupported_reason_with_network_impl(transport, false)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn local_vertex_gemini_transport_unsupported_reason_with_network_impl(
|
||||||
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
require_api_key: bool,
|
||||||
) -> Option<&'static str> {
|
) -> Option<&'static str> {
|
||||||
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
|
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
|
||||||
return if !transport.provider.is_active {
|
return if !transport.provider.is_active {
|
||||||
@@ -41,7 +57,15 @@ pub fn local_vertex_api_key_gemini_transport_unsupported_reason_with_network(
|
|||||||
if !body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref()) {
|
if !body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref()) {
|
||||||
return Some("transport_body_rules_unsupported");
|
return Some("transport_body_rules_unsupported");
|
||||||
}
|
}
|
||||||
if resolve_local_vertex_api_key_query_auth(transport).is_none() {
|
let has_api_key_auth = resolve_local_vertex_api_key_query_auth(transport).is_some();
|
||||||
|
let has_service_account_auth =
|
||||||
|
supports_local_vertex_service_account_auth_resolution(transport)
|
||||||
|
&& supports_local_oauth_request_auth_resolution(transport);
|
||||||
|
if require_api_key {
|
||||||
|
if !has_api_key_auth {
|
||||||
|
return Some("transport_auth_unavailable");
|
||||||
|
}
|
||||||
|
} else if !has_api_key_auth && !has_service_account_auth {
|
||||||
return Some("transport_auth_unavailable");
|
return Some("transport_auth_unavailable");
|
||||||
}
|
}
|
||||||
if !transport_proxy_is_locally_supported(transport) {
|
if !transport_proxy_is_locally_supported(transport) {
|
||||||
@@ -71,6 +95,12 @@ pub fn supports_local_vertex_api_key_gemini_transport_with_network(
|
|||||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport).is_none()
|
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(transport).is_none()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn supports_local_vertex_gemini_transport_with_network(
|
||||||
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
|
) -> bool {
|
||||||
|
local_vertex_gemini_transport_unsupported_reason_with_network(transport).is_none()
|
||||||
|
}
|
||||||
|
|
||||||
pub fn supports_local_vertex_api_key_imagen_transport(
|
pub fn supports_local_vertex_api_key_imagen_transport(
|
||||||
transport: &GatewayProviderTransportSnapshot,
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
@@ -159,8 +189,10 @@ mod tests {
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
|
local_vertex_api_key_gemini_transport_unsupported_reason_with_network,
|
||||||
|
local_vertex_gemini_transport_unsupported_reason_with_network,
|
||||||
supports_local_vertex_api_key_gemini_transport,
|
supports_local_vertex_api_key_gemini_transport,
|
||||||
supports_local_vertex_api_key_gemini_transport_with_network,
|
supports_local_vertex_api_key_gemini_transport_with_network,
|
||||||
|
supports_local_vertex_gemini_transport_with_network,
|
||||||
};
|
};
|
||||||
|
|
||||||
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
fn sample_transport() -> GatewayProviderTransportSnapshot {
|
||||||
@@ -242,13 +274,35 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rejects_vertex_service_account_subset() {
|
fn rejects_vertex_service_account_from_api_key_subset() {
|
||||||
let mut transport = sample_transport();
|
let mut transport = sample_transport();
|
||||||
transport.key.auth_type = "service_account".to_string();
|
transport.key.auth_type = "service_account".to_string();
|
||||||
transport.key.decrypted_auth_config = Some("{\"project_id\":\"demo-project\"}".to_string());
|
transport.key.decrypted_auth_config = Some("{\"project_id\":\"demo-project\"}".to_string());
|
||||||
assert!(!supports_local_vertex_api_key_gemini_transport(&transport));
|
assert!(!supports_local_vertex_api_key_gemini_transport(&transport));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn supports_vertex_service_account_gemini_transport_with_network() {
|
||||||
|
let mut transport = sample_transport();
|
||||||
|
transport.key.auth_type = "service_account".to_string();
|
||||||
|
transport.key.decrypted_api_key = "__placeholder__".to_string();
|
||||||
|
transport.key.decrypted_auth_config = Some(
|
||||||
|
r#"{
|
||||||
|
"client_email":"svc@example.iam.gserviceaccount.com",
|
||||||
|
"private_key":"TEST-PRIVATE-KEY",
|
||||||
|
"project_id":"demo-project"
|
||||||
|
}"#
|
||||||
|
.to_string(),
|
||||||
|
);
|
||||||
|
|
||||||
|
assert!(!supports_local_vertex_api_key_gemini_transport_with_network(
|
||||||
|
&transport
|
||||||
|
));
|
||||||
|
assert!(supports_local_vertex_gemini_transport_with_network(
|
||||||
|
&transport
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn allows_network_passthrough_for_custom_path_with_local_proxy_support() {
|
fn allows_network_passthrough_for_custom_path_with_local_proxy_support() {
|
||||||
let mut transport = sample_transport();
|
let mut transport = sample_transport();
|
||||||
@@ -272,5 +326,9 @@ mod tests {
|
|||||||
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(&transport),
|
local_vertex_api_key_gemini_transport_unsupported_reason_with_network(&transport),
|
||||||
Some("transport_auth_unavailable")
|
Some("transport_auth_unavailable")
|
||||||
);
|
);
|
||||||
|
assert_eq!(
|
||||||
|
local_vertex_gemini_transport_unsupported_reason_with_network(&transport),
|
||||||
|
Some("transport_auth_unavailable")
|
||||||
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
|
|||||||
use url::form_urlencoded;
|
use url::form_urlencoded;
|
||||||
|
|
||||||
use super::super::url::build_passthrough_path_url;
|
use super::super::url::build_passthrough_path_url;
|
||||||
|
use super::auth::VertexServiceAccountAuthConfig;
|
||||||
|
|
||||||
pub const VERTEX_API_KEY_BASE_URL: &str = "https://aiplatform.googleapis.com";
|
pub const VERTEX_API_KEY_BASE_URL: &str = "https://aiplatform.googleapis.com";
|
||||||
|
|
||||||
@@ -24,6 +25,15 @@ pub fn build_vertex_api_key_imagen_content_url(
|
|||||||
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
|
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn build_vertex_service_account_gemini_content_url(
|
||||||
|
model: &str,
|
||||||
|
stream: bool,
|
||||||
|
auth_config: &VertexServiceAccountAuthConfig,
|
||||||
|
request_query: Option<&str>,
|
||||||
|
) -> Option<String> {
|
||||||
|
build_vertex_service_account_google_model_url(model, stream, auth_config, request_query)
|
||||||
|
}
|
||||||
|
|
||||||
fn build_vertex_api_key_google_model_url(
|
fn build_vertex_api_key_google_model_url(
|
||||||
model: &str,
|
model: &str,
|
||||||
stream: bool,
|
stream: bool,
|
||||||
@@ -46,6 +56,89 @@ fn build_vertex_api_key_google_model_url(
|
|||||||
build_passthrough_path_url(VERTEX_API_KEY_BASE_URL, &path, merged_query.as_deref(), &[])
|
build_passthrough_path_url(VERTEX_API_KEY_BASE_URL, &path, merged_query.as_deref(), &[])
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn build_vertex_service_account_google_model_url(
|
||||||
|
model: &str,
|
||||||
|
stream: bool,
|
||||||
|
auth_config: &VertexServiceAccountAuthConfig,
|
||||||
|
request_query: Option<&str>,
|
||||||
|
) -> Option<String> {
|
||||||
|
let trimmed_model = model.trim();
|
||||||
|
let project_id = auth_config.project_id.trim();
|
||||||
|
if trimmed_model.is_empty() || project_id.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
let region = resolve_vertex_service_account_region(trimmed_model, auth_config);
|
||||||
|
let action = if stream {
|
||||||
|
"streamGenerateContent"
|
||||||
|
} else {
|
||||||
|
"generateContent"
|
||||||
|
};
|
||||||
|
let base_url = if region == "global" {
|
||||||
|
VERTEX_API_KEY_BASE_URL.to_string()
|
||||||
|
} else {
|
||||||
|
format!("https://{region}-aiplatform.googleapis.com")
|
||||||
|
};
|
||||||
|
let path = format!(
|
||||||
|
"/v1/projects/{project_id}/locations/{region}/publishers/google/models/{trimmed_model}:{action}"
|
||||||
|
);
|
||||||
|
let merged_query = build_vertex_service_account_query(request_query, stream);
|
||||||
|
build_passthrough_path_url(&base_url, &path, merged_query.as_deref(), &[])
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn resolve_vertex_service_account_region(
|
||||||
|
model: &str,
|
||||||
|
auth_config: &VertexServiceAccountAuthConfig,
|
||||||
|
) -> String {
|
||||||
|
let trimmed_model = model.trim();
|
||||||
|
if let Some(region) = auth_config
|
||||||
|
.model_regions
|
||||||
|
.get(trimmed_model)
|
||||||
|
.map(String::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
{
|
||||||
|
return region.to_string();
|
||||||
|
}
|
||||||
|
if let Some(region) = default_vertex_model_region(trimmed_model) {
|
||||||
|
return region.to_string();
|
||||||
|
}
|
||||||
|
if let Some(region) = auth_config
|
||||||
|
.region
|
||||||
|
.as_deref()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
{
|
||||||
|
return region.to_string();
|
||||||
|
}
|
||||||
|
"global".to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_vertex_model_region(model: &str) -> Option<&'static str> {
|
||||||
|
if model.starts_with("gemini-3.") || model == "gemini-3-pro-image-preview" {
|
||||||
|
return Some("global");
|
||||||
|
}
|
||||||
|
if matches!(
|
||||||
|
model,
|
||||||
|
"gemini-2.0-flash"
|
||||||
|
| "gemini-2.0-flash-exp"
|
||||||
|
| "gemini-2.0-flash-001"
|
||||||
|
| "gemini-2.0-pro-exp"
|
||||||
|
| "gemini-2.0-flash-exp-image-generation"
|
||||||
|
| "gemini-1.5-pro"
|
||||||
|
| "gemini-1.5-pro-001"
|
||||||
|
| "gemini-1.5-pro-002"
|
||||||
|
| "gemini-1.5-flash"
|
||||||
|
| "gemini-1.5-flash-001"
|
||||||
|
| "gemini-1.5-flash-002"
|
||||||
|
| "imagen-3.0-generate-001"
|
||||||
|
| "imagen-3.0-fast-generate-001"
|
||||||
|
) {
|
||||||
|
return Some("us-central1");
|
||||||
|
}
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
fn build_vertex_api_key_query(
|
fn build_vertex_api_key_query(
|
||||||
api_key: &str,
|
api_key: &str,
|
||||||
request_query: Option<&str>,
|
request_query: Option<&str>,
|
||||||
@@ -73,6 +166,29 @@ fn build_vertex_api_key_query(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn build_vertex_service_account_query(request_query: Option<&str>, stream: bool) -> Option<String> {
|
||||||
|
let mut merged = BTreeMap::new();
|
||||||
|
merge_query_string(&mut merged, request_query);
|
||||||
|
merged.remove("beta");
|
||||||
|
merged.remove("key");
|
||||||
|
if stream {
|
||||||
|
merged
|
||||||
|
.entry("alt".to_string())
|
||||||
|
.or_insert_with(|| "sse".to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut serializer = form_urlencoded::Serializer::new(String::new());
|
||||||
|
for (key, value) in merged {
|
||||||
|
serializer.append_pair(&key, &value);
|
||||||
|
}
|
||||||
|
let query = serializer.finish();
|
||||||
|
if query.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(query)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn merge_query_string(out: &mut BTreeMap<String, String>, query: Option<&str>) {
|
fn merge_query_string(out: &mut BTreeMap<String, String>, query: Option<&str>) {
|
||||||
let Some(query) = query.map(str::trim).filter(|value| !value.is_empty()) else {
|
let Some(query) = query.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||||
return;
|
return;
|
||||||
@@ -85,7 +201,13 @@ fn merge_query_string(out: &mut BTreeMap<String, String>, query: Option<&str>) {
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url};
|
use std::collections::BTreeMap;
|
||||||
|
|
||||||
|
use super::{
|
||||||
|
build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url,
|
||||||
|
build_vertex_service_account_gemini_content_url,
|
||||||
|
};
|
||||||
|
use crate::vertex::VertexServiceAccountAuthConfig;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn builds_vertex_gemini_api_key_stream_url() {
|
fn builds_vertex_gemini_api_key_stream_url() {
|
||||||
@@ -118,4 +240,57 @@ mod tests {
|
|||||||
)
|
)
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn builds_vertex_service_account_gemini_sync_url() {
|
||||||
|
let auth_config = VertexServiceAccountAuthConfig {
|
||||||
|
client_email: "svc@example.iam.gserviceaccount.com".to_string(),
|
||||||
|
private_key: "not-used".to_string(),
|
||||||
|
project_id: "demo-project".to_string(),
|
||||||
|
token_uri: "https://oauth2.googleapis.com/token".to_string(),
|
||||||
|
region: None,
|
||||||
|
model_regions: BTreeMap::new(),
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
build_vertex_service_account_gemini_content_url(
|
||||||
|
"gemini-3.1-pro-preview",
|
||||||
|
false,
|
||||||
|
&auth_config,
|
||||||
|
Some("foo=bar&beta=1&key=client-key")
|
||||||
|
)
|
||||||
|
.as_deref(),
|
||||||
|
Some(
|
||||||
|
"https://aiplatform.googleapis.com/v1/projects/demo-project/locations/global/publishers/google/models/gemini-3.1-pro-preview:generateContent?foo=bar"
|
||||||
|
)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn builds_vertex_service_account_gemini_stream_url_with_model_region_override() {
|
||||||
|
let auth_config = VertexServiceAccountAuthConfig {
|
||||||
|
client_email: "svc@example.iam.gserviceaccount.com".to_string(),
|
||||||
|
private_key: "not-used".to_string(),
|
||||||
|
project_id: "demo-project".to_string(),
|
||||||
|
token_uri: "https://oauth2.googleapis.com/token".to_string(),
|
||||||
|
region: Some("global".to_string()),
|
||||||
|
model_regions: BTreeMap::from([(
|
||||||
|
"gemini-2.0-flash".to_string(),
|
||||||
|
"us-central1".to_string(),
|
||||||
|
)]),
|
||||||
|
};
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
build_vertex_service_account_gemini_content_url(
|
||||||
|
"gemini-2.0-flash",
|
||||||
|
true,
|
||||||
|
&auth_config,
|
||||||
|
Some("foo=bar")
|
||||||
|
)
|
||||||
|
.as_deref(),
|
||||||
|
Some(
|
||||||
|
"https://us-central1-aiplatform.googleapis.com/v1/projects/demo-project/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent?alt=sse&foo=bar"
|
||||||
|
)
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user