@genkit-ai/firebase/lib/firestore-retriever.mjs.map

Genkit AI framework plugin for Firebase including Firestore trace/state store and deployment helpers for Cloud Functions for Firebase.

{"version":3,"sources":["../src/firestore-retriever.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport type {\n  Firestore,\n  Query,\n  QueryDocumentSnapshot,\n  VectorQuerySnapshot,\n} from '@google-cloud/firestore';\nimport {\n  z,\n  type EmbedderArgument,\n  type Genkit,\n  type RetrieverAction,\n} from 'genkit';\nimport type { DocumentData, Part } from 'genkit/retriever';\n\nfunction toContent(\n  d: QueryDocumentSnapshot,\n  contentField: string | ((snap: QueryDocumentSnapshot) => Part[])\n): Part[] {\n  if (typeof contentField === 'function') {\n    return contentField(d);\n  }\n\n  return [{ text: d.get(contentField) }];\n}\n\nfunction toDocuments(\n  result: VectorQuerySnapshot,\n  vectorField: string,\n  contentField: string | ((snap: QueryDocumentSnapshot) => Part[]),\n  metadataFields?:\n    | string[]\n    | ((snap: QueryDocumentSnapshot) => Record<string, any>)\n): DocumentData[] {\n  return result.docs.map((d) => {\n    const out: DocumentData = { content: toContent(d, contentField) };\n    if (typeof metadataFields === 'function') {\n      out.metadata = metadataFields(d);\n      return out;\n    }\n\n    out.metadata = { id: d.id };\n    if (metadataFields) {\n      for (const field of metadataFields) {\n        out.metadata[field] = d.get(field);\n      }\n      return out;\n    }\n\n    out.metadata = d.data();\n    delete out.metadata[vectorField];\n    if (typeof contentField === 'string') delete out.metadata[contentField];\n    return out;\n  });\n}\n\n/**\n * Define a retriever that uses vector similarity search to retrieve documents from Firestore.\n * You must create a vector index on the associated field before you can perform nearest-neighbor\n * search.\n **/\nexport function defineFirestoreRetriever(\n  ai: Genkit,\n  config: {\n    /** The name of the retriever. */\n    name: string;\n    /** Optional label for display in Developer UI. */\n    label?: string;\n    /** The Firestore database instance from which to query. */\n    firestore: Firestore;\n    /** The name of the collection from which to query. */\n    collection?: string;\n    /** The embedder to use with this retriever. */\n    embedder: EmbedderArgument;\n    /** The name of the field within the collection containing the vector data. */\n    vectorField: string;\n    /** The name of the field containing the document content you wish to return. */\n    contentField: string | ((snap: QueryDocumentSnapshot) => Part[]);\n    /** The distance measure to use when comparing vectors. Defaults to 'COSINE'. */\n    distanceMeasure?: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT';\n    /**\n     * Specifies a threshold for which no less similar documents will be returned. The behavior\n     * of the specified `distanceMeasure` will affect the meaning of the distance threshold.\n     *\n     *  - For `distanceMeasure: \"EUCLIDEAN\"`, the meaning of `distanceThreshold` is:\n     *     SELECT docs WHERE euclidean_distance <= distanceThreshold\n     *  - For `distanceMeasure: \"COSINE\"`, the meaning of `distanceThreshold` is:\n     *     SELECT docs WHERE cosine_distance <= distanceThreshold\n     *  - For `distanceMeasure: \"DOT_PRODUCT\"`, the meaning of `distanceThreshold` is:\n     *     SELECT docs WHERE dot_product_distance >= distanceThreshold\n     */\n    distanceThreshold?: number;\n    /**\n     * Optionally specifies the name of a metadata field that will be set on each returned Document,\n     * which will contain the computed distance for the document.\n     */\n    distanceResultField?: string;\n    /**\n     * A list of fields to include in the returned document metadata. If not supplied, all fields other\n     * than the vector are included. Alternatively, provide a transform function to extract the desired\n     * metadata fields from a snapshot.\n     **/\n    metadataFields?:\n      | string[]\n      | ((snap: QueryDocumentSnapshot) => Record<string, any>);\n  }\n): RetrieverAction {\n  const {\n    name,\n    label,\n    firestore,\n    embedder,\n    collection,\n    vectorField,\n    metadataFields,\n    contentField,\n    distanceMeasure,\n    distanceThreshold,\n    distanceResultField,\n  } = config;\n  return ai.defineRetriever(\n    {\n      name,\n      info: {\n        label: label || `Firestore - ${name}`,\n      },\n      configSchema: z.object({\n        where: z.record(z.any()).optional(),\n        /** Max number of results to return. Defaults to 10. */\n        limit: z.number().optional(),\n        /* Supply or override the distanceMeasure */\n        distanceMeasure: z\n          .enum(['COSINE', 'DOT_PRODUCT', 'EUCLIDEAN'])\n          .optional(),\n        /* Supply or override the distanceThreshold */\n        distanceThreshold: z.number().optional(),\n        /* Supply or override the metadata field where distances are stored. */\n        distanceResultField: z.string().optional(),\n        /* Supply or override the collection for retrieval. */\n        collection: z.string().optional(),\n      }),\n    },\n    async (content, options) => {\n      options = options || {};\n      if (!options.collection && !collection) {\n        throw new Error(\n          'Must specify a collection to query in Firestore retriever.'\n        );\n      }\n      let query: Query = firestore.collection(\n        options.collection || collection!\n      );\n      for (const field in options.where || {}) {\n        query = query.where(field, '==', options.where![field]);\n      }\n      // Single embedding for text input\n      const queryVector = (await ai.embed({ embedder, content }))[0].embedding;\n\n      const result = await query\n        .findNearest({\n          vectorField,\n          queryVector,\n          limit: options.limit || 10,\n          distanceMeasure:\n            options.distanceMeasure || distanceMeasure || 'COSINE',\n          distanceResultField:\n            options.distanceResultField || distanceResultField,\n          distanceThreshold: options.distanceThreshold || distanceThreshold,\n        })\n        .get();\n\n      return {\n        documents: toDocuments(\n          result,\n          vectorField,\n          contentField,\n          metadataFields\n        ),\n      };\n    }\n  );\n}\n"],"mappings":"AAsBA;AAAA,EACE;AAAA,OAIK;AAGP,SAAS,UACP,GACA,cACQ;AACR,MAAI,OAAO,iBAAiB,YAAY;AACtC,WAAO,aAAa,CAAC;AAAA,EACvB;AAEA,SAAO,CAAC,EAAE,MAAM,EAAE,IAAI,YAAY,EAAE,CAAC;AACvC;AAEA,SAAS,YACP,QACA,aACA,cACA,gBAGgB;AAChB,SAAO,OAAO,KAAK,IAAI,CAAC,MAAM;AAC5B,UAAM,MAAoB,EAAE,SAAS,UAAU,GAAG,YAAY,EAAE;AAChE,QAAI,OAAO,mBAAmB,YAAY;AACxC,UAAI,WAAW,eAAe,CAAC;AAC/B,aAAO;AAAA,IACT;AAEA,QAAI,WAAW,EAAE,IAAI,EAAE,GAAG;AAC1B,QAAI,gBAAgB;AAClB,iBAAW,SAAS,gBAAgB;AAClC,YAAI,SAAS,KAAK,IAAI,EAAE,IAAI,KAAK;AAAA,MACnC;AACA,aAAO;AAAA,IACT;AAEA,QAAI,WAAW,EAAE,KAAK;AACtB,WAAO,IAAI,SAAS,WAAW;AAC/B,QAAI,OAAO,iBAAiB,SAAU,QAAO,IAAI,SAAS,YAAY;AACtE,WAAO;AAAA,EACT,CAAC;AACH;AAOO,SAAS,yBACd,IACA,QA2CiB;AACjB,QAAM;AAAA,IACJ;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,EACF,IAAI;AACJ,SAAO,GAAG;AAAA,IACR;AAAA,MACE;AAAA,MACA,MAAM;AAAA,QACJ,OAAO,SAAS,eAAe,IAAI;AAAA,MACrC;AAAA,MACA,cAAc,EAAE,OAAO;AAAA,QACrB,OAAO,EAAE,OAAO,EAAE,IAAI,CAAC,EAAE,SAAS;AAAA;AAAA,QAElC,OAAO,EAAE,OAAO,EAAE,SAAS;AAAA;AAAA,QAE3B,iBAAiB,EACd,KAAK,CAAC,UAAU,eAAe,WAAW,CAAC,EAC3C,SAAS;AAAA;AAAA,QAEZ,mBAAmB,EAAE,OAAO,EAAE,SAAS;AAAA;AAAA,QAEvC,qBAAqB,EAAE,OAAO,EAAE,SAAS;AAAA;AAAA,QAEzC,YAAY,EAAE,OAAO,EAAE,SAAS;AAAA,MAClC,CAAC;AAAA,IACH;AAAA,IACA,OAAO,SAAS,YAAY;AAC1B,gBAAU,WAAW,CAAC;AACtB,UAAI,CAAC,QAAQ,cAAc,CAAC,YAAY;AACtC,cAAM,IAAI;AAAA,UACR;AAAA,QACF;AAAA,MACF;AACA,UAAI,QAAe,UAAU;AAAA,QAC3B,QAAQ,cAAc;AAAA,MACxB;AACA,iBAAW,SAAS,QAAQ,SAAS,CAAC,GAAG;AACvC,gBAAQ,MAAM,MAAM,OAAO,MAAM,QAAQ,MAAO,KAAK,CAAC;AAAA,MACxD;AAEA,YAAM,eAAe,MAAM,GAAG,MAAM,EAAE,UAAU,QAAQ,CAAC,GAAG,CAAC,EAAE;AAE/D,YAAM,SAAS,MAAM,MAClB,YAAY;AAAA,QACX;AAAA,QACA;AAAA,QACA,OAAO,QAAQ,SAAS;AAAA,QACxB,iBACE,QAAQ,mBAAmB,mBAAmB;AAAA,QAChD,qBACE,QAAQ,uBAAuB;AAAA,QACjC,mBAAmB,QAAQ,qBAAqB;AAAA,MAClD,CAAC,EACA,IAAI;AAEP,aAAO;AAAA,QACL,WAAW;AAAA,UACT;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,QACF;AAAA,MACF;AAAA,IACF;AAAA,EACF;AACF;","names":[]}