diff --git a/kits/firestore-vector-search/CHANGELOG.md b/kits/firestore-vector-search/CHANGELOG.md index 711eb60d36..f1793ff488 100644 --- a/kits/firestore-vector-search/CHANGELOG.md +++ b/kits/firestore-vector-search/CHANGELOG.md @@ -1 +1,5 @@ - Initial release of kit, see README for differences between the legacy extension and this kit +- Restored the extension's batched backfill: the backfill and update triggers enumerate the collection by document reference instead of loading it into memory, chunk it into 50 document ids per Cloud Task, run one task at a time, and embed each chunk in provider-sized batches with a single API call per batch. Documents are marked `BACKFILLED` or `FAILED_BACKFILL` in one batched write, and a failed batch no longer fails the task. +- Restored the index metadata gate at `_/index`: a backfill or update pass runs only when the embedding provider, the vector dimension or the input/output field names differ from the last recorded pass, so a redeploy that changes none of them no longer re-embeds the whole collection. Unlike the extension, the progress counters are merged into that document rather than replacing it, so the gate survives its own first pass. +- Raised the OpenAI embedding client's batch size from 1 to the extension's 16, so a backfill chunk is embedded in one request per 16 documents. +- With both `DO_BACKFILL` and `UPDATE_ON_CONFIGURE` set, only the backfill pass is enqueued. The two passes share one task thread on `_/index`, and the backfill pass covers every document the update pass would. diff --git a/kits/firestore-vector-search/README.md b/kits/firestore-vector-search/README.md index 0552ca77b5..500394c61d 100644 --- a/kits/firestore-vector-search/README.md +++ b/kits/firestore-vector-search/README.md @@ -156,9 +156,9 @@ 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 -written onto your documents all changed, so read this before you point the kit at -a collection an installed instance has already embedded. +changes. The embedding providers and the shape of the status field written onto +your documents both changed, so read this before you point the kit at a +collection an installed instance has already embedded. ### `EMBEDDING_PROVIDER: multimodal` is not implemented @@ -182,8 +182,9 @@ delete the old vector index first if it was created with a different dimension. ### You set `INSTANCE_ID` yourself, and it names the query collection The extension derived its instance id at install and used it for the query -collection (`_/index/queries`), the index metadata document -(`_/index`) and its task queues. Here `INSTANCE_ID` is a setting you +collection (`_/index/queries`), the index metadata and backfill +progress document (`_/index`, with its `enqueues` subcollection) +and its task queues. Here `INSTANCE_ID` is a setting you provide, and it must match this instance's key in the `instances` map in `firebase.json`. To keep serving the query documents your clients already write to, set it to your installed instance's id. The four task queue names can also be @@ -201,35 +202,40 @@ whatever `EMBEDDING_PROVIDER` is set to. If either does not exist, `firebase deploy` prompts you for a value and fails outright when running non-interactively (CI). Create the one you do not need with a placeholder value. -### `UPDATE_ON_CONFIGURE` now re-embeds on every deploy - -This setting was declared by the extension but never read. Reconfiguring an -installed instance re-embedded documents only when the provider, the vector -dimension or the input/output field names had actually changed, which the -extension tracked in its index metadata document. - -The kit keeps no such metadata and does no comparison. `UPDATE_ON_CONFIGURE: true` -enqueues a full re-embed of every document that already has an embedding after -*every* `firebase deploy`, whether anything relevant changed or not, and -`DO_BACKFILL: true` embeds the whole collection after the first deploy. On a large -collection that is a large Vertex AI or OpenAI bill per deploy. Set -`UPDATE_ON_CONFIGURE: false` and re-embed deliberately when you change providers. - -### Backfill is one task per document, and reads the collection in one go - -The extension chunked the collection into batches sized to the provider (16 -documents per OpenAI call), embedded each batch in a single API call, and tracked -progress in its metadata document. The kit reads the entire collection with one -`get()` and enqueues one Cloud Task per document, each of which embeds one -document with one API call. - -Two consequences. A collection large enough that a single `get()` does not fit in -the trigger's 512 MiB will fail the backfill outright, and there is no -resume-from-progress. Backfilling *n* documents now costs *n* task invocations and -*n* embedding calls rather than *n*/batch size. - -There is also no install-time progress reporting, since there is no extension -install UI to report into. Watch the function logs instead. +### `UPDATE_ON_CONFIGURE` is read, and the backfill gate is stricter than the extension's + +This setting was declared by the extension but never read: its update pass was +gated on `DO_BACKFILL` instead. The kit reads `UPDATE_ON_CONFIGURE`, so the two +passes are controlled independently — `DO_BACKFILL` after the first deploy, +`UPDATE_ON_CONFIGURE` after every redeploy. + +Both passes are then gated on the index metadata document at +`_/index`, as the extension's were: a pass runs only when the +embedding provider, the vector dimension or the input/output field names differ +from what the last pass recorded there. Redeploying without changing any of them +enqueues nothing and costs nothing. + +The two passes share that document as their task thread, so only one of them +runs per deploy: with both settings on, the backfill pass runs, which covers +every document the update pass would have (the update pass is the same +eligibility rule plus "and already has an embedding"). + +The extension's gate did not survive its own first pass, because the progress +counters it wrote to the same document replaced the recorded configuration. The +kit merges instead, so the comparison fields persist and the gate holds on every +later deploy. To force a full re-embed without changing any setting, delete the +`_/index` document; its `queries` subcollection is untouched, so +the query documents your clients write to survive. + +### There is no install-time progress reporting + +The extension reported backfill progress and failures through the extension +install UI (`setProcessingState`). There is no such surface for a kit, so +progress is visible in the function logs and in the progress fields on +`_/index` (`backfillJobsTotal`, `backfillJobsProcessed`, +`backfillJobsSkipped`, `backfillJobsFailed`, `backfillStatus`) instead. One +document per chunk is written under `_/index/enqueues`, as the +extension did, each carrying its chunk of document ids and its own status. ### The `status` field on your documents is a different shape @@ -246,10 +252,12 @@ status: { state: "COMPLETED" } status: { state: "ERROR", message: "" } ``` -The states themselves are narrower too: `PROCESSING` and `BACKFILLED` are no -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`. +The states themselves are narrower too: `PROCESSING` is no longer written. The +write triggers write `COMPLETED` or `ERROR`, and the backfill and update passes +write `BACKFILLED` or `FAILED_BACKFILL` alongside a `completeTime`, as the +extension did. 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`. 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, @@ -302,6 +310,13 @@ for; the Firebase CLI grants these for you. ### Unchanged +- The backfill and update passes still enumerate the collection by document + reference, chunk it into 50 document ids per Cloud Task, run one task at a + time, and embed each chunk in provider-sized batches with a single API call per + batch (16 documents per OpenAI call). A document whose input is not a string is + skipped, as is one whose status is already set to anything other than + `BACKFILLED`. A failed batch marks its documents `FAILED_BACKFILL` and the task + still succeeds. - The indexed collection is still `COLLECTION_NAME` (default `products`), the input, output and status fields still default to `input`, `embedding` and `status`, and embeddings are still written as native Firestore vectors. diff --git a/kits/firestore-vector-search/src/backfill.ts b/kits/firestore-vector-search/src/backfill.ts new file mode 100644 index 0000000000..4ce1f76cca --- /dev/null +++ b/kits/firestore-vector-search/src/backfill.ts @@ -0,0 +1,521 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { DocumentSnapshot, Firestore } from "firebase-admin/firestore"; +import { FieldValue, Timestamp } from "firebase-admin/firestore"; +import type { TaskQueue } from "firebase-admin/functions"; +import { logger } from "firebase-functions"; + +/** Document ids carried by a single backfill task. */ +export const TASK_CHUNK_SIZE = 50; + +/** Embedding batch size used when a process declares none. */ +export const DEFAULT_BATCH_SIZE = 50; + +export type BackfillJobStatus = + | "PENDING" + | "RUNNING" + | "PROCESSING" + | "DONE" + | "FAILED"; + +/** The state written onto a document once its backfill pass has finished. */ +export const BACKFILLED_STATE = "BACKFILLED"; +export const FAILED_BACKFILL_STATE = "FAILED_BACKFILL"; + +export type BackfillDocumentData = Record; + +/** Payload of a `backfillTask` / `updateTask` dispatch. */ +export interface BackfillTaskData { + taskId: string; + chunk: string[]; + tasksDoc: string; +} + +/** + * The fields compared against the index metadata document to decide whether a + * backfill pass is required. + */ +export interface BackfillMetadata { + collectionName: string; + instanceId: string; + embeddingProvider: string; + dimension: number; + inputField: string; + outputField: string; +} + +/** + * One backfill pass: how to decide a document is eligible, and how to produce + * the fields written back onto it. + */ +export interface BackfillProcess { + id: string; + batchSize: number; + shouldBackfill(data: BackfillDocumentData): boolean; + processFn(data: BackfillDocumentData): Promise; + /** + * Embeds a whole batch in one call. Processes without one fall back to + * `processFn` per document, and a single document failing then fails only + * that document. + */ + batchFn?(data: BackfillDocumentData[]): Promise; +} + +export interface BackfillOptions { + firestore: Firestore; + collectionName: string; + statusField: string; +} + +export interface ChunkResult { + success: number; + failed: number; + skipped: number; +} + +export function chunkArray(array: readonly T[], chunkSize: number): T[][] { + const result: T[][] = []; + for (let i = 0; i < array.length; i += chunkSize) { + result.push(array.slice(i, i + chunkSize)); + } + return result; +} + +function isRecord(value: unknown): value is BackfillDocumentData { + return typeof value === "object" && value !== null; +} + +function metadataChanged( + current: BackfillDocumentData, + metadata: BackfillMetadata +): boolean { + return ( + current.embeddingProvider !== metadata.embeddingProvider || + current.dimension !== metadata.dimension || + current.inputField !== metadata.inputField || + current.outputField !== metadata.outputField + ); +} + +/** + * Reads the index metadata document, decides whether the embedding + * configuration has changed since the last pass, and records the current + * configuration when it has. + * + * The metadata document doubles as the task-thread progress document, so every + * write merges rather than replaces. The extension replaced it, which wiped the + * comparison fields on the first pass and made the gate a no-op from then on. + */ +export async function updateOrCreateMetadataDoc( + firestore: Firestore, + metadataDocumentPath: string, + metadata: BackfillMetadata +): Promise<{ path: string; shouldBackfill: boolean }> { + logger.info( + `Fetching existing metadata doc for ${metadata.collectionName} 📝` + ); + const ref = firestore.doc(metadataDocumentPath); + const snapshot = await ref.get(); + const record = { ...metadata, createdAt: Timestamp.now() }; + + if (!snapshot.exists) { + logger.info( + `No existing metadata doc found for ${metadata.collectionName} 📝` + ); + logger.info("Creating a new metadata doc"); + await ref.set(record, { merge: true }); + return { path: ref.path, shouldBackfill: true }; + } + + const shouldBackfill = metadataChanged(snapshot.data() ?? {}, metadata); + if (shouldBackfill) { + logger.info("Updating existing metadata doc"); + await ref.set(record, { merge: true }); + } + return { path: ref.path, shouldBackfill }; +} + +/** + * Writes the progress document, records one enqueue document per chunk of + * document ids, and dispatches the first task. Each task enqueues its successor + * once it completes, so only one task is in flight at a time. + */ +export async function enqueueTaskThread(params: { + firestore: Firestore; + tasksDoc: string; + queue: TaskQueue; + taskParams: readonly string[]; + instanceId: string; +}): Promise { + const { firestore, tasksDoc, queue, taskParams, instanceId } = params; + + await firestore.doc(tasksDoc).set( + { + backfillJobsTotal: taskParams.length, + backfillJobsProcessed: 0, + backfillJobsSkipped: 0, + backfillJobsFailed: 0, + backfillStatus: "PENDING" satisfies BackfillJobStatus, + }, + { merge: true } + ); + + const chunks = chunkArray(taskParams, TASK_CHUNK_SIZE); + if (chunks.length === 0) { + return; + } + + // Record every chunk before dispatching anything. The first task enqueues its + // successor as soon as it finishes, so a successor that has not been written + // yet costs the thread a retry. + let writer = firestore.batch(); + let pendingWrites = 0; + + for (const [index, chunk] of chunks.entries()) { + const taskId = taskIdFor(instanceId, index + 1); + writer.set(firestore.doc(`${tasksDoc}/enqueues/${taskId}`), { + taskId, + status: "PENDING" satisfies BackfillJobStatus, + chunk, + }); + pendingWrites++; + + if (pendingWrites === TASK_CHUNK_SIZE) { + logger.info("Committing the batch..."); + await writer.commit(); + writer = firestore.batch(); + pendingWrites = 0; + } + } + + if (pendingWrites > 0) { + logger.info("Committing the batch..."); + await writer.commit(); + } + + const firstTaskId = taskIdFor(instanceId, 1); + logger.info(`Enqueuing the first task ${firstTaskId} 🚀`); + await queue.enqueue({ + taskId: firstTaskId, + chunk: chunks[0], + tasksDoc, + }); + await firestore.doc(tasksDoc).update({ + backfillStatus: "RUNNING" satisfies BackfillJobStatus, + }); + + logger.info(`${chunks.length} tasks enqueued successfully 🚀`); +} + +function taskIdFor(instanceId: string, counter: number): string { + return `kit-${instanceId}-task-${counter}`; +} + +export function getNextTaskId(prevId: string, instanceId: string): string { + // Captured rather than split on "task-", because an instance id may contain + // that substring itself. + const match = prevId.match(new RegExp(`^kit-${instanceId}-task-([0-9]+)$`)); + if (!match) { + throw new Error(`Invalid task ID format: ${prevId}`); + } + return taskIdFor(instanceId, Number.parseInt(match[1], 10) + 1); +} + +/** + * Runs one dispatched chunk: marks the enqueue document, embeds the chunk, + * updates the progress counters, and either finishes the thread or dispatches + * the next task. + */ +export async function runBackfillTask(params: { + data: BackfillTaskData; + process: BackfillProcess; + options: BackfillOptions; + queue: TaskQueue; + instanceId: string; +}): Promise { + const { data, process, options, queue, instanceId } = params; + const { firestore } = options; + const { taskId, chunk, tasksDoc } = data; + + if (!chunk || chunk.length === 0) { + logger.info("No data to handle, skipping..."); + return; + } + logger.info(`Handling ${chunk.length} documents`); + + const taskRef = firestore.doc(`${tasksDoc}/enqueues/${taskId}`); + await taskRef.update({ + status: "PROCESSING" satisfies BackfillJobStatus, + }); + + const { success, failed, skipped } = await runChunk(process, chunk, options); + + await taskRef.update({ status: "DONE" satisfies BackfillJobStatus }); + logger.info(`Task ${taskId} completed with ${success} success(es)`); + + const tasksDocSnapshot = await firestore.doc(tasksDoc).get(); + const progress = tasksDocSnapshot.data() ?? {}; + const totalTasks = progress.backfillJobsTotal; + const processedTasks = progress.backfillJobsProcessed; + const skippedTasks = progress.backfillJobsSkipped; + const failedTasks = progress.backfillJobsFailed; + + if ( + [totalTasks, processedTasks, skippedTasks, failedTasks].some( + (value) => typeof value !== "number" + ) + ) { + throw new Error("Invalid task document"); + } + + await firestore.doc(tasksDoc).update({ + backfillJobsFailed: FieldValue.increment(failed), + backfillJobsSkipped: FieldValue.increment(skipped), + backfillJobsProcessed: FieldValue.increment(success), + }); + + const processed = (processedTasks as number) + success; + const totalSkipped = (skippedTasks as number) + skipped; + const totalFailed = (failedTasks as number) + failed; + + logger.info( + `Current state: ${processed} processed, ${totalSkipped} skipped, ${totalFailed} failed out of ${totalTasks} total tasks` + ); + + if (processed + totalSkipped + totalFailed === totalTasks) { + await firestore.doc(tasksDoc).update({ + backfillStatus: "DONE" satisfies BackfillJobStatus, + }); + return; + } + + await enqueueNextTask({ + firestore, + prevId: taskId, + tasksDoc, + queue, + instanceId, + }); +} + +async function enqueueNextTask(params: { + firestore: Firestore; + prevId: string; + tasksDoc: string; + queue: TaskQueue; + instanceId: string; +}): Promise { + const { firestore, prevId, tasksDoc, queue, instanceId } = params; + const nextId = getNextTaskId(prevId, instanceId); + + const nextTask = await firestore.doc(`${tasksDoc}/enqueues/${nextId}`).get(); + if (!nextTask.exists) { + logger.error(`Next task document ${nextId} not found.`); + throw new Error(`Next task document ${nextId} does not exist.`); + } + + const chunk = nextTask.data()?.chunk; + if (!Array.isArray(chunk) || chunk.length === 0) { + logger.error(`Next task ${nextId} has an invalid or empty chunk.`); + throw new Error(`Next task ${nextId} does not have valid chunk data.`); + } + + await queue.enqueue({ taskId: nextId, chunk, tasksDoc }); + logger.info(`Successfully enqueued task ${nextId}`); +} + +async function runChunk( + process: BackfillProcess, + chunk: readonly string[], + options: BackfillOptions +): Promise { + const { validDocuments, skippedDocuments } = await getValidDocs( + process, + chunk, + options + ); + + if (validDocuments.length === 0) { + logger.info("No data to handle, skipping..."); + return { success: 0, failed: 0, skipped: skippedDocuments.length }; + } + + logger.info(`Handling ${validDocuments.length} documents`); + + if (validDocuments.length === 1) { + return handleSingleDocument( + process, + validDocuments[0], + skippedDocuments.length, + options + ); + } + + const batches = chunkArray( + validDocuments, + process.batchSize || DEFAULT_BATCH_SIZE + ); + const results = await Promise.allSettled( + batches.map((batch) => batchProcess(process, batch)) + ); + + const writer = options.firestore.batch(); + let failedDocumentsCount = 0; + + results.forEach((result, index) => { + const batch = batches[index]; + + if (result.status === "rejected") { + // A failed batch means all its documents are considered failed. + failedDocumentsCount += batch.length; + logger.error(`Batch ${index + 1} failed`, result.reason); + for (const doc of batch) { + writer.update(doc.ref, failedPayload(options)); + } + return; + } + + batch.forEach((doc, i) => { + const fields = result.value[i]; + if (!fields) { + failedDocumentsCount++; + writer.update(doc.ref, failedPayload(options)); + return; + } + writer.update(doc.ref, { + ...fields, + ...backfilledPayload(options), + }); + }); + }); + + await writer.commit(); + + return { + success: validDocuments.length - failedDocumentsCount, + failed: failedDocumentsCount, + skipped: skippedDocuments.length, + }; +} + +/** + * Embeds one batch. A process with a `batchFn` embeds the whole batch in a + * single call, so the batch succeeds or fails as a unit; otherwise each + * document is embedded on its own and failures are reported per document by + * leaving that slot empty. + */ +async function batchProcess( + process: BackfillProcess, + batch: readonly DocumentSnapshot[] +): Promise<(BackfillDocumentData | undefined)[]> { + const data = batch.map((doc) => doc.data() as BackfillDocumentData); + + if (process.batchFn) { + return process.batchFn(data); + } + + const results = await Promise.allSettled(data.map(process.processFn)); + return results.map((result) => { + if (result.status === "fulfilled") { + return result.value; + } + logger.error(result.reason); + return undefined; + }); +} + +export async function getValidDocs( + process: BackfillProcess, + documentIds: readonly string[], + options: BackfillOptions +): Promise<{ + validDocuments: DocumentSnapshot[]; + skippedDocuments: DocumentSnapshot[]; +}> { + const collection = options.firestore.collection(options.collectionName); + + // Collected inside the transaction: Firestore may run the callback more than + // once, and arrays held outside it would accumulate duplicates on a retry. + return options.firestore.runTransaction(async (transaction) => { + const validDocuments: DocumentSnapshot[] = []; + const skippedDocuments: DocumentSnapshot[] = []; + const refs = documentIds.map((id) => collection.doc(id)); + const docs = await transaction.getAll(...refs); + + for (const doc of docs) { + const data = doc.data(); + + if (!data || !process.shouldBackfill(data)) { + skippedDocuments.push(doc); + logger.warn( + `Document ${doc.ref.path} is not valid for ${process.id} process` + ); + continue; + } + + const status = data[options.statusField]; + const state = isRecord(status) ? status.state : undefined; + if (state && state !== BACKFILLED_STATE) { + skippedDocuments.push(doc); + logger.warn( + `Document ${doc.ref.path} is not in the correct state to be backfilled` + ); + continue; + } + + validDocuments.push(doc); + } + + return { validDocuments, skippedDocuments }; + }); +} + +async function handleSingleDocument( + process: BackfillProcess, + document: DocumentSnapshot, + skipped: number, + options: BackfillOptions +): Promise { + try { + const result = await process.processFn( + document.data() as BackfillDocumentData + ); + await document.ref.update({ + ...result, + ...backfilledPayload(options), + }); + return { success: 1, failed: 0, skipped }; + } catch (err) { + logger.error(err); + await document.ref.update(failedPayload(options)); + return { success: 0, failed: 1, skipped }; + } +} + +function backfilledPayload(options: BackfillOptions): BackfillDocumentData { + return { + [`${options.statusField}.state`]: BACKFILLED_STATE, + [`${options.statusField}.completeTime`]: FieldValue.serverTimestamp(), + }; +} + +function failedPayload(options: BackfillOptions): BackfillDocumentData { + return { + [`${options.statusField}.state`]: FAILED_BACKFILL_STATE, + [`${options.statusField}.completeTime`]: FieldValue.serverTimestamp(), + }; +} 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..c52bd3311e 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,8 @@ export class OpenAiEmbedClient extends BaseEmbedClient { private readonly client: OpenAI; constructor(config: ResolvedVectorSearchConfig) { - super(1); + // The extension's OpenAI client batched 16 inputs per request. + super(16); if (!config.openAiApiKey) { throw new Error("OpenAI embeddings require OPENAI_API_KEY"); } diff --git a/kits/firestore-vector-search/src/handlers.ts b/kits/firestore-vector-search/src/handlers.ts index d62061ac9e..20d357fbdf 100644 --- a/kits/firestore-vector-search/src/handlers.ts +++ b/kits/firestore-vector-search/src/handlers.ts @@ -20,6 +20,17 @@ import type { Change, FirestoreEvent } from "firebase-functions/v2/firestore"; import type { CallableRequest } from "firebase-functions/v2/https"; import { HttpsError } from "firebase-functions/v2/https"; import type { Request } from "firebase-functions/v2/tasks"; +import { logger } from "firebase-functions"; +import { + type BackfillMetadata, + type BackfillProcess, + type BackfillTaskData, + DEFAULT_BATCH_SIZE, + type BackfillDocumentData, + enqueueTaskThread, + runBackfillTask, + updateOrCreateMetadataDoc, +} from "./backfill"; import { createEmbedClient } from "./embeddings"; import * as events from "./events"; import type { ResolvedVectorSearchConfig } from "./export-config"; @@ -38,9 +49,7 @@ export interface HandlerContext { config: ResolvedVectorSearchConfig; } -export interface VectorTaskData { - path: string; -} +export type VectorTaskData = BackfillTaskData; export type VectorWriteEvent = FirestoreEvent< Change | undefined, @@ -174,6 +183,10 @@ export async function handleInit(ctx: HandlerContext): Promise { if (ctx.config.doBackfill) { await enqueueBackfillTrigger(ctx); + // The two passes share one task thread on the index metadata document, and + // the backfill pass covers every document the update pass would. Running + // both at once would have them overwrite each other's progress. + return; } if (ctx.config.updateOnConfigure) { await enqueueUpdateTrigger(ctx); @@ -184,69 +197,172 @@ export async function handleBackfillTrigger( _request: Request, ctx: HandlerContext ): Promise { - const snapshot = await ctx.firestore - .collection(ctx.config.collectionPath) - .get(); - const queue = getFunctions().taskQueue( - queuePath(ctx.config, ctx.config.queueNames.backfillTask) - ); - await Promise.all( - snapshot.docs.map((doc) => queue.enqueue({ path: doc.ref.path })) - ); + await runTrigger(ctx, ctx.config.queueNames.backfillTask); } export async function handleUpdateTrigger( _request: Request, ctx: HandlerContext ): Promise { - const snapshot = await ctx.firestore - .collection(ctx.config.collectionPath) - .get(); - const queue = getFunctions().taskQueue( - queuePath(ctx.config, ctx.config.queueNames.updateTask) - ); - await Promise.all( - snapshot.docs.map((doc) => queue.enqueue({ path: doc.ref.path })) - ); + await runTrigger(ctx, ctx.config.queueNames.updateTask); } export async function handleBackfillTask( request: Request, ctx: HandlerContext ): Promise { - await embedPath(request.data.path, ctx, false); + await runBackfillTask({ + data: request.data, + process: embedProcess(ctx), + options: backfillOptions(ctx), + queue: taskQueue(ctx, ctx.config.queueNames.backfillTask), + instanceId: ctx.config.instanceId, + }); } export async function handleUpdateTask( request: Request, ctx: HandlerContext ): Promise { - await embedPath(request.data.path, ctx, true); + await runBackfillTask({ + data: request.data, + process: updateEmbedProcess(ctx), + options: backfillOptions(ctx), + queue: taskQueue(ctx, ctx.config.queueNames.updateTask), + instanceId: ctx.config.instanceId, + }); } -async function embedPath( - path: string, +/** + * Gates the pass on the index metadata document, enumerates the collection by + * reference, and hands the document ids to the task thread. + */ +async function runTrigger( ctx: HandlerContext, - requireExistingEmbedding: boolean + taskQueueName: string ): Promise { - const ref = ctx.firestore.doc(path); - const snapshot = await ref.get(); - if (!snapshot.exists) return; - const input = snapshot.get(ctx.config.inputFieldName); - if (typeof input !== "string") return; - if (requireExistingEmbedding && !snapshot.get(ctx.config.outputFieldName)) { + const { path, shouldBackfill } = await updateOrCreateMetadataDoc( + ctx.firestore, + ctx.config.indexMetadataDocumentPath, + metadataFor(ctx) + ); + + if (!shouldBackfill) { + logger.info( + `Embedding configuration is unchanged for ${ctx.config.collectionPath}, no pass required.` + ); return; } - const embedding = await embedClient(ctx).getSingleEmbedding(input); - await ref.set( - { - [ctx.config.outputFieldName]: FieldValue.vector(embedding), - [ctx.config.statusFieldName]: { state: "COMPLETED" }, - }, - { merge: true } + + try { + const refs = await ctx.firestore + .collection(ctx.config.collectionPath) + .listDocuments(); + + if (refs.length === 0) { + logger.info( + `No documents found in the collection ${ctx.config.collectionPath} 📚` + ); + return; + } + + logger.info( + `Found ${refs.length} documents in the collection ${ctx.config.collectionPath} 📚` + ); + logger.info("Enqueuing backfill tasks 🚀"); + + await enqueueTaskThread({ + firestore: ctx.firestore, + tasksDoc: path, + queue: taskQueue(ctx, taskQueueName), + taskParams: refs.map((ref) => ref.id), + instanceId: ctx.config.instanceId, + }); + } catch (err) { + logger.error("Error with backfill trigger"); + logger.error(err); + } +} + +function taskQueue(ctx: HandlerContext, queueName: string) { + return getFunctions().taskQueue( + queuePath(ctx.config, queueName) ); } +function backfillOptions(ctx: HandlerContext) { + return { + firestore: ctx.firestore, + collectionName: ctx.config.collectionPath, + statusField: ctx.config.statusFieldName, + }; +} + +function metadataFor(ctx: HandlerContext): BackfillMetadata { + return { + collectionName: ctx.config.collectionPath, + instanceId: ctx.config.instanceId, + embeddingProvider: ctx.config.embeddingProvider, + dimension: ctx.config.dimension, + inputField: ctx.config.inputFieldName, + outputField: ctx.config.outputFieldName, + }; +} + +function hasStringInput(data: BackfillDocumentData, field: string): boolean { + const value = data[field]; + return !!value && typeof value === "string"; +} + +/** The backfill pass: embeds a whole batch of documents in one call. */ +function embedProcess(ctx: HandlerContext): BackfillProcess { + const client = embedClient(ctx); + const { inputFieldName, outputFieldName } = ctx.config; + const embedOne = async ( + data: BackfillDocumentData + ): Promise => ({ + [outputFieldName]: FieldValue.vector( + await client.getSingleEmbedding(data[inputFieldName] as string) + ), + }); + + return { + id: ctx.config.instanceId, + batchSize: client.batchSize, + shouldBackfill: (data) => hasStringInput(data, inputFieldName), + processFn: embedOne, + batchFn: async (docs) => { + const embeddings = await client.getEmbeddings( + docs.map((doc) => doc[inputFieldName] as string) + ); + return embeddings.map((embedding) => ({ + [outputFieldName]: FieldValue.vector(embedding), + })); + }, + }; +} + +/** + * The update pass: only documents that already carry an embedding, and one + * embedding call per document, as the extension's update process did. + */ +function updateEmbedProcess(ctx: HandlerContext): BackfillProcess { + const client = embedClient(ctx); + const { inputFieldName, outputFieldName } = ctx.config; + + return { + id: ctx.config.instanceId, + batchSize: DEFAULT_BATCH_SIZE, + shouldBackfill: (data) => + hasStringInput(data, inputFieldName) && !!data[outputFieldName], + processFn: async (data) => ({ + [outputFieldName]: FieldValue.vector( + await client.getSingleEmbedding(data[inputFieldName] as string) + ), + }), + }; +} + async function enqueueBackfillTrigger(ctx: HandlerContext): Promise { await getFunctions() .taskQueue(queuePath(ctx.config, ctx.config.queueNames.backfillTrigger)) diff --git a/kits/firestore-vector-search/src/lib.ts b/kits/firestore-vector-search/src/lib.ts index a40457fc47..5fd57e095f 100644 --- a/kits/firestore-vector-search/src/lib.ts +++ b/kits/firestore-vector-search/src/lib.ts @@ -14,6 +14,13 @@ * limitations under the License. */ +export { + type BackfillDocumentData, + type BackfillMetadata, + type BackfillProcess, + type BackfillTaskData, + type ChunkResult, +} from "./backfill"; export { configFromEnv, geminiApiKey, openAiApiKey } from "./config"; export { createEmbedClient, type EmbedClient } from "./embeddings"; export { diff --git a/kits/firestore-vector-search/tests/backfill.test.ts b/kits/firestore-vector-search/tests/backfill.test.ts new file mode 100644 index 0000000000..5d510db808 --- /dev/null +++ b/kits/firestore-vector-search/tests/backfill.test.ts @@ -0,0 +1,914 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { Request } from "firebase-functions/v2/tasks"; +import { beforeEach, describe, expect, test, vi } from "vitest"; + +const { getSingleEmbedding, getEmbeddings, batchSize, enqueue } = vi.hoisted( + () => ({ + getSingleEmbedding: vi.fn(), + getEmbeddings: vi.fn(), + batchSize: { value: 2 }, + enqueue: vi.fn(), + }) +); + +vi.mock("../src/embeddings", () => ({ + createEmbedClient: vi.fn(() => ({ + get batchSize() { + return batchSize.value; + }, + getEmbeddings, + getSingleEmbedding, + })), +})); + +// `queries/setup` builds a FirestoreAdminClient at module scope; the backfill +// handlers never need it. +vi.mock("../src/queries/setup", () => ({ createIndex: vi.fn() })); + +vi.mock("firebase-admin/functions", () => ({ + getFunctions: () => ({ taskQueue: () => ({ enqueue }) }), +})); + +vi.mock("firebase-functions", () => ({ + logger: { + log: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + }, +})); + +import { + type BackfillProcess, + type BackfillTaskData, + chunkArray, + enqueueTaskThread, + getNextTaskId, + getValidDocs, + updateOrCreateMetadataDoc, +} from "../src/backfill"; +import { resolveVectorSearchConfig } from "../src/export-config"; +import { + type HandlerContext, + handleBackfillTask, + handleBackfillTrigger, + handleInit, + handleUpdateTask, + handleUpdateTrigger, +} from "../src/handlers"; + +const config = resolveVectorSearchConfig({ + projectId: "test-project", + instanceId: "test-instance", + region: "us-central1", +}); + +const METADATA_PATH = config.indexMetadataDocumentPath; +const COLLECTION = config.collectionPath; +const EMBEDDING = [0.1, 0.2, 0.3]; + +const METADATA = { + collectionName: COLLECTION, + instanceId: config.instanceId, + embeddingProvider: config.embeddingProvider, + dimension: config.dimension, + inputField: config.inputFieldName, + outputField: config.outputFieldName, +}; + +interface Write { + op: "set" | "update"; + path: string; + data: Record; + merge?: boolean; +} + +/** + * A minimal in-memory Firestore that records every write, so the tests can + * assert on the exact payloads the backfill writes back onto documents. + */ +function makeFirestore(seed: Record> = {}) { + const store = new Map(Object.entries(seed)); + const writes: Write[] = []; + const commits: Write[][] = []; + + const merge = (path: string, data: Record) => { + store.set(path, { ...(store.get(path) ?? {}), ...data }); + }; + + const ref = (path: string) => ({ + path, + id: path.split("/").pop() as string, + get: async () => snapshot(path), + set: async (data: Record, opts?: { merge?: boolean }) => { + writes.push({ op: "set", path, data, merge: opts?.merge }); + if (opts?.merge) merge(path, data); + else store.set(path, data); + }, + update: async (data: Record) => { + writes.push({ op: "update", path, data }); + merge(path, data); + }, + }); + + const snapshot = (path: string) => { + const data = store.get(path); + return { + exists: data !== undefined, + data: () => data, + get: (field: string) => data?.[field], + ref: ref(path), + }; + }; + + const firestore = { + doc: (path: string) => ref(path), + collection: (name: string) => ({ + doc: (id: string) => ref(`${name}/${id}`), + listDocuments: async () => + [...store.keys()] + .filter( + (key) => + key.startsWith(`${name}/`) && + key.slice(name.length + 1).includes("/") === false + ) + .map((key) => ref(key)), + get: async () => { + throw new Error("the backfill must not read the whole collection"); + }, + }), + runTransaction: async ( + fn: (tx: { + getAll: ( + ...refs: { path: string }[] + ) => Promise[]>; + }) => Promise + ) => + fn({ + getAll: async (...refs) => refs.map((r) => snapshot(r.path)), + }), + batch: () => { + const ops: Write[] = []; + return { + set: ( + target: { path: string }, + data: Record, + opts?: { merge?: boolean } + ) => { + ops.push({ op: "set", path: target.path, data, merge: opts?.merge }); + }, + update: (target: { path: string }, data: Record) => { + ops.push({ op: "update", path: target.path, data }); + }, + commit: async () => { + commits.push([...ops]); + for (const op of ops) { + writes.push(op); + merge(op.path, op.data); + } + ops.length = 0; + }, + }; + }, + }; + + return { firestore, store, writes, commits }; +} + +function makeCtx(seed: Record> = {}) { + const fake = makeFirestore(seed); + const ctx = { + firestore: fake.firestore, + config, + } as unknown as HandlerContext; + return { ...fake, ctx }; +} + +function taskRequest(data: BackfillTaskData) { + return { data } as unknown as Request; +} + +/** A progress document that still has one chunk outstanding. */ +function progress(overrides: Record = {}) { + return { + ...METADATA, + backfillJobsTotal: 4, + backfillJobsProcessed: 0, + backfillJobsSkipped: 0, + backfillJobsFailed: 0, + backfillStatus: "RUNNING", + ...overrides, + }; +} + +function docWrites(writes: Write[], id: string) { + return writes.filter((write) => write.path === `${COLLECTION}/${id}`); +} + +function stateOf(writes: Write[], id: string) { + const last = docWrites(writes, id).at(-1); + return last?.data[`${config.statusFieldName}.state`]; +} + +beforeEach(() => { + vi.clearAllMocks(); + // `clearAllMocks` keeps implementations, and some tests install a rejecting + // or recording `enqueue`. + enqueue.mockReset(); + batchSize.value = 2; + getSingleEmbedding.mockResolvedValue(EMBEDDING); + getEmbeddings.mockImplementation(async (inputs: string[]) => + inputs.map(() => EMBEDDING) + ); +}); + +describe("chunkArray", () => { + test("splits into chunks of at most the given size", () => { + expect(chunkArray([1, 2, 3, 4, 5], 2)).toEqual([[1, 2], [3, 4], [5]]); + }); + + test("returns nothing for an empty array", () => { + expect(chunkArray([], 2)).toEqual([]); + }); +}); + +describe("getNextTaskId", () => { + test("increments the counter", () => { + expect(getNextTaskId("kit-test-instance-task-1", "test-instance")).toBe( + "kit-test-instance-task-2" + ); + expect(getNextTaskId("kit-test-instance-task-49", "test-instance")).toBe( + "kit-test-instance-task-50" + ); + }); + + test('reads the counter when the instance id itself contains "task-"', () => { + expect(getNextTaskId("kit-my-task-force-task-3", "my-task-force")).toBe( + "kit-my-task-force-task-4" + ); + }); + + test("rejects an id that is not part of this instance's thread", () => { + expect(() => getNextTaskId("task-1", "test-instance")).toThrow( + "Invalid task ID format: task-1" + ); + expect(() => + getNextTaskId("kit-other-instance-task-1", "test-instance") + ).toThrow("Invalid task ID format"); + }); +}); + +describe("updateOrCreateMetadataDoc", () => { + test("creates the metadata document and requires a pass", async () => { + const { firestore, writes } = makeFirestore(); + + const result = await updateOrCreateMetadataDoc( + firestore as never, + METADATA_PATH, + METADATA + ); + + expect(result).toEqual({ path: METADATA_PATH, shouldBackfill: true }); + expect(writes).toHaveLength(1); + expect(writes[0].merge).toBe(true); + expect(writes[0].data).toMatchObject(METADATA); + }); + + test("skips the pass when the embedding configuration is unchanged", async () => { + const { firestore, writes } = makeFirestore({ [METADATA_PATH]: METADATA }); + + const result = await updateOrCreateMetadataDoc( + firestore as never, + METADATA_PATH, + METADATA + ); + + expect(result.shouldBackfill).toBe(false); + expect(writes).toHaveLength(0); + }); + + test.each([ + ["embeddingProvider", { embeddingProvider: "openai" }], + ["dimension", { dimension: 512 }], + ["inputField", { inputField: "text" }], + ["outputField", { outputField: "vector" }], + ])("requires a pass when %s changed", async (_field, previous) => { + const { firestore, writes } = makeFirestore({ + [METADATA_PATH]: { ...METADATA, ...previous }, + }); + + const result = await updateOrCreateMetadataDoc( + firestore as never, + METADATA_PATH, + METADATA + ); + + expect(result.shouldBackfill).toBe(true); + expect(writes).toHaveLength(1); + expect(writes[0].data).toMatchObject(METADATA); + }); + + test("merges so the progress counters do not replace the comparison fields", async () => { + const { firestore, store } = makeFirestore(); + + await updateOrCreateMetadataDoc( + firestore as never, + METADATA_PATH, + METADATA + ); + await enqueueTaskThread({ + firestore: firestore as never, + tasksDoc: METADATA_PATH, + queue: { enqueue } as never, + taskParams: ["doc-1"], + instanceId: config.instanceId, + }); + + // The extension replaced the document here, which lost these fields and + // made every later deploy re-embed the whole collection. + expect(store.get(METADATA_PATH)).toMatchObject(METADATA); + + const second = await updateOrCreateMetadataDoc( + firestore as never, + METADATA_PATH, + METADATA + ); + expect(second.shouldBackfill).toBe(false); + }); +}); + +describe("enqueueTaskThread", () => { + test("chunks ids, records every chunk, and dispatches only the first task", async () => { + const { firestore, writes, store } = makeFirestore(); + const ids = Array.from({ length: 120 }, (_, i) => `doc-${i}`); + + await enqueueTaskThread({ + firestore: firestore as never, + tasksDoc: METADATA_PATH, + queue: { enqueue } as never, + taskParams: ids, + instanceId: config.instanceId, + }); + + expect(enqueue).toHaveBeenCalledTimes(1); + expect(enqueue).toHaveBeenCalledWith({ + taskId: "kit-test-instance-task-1", + chunk: ids.slice(0, 50), + tasksDoc: METADATA_PATH, + }); + + const enqueueDocs = writes.filter((write) => + write.path.startsWith(`${METADATA_PATH}/enqueues/`) + ); + expect(enqueueDocs).toHaveLength(3); + expect(enqueueDocs[2].data).toEqual({ + taskId: "kit-test-instance-task-3", + status: "PENDING", + chunk: ids.slice(100), + }); + + expect(store.get(METADATA_PATH)).toMatchObject({ + backfillJobsTotal: 120, + backfillJobsProcessed: 0, + backfillJobsSkipped: 0, + backfillJobsFailed: 0, + backfillStatus: "RUNNING", + }); + expect(writes[0].merge).toBe(true); + }); + + test("records every chunk before dispatching the first task", async () => { + const { firestore, writes } = makeFirestore(); + const ids = Array.from({ length: 120 }, (_, i) => `doc-${i}`); + const dispatchedAfter: number[] = []; + enqueue.mockImplementation(async () => { + dispatchedAfter.push( + writes.filter((write) => + write.path.startsWith(`${METADATA_PATH}/enqueues/`) + ).length + ); + }); + + await enqueueTaskThread({ + firestore: firestore as never, + tasksDoc: METADATA_PATH, + queue: { enqueue } as never, + taskParams: ids, + instanceId: config.instanceId, + }); + + // All three enqueue documents exist by the time task-1 runs, so it can + // always find its successor. + expect(dispatchedAfter).toEqual([3]); + }); + + test("commits the trailing chunks when there are more than 50 of them", async () => { + const { firestore, writes } = makeFirestore(); + const ids = Array.from({ length: 2600 }, (_, i) => `doc-${i}`); + + await enqueueTaskThread({ + firestore: firestore as never, + tasksDoc: METADATA_PATH, + queue: { enqueue } as never, + taskParams: ids, + instanceId: config.instanceId, + }); + + // 52 chunks. The extension only committed on every 50th, so the last two + // enqueue documents were never written and the thread stalled on them. + const enqueueDocs = writes.filter((write) => + write.path.startsWith(`${METADATA_PATH}/enqueues/`) + ); + expect(enqueueDocs).toHaveLength(52); + expect(enqueueDocs.at(-1)?.data).toMatchObject({ + taskId: "kit-test-instance-task-52", + }); + }); + + test("writes nothing but the progress document for an empty id list", async () => { + const { firestore, writes } = makeFirestore(); + + await enqueueTaskThread({ + firestore: firestore as never, + tasksDoc: METADATA_PATH, + queue: { enqueue } as never, + taskParams: [], + instanceId: config.instanceId, + }); + + expect(enqueue).not.toHaveBeenCalled(); + expect(writes).toHaveLength(1); + expect(writes[0].data).toMatchObject({ backfillJobsTotal: 0 }); + }); +}); + +describe("handleBackfillTask", () => { + const TASK = { + taskId: "kit-test-instance-task-1", + chunk: ["doc-1", "doc-2"], + tasksDoc: METADATA_PATH, + }; + + test("embeds a batch of documents in a single call", async () => { + const { ctx, writes, commits } = makeCtx({ + [METADATA_PATH]: progress(), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${METADATA_PATH}/enqueues/kit-test-instance-task-2`]: { + chunk: ["doc-3", "doc-4"], + }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + }); + + await handleBackfillTask(taskRequest(TASK), ctx); + + expect(getEmbeddings).toHaveBeenCalledTimes(1); + expect(getEmbeddings).toHaveBeenCalledWith(["one", "two"]); + expect(getSingleEmbedding).not.toHaveBeenCalled(); + + expect(stateOf(writes, "doc-1")).toBe("BACKFILLED"); + expect(stateOf(writes, "doc-2")).toBe("BACKFILLED"); + expect(docWrites(writes, "doc-1")[0].data).toHaveProperty( + config.outputFieldName + ); + // Both documents are written in one committed batch. + expect( + commits.some( + (ops) => + ops.length === 2 && ops.every((op) => op.path.startsWith(COLLECTION)) + ) + ).toBe(true); + }); + + test("marks the enqueue document and dispatches the next task", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress(), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${METADATA_PATH}/enqueues/kit-test-instance-task-2`]: { + chunk: ["doc-3", "doc-4"], + }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + }); + + await handleBackfillTask(taskRequest(TASK), ctx); + + const taskDocWrites = writes.filter( + (write) => write.path === `${METADATA_PATH}/enqueues/${TASK.taskId}` + ); + expect(taskDocWrites.map((write) => write.data.status)).toEqual([ + "PROCESSING", + "DONE", + ]); + expect(enqueue).toHaveBeenCalledWith({ + taskId: "kit-test-instance-task-2", + chunk: ["doc-3", "doc-4"], + tasksDoc: METADATA_PATH, + }); + expect( + writes.find( + (write) => + write.path === METADATA_PATH && "backfillJobsProcessed" in write.data + ) + ).toBeDefined(); + }); + + test("finishes the thread instead of dispatching when every job is accounted for", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ + backfillJobsTotal: 2, + }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + }); + + await handleBackfillTask(taskRequest(TASK), ctx); + + expect(enqueue).not.toHaveBeenCalled(); + expect( + writes.some( + (write) => + write.path === METADATA_PATH && write.data.backfillStatus === "DONE" + ) + ).toBe(true); + }); + + test("uses the single-document path for a chunk with one eligible document", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 1 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: ["doc-1"] }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + }); + + await handleBackfillTask(taskRequest({ ...TASK, chunk: ["doc-1"] }), ctx); + + expect(getSingleEmbedding).toHaveBeenCalledWith("one"); + expect(getEmbeddings).not.toHaveBeenCalled(); + expect(stateOf(writes, "doc-1")).toBe("BACKFILLED"); + }); + + test("skips documents without a usable input string", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 2 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${COLLECTION}/doc-1`]: { input: 42 }, + [`${COLLECTION}/doc-2`]: { input: "" }, + }); + + await handleBackfillTask(taskRequest(TASK), ctx); + + expect(getEmbeddings).not.toHaveBeenCalled(); + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(docWrites(writes, "doc-1")).toHaveLength(0); + expect(docWrites(writes, "doc-2")).toHaveLength(0); + }); + + test("skips a document that is missing entirely", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 2 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + }); + + await handleBackfillTask(taskRequest(TASK), ctx); + + expect(getEmbeddings).not.toHaveBeenCalled(); + expect(writes.filter((write) => write.path.startsWith(COLLECTION))).toEqual( + [] + ); + }); + + test("skips documents whose status is already in a non-backfill state", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 2 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${COLLECTION}/doc-1`]: { input: "one", status: { state: "COMPLETED" } }, + [`${COLLECTION}/doc-2`]: { input: "two", status: { state: "ERROR" } }, + }); + + await handleBackfillTask(taskRequest(TASK), ctx); + + expect(getEmbeddings).not.toHaveBeenCalled(); + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(docWrites(writes, "doc-1")).toHaveLength(0); + }); + + test("re-embeds a document that was previously backfilled", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 1 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: ["doc-1"] }, + [`${COLLECTION}/doc-1`]: { + input: "one", + status: { state: "BACKFILLED" }, + }, + }); + + await handleBackfillTask(taskRequest({ ...TASK, chunk: ["doc-1"] }), ctx); + + expect(getSingleEmbedding).toHaveBeenCalledWith("one"); + expect(stateOf(writes, "doc-1")).toBe("BACKFILLED"); + }); + + test("marks a failed batch as FAILED_BACKFILL without failing the task", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 2 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + }); + getEmbeddings.mockRejectedValue(new Error("provider is down")); + + await expect( + handleBackfillTask(taskRequest(TASK), ctx) + ).resolves.toBeUndefined(); + + expect(stateOf(writes, "doc-1")).toBe("FAILED_BACKFILL"); + expect(stateOf(writes, "doc-2")).toBe("FAILED_BACKFILL"); + expect(docWrites(writes, "doc-1")[0].data).not.toHaveProperty( + config.outputFieldName + ); + }); + + test("marks a failed single document as FAILED_BACKFILL without failing the task", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 1 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: ["doc-1"] }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + }); + getSingleEmbedding.mockRejectedValue(new Error("provider is down")); + + await expect( + handleBackfillTask(taskRequest({ ...TASK, chunk: ["doc-1"] }), ctx) + ).resolves.toBeUndefined(); + + expect(stateOf(writes, "doc-1")).toBe("FAILED_BACKFILL"); + }); + + test("splits a chunk into provider-sized embedding calls", async () => { + batchSize.value = 2; + const chunk = ["doc-1", "doc-2", "doc-3"]; + const { ctx } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 3 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + [`${COLLECTION}/doc-3`]: { input: "three" }, + }); + + await handleBackfillTask(taskRequest({ ...TASK, chunk }), ctx); + + expect(getEmbeddings.mock.calls).toEqual([[["one", "two"]], [["three"]]]); + }); + + test("does nothing for an empty chunk", async () => { + const { ctx, writes } = makeCtx({ [METADATA_PATH]: progress() }); + + await handleBackfillTask(taskRequest({ ...TASK, chunk: [] }), ctx); + + expect(writes).toEqual([]); + expect(enqueue).not.toHaveBeenCalled(); + }); + + test("rejects a progress document without counters", async () => { + const { ctx } = makeCtx({ + [METADATA_PATH]: { ...METADATA }, + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + }); + + await expect(handleBackfillTask(taskRequest(TASK), ctx)).rejects.toThrow( + "Invalid task document" + ); + }); + + test("fails when the next enqueue document is missing", async () => { + const { ctx } = makeCtx({ + [METADATA_PATH]: progress(), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + }); + + await expect(handleBackfillTask(taskRequest(TASK), ctx)).rejects.toThrow( + "Next task document kit-test-instance-task-2 does not exist." + ); + }); +}); + +describe("handleUpdateTask", () => { + const TASK = { + taskId: "kit-test-instance-task-1", + chunk: ["doc-1", "doc-2"], + tasksDoc: METADATA_PATH, + }; + + test("only re-embeds documents that already carry an embedding", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 2 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${COLLECTION}/doc-1`]: { input: "one", embedding: [0, 0, 0] }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + }); + + await handleUpdateTask(taskRequest(TASK), ctx); + + expect(getSingleEmbedding).toHaveBeenCalledTimes(1); + expect(getSingleEmbedding).toHaveBeenCalledWith("one"); + expect(stateOf(writes, "doc-1")).toBe("BACKFILLED"); + expect(docWrites(writes, "doc-2")).toHaveLength(0); + }); + + test("embeds one document per call and fails only the documents that failed", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: progress({ backfillJobsTotal: 2 }), + [`${METADATA_PATH}/enqueues/${TASK.taskId}`]: { chunk: TASK.chunk }, + [`${COLLECTION}/doc-1`]: { input: "one", embedding: [0, 0, 0] }, + [`${COLLECTION}/doc-2`]: { input: "two", embedding: [0, 0, 0] }, + }); + getSingleEmbedding.mockImplementation(async (input: string) => { + if (input === "two") throw new Error("provider is down"); + return EMBEDDING; + }); + + await expect( + handleUpdateTask(taskRequest(TASK), ctx) + ).resolves.toBeUndefined(); + + expect(getEmbeddings).not.toHaveBeenCalled(); + expect(stateOf(writes, "doc-1")).toBe("BACKFILLED"); + expect(docWrites(writes, "doc-1")[0].data).toHaveProperty( + config.outputFieldName + ); + expect(stateOf(writes, "doc-2")).toBe("FAILED_BACKFILL"); + expect(docWrites(writes, "doc-2")[0].data).not.toHaveProperty( + config.outputFieldName + ); + }); +}); + +describe.each([ + ["handleBackfillTrigger", handleBackfillTrigger], + ["handleUpdateTrigger", handleUpdateTrigger], +])("%s", (_name, handler) => { + const request = {} as Request; + + test("enumerates the collection by reference and enqueues the first task", async () => { + const { ctx, writes } = makeCtx({ + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "two" }, + }); + + await handler(request, ctx); + + expect(enqueue).toHaveBeenCalledTimes(1); + expect(enqueue).toHaveBeenCalledWith({ + taskId: "kit-test-instance-task-1", + chunk: ["doc-1", "doc-2"], + tasksDoc: METADATA_PATH, + }); + expect(writes[0].data).toMatchObject(METADATA); + }); + + test("skips the pass when the embedding configuration is unchanged", async () => { + const { ctx, writes } = makeCtx({ + [METADATA_PATH]: { ...METADATA }, + [`${COLLECTION}/doc-1`]: { input: "one" }, + }); + + await handler(request, ctx); + + expect(enqueue).not.toHaveBeenCalled(); + expect(writes).toEqual([]); + }); + + test("enqueues nothing for an empty collection", async () => { + const { ctx } = makeCtx(); + + await handler(request, ctx); + + expect(enqueue).not.toHaveBeenCalled(); + }); + + test("swallows an enqueue failure so the trigger task is not retried", async () => { + const { ctx } = makeCtx({ [`${COLLECTION}/doc-1`]: { input: "one" } }); + enqueue.mockRejectedValue(new Error("queue not found")); + + await expect(handler(request, ctx)).resolves.toBeUndefined(); + }); +}); + +describe("handleInit", () => { + function ctxWith(overrides: { + doBackfill: boolean; + updateOnConfigure: boolean; + }) { + const { firestore } = makeFirestore(); + return { + firestore, + config: resolveVectorSearchConfig({ + projectId: "test-project", + instanceId: "test-instance", + region: "us-central1", + ...overrides, + }), + } as unknown as HandlerContext; + } + + test("enqueues only the backfill trigger when both passes are enabled", async () => { + await handleInit(ctxWith({ doBackfill: true, updateOnConfigure: true })); + + // Both passes share one task thread on the metadata document, so running + // them together would have them overwrite each other's progress. + expect(enqueue).toHaveBeenCalledTimes(1); + }); + + test("enqueues the backfill trigger on its own", async () => { + await handleInit(ctxWith({ doBackfill: true, updateOnConfigure: false })); + + expect(enqueue).toHaveBeenCalledTimes(1); + }); + + test("enqueues the update trigger on its own", async () => { + await handleInit(ctxWith({ doBackfill: false, updateOnConfigure: true })); + + expect(enqueue).toHaveBeenCalledTimes(1); + }); + + test("enqueues nothing when neither pass is enabled", async () => { + await handleInit(ctxWith({ doBackfill: false, updateOnConfigure: false })); + + expect(enqueue).not.toHaveBeenCalled(); + }); +}); + +describe("getValidDocs", () => { + const process = { + id: "test-instance", + batchSize: 2, + shouldBackfill: (data: Record) => + typeof data.input === "string" && data.input.length > 0, + processFn: async () => ({}), + } satisfies BackfillProcess; + + /** A Firestore whose transaction callback runs twice, as a retry would. */ + function retryingFirestore(seed: Record>) { + let attempts = 0; + const snapshot = (path: string) => ({ + exists: seed[path] !== undefined, + data: () => seed[path], + ref: { path, id: path.split("/").pop() as string }, + }); + return { + attempts: () => attempts, + firestore: { + collection: (name: string) => ({ + doc: (id: string) => ({ path: `${name}/${id}` }), + }), + runTransaction: async (fn: (tx: unknown) => Promise) => { + const tx = { + getAll: async (...refs: { path: string }[]) => + refs.map((r) => snapshot(r.path)), + }; + attempts++; + await fn(tx); + attempts++; + return fn(tx); + }, + }, + }; + } + + test("does not double-count documents when the transaction retries", async () => { + const { firestore, attempts } = retryingFirestore({ + [`${COLLECTION}/doc-1`]: { input: "one" }, + [`${COLLECTION}/doc-2`]: { input: "" }, + }); + + const result = await getValidDocs(process, ["doc-1", "doc-2"], { + firestore: firestore as never, + collectionName: COLLECTION, + statusField: config.statusFieldName, + }); + + expect(attempts()).toBe(2); + expect(result.validDocuments.map((d) => d.ref.id)).toEqual(["doc-1"]); + expect(result.skippedDocuments.map((d) => d.ref.id)).toEqual(["doc-2"]); + }); +}); diff --git a/kits/firestore-vector-search/tests/embeddings.test.ts b/kits/firestore-vector-search/tests/embeddings.test.ts index ee80ea15a7..a318c4ff18 100644 --- a/kits/firestore-vector-search/tests/embeddings.test.ts +++ b/kits/firestore-vector-search/tests/embeddings.test.ts @@ -16,7 +16,10 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; -const { embedMany } = vi.hoisted(() => ({ embedMany: vi.fn() })); +const { embedMany, createEmbeddings } = vi.hoisted(() => ({ + embedMany: vi.fn(), + createEmbeddings: vi.fn(), +})); vi.mock("genkit", () => ({ genkit: vi.fn(() => ({ embedMany })), @@ -31,11 +34,18 @@ vi.mock("@genkit-ai/google-genai", () => ({ }), })); +vi.mock("openai", () => ({ + default: class { + embeddings = { create: createEmbeddings }; + }, +})); + 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, @@ -257,3 +267,47 @@ describe("CustomEndpointClient", () => { ); }); }); + +describe("OpenAiEmbedClient", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + function openAiConfig() { + return config({ embeddingProvider: "openai", openAiApiKey: "test-key" }); + } + + test("requires an API key", () => { + expect( + () => + new OpenAiEmbedClient( + config({ embeddingProvider: "openai", openAiApiKey: undefined }) + ) + ).toThrow("OpenAI embeddings require OPENAI_API_KEY"); + }); + + test("embeds 16 inputs per batch, as the extension did", () => { + expect(new OpenAiEmbedClient(openAiConfig()).batchSize).toBe(16); + }); + + test("sends a whole batch in one request", async () => { + createEmbeddings.mockResolvedValueOnce({ + data: [{ embedding: [1, 2, 3] }, { embedding: [4, 5, 6] }], + }); + + const embeddings = await new OpenAiEmbedClient( + openAiConfig() + ).getEmbeddings(["one", "two"]); + + expect(createEmbeddings).toHaveBeenCalledTimes(1); + expect(createEmbeddings).toHaveBeenCalledWith({ + model: "text-embedding-3-small", + input: ["one", "two"], + dimensions: 512, + }); + expect(embeddings).toEqual([ + [1, 2, 3], + [4, 5, 6], + ]); + }); +});