{"version":3,"sources":["../../../src/vectorsearch/vector_search/retrievers.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 {\n retrieverRef,\n type Genkit,\n type RetrieverAction,\n type z,\n} from 'genkit';\nimport { queryPublicEndpoint } from './query_public_endpoint';\nimport {\n VertexAIVectorRetrieverOptionsSchema,\n type VertexVectorSearchOptions,\n} from './types';\nimport { getProjectNumber } from './utils';\n\nconst DEFAULT_K = 10;\n\n/**\n * Creates Vertex AI retrievers.\n *\n * This function returns a list of retriever actions for Vertex AI based on the provided\n * vector search options and embedder configurations.\n *\n * @param {VertexVectorSearchOptions} params - The parameters for creating the retrievers.\n * @returns {RetrieverAction[]} - An array of retriever actions.\n */\nexport function vertexAiRetrievers(\n ai: Genkit,\n params: VertexVectorSearchOptions\n): RetrieverAction[] {\n const vectorSearchOptions = params.pluginOptions.vectorSearchOptions;\n const defaultEmbedder = params.defaultEmbedder;\n\n const retrievers: RetrieverAction[] = [];\n\n if (!vectorSearchOptions || vectorSearchOptions.length === 0) {\n return retrievers;\n }\n\n for (const vectorSearchOption of vectorSearchOptions) {\n const { documentRetriever, indexId, publicDomainName } = vectorSearchOption;\n const embedderOptions = vectorSearchOption.embedderOptions;\n\n const retriever = ai.defineRetriever(\n {\n name: `vertexai/${indexId}`,\n configSchema: VertexAIVectorRetrieverOptionsSchema.optional(),\n },\n async (content, options) => {\n const embedderReference =\n vectorSearchOption.embedder ?? defaultEmbedder;\n\n if (!embedderReference) {\n throw new Error(\n 'Embedder reference is required to define Vertex AI retriever'\n );\n }\n\n const queryEmbedding = (\n await ai.embed({\n embedder: embedderReference,\n options: embedderOptions,\n content,\n })\n )[0].embedding; // Single embedding for text\n\n const accessToken = await params.authClient.getAccessToken();\n\n if (!accessToken) {\n throw new Error(\n 'Error generating access token when defining Vertex AI retriever'\n );\n }\n\n const projectId = params.pluginOptions.projectId;\n if (!projectId) {\n throw new Error(\n 'Project ID is required to define Vertex AI retriever'\n );\n }\n const projectNumber = await getProjectNumber(projectId);\n const location = params.pluginOptions.location;\n if (!location) {\n throw new Error('Location is required to define Vertex AI retriever');\n }\n\n const res = await queryPublicEndpoint({\n featureVector: queryEmbedding,\n neighborCount: options?.k || DEFAULT_K,\n accessToken,\n projectId,\n location,\n publicDomainName,\n projectNumber,\n indexEndpointId: vectorSearchOption.indexEndpointId,\n deployedIndexId: vectorSearchOption.deployedIndexId,\n restricts: content.metadata?.restricts,\n numericRestricts: content.metadata?.numericRestricts,\n });\n const nearestNeighbors = res.nearestNeighbors;\n\n const queryRes = nearestNeighbors ? nearestNeighbors[0] : null;\n const neighbors = queryRes ? queryRes.neighbors : null;\n if (!neighbors) {\n return { documents: [] };\n }\n\n const documents = await documentRetriever(neighbors, options);\n\n return { documents };\n }\n );\n\n retrievers.push(retriever);\n }\n\n return retrievers;\n}\n\n/**\n * Creates a reference to a Vertex AI retriever.\n *\n * @param {Object} params - The parameters for the retriever reference.\n * @param {string} params.indexId - The ID of the Vertex AI index.\n * @param {string} [params.displayName] - An optional display name for the retriever.\n * @returns {Object} - The retriever reference object.\n */\nexport const vertexAiRetrieverRef = (params: {\n indexId: string;\n displayName?: string;\n}) => {\n return retrieverRef({\n name: `vertexai/${params.indexId}`,\n info: {\n label: params.displayName ?? `ertex AI - ${params.indexId}`,\n },\n configSchema: VertexAIVectorRetrieverOptionsSchema.optional(),\n });\n};\n"],"mappings":"AAgBA;AAAA,EACE;AAAA,OAIK;AACP,SAAS,2BAA2B;AACpC;AAAA,EACE;AAAA,OAEK;AACP,SAAS,wBAAwB;AAEjC,MAAM,YAAY;AAWX,SAAS,mBACd,IACA,QACiC;AACjC,QAAM,sBAAsB,OAAO,cAAc;AACjD,QAAM,kBAAkB,OAAO;AAE/B,QAAM,aAA8C,CAAC;AAErD,MAAI,CAAC,uBAAuB,oBAAoB,WAAW,GAAG;AAC5D,WAAO;AAAA,EACT;AAEA,aAAW,sBAAsB,qBAAqB;AACpD,UAAM,EAAE,mBAAmB,SAAS,iBAAiB,IAAI;AACzD,UAAM,kBAAkB,mBAAmB;AAE3C,UAAM,YAAY,GAAG;AAAA,MACnB;AAAA,QACE,MAAM,YAAY,OAAO;AAAA,QACzB,cAAc,qCAAqC,SAAS;AAAA,MAC9D;AAAA,MACA,OAAO,SAAS,YAAY;AAC1B,cAAM,oBACJ,mBAAmB,YAAY;AAEjC,YAAI,CAAC,mBAAmB;AACtB,gBAAM,IAAI;AAAA,YACR;AAAA,UACF;AAAA,QACF;AAEA,cAAM,kBACJ,MAAM,GAAG,MAAM;AAAA,UACb,UAAU;AAAA,UACV,SAAS;AAAA,UACT;AAAA,QACF,CAAC,GACD,CAAC,EAAE;AAEL,cAAM,cAAc,MAAM,OAAO,WAAW,eAAe;AAE3D,YAAI,CAAC,aAAa;AAChB,gBAAM,IAAI;AAAA,YACR;AAAA,UACF;AAAA,QACF;AAEA,cAAM,YAAY,OAAO,cAAc;AACvC,YAAI,CAAC,WAAW;AACd,gBAAM,IAAI;AAAA,YACR;AAAA,UACF;AAAA,QACF;AACA,cAAM,gBAAgB,MAAM,iBAAiB,SAAS;AACtD,cAAM,WAAW,OAAO,cAAc;AACtC,YAAI,CAAC,UAAU;AACb,gBAAM,IAAI,MAAM,oDAAoD;AAAA,QACtE;AAEA,cAAM,MAAM,MAAM,oBAAoB;AAAA,UACpC,eAAe;AAAA,UACf,eAAe,SAAS,KAAK;AAAA,UAC7B;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA;AAAA,UACA,iBAAiB,mBAAmB;AAAA,UACpC,iBAAiB,mBAAmB;AAAA,UACpC,WAAW,QAAQ,UAAU;AAAA,UAC7B,kBAAkB,QAAQ,UAAU;AAAA,QACtC,CAAC;AACD,cAAM,mBAAmB,IAAI;AAE7B,cAAM,WAAW,mBAAmB,iBAAiB,CAAC,IAAI;AAC1D,cAAM,YAAY,WAAW,SAAS,YAAY;AAClD,YAAI,CAAC,WAAW;AACd,iBAAO,EAAE,WAAW,CAAC,EAAE;AAAA,QACzB;AAEA,cAAM,YAAY,MAAM,kBAAkB,WAAW,OAAO;AAE5D,eAAO,EAAE,UAAU;AAAA,MACrB;AAAA,IACF;AAEA,eAAW,KAAK,SAAS;AAAA,EAC3B;AAEA,SAAO;AACT;AAUO,MAAM,uBAAuB,CAAC,WAG/B;AACJ,SAAO,aAAa;AAAA,IAClB,MAAM,YAAY,OAAO,OAAO;AAAA,IAChC,MAAM;AAAA,MACJ,OAAO,OAAO,eAAe,cAAc,OAAO,OAAO;AAAA,IAC3D;AAAA,IACA,cAAc,qCAAqC,SAAS;AAAA,EAC9D,CAAC;AACH;","names":[]}