From 4b4def3a1141e9169f0d29927af95371b175065d Mon Sep 17 00:00:00 2001 From: DarkSky <25152247+darkskygit@users.noreply.github.com> Date: Thu, 14 May 2026 04:12:49 +0800 Subject: [PATCH] feat(server): gemini embedding 2 support (#14956) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #### PR Dependency Tree * **PR #14956** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal) ## Summary by CodeRabbit ## Release Notes * **Bug Fixes** * Improved Gemini Vertex provider configuration validation logic for enhanced reliability. * Refined Google Vertex publisher base URL construction handling. * **Tests** * Added test coverage for Gemini Embedding 2 model resolution. * Added test coverage for Gemini Vertex provider Google Cloud integration. [![Review Change Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/toeverything/AFFiNE/pull/14956) --- Cargo.lock | 4 ++-- .../native/src/llm/core/model_registry.rs | 15 +++++++++++++++ .../__tests__/copilot/native-provider.spec.ts | 12 ++++++++++++ .../plugins/copilot/providers/gemini/vertex.ts | 8 ++++++-- .../src/plugins/copilot/providers/utils.ts | 18 +++++++++++------- 5 files changed, 46 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d105277554..28994bf484 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3625,9 +3625,9 @@ dependencies = [ [[package]] name = "llm_adapter" -version = "0.2.5" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6e139f0a1609d6078293140fb7e281cf2bd5a45a7a29ef39f8606c803be7822" +checksum = "8ca30267ba36e247d1ff7a916a03db2ceb1de7f0bfcab7250cde006cdda68c19" dependencies = [ "base64", "jsonschema", diff --git a/packages/backend/native/src/llm/core/model_registry.rs b/packages/backend/native/src/llm/core/model_registry.rs index 227dc96e1a..237007e613 100644 --- a/packages/backend/native/src/llm/core/model_registry.rs +++ b/packages/backend/native/src/llm/core/model_registry.rs @@ -111,6 +111,21 @@ mod tests { assert_eq!(response.variant.unwrap().raw_model_id, "gemini-embedding-001"); } + #[test] + fn should_resolve_gemini_embedding_2() { + let response = llm_resolve_model_registry_variant(ModelRegistryResolveRequest { + backend_kind: Some("gemini_api".to_string()), + model_id: "gemini-embedding-2".to_string(), + }) + .unwrap(); + let variant = response.variant.unwrap(); + + assert_eq!(variant.raw_model_id, "gemini-embedding-2"); + assert_eq!(variant.protocol.as_deref(), Some("gemini")); + assert_eq!(variant.request_layer.as_deref(), Some("gemini_api")); + assert_eq!(variant.display_name.as_deref(), Some("Gemini Embedding 2")); + } + #[test] fn should_keep_same_raw_id_as_two_backend_variants() { let api_variant = llm_resolve_model_registry_variant(ModelRegistryResolveRequest { diff --git a/packages/backend/server/src/__tests__/copilot/native-provider.spec.ts b/packages/backend/server/src/__tests__/copilot/native-provider.spec.ts index 1637d0e994..413c873af1 100644 --- a/packages/backend/server/src/__tests__/copilot/native-provider.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/native-provider.spec.ts @@ -25,6 +25,7 @@ import { type PromptMessage, type StreamObject, } from '../../plugins/copilot/providers/types'; +import { getVertexGoogleBaseUrl } from '../../plugins/copilot/providers/utils'; import { buildPromptStructuredResponseFromFields, buildStructuredResponseContract, @@ -1823,6 +1824,17 @@ test('GeminiVertexProvider should prefetch bearer token for native config', asyn t.snapshot(config); }); +test('GeminiVertexProvider should build project scoped Vertex base URL', t => { + t.is( + getVertexGoogleBaseUrl({ + project: 'p1', + location: 'us-central1', + googleAuthOptions: {}, + }), + 'https://us-central1-aiplatform.googleapis.com/v1/projects/p1/locations/us-central1/publishers/google' + ); +}); + test('GeminiVertexProvider should materialize remote attachments before native text path', async t => { const cases = [ { diff --git a/packages/backend/server/src/plugins/copilot/providers/gemini/vertex.ts b/packages/backend/server/src/plugins/copilot/providers/gemini/vertex.ts index 5baf6013e7..70bf196dc4 100644 --- a/packages/backend/server/src/plugins/copilot/providers/gemini/vertex.ts +++ b/packages/backend/server/src/plugins/copilot/providers/gemini/vertex.ts @@ -1,7 +1,11 @@ import type { LlmBackendConfig } from '../../../../native'; import type { CopilotProviderExecution } from '../provider-runtime-contract'; import { CopilotProviderType } from '../types'; -import { getGoogleAuth, type VertexProviderConfig } from '../utils'; +import { + getGoogleAuth, + getVertexGoogleBaseUrl, + type VertexProviderConfig, +} from '../utils'; import { GeminiProvider } from './gemini'; export type GeminiVertexConfig = VertexProviderConfig; @@ -10,7 +14,7 @@ export class GeminiVertexProvider extends GeminiProvider { override readonly type = CopilotProviderType.GeminiVertex; override configured(execution?: CopilotProviderExecution): boolean { const config = this.getConfig(execution); - return !!config.location && !!config.googleAuthOptions; + return !!getVertexGoogleBaseUrl(config) && !!config.googleAuthOptions; } protected async resolveVertexAuth(execution?: CopilotProviderExecution) { return await getGoogleAuth(this.getConfig(execution), 'google'); diff --git a/packages/backend/server/src/plugins/copilot/providers/utils.ts b/packages/backend/server/src/plugins/copilot/providers/utils.ts index 7fffddee44..3b7859cbc6 100644 --- a/packages/backend/server/src/plugins/copilot/providers/utils.ts +++ b/packages/backend/server/src/plugins/copilot/providers/utils.ts @@ -403,18 +403,22 @@ export function getVertexAnthropicBaseUrl(options: VertexProviderConfig) { return `https://${location}-aiplatform.googleapis.com/v1/projects/${project}/locations/${location}/publishers/anthropic`; } +export function getVertexGoogleBaseUrl(options: VertexProviderConfig) { + const normalizedBaseUrl = normalizeUrl(options.baseURL); + if (normalizedBaseUrl) return normalizedBaseUrl; + const { location, project } = options; + if (!location || !project) return undefined; + return `https://${location}-aiplatform.googleapis.com/v1/projects/${project}/locations/${location}/publishers/google`; +} + export async function getGoogleAuth( options: VertexProviderConfig, publisher: 'anthropic' | 'google' ) { function getBaseUrl() { - const normalizedBaseUrl = normalizeUrl(options.baseURL); - if (normalizedBaseUrl) return normalizedBaseUrl; - const { location } = options; - if (location) { - return `https://${location}-aiplatform.googleapis.com/v1beta1/publishers/${publisher}`; - } - return undefined; + return publisher === 'google' + ? getVertexGoogleBaseUrl(options) + : getVertexAnthropicBaseUrl(options); } async function generateAuthToken() {