UNPKG

@genkit-ai/firebase

Version:

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

69 lines (65 loc) 3.19 kB
import { Firestore, QueryDocumentSnapshot } from '@google-cloud/firestore'; import { Genkit, EmbedderArgument, RetrieverAction } from 'genkit'; import { Part } from 'genkit/retriever'; /** * Copyright 2024 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 * * http://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. */ /** * Define a retriever that uses vector similarity search to retrieve documents from Firestore. * You must create a vector index on the associated field before you can perform nearest-neighbor * search. **/ declare function defineFirestoreRetriever(ai: Genkit, config: { /** The name of the retriever. */ name: string; /** Optional label for display in Developer UI. */ label?: string; /** The Firestore database instance from which to query. */ firestore: Firestore; /** The name of the collection from which to query. */ collection?: string; /** The embedder to use with this retriever. */ embedder: EmbedderArgument; /** The name of the field within the collection containing the vector data. */ vectorField: string; /** The name of the field containing the document content you wish to return. */ contentField: string | ((snap: QueryDocumentSnapshot) => Part[]); /** The distance measure to use when comparing vectors. Defaults to 'COSINE'. */ distanceMeasure?: 'EUCLIDEAN' | 'COSINE' | 'DOT_PRODUCT'; /** * Specifies a threshold for which no less similar documents will be returned. The behavior * of the specified `distanceMeasure` will affect the meaning of the distance threshold. * * - For `distanceMeasure: "EUCLIDEAN"`, the meaning of `distanceThreshold` is: * SELECT docs WHERE euclidean_distance <= distanceThreshold * - For `distanceMeasure: "COSINE"`, the meaning of `distanceThreshold` is: * SELECT docs WHERE cosine_distance <= distanceThreshold * - For `distanceMeasure: "DOT_PRODUCT"`, the meaning of `distanceThreshold` is: * SELECT docs WHERE dot_product_distance >= distanceThreshold */ distanceThreshold?: number; /** * Optionally specifies the name of a metadata field that will be set on each returned Document, * which will contain the computed distance for the document. */ distanceResultField?: string; /** * A list of fields to include in the returned document metadata. If not supplied, all fields other * than the vector are included. Alternatively, provide a transform function to extract the desired * metadata fields from a snapshot. **/ metadataFields?: string[] | ((snap: QueryDocumentSnapshot) => Record<string, any>); }): RetrieverAction; export { defineFirestoreRetriever };