mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
fix(gateway): harden Gemini endpoint routing
This commit is contained in:
@@ -66,6 +66,10 @@ pub fn from_raw(body_json: &Value) -> Option<CanonicalResponse> {
|
||||
extensions: Default::default(),
|
||||
});
|
||||
}
|
||||
outputs.retain(gemini_response_output_has_visible_content);
|
||||
if outputs.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let content = outputs
|
||||
.first()
|
||||
.map(|output| output.content.clone())
|
||||
@@ -108,6 +112,18 @@ pub fn from_raw(body_json: &Value) -> Option<CanonicalResponse> {
|
||||
Some(canonical)
|
||||
}
|
||||
|
||||
fn gemini_response_output_has_visible_content(output: &CanonicalResponseOutput) -> bool {
|
||||
output.content.iter().any(|block| match block {
|
||||
CanonicalContentBlock::Text { text, .. } => !text.trim().is_empty(),
|
||||
CanonicalContentBlock::ToolUse { .. }
|
||||
| CanonicalContentBlock::ToolResult { .. }
|
||||
| CanonicalContentBlock::Image { .. }
|
||||
| CanonicalContentBlock::File { .. }
|
||||
| CanonicalContentBlock::Audio { .. } => true,
|
||||
CanonicalContentBlock::Thinking { .. } | CanonicalContentBlock::Unknown { .. } => false,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value) -> Option<Value> {
|
||||
let mut response = canonical_to_gemini_response(canonical, report_context)?;
|
||||
if let Some(object) = response.as_object_mut() {
|
||||
@@ -353,3 +369,72 @@ fn canonical_usage_to_gemini_usage_metadata(usage: &CanonicalUsage) -> Value {
|
||||
}
|
||||
Value::Object(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::CanonicalContentBlock;
|
||||
|
||||
#[test]
|
||||
fn gemini_response_without_visible_parts_is_not_success() {
|
||||
let body = json!({
|
||||
"candidates": [{
|
||||
"content": {"role": "model"},
|
||||
"finishReason": "MAX_TOKENS"
|
||||
}],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 8,
|
||||
"candidatesTokenCount": 1,
|
||||
"thoughtsTokenCount": 25,
|
||||
"totalTokenCount": 34
|
||||
},
|
||||
"modelVersion": "gemini-3-flash-preview",
|
||||
"responseId": "resp-empty"
|
||||
});
|
||||
|
||||
assert!(from_raw(&body).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_response_with_only_thought_parts_is_not_success() {
|
||||
let body = json!({
|
||||
"candidates": [{
|
||||
"content": {
|
||||
"role": "model",
|
||||
"parts": [{"text": "hidden plan", "thought": true}]
|
||||
},
|
||||
"finishReason": "MAX_TOKENS"
|
||||
}],
|
||||
"modelVersion": "gemini-3-flash-preview",
|
||||
"responseId": "resp-thought-only"
|
||||
});
|
||||
|
||||
assert!(from_raw(&body).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_response_with_function_call_is_visible_output() {
|
||||
let body = json!({
|
||||
"candidates": [{
|
||||
"content": {
|
||||
"role": "model",
|
||||
"parts": [{
|
||||
"functionCall": {
|
||||
"name": "lookup",
|
||||
"args": {"query": "weather"}
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finishReason": "STOP"
|
||||
}],
|
||||
"modelVersion": "gemini-3-flash-preview",
|
||||
"responseId": "resp-tool"
|
||||
});
|
||||
|
||||
let canonical = from_raw(&body).expect("function call should be visible output");
|
||||
assert!(matches!(
|
||||
canonical.content.first(),
|
||||
Some(CanonicalContentBlock::ToolUse { name, .. }) if name == "lookup"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -317,11 +317,15 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
|
||||
auth_type == "oauth" && api_format == "gemini:generate_content"
|
||||
}
|
||||
"vertex_ai" => {
|
||||
(auth_type == "api_key" && api_format == "gemini:generate_content")
|
||||
(auth_type == "api_key"
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"gemini:generate_content" | "gemini:embedding"
|
||||
))
|
||||
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"claude:messages" | "gemini:generate_content"
|
||||
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
|
||||
))
|
||||
}
|
||||
_ => auth_type != "oauth",
|
||||
|
||||
@@ -446,11 +446,15 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
|
||||
auth_type == "oauth" && api_format == "gemini:generate_content"
|
||||
}
|
||||
"vertex_ai" => {
|
||||
(auth_type == "api_key" && api_format == "gemini:generate_content")
|
||||
(auth_type == "api_key"
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"gemini:generate_content" | "gemini:embedding"
|
||||
))
|
||||
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"claude:messages" | "gemini:generate_content"
|
||||
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
|
||||
))
|
||||
}
|
||||
_ => auth_type != "oauth",
|
||||
|
||||
@@ -113,11 +113,11 @@ WHERE p.is_active = TRUE
|
||||
AND (
|
||||
(
|
||||
LOWER(BTRIM(pak.auth_type)) = 'api_key'
|
||||
AND LOWER($3) = 'gemini:generate_content'
|
||||
AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding')
|
||||
)
|
||||
OR (
|
||||
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
|
||||
AND LOWER($3) IN ('claude:messages', 'gemini:generate_content')
|
||||
AND LOWER($3) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
|
||||
)
|
||||
)
|
||||
)
|
||||
@@ -296,11 +296,11 @@ WHERE p.is_active = TRUE
|
||||
AND (
|
||||
(
|
||||
LOWER(BTRIM(pak.auth_type)) = 'api_key'
|
||||
AND LOWER($4) = 'gemini:generate_content'
|
||||
AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding')
|
||||
)
|
||||
OR (
|
||||
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
|
||||
AND LOWER($4) IN ('claude:messages', 'gemini:generate_content')
|
||||
AND LOWER($4) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
|
||||
)
|
||||
)
|
||||
)
|
||||
@@ -478,11 +478,11 @@ WHERE p.is_active = TRUE
|
||||
AND (
|
||||
(
|
||||
LOWER(BTRIM(pak.auth_type)) = 'api_key'
|
||||
AND LOWER($6) = 'gemini:generate_content'
|
||||
AND LOWER($6) IN ('gemini:generate_content', 'gemini:embedding')
|
||||
)
|
||||
OR (
|
||||
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
|
||||
AND LOWER($6) IN ('claude:messages', 'gemini:generate_content')
|
||||
AND LOWER($6) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
|
||||
)
|
||||
)
|
||||
)
|
||||
@@ -1287,6 +1287,22 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_selection_sql_allows_vertex_embedding_auth() {
|
||||
let requested_model_sql = requested_model_selection_sql();
|
||||
for sql in [
|
||||
LIST_FOR_EXACT_API_FORMAT_SQL,
|
||||
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
|
||||
LIST_POOL_KEYS_FOR_GROUP_SQL,
|
||||
requested_model_sql.as_str(),
|
||||
] {
|
||||
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'vertex_ai'"));
|
||||
assert!(sql.contains("gemini:embedding"));
|
||||
assert!(sql.contains("gemini:generate_content"));
|
||||
assert!(sql.contains("claude:messages"));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requested_model_selection_page_sql_adds_limit_and_offset() {
|
||||
let sql = requested_model_selection_page_sql();
|
||||
|
||||
@@ -446,11 +446,15 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
|
||||
auth_type == "oauth" && api_format == "gemini:generate_content"
|
||||
}
|
||||
"vertex_ai" => {
|
||||
(auth_type == "api_key" && api_format == "gemini:generate_content")
|
||||
(auth_type == "api_key"
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"gemini:generate_content" | "gemini:embedding"
|
||||
))
|
||||
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
|
||||
&& matches!(
|
||||
api_format.as_str(),
|
||||
"claude:messages" | "gemini:generate_content"
|
||||
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
|
||||
))
|
||||
}
|
||||
_ => auth_type != "oauth",
|
||||
|
||||
@@ -423,7 +423,7 @@ mod tests {
|
||||
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
|
||||
request_conversion_direct_auth, request_conversion_enabled_for_transport,
|
||||
request_conversion_transport_supported, request_pair_allowed_for_transport,
|
||||
CandidateTransportPolicyFacts,
|
||||
request_pair_direct_auth, CandidateTransportPolicyFacts,
|
||||
};
|
||||
use aether_ai_formats::formats::matrix::RequestConversionKind;
|
||||
use serde_json::json;
|
||||
@@ -593,6 +593,21 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_gemini_embedding_transport_supports_openai_embedding_conversion() {
|
||||
let transport = transport_snapshot("vertex_ai", "gemini:embedding", "api_key", true, None);
|
||||
|
||||
assert!(request_pair_allowed_for_transport(
|
||||
&transport,
|
||||
"openai:embedding",
|
||||
"gemini:embedding"
|
||||
));
|
||||
assert_eq!(
|
||||
request_pair_direct_auth(&transport, "gemini:embedding"),
|
||||
Some(("key".to_string(), "secret".to_string()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kiro_claude_messages_transport_supports_cross_format_conversion_via_envelope() {
|
||||
let transport = transport_snapshot("kiro", "claude:messages", "bearer", true, None);
|
||||
|
||||
@@ -47,6 +47,7 @@ pub enum ProviderApiFormatInheritance {
|
||||
None,
|
||||
OAuth,
|
||||
OAuthOrBearer,
|
||||
OAuthOrServiceAccount,
|
||||
OAuthOrConfiguredBearer,
|
||||
}
|
||||
|
||||
@@ -61,6 +62,9 @@ impl ProviderApiFormatInheritance {
|
||||
Self::None => false,
|
||||
Self::OAuth => auth_type == "oauth",
|
||||
Self::OAuthOrBearer => auth_type == "oauth" || auth_type == "bearer",
|
||||
Self::OAuthOrServiceAccount => {
|
||||
auth_type == "oauth" || auth_type == "service_account" || auth_type == "vertex_ai"
|
||||
}
|
||||
Self::OAuthOrConfiguredBearer => {
|
||||
auth_type == "oauth"
|
||||
|| auth_type == "bearer"
|
||||
@@ -221,11 +225,12 @@ const GEMINI_CLI_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy {
|
||||
};
|
||||
const VERTEX_AI_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy {
|
||||
fixed_provider: true,
|
||||
api_format_inheritance: ProviderApiFormatInheritance::OAuth,
|
||||
api_format_inheritance: ProviderApiFormatInheritance::OAuthOrServiceAccount,
|
||||
enable_format_conversion_by_default: true,
|
||||
supports_model_fetch: false,
|
||||
supports_local_openai_chat_transport: false,
|
||||
supports_local_same_format_transport: false,
|
||||
local_embedding_support: ProviderLocalEmbeddingSupport::Gemini,
|
||||
..STANDARD_RUNTIME_POLICY
|
||||
};
|
||||
const ANTIGRAVITY_RUNTIME_POLICY: ProviderRuntimePolicy = ProviderRuntimePolicy {
|
||||
@@ -329,6 +334,12 @@ const VERTEX_AI_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTe
|
||||
custom_path: None,
|
||||
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
|
||||
},
|
||||
FixedProviderEndpointTemplate {
|
||||
item_key: "gemini:embedding",
|
||||
api_format: "gemini:embedding",
|
||||
custom_path: None,
|
||||
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
|
||||
},
|
||||
FixedProviderEndpointTemplate {
|
||||
item_key: "claude:messages",
|
||||
api_format: "claude:messages",
|
||||
@@ -613,6 +624,11 @@ mod tests {
|
||||
"bearer",
|
||||
Some("{}")
|
||||
));
|
||||
assert!(fixed_provider_key_inherits_api_formats(
|
||||
"vertex_ai",
|
||||
"service_account",
|
||||
None
|
||||
));
|
||||
assert!(!fixed_provider_key_inherits_api_formats(
|
||||
"kiro", "bearer", None
|
||||
));
|
||||
@@ -681,6 +697,7 @@ mod tests {
|
||||
("custom", "openai:embedding"),
|
||||
("gemini", "gemini:embedding"),
|
||||
("google", "gemini:embedding"),
|
||||
("vertex_ai", "gemini:embedding"),
|
||||
("jina", "jina:embedding"),
|
||||
("doubao", "doubao:embedding"),
|
||||
("volcengine", "doubao:embedding"),
|
||||
@@ -694,6 +711,7 @@ mod tests {
|
||||
for (provider_type, api_format) in [
|
||||
("openai", "gemini:embedding"),
|
||||
("gemini", "openai:embedding"),
|
||||
("vertex_ai", "openai:embedding"),
|
||||
("jina", "doubao:embedding"),
|
||||
("doubao", "jina:embedding"),
|
||||
("claude_code", "openai:embedding"),
|
||||
@@ -710,4 +728,28 @@ mod tests {
|
||||
"GEMINI:EMBEDDING"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_fixed_provider_template_includes_gemini_embedding_endpoint() {
|
||||
let template =
|
||||
fixed_provider_template("vertex_ai").expect("vertex_ai template should exist");
|
||||
|
||||
assert_eq!(
|
||||
template
|
||||
.endpoints
|
||||
.iter()
|
||||
.map(|item| item.api_format)
|
||||
.collect::<Vec<_>>(),
|
||||
vec![
|
||||
"gemini:generate_content",
|
||||
"gemini:embedding",
|
||||
"claude:messages",
|
||||
]
|
||||
);
|
||||
|
||||
assert!(
|
||||
fixed_provider_endpoint_template_by_api_format("vertex_ai", "gemini:embedding")
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,9 @@ use crate::url::{
|
||||
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
|
||||
};
|
||||
use crate::vertex::{
|
||||
build_vertex_api_key_gemini_content_url, build_vertex_service_account_gemini_content_url,
|
||||
build_vertex_api_key_gemini_content_url, build_vertex_api_key_gemini_embedding_url,
|
||||
build_vertex_service_account_gemini_content_url,
|
||||
build_vertex_service_account_gemini_embedding_url, is_vertex_transport_context,
|
||||
resolve_local_vertex_api_key_query_auth, resolve_local_vertex_service_account_auth_config,
|
||||
};
|
||||
|
||||
@@ -62,13 +64,20 @@ fn build_transport_request_url_inner(
|
||||
params: TransportRequestUrlParams<'_>,
|
||||
gemini_embedding_batch: bool,
|
||||
) -> Option<String> {
|
||||
let provider_api_format = params.provider_api_format.trim().to_ascii_lowercase();
|
||||
let normalized_provider_api_format =
|
||||
aether_ai_formats::normalize_api_format_alias(&provider_api_format);
|
||||
if normalized_provider_api_format == "gemini:embedding"
|
||||
&& gemini_embedding_batch
|
||||
&& is_vertex_transport_context(transport)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(url) = build_transport_hook_url(transport, params) {
|
||||
return Some(url);
|
||||
}
|
||||
|
||||
let provider_api_format = params.provider_api_format.trim().to_ascii_lowercase();
|
||||
let normalized_provider_api_format =
|
||||
aether_ai_formats::normalize_api_format_alias(&provider_api_format);
|
||||
let custom_path = transport
|
||||
.endpoint
|
||||
.custom_path
|
||||
@@ -281,25 +290,42 @@ fn build_transport_hook_url(
|
||||
));
|
||||
}
|
||||
|
||||
if aether_ai_formats::normalize_api_format_alias(params.provider_api_format)
|
||||
== "gemini:generate_content"
|
||||
{
|
||||
if let Some(auth) = resolve_local_vertex_api_key_query_auth(transport) {
|
||||
return build_vertex_api_key_gemini_content_url(
|
||||
params.mapped_model?,
|
||||
params.upstream_is_stream,
|
||||
&auth.value,
|
||||
params.request_query,
|
||||
);
|
||||
match aether_ai_formats::normalize_api_format_alias(params.provider_api_format).as_str() {
|
||||
"gemini:generate_content" => {
|
||||
if let Some(auth) = resolve_local_vertex_api_key_query_auth(transport) {
|
||||
return build_vertex_api_key_gemini_content_url(
|
||||
params.mapped_model?,
|
||||
params.upstream_is_stream,
|
||||
&auth.value,
|
||||
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 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,
|
||||
);
|
||||
"gemini:embedding" => {
|
||||
if let Some(auth) = resolve_local_vertex_api_key_query_auth(transport) {
|
||||
return build_vertex_api_key_gemini_embedding_url(
|
||||
params.mapped_model?,
|
||||
&auth.value,
|
||||
params.request_query,
|
||||
);
|
||||
}
|
||||
if let Some(auth_config) = resolve_local_vertex_service_account_auth_config(transport) {
|
||||
return build_vertex_service_account_gemini_embedding_url(
|
||||
params.mapped_model?,
|
||||
&auth_config,
|
||||
params.request_query,
|
||||
);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
if is_antigravity_provider_transport(transport) {
|
||||
@@ -629,6 +655,91 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn uses_vertex_service_account_hook_for_gemini_embedding_url() {
|
||||
let mut transport = sample_transport(
|
||||
"vertex_ai",
|
||||
"gemini:embedding",
|
||||
"https://aiplatform.googleapis.com",
|
||||
None,
|
||||
);
|
||||
transport.endpoint.endpoint_kind = Some("embedding".to_string());
|
||||
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 provider_request_body = json!({
|
||||
"content": {"parts": [{"text": "hello"}]}
|
||||
});
|
||||
let url = build_transport_request_url_for_request_body(
|
||||
&transport,
|
||||
TransportRequestUrlParams {
|
||||
provider_api_format: "gemini:embedding",
|
||||
mapped_model: Some("gemini-embedding-2"),
|
||||
upstream_is_stream: false,
|
||||
request_query: Some("foo=bar&beta=1"),
|
||||
kiro_api_region: None,
|
||||
},
|
||||
Some(&provider_request_body),
|
||||
)
|
||||
.expect("vertex embedding service account hook url");
|
||||
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://aiplatform.googleapis.com/v1/projects/demo-project/locations/global/publishers/google/models/gemini-embedding-2:embedContent?foo=bar"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn vertex_gemini_embedding_batch_request_does_not_use_gemini_api_batch_endpoint() {
|
||||
let mut transport = sample_transport(
|
||||
"vertex_ai",
|
||||
"gemini:embedding",
|
||||
"https://aiplatform.googleapis.com",
|
||||
None,
|
||||
);
|
||||
transport.endpoint.endpoint_kind = Some("embedding".to_string());
|
||||
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 batch_body = json!({
|
||||
"requests": [
|
||||
{
|
||||
"model": "models/gemini-embedding-2",
|
||||
"content": {"parts": [{"text": "alpha"}]}
|
||||
}
|
||||
]
|
||||
});
|
||||
|
||||
assert!(build_transport_request_url_for_request_body(
|
||||
&transport,
|
||||
TransportRequestUrlParams {
|
||||
provider_api_format: "gemini:embedding",
|
||||
mapped_model: Some("gemini-embedding-2"),
|
||||
upstream_is_stream: false,
|
||||
request_query: None,
|
||||
kiro_api_region: None,
|
||||
},
|
||||
Some(&batch_body),
|
||||
)
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_openai_responses_url_for_formal_format_name() {
|
||||
let transport = sample_transport(
|
||||
|
||||
@@ -24,8 +24,9 @@ pub use policy::{
|
||||
supports_local_vertex_gemini_transport_with_network,
|
||||
};
|
||||
pub use 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,
|
||||
build_vertex_api_key_gemini_content_url, build_vertex_api_key_gemini_embedding_url,
|
||||
build_vertex_api_key_imagen_content_url, build_vertex_service_account_gemini_content_url,
|
||||
build_vertex_service_account_gemini_embedding_url, resolve_vertex_service_account_region,
|
||||
VERTEX_API_KEY_BASE_URL,
|
||||
};
|
||||
|
||||
|
||||
@@ -42,9 +42,12 @@ fn local_vertex_gemini_transport_unsupported_reason_with_network_impl(
|
||||
Some("key_inactive")
|
||||
};
|
||||
}
|
||||
if aether_ai_formats::normalize_api_format_alias(&transport.endpoint.api_format)
|
||||
!= "gemini:generate_content"
|
||||
{
|
||||
let endpoint_api_format =
|
||||
aether_ai_formats::normalize_api_format_alias(&transport.endpoint.api_format);
|
||||
if !matches!(
|
||||
endpoint_api_format.as_str(),
|
||||
"gemini:generate_content" | "gemini:embedding"
|
||||
) {
|
||||
return Some("transport_api_format_mismatch");
|
||||
}
|
||||
if !is_vertex_transport_family(transport) {
|
||||
@@ -299,6 +302,28 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supports_vertex_service_account_gemini_embedding_transport_with_network() {
|
||||
let mut transport = sample_transport();
|
||||
transport.endpoint.api_format = "gemini:embedding".to_string();
|
||||
transport.endpoint.endpoint_kind = Some("embedding".to_string());
|
||||
transport.key.api_formats = Some(vec!["gemini:embedding".to_string()]);
|
||||
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_gemini_transport_with_network(
|
||||
&transport
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_network_passthrough_for_custom_path_with_local_proxy_support() {
|
||||
let mut transport = sample_transport();
|
||||
|
||||
@@ -13,7 +13,12 @@ pub fn build_vertex_api_key_gemini_content_url(
|
||||
api_key: &str,
|
||||
request_query: Option<&str>,
|
||||
) -> Option<String> {
|
||||
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
|
||||
let action = if stream {
|
||||
"streamGenerateContent"
|
||||
} else {
|
||||
"generateContent"
|
||||
};
|
||||
build_vertex_api_key_google_model_url(model, action, stream, api_key, request_query)
|
||||
}
|
||||
|
||||
pub fn build_vertex_api_key_imagen_content_url(
|
||||
@@ -22,7 +27,20 @@ pub fn build_vertex_api_key_imagen_content_url(
|
||||
api_key: &str,
|
||||
request_query: Option<&str>,
|
||||
) -> Option<String> {
|
||||
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
|
||||
let action = if stream {
|
||||
"streamGenerateContent"
|
||||
} else {
|
||||
"generateContent"
|
||||
};
|
||||
build_vertex_api_key_google_model_url(model, action, stream, api_key, request_query)
|
||||
}
|
||||
|
||||
pub fn build_vertex_api_key_gemini_embedding_url(
|
||||
model: &str,
|
||||
api_key: &str,
|
||||
request_query: Option<&str>,
|
||||
) -> Option<String> {
|
||||
build_vertex_api_key_google_model_url(model, "embedContent", false, api_key, request_query)
|
||||
}
|
||||
|
||||
pub fn build_vertex_service_account_gemini_content_url(
|
||||
@@ -31,56 +49,69 @@ pub fn build_vertex_service_account_gemini_content_url(
|
||||
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(
|
||||
model: &str,
|
||||
stream: bool,
|
||||
api_key: &str,
|
||||
request_query: Option<&str>,
|
||||
) -> Option<String> {
|
||||
let trimmed_model = model.trim();
|
||||
let trimmed_api_key = api_key.trim();
|
||||
if trimmed_model.is_empty() || trimmed_api_key.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let action = if stream {
|
||||
"streamGenerateContent"
|
||||
} else {
|
||||
"generateContent"
|
||||
};
|
||||
let path = format!("/v1/publishers/google/models/{trimmed_model}:{action}");
|
||||
build_vertex_service_account_google_model_url(model, action, stream, auth_config, request_query)
|
||||
}
|
||||
|
||||
pub fn build_vertex_service_account_gemini_embedding_url(
|
||||
model: &str,
|
||||
auth_config: &VertexServiceAccountAuthConfig,
|
||||
request_query: Option<&str>,
|
||||
) -> Option<String> {
|
||||
build_vertex_service_account_google_model_url(
|
||||
model,
|
||||
"embedContent",
|
||||
false,
|
||||
auth_config,
|
||||
request_query,
|
||||
)
|
||||
}
|
||||
|
||||
fn build_vertex_api_key_google_model_url(
|
||||
model: &str,
|
||||
action: &str,
|
||||
stream: bool,
|
||||
api_key: &str,
|
||||
request_query: Option<&str>,
|
||||
) -> Option<String> {
|
||||
let trimmed_model = model.trim();
|
||||
let trimmed_action = action.trim();
|
||||
let trimmed_api_key = api_key.trim();
|
||||
if trimmed_model.is_empty() || trimmed_action.is_empty() || trimmed_api_key.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let path = format!("/v1/publishers/google/models/{trimmed_model}:{trimmed_action}");
|
||||
let merged_query = build_vertex_api_key_query(trimmed_api_key, request_query, stream);
|
||||
build_passthrough_path_url(VERTEX_API_KEY_BASE_URL, &path, merged_query.as_deref(), &[])
|
||||
}
|
||||
|
||||
fn build_vertex_service_account_google_model_url(
|
||||
model: &str,
|
||||
action: &str,
|
||||
stream: bool,
|
||||
auth_config: &VertexServiceAccountAuthConfig,
|
||||
request_query: Option<&str>,
|
||||
) -> Option<String> {
|
||||
let trimmed_model = model.trim();
|
||||
let trimmed_action = action.trim();
|
||||
let project_id = auth_config.project_id.trim();
|
||||
if trimmed_model.is_empty() || project_id.is_empty() {
|
||||
if trimmed_model.is_empty() || trimmed_action.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}"
|
||||
"/v1/projects/{project_id}/locations/{region}/publishers/google/models/{trimmed_model}:{trimmed_action}"
|
||||
);
|
||||
let merged_query = build_vertex_service_account_query(request_query, stream);
|
||||
build_passthrough_path_url(&base_url, &path, merged_query.as_deref(), &[])
|
||||
|
||||
Reference in New Issue
Block a user