diff --git a/kits/firestore-vector-search/README.md b/kits/firestore-vector-search/README.md index 0552ca77b5..ec98ffb19e 100644 --- a/kits/firestore-vector-search/README.md +++ b/kits/firestore-vector-search/README.md @@ -251,23 +251,14 @@ longer written, only `COMPLETED` and `ERROR`. Anything reading `status..state`, or a security rule or index keyed to it, needs updating. The field name is still `STATUS_FIELD_NAME`, defaulting to `status`. +`embedOnWrite` still reads the same four states the extension treated as final — +`PROCESSING`, `COMPLETED`, `ERROR` and `BACKFILLED` — so documents an installed +instance already embedded are still skipped once you flatten their status field. + Query documents no longer get a status field at all. They previously carried `status.textQuery`, so if you were waiting on that to know a query had finished, wait for `result` instead. -### Editing a document's input re-embeds it - -The extension embedded each document once. Its skip rule was "this document's -status is already in a final state", so once a document reached `COMPLETED` (or -`ERROR`), changing its input field never produced a new embedding and a failure -was never retried. - -The kit compares the input instead: it re-embeds when the input field changes, and -skips only when the input is unchanged and an embedding is already present. This -is usually what you wanted, but it means editing inputs in bulk now costs -embedding calls, and a document that previously sat stale will be brought up to -date on its next write. - ### The lifecycle hooks and the function region Install and reconfigure hooks are replaced by an `initVectorSearch` task that the diff --git a/kits/firestore-vector-search/src/handlers.ts b/kits/firestore-vector-search/src/handlers.ts index d62061ac9e..69bb732aea 100644 --- a/kits/firestore-vector-search/src/handlers.ts +++ b/kits/firestore-vector-search/src/handlers.ts @@ -47,6 +47,27 @@ export type VectorWriteEvent = FirestoreEvent< Record >; +/** + * States the extension's `FirestoreOnWriteProcessor` treated as final. A + * document that has reached one of these is never processed again, so each + * document is embedded once and a failure is never retried. + */ +const TERMINAL_STATES = new Set([ + "PROCESSING", + "COMPLETED", + "ERROR", + "BACKFILLED", +]); + +function isInTerminalState( + data: FirebaseFirestore.DocumentData, + statusFieldName: string +): boolean { + const status = data[statusFieldName] as { state?: unknown } | undefined; + const state = status?.state; + return typeof state === "string" && TERMINAL_STATES.has(state); +} + function queuePath( config: ResolvedVectorSearchConfig, queueName: string @@ -78,12 +99,9 @@ export async function handleEmbedOnWrite( logs.start("embedOnWrite"); const data = event.data.after.data() ?? {}; + if (isInTerminalState(data, ctx.config.statusFieldName)) return; const input = data[ctx.config.inputFieldName]; if (typeof input !== "string") return; - const beforeInput = event.data.before.exists - ? event.data.before.get(ctx.config.inputFieldName) - : undefined; - if (beforeInput === input && data[ctx.config.outputFieldName]) return; try { const embedding = await embedClient(ctx).getSingleEmbedding(input); diff --git a/kits/firestore-vector-search/tests/handlers.test.ts b/kits/firestore-vector-search/tests/handlers.test.ts index edcfc10245..e16426dcb0 100644 --- a/kits/firestore-vector-search/tests/handlers.test.ts +++ b/kits/firestore-vector-search/tests/handlers.test.ts @@ -14,6 +14,7 @@ * limitations under the License. */ +import { FieldValue } from "firebase-admin/firestore"; import type { CallableRequest } from "firebase-functions/v2/https"; import { HttpsError } from "firebase-functions/v2/https"; import { beforeEach, describe, expect, test, vi } from "vitest"; @@ -36,7 +37,12 @@ vi.mock("../src/embeddings", () => ({ // handler never needs it. vi.mock("../src/queries/setup", () => ({ createIndex: vi.fn() })); -import { type HandlerContext, handleQueryCall } from "../src/handlers"; +import { + type HandlerContext, + type VectorWriteEvent, + handleEmbedOnWrite, + handleQueryCall, +} from "../src/handlers"; import { resolveVectorSearchConfig } from "../src/export-config"; const config = resolveVectorSearchConfig({ @@ -205,3 +211,184 @@ describe("handleQueryCall", () => { expect((err as Error).message).toBe("Query failed"); }); }); + +/** A minimal `DocumentSnapshot` stand-in backed by a plain object. */ +function snapshot(data: Record | null) { + const set = vi.fn().mockResolvedValue(undefined); + return { + snap: { + exists: data !== null, + data: () => data ?? undefined, + get: (field: string) => data?.[field], + ref: { path: `${config.collectionPath}/doc-1`, set }, + }, + set, + }; +} + +function writeEvent( + before: Record | null, + after: Record | null +) { + const beforeSnap = snapshot(before); + const afterSnap = snapshot(after); + const event = { + params: { docId: "doc-1" }, + data: { before: beforeSnap.snap, after: afterSnap.snap }, + } as unknown as VectorWriteEvent; + return { event, set: afterSnap.set }; +} + +function embedCtx() { + return { firestore: {}, config } as unknown as HandlerContext; +} + +describe("handleEmbedOnWrite", () => { + beforeEach(() => { + vi.clearAllMocks(); + getSingleEmbedding.mockResolvedValue(EMBEDDING); + }); + + test("embeds a new document and marks it COMPLETED", async () => { + const { event, set } = writeEvent(null, { input: "hello" }); + + await handleEmbedOnWrite(event, embedCtx()); + + expect(getSingleEmbedding).toHaveBeenCalledWith("hello"); + expect(set).toHaveBeenCalledWith( + { + [config.outputFieldName]: FieldValue.vector(EMBEDDING), + [config.statusFieldName]: { state: "COMPLETED" }, + }, + { merge: true } + ); + }); + + test("embeds a document that already has an embedding but no status", async () => { + const { event } = writeEvent(null, { + input: "hello", + [config.outputFieldName]: FieldValue.vector(EMBEDDING), + }); + + await handleEmbedOnWrite(event, embedCtx()); + + expect(getSingleEmbedding).toHaveBeenCalledWith("hello"); + }); + + test("marks the document ERROR and rethrows when embedding fails", async () => { + const { event, set } = writeEvent(null, { input: "hello" }); + getSingleEmbedding.mockRejectedValue(new Error("Embedding failed")); + + await expect(handleEmbedOnWrite(event, embedCtx())).rejects.toThrow( + "Embedding failed" + ); + expect(set).toHaveBeenCalledWith( + { + [config.statusFieldName]: { + state: "ERROR", + message: "Embedding failed", + }, + }, + { merge: true } + ); + }); + + test("skips a deleted document", async () => { + const { event, set } = writeEvent({ input: "hello" }, null); + + await handleEmbedOnWrite(event, embedCtx()); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(set).not.toHaveBeenCalled(); + }); + + test("skips a document whose input is not a string", async () => { + const { event, set } = writeEvent(null, { input: 42 }); + + await handleEmbedOnWrite(event, embedCtx()); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(set).not.toHaveBeenCalled(); + }); + + // Parity with the extension: `FirestoreOnWriteProcessor` skipped any document + // already in a final state, so an edited input never produced a new embedding + // and a failure was never retried. + for (const state of ["PROCESSING", "COMPLETED", "ERROR", "BACKFILLED"]) { + test(`does not re-embed a document in the ${state} state`, async () => { + const { event, set } = writeEvent( + { input: "hello", [config.statusFieldName]: { state } }, + { input: "goodbye", [config.statusFieldName]: { state } } + ); + + await handleEmbedOnWrite(event, embedCtx()); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(set).not.toHaveBeenCalled(); + }); + } + + test("embeds a document in an unrecognised state", async () => { + const { event } = writeEvent(null, { + input: "hello", + [config.statusFieldName]: { state: "SOMETHING_ELSE" }, + }); + + await handleEmbedOnWrite(event, embedCtx()); + + expect(getSingleEmbedding).toHaveBeenCalledWith("hello"); + }); + + // The extension's skip rule was the status state alone, with no comparison + // against the previous input, so an unchanged document with an embedding but + // no status was still processed. + test("embeds an unchanged document that has an embedding but no status", async () => { + const doc = { + input: "hello", + [config.outputFieldName]: FieldValue.vector(EMBEDDING), + }; + const { event } = writeEvent({ ...doc }, { ...doc }); + + await handleEmbedOnWrite(event, embedCtx()); + + expect(getSingleEmbedding).toHaveBeenCalledWith("hello"); + }); + + describe("with a custom status field name", () => { + const customConfig = resolveVectorSearchConfig({ + projectId: "test-project", + instanceId: "test-instance", + statusFieldName: "embedStatus", + }); + + function customCtx() { + return { + firestore: {}, + config: customConfig, + } as unknown as HandlerContext; + } + + test("skips on the configured field", async () => { + const { event, set } = writeEvent(null, { + input: "hello", + embedStatus: { state: "COMPLETED" }, + }); + + await handleEmbedOnWrite(event, customCtx()); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(set).not.toHaveBeenCalled(); + }); + + test("ignores a terminal state on the default field", async () => { + const { event } = writeEvent(null, { + input: "hello", + status: { state: "COMPLETED" }, + }); + + await handleEmbedOnWrite(event, customCtx()); + + expect(getSingleEmbedding).toHaveBeenCalledWith("hello"); + }); + }); +});