Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions kits/firestore-vector-search/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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
32 changes: 21 additions & 11 deletions kits/firestore-vector-search/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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

Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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");
}
Expand All @@ -31,9 +31,8 @@ export class OpenAiEmbedClient extends BaseEmbedClient {

async getEmbeddings(inputs: ReadonlyArray<string>): Promise<number[][]> {
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);
}
Comment thread
CorieW marked this conversation as resolved.
Expand Down
96 changes: 96 additions & 0 deletions kits/firestore-vector-search/tests/embeddings.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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({
Expand Down Expand Up @@ -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({
Expand Down
Loading