fix(gateway): route Gemini embedding batches correctly

2026-10-09 18:59:50 +08:00 · 2026-05-17 14:24:32 +00:00
parent a2f91b4108
commit 81ff375bfd
37 changed files with 1353 additions and 120 deletions
@@ -78,6 +78,87 @@ fn embedding_execution_runtime() -> Router {
    )
 }

+fn gemini_embedding_success_state(
+    execution_runtime_url: String,
+    client_api_format: &str,
+) -> AppState {
+    let mut snapshot = sample_currently_usable_auth_snapshot(
+        "key-gemini-embedding-success",
+        "user-gemini-embedding-success",
+    );
+    snapshot.user_allowed_providers = None;
+    snapshot.api_key_allowed_providers = None;
+    snapshot.user_allowed_api_formats = Some(vec![client_api_format.to_string()]);
+    snapshot.api_key_allowed_api_formats = Some(vec![client_api_format.to_string()]);
+    snapshot.user_allowed_models = Some(vec!["gemini-embedding-2-preview".to_string()]);
+    snapshot.api_key_allowed_models = Some(vec!["gemini-embedding-2-preview".to_string()]);
+    let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
+        Some(hash_api_key("sk-gemini-embedding-success")),
+        snapshot,
+    )]));
+    let candidate_repository =
+        Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
+            gemini_embedding_candidate_row(),
+        ]));
+    let mut provider = sample_provider("provider-gemini-embedding", "Gemini Embeddings", 1);
+    provider.provider_type = "gemini".to_string();
+    let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+        vec![provider],
+        vec![sample_endpoint(
+            "endpoint-gemini-embedding",
+            "provider-gemini-embedding",
+            "gemini:embedding",
+            "https://generativelanguage.googleapis.com/v1beta",
+        )],
+        vec![sample_key(
+            "key-upstream-gemini-embedding",
+            "provider-gemini-embedding",
+            "gemini:embedding",
+            "sk-upstream-gemini-embedding",
+        )],
+    ));
+    let data_state =
+        GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
+            provider_catalog_repository,
+            candidate_repository,
+        )
+        .with_auth_api_key_reader(auth_repository)
+        .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
+
+    build_state_with_execution_runtime_override(execution_runtime_url)
+        .with_data_state_for_tests(data_state)
+}
+
+fn gemini_embedding_conversion_execution_runtime() -> Router {
+    Router::new().route(
+        "/v1/execute/sync",
+        any(|Json(plan): Json<ExecutionPlan>| async move {
+            assert_openai_to_gemini_embedding_execution_plan(&plan);
+            Json(gemini_embedding_execution_result(&plan))
+        }),
+    )
+}
+
+fn gemini_embedding_batch_conversion_execution_runtime() -> Router {
+    Router::new().route(
+        "/v1/execute/sync",
+        any(|Json(plan): Json<ExecutionPlan>| async move {
+            assert_openai_to_gemini_batch_embedding_execution_plan(&plan);
+            Json(gemini_batch_embedding_execution_result(&plan))
+        }),
+    )
+}
+
+fn gemini_embedding_native_execution_runtime() -> Router {
+    Router::new().route(
+        "/v1/execute/sync",
+        any(|Json(plan): Json<ExecutionPlan>| async move {
+            assert_native_gemini_embedding_execution_plan(&plan);
+            Json(gemini_embedding_execution_result(&plan))
+        }),
+    )
+}
+
 fn embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
    StoredMinimalCandidateSelectionRow {
        provider_id: "provider-embedding".to_string(),
@@ -112,6 +193,40 @@ fn embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
    }
 }

+fn gemini_embedding_candidate_row() -> StoredMinimalCandidateSelectionRow {
+    StoredMinimalCandidateSelectionRow {
+        provider_id: "provider-gemini-embedding".to_string(),
+        provider_name: "Gemini Embeddings".to_string(),
+        provider_type: "gemini".to_string(),
+        provider_priority: 1,
+        provider_is_active: true,
+        endpoint_id: "endpoint-gemini-embedding".to_string(),
+        endpoint_api_format: "gemini:embedding".to_string(),
+        endpoint_api_family: Some("gemini".to_string()),
+        endpoint_kind: Some("embedding".to_string()),
+        endpoint_is_active: true,
+        key_id: "key-upstream-gemini-embedding".to_string(),
+        key_name: "default".to_string(),
+        key_auth_type: "api_key".to_string(),
+        key_is_active: true,
+        key_api_formats: Some(vec!["gemini:embedding".to_string()]),
+        key_allowed_models: None,
+        key_capabilities: None,
+        key_internal_priority: 50,
+        key_global_priority_by_format: None,
+        model_id: "model-gemini-embedding-preview".to_string(),
+        global_model_id: "global-gemini-embedding-preview".to_string(),
+        global_model_name: "gemini-embedding-2-preview".to_string(),
+        global_model_mappings: None,
+        global_model_supports_streaming: Some(false),
+        model_provider_model_name: "gemini-embedding-2-preview".to_string(),
+        model_provider_model_mappings: None,
+        model_supports_streaming: Some(false),
+        model_is_active: true,
+        model_is_available: true,
+    }
+}
+
 fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
    assert_eq!(plan.client_api_format, "openai:embedding");
    assert_eq!(plan.provider_api_format, "openai:embedding");
@@ -123,6 +238,78 @@ fn assert_embedding_execution_plan(plan: &ExecutionPlan) {
    assert!(body.get("input").is_some());
 }

+fn assert_openai_to_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
+    assert_eq!(plan.client_api_format, "openai:embedding");
+    assert_eq!(plan.provider_api_format, "gemini:embedding");
+    assert_eq!(plan.method, "POST");
+    assert_eq!(
+        plan.url,
+        "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent"
+    );
+    assert_eq!(
+        plan.headers.get("x-goog-api-key").map(String::as_str),
+        Some("sk-upstream-gemini-embedding")
+    );
+    assert_eq!(
+        plan.model_name.as_deref(),
+        Some("gemini-embedding-2-preview")
+    );
+    assert!(!plan.stream);
+    let body = plan.body.json_body.as_ref().expect("json request body");
+    assert_eq!(body["model"], "gemini-embedding-2-preview");
+    assert_eq!(body["content"]["parts"][0]["text"], "hello");
+    assert!(body.get("input").is_none());
+    assert!(body.get("messages").is_none());
+}
+
+fn assert_openai_to_gemini_batch_embedding_execution_plan(plan: &ExecutionPlan) {
+    assert_eq!(plan.client_api_format, "openai:embedding");
+    assert_eq!(plan.provider_api_format, "gemini:embedding");
+    assert_eq!(plan.method, "POST");
+    assert_eq!(
+        plan.url,
+        "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:batchEmbedContents"
+    );
+    assert_eq!(
+        plan.headers.get("x-goog-api-key").map(String::as_str),
+        Some("sk-upstream-gemini-embedding")
+    );
+    assert!(!plan.stream);
+    let body = plan.body.json_body.as_ref().expect("json request body");
+    assert!(body.get("model").is_none());
+    let requests = body["requests"].as_array().expect("batch requests");
+    assert_eq!(requests.len(), 2);
+    assert_eq!(requests[0]["model"], "models/gemini-embedding-2-preview");
+    assert_eq!(requests[0]["content"]["parts"][0]["text"], "hello");
+    assert_eq!(requests[1]["model"], "models/gemini-embedding-2-preview");
+    assert_eq!(requests[1]["content"]["parts"][0]["text"], "world");
+    assert!(body.get("input").is_none());
+    assert!(body.get("messages").is_none());
+}
+
+fn assert_native_gemini_embedding_execution_plan(plan: &ExecutionPlan) {
+    assert_eq!(plan.client_api_format, "gemini:embedding");
+    assert_eq!(plan.provider_api_format, "gemini:embedding");
+    assert_eq!(plan.method, "POST");
+    assert_eq!(
+        plan.url,
+        "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent"
+    );
+    assert_eq!(
+        plan.headers.get("x-goog-api-key").map(String::as_str),
+        Some("sk-upstream-gemini-embedding")
+    );
+    assert_eq!(
+        plan.model_name.as_deref(),
+        Some("gemini-embedding-2-preview")
+    );
+    assert!(!plan.stream);
+    let body = plan.body.json_body.as_ref().expect("json request body");
+    assert_eq!(body["content"]["parts"][0]["text"], "hello");
+    assert!(body.get("input").is_none());
+    assert!(body.get("messages").is_none());
+}
+
 fn embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
    ExecutionResult {
        request_id: plan.request_id.clone(),
@@ -145,6 +332,55 @@ fn embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
    }
 }

+fn gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
+    ExecutionResult {
+        request_id: plan.request_id.clone(),
+        candidate_id: plan.candidate_id.clone(),
+        status_code: 200,
+        headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
+        body: Some(ResponseBody {
+            json_body: Some(json!({
+                "model": "gemini-embedding-2-preview",
+                "embedding": {
+                    "values": [0.1, 0.2, 0.3]
+                },
+                "usageMetadata": {
+                    "promptTokenCount": 4,
+                    "totalTokenCount": 4
+                }
+            })),
+            body_bytes_b64: None,
+        }),
+        telemetry: None,
+        error: None,
+    }
+}
+
+fn gemini_batch_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
+    ExecutionResult {
+        request_id: plan.request_id.clone(),
+        candidate_id: plan.candidate_id.clone(),
+        status_code: 200,
+        headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
+        body: Some(ResponseBody {
+            json_body: Some(json!({
+                "model": "gemini-embedding-2-preview",
+                "embeddings": [
+                    {"values": [0.1, 0.2, 0.3]},
+                    {"values": [0.4, 0.5, 0.6]}
+                ],
+                "usageMetadata": {
+                    "promptTokenCount": 8,
+                    "totalTokenCount": 8
+                }
+            })),
+            body_bytes_b64: None,
+        }),
+        telemetry: None,
+        error: None,
+    }
+}
+
 #[tokio::test]
 async fn embeddings_route_accepts_openai_payload() {
    let (execution_runtime_url, execution_runtime_handle) =
@@ -215,6 +451,157 @@ async fn embeddings_route_accepts_openai_payload() {
    execution_runtime_handle.abort();
 }

+#[tokio::test]
+async fn embeddings_route_converts_openai_payload_to_gemini_embedding_provider() {
+    let (execution_runtime_url, execution_runtime_handle) =
+        start_server(gemini_embedding_conversion_execution_runtime()).await;
+    let gateway = build_router_with_state(gemini_embedding_success_state(
+        execution_runtime_url,
+        "openai:embedding",
+    ));
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+
+    let response = reqwest::Client::new()
+        .post(format!("{gateway_url}/v1/embeddings"))
+        .header(
+            http::header::AUTHORIZATION,
+            "Bearer sk-gemini-embedding-success",
+        )
+        .json(&json!({
+            "model": "gemini-embedding-2-preview",
+            "input": "hello"
+        }))
+        .send()
+        .await
+        .expect("request should succeed");
+
+    assert_eq!(response.status(), StatusCode::OK);
+    assert_eq!(
+        response
+            .headers()
+            .get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
+            .and_then(|value| value.to_str().ok()),
+        Some("openai:embedding")
+    );
+    assert_eq!(
+        response
+            .headers()
+            .get(CONTROL_EXECUTION_RUNTIME_HEADER)
+            .and_then(|value| value.to_str().ok()),
+        Some("true")
+    );
+    let payload: serde_json::Value = response.json().await.expect("body should parse");
+    assert_eq!(payload["object"], "list");
+    assert_eq!(payload["model"], "gemini-embedding-2-preview");
+    assert_eq!(payload["data"][0]["object"], "embedding");
+    assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
+    assert_eq!(payload["usage"]["prompt_tokens"], json!(4));
+    assert_eq!(payload["usage"]["total_tokens"], json!(4));
+
+    gateway_handle.abort();
+    execution_runtime_handle.abort();
+}
+
+#[tokio::test]
+async fn embeddings_route_converts_openai_batch_payload_to_gemini_batch_endpoint() {
+    let (execution_runtime_url, execution_runtime_handle) =
+        start_server(gemini_embedding_batch_conversion_execution_runtime()).await;
+    let gateway = build_router_with_state(gemini_embedding_success_state(
+        execution_runtime_url,
+        "openai:embedding",
+    ));
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+
+    let response = reqwest::Client::new()
+        .post(format!("{gateway_url}/v1/embeddings"))
+        .header(
+            http::header::AUTHORIZATION,
+            "Bearer sk-gemini-embedding-success",
+        )
+        .json(&json!({
+            "model": "gemini-embedding-2-preview",
+            "input": ["hello", "world"]
+        }))
+        .send()
+        .await
+        .expect("request should succeed");
+
+    assert_eq!(response.status(), StatusCode::OK);
+    assert_eq!(
+        response
+            .headers()
+            .get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
+            .and_then(|value| value.to_str().ok()),
+        Some("openai:embedding")
+    );
+    let payload: serde_json::Value = response.json().await.expect("body should parse");
+    assert_eq!(payload["object"], "list");
+    assert_eq!(payload["data"].as_array().map(Vec::len), Some(2));
+    assert_eq!(payload["data"][0]["index"], json!(0));
+    assert_eq!(payload["data"][0]["embedding"], json!([0.1, 0.2, 0.3]));
+    assert_eq!(payload["data"][1]["index"], json!(1));
+    assert_eq!(payload["data"][1]["embedding"], json!([0.4, 0.5, 0.6]));
+    assert_eq!(payload["usage"]["prompt_tokens"], json!(8));
+    assert_eq!(payload["usage"]["total_tokens"], json!(8));
+
+    gateway_handle.abort();
+    execution_runtime_handle.abort();
+}
+
+#[tokio::test]
+async fn gemini_embed_content_route_uses_native_gemini_embedding_provider() {
+    let (execution_runtime_url, execution_runtime_handle) =
+        start_server(gemini_embedding_native_execution_runtime()).await;
+    let gateway = build_router_with_state(gemini_embedding_success_state(
+        execution_runtime_url,
+        "gemini:embedding",
+    ));
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+
+    let response = reqwest::Client::new()
+        .post(format!(
+            "{gateway_url}/v1beta/models/gemini-embedding-2-preview:embedContent"
+        ))
+        .header("x-goog-api-key", "sk-gemini-embedding-success")
+        .json(&json!({
+            "content": {
+                "parts": [{"text": "hello"}]
+            }
+        }))
+        .send()
+        .await
+        .expect("request should succeed");
+
+    assert_eq!(response.status(), StatusCode::OK);
+    assert_eq!(
+        response
+            .headers()
+            .get(CONTROL_ROUTE_FAMILY_HEADER)
+            .and_then(|value| value.to_str().ok()),
+        Some("gemini")
+    );
+    assert_eq!(
+        response
+            .headers()
+            .get(CONTROL_ROUTE_KIND_HEADER)
+            .and_then(|value| value.to_str().ok()),
+        Some("embedding")
+    );
+    assert_eq!(
+        response
+            .headers()
+            .get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
+            .and_then(|value| value.to_str().ok()),
+        Some("gemini:embedding")
+    );
+    let payload: serde_json::Value = response.json().await.expect("body should parse");
+    assert_eq!(payload["embedding"]["values"], json!([0.1, 0.2, 0.3]));
+    assert_eq!(payload["model"], "gemini-embedding-2-preview");
+
+    gateway_handle.abort();
+    execution_runtime_handle.abort();
+}
+
 #[tokio::test]
 async fn embeddings_route_accepts_all_canonical_input_shapes() {
    let (execution_runtime_url, execution_runtime_handle) =