diff --git a/kits/firestore-vector-search/CHANGELOG.md b/kits/firestore-vector-search/CHANGELOG.md index 711eb60d36..4efea603df 100644 --- a/kits/firestore-vector-search/CHANGELOG.md +++ b/kits/firestore-vector-search/CHANGELOG.md @@ -1 +1,2 @@ +- OpenAI embeddings are back on the extension's model and size: `EMBEDDING_PROVIDER: openai` requests `text-embedding-ada-002` at its native 1536 dimensions with a batch size of 16, replacing `text-embedding-3-small` pinned at 512 with a batch size of 1. Vectors written by an earlier version of the kit are not comparable with the ones it writes now, so re-embed the collection after upgrading. The vector index the kit creates for OpenAI is still declared with 512 dimensions, exactly as the extension declared it, so it does not cover the 1536-dimension vectors and `findNearest` fails against it; create the 1536-dimension index yourself if you query an OpenAI-embedded collection. - Initial release of kit, see README for differences between the legacy extension and this kit diff --git a/kits/firestore-vector-search/README.md b/kits/firestore-vector-search/README.md index 0552ca77b5..b93deea357 100644 --- a/kits/firestore-vector-search/README.md +++ b/kits/firestore-vector-search/README.md @@ -156,7 +156,7 @@ This kit is version 0.1.3 of the extension repackaged as an npm package, and it the least literal of the ports. The seven functions, the Firestore vector index, the query document collection and the callable all survive with their names and settings intact, so a `.env` copied from your installed instance needs no value -changes. The embedding providers, the backfill, and the shape of the status field +changes. Multimodal embedding, the backfill, and the shape of the status field written onto your documents all changed, so read this before you point the kit at a collection an installed instance has already embedded. @@ -167,17 +167,19 @@ Selecting `multimodal` deploys, and then every embedding attempt throws multimodal image embedding, including reading images out of Cloud Storage, has no equivalent here. If you use it, stay on the extension. -### OpenAI embeddings are a different model and a different size +### Gemini and Vertex AI embeddings are truncated to 768 dimensions -`EMBEDDING_PROVIDER: openai` used `text-embedding-ada-002` and stored the full -1536-dimension vector, while the Firestore index it created was declared with 512 -dimensions. The kit uses `text-embedding-3-small` at 512 dimensions, which matches -the index. +Both the extension and the kit ask for `gemini-embedding-001` with +`outputDimensionality: 768`, set on the embedder reference. Genkit does not apply +it there, so the model returns its full 3072-dimension vector. Firestore refuses +any vector above 2048 dimensions, so on the extension every gemini and vertex +embed fails with `Vectors must be at most 2048 dimensions` and the document is +marked `ERROR`. -Vectors from the two models are not comparable, and the existing index is reused -as-is because the "does this index already exist" check only looks at the field -path, not the dimension. Re-embed the whole collection after you switch, and -delete the old vector index first if it was created with a different dimension. +The kit truncates each returned embedding to 768 before writing it, which is the +dimension both the extension and the kit declare their vector index with, so the +default provider works. This is the one place the kit deliberately does not match +the extension's behaviour, because matching it means writing nothing at all. ### You set `INSTANCE_ID` yourself, and it names the query collection @@ -314,7 +316,15 @@ for; the Firebase CLI grants these for you. - `DEFAULT_QUERY_LIMIT` (default 3) and `DISTANCE_MEASURE` (`COSINE`, `EUCLIDEAN`, `DOT_PRODUCT`, default `COSINE`) behave as before. - Gemini and Vertex AI embeddings are still `gemini-embedding-001` at 768 - dimensions. + dimensions, though the kit has to truncate the model's response to get there. + See *Gemini and Vertex AI embeddings are truncated to 768 dimensions* above. +- OpenAI embeddings are still `text-embedding-ada-002` at its native 1536 + dimensions, so vectors already written by an installed instance stay + comparable with the ones the kit writes. The vector index the kit creates for + `EMBEDDING_PROVIDER: openai` is still declared with 512 dimensions, exactly as + the extension declared it, so it does not cover those 1536-dimension vectors + and `findNearest` fails against it. Create the 1536-dimension index yourself if + you query an OpenAI-embedded collection. - A custom endpoint still receives `{ batch: [...] }` and must return `{ embeddings: [[...]] }`, and still requires all three of `CUSTOM_EMBEDDINGS_ENDPOINT`, `CUSTOM_EMBEDDINGS_BATCH_SIZE` and diff --git a/kits/firestore-vector-search/src/embeddings/client/text/open_ai.ts b/kits/firestore-vector-search/src/embeddings/client/text/open_ai.ts index 978c5df4ae..62a94b125f 100644 --- a/kits/firestore-vector-search/src/embeddings/client/text/open_ai.ts +++ b/kits/firestore-vector-search/src/embeddings/client/text/open_ai.ts @@ -22,7 +22,7 @@ export class OpenAiEmbedClient extends BaseEmbedClient { private readonly client: OpenAI; constructor(config: ResolvedVectorSearchConfig) { - super(1); + super(16); if (!config.openAiApiKey) { throw new Error("OpenAI embeddings require OPENAI_API_KEY"); } @@ -31,9 +31,8 @@ export class OpenAiEmbedClient extends BaseEmbedClient { async getEmbeddings(inputs: ReadonlyArray): Promise { const results = await this.client.embeddings.create({ - model: "text-embedding-3-small", + model: "text-embedding-ada-002", input: [...inputs], - dimensions: 512, }); return results.data.map((result) => result.embedding); } diff --git a/kits/firestore-vector-search/tests/embeddings.test.ts b/kits/firestore-vector-search/tests/embeddings.test.ts index ee80ea15a7..cea9f3bbd6 100644 --- a/kits/firestore-vector-search/tests/embeddings.test.ts +++ b/kits/firestore-vector-search/tests/embeddings.test.ts @@ -31,11 +31,26 @@ vi.mock("@genkit-ai/google-genai", () => ({ }), })); +const { embeddingsCreate, openAiConstructor } = vi.hoisted(() => ({ + embeddingsCreate: vi.fn(), + openAiConstructor: vi.fn(), +})); + +vi.mock("openai", () => ({ + default: class { + embeddings = { create: embeddingsCreate }; + constructor(options: unknown) { + openAiConstructor(options); + } + }, +})); + import { googleAI, vertexAI } from "@genkit-ai/google-genai"; import { genkit } from "genkit"; import { GenkitEmbedClient } from "../src/embeddings/client/genkit"; import { CustomEndpointClient } from "../src/embeddings/client/text/custom_function"; +import { OpenAiEmbedClient } from "../src/embeddings/client/text/open_ai"; import { type ResolvedVectorSearchConfig, resolveVectorSearchConfig, @@ -118,6 +133,8 @@ describe("GenkitEmbedClient", () => { ]); }); + // gemini-embedding-001 ignores the embedder's outputDimensionality and + // returns its full 3072-dimension vector, which Firestore rejects outright. test("truncates embeddings longer than the configured dimension", async () => { const client = new GenkitEmbedClient( config({ @@ -178,6 +195,85 @@ describe("GenkitEmbedClient", () => { }); }); +describe("OpenAiEmbedClient", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + function openAiConfig(): ResolvedVectorSearchConfig { + return config({ + embeddingProvider: "openai", + openAiApiKey: "test-openai-key", + }); + } + + describe("constructor", () => { + test("builds the client with the configured API key", () => { + new OpenAiEmbedClient(openAiConfig()); + + expect(openAiConstructor).toHaveBeenCalledWith({ + apiKey: "test-openai-key", + }); + }); + + test("throws when no API key is configured", () => { + expect( + () => new OpenAiEmbedClient(config({ embeddingProvider: "openai" })) + ).toThrow("OpenAI embeddings require OPENAI_API_KEY"); + }); + + test("embeds sixteen inputs per batch", () => { + expect(new OpenAiEmbedClient(openAiConfig()).batchSize).toBe(16); + }); + }); + + describe("getEmbeddings", () => { + test("requests text-embedding-ada-002 at its native dimension", async () => { + const client = new OpenAiEmbedClient(openAiConfig()); + embeddingsCreate.mockResolvedValueOnce({ + data: [{ embedding: [1, 2, 3] }, { embedding: [4, 5, 6] }], + }); + + const inputs = ["input1", "input2"]; + + await expect(client.getEmbeddings(inputs)).resolves.toEqual([ + [1, 2, 3], + [4, 5, 6], + ]); + expect(embeddingsCreate).toHaveBeenCalledWith({ + model: "text-embedding-ada-002", + input: inputs, + }); + }); + + test("throws when embedding fails", async () => { + const client = new OpenAiEmbedClient(openAiConfig()); + embeddingsCreate.mockRejectedValueOnce(new Error("Embedding failed")); + + await expect(client.getEmbeddings(["input"])).rejects.toThrow( + "Embedding failed" + ); + }); + }); + + describe("getSingleEmbedding", () => { + test("returns a single embedding for an input", async () => { + const client = new OpenAiEmbedClient(openAiConfig()); + embeddingsCreate.mockResolvedValueOnce({ + data: [{ embedding: [7, 8, 9] }], + }); + + await expect(client.getSingleEmbedding("input1")).resolves.toEqual([ + 7, 8, 9, + ]); + expect(embeddingsCreate).toHaveBeenCalledWith({ + model: "text-embedding-ada-002", + input: ["input1"], + }); + }); + }); +}); + describe("CustomEndpointClient", () => { const customConfig = () => config({