@genkit-ai/vertexai/lib/rerankers/v2/reranker.mjs.map

Genkit AI framework plugin for Google Cloud Vertex AI APIs including Gemini APIs, Imagen, and more.

{"version":3,"sources":["../../../src/rerankers/v2/reranker.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  RerankerInfo,\n  RerankerReference,\n  z,\n  type RerankerAction,\n} from 'genkit';\nimport { reranker as pluginReranker } from 'genkit/plugin';\nimport { RankedDocument, rerankerRef } from 'genkit/reranker';\nimport { checkModelName } from '../../common/utils.mjs';\nimport { rerankerRank } from './client.mjs';\nimport {\n  RerankRequest,\n  RerankRequestRecord,\n  VertexRerankerClientOptions,\n} from './types.mjs';\n\n// Define the schema for the options used in the Vertex AI reranker\nexport const VertexRerankerConfigSchema = z\n  .object({\n    topN: z.number().optional().describe('Number of top documents to rerank'), // Optional: Number of documents to rerank\n    ignoreRecordDetailsInResponse: z\n      .boolean()\n      .optional()\n      .describe(\n        'If true, the response will contain only record ID and score. By default, it is false, the response will contain record details.'\n      ),\n    location: z\n      .string()\n      .optional()\n      .describe('Google Cloud location, e.g., \"us-central1\"'), // Optional: Location of the reranking model\n  })\n  .passthrough();\nexport type VertexRerankerConfigSchemaType = typeof VertexRerankerConfigSchema;\nexport type VertexRerankerConfig = z.infer<VertexRerankerConfigSchemaType>;\n\ntype ConfigSchemaType = VertexRerankerConfigSchemaType;\n\nfunction commonRef(\n  name: string,\n  info?: RerankerInfo,\n  configSchema: ConfigSchemaType = VertexRerankerConfigSchema\n): RerankerReference<ConfigSchemaType> {\n  return rerankerRef({\n    name: `vertex-rerankers/${name}`,\n    configSchema,\n    info: info ?? {\n      supports: {\n        media: false,\n      },\n    },\n  });\n}\n\nexport const GENERIC_MODEL = commonRef('reranker');\n\nexport const DEFAULT_MODEL_NAME = 'semantic-ranker-default@latest';\n\nexport const KNOWN_MODELS = {\n  'semantic-ranker-default@latest': commonRef('semantic-ranker-default@latest'),\n  'semantic-ranker-default-004': commonRef('semantic-ranker-default-004'),\n  'semantic-ranker-fast-004': commonRef('semantic-ranker-fast-004'),\n  'semantic-ranker-default-003': commonRef('semantic-ranker-default-003'),\n  'semantic-ranker-default-002': commonRef('semantic-ranker-default-002'),\n} as const;\nexport type KnownModels = keyof typeof KNOWN_MODELS;\nexport type RerankerModelName = `semantic-ranker-${string}`;\nexport function isRerankerModelName(\n  value?: string\n): value is RerankerModelName {\n  return !!value?.startsWith('semantic-ranker-');\n}\n\nexport function reranker(\n  version: string,\n  options: VertexRerankerConfig = {}\n): RerankerReference<VertexRerankerConfigSchemaType> {\n  const name = checkModelName(version);\n  return rerankerRef({\n    name: `vertex-rerankers/${name}`,\n    configSchema: VertexRerankerConfigSchema,\n    info: {\n      ...GENERIC_MODEL.info,\n    },\n  });\n}\n\nexport function listKnownRerankers(clientOptions: VertexRerankerClientOptions) {\n  return Object.keys(KNOWN_MODELS).map((name) =>\n    defineReranker(name, clientOptions)\n  );\n}\n\nexport function defineReranker(\n  name: string,\n  clientOptions: VertexRerankerClientOptions\n): RerankerAction {\n  const ref = reranker(name);\n\n  return pluginReranker(\n    {\n      name: ref.name,\n      ...ref.info,\n      configSchema: ref.configSchema,\n    },\n    async (query, documents, options) => {\n      const rerankRequest: RerankRequest = {\n        // Note that model silently falls back to default if it's not recognized\n        // This happens in the vertexai reranker service backend.\n        model: checkModelName(ref.name),\n        query: query.text,\n        records: documents.map(toRerankerDoc),\n        ...options,\n      };\n      const response = await rerankerRank(\n        ref.name,\n        rerankRequest,\n        clientOptions\n      );\n      return { documents: fromRerankResponse(response, documents) };\n    }\n  );\n}\n\nfunction toRerankerDoc(doc, idx): RerankRequestRecord {\n  return {\n    id: `${idx}`,\n    content: doc.text,\n  };\n}\n\nfunction fromRerankResponse(response, documents): RankedDocument[] {\n  return response.records.map((record) => {\n    const doc = documents[record.id];\n    return new RankedDocument({\n      content: doc.content,\n      metadata: {\n        ...doc.metadata,\n        score: record.score,\n      },\n    });\n  });\n}\n"],"mappings":"AAgBA;AAAA,EAGE;AAAA,OAEK;AACP,SAAS,YAAY,sBAAsB;AAC3C,SAAS,gBAAgB,mBAAmB;AAC5C,SAAS,sBAAsB;AAC/B,SAAS,oBAAoB;AAQtB,MAAM,6BAA6B,EACvC,OAAO;AAAA,EACN,MAAM,EAAE,OAAO,EAAE,SAAS,EAAE,SAAS,mCAAmC;AAAA;AAAA,EACxE,+BAA+B,EAC5B,QAAQ,EACR,SAAS,EACT;AAAA,IACC;AAAA,EACF;AAAA,EACF,UAAU,EACP,OAAO,EACP,SAAS,EACT,SAAS,4CAA4C;AAAA;AAC1D,CAAC,EACA,YAAY;AAMf,SAAS,UACP,MACA,MACA,eAAiC,4BACI;AACrC,SAAO,YAAY;AAAA,IACjB,MAAM,oBAAoB,IAAI;AAAA,IAC9B;AAAA,IACA,MAAM,QAAQ;AAAA,MACZ,UAAU;AAAA,QACR,OAAO;AAAA,MACT;AAAA,IACF;AAAA,EACF,CAAC;AACH;AAEO,MAAM,gBAAgB,UAAU,UAAU;AAE1C,MAAM,qBAAqB;AAE3B,MAAM,eAAe;AAAA,EAC1B,kCAAkC,UAAU,gCAAgC;AAAA,EAC5E,+BAA+B,UAAU,6BAA6B;AAAA,EACtE,4BAA4B,UAAU,0BAA0B;AAAA,EAChE,+BAA+B,UAAU,6BAA6B;AAAA,EACtE,+BAA+B,UAAU,6BAA6B;AACxE;AAGO,SAAS,oBACd,OAC4B;AAC5B,SAAO,CAAC,CAAC,OAAO,WAAW,kBAAkB;AAC/C;AAEO,SAAS,SACd,SACA,UAAgC,CAAC,GACkB;AACnD,QAAM,OAAO,eAAe,OAAO;AACnC,SAAO,YAAY;AAAA,IACjB,MAAM,oBAAoB,IAAI;AAAA,IAC9B,cAAc;AAAA,IACd,MAAM;AAAA,MACJ,GAAG,cAAc;AAAA,IACnB;AAAA,EACF,CAAC;AACH;AAEO,SAAS,mBAAmB,eAA4C;AAC7E,SAAO,OAAO,KAAK,YAAY,EAAE;AAAA,IAAI,CAAC,SACpC,eAAe,MAAM,aAAa;AAAA,EACpC;AACF;AAEO,SAAS,eACd,MACA,eACgB;AAChB,QAAM,MAAM,SAAS,IAAI;AAEzB,SAAO;AAAA,IACL;AAAA,MACE,MAAM,IAAI;AAAA,MACV,GAAG,IAAI;AAAA,MACP,cAAc,IAAI;AAAA,IACpB;AAAA,IACA,OAAO,OAAO,WAAW,YAAY;AACnC,YAAM,gBAA+B;AAAA;AAAA;AAAA,QAGnC,OAAO,eAAe,IAAI,IAAI;AAAA,QAC9B,OAAO,MAAM;AAAA,QACb,SAAS,UAAU,IAAI,aAAa;AAAA,QACpC,GAAG;AAAA,MACL;AACA,YAAM,WAAW,MAAM;AAAA,QACrB,IAAI;AAAA,QACJ;AAAA,QACA;AAAA,MACF;AACA,aAAO,EAAE,WAAW,mBAAmB,UAAU,SAAS,EAAE;AAAA,IAC9D;AAAA,EACF;AACF;AAEA,SAAS,cAAc,KAAK,KAA0B;AACpD,SAAO;AAAA,IACL,IAAI,GAAG,GAAG;AAAA,IACV,SAAS,IAAI;AAAA,EACf;AACF;AAEA,SAAS,mBAAmB,UAAU,WAA6B;AACjE,SAAO,SAAS,QAAQ,IAAI,CAAC,WAAW;AACtC,UAAM,MAAM,UAAU,OAAO,EAAE;AAC/B,WAAO,IAAI,eAAe;AAAA,MACxB,SAAS,IAAI;AAAA,MACb,UAAU;AAAA,QACR,GAAG,IAAI;AAAA,QACP,OAAO,OAAO;AAAA,MAChB;AAAA,IACF,CAAC;AAAA,EACH,CAAC;AACH;","names":[]}