This commit is contained in:
fawney19
2026-05-19 03:16:35 +08:00
77 changed files with 4051 additions and 333 deletions

View File

@@ -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",

View File

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

View File

@@ -219,6 +219,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(&[]);