export async function backfillEmbeddingsTaskHandler()

in storage-reverse-image-search/functions/src/functions/backfill_task.ts [27:104]


export async function backfillEmbeddingsTaskHandler(data: any) {
  const {id, objects} = data;

  const taskRef = admin.firestore().doc(`${config.tasksDoc}/enqueues/${id}`);

  await taskRef.update({
    status: BackfillStatus.PROCESSING,
  });

  const datapoints: any = [];

  const featureVectors = await getFeatureVectors(objects);

  if (featureVectors) {
    for (let i = 0; i < featureVectors.length; i++) {
      datapoints.push({
        id: objects[i],
        embedding: featureVectors[i],
      });
    }
  }

  if (datapoints.length === 0) {
    functions.logger.info('No datapoints found, skipping...');
    return;
  }

  const outputShape = featureVectors[0].length;

  const localFilePath = await utils.saveEmbeddingsToTmpFile(datapoints);

  const destinationPath = `datapoints/${id}.json`;

  functions.logger.info(
    `Embeddings will be saved to ${destinationPath} 📝, uploading to the bucket...`
  );

  await utils.uploadToCloudStorage(localFilePath, destinationPath);

  functions.logger.info(
    `Embeddings uploaded to the bucket ${config.bucketName} in ${destinationPath} 🚀`
  );

  await taskRef.update({
    status: 'DONE',
    filePath: `gs://${config.bucketName}/${destinationPath}`,
  });

  await utils.deleteTempFiles();

  const tasksDoc = await admin.firestore().doc(config.tasksDoc).get();
  const {totalLength} = tasksDoc.data() as any;
  let {processedLength} = tasksDoc.data() as any;

  processedLength += objects.length;
  await admin
    .firestore()
    .doc(config.tasksDoc)
    .update({
      processedLength: admin.firestore.FieldValue.increment(objects.length),
    });

  if (processedLength === totalLength) {
    // Update the metadata doc with the output shape to be used in the index creation.
    await admin.firestore().doc(config.metadataDoc).set(
      {
        outputShape,
      },
      {merge: true}
    );

    await admin.firestore().doc(config.tasksDoc).update({
      status: BackfillStatus.DONE,
    });
  } else {
    await _createNextTask(id);
  }
}