@genkit-ai/core/lib/registry.mjs.map

Genkit AI framework core libraries.

{"version":3,"sources":["../src/registry.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 { Dotprompt } from 'dotprompt';\nimport { AsyncLocalStorage } from 'node:async_hooks';\nimport * as z from 'zod';\nimport { Action, runOutsideActionRuntimeContext } from './action.js';\nimport { GenkitError } from './error.js';\nimport { logger } from './logging.js';\nimport { PluginProvider } from './plugin.js';\nimport { JSONSchema, toJsonSchema } from './schema.js';\n\nexport type AsyncProvider<T> = () => Promise<T>;\n\n/**\n * Type of a runnable action.\n */\nexport type ActionType =\n  | 'custom'\n  | 'embedder'\n  | 'evaluator'\n  | 'executable-prompt'\n  | 'flow'\n  | 'indexer'\n  | 'model'\n  | 'prompt'\n  | 'reranker'\n  | 'retriever'\n  | 'tool'\n  | 'util';\n\n/**\n * A schema is either a Zod schema or a JSON schema.\n */\nexport interface Schema {\n  schema?: z.ZodTypeAny;\n  jsonSchema?: JSONSchema;\n}\n\nfunction parsePluginName(registryKey: string) {\n  const tokens = registryKey.split('/');\n  if (tokens.length >= 4) {\n    return tokens[2];\n  }\n  return undefined;\n}\n\ntype ActionsRecord = Record<string, Action<z.ZodTypeAny, z.ZodTypeAny>>;\n\n/**\n * The registry is used to store and lookup actions, trace stores, flow state stores, plugins, and schemas.\n */\nexport class Registry {\n  private actionsById: Record<\n    string,\n    | Action<z.ZodTypeAny, z.ZodTypeAny>\n    | PromiseLike<Action<z.ZodTypeAny, z.ZodTypeAny>>\n  > = {};\n  private pluginsByName: Record<string, PluginProvider> = {};\n  private schemasByName: Record<string, Schema> = {};\n  private valueByTypeAndName: Record<string, Record<string, any>> = {};\n  private allPluginsInitialized = false;\n  public apiStability: 'stable' | 'beta' = 'stable';\n\n  readonly asyncStore = new AsyncStore();\n  readonly dotprompt = new Dotprompt({\n    schemaResolver: async (name) => {\n      const resolvedSchema = await this.lookupSchema(name);\n      if (!resolvedSchema) {\n        throw new GenkitError({\n          message: `Schema '${name}' not found`,\n          status: 'NOT_FOUND',\n        });\n      }\n      return toJsonSchema(resolvedSchema);\n    },\n  });\n\n  constructor(public parent?: Registry) {}\n\n  /**\n   * Creates a new registry overlaid onto the provided registry.\n   * @param parent The parent registry.\n   * @returns The new overlaid registry.\n   */\n  static withParent(parent: Registry) {\n    return new Registry(parent);\n  }\n\n  /**\n   * Looks up an action in the registry.\n   * @param key The key of the action to lookup.\n   * @returns The action.\n   */\n  async lookupAction<\n    I extends z.ZodTypeAny,\n    O extends z.ZodTypeAny,\n    R extends Action<I, O>,\n  >(key: string): Promise<R> {\n    // If we don't see the key in the registry we try to initialize the plugin first.\n    const pluginName = parsePluginName(key);\n    if (!this.actionsById[key] && pluginName) {\n      await this.initializePlugin(pluginName);\n    }\n    return (\n      ((await this.actionsById[key]) as R) || this.parent?.lookupAction(key)\n    );\n  }\n\n  /**\n   * Registers an action in the registry.\n   * @param type The type of the action to register.\n   * @param action The action to register.\n   */\n  registerAction<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(\n    type: ActionType,\n    action: Action<I, O>\n  ) {\n    const key = `/${type}/${action.__action.name}`;\n    logger.debug(`registering ${key}`);\n    if (this.actionsById.hasOwnProperty(key)) {\n      // TODO: Make this an error!\n      logger.warn(\n        `WARNING: ${key} already has an entry in the registry. Overwriting.`\n      );\n    }\n    this.actionsById[key] = action;\n  }\n\n  /**\n   * Registers an action promise in the registry.\n   */\n  registerActionAsync<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(\n    type: ActionType,\n    name: string,\n    action: PromiseLike<Action<I, O>>\n  ) {\n    const key = `/${type}/${name}`;\n    logger.debug(`registering ${key} (async)`);\n    if (this.actionsById.hasOwnProperty(key)) {\n      // TODO: Make this an error!\n      logger.warn(\n        `WARNING: ${key} already has an entry in the registry. Overwriting.`\n      );\n    }\n    this.actionsById[key] = action;\n  }\n\n  /**\n   * Returns all actions in the registry.\n   * @returns All actions in the registry.\n   */\n  async listActions(): Promise<ActionsRecord> {\n    await this.initializeAllPlugins();\n    const actions: Record<string, Action<z.ZodTypeAny, z.ZodTypeAny>> = {};\n    await Promise.all(\n      Object.entries(this.actionsById).map(async ([key, action]) => {\n        actions[key] = await action;\n      })\n    );\n    return {\n      ...(await this.parent?.listActions()),\n      ...actions,\n    };\n  }\n\n  /**\n   * Initializes all plugins in the registry.\n   */\n  async initializeAllPlugins() {\n    if (this.allPluginsInitialized) {\n      return;\n    }\n    for (const pluginName of Object.keys(this.pluginsByName)) {\n      await this.initializePlugin(pluginName);\n    }\n    this.allPluginsInitialized = true;\n  }\n\n  /**\n   * Registers a plugin provider. This plugin must be initialized before it can be used by calling {@link initializePlugin} or {@link initializeAllPlugins}.\n   * @param name The name of the plugin to register.\n   * @param provider The plugin provider.\n   */\n  registerPluginProvider(name: string, provider: PluginProvider) {\n    if (this.pluginsByName[name]) {\n      throw new Error(`Plugin ${name} already registered`);\n    }\n    this.allPluginsInitialized = false;\n    let cached;\n    let isInitialized = false;\n    this.pluginsByName[name] = {\n      name: provider.name,\n      initializer: () => {\n        if (!isInitialized) {\n          cached = provider.initializer();\n          isInitialized = true;\n        }\n        return cached;\n      },\n    };\n  }\n\n  /**\n   * Looks up a plugin.\n   * @param name The name of the plugin to lookup.\n   * @returns The plugin provider.\n   */\n  lookupPlugin(name: string): PluginProvider | undefined {\n    return this.pluginsByName[name] || this.parent?.lookupPlugin(name);\n  }\n\n  /**\n   * Initializes a plugin already registered with {@link registerPluginProvider}.\n   * @param name The name of the plugin to initialize.\n   * @returns The plugin.\n   */\n  async initializePlugin(name: string) {\n    if (this.pluginsByName[name]) {\n      return await runOutsideActionRuntimeContext(this, () =>\n        this.pluginsByName[name].initializer()\n      );\n    }\n  }\n\n  /**\n   * Registers a schema.\n   * @param name The name of the schema to register.\n   * @param data The schema to register (either a Zod schema or a JSON schema).\n   */\n  registerSchema(name: string, data: Schema) {\n    if (this.schemasByName[name]) {\n      throw new Error(`Schema ${name} already registered`);\n    }\n    this.schemasByName[name] = data;\n  }\n\n  registerValue(type: string, name: string, value: any) {\n    if (!this.valueByTypeAndName[type]) {\n      this.valueByTypeAndName[type] = {};\n    }\n    this.valueByTypeAndName[type][name] = value;\n  }\n\n  async lookupValue<T = unknown>(\n    type: string,\n    key: string\n  ): Promise<T | undefined> {\n    const pluginName = parsePluginName(key);\n    if (!this.valueByTypeAndName[type]?.[key] && pluginName) {\n      await this.initializePlugin(pluginName);\n    }\n    return (\n      (this.valueByTypeAndName[type]?.[key] as T) ||\n      this.parent?.lookupValue(type, key)\n    );\n  }\n\n  async listValues<T>(type: string): Promise<Record<string, T>> {\n    await this.initializeAllPlugins();\n    return {\n      ...((await this.parent?.listValues(type)) || {}),\n      ...(this.valueByTypeAndName[type] || {}),\n    } as Record<string, T>;\n  }\n\n  /**\n   * Looks up a schema.\n   * @param name The name of the schema to lookup.\n   * @returns The schema.\n   */\n  lookupSchema(name: string): Schema | undefined {\n    return this.schemasByName[name] || this.parent?.lookupSchema(name);\n  }\n}\n\n/**\n * Manages AsyncLocalStorage instances in a single place.\n */\nexport class AsyncStore {\n  private asls: Record<string, AsyncLocalStorage<any>> = {};\n\n  getStore<T>(key: string): T | undefined {\n    return this.asls[key]?.getStore();\n  }\n\n  run<T, R>(key: string, store: T, callback: () => R): R {\n    if (!this.asls[key]) {\n      this.asls[key] = new AsyncLocalStorage<T>();\n    }\n    return this.asls[key].run(store, callback);\n  }\n}\n\n/**\n * An object that has a reference to Genkit Registry.\n */\nexport interface HasRegistry {\n  get registry(): Registry;\n}\n"],"mappings":"AAgBA,SAAS,iBAAiB;AAC1B,SAAS,yBAAyB;AAElC,SAAiB,sCAAsC;AACvD,SAAS,mBAAmB;AAC5B,SAAS,cAAc;AAEvB,SAAqB,oBAAoB;AA6BzC,SAAS,gBAAgB,aAAqB;AAC5C,QAAM,SAAS,YAAY,MAAM,GAAG;AACpC,MAAI,OAAO,UAAU,GAAG;AACtB,WAAO,OAAO,CAAC;AAAA,EACjB;AACA,SAAO;AACT;AAOO,MAAM,SAAS;AAAA,EA0BpB,YAAmB,QAAmB;AAAnB;AAAA,EAAoB;AAAA,EAzB/B,cAIJ,CAAC;AAAA,EACG,gBAAgD,CAAC;AAAA,EACjD,gBAAwC,CAAC;AAAA,EACzC,qBAA0D,CAAC;AAAA,EAC3D,wBAAwB;AAAA,EACzB,eAAkC;AAAA,EAEhC,aAAa,IAAI,WAAW;AAAA,EAC5B,YAAY,IAAI,UAAU;AAAA,IACjC,gBAAgB,OAAO,SAAS;AAC9B,YAAM,iBAAiB,MAAM,KAAK,aAAa,IAAI;AACnD,UAAI,CAAC,gBAAgB;AACnB,cAAM,IAAI,YAAY;AAAA,UACpB,SAAS,WAAW,IAAI;AAAA,UACxB,QAAQ;AAAA,QACV,CAAC;AAAA,MACH;AACA,aAAO,aAAa,cAAc;AAAA,IACpC;AAAA,EACF,CAAC;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EASD,OAAO,WAAW,QAAkB;AAClC,WAAO,IAAI,SAAS,MAAM;AAAA,EAC5B;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,MAAM,aAIJ,KAAyB;AAEzB,UAAM,aAAa,gBAAgB,GAAG;AACtC,QAAI,CAAC,KAAK,YAAY,GAAG,KAAK,YAAY;AACxC,YAAM,KAAK,iBAAiB,UAAU;AAAA,IACxC;AACA,WACI,MAAM,KAAK,YAAY,GAAG,KAAY,KAAK,QAAQ,aAAa,GAAG;AAAA,EAEzE;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,eACE,MACA,QACA;AACA,UAAM,MAAM,IAAI,IAAI,IAAI,OAAO,SAAS,IAAI;AAC5C,WAAO,MAAM,eAAe,GAAG,EAAE;AACjC,QAAI,KAAK,YAAY,eAAe,GAAG,GAAG;AAExC,aAAO;AAAA,QACL,YAAY,GAAG;AAAA,MACjB;AAAA,IACF;AACA,SAAK,YAAY,GAAG,IAAI;AAAA,EAC1B;AAAA;AAAA;AAAA;AAAA,EAKA,oBACE,MACA,MACA,QACA;AACA,UAAM,MAAM,IAAI,IAAI,IAAI,IAAI;AAC5B,WAAO,MAAM,eAAe,GAAG,UAAU;AACzC,QAAI,KAAK,YAAY,eAAe,GAAG,GAAG;AAExC,aAAO;AAAA,QACL,YAAY,GAAG;AAAA,MACjB;AAAA,IACF;AACA,SAAK,YAAY,GAAG,IAAI;AAAA,EAC1B;AAAA;AAAA;AAAA;AAAA;AAAA,EAMA,MAAM,cAAsC;AAC1C,UAAM,KAAK,qBAAqB;AAChC,UAAM,UAA8D,CAAC;AACrE,UAAM,QAAQ;AAAA,MACZ,OAAO,QAAQ,KAAK,WAAW,EAAE,IAAI,OAAO,CAAC,KAAK,MAAM,MAAM;AAC5D,gBAAQ,GAAG,IAAI,MAAM;AAAA,MACvB,CAAC;AAAA,IACH;AACA,WAAO;AAAA,MACL,GAAI,MAAM,KAAK,QAAQ,YAAY;AAAA,MACnC,GAAG;AAAA,IACL;AAAA,EACF;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,uBAAuB;AAC3B,QAAI,KAAK,uBAAuB;AAC9B;AAAA,IACF;AACA,eAAW,cAAc,OAAO,KAAK,KAAK,aAAa,GAAG;AACxD,YAAM,KAAK,iBAAiB,UAAU;AAAA,IACxC;AACA,SAAK,wBAAwB;AAAA,EAC/B;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,uBAAuB,MAAc,UAA0B;AAC7D,QAAI,KAAK,cAAc,IAAI,GAAG;AAC5B,YAAM,IAAI,MAAM,UAAU,IAAI,qBAAqB;AAAA,IACrD;AACA,SAAK,wBAAwB;AAC7B,QAAI;AACJ,QAAI,gBAAgB;AACpB,SAAK,cAAc,IAAI,IAAI;AAAA,MACzB,MAAM,SAAS;AAAA,MACf,aAAa,MAAM;AACjB,YAAI,CAAC,eAAe;AAClB,mBAAS,SAAS,YAAY;AAC9B,0BAAgB;AAAA,QAClB;AACA,eAAO;AAAA,MACT;AAAA,IACF;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,aAAa,MAA0C;AACrD,WAAO,KAAK,cAAc,IAAI,KAAK,KAAK,QAAQ,aAAa,IAAI;AAAA,EACnE;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,MAAM,iBAAiB,MAAc;AACnC,QAAI,KAAK,cAAc,IAAI,GAAG;AAC5B,aAAO,MAAM;AAAA,QAA+B;AAAA,QAAM,MAChD,KAAK,cAAc,IAAI,EAAE,YAAY;AAAA,MACvC;AAAA,IACF;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,eAAe,MAAc,MAAc;AACzC,QAAI,KAAK,cAAc,IAAI,GAAG;AAC5B,YAAM,IAAI,MAAM,UAAU,IAAI,qBAAqB;AAAA,IACrD;AACA,SAAK,cAAc,IAAI,IAAI;AAAA,EAC7B;AAAA,EAEA,cAAc,MAAc,MAAc,OAAY;AACpD,QAAI,CAAC,KAAK,mBAAmB,IAAI,GAAG;AAClC,WAAK,mBAAmB,IAAI,IAAI,CAAC;AAAA,IACnC;AACA,SAAK,mBAAmB,IAAI,EAAE,IAAI,IAAI;AAAA,EACxC;AAAA,EAEA,MAAM,YACJ,MACA,KACwB;AACxB,UAAM,aAAa,gBAAgB,GAAG;AACtC,QAAI,CAAC,KAAK,mBAAmB,IAAI,IAAI,GAAG,KAAK,YAAY;AACvD,YAAM,KAAK,iBAAiB,UAAU;AAAA,IACxC;AACA,WACG,KAAK,mBAAmB,IAAI,IAAI,GAAG,KACpC,KAAK,QAAQ,YAAY,MAAM,GAAG;AAAA,EAEtC;AAAA,EAEA,MAAM,WAAc,MAA0C;AAC5D,UAAM,KAAK,qBAAqB;AAChC,WAAO;AAAA,MACL,GAAK,MAAM,KAAK,QAAQ,WAAW,IAAI,KAAM,CAAC;AAAA,MAC9C,GAAI,KAAK,mBAAmB,IAAI,KAAK,CAAC;AAAA,IACxC;AAAA,EACF;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAOA,aAAa,MAAkC;AAC7C,WAAO,KAAK,cAAc,IAAI,KAAK,KAAK,QAAQ,aAAa,IAAI;AAAA,EACnE;AACF;AAKO,MAAM,WAAW;AAAA,EACd,OAA+C,CAAC;AAAA,EAExD,SAAY,KAA4B;AACtC,WAAO,KAAK,KAAK,GAAG,GAAG,SAAS;AAAA,EAClC;AAAA,EAEA,IAAU,KAAa,OAAU,UAAsB;AACrD,QAAI,CAAC,KAAK,KAAK,GAAG,GAAG;AACnB,WAAK,KAAK,GAAG,IAAI,IAAI,kBAAqB;AAAA,IAC5C;AACA,WAAO,KAAK,KAAK,GAAG,EAAE,IAAI,OAAO,QAAQ;AAAA,EAC3C;AACF;","names":[]}