i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge PR #488 and #489

Browse Source
This commit is contained in:
fawney19
2026-05-19 03:16:35 +08:00
parent cd1aa92931 be2e8e594c
commit 69af74b1e0
77 changed files with 4051 additions and 333 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/aether-admin/src/system.rs
View File

@@ -736,7 +736,7 @@ const ADMIN_API_FORMAT_DEFINITIONS: &[AdminApiFormatDefinition] = &[
AdminApiFormatDefinition {
value: "gemini:embedding",
label: "Gemini Embedding",
default_path: "/v1/embeddings",
default_path: "/v1beta/models/{model}:{action}",
aliases: &["gemini_embedding"],
},
AdminApiFormatDefinition {

43
crates/aether-ai-formats/src/api.rs
View File

@@ -16,18 +16,21 @@ pub use crate::contracts::{
GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_PLAN_KIND,
GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND, GEMINI_CLI_SYNC_ERROR_REPORT_KIND,
GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND, GEMINI_CLI_SYNC_PLAN_KIND,
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND,
OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_ERROR_REPORT_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_EMBEDDING_SYNC_ERROR_REPORT_KIND, OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND,
OPENAI_IMAGE_SYNC_ERROR_REPORT_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND, OPENAI_RESPONSES_STREAM_PLAN_KIND,
@@ -139,18 +142,22 @@ pub use crate::formats::{
resolve_stream_spec as resolve_gemini_stream_spec,
resolve_sync_spec as resolve_gemini_sync_spec,
},
openai::responses::{
codex::{
apply_codex_openai_responses_chat_body_edits,
apply_codex_openai_responses_special_body_edits,
apply_codex_openai_responses_special_headers,
apply_openai_responses_compact_special_body_edits, CODEX_OPENAI_IMAGE_DEFAULT_MODEL,
CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT, CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL,
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
},
spec::{
resolve_stream_spec as resolve_openai_responses_stream_spec,
resolve_sync_spec as resolve_openai_responses_sync_spec, LocalOpenAiResponsesSpec,
openai::{
embedding::spec::resolve_sync_spec as resolve_openai_embedding_sync_spec,
responses::{
codex::{
apply_codex_openai_responses_chat_body_edits,
apply_codex_openai_responses_special_body_edits,
apply_codex_openai_responses_special_headers,
apply_openai_responses_compact_special_body_edits,
CODEX_OPENAI_IMAGE_DEFAULT_MODEL, CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT,
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_MODEL,
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
},
spec::{
resolve_stream_spec as resolve_openai_responses_stream_spec,
resolve_sync_spec as resolve_openai_responses_sync_spec, LocalOpenAiResponsesSpec,
},
},
},
shared::{

14
crates/aether-ai-formats/src/contracts/mod.rs
View File

@@ -14,9 +14,9 @@ pub use plan_kinds::{
is_openai_responses_stream_plan_kind, is_openai_responses_sync_plan_kind,
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
@@ -36,9 +36,11 @@ pub use report_kinds::{
GEMINI_CHAT_SYNC_ERROR_REPORT_KIND, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND,
GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND, GEMINI_CLI_STREAM_SUCCESS_REPORT_KIND,
GEMINI_CLI_SYNC_ERROR_REPORT_KIND, GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND,
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND,
OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND, OPENAI_CHAT_SYNC_ERROR_REPORT_KIND,
OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND,
GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND, GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND,
GEMINI_VIDEO_CREATE_SYNC_FINALIZE_REPORT_KIND, OPENAI_CHAT_STREAM_SUCCESS_REPORT_KIND,
OPENAI_CHAT_SYNC_ERROR_REPORT_KIND, OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND,
OPENAI_CHAT_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_ERROR_REPORT_KIND,
OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND, OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND,
OPENAI_IMAGE_STREAM_SUCCESS_REPORT_KIND, OPENAI_IMAGE_SYNC_ERROR_REPORT_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND, OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_SUCCESS_REPORT_KIND,

1
crates/aether-ai-formats/src/contracts/plan_kinds.rs
View File

@@ -22,6 +22,7 @@ pub const OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND: &str = "openai_video_create_sync";
pub const OPENAI_CHAT_SYNC_PLAN_KIND: &str = "openai_chat_sync";
pub const OPENAI_EMBEDDING_SYNC_PLAN_KIND: &str = "openai_embedding_sync";
pub const OPENAI_RERANK_SYNC_PLAN_KIND: &str = "openai_rerank_sync";
pub const GEMINI_EMBEDDING_SYNC_PLAN_KIND: &str = "gemini_embedding_sync";
pub const OPENAI_RESPONSES_SYNC_PLAN_KIND: &str = "openai_responses_sync";
pub const OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND: &str = "openai_responses_compact_sync";
pub const CLAUDE_CHAT_SYNC_PLAN_KIND: &str = "claude_chat_sync";

16
crates/aether-ai-formats/src/contracts/report_kinds.rs
View File

@@ -1,8 +1,8 @@
use crate::contracts::{
CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_SYNC_PLAN_KIND,
GEMINI_CLI_SYNC_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
};
pub const OPENAI_CHAT_SYNC_FINALIZE_REPORT_KIND: &str = "openai_chat_sync_finalize";
@@ -11,6 +11,7 @@ pub const GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND: &str = "gemini_chat_sync_finali
pub const OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND: &str = "openai_responses_sync_finalize";
pub const OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND: &str =
"openai_responses_compact_sync_finalize";
pub const OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND: &str = "openai_embedding_sync_finalize";
pub const OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND: &str = "openai_image_sync_finalize";
pub const CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND: &str = "claude_cli_sync_finalize";
pub const GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND: &str = "gemini_cli_sync_finalize";
@@ -25,6 +26,8 @@ pub const GEMINI_CHAT_SYNC_SUCCESS_REPORT_KIND: &str = "gemini_chat_sync_success
pub const OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND: &str = "openai_responses_sync_success";
pub const OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND: &str =
"openai_responses_compact_sync_success";
pub const OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND: &str = "openai_embedding_sync_success";
pub const GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND: &str = "gemini_embedding_sync_success";
pub const OPENAI_IMAGE_SYNC_SUCCESS_REPORT_KIND: &str = "openai_image_sync_success";
pub const CLAUDE_CLI_SYNC_SUCCESS_REPORT_KIND: &str = "claude_cli_sync_success";
pub const GEMINI_CLI_SYNC_SUCCESS_REPORT_KIND: &str = "gemini_cli_sync_success";
@@ -45,6 +48,7 @@ pub const GEMINI_CHAT_SYNC_ERROR_REPORT_KIND: &str = "gemini_chat_sync_error";
pub const OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND: &str = "openai_responses_sync_error";
pub const OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND: &str =
"openai_responses_compact_sync_error";
pub const OPENAI_EMBEDDING_SYNC_ERROR_REPORT_KIND: &str = "openai_embedding_sync_error";
pub const OPENAI_IMAGE_SYNC_ERROR_REPORT_KIND: &str = "openai_image_sync_error";
pub const CLAUDE_CLI_SYNC_ERROR_REPORT_KIND: &str = "claude_cli_sync_error";
pub const GEMINI_CLI_SYNC_ERROR_REPORT_KIND: &str = "gemini_cli_sync_error";
@@ -58,6 +62,7 @@ pub fn implicit_sync_finalize_report_kind(plan_kind: &str) -> Option<&'static st
OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND => {
Some(OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND)
}
OPENAI_EMBEDDING_SYNC_PLAN_KIND => Some(OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND),
OPENAI_IMAGE_SYNC_PLAN_KIND => Some(OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND),
CLAUDE_CLI_SYNC_PLAN_KIND => Some(CLAUDE_CLI_SYNC_FINALIZE_REPORT_KIND),
GEMINI_CLI_SYNC_PLAN_KIND => Some(GEMINI_CLI_SYNC_FINALIZE_REPORT_KIND),
@@ -72,6 +77,7 @@ pub fn core_error_default_client_api_format(report_kind: &str) -> Option<&'stati
GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND => Some("gemini:generate_content"),
OPENAI_RESPONSES_SYNC_FINALIZE_REPORT_KIND => Some("openai:responses"),
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND => Some("openai:responses:compact"),
OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND => Some("openai:embedding"),
LEGACY_OPENAI_CLI_SYNC_FINALIZE_REPORT_KIND => Some("openai:responses"),
LEGACY_OPENAI_COMPACT_SYNC_FINALIZE_REPORT_KIND => Some("openai:responses:compact"),
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND => Some("openai:image"),
@@ -90,6 +96,7 @@ pub fn core_error_background_report_kind(report_kind: &str) -> Option<&'static s
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND => {
Some(OPENAI_RESPONSES_COMPACT_SYNC_ERROR_REPORT_KIND)
}
OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND => Some(OPENAI_EMBEDDING_SYNC_ERROR_REPORT_KIND),
LEGACY_OPENAI_CLI_SYNC_FINALIZE_REPORT_KIND => {
Some(OPENAI_RESPONSES_SYNC_ERROR_REPORT_KIND)
}
@@ -115,6 +122,9 @@ pub fn core_success_background_report_kind(report_kind: &str) -> Option<&'static
OPENAI_RESPONSES_COMPACT_SYNC_FINALIZE_REPORT_KIND => {
Some(OPENAI_RESPONSES_COMPACT_SYNC_SUCCESS_REPORT_KIND)
}
OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND => {
Some(OPENAI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND)
}
LEGACY_OPENAI_CLI_SYNC_FINALIZE_REPORT_KIND => {
Some(OPENAI_RESPONSES_SYNC_SUCCESS_REPORT_KIND)
}

1
crates/aether-ai-formats/src/formats/gemini/embedding/mod.rs
View File

@@ -1 +1,2 @@
pub mod request;
pub mod response;

186
crates/aether-ai-formats/src/formats/gemini/embedding/request.rs
View File

@@ -1,9 +1,8 @@
use serde_json::json;
use serde_json::Value;
use serde_json::{json, Map, Value};
use crate::formats::context::FormatContext;
use crate::formats::openai::embedding::request::mapped_embedding_model;
use crate::protocol::canonical::CanonicalRequest;
use crate::protocol::canonical::{CanonicalEmbeddingRequest, CanonicalRequest};
pub fn to(request: &CanonicalRequest, ctx: &FormatContext) -> Option<Value> {
let embedding = request.embedding.as_ref()?;
@@ -13,22 +12,169 @@ pub fn to(request: &CanonicalRequest, ctx: &FormatContext) -> Option<Value> {
}
let model = mapped_embedding_model(request, ctx.mapped_model_or(request.model.as_str()));
if items.len() == 1 {
return Some(json!({
"model": model,
"content": {
"parts": [{"text": items[0]}]
}
}));
return Some(Value::Object(gemini_embedding_request_object(
&model, items[0], embedding,
)));
}
let model_resource = gemini_embedding_model_resource_name(&model);
let requests = items
.into_iter()
.map(|text| {
Value::Object(gemini_embedding_request_object(
&model_resource,
text,
embedding,
))
})
.collect::<Vec<_>>();
Some(json!({ "requests": requests }))
}
fn gemini_embedding_request_object(
model: &str,
text: &str,
embedding: &CanonicalEmbeddingRequest,
) -> Map<String, Value> {
let mut object = Map::new();
object.insert("model".to_string(), Value::String(model.to_string()));
object.insert(
"content".to_string(),
json!({
"parts": [{"text": text}]
}),
);
insert_gemini_embedding_options(&mut object, embedding);
object
}
fn gemini_embedding_model_resource_name(model: &str) -> String {
let trimmed = model.trim();
if trimmed.starts_with("models/") {
trimmed.to_string()
} else {
format!("models/{trimmed}")
}
}
fn insert_gemini_embedding_options(
object: &mut Map<String, Value>,
embedding: &CanonicalEmbeddingRequest,
) {
if let Some(dimensions) = embedding.dimensions {
object.insert("outputDimensionality".to_string(), Value::from(dimensions));
}
if let Some(task_type) = embedding
.task
.as_deref()
.and_then(normalize_gemini_embedding_task_type)
{
object.insert("taskType".to_string(), Value::String(task_type));
}
}
fn normalize_gemini_embedding_task_type(value: &str) -> Option<String> {
let normalized = value.trim();
if normalized.is_empty() {
return None;
}
let key = normalized.replace(['-', ' '], "_").to_ascii_uppercase();
let task_type = match key.as_str() {
"QUERY" | "RETRIEVAL_QUERY" => "RETRIEVAL_QUERY",
"DOCUMENT" | "RETRIEVAL_DOCUMENT" => "RETRIEVAL_DOCUMENT",
"TEXT_MATCHING" | "SEMANTIC_SIMILARITY" => "SEMANTIC_SIMILARITY",
"CLASSIFICATION" => "CLASSIFICATION",
"CLUSTERING" => "CLUSTERING",
"QUESTION_ANSWERING" => "QUESTION_ANSWERING",
"FACT_VERIFICATION" => "FACT_VERIFICATION",
"CODE_RETRIEVAL_QUERY" => "CODE_RETRIEVAL_QUERY",
_ => key.as_str(),
};
Some(task_type.to_string())
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use super::to;
use crate::formats::context::FormatContext;
use crate::protocol::canonical::{
CanonicalEmbeddingInput, CanonicalEmbeddingRequest, CanonicalRequest,
};
fn canonical_embedding(input: CanonicalEmbeddingInput) -> CanonicalRequest {
CanonicalRequest {
model: "text-embedding-3-small".to_string(),
embedding: Some(CanonicalEmbeddingRequest {
input,
encoding_format: None,
dimensions: None,
task: None,
user: None,
extensions: BTreeMap::new(),
}),
..CanonicalRequest::default()
}
}
#[test]
fn single_string_array_item_uses_single_embed_content_body() {
let request = canonical_embedding(CanonicalEmbeddingInput::StringArray(vec![
"hello".to_string()
]));
let body = to(
&request,
&FormatContext::default().with_mapped_model("gemini-embedding-2-preview"),
)
.expect("gemini embedding request");
assert_eq!(body["model"], "gemini-embedding-2-preview");
assert_eq!(body["content"]["parts"][0]["text"], "hello");
assert!(body.get("requests").is_none());
}
#[test]
fn multiple_string_items_use_gemini_batch_request_body() {
let request = canonical_embedding(CanonicalEmbeddingInput::StringArray(vec![
"alpha".to_string(),
"beta".to_string(),
]));
let body = to(
&request,
&FormatContext::default().with_mapped_model("gemini-embedding-2-preview"),
)
.expect("gemini embedding request");
assert!(body.get("model").is_none());
assert_eq!(body["requests"].as_array().map(Vec::len), Some(2));
assert_eq!(
body["requests"][0]["model"],
"models/gemini-embedding-2-preview"
);
assert_eq!(body["requests"][0]["content"]["parts"][0]["text"], "alpha");
assert_eq!(
body["requests"][1]["model"],
"models/gemini-embedding-2-preview"
);
assert_eq!(body["requests"][1]["content"]["parts"][0]["text"], "beta");
}
#[test]
fn explicit_embedding_options_are_preserved_without_defaults() {
let mut request = canonical_embedding(CanonicalEmbeddingInput::String("query".to_string()));
let embedding = request.embedding.as_mut().expect("embedding request");
embedding.dimensions = Some(768);
embedding.task = Some("retrieval_query".to_string());
let body = to(
&request,
&FormatContext::default().with_mapped_model("gemini-embedding-2-preview"),
)
.expect("gemini embedding request");
assert_eq!(body["outputDimensionality"], json!(768));
assert_eq!(body["taskType"], "RETRIEVAL_QUERY");
}
Some(json!({
"model": model,
"requests": items.into_iter().map(|text| {
json!({
"model": model,
"content": {
"parts": [{"text": text}]
}
})
}).collect::<Vec<_>>()
}))
}

158
crates/aether-ai-formats/src/formats/gemini/embedding/response.rs Normal file
View File

@@ -0,0 +1,158 @@
use serde_json::Value;
use crate::formats::openai::embedding::request::namespace_extensions;
use crate::protocol::canonical::{
gemini_usage_to_canonical, CanonicalEmbedding, CanonicalEmbeddingResponse,
};
pub fn from(body_json: &Value) -> Option<CanonicalEmbeddingResponse> {
let body = body_json.as_object()?;
if body.contains_key("error") {
return None;
}
let embeddings = if let Some(raw_embeddings) = vertex_predict_embeddings(body) {
raw_embeddings
} else if let Some(values) = body
.get("embedding")
.and_then(Value::as_object)
.and_then(|embedding| embedding.get("values"))
.and_then(Value::as_array)
{
vec![CanonicalEmbedding {
index: 0,
embedding: embedding_values(values)?,
extensions: Default::default(),
}]
} else {
let raw_embeddings = body.get("embeddings")?.as_array()?;
raw_embeddings
.iter()
.enumerate()
.map(|(index, item)| {
let item_object = item.as_object()?;
let values = item_object.get("values")?.as_array()?;
Some(CanonicalEmbedding {
index,
embedding: embedding_values(values)?,
extensions: namespace_extensions("gemini", item_object, &["values"]),
})
})
.collect::<Option<Vec<_>>>()?
};
if embeddings.is_empty()
|| embeddings
.iter()
.any(|embedding| embedding.embedding.is_empty())
{
return None;
}
Some(CanonicalEmbeddingResponse {
id: body
.get("id")
.or_else(|| body.get("responseId"))
.and_then(Value::as_str)
.unwrap_or("embd-gemini-unknown")
.to_string(),
model: body
.get("model")
.or_else(|| body.get("modelVersion"))
.or_else(|| body.get("deployedModelId"))
.and_then(Value::as_str)
.unwrap_or("unknown")
.to_string(),
embeddings,
usage: gemini_usage_to_canonical(body.get("usageMetadata")),
extensions: namespace_extensions(
"gemini",
body,
&[
"id",
"responseId",
"model",
"modelVersion",
"deployedModelId",
"embedding",
"embeddings",
"predictions",
"usageMetadata",
],
),
})
}
fn vertex_predict_embeddings(
body: &serde_json::Map<String, Value>,
) -> Option<Vec<CanonicalEmbedding>> {
let predictions = body.get("predictions")?.as_array()?;
predictions
.iter()
.enumerate()
.map(|(index, item)| {
let item_object = item.as_object()?;
let embedding_object = item_object.get("embeddings")?.as_object()?;
let values = embedding_object.get("values")?.as_array()?;
Some(CanonicalEmbedding {
index,
embedding: embedding_values(values)?,
extensions: namespace_extensions("vertex", item_object, &["embeddings"]),
})
})
.collect()
}
fn embedding_values(values: &[Value]) -> Option<Vec<f64>> {
values.iter().map(Value::as_f64).collect()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::from;
#[test]
fn parses_gemini_single_embedding_response() {
let body = json!({
"embedding": {"values": [0.1, 0.2, 0.3]},
"usageMetadata": {
"promptTokenCount": 4,
"totalTokenCount": 4
}
});
let parsed = from(&body).expect("response should parse");
assert_eq!(parsed.model, "unknown");
assert_eq!(parsed.embeddings[0].embedding, vec![0.1, 0.2, 0.3]);
let usage = parsed.usage.expect("usage should parse");
assert_eq!(usage.input_tokens, 4);
assert_eq!(usage.total_tokens, 4);
}
#[test]
fn parses_vertex_predict_embedding_response() {
let body = json!({
"predictions": [
{
"embeddings": {
"values": [0.1, 0.2, 0.3]
}
},
{
"embeddings": {
"values": [0.4, 0.5, 0.6]
}
}
],
"deployedModelId": "gemini-embedding-2"
});
let parsed = from(&body).expect("response should parse");
assert_eq!(parsed.model, "gemini-embedding-2");
assert_eq!(parsed.embeddings[0].embedding, vec![0.1, 0.2, 0.3]);
assert_eq!(parsed.embeddings[1].embedding, vec![0.4, 0.5, 0.6]);
}
}

85
crates/aether-ai-formats/src/formats/gemini/generate_content/response.rs
View File

@@ -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"
));
}
}

1
crates/aether-ai-formats/src/formats/openai/embedding/mod.rs
View File

@@ -1,2 +1,3 @@
pub mod request;
pub mod response;
pub mod spec;

35
crates/aether-ai-formats/src/formats/openai/embedding/spec.rs Normal file
View File

@@ -0,0 +1,35 @@
use crate::contracts::{
OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
};
use crate::formats::shared::family::{
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
};
pub fn resolve_sync_spec(plan_kind: &str) -> Option<LocalStandardSpec> {
match plan_kind {
OPENAI_EMBEDDING_SYNC_PLAN_KIND => Some(LocalStandardSpec {
api_format: "openai:embedding",
decision_kind: OPENAI_EMBEDDING_SYNC_PLAN_KIND,
report_kind: OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND,
family: LocalStandardSourceFamily::Standard,
mode: LocalStandardSourceMode::Embedding,
require_streaming: false,
}),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::resolve_sync_spec;
use crate::formats::shared::family::LocalStandardSourceMode;
#[test]
fn resolves_openai_embedding_sync_standard_spec() {
let spec = resolve_sync_spec("openai_embedding_sync").expect("spec");
assert_eq!(spec.api_format, "openai:embedding");
assert_eq!(spec.report_kind, "openai_embedding_sync_finalize");
assert_eq!(spec.mode, LocalStandardSourceMode::Embedding);
assert!(!spec.require_streaming);
}
}

7
crates/aether-ai-formats/src/formats/registry.rs
View File

@@ -228,11 +228,16 @@ mod tests {
&FormatContext::default().with_mapped_model("gemini-embedding-001"),
)
.expect("gemini embedding conversion should succeed");
assert_eq!(gemini["model"], "gemini-embedding-001");
assert!(gemini.get("model").is_none());
assert_eq!(
gemini["requests"][0]["model"],
"models/gemini-embedding-001"
);
assert_eq!(
gemini["requests"][0]["content"]["parts"][0]["text"],
"alpha"
);
assert_eq!(gemini["requests"][0]["outputDimensionality"], 2);
assert!(gemini.get("messages").is_none());
let doubao = convert_request(

2
crates/aether-ai-formats/src/formats/shared/error_body.rs
View File

@@ -38,7 +38,7 @@ pub fn build_core_error_body_for_client_format(
error_object.insert("message".to_string(), Value::String(message.to_string()));
match aether_ai_formats::normalize_api_format_alias(client_api_format).as_str() {
"openai:chat" | "openai:responses" | "openai:responses:compact" => {
"openai:chat" | "openai:responses" | "openai:responses:compact" | "openai:embedding" => {
error_object.insert(
"type".to_string(),
Value::String(map_local_sync_error_kind_to_openai_type(kind).to_string()),

19
crates/aether-ai-formats/src/formats/shared/passthrough.rs
View File

@@ -1,7 +1,8 @@
use crate::contracts::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND,
};
@@ -50,6 +51,13 @@ pub fn resolve_sync_spec(plan_kind: &str) -> Option<LocalSameFormatProviderSpec>
family: LocalSameFormatProviderFamily::Gemini,
require_streaming: false,
}),
GEMINI_EMBEDDING_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "gemini:embedding",
decision_kind: GEMINI_EMBEDDING_SYNC_PLAN_KIND,
report_kind: GEMINI_EMBEDDING_SYNC_SUCCESS_REPORT_KIND,
family: LocalSameFormatProviderFamily::Gemini,
require_streaming: false,
}),
OPENAI_EMBEDDING_SYNC_PLAN_KIND => Some(LocalSameFormatProviderSpec {
api_format: "openai:embedding",
decision_kind: OPENAI_EMBEDDING_SYNC_PLAN_KIND,
@@ -130,6 +138,15 @@ mod tests {
assert!(!spec.require_streaming);
}
#[test]
fn resolves_gemini_embedding_sync_same_format_spec() {
let spec = resolve_sync_spec("gemini_embedding_sync").expect("spec");
assert_eq!(spec.api_format, "gemini:embedding");
assert_eq!(spec.report_kind, "gemini_embedding_sync_success");
assert_eq!(spec.family, super::LocalSameFormatProviderFamily::Gemini);
assert!(!spec.require_streaming);
}
#[test]
fn resolves_openai_rerank_sync_same_format_spec() {
let spec = resolve_sync_spec("openai_rerank_sync").expect("spec");

50
crates/aether-ai-formats/src/formats/shared/routing.rs
View File

@@ -4,9 +4,9 @@ use url::form_urlencoded;
use crate::contracts::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_FILES_DELETE_PLAN_KIND,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND, GEMINI_FILES_LIST_PLAN_KIND,
GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
GEMINI_FILES_DELETE_PLAN_KIND, GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_FILES_GET_PLAN_KIND,
GEMINI_FILES_LIST_PLAN_KIND, GEMINI_FILES_UPLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND,
OPENAI_RERANK_SYNC_PLAN_KIND, OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND,
@@ -166,6 +166,14 @@ pub fn resolve_execution_runtime_sync_plan_kind(
return Some(GEMINI_VIDEO_CREATE_SYNC_PLAN_KIND);
}
if route_family == Some("gemini")
&& route_kind == Some("embedding")
&& *method == Method::POST
&& (path.ends_with(":embedContent") || path.ends_with(":batchEmbedContents"))
{
return Some(GEMINI_EMBEDDING_SYNC_PLAN_KIND);
}
if route_family == Some("openai")
&& route_kind == Some("chat")
&& *method == Method::POST
@@ -411,6 +419,7 @@ pub fn supports_sync_execution_decision_kind(plan_kind: &str) -> bool {
| CLAUDE_CLI_SYNC_PLAN_KIND
| GEMINI_CHAT_SYNC_PLAN_KIND
| GEMINI_CLI_SYNC_PLAN_KIND
| GEMINI_EMBEDDING_SYNC_PLAN_KIND
| GEMINI_FILES_UPLOAD_PLAN_KIND
| OPENAI_VIDEO_CREATE_SYNC_PLAN_KIND
| OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND
@@ -455,9 +464,9 @@ mod tests {
use crate::contracts::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
CLAUDE_CLI_SYNC_PLAN_KIND, GEMINI_CHAT_STREAM_PLAN_KIND, GEMINI_CHAT_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, OPENAI_CHAT_STREAM_PLAN_KIND,
OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND, OPENAI_IMAGE_STREAM_PLAN_KIND,
OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
GEMINI_CLI_STREAM_PLAN_KIND, GEMINI_CLI_SYNC_PLAN_KIND, GEMINI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_PLAN_KIND, OPENAI_RERANK_SYNC_PLAN_KIND,
OPENAI_RESPONSES_COMPACT_STREAM_PLAN_KIND, OPENAI_RESPONSES_COMPACT_SYNC_PLAN_KIND,
OPENAI_RESPONSES_STREAM_PLAN_KIND, OPENAI_RESPONSES_SYNC_PLAN_KIND,
};
@@ -783,6 +792,35 @@ mod tests {
));
}
#[test]
fn resolves_gemini_embedding_sync_plan_kind() {
assert_eq!(
resolve_execution_runtime_sync_plan_kind(
Some("ai_public"),
Some("gemini"),
Some("embedding"),
Some("api_key"),
&Method::POST,
"/v1beta/models/gemini-embedding-2-preview:embedContent",
),
Some(GEMINI_EMBEDDING_SYNC_PLAN_KIND)
);
assert_eq!(
resolve_execution_runtime_sync_plan_kind(
Some("ai_public"),
Some("gemini"),
Some("embedding"),
Some("api_key"),
&Method::POST,
"/v1beta/models/gemini-embedding-2-preview:batchEmbedContents",
),
Some(GEMINI_EMBEDDING_SYNC_PLAN_KIND)
);
assert!(supports_sync_execution_decision_kind(
GEMINI_EMBEDDING_SYNC_PLAN_KIND
));
}
#[test]
fn resolves_openai_rerank_sync_plan_kind() {
assert_eq!(

107
crates/aether-ai-formats/src/formats/shared/sync_products.rs
View File

@@ -9,9 +9,10 @@ use aether_ai_formats::formats::conversion::response::{
};
use aether_ai_formats::formats::registry::{convert_response, FormatContext};
use aether_ai_formats::{
canonical_to_claude_response, canonical_to_gemini_response, canonical_to_openai_chat_response,
canonical_to_openai_responses_compact_response, canonical_to_openai_responses_response,
from_claude_to_canonical_response, from_gemini_to_canonical_response,
canonical_to_claude_response, canonical_to_embedding_response, canonical_to_gemini_response,
canonical_to_openai_chat_response, canonical_to_openai_responses_compact_response,
canonical_to_openai_responses_response, from_claude_to_canonical_response,
from_embedding_to_canonical_response, from_gemini_to_canonical_response,
from_openai_chat_to_canonical_response, from_openai_responses_to_canonical_response,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
};
@@ -330,6 +331,18 @@ pub fn maybe_build_standard_sync_finalize_product_from_normalized_payload(
)));
}
if let Some(product) = maybe_build_embedding_cross_format_sync_product_from_normalized_payload(
report_kind,
status_code,
report_context,
body_json,
body_base64,
)? {
return Ok(Some(StandardSyncFinalizeNormalizedProduct::CrossFormat(
product,
)));
}
Ok(
maybe_build_standard_cross_format_sync_product_from_normalized_payload(
report_kind,
@@ -342,6 +355,90 @@ pub fn maybe_build_standard_sync_finalize_product_from_normalized_payload(
)
}
pub fn maybe_build_embedding_cross_format_sync_product_from_normalized_payload(
report_kind: &str,
status_code: u16,
report_context: Option<&Value>,
body_json: Option<&Value>,
body_base64: Option<&str>,
) -> Result<Option<StandardCrossFormatSyncProduct>, AiSurfaceFinalizeError> {
if report_kind != crate::contracts::OPENAI_EMBEDDING_SYNC_FINALIZE_REPORT_KIND
|| status_code >= 400
{
return Ok(None);
}
let Some(report_context) = report_context else {
return Ok(None);
};
let provider_api_format = report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
let client_api_format = report_context
.get("client_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
if provider_api_format == client_api_format {
return Ok(None);
}
let Some(provider_namespace) =
embedding_response_namespace_for_api_format(&provider_api_format)
else {
return Ok(None);
};
let Some(client_namespace) = embedding_response_namespace_for_api_format(&client_api_format)
else {
return Ok(None);
};
let provider_body_json = match body_base64 {
Some(body_base64) => {
let body_bytes = base64::engine::general_purpose::STANDARD.decode(body_base64)?;
serde_json::from_slice::<Value>(&body_bytes).ok()
}
None => body_json.cloned(),
};
let Some(provider_body_json) =
provider_body_json.filter(|value| !is_error_like_sync_body(value))
else {
return Ok(None);
};
let mut canonical =
match from_embedding_to_canonical_response(&provider_body_json, provider_namespace) {
Some(canonical) => canonical,
None => return Ok(None),
};
apply_report_context_model_fallback(&mut canonical.model, report_context);
let Some(client_body_json) = canonical_to_embedding_response(&canonical, client_namespace)
else {
return Ok(None);
};
let client_body_json =
client_body_with_report_context_model(client_body_json, report_context, &client_api_format);
Ok(Some(StandardCrossFormatSyncProduct {
client_body_json,
provider_body_json,
}))
}
fn embedding_response_namespace_for_api_format(api_format: &str) -> Option<&'static str> {
match aether_ai_formats::normalize_api_format_alias(api_format).as_str() {
"openai:embedding" => Some("openai"),
"jina:embedding" => Some("jina"),
"gemini:embedding" => Some("gemini"),
_ => None,
}
}
fn maybe_build_standard_same_format_sync_body(
report_kind: &str,
status_code: u16,
@@ -829,6 +926,9 @@ fn client_body_with_report_context_model(
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages" => {
object.insert("model".to_string(), Value::String(display_model));
}
"openai:embedding" | "jina:embedding" => {
object.insert("model".to_string(), Value::String(display_model));
}
"gemini:generate_content" => {
object.insert("modelVersion".to_string(), Value::String(display_model));
}
@@ -1310,6 +1410,7 @@ fn standard_same_format_api_format(report_kind: &str) -> Option<&'static str> {
"openai_chat_sync_finalize" => Some("openai:chat"),
"claude_chat_sync_finalize" => Some("claude:messages"),
"gemini_chat_sync_finalize" => Some("gemini:generate_content"),
"openai_embedding_sync_finalize" => Some("openai:embedding"),
"claude_cli_sync_finalize" => Some("claude:messages"),
"gemini_cli_sync_finalize" => Some("gemini:generate_content"),
_ => None,

7
crates/aether-ai-formats/src/protocol/canonical.rs
View File

@@ -684,6 +684,7 @@ pub fn from_embedding_to_canonical_response(
crate::formats::openai::embedding::response::from_namespace(body_json, "openai")
}
"jina" => crate::formats::openai::embedding::response::from_namespace(body_json, "jina"),
"gemini" => crate::formats::gemini::embedding::response::from(body_json),
_ => None,
}
}
@@ -4830,10 +4831,16 @@ mod tests {
let gemini =
super::canonical_to_embedding_request(&canonical, "gemini-embedding-001", "gemini")
.expect("gemini embedding request");
assert!(gemini.get("model").is_none());
assert_eq!(
gemini["requests"][0]["model"],
"models/gemini-embedding-001"
);
assert_eq!(
gemini["requests"][0]["content"]["parts"][0]["text"],
"alpha"
);
assert_eq!(gemini["requests"][0]["outputDimensionality"], 2);
assert!(gemini.get("messages").is_none());
let doubao = super::canonical_to_embedding_request(

15
crates/aether-ai-serving/src/failure_diagnostic.rs
View File

@@ -162,6 +162,21 @@ impl CandidateFailureDiagnostic {
.source(source)
}
pub fn request_conversion_failed(
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,
source: impl Into<String>,
message: impl Into<String>,
) -> Self {
Self::new(
CandidateFailureDiagnosticKind::RequestConversion,
"$",
message,
)
.formats(client_api_format, provider_api_format)
.source(source)
}
pub fn envelope_build_failed(
client_api_format: impl Into<String>,
provider_api_format: impl Into<String>,

8
crates/aether-data/src/repository/candidate_selection/memory.rs
View File

@@ -324,11 +324,15 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
)
}
"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",

8
crates/aether-data/src/repository/candidate_selection/mysql.rs
View File

@@ -453,11 +453,15 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
)
}
"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",

33
crates/aether-data/src/repository/candidate_selection/postgres.rs
View File

@@ -118,11 +118,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')
)
)
)
@@ -307,11 +307,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')
)
)
)
@@ -495,11 +495,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')
)
)
)
@@ -1316,12 +1316,29 @@ mod tests {
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'grok'"));
assert!(sql.contains("LOWER(BTRIM(pak.auth_type)) = 'oauth'"));
assert!(sql
.contains("'openai:chat', 'openai:responses', 'claude:messages', 'openai:image'"));
assert!(
sql.contains("'openai:chat', 'openai:responses', 'claude:messages', 'openai:image'")
);
assert!(sql.contains("'grok',"));
}
}
#[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();

8
crates/aether-data/src/repository/candidate_selection/sqlite.rs
View File

@@ -832,11 +832,15 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
)
}
"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",

87
crates/aether-model-fetch/src/strategy.rs
View File

@@ -28,7 +28,6 @@ const ANTIGRAVITY_DAILY_BASE_URL: &str = "https://daily-cloudcode-pa.googleapis.
const ANTIGRAVITY_PROD_BASE_URL: &str = "https://cloudcode-pa.googleapis.com";
const ANTIGRAVITY_BLOCKED_MODELS: &[&str] = &["chat_23310", "chat_20706"];
const VERTEX_API_BASE_URL: &str = "https://aiplatform.googleapis.com";
const VERTEX_MODEL_GARDEN_LIST_API_VERSION: &str = "v1beta1";
const VERTEX_PAGE_SIZE: &str = "100";
const VERTEX_MAX_PAGES: usize = 20;
const GOOGLE_OAUTH_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
@@ -465,8 +464,6 @@ async fn fetch_vertex_service_account_models(
});
};
let token = exchange_vertex_service_account_token(runtime, &transports[0], auth_config).await?;
let project_id = json_string(auth_config.get("project_id"))
.ok_or_else(|| "vertex_ai(service_account): missing project_id".to_string())?;
let gemini_transport =
select_transport_for_api_format(transports, "gemini:").unwrap_or(&transports[0]);
let claude_transport =
@@ -477,18 +474,12 @@ async fn fetch_vertex_service_account_models(
let mut soft_errors = Vec::new();
let mut has_success = false;
for region in vertex_regions(auth_config) {
let base = if region == "global" {
VERTEX_API_BASE_URL.to_string()
} else {
format!("https://{region}-aiplatform.googleapis.com")
};
for base in iter_vertex_base_urls(transports) {
for (publisher, transport, api_format) in [
("google", gemini_transport, "gemini:generate_content"),
("anthropic", claude_transport, "claude:messages"),
] {
let url =
build_vertex_service_account_list_url(&base, &project_id, &region, publisher, None);
let url = build_vertex_service_account_list_url(&base, publisher, None);
let outcome = fetch_vertex_models_from_url(
runtime,
transport,
@@ -930,7 +921,7 @@ fn iter_vertex_base_urls(transports: &[GatewayProviderTransportSnapshot]) -> Vec
}
fn build_vertex_google_list_url(base_url: &str, api_key: &str, page_token: Option<&str>) -> String {
let url = build_vertex_model_garden_list_url(base_url, "google");
let url = build_vertex_publisher_models_list_base_url(base_url, "google");
let mut url = append_query_param(url, "key", api_key);
url = append_query_param(url, "pageSize", VERTEX_PAGE_SIZE);
if let Some(page_token) = page_token {
@@ -941,12 +932,10 @@ fn build_vertex_google_list_url(base_url: &str, api_key: &str, page_token: Optio
fn build_vertex_service_account_list_url(
base_url: &str,
_project_id: &str,
_region: &str,
publisher: &str,
page_token: Option<&str>,
) -> String {
let mut url = build_vertex_model_garden_list_url(base_url, publisher);
let mut url = build_vertex_publisher_models_list_base_url(base_url, publisher);
url = append_query_param(url, "pageSize", VERTEX_PAGE_SIZE);
if let Some(page_token) = page_token {
url = append_query_param(url, "pageToken", page_token);
@@ -954,14 +943,14 @@ fn build_vertex_service_account_list_url(
url
}
fn build_vertex_model_garden_list_url(base_url: &str, publisher: &str) -> String {
fn build_vertex_publisher_models_list_base_url(base_url: &str, publisher: &str) -> String {
let trimmed_base = base_url.trim().trim_end_matches('/');
let unversioned_base = trimmed_base
.strip_suffix("/v1beta1")
.or_else(|| trimmed_base.strip_suffix("/v1"))
.unwrap_or(trimmed_base);
let path = format!("/{VERTEX_MODEL_GARDEN_LIST_API_VERSION}/publishers/{publisher}/models");
build_simple_path_url(unversioned_base, &path)
let path = if trimmed_base.ends_with("/v1") || trimmed_base.ends_with("/v1beta1") {
format!("/publishers/{publisher}/models")
} else {
format!("/v1beta1/publishers/{publisher}/models")
};
build_simple_path_url(trimmed_base, &path)
}
fn build_simple_path_url(base_url: &str, path: &str) -> String {
@@ -1089,38 +1078,6 @@ fn vertex_effective_format(model_id: &str, auth_config: Option<&Value>) -> Strin
}
}
fn vertex_regions(auth_config: &Value) -> Vec<String> {
let mut seen = BTreeSet::new();
let mut regions = Vec::new();
let auth_config = auth_config.as_object();
let mut push_region = |region: Option<&str>| {
let Some(region) = region.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
if seen.insert(region.to_string()) {
regions.push(region.to_string());
}
};
push_region(
auth_config
.and_then(|value| value.get("region"))
.and_then(Value::as_str),
);
if let Some(model_regions) = auth_config
.and_then(|value| value.get("model_regions"))
.and_then(Value::as_object)
{
for value in model_regions.values() {
push_region(value.as_str());
}
}
push_region(Some("global"));
push_region(Some("us-central1"));
regions
}
fn is_soft_not_found(error: &str) -> bool {
error.trim().starts_with("HTTP 404:")
}
@@ -1517,7 +1474,7 @@ mod tests {
fn vertex_model_fetch_uses_model_garden_list_endpoint() {
assert_eq!(
build_vertex_google_list_url(
"https://aiplatform.googleapis.com/v1",
"https://aiplatform.googleapis.com",
"vertex-secret",
None,
),
@@ -1526,8 +1483,6 @@ mod tests {
assert_eq!(
build_vertex_service_account_list_url(
"https://aiplatform.googleapis.com",
"project-1",
"global",
"google",
Some("page-2"),
),
@@ -1535,6 +1490,24 @@ mod tests {
);
}
#[test]
fn vertex_publisher_models_list_url_uses_model_garden_resource_not_runtime_resource() {
let url = super::build_vertex_service_account_list_url(
"https://aiplatform.googleapis.com",
"google",
None,
);
assert_eq!(
url,
"https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=100"
);
assert!(
!url.contains("/projects/") && !url.contains("/locations/"),
"Model Garden publisher list must not use Vertex runtime project/location path"
);
}
#[tokio::test]
async fn codex_transport_fetches_upstream_models_instead_of_preset_catalog() {
let executed_urls = Arc::new(Mutex::new(Vec::new()));

118
crates/aether-provider-transport/src/conversion.rs
View File

@@ -1,5 +1,6 @@
use aether_ai_formats::formats::matrix::{
request_conversion_kind, request_conversion_requires_enable_flag, RequestConversionKind,
api_data_format_id, request_conversion_kind, request_conversion_requires_enable_flag,
RequestConversionKind,
};
use aether_ai_formats::normalize_api_format_alias;
@@ -39,7 +40,14 @@ pub fn request_conversion_enabled_for_transport(
if client_api_format == provider_api_format {
return true;
}
if request_conversion_kind(client_api_format.as_str(), provider_api_format.as_str()).is_none() {
let conversion_kind =
request_conversion_kind(client_api_format.as_str(), provider_api_format.as_str());
if conversion_kind.is_none()
&& !same_data_format_transport_pair(
client_api_format.as_str(),
provider_api_format.as_str(),
)
{
return false;
}
if !request_conversion_requires_enable_flag(
@@ -62,7 +70,14 @@ pub fn request_pair_allowed_for_transport(
if client_api_format == provider_api_format {
return true;
}
if request_conversion_kind(client_api_format.as_str(), provider_api_format.as_str()).is_none() {
let conversion_kind =
request_conversion_kind(client_api_format.as_str(), provider_api_format.as_str());
if conversion_kind.is_none()
&& !same_data_format_transport_pair(
client_api_format.as_str(),
provider_api_format.as_str(),
)
{
return false;
}
if is_kiro_claude_messages_transport(transport, &provider_api_format) {
@@ -80,6 +95,19 @@ pub fn request_pair_allowed_for_transport(
)
}
fn same_data_format_transport_pair(client_api_format: &str, provider_api_format: &str) -> bool {
if aether_ai_formats::api_format_alias_matches(client_api_format, provider_api_format) {
return false;
}
matches!(
(
api_data_format_id(client_api_format),
api_data_format_id(provider_api_format)
),
(Some("embedding"), Some("embedding")) | (Some("rerank"), Some("rerank"))
)
}
pub fn request_conversion_transport_supported(
transport: &GatewayProviderTransportSnapshot,
kind: RequestConversionKind,
@@ -117,15 +145,72 @@ pub fn request_conversion_transport_unsupported_reason(
}
}
pub fn request_pair_transport_unsupported_reason(
transport: &GatewayProviderTransportSnapshot,
client_api_format: &str,
provider_api_format: &str,
) -> Option<&'static str> {
let client_api_format = normalize_api_format_alias(client_api_format);
let provider_api_format = normalize_api_format_alias(provider_api_format);
if let Some(kind) =
request_conversion_kind(client_api_format.as_str(), provider_api_format.as_str())
{
return request_conversion_transport_unsupported_reason(transport, kind);
}
if !same_data_format_transport_pair(client_api_format.as_str(), provider_api_format.as_str()) {
return Some("transport_api_format_unsupported");
}
match provider_api_format.as_str() {
"gemini:embedding" => {
if is_vertex_transport_context(transport) {
local_vertex_gemini_transport_unsupported_reason_with_network(transport)
} else {
local_gemini_transport_unsupported_reason_with_network(
transport,
"gemini:embedding",
)
}
}
"openai:embedding" | "jina:embedding" | "doubao:embedding" | "openai:rerank"
| "jina:rerank" => local_standard_transport_unsupported_reason_with_network(
transport,
provider_api_format.as_str(),
),
_ => Some("transport_api_format_unsupported"),
}
}
pub fn request_conversion_direct_auth(
transport: &GatewayProviderTransportSnapshot,
_kind: RequestConversionKind,
) -> Option<(String, String)> {
match normalize_api_format_alias(&transport.endpoint.api_format).as_str() {
"openai:chat" | "openai:responses" | "openai:responses:compact" => {
resolve_local_openai_bearer_auth(transport)
}
"gemini:generate_content" => {
request_direct_auth_for_provider_format(transport, transport.endpoint.api_format.as_str())
}
pub fn request_pair_direct_auth(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
) -> Option<(String, String)> {
request_direct_auth_for_provider_format(transport, provider_api_format)
}
fn request_direct_auth_for_provider_format(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
) -> Option<(String, String)> {
match normalize_api_format_alias(provider_api_format).as_str() {
"openai:chat"
| "openai:responses"
| "openai:responses:compact"
| "openai:embedding"
| "jina:embedding"
| "doubao:embedding"
| "openai:rerank"
| "jina:rerank" => resolve_local_openai_bearer_auth(transport),
"gemini:generate_content" | "gemini:embedding" => {
if is_vertex_api_key_transport_context(transport) {
resolve_local_vertex_api_key_query_auth(transport)
.map(|auth| (VERTEX_API_KEY_QUERY_PARAM.to_string(), auth.value))
@@ -338,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;
@@ -508,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);

8
crates/aether-provider-transport/src/lib.rs
View File

@@ -15,6 +15,7 @@ pub mod oauth_refresh;
mod openai_image;
pub mod policy;
pub mod provider_types;
mod request_body;
mod request_url;
pub mod rules;
pub mod same_format_provider;
@@ -31,7 +32,8 @@ pub use conversion::{
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
request_conversion_direct_auth, request_conversion_enabled_for_transport,
request_conversion_transport_supported, request_conversion_transport_unsupported_reason,
request_pair_allowed_for_transport, CandidateTransportPolicyFacts,
request_pair_allowed_for_transport, request_pair_direct_auth,
request_pair_transport_unsupported_reason, CandidateTransportPolicyFacts,
};
pub use diagnostics::{
append_transport_diagnostics_to_value, build_request_trace_proxy_value,
@@ -80,10 +82,14 @@ pub use policy::{
local_standard_transport_unsupported_reason_with_network, supports_local_gemini_transport,
supports_local_gemini_transport_with_network, supports_local_standard_transport,
};
pub use request_body::{
apply_transport_request_body_semantics, TransportRequestBodySemanticsError,
};
pub use request_url::{
build_cross_format_openai_chat_upstream_url, build_cross_format_openai_responses_upstream_url,
build_kiro_cross_format_upstream_url, build_local_openai_chat_upstream_url,
build_local_openai_responses_upstream_url, build_transport_request_url,
build_transport_request_url_for_request_body, gemini_embedding_request_body_uses_batch,
TransportRequestUrlParams,
};
pub use rules::{

44
crates/aether-provider-transport/src/provider_types.rs
View File

@@ -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 {
@@ -338,6 +343,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",
@@ -679,6 +690,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
));
@@ -747,6 +763,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"),
@@ -760,6 +777,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"),
@@ -776,4 +794,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()
);
}
}

415
crates/aether-provider-transport/src/request_body.rs Normal file
View File

@@ -0,0 +1,415 @@
use serde_json::{Map, Value};
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::vertex::is_vertex_transport_context;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TransportRequestBodySemanticsError {
message: &'static str,
}
impl TransportRequestBodySemanticsError {
const fn new(message: &'static str) -> Self {
Self { message }
}
pub const fn message(&self) -> &'static str {
self.message
}
}
impl std::fmt::Display for TransportRequestBodySemanticsError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.message)
}
}
impl std::error::Error for TransportRequestBodySemanticsError {}
pub fn apply_transport_request_body_semantics(
provider_request_body: &mut Value,
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
) -> Result<(), TransportRequestBodySemanticsError> {
let provider_api_format = aether_ai_formats::normalize_api_format_alias(provider_api_format);
if provider_api_format == "gemini:embedding" && is_vertex_transport_context(transport) {
apply_vertex_gemini_embedding_body_semantics(provider_request_body)?;
}
Ok(())
}
fn apply_vertex_gemini_embedding_body_semantics(
provider_request_body: &mut Value,
) -> Result<(), TransportRequestBodySemanticsError> {
let object = provider_request_body.as_object_mut().ok_or_else(|| {
TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding request body must be a JSON object",
)
})?;
if object.contains_key("instances") {
validate_existing_vertex_predict_body(object)?;
object.remove("model");
return Ok(());
}
let next = build_vertex_predict_body_from_gemini_embedding_object(object)?;
*object = next;
Ok(())
}
fn build_vertex_predict_body_from_gemini_embedding_object(
object: &Map<String, Value>,
) -> Result<Map<String, Value>, TransportRequestBodySemanticsError> {
if let Some(requests) = object.get("requests") {
if object.keys().any(|key| key != "requests") {
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding batch body cannot mix requests with other top-level fields",
));
}
let request_items = requests.as_array().ok_or_else(|| {
TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding requests must be an array",
)
})?;
let request_objects = request_items
.iter()
.map(Value::as_object)
.collect::<Option<Vec<_>>>()
.ok_or_else(|| {
TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding requests must be an array of objects",
)
})?;
if request_objects.is_empty() {
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding requests must contain at least one item",
));
}
return build_vertex_predict_body_from_gemini_embedding_items(&request_objects);
}
build_vertex_predict_body_from_gemini_embedding_items(&[object])
}
fn build_vertex_predict_body_from_gemini_embedding_items(
items: &[&Map<String, Value>],
) -> Result<Map<String, Value>, TransportRequestBodySemanticsError> {
if items.iter().any(|item| {
item.keys().any(|key| {
!matches!(
key.as_str(),
"model"
| "content"
| "taskType"
| "title"
| "outputDimensionality"
| "autoTruncate"
)
})
}) {
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding body contains fields that cannot be mapped to predict instances",
));
}
let instances = items
.iter()
.map(|item| build_vertex_predict_instance(item))
.collect::<Option<Vec<_>>>()
.ok_or_else(|| {
TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding body must contain text content parts",
)
})?;
if instances.is_empty() {
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding body must contain at least one instance",
));
}
let mut output = Map::new();
output.insert("instances".to_string(), Value::Array(instances));
let mut parameters = Map::new();
insert_shared_parameter(items, &mut parameters, "outputDimensionality")?;
insert_shared_parameter(items, &mut parameters, "autoTruncate")?;
if !parameters.is_empty() {
output.insert("parameters".to_string(), Value::Object(parameters));
}
Ok(output)
}
fn validate_existing_vertex_predict_body(
object: &Map<String, Value>,
) -> Result<(), TransportRequestBodySemanticsError> {
if object
.keys()
.any(|key| !matches!(key.as_str(), "model" | "instances" | "parameters"))
{
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding predict body contains unsupported top-level fields",
));
}
let Some(instances) = object.get("instances").and_then(Value::as_array) else {
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding predict body must contain an instances array",
));
};
if instances.is_empty() {
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding predict body must contain at least one instance",
));
}
if object
.get("parameters")
.is_some_and(|parameters| !parameters.is_object())
{
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding predict parameters must be an object",
));
}
Ok(())
}
fn build_vertex_predict_instance(item: &Map<String, Value>) -> Option<Value> {
let content = gemini_embedding_content_text(item.get("content")?)?;
let mut instance = Map::new();
instance.insert("content".to_string(), Value::String(content));
if let Some(task_type) = item.get("taskType") {
instance.insert(
"task_type".to_string(),
Value::String(task_type.as_str()?.to_string()),
);
}
if let Some(title) = item.get("title") {
instance.insert(
"title".to_string(),
Value::String(title.as_str()?.to_string()),
);
}
Some(Value::Object(instance))
}
fn gemini_embedding_content_text(content: &Value) -> Option<String> {
let parts = content
.as_object()?
.get("parts")?
.as_array()?
.iter()
.filter_map(|part| part.as_object()?.get("text")?.as_str())
.filter(|text| !text.trim().is_empty())
.collect::<Vec<_>>();
if parts.is_empty() {
return None;
}
Some(parts.join(""))
}
fn insert_shared_parameter(
items: &[&Map<String, Value>],
parameters: &mut Map<String, Value>,
key: &str,
) -> Result<(), TransportRequestBodySemanticsError> {
let mut value: Option<Value> = None;
for item in items {
let Some(next) = item.get(key) else {
continue;
};
match &value {
Some(current) if current != next => {
return Err(TransportRequestBodySemanticsError::new(
"Vertex Gemini embedding batch items must use the same shared parameters",
));
}
None => value = Some(next.clone()),
_ => {}
}
}
if let Some(value) = value {
parameters.insert(key.to_string(), value);
}
Ok(())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::apply_transport_request_body_semantics;
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
fn sample_transport(provider_type: &str, base_url: &str) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "provider".to_string(),
provider_type: provider_type.to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "gemini:embedding".to_string(),
api_family: Some("gemini".to_string()),
endpoint_kind: Some("embedding".to_string()),
is_active: true,
base_url: base_url.to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:embedding".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: "secret".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn vertex_gemini_embedding_single_body_uses_predict_contract() {
let transport = sample_transport("vertex_ai", "https://aiplatform.googleapis.com");
let mut body = json!({
"model": "gemini-embedding-2",
"content": {"parts": [{"text": "hello"}]},
"taskType": "RETRIEVAL_QUERY",
"outputDimensionality": 768
});
apply_transport_request_body_semantics(&mut body, &transport, "gemini:embedding")
.expect("body semantics should apply");
assert!(body.get("model").is_none());
assert!(body.get("content").is_none());
assert_eq!(body["instances"][0]["content"], "hello");
assert_eq!(body["instances"][0]["task_type"], "RETRIEVAL_QUERY");
assert_eq!(body["parameters"]["outputDimensionality"], 768);
}
#[test]
fn gemini_api_embedding_single_body_keeps_model_for_developer_api() {
let transport =
sample_transport("gemini", "https://generativelanguage.googleapis.com/v1beta");
let mut body = json!({
"model": "gemini-embedding-2",
"content": {"parts": [{"text": "hello"}]}
});
apply_transport_request_body_semantics(&mut body, &transport, "gemini:embedding")
.expect("developer API body should pass through");
assert_eq!(body["model"], "gemini-embedding-2");
}
#[test]
fn vertex_gemini_embedding_batch_body_uses_predict_instances() {
let transport = sample_transport("vertex_ai", "https://aiplatform.googleapis.com");
let mut body = json!({
"requests": [
{
"model": "models/gemini-embedding-2",
"content": {"parts": [{"text": "hello"}]}
},
{
"model": "models/gemini-embedding-2",
"content": {"parts": [{"text": "world"}]}
}
]
});
apply_transport_request_body_semantics(&mut body, &transport, "gemini:embedding")
.expect("batch body semantics should apply");
assert!(body.get("requests").is_none());
assert_eq!(body["instances"][0]["content"], "hello");
assert_eq!(body["instances"][1]["content"], "world");
}
#[test]
fn vertex_gemini_embedding_existing_predict_body_removes_duplicate_model() {
let transport = sample_transport("vertex_ai", "https://aiplatform.googleapis.com");
let mut body = json!({
"model": "gemini-embedding-2",
"instances": [
{"content": "hello", "task_type": "RETRIEVAL_QUERY"}
],
"parameters": {
"outputDimensionality": 768
}
});
apply_transport_request_body_semantics(&mut body, &transport, "gemini:embedding")
.expect("existing predict body should be accepted");
assert!(body.get("model").is_none());
assert_eq!(body["instances"][0]["content"], "hello");
assert_eq!(body["parameters"]["outputDimensionality"], 768);
}
#[test]
fn vertex_gemini_embedding_existing_predict_body_rejects_unconsumed_fields() {
let transport = sample_transport("vertex_ai", "https://aiplatform.googleapis.com");
let mut body = json!({
"model": "gemini-embedding-2",
"instances": [
{"content": "hello"}
],
"input": "this field would not be consumed by Vertex predict"
});
let error =
apply_transport_request_body_semantics(&mut body, &transport, "gemini:embedding")
.expect_err("predict body must not carry unconsumed OpenAI fields");
assert!(error.message().contains("unsupported top-level fields"));
assert!(body.get("model").is_some());
}
#[test]
fn vertex_gemini_embedding_rejects_unconverted_openai_body() {
let transport = sample_transport("vertex_ai", "https://aiplatform.googleapis.com");
let mut body = json!({
"model": "gemini-embedding-2",
"input": "hello"
});
let error =
apply_transport_request_body_semantics(&mut body, &transport, "gemini:embedding")
.expect_err("OpenAI embedding body must not be sent to Vertex native predict");
assert!(error.message().contains("cannot be mapped"));
assert!(body.get("input").is_some());
}
}

358
crates/aether-provider-transport/src/request_url/mod.rs
View File

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
use std::sync::OnceLock;
use regex::Regex;
use serde_json::Value;
use url::form_urlencoded;
use crate::antigravity::{
@@ -12,11 +13,14 @@ use crate::claude_code::build_claude_code_messages_url;
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
build_openai_responses_url, build_passthrough_path_url,
google_openai_compat_base_includes_api_root, normalize_gemini_content_action_path,
};
use crate::vertex::{
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,
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, resolve_local_vertex_api_key_query_auth,
resolve_local_vertex_service_account_auth_config,
};
#[derive(Debug, Clone, Copy)]
@@ -32,20 +36,51 @@ pub fn build_transport_request_url(
transport: &GatewayProviderTransportSnapshot,
params: TransportRequestUrlParams<'_>,
) -> Option<String> {
build_transport_request_url_inner(transport, params, false)
}
pub fn build_transport_request_url_for_request_body(
transport: &GatewayProviderTransportSnapshot,
params: TransportRequestUrlParams<'_>,
provider_request_body: Option<&Value>,
) -> Option<String> {
let gemini_embedding_batch =
gemini_embedding_request_body_uses_batch(params.provider_api_format, provider_request_body);
build_transport_request_url_inner(transport, params, gemini_embedding_batch)
}
pub fn gemini_embedding_request_body_uses_batch(
provider_api_format: &str,
provider_request_body: Option<&Value>,
) -> bool {
aether_ai_formats::normalize_api_format_alias(provider_api_format) == "gemini:embedding"
&& provider_request_body
.and_then(|body| body.get("requests"))
.and_then(Value::as_array)
.is_some_and(|requests| !requests.is_empty())
}
fn build_transport_request_url_inner(
transport: &GatewayProviderTransportSnapshot,
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 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
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|path| expand_custom_path_template(path, build_path_params(params)));
.map(|path| {
expand_custom_path_template(path, build_path_params(params, gemini_embedding_batch))
});
if let Some(path) = custom_path.as_deref() {
let blocked_keys = if normalized_provider_api_format.starts_with("gemini:") {
@@ -55,6 +90,8 @@ pub fn build_transport_request_url(
};
let normalized_path = if normalized_provider_api_format == "gemini:generate_content" {
normalize_gemini_content_action_path(path, params.upstream_is_stream)
} else if normalized_provider_api_format == "gemini:embedding" {
normalize_gemini_embedding_action_path(path, gemini_embedding_batch)
} else {
path.to_string()
};
@@ -106,6 +143,7 @@ pub fn build_transport_request_url(
&transport.endpoint.base_url,
params.mapped_model?,
params.request_query,
gemini_embedding_batch,
),
"doubao:embedding" => build_passthrough_path_url(
&transport.endpoint.base_url,
@@ -246,25 +284,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) {
@@ -287,7 +342,10 @@ fn build_transport_hook_url(
None
}
fn build_path_params(params: TransportRequestUrlParams<'_>) -> BTreeMap<&'static str, &str> {
fn build_path_params(
params: TransportRequestUrlParams<'_>,
gemini_embedding_batch: bool,
) -> BTreeMap<&'static str, &str> {
let mut path_params = BTreeMap::new();
if let Some(model) = params
.mapped_model
@@ -302,7 +360,11 @@ fn build_path_params(params: TransportRequestUrlParams<'_>) -> BTreeMap<&'static
path_params.insert(
"action",
if provider_api_format == "gemini:embedding" {
"embedContent"
if gemini_embedding_batch {
"batchEmbedContents"
} else {
"embedContent"
}
} else if params.upstream_is_stream {
"streamGenerateContent"
} else {
@@ -313,6 +375,14 @@ fn build_path_params(params: TransportRequestUrlParams<'_>) -> BTreeMap<&'static
path_params
}
fn normalize_gemini_embedding_action_path(path: &str, batch: bool) -> String {
if batch {
path.replace(":embedContent", ":batchEmbedContents")
} else {
path.replace(":batchEmbedContents", ":embedContent")
}
}
fn build_provider_embedding_v1_url(upstream_base_url: &str, query: Option<&str>) -> Option<String> {
build_provider_v1_url(upstream_base_url, "/embeddings", "/v1/embeddings", query)
}
@@ -333,7 +403,9 @@ fn build_provider_v1_url(
.map(|(base, _)| base)
.unwrap_or_else(|| upstream_base_url.trim())
.trim_end_matches('/');
let path = if base_without_query.ends_with("/v1") {
let path = if base_without_query.ends_with("/v1")
|| google_openai_compat_base_includes_api_root(base_without_query)
{
v1_path
} else {
default_path
@@ -345,6 +417,7 @@ fn build_gemini_embedding_url(
upstream_base_url: &str,
model: &str,
query: Option<&str>,
batch: bool,
) -> Option<String> {
let trimmed_base_url = upstream_base_url
.trim()
@@ -357,12 +430,17 @@ fn build_gemini_embedding_url(
return None;
}
let path = if trimmed_base_url.ends_with("/v1beta") {
format!("/models/{trimmed_model}:embedContent")
} else if trimmed_base_url.contains("/v1beta/models/") {
":embedContent".to_string()
let action = if batch {
"batchEmbedContents"
} else {
format!("/v1beta/models/{trimmed_model}:embedContent")
"embedContent"
};
let path = if trimmed_base_url.ends_with("/v1beta") {
format!("/models/{trimmed_model}:{action}")
} else if trimmed_base_url.contains("/v1beta/models/") {
format!(":{action}")
} else {
format!("/v1beta/models/{trimmed_model}:{action}")
};
build_passthrough_path_url(upstream_base_url, &path, query, &["key"])
}
@@ -440,12 +518,13 @@ fn custom_path_template_regex() -> &'static Regex {
mod tests {
use super::{
build_kiro_cross_format_upstream_url, build_transport_request_url,
TransportRequestUrlParams,
build_transport_request_url_for_request_body, TransportRequestUrlParams,
};
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use serde_json::json;
fn sample_transport(
provider_type: &str,
@@ -572,6 +651,96 @@ 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:predict?foo=bar"
);
}
#[test]
fn vertex_gemini_embedding_batch_request_uses_vertex_predict_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_eq!(
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),
)
.as_deref(),
Some(
"https://aiplatform.googleapis.com/v1/projects/demo-project/locations/global/publishers/google/models/gemini-embedding-2:predict"
)
);
}
#[test]
fn builds_openai_responses_url_for_formal_format_name() {
let transport = sample_transport(
@@ -829,6 +998,129 @@ mod tests {
);
}
#[test]
fn embedding_request_url_preserves_google_openai_compat_roots() {
let developer_api_openai = sample_transport(
"custom",
"openai:embedding",
"https://generativelanguage.googleapis.com/v1beta/openai",
None,
);
let vertex_openai = sample_transport(
"custom",
"openai:embedding",
"https://aiplatform.googleapis.com/v1/projects/project-1/locations/global/endpoints/openapi",
None,
);
assert_eq!(
build_transport_request_url(
&developer_api_openai,
TransportRequestUrlParams {
provider_api_format: "openai:embedding",
mapped_model: Some("gemini-embedding-001"),
upstream_is_stream: false,
request_query: Some("trace=1"),
kiro_api_region: None,
},
)
.as_deref(),
Some("https://generativelanguage.googleapis.com/v1beta/openai/embeddings?trace=1")
);
assert_eq!(
build_transport_request_url(
&vertex_openai,
TransportRequestUrlParams {
provider_api_format: "openai:embedding",
mapped_model: Some("gemini-embedding-001"),
upstream_is_stream: false,
request_query: None,
kiro_api_region: None,
},
)
.as_deref(),
Some(
"https://aiplatform.googleapis.com/v1/projects/project-1/locations/global/endpoints/openapi/embeddings"
)
);
}
#[test]
fn gemini_embedding_batch_body_uses_batch_endpoint() {
let gemini = sample_transport(
"gemini",
"gemini:embedding",
"https://generativelanguage.googleapis.com/v1beta",
None,
);
let batch_body = json!({
"requests": [
{
"model": "models/gemini-embedding-001",
"content": {"parts": [{"text": "alpha"}]}
},
{
"model": "models/gemini-embedding-001",
"content": {"parts": [{"text": "beta"}]}
}
]
});
assert_eq!(
build_transport_request_url_for_request_body(
&gemini,
TransportRequestUrlParams {
provider_api_format: "gemini:embedding",
mapped_model: Some("gemini-embedding-001"),
upstream_is_stream: false,
request_query: Some("key=client-key&foo=bar"),
kiro_api_region: None,
},
Some(&batch_body),
)
.as_deref(),
Some(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents?foo=bar"
)
);
}
#[test]
fn gemini_embedding_custom_action_template_follows_batch_body() {
let gemini = sample_transport(
"gemini",
"gemini:embedding",
"https://generativelanguage.googleapis.com",
Some("/v1beta/models/{model}:{action}"),
);
let batch_body = json!({
"requests": [
{
"model": "models/gemini-embedding-001",
"content": {"parts": [{"text": "alpha"}]}
}
]
});
assert_eq!(
build_transport_request_url_for_request_body(
&gemini,
TransportRequestUrlParams {
provider_api_format: "gemini:embedding",
mapped_model: Some("gemini-embedding-001"),
upstream_is_stream: false,
request_query: None,
kiro_api_region: None,
},
Some(&batch_body),
)
.as_deref(),
Some(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents"
)
);
}
#[test]
fn rerank_request_url_builds_provider_default_paths() {
let openai = sample_transport(

9
crates/aether-provider-transport/src/same_format_provider/mod.rs
View File

@@ -27,7 +27,10 @@ use crate::vertex::{
is_vertex_service_account_transport_context, is_vertex_transport_context,
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_for_request_body, ensure_upstream_auth_header,
TransportRequestUrlParams,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SameFormatProviderFamily {
@@ -77,6 +80,7 @@ pub struct SameFormatProviderUpstreamUrlParams<'a> {
pub upstream_is_stream: bool,
pub request_query: Option<&'a str>,
pub kiro_api_region: Option<&'a str>,
pub provider_request_body: Option<&'a Value>,
}
#[derive(Debug, Clone, Copy)]
@@ -257,7 +261,7 @@ pub fn build_same_format_provider_upstream_url(
transport: &GatewayProviderTransportSnapshot,
params: SameFormatProviderUpstreamUrlParams<'_>,
) -> Option<String> {
build_transport_request_url(
build_transport_request_url_for_request_body(
transport,
TransportRequestUrlParams {
provider_api_format: params.provider_api_format,
@@ -266,6 +270,7 @@ pub fn build_same_format_provider_upstream_url(
request_query: params.request_query,
kiro_api_region: params.kiro_api_region,
},
params.provider_request_body,
)
}

57
crates/aether-provider-transport/src/url.rs
View File

@@ -2,15 +2,17 @@ use std::collections::BTreeMap;
use super::provider_types::is_codex_cli_backend_url;
use url::form_urlencoded;
use url::Url;
pub fn build_openai_chat_url(upstream_base_url: &str, query: Option<&str>) -> String {
let (trimmed, base_query) = split_base_url_query(upstream_base_url);
let trimmed = trimmed.trim_end_matches('/');
let mut url = if trimmed.ends_with("/v1") {
format!("{trimmed}/chat/completions")
} else {
format!("{trimmed}/v1/chat/completions")
};
let mut url =
if trimmed.ends_with("/v1") || google_openai_compat_base_includes_api_root(trimmed) {
format!("{trimmed}/chat/completions")
} else {
format!("{trimmed}/v1/chat/completions")
};
append_merged_query(&mut url, base_query, None, query, &[]);
url
}
@@ -195,6 +197,33 @@ fn split_base_url_query(base_url: &str) -> (&str, Option<&str>) {
.unwrap_or((trimmed, None))
}
pub(crate) fn google_openai_compat_base_includes_api_root(base_url: &str) -> bool {
let Ok(parsed) = Url::parse(base_url.trim()) else {
return false;
};
let Some(host) = parsed.host_str().map(|value| value.to_ascii_lowercase()) else {
return false;
};
let path = parsed.path().trim_end_matches('/');
if host == "generativelanguage.googleapis.com" {
return path == "/v1beta/openai" || path == "/v1/openai";
}
if looks_like_vertex_ai_host(&host) {
return path.ends_with("/endpoints/openapi");
}
false
}
fn looks_like_vertex_ai_host(host: &str) -> bool {
const VERTEX_AI_HOST: &str = "aiplatform.googleapis.com";
host == VERTEX_AI_HOST
|| host.ends_with(&format!(".{VERTEX_AI_HOST}"))
|| host.ends_with(&format!("-{VERTEX_AI_HOST}"))
}
fn split_path_query(path: &str) -> (&str, Option<&str>) {
path.split_once('?')
.map(|(path, query)| (path, Some(query)))
@@ -294,6 +323,24 @@ mod tests {
);
}
#[test]
fn openai_chat_url_preserves_google_openai_compat_roots() {
assert_eq!(
build_openai_chat_url(
"https://generativelanguage.googleapis.com/v1beta/openai",
Some("trace=1")
),
"https://generativelanguage.googleapis.com/v1beta/openai/chat/completions?trace=1"
);
assert_eq!(
build_openai_chat_url(
"https://aiplatform.googleapis.com/v1/projects/project-1/locations/global/endpoints/openapi",
None,
),
"https://aiplatform.googleapis.com/v1/projects/project-1/locations/global/endpoints/openapi/chat/completions"
);
}
#[test]
fn openai_responses_url_preserves_codex_path_prefix() {
assert_eq!(

39
crates/aether-provider-transport/src/vertex/context.rs
View File

@@ -82,7 +82,20 @@ fn is_vertex_host_format_context(transport: &GatewayProviderTransportSnapshot) -
}
let endpoint_api_format = transport.endpoint.api_format.trim().to_ascii_lowercase();
endpoint_api_format.starts_with("gemini:") || endpoint_api_format.starts_with("claude:")
endpoint_api_format.starts_with("gemini:")
|| endpoint_api_format.starts_with("claude:")
|| (endpoint_api_format.starts_with("openai:")
&& looks_like_vertex_openai_compat_base(&transport.endpoint.base_url))
}
fn looks_like_vertex_openai_compat_base(base_url: &str) -> bool {
let Ok(parsed) = Url::parse(base_url.trim()) else {
return false;
};
parsed
.path()
.trim_end_matches('/')
.ends_with("/endpoints/openapi")
}
#[cfg(test)]
@@ -198,4 +211,28 @@ mod tests {
"claude:messages"
));
}
#[test]
fn infers_vertex_service_account_context_for_openai_compat_endpoint_root() {
let mut transport = sample_transport();
transport.endpoint.api_format = "openai:chat".to_string();
transport.endpoint.base_url =
"https://aiplatform.googleapis.com/v1/projects/project-1/locations/global/endpoints/openapi"
.to_string();
transport.key.auth_type = "service_account".to_string();
assert!(is_vertex_service_account_transport_context(&transport));
assert!(is_vertex_transport_context(&transport));
}
#[test]
fn does_not_infer_vertex_context_for_generic_openai_format_on_aiplatform_root() {
let mut transport = sample_transport();
transport.endpoint.api_format = "openai:chat".to_string();
transport.endpoint.base_url = "https://aiplatform.googleapis.com".to_string();
transport.key.auth_type = "service_account".to_string();
assert!(!is_vertex_service_account_transport_context(&transport));
assert!(!is_vertex_transport_context(&transport));
}
}

5
crates/aether-provider-transport/src/vertex/mod.rs
View File

@@ -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,
};

31
crates/aether-provider-transport/src/vertex/policy.rs
View File

@@ -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();

81
crates/aether-provider-transport/src/vertex/url.rs
View File

@@ -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, "predict", 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,
"predict",
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(), &[])

10
crates/aether-usage-runtime/src/report.rs
View File

@@ -234,6 +234,9 @@ pub fn is_local_ai_sync_report_kind(report_kind: &str) -> bool {
| "openai_cli_sync_success"
| "openai_image_sync_success"
| "openai_image_sync_error"
| "openai_embedding_sync_success"
| "openai_embedding_sync_error"
| "gemini_embedding_sync_success"
| "claude_cli_sync_success"
| "gemini_cli_sync_success"
| "openai_cli_sync_error"
@@ -426,6 +429,13 @@ mod tests {
));
assert!(is_local_ai_sync_report_kind("openai_image_sync_success"));
assert!(is_local_ai_sync_report_kind("openai_image_sync_error"));
assert!(is_local_ai_sync_report_kind(
"openai_embedding_sync_success"
));
assert!(is_local_ai_sync_report_kind("openai_embedding_sync_error"));
assert!(is_local_ai_sync_report_kind(
"gemini_embedding_sync_success"
));
assert!(is_local_ai_sync_report_kind("gemini_files_delete_mapping"));
assert!(!is_local_ai_sync_report_kind("unknown_sync_kind"));
}