retriever.ts

/**
 * 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.
 */

import { Action, GenkitError, defineAction, z } from '@genkit-ai/core';
import { Registry } from '@genkit-ai/core/registry';
import { Document, DocumentData, DocumentDataSchema } from './document.js';
import { EmbedderInfo } from './embedder.js';

export {
  Document,
  DocumentDataSchema,
  type DocumentData,
  type MediaPart,
  type Part,
  type TextPart,
} from './document.js';

/**
 * Retriever implementation function signature.
 */
export type RetrieverFn<RetrieverOptions extends z.ZodTypeAny> = (
  query: Document,
  queryOpts: z.infer<RetrieverOptions>
) => Promise<RetrieverResponse>;

/**
 * Indexer implementation function signature.
 */
export type IndexerFn<IndexerOptions extends z.ZodTypeAny> = (
  docs: Array<Document>,
  indexerOpts: z.infer<IndexerOptions>
) => Promise<void>;

const RetrieverRequestSchema = z.object({
  query: DocumentDataSchema,
  options: z.any().optional(),
});

const RetrieverResponseSchema = z.object({
  documents: z.array(DocumentDataSchema),
  // TODO: stats, etc.
});
type RetrieverResponse = z.infer<typeof RetrieverResponseSchema>;

const IndexerRequestSchema = z.object({
  documents: z.array(DocumentDataSchema),
  options: z.any().optional(),
});

/**
 * Zod schema of retriever info metadata.
 */
export const RetrieverInfoSchema = z.object({
  label: z.string().optional(),
  /** Supported model capabilities. */
  supports: z
    .object({
      /** Model can process media as part of the prompt (multimodal input). */
      media: z.boolean().optional(),
    })
    .optional(),
});
export type RetrieverInfo = z.infer<typeof RetrieverInfoSchema>;

/**
 * A retriever action type.
 */
export type RetrieverAction<CustomOptions extends z.ZodTypeAny = z.ZodTypeAny> =
  Action<typeof RetrieverRequestSchema, typeof RetrieverResponseSchema> & {
    __configSchema?: CustomOptions;
  };

/**
 * An indexer action type.
 */
export type IndexerAction<IndexerOptions extends z.ZodTypeAny = z.ZodTypeAny> =
  Action<typeof IndexerRequestSchema, z.ZodVoid> & {
    __configSchema?: IndexerOptions;
  };

function retrieverWithMetadata<
  RetrieverOptions extends z.ZodTypeAny = z.ZodTypeAny,
>(
  retriever: Action<
    typeof RetrieverRequestSchema,
    typeof RetrieverResponseSchema
  >,
  configSchema?: RetrieverOptions
): RetrieverAction<RetrieverOptions> {
  const withMeta = retriever as RetrieverAction<RetrieverOptions>;
  withMeta.__configSchema = configSchema;
  return withMeta;
}

function indexerWithMetadata<
  IndexerOptions extends z.ZodTypeAny = z.ZodTypeAny,
>(
  indexer: Action<typeof IndexerRequestSchema, z.ZodVoid>,
  configSchema?: IndexerOptions
): IndexerAction<IndexerOptions> {
  const withMeta = indexer as IndexerAction<IndexerOptions>;
  withMeta.__configSchema = configSchema;
  return withMeta;
}

/**
 *  Creates a retriever action for the provided {@link RetrieverFn} implementation.
 */
export function defineRetriever<
  OptionsType extends z.ZodTypeAny = z.ZodTypeAny,
>(
  registry: Registry,
  options: {
    name: string;
    configSchema?: OptionsType;
    info?: RetrieverInfo;
  },
  runner: RetrieverFn<OptionsType>
) {
  const retriever = defineAction(
    registry,
    {
      actionType: 'retriever',
      name: options.name,
      inputSchema: options.configSchema
        ? RetrieverRequestSchema.extend({
            options: options.configSchema.optional(),
          })
        : RetrieverRequestSchema,
      outputSchema: RetrieverResponseSchema,
      metadata: {
        type: 'retriever',
        info: options.info,
      },
    },
    (i) => runner(new Document(i.query), i.options)
  );
  const rwm = retrieverWithMetadata(
    retriever as Action<
      typeof RetrieverRequestSchema,
      typeof RetrieverResponseSchema
    >,
    options.configSchema
  );
  return rwm;
}

/**
 *  Creates an indexer action for the provided {@link IndexerFn} implementation.
 */
export function defineIndexer<IndexerOptions extends z.ZodTypeAny>(
  registry: Registry,
  options: {
    name: string;
    embedderInfo?: EmbedderInfo;
    configSchema?: IndexerOptions;
  },
  runner: IndexerFn<IndexerOptions>
) {
  const indexer = defineAction(
    registry,
    {
      actionType: 'indexer',
      name: options.name,
      inputSchema: options.configSchema
        ? IndexerRequestSchema.extend({
            options: options.configSchema.optional(),
          })
        : IndexerRequestSchema,
      outputSchema: z.void(),
      metadata: {
        type: 'indexer',
        embedderInfo: options.embedderInfo,
      },
    },
    (i) =>
      runner(
        i.documents.map((dd) => new Document(dd)),
        i.options
      )
  );
  const iwm = indexerWithMetadata(
    indexer as Action<typeof IndexerRequestSchema, z.ZodVoid>,
    options.configSchema
  );
  return iwm;
}

export interface RetrieverParams<
  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
> {
  retriever: RetrieverArgument<CustomOptions>;
  query: string | DocumentData;
  options?: z.infer<CustomOptions>;
}

/**
 * A type that can be used to pass a retriever as an argument, either using a reference or an action.
 */
export type RetrieverArgument<
  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
> = RetrieverAction<CustomOptions> | RetrieverReference<CustomOptions> | string;

/**
 * Retrieves documents from a {@link RetrieverArgument} based on the provided query.
 */
export async function retrieve<CustomOptions extends z.ZodTypeAny>(
  registry: Registry,
  params: RetrieverParams<CustomOptions>
): Promise<Array<Document>> {
  let retriever: RetrieverAction<CustomOptions>;
  if (typeof params.retriever === 'string') {
    retriever = await registry.lookupAction(`/retriever/${params.retriever}`);
  } else if (Object.hasOwnProperty.call(params.retriever, 'info')) {
    retriever = await registry.lookupAction(
      `/retriever/${params.retriever.name}`
    );
  } else {
    retriever = params.retriever as RetrieverAction<CustomOptions>;
  }
  if (!retriever) {
    throw new Error('Unable to resolve the retriever');
  }
  const response = await retriever({
    query:
      typeof params.query === 'string'
        ? Document.fromText(params.query)
        : params.query,
    options: params.options,
  });

  return response.documents.map((d) => new Document(d));
}

/**
 * A type that can be used to pass an indexer as an argument, either using a reference or an action.
 */
export type IndexerArgument<CustomOptions extends z.ZodTypeAny = z.ZodTypeAny> =
  IndexerReference<CustomOptions> | IndexerAction<CustomOptions> | string;

/**
 * Options passed to the index function.
 */
export interface IndexerParams<
  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
> {
  indexer: IndexerArgument<CustomOptions>;
  documents: Array<DocumentData>;
  options?: z.infer<CustomOptions>;
}

/**
 * Indexes documents using a {@link IndexerArgument}.
 */
export async function index<CustomOptions extends z.ZodTypeAny>(
  registry: Registry,
  params: IndexerParams<CustomOptions>
): Promise<void> {
  let indexer: IndexerAction<CustomOptions>;
  if (typeof params.indexer === 'string') {
    indexer = await registry.lookupAction(`/indexer/${params.indexer}`);
  } else if (Object.hasOwnProperty.call(params.indexer, 'info')) {
    indexer = await registry.lookupAction(`/indexer/${params.indexer.name}`);
  } else {
    indexer = params.indexer as IndexerAction<CustomOptions>;
  }
  if (!indexer) {
    throw new Error('Unable to utilize the provided indexer');
  }
  return await indexer({
    documents: params.documents,
    options: params.options,
  });
}

/**
 * Zod schema of common retriever options.
 */
export const CommonRetrieverOptionsSchema = z.object({
  k: z.number().describe('Number of documents to retrieve').optional(),
});

/**
 * A retriver reference object.
 */
export interface RetrieverReference<CustomOptions extends z.ZodTypeAny> {
  name: string;
  configSchema?: CustomOptions;
  info?: RetrieverInfo;
}

/**
 * Helper method to configure a {@link RetrieverReference} to a plugin.
 */
export function retrieverRef<
  CustomOptionsSchema extends z.ZodTypeAny = z.ZodTypeAny,
>(
  options: RetrieverReference<CustomOptionsSchema>
): RetrieverReference<CustomOptionsSchema> {
  return { ...options };
}

// Reuse the same schema for both indexers and retrievers -- for now.
export const IndexerInfoSchema = RetrieverInfoSchema;

/**
 * Indexer metadata.
 */
export type IndexerInfo = z.infer<typeof IndexerInfoSchema>;

export interface IndexerReference<CustomOptions extends z.ZodTypeAny> {
  name: string;
  configSchema?: CustomOptions;
  info?: IndexerInfo;
}

/**
 * Helper method to configure a {@link IndexerReference} to a plugin.
 */
export function indexerRef<
  CustomOptionsSchema extends z.ZodTypeAny = z.ZodTypeAny,
>(
  options: IndexerReference<CustomOptionsSchema>
): IndexerReference<CustomOptionsSchema> {
  return { ...options };
}

function itemToDocument<R>(
  item: any,
  options: SimpleRetrieverOptions
): Document {
  if (!item)
    throw new GenkitError({
      status: 'INVALID_ARGUMENT',
      message: `Items returned from simple retriever must be non-null.`,
    });
  if (typeof item === 'string') return Document.fromText(item);
  if (typeof options.content === 'function') {
    const transformed = options.content(item);
    return typeof transformed === 'string'
      ? Document.fromText(transformed)
      : new Document({ content: transformed });
  }
  if (typeof options.content === 'string' && typeof item === 'object')
    return Document.fromText(item[options.content]);
  throw new GenkitError({
    status: 'INVALID_ARGUMENT',
    message: `Cannot convert item to document without content option. Item: ${JSON.stringify(item)}`,
  });
}

function itemToMetadata(
  item: any,
  options: SimpleRetrieverOptions
): Document['metadata'] {
  if (typeof item === 'string') return undefined;
  if (Array.isArray(options.metadata) && typeof item === 'object') {
    const out: Record<string, any> = {};
    options.metadata.forEach((key) => (out[key] = item[key]));
    return out;
  }
  if (typeof options.metadata === 'function') return options.metadata(item);
  if (!options.metadata && typeof item === 'object') {
    const out = { ...item };
    if (typeof options.content === 'string') delete out[options.content];
    return out;
  }
  throw new GenkitError({
    status: 'INVALID_ARGUMENT',
    message: `Unable to extract metadata from item with supplied options. Item: ${JSON.stringify(item)}`,
  });
}

/**
 * Simple retriever options.
 */
export interface SimpleRetrieverOptions<
  C extends z.ZodTypeAny = z.ZodTypeAny,
  R = any,
> {
  /** The name of the retriever you're creating. */
  name: string;
  /** A Zod schema containing any configuration info available beyond the query. */
  configSchema?: C;
  /**
   * Specifies how to extract content from the returned items.
   *
   * - If a string, specifies the key of the returned item to extract as content.
   * - If a function, allows you to extract content as text or a document part.
   **/
  content?: string | ((item: R) => Document['content'] | string);
  /**
   * Specifies how to extract metadata from the returned items.
   *
   * - If an array of strings, specifies list of keys to extract from returned objects.
   * - If a function, allows you to use custom behavior to extract metadata from returned items.
   */
  metadata?: string[] | ((item: R) => Document['metadata']);
}

/**
 * defineSimpleRetriever makes it easy to map existing data into documents that
 * can be used for prompt augmentation.
 *
 * @param options Configuration options for the retriever.
 * @param handler A function that queries a datastore and returns items from which to extract documents.
 * @returns A Genkit retriever.
 */
export function defineSimpleRetriever<
  C extends z.ZodTypeAny = z.ZodTypeAny,
  R = any,
>(
  registry: Registry,
  options: SimpleRetrieverOptions<C, R>,
  handler: (query: Document, config: z.infer<C>) => Promise<R[]>
) {
  return defineRetriever(
    registry,
    {
      name: options.name,
      configSchema: options.configSchema,
    },
    async (query, config) => {
      const result = await handler(query, config);
      return {
        documents: result.map((item) => {
          const doc = itemToDocument(item, options);
          if (typeof item !== 'string')
            doc.metadata = itemToMetadata(item, options);
          return doc;
        }),
      };
    }
  );
}