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

Genkit AI framework generative AI APIs.

{"version":3,"sources":["../src/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 { Action, GenkitError, defineAction, z } from '@genkit-ai/core';\nimport { Registry } from '@genkit-ai/core/registry';\nimport { Document, DocumentData, DocumentDataSchema } from './document.js';\nimport { EmbedderInfo } from './embedder.js';\n\nexport {\n  Document,\n  DocumentDataSchema,\n  type DocumentData,\n  type MediaPart,\n  type Part,\n  type TextPart,\n} from './document.js';\n\n/**\n * Retriever implementation function signature.\n */\nexport type RetrieverFn<RetrieverOptions extends z.ZodTypeAny> = (\n  query: Document,\n  queryOpts: z.infer<RetrieverOptions>\n) => Promise<RetrieverResponse>;\n\n/**\n * Indexer implementation function signature.\n */\nexport type IndexerFn<IndexerOptions extends z.ZodTypeAny> = (\n  docs: Array<Document>,\n  indexerOpts: z.infer<IndexerOptions>\n) => Promise<void>;\n\nconst RetrieverRequestSchema = z.object({\n  query: DocumentDataSchema,\n  options: z.any().optional(),\n});\n\nconst RetrieverResponseSchema = z.object({\n  documents: z.array(DocumentDataSchema),\n  // TODO: stats, etc.\n});\ntype RetrieverResponse = z.infer<typeof RetrieverResponseSchema>;\n\nconst IndexerRequestSchema = z.object({\n  documents: z.array(DocumentDataSchema),\n  options: z.any().optional(),\n});\n\n/**\n * Zod schema of retriever info metadata.\n */\nexport const RetrieverInfoSchema = z.object({\n  label: z.string().optional(),\n  /** Supported model capabilities. */\n  supports: z\n    .object({\n      /** Model can process media as part of the prompt (multimodal input). */\n      media: z.boolean().optional(),\n    })\n    .optional(),\n});\nexport type RetrieverInfo = z.infer<typeof RetrieverInfoSchema>;\n\n/**\n * A retriever action type.\n */\nexport type RetrieverAction<CustomOptions extends z.ZodTypeAny = z.ZodTypeAny> =\n  Action<typeof RetrieverRequestSchema, typeof RetrieverResponseSchema> & {\n    __configSchema?: CustomOptions;\n  };\n\n/**\n * An indexer action type.\n */\nexport type IndexerAction<IndexerOptions extends z.ZodTypeAny = z.ZodTypeAny> =\n  Action<typeof IndexerRequestSchema, z.ZodVoid> & {\n    __configSchema?: IndexerOptions;\n  };\n\nfunction retrieverWithMetadata<\n  RetrieverOptions extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  retriever: Action<\n    typeof RetrieverRequestSchema,\n    typeof RetrieverResponseSchema\n  >,\n  configSchema?: RetrieverOptions\n): RetrieverAction<RetrieverOptions> {\n  const withMeta = retriever as RetrieverAction<RetrieverOptions>;\n  withMeta.__configSchema = configSchema;\n  return withMeta;\n}\n\nfunction indexerWithMetadata<\n  IndexerOptions extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  indexer: Action<typeof IndexerRequestSchema, z.ZodVoid>,\n  configSchema?: IndexerOptions\n): IndexerAction<IndexerOptions> {\n  const withMeta = indexer as IndexerAction<IndexerOptions>;\n  withMeta.__configSchema = configSchema;\n  return withMeta;\n}\n\n/**\n *  Creates a retriever action for the provided {@link RetrieverFn} implementation.\n */\nexport function defineRetriever<\n  OptionsType extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  registry: Registry,\n  options: {\n    name: string;\n    configSchema?: OptionsType;\n    info?: RetrieverInfo;\n  },\n  runner: RetrieverFn<OptionsType>\n) {\n  const retriever = defineAction(\n    registry,\n    {\n      actionType: 'retriever',\n      name: options.name,\n      inputSchema: options.configSchema\n        ? RetrieverRequestSchema.extend({\n            options: options.configSchema.optional(),\n          })\n        : RetrieverRequestSchema,\n      outputSchema: RetrieverResponseSchema,\n      metadata: {\n        type: 'retriever',\n        info: options.info,\n      },\n    },\n    (i) => runner(new Document(i.query), i.options)\n  );\n  const rwm = retrieverWithMetadata(\n    retriever as Action<\n      typeof RetrieverRequestSchema,\n      typeof RetrieverResponseSchema\n    >,\n    options.configSchema\n  );\n  return rwm;\n}\n\n/**\n *  Creates an indexer action for the provided {@link IndexerFn} implementation.\n */\nexport function defineIndexer<IndexerOptions extends z.ZodTypeAny>(\n  registry: Registry,\n  options: {\n    name: string;\n    embedderInfo?: EmbedderInfo;\n    configSchema?: IndexerOptions;\n  },\n  runner: IndexerFn<IndexerOptions>\n) {\n  const indexer = defineAction(\n    registry,\n    {\n      actionType: 'indexer',\n      name: options.name,\n      inputSchema: options.configSchema\n        ? IndexerRequestSchema.extend({\n            options: options.configSchema.optional(),\n          })\n        : IndexerRequestSchema,\n      outputSchema: z.void(),\n      metadata: {\n        type: 'indexer',\n        embedderInfo: options.embedderInfo,\n      },\n    },\n    (i) =>\n      runner(\n        i.documents.map((dd) => new Document(dd)),\n        i.options\n      )\n  );\n  const iwm = indexerWithMetadata(\n    indexer as Action<typeof IndexerRequestSchema, z.ZodVoid>,\n    options.configSchema\n  );\n  return iwm;\n}\n\nexport interface RetrieverParams<\n  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,\n> {\n  retriever: RetrieverArgument<CustomOptions>;\n  query: string | DocumentData;\n  options?: z.infer<CustomOptions>;\n}\n\n/**\n * A type that can be used to pass a retriever as an argument, either using a reference or an action.\n */\nexport type RetrieverArgument<\n  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,\n> = RetrieverAction<CustomOptions> | RetrieverReference<CustomOptions> | string;\n\n/**\n * Retrieves documents from a {@link RetrieverArgument} based on the provided query.\n */\nexport async function retrieve<CustomOptions extends z.ZodTypeAny>(\n  registry: Registry,\n  params: RetrieverParams<CustomOptions>\n): Promise<Array<Document>> {\n  let retriever: RetrieverAction<CustomOptions>;\n  if (typeof params.retriever === 'string') {\n    retriever = await registry.lookupAction(`/retriever/${params.retriever}`);\n  } else if (Object.hasOwnProperty.call(params.retriever, 'info')) {\n    retriever = await registry.lookupAction(\n      `/retriever/${params.retriever.name}`\n    );\n  } else {\n    retriever = params.retriever as RetrieverAction<CustomOptions>;\n  }\n  if (!retriever) {\n    throw new Error('Unable to resolve the retriever');\n  }\n  const response = await retriever({\n    query:\n      typeof params.query === 'string'\n        ? Document.fromText(params.query)\n        : params.query,\n    options: params.options,\n  });\n\n  return response.documents.map((d) => new Document(d));\n}\n\n/**\n * A type that can be used to pass an indexer as an argument, either using a reference or an action.\n */\nexport type IndexerArgument<CustomOptions extends z.ZodTypeAny = z.ZodTypeAny> =\n  IndexerReference<CustomOptions> | IndexerAction<CustomOptions> | string;\n\n/**\n * Options passed to the index function.\n */\nexport interface IndexerParams<\n  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,\n> {\n  indexer: IndexerArgument<CustomOptions>;\n  documents: Array<DocumentData>;\n  options?: z.infer<CustomOptions>;\n}\n\n/**\n * Indexes documents using a {@link IndexerArgument}.\n */\nexport async function index<CustomOptions extends z.ZodTypeAny>(\n  registry: Registry,\n  params: IndexerParams<CustomOptions>\n): Promise<void> {\n  let indexer: IndexerAction<CustomOptions>;\n  if (typeof params.indexer === 'string') {\n    indexer = await registry.lookupAction(`/indexer/${params.indexer}`);\n  } else if (Object.hasOwnProperty.call(params.indexer, 'info')) {\n    indexer = await registry.lookupAction(`/indexer/${params.indexer.name}`);\n  } else {\n    indexer = params.indexer as IndexerAction<CustomOptions>;\n  }\n  if (!indexer) {\n    throw new Error('Unable to utilize the provided indexer');\n  }\n  return await indexer({\n    documents: params.documents,\n    options: params.options,\n  });\n}\n\n/**\n * Zod schema of common retriever options.\n */\nexport const CommonRetrieverOptionsSchema = z.object({\n  k: z.number().describe('Number of documents to retrieve').optional(),\n});\n\n/**\n * A retriver reference object.\n */\nexport interface RetrieverReference<CustomOptions extends z.ZodTypeAny> {\n  name: string;\n  configSchema?: CustomOptions;\n  info?: RetrieverInfo;\n}\n\n/**\n * Helper method to configure a {@link RetrieverReference} to a plugin.\n */\nexport function retrieverRef<\n  CustomOptionsSchema extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  options: RetrieverReference<CustomOptionsSchema>\n): RetrieverReference<CustomOptionsSchema> {\n  return { ...options };\n}\n\n// Reuse the same schema for both indexers and retrievers -- for now.\nexport const IndexerInfoSchema = RetrieverInfoSchema;\n\n/**\n * Indexer metadata.\n */\nexport type IndexerInfo = z.infer<typeof IndexerInfoSchema>;\n\nexport interface IndexerReference<CustomOptions extends z.ZodTypeAny> {\n  name: string;\n  configSchema?: CustomOptions;\n  info?: IndexerInfo;\n}\n\n/**\n * Helper method to configure a {@link IndexerReference} to a plugin.\n */\nexport function indexerRef<\n  CustomOptionsSchema extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  options: IndexerReference<CustomOptionsSchema>\n): IndexerReference<CustomOptionsSchema> {\n  return { ...options };\n}\n\nfunction itemToDocument<R>(\n  item: any,\n  options: SimpleRetrieverOptions\n): Document {\n  if (!item)\n    throw new GenkitError({\n      status: 'INVALID_ARGUMENT',\n      message: `Items returned from simple retriever must be non-null.`,\n    });\n  if (typeof item === 'string') return Document.fromText(item);\n  if (typeof options.content === 'function') {\n    const transformed = options.content(item);\n    return typeof transformed === 'string'\n      ? Document.fromText(transformed)\n      : new Document({ content: transformed });\n  }\n  if (typeof options.content === 'string' && typeof item === 'object')\n    return Document.fromText(item[options.content]);\n  throw new GenkitError({\n    status: 'INVALID_ARGUMENT',\n    message: `Cannot convert item to document without content option. Item: ${JSON.stringify(item)}`,\n  });\n}\n\nfunction itemToMetadata(\n  item: any,\n  options: SimpleRetrieverOptions\n): Document['metadata'] {\n  if (typeof item === 'string') return undefined;\n  if (Array.isArray(options.metadata) && typeof item === 'object') {\n    const out: Record<string, any> = {};\n    options.metadata.forEach((key) => (out[key] = item[key]));\n    return out;\n  }\n  if (typeof options.metadata === 'function') return options.metadata(item);\n  if (!options.metadata && typeof item === 'object') {\n    const out = { ...item };\n    if (typeof options.content === 'string') delete out[options.content];\n    return out;\n  }\n  throw new GenkitError({\n    status: 'INVALID_ARGUMENT',\n    message: `Unable to extract metadata from item with supplied options. Item: ${JSON.stringify(item)}`,\n  });\n}\n\n/**\n * Simple retriever options.\n */\nexport interface SimpleRetrieverOptions<\n  C extends z.ZodTypeAny = z.ZodTypeAny,\n  R = any,\n> {\n  /** The name of the retriever you're creating. */\n  name: string;\n  /** A Zod schema containing any configuration info available beyond the query. */\n  configSchema?: C;\n  /**\n   * Specifies how to extract content from the returned items.\n   *\n   * - If a string, specifies the key of the returned item to extract as content.\n   * - If a function, allows you to extract content as text or a document part.\n   **/\n  content?: string | ((item: R) => Document['content'] | string);\n  /**\n   * Specifies how to extract metadata from the returned items.\n   *\n   * - If an array of strings, specifies list of keys to extract from returned objects.\n   * - If a function, allows you to use custom behavior to extract metadata from returned items.\n   */\n  metadata?: string[] | ((item: R) => Document['metadata']);\n}\n\n/**\n * defineSimpleRetriever makes it easy to map existing data into documents that\n * can be used for prompt augmentation.\n *\n * @param options Configuration options for the retriever.\n * @param handler A function that queries a datastore and returns items from which to extract documents.\n * @returns A Genkit retriever.\n */\nexport function defineSimpleRetriever<\n  C extends z.ZodTypeAny = z.ZodTypeAny,\n  R = any,\n>(\n  registry: Registry,\n  options: SimpleRetrieverOptions<C, R>,\n  handler: (query: Document, config: z.infer<C>) => Promise<R[]>\n) {\n  return defineRetriever(\n    registry,\n    {\n      name: options.name,\n      configSchema: options.configSchema,\n    },\n    async (query, config) => {\n      const result = await handler(query, config);\n      return {\n        documents: result.map((item) => {\n          const doc = itemToDocument(item, options);\n          if (typeof item !== 'string')\n            doc.metadata = itemToMetadata(item, options);\n          return doc;\n        }),\n      };\n    }\n  );\n}\n"],"mappings":"AAgBA,SAAiB,aAAa,cAAc,SAAS;AAErD,SAAS,UAAwB,0BAA0B;AAG3D;AAAA,EACE,YAAAA;AAAA,EACA,sBAAAC;AAAA,OAKK;AAkBP,MAAM,yBAAyB,EAAE,OAAO;AAAA,EACtC,OAAO;AAAA,EACP,SAAS,EAAE,IAAI,EAAE,SAAS;AAC5B,CAAC;AAED,MAAM,0BAA0B,EAAE,OAAO;AAAA,EACvC,WAAW,EAAE,MAAM,kBAAkB;AAAA;AAEvC,CAAC;AAGD,MAAM,uBAAuB,EAAE,OAAO;AAAA,EACpC,WAAW,EAAE,MAAM,kBAAkB;AAAA,EACrC,SAAS,EAAE,IAAI,EAAE,SAAS;AAC5B,CAAC;AAKM,MAAM,sBAAsB,EAAE,OAAO;AAAA,EAC1C,OAAO,EAAE,OAAO,EAAE,SAAS;AAAA;AAAA,EAE3B,UAAU,EACP,OAAO;AAAA;AAAA,IAEN,OAAO,EAAE,QAAQ,EAAE,SAAS;AAAA,EAC9B,CAAC,EACA,SAAS;AACd,CAAC;AAmBD,SAAS,sBAGP,WAIA,cACmC;AACnC,QAAM,WAAW;AACjB,WAAS,iBAAiB;AAC1B,SAAO;AACT;AAEA,SAAS,oBAGP,SACA,cAC+B;AAC/B,QAAM,WAAW;AACjB,WAAS,iBAAiB;AAC1B,SAAO;AACT;AAKO,SAAS,gBAGd,UACA,SAKA,QACA;AACA,QAAM,YAAY;AAAA,IAChB;AAAA,IACA;AAAA,MACE,YAAY;AAAA,MACZ,MAAM,QAAQ;AAAA,MACd,aAAa,QAAQ,eACjB,uBAAuB,OAAO;AAAA,QAC5B,SAAS,QAAQ,aAAa,SAAS;AAAA,MACzC,CAAC,IACD;AAAA,MACJ,cAAc;AAAA,MACd,UAAU;AAAA,QACR,MAAM;AAAA,QACN,MAAM,QAAQ;AAAA,MAChB;AAAA,IACF;AAAA,IACA,CAAC,MAAM,OAAO,IAAI,SAAS,EAAE,KAAK,GAAG,EAAE,OAAO;AAAA,EAChD;AACA,QAAM,MAAM;AAAA,IACV;AAAA,IAIA,QAAQ;AAAA,EACV;AACA,SAAO;AACT;AAKO,SAAS,cACd,UACA,SAKA,QACA;AACA,QAAM,UAAU;AAAA,IACd;AAAA,IACA;AAAA,MACE,YAAY;AAAA,MACZ,MAAM,QAAQ;AAAA,MACd,aAAa,QAAQ,eACjB,qBAAqB,OAAO;AAAA,QAC1B,SAAS,QAAQ,aAAa,SAAS;AAAA,MACzC,CAAC,IACD;AAAA,MACJ,cAAc,EAAE,KAAK;AAAA,MACrB,UAAU;AAAA,QACR,MAAM;AAAA,QACN,cAAc,QAAQ;AAAA,MACxB;AAAA,IACF;AAAA,IACA,CAAC,MACC;AAAA,MACE,EAAE,UAAU,IAAI,CAAC,OAAO,IAAI,SAAS,EAAE,CAAC;AAAA,MACxC,EAAE;AAAA,IACJ;AAAA,EACJ;AACA,QAAM,MAAM;AAAA,IACV;AAAA,IACA,QAAQ;AAAA,EACV;AACA,SAAO;AACT;AAoBA,eAAsB,SACpB,UACA,QAC0B;AAC1B,MAAI;AACJ,MAAI,OAAO,OAAO,cAAc,UAAU;AACxC,gBAAY,MAAM,SAAS,aAAa,cAAc,OAAO,SAAS,EAAE;AAAA,EAC1E,WAAW,OAAO,eAAe,KAAK,OAAO,WAAW,MAAM,GAAG;AAC/D,gBAAY,MAAM,SAAS;AAAA,MACzB,cAAc,OAAO,UAAU,IAAI;AAAA,IACrC;AAAA,EACF,OAAO;AACL,gBAAY,OAAO;AAAA,EACrB;AACA,MAAI,CAAC,WAAW;AACd,UAAM,IAAI,MAAM,iCAAiC;AAAA,EACnD;AACA,QAAM,WAAW,MAAM,UAAU;AAAA,IAC/B,OACE,OAAO,OAAO,UAAU,WACpB,SAAS,SAAS,OAAO,KAAK,IAC9B,OAAO;AAAA,IACb,SAAS,OAAO;AAAA,EAClB,CAAC;AAED,SAAO,SAAS,UAAU,IAAI,CAAC,MAAM,IAAI,SAAS,CAAC,CAAC;AACtD;AAsBA,eAAsB,MACpB,UACA,QACe;AACf,MAAI;AACJ,MAAI,OAAO,OAAO,YAAY,UAAU;AACtC,cAAU,MAAM,SAAS,aAAa,YAAY,OAAO,OAAO,EAAE;AAAA,EACpE,WAAW,OAAO,eAAe,KAAK,OAAO,SAAS,MAAM,GAAG;AAC7D,cAAU,MAAM,SAAS,aAAa,YAAY,OAAO,QAAQ,IAAI,EAAE;AAAA,EACzE,OAAO;AACL,cAAU,OAAO;AAAA,EACnB;AACA,MAAI,CAAC,SAAS;AACZ,UAAM,IAAI,MAAM,wCAAwC;AAAA,EAC1D;AACA,SAAO,MAAM,QAAQ;AAAA,IACnB,WAAW,OAAO;AAAA,IAClB,SAAS,OAAO;AAAA,EAClB,CAAC;AACH;AAKO,MAAM,+BAA+B,EAAE,OAAO;AAAA,EACnD,GAAG,EAAE,OAAO,EAAE,SAAS,iCAAiC,EAAE,SAAS;AACrE,CAAC;AAcM,SAAS,aAGd,SACyC;AACzC,SAAO,EAAE,GAAG,QAAQ;AACtB;AAGO,MAAM,oBAAoB;AAgB1B,SAAS,WAGd,SACuC;AACvC,SAAO,EAAE,GAAG,QAAQ;AACtB;AAEA,SAAS,eACP,MACA,SACU;AACV,MAAI,CAAC;AACH,UAAM,IAAI,YAAY;AAAA,MACpB,QAAQ;AAAA,MACR,SAAS;AAAA,IACX,CAAC;AACH,MAAI,OAAO,SAAS,SAAU,QAAO,SAAS,SAAS,IAAI;AAC3D,MAAI,OAAO,QAAQ,YAAY,YAAY;AACzC,UAAM,cAAc,QAAQ,QAAQ,IAAI;AACxC,WAAO,OAAO,gBAAgB,WAC1B,SAAS,SAAS,WAAW,IAC7B,IAAI,SAAS,EAAE,SAAS,YAAY,CAAC;AAAA,EAC3C;AACA,MAAI,OAAO,QAAQ,YAAY,YAAY,OAAO,SAAS;AACzD,WAAO,SAAS,SAAS,KAAK,QAAQ,OAAO,CAAC;AAChD,QAAM,IAAI,YAAY;AAAA,IACpB,QAAQ;AAAA,IACR,SAAS,iEAAiE,KAAK,UAAU,IAAI,CAAC;AAAA,EAChG,CAAC;AACH;AAEA,SAAS,eACP,MACA,SACsB;AACtB,MAAI,OAAO,SAAS,SAAU,QAAO;AACrC,MAAI,MAAM,QAAQ,QAAQ,QAAQ,KAAK,OAAO,SAAS,UAAU;AAC/D,UAAM,MAA2B,CAAC;AAClC,YAAQ,SAAS,QAAQ,CAAC,QAAS,IAAI,GAAG,IAAI,KAAK,GAAG,CAAE;AACxD,WAAO;AAAA,EACT;AACA,MAAI,OAAO,QAAQ,aAAa,WAAY,QAAO,QAAQ,SAAS,IAAI;AACxE,MAAI,CAAC,QAAQ,YAAY,OAAO,SAAS,UAAU;AACjD,UAAM,MAAM,EAAE,GAAG,KAAK;AACtB,QAAI,OAAO,QAAQ,YAAY,SAAU,QAAO,IAAI,QAAQ,OAAO;AACnE,WAAO;AAAA,EACT;AACA,QAAM,IAAI,YAAY;AAAA,IACpB,QAAQ;AAAA,IACR,SAAS,qEAAqE,KAAK,UAAU,IAAI,CAAC;AAAA,EACpG,CAAC;AACH;AAqCO,SAAS,sBAId,UACA,SACA,SACA;AACA,SAAO;AAAA,IACL;AAAA,IACA;AAAA,MACE,MAAM,QAAQ;AAAA,MACd,cAAc,QAAQ;AAAA,IACxB;AAAA,IACA,OAAO,OAAO,WAAW;AACvB,YAAM,SAAS,MAAM,QAAQ,OAAO,MAAM;AAC1C,aAAO;AAAA,QACL,WAAW,OAAO,IAAI,CAAC,SAAS;AAC9B,gBAAM,MAAM,eAAe,MAAM,OAAO;AACxC,cAAI,OAAO,SAAS;AAClB,gBAAI,WAAW,eAAe,MAAM,OAAO;AAC7C,iBAAO;AAAA,QACT,CAAC;AAAA,MACH;AAAA,IACF;AAAA,EACF;AACF;","names":["Document","DocumentDataSchema"]}