Skip to content

Commit 6004c1c

Browse files
gumadeirasgumadeiras
committed
memory: normalize Gemini embeddings
1 parent 60aed95 commit 6004c1c

File tree

7 files changed

+94
-23
lines changed

7 files changed

+94
-23
lines changed

‎src/memory/batch-gemini.test.ts‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
import { afterEach, beforeAll, describe, expect, it, vi } from "vitest";
22
import type { GeminiEmbeddingClient } from "./embeddings-gemini.js";
33

4+
function magnitude(values: number[]) {
5+
return Math.sqrt(values.reduce((sum, value) => sum + value * value, 0));
6+
}
7+
48
describe("runGeminiEmbeddingBatches", () => {
59
let runGeminiEmbeddingBatches: typeof import("./batch-gemini.js").runGeminiEmbeddingBatches;
610

@@ -56,7 +60,7 @@ describe("runGeminiEmbeddingBatches", () => {
5660
return new Response(
5761
JSON.stringify({
5862
key: "req-1",
59-
response: { embedding: { values: [0.1, 0.2, 0.3] } },
63+
response: { embedding: { values: [3, 4] } },
6064
}),
6165
{
6266
status: 200,
@@ -88,7 +92,11 @@ describe("runGeminiEmbeddingBatches", () => {
8892
concurrency: 1,
8993
});
9094

91-
expect(results.get("req-1")).toEqual([0.1, 0.2, 0.3]);
95+
const embedding = results.get("req-1");
96+
expect(embedding).toBeDefined();
97+
expect(embedding?.[0]).toBeCloseTo(0.6, 5);
98+
expect(embedding?.[1]).toBeCloseTo(0.8, 5);
99+
expect(magnitude(embedding ?? [])).toBeCloseTo(1, 5);
92100
expect(fetchMock).toHaveBeenCalledTimes(3);
93101
});
94102
});

‎src/memory/batch-gemini.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import {
44
type EmbeddingBatchExecutionParams,
55
} from "./batch-runner.js";
66
import { buildBatchHeaders, normalizeBatchBaseUrl } from "./batch-utils.js";
7+
import { sanitizeAndNormalizeEmbedding } from "./embedding-vectors.js";
78
import { debugEmbeddingsLog } from "./embeddings-debug.js";
89
import type { GeminiEmbeddingClient, GeminiTextEmbeddingRequest } from "./embeddings-gemini.js";
910
import { hashText } from "./internal.js";
@@ -346,7 +347,9 @@ export async function runGeminiEmbeddingBatches(
346347
errors.push(`${customId}: ${line.response.error.message}`);
347348
continue;
348349
}
349-
const embedding = line.embedding?.values ?? line.response?.embedding?.values ?? [];
350+
const embedding = sanitizeAndNormalizeEmbedding(
351+
line.embedding?.values ?? line.response?.embedding?.values ?? [],
352+
);
350353
if (embedding.length === 0) {
351354
errors.push(`${customId}: empty embedding`);
352355
continue;

‎src/memory/embedding-vectors.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
export function sanitizeAndNormalizeEmbedding(vec: number[]): number[] {
2+
const sanitized = vec.map((value) => (Number.isFinite(value) ? value : 0));
3+
const magnitude = Math.sqrt(sanitized.reduce((sum, value) => sum + value * value, 0));
4+
if (magnitude < 1e-10) {
5+
return sanitized;
6+
}
7+
return sanitized.map((value) => value / magnitude);
8+
}

‎src/memory/embeddings-gemini.test.ts‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,10 @@ function parseFetchBody(fetchMock: { mock: { calls: unknown[][] } }, callIndex =
4444
return JSON.parse((init?.body as string) ?? "{}") as Record<string, unknown>;
4545
}
4646

47+
function magnitude(values: number[]) {
48+
return Math.sqrt(values.reduce((sum, value) => sum + value * value, 0));
49+
}
50+
4751
afterEach(() => {
4852
vi.resetAllMocks();
4953
vi.unstubAllGlobals();
@@ -224,6 +228,25 @@ describe("gemini-embedding-2-preview provider", () => {
224228
expect(body.content).toEqual({ parts: [{ text: "test query" }] });
225229
});
226230

231+
it("normalizes embedQuery response vectors", async () => {
232+
const fetchMock = createGeminiFetchMock([3, 4]);
233+
vi.stubGlobal("fetch", fetchMock);
234+
mockResolvedProviderKey();
235+
236+
const { provider } = await createGeminiEmbeddingProvider({
237+
config: {} as never,
238+
provider: "gemini",
239+
model: "gemini-embedding-2-preview",
240+
fallback: "none",
241+
});
242+
243+
const embedding = await provider.embedQuery("test query");
244+
245+
expect(embedding[0]).toBeCloseTo(0.6, 5);
246+
expect(embedding[1]).toBeCloseTo(0.8, 5);
247+
expect(magnitude(embedding)).toBeCloseTo(1, 5);
248+
});
249+
227250
it("includes outputDimensionality in embedBatch request", async () => {
228251
const fetchMock = createGeminiBatchFetchMock(2);
229252
vi.stubGlobal("fetch", fetchMock);
@@ -255,6 +278,28 @@ describe("gemini-embedding-2-preview provider", () => {
255278
]);
256279
});
257280

281+
it("normalizes embedBatch response vectors", async () => {
282+
const fetchMock = createGeminiBatchFetchMock(2, [3, 4]);
283+
vi.stubGlobal("fetch", fetchMock);
284+
mockResolvedProviderKey();
285+
286+
const { provider } = await createGeminiEmbeddingProvider({
287+
config: {} as never,
288+
provider: "gemini",
289+
model: "gemini-embedding-2-preview",
290+
fallback: "none",
291+
});
292+
293+
const embeddings = await provider.embedBatch(["text1", "text2"]);
294+
295+
expect(embeddings).toHaveLength(2);
296+
for (const embedding of embeddings) {
297+
expect(embedding[0]).toBeCloseTo(0.6, 5);
298+
expect(embedding[1]).toBeCloseTo(0.8, 5);
299+
expect(magnitude(embedding)).toBeCloseTo(1, 5);
300+
}
301+
});
302+
258303
it("respects custom outputDimensionality", async () => {
259304
const fetchMock = createGeminiFetchMock();
260305
vi.stubGlobal("fetch", fetchMock);
@@ -310,6 +355,28 @@ describe("gemini-embedding-2-preview provider", () => {
310355
).rejects.toThrow(/Invalid outputDimensionality 512/);
311356
});
312357

358+
it("sanitizes non-finite values before normalization", async () => {
359+
const fetchMock = createGeminiFetchMock([
360+
1,
361+
Number.NaN,
362+
Number.POSITIVE_INFINITY,
363+
Number.NEGATIVE_INFINITY,
364+
]);
365+
vi.stubGlobal("fetch", fetchMock);
366+
mockResolvedProviderKey();
367+
368+
const { provider } = await createGeminiEmbeddingProvider({
369+
config: {} as never,
370+
provider: "gemini",
371+
model: "gemini-embedding-2-preview",
372+
fallback: "none",
373+
});
374+
375+
const embedding = await provider.embedQuery("test");
376+
377+
expect(embedding).toEqual([1, 0, 0, 0]);
378+
});
379+
313380
it("uses correct endpoint URL", async () => {
314381
const fetchMock = createGeminiFetchMock();
315382
vi.stubGlobal("fetch", fetchMock);

‎src/memory/embeddings-gemini.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import {
55
import { requireApiKey, resolveApiKeyForProvider } from "../agents/model-auth.js";
66
import { parseGeminiAuth } from "../infra/gemini-auth.js";
77
import type { SsrFPolicy } from "../infra/net/ssrf.js";
8+
import { sanitizeAndNormalizeEmbedding } from "./embedding-vectors.js";
89
import { debugEmbeddingsLog } from "./embeddings-debug.js";
910
import type { EmbeddingProvider, EmbeddingProviderOptions } from "./embeddings.js";
1011
import { buildRemoteBaseUrlPolicy, withRemoteHttpResponse } from "./remote-http.js";
@@ -222,7 +223,7 @@ export async function createGeminiEmbeddingProvider(
222223
apiKeys: client.apiKeys,
223224
execute: (apiKey) => fetchWithGeminiAuth(apiKey, embedUrl, body),
224225
});
225-
return payload.embedding?.values ?? [];
226+
return sanitizeAndNormalizeEmbedding(payload.embedding?.values ?? []);
226227
};
227228

228229
const embedBatch = async (texts: string[]): Promise<number[][]> => {
@@ -244,7 +245,7 @@ export async function createGeminiEmbeddingProvider(
244245
execute: (apiKey) => fetchWithGeminiAuth(apiKey, batchUrl, batchBody),
245246
});
246247
const embeddings = Array.isArray(payload.embeddings) ? payload.embeddings : [];
247-
return texts.map((_, index) => embeddings[index]?.values ?? []);
248+
return texts.map((_, index) => sanitizeAndNormalizeEmbedding(embeddings[index]?.values ?? []));
248249
};
249250

250251
return {

‎src/memory/embeddings-ollama.ts‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { resolveOllamaApiBase } from "../agents/ollama-models.js";
33
import { formatErrorMessage } from "../infra/errors.js";
44
import type { SsrFPolicy } from "../infra/net/ssrf.js";
55
import { normalizeOptionalSecretInput } from "../utils/normalize-secret-input.js";
6+
import { sanitizeAndNormalizeEmbedding } from "./embedding-vectors.js";
67
import { normalizeEmbeddingModelWithPrefixes } from "./embeddings-model-normalize.js";
78
import type { EmbeddingProvider, EmbeddingProviderOptions } from "./embeddings.js";
89
import { buildRemoteBaseUrlPolicy, withRemoteHttpResponse } from "./remote-http.js";
@@ -19,15 +20,6 @@ type OllamaEmbeddingClientConfig = Omit<OllamaEmbeddingClient, "embedBatch">;
1920

2021
export const DEFAULT_OLLAMA_EMBEDDING_MODEL = "nomic-embed-text";
2122

22-
function sanitizeAndNormalizeEmbedding(vec: number[]): number[] {
23-
const sanitized = vec.map((value) => (Number.isFinite(value) ? value : 0));
24-
const magnitude = Math.sqrt(sanitized.reduce((sum, value) => sum + value * value, 0));
25-
if (magnitude < 1e-10) {
26-
return sanitized;
27-
}
28-
return sanitized.map((value) => value / magnitude);
29-
}
30-
3123
function normalizeOllamaModel(model: string): string {
3224
return normalizeEmbeddingModelWithPrefixes({
3325
model,

‎src/memory/embeddings.ts‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type { OpenClawConfig } from "../config/config.js";
44
import type { SecretInput } from "../config/types.secrets.js";
55
import { formatErrorMessage } from "../infra/errors.js";
66
import { resolveUserPath } from "../utils.js";
7+
import { sanitizeAndNormalizeEmbedding } from "./embedding-vectors.js";
78
import {
89
createGeminiEmbeddingProvider,
910
type GeminiEmbeddingClient,
@@ -18,15 +19,6 @@ import { createOpenAiEmbeddingProvider, type OpenAiEmbeddingClient } from "./emb
1819
import { createVoyageEmbeddingProvider, type VoyageEmbeddingClient } from "./embeddings-voyage.js";
1920
import { importNodeLlamaCpp } from "./node-llama.js";
2021

21-
function sanitizeAndNormalizeEmbedding(vec: number[]): number[] {
22-
const sanitized = vec.map((value) => (Number.isFinite(value) ? value : 0));
23-
const magnitude = Math.sqrt(sanitized.reduce((sum, value) => sum + value * value, 0));
24-
if (magnitude < 1e-10) {
25-
return sanitized;
26-
}
27-
return sanitized.map((value) => value / magnitude);
28-
}
29-
3022
export type { GeminiEmbeddingClient } from "./embeddings-gemini.js";
3123
export type { MistralEmbeddingClient } from "./embeddings-mistral.js";
3224
export type { OpenAiEmbeddingClient } from "./embeddings-openai.js";

0 commit comments

Comments
 (0)