feat(server): gemini embedding 2 support (#14956)

#### PR Dependency Tree


* **PR #14956** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## 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_stack_entry_start -->

[![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)

<!-- review_stack_entry_end -->

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
2026-05-14 04:12:49 +08:00
committed by GitHub
parent 2b22fe4692
commit 4b4def3a11
5 changed files with 46 additions and 11 deletions
Generated
+2 -2
View File
@@ -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",
@@ -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 {
@@ -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 = [
{
@@ -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<GeminiVertexConfig> {
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');
@@ -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() {