mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-02 02:00:49 +08:00
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 --> [](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:
Generated
+2
-2
@@ -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() {
|
||||
|
||||
Reference in New Issue
Block a user