mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
fix(gateway): route Gemini embedding batches correctly
This commit is contained in:
@@ -104,6 +104,17 @@ pub(super) fn classify_ai_public_route(
|
||||
"gemini:video",
|
||||
true,
|
||||
))
|
||||
} else if normalized_path.ends_with(":embedContent")
|
||||
|| normalized_path.ends_with(":batchEmbedContents")
|
||||
{
|
||||
Some(classified_with_request_auth_channel(
|
||||
"ai_public",
|
||||
"gemini",
|
||||
"embedding",
|
||||
"api_key",
|
||||
"gemini:embedding",
|
||||
true,
|
||||
))
|
||||
} else if is_gemini_cli_request(headers) {
|
||||
Some(classified_with_request_auth_channel(
|
||||
"ai_public",
|
||||
|
||||
@@ -229,6 +229,8 @@ pub(super) fn is_gemini_models_route(path: &str) -> bool {
|
||||
(path.starts_with("/v1/models/") || path.starts_with("/v1beta/models/"))
|
||||
&& (path.contains(":generateContent")
|
||||
|| path.contains(":streamGenerateContent")
|
||||
|| path.contains(":embedContent")
|
||||
|| path.contains(":batchEmbedContents")
|
||||
|| path.contains(":predictLongRunning"))
|
||||
}
|
||||
|
||||
|
||||
@@ -197,6 +197,44 @@ fn classifies_gemini_generate_content_api_key_without_cli_marker() {
|
||||
assert!(decision.is_execution_runtime_candidate());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_gemini_embed_content_as_embedding_route() {
|
||||
let headers = headers(&[("x-goog-api-key", "gemini-key")]);
|
||||
let uri: Uri = "/v1beta/models/gemini-embedding-2-preview:embedContent"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let decision =
|
||||
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
|
||||
|
||||
assert_eq!(decision.route_family.as_deref(), Some("gemini"));
|
||||
assert_eq!(decision.route_kind.as_deref(), Some("embedding"));
|
||||
assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key"));
|
||||
assert_eq!(
|
||||
decision.auth_endpoint_signature.as_deref(),
|
||||
Some("gemini:embedding")
|
||||
);
|
||||
assert!(decision.is_execution_runtime_candidate());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_gemini_batch_embed_contents_as_embedding_route() {
|
||||
let headers = headers(&[("x-goog-api-key", "gemini-key")]);
|
||||
let uri: Uri = "/v1beta/models/gemini-embedding-2-preview:batchEmbedContents"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let decision =
|
||||
classify_control_route(&http::Method::POST, &uri, &headers).expect("route should classify");
|
||||
|
||||
assert_eq!(decision.route_family.as_deref(), Some("gemini"));
|
||||
assert_eq!(decision.route_kind.as_deref(), Some("embedding"));
|
||||
assert_eq!(decision.request_auth_channel.as_deref(), Some("api_key"));
|
||||
assert_eq!(
|
||||
decision.auth_endpoint_signature.as_deref(),
|
||||
Some("gemini:embedding")
|
||||
);
|
||||
assert!(decision.is_execution_runtime_candidate());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_gemini_predict_long_running_as_video_route() {
|
||||
let headers = headers(&[]);
|
||||
|
||||
Reference in New Issue
Block a user