@genkit-ai/core
Version:
Genkit AI framework core libraries.
342 lines • 10.2 kB
JavaScript
import { Dotprompt } from "dotprompt";
import { AsyncLocalStorage } from "node:async_hooks";
import {
runOutsideActionRuntimeContext
} from "./action.js";
import {
lookupBackgroundAction
} from "./background-action.js";
import { GenkitError } from "./error.js";
import { logger } from "./logging.js";
import { toJsonSchema } from "./schema.js";
function parsePluginName(registryKey) {
const tokens = registryKey.split("/");
if (tokens.length >= 4) {
return tokens[2];
}
return void 0;
}
function parseRegistryKey(registryKey) {
const tokens = registryKey.split("/");
if (tokens.length < 3) {
return void 0;
}
if (tokens.length >= 4) {
return {
actionType: tokens[1],
pluginName: tokens[2],
actionName: tokens.slice(3).join("/")
};
}
return {
actionType: tokens[1],
actionName: tokens[2]
};
}
class Registry {
actionsById = {};
pluginsByName = {};
schemasByName = {};
valueByTypeAndName = {};
allPluginsInitialized = false;
apiStability = "stable";
asyncStore;
dotprompt;
parent;
/** Additional runtime context data for flows and tools. */
context;
constructor(parent) {
if (parent) {
this.parent = parent;
this.apiStability = parent?.apiStability;
this.asyncStore = parent.asyncStore;
this.dotprompt = parent.dotprompt;
} else {
this.asyncStore = new AsyncStore();
this.dotprompt = new Dotprompt({
schemaResolver: async (name) => {
const resolvedSchema = await this.lookupSchema(name);
if (!resolvedSchema) {
throw new GenkitError({
message: `Schema '${name}' not found`,
status: "NOT_FOUND"
});
}
return toJsonSchema(resolvedSchema);
}
});
}
}
/**
* Creates a new registry overlaid onto the provided registry.
* @param parent The parent registry.
* @returns The new overlaid registry.
*/
static withParent(parent) {
return new Registry(parent);
}
/**
* Looks up an action in the registry.
* @param key The key of the action to lookup.
* @returns The action.
*/
async lookupAction(key) {
const parsedKey = parseRegistryKey(key);
if (parsedKey?.pluginName && this.pluginsByName[parsedKey.pluginName]) {
await this.initializePlugin(parsedKey.pluginName);
if (!this.actionsById[key]) {
await this.resolvePluginAction(
parsedKey.pluginName,
parsedKey.actionType,
parsedKey.actionName
);
}
}
return await this.actionsById[key] || this.parent?.lookupAction(key);
}
/**
* Looks up a background action from the registry.
* @param key The key of the action to lookup.
* @returns The action.
*/
async lookupBackgroundAction(key) {
return lookupBackgroundAction(this, key);
}
/**
* Registers an action in the registry.
* @param type The type of the action to register.
* @param action The action to register.
*/
registerAction(type, action) {
if (type !== action.__action.actionType) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `action type (${type}) does not match type on action (${action.__action.actionType})`
});
}
const key = `/${type}/${action.__action.name}`;
logger.debug(`registering ${key}`);
if (this.actionsById.hasOwnProperty(key)) {
logger.warn(
`WARNING: ${key} already has an entry in the registry. Overwriting.`
);
}
this.actionsById[key] = action;
}
/**
* Registers an action promise in the registry.
*/
registerActionAsync(type, name, action) {
const key = `/${type}/${name}`;
logger.debug(`registering ${key} (async)`);
if (this.actionsById.hasOwnProperty(key)) {
logger.warn(
`WARNING: ${key} already has an entry in the registry. Overwriting.`
);
}
this.actionsById[key] = action;
}
/**
* Returns all actions that have been registered in the registry.
* @returns All actions in the registry as a map of <key, action>.
*/
async listActions() {
await this.initializeAllPlugins();
const actions = {};
await Promise.all(
Object.entries(this.actionsById).map(async ([key, action]) => {
actions[key] = await action;
})
);
return {
...await this.parent?.listActions(),
...actions
};
}
/**
* Returns all actions that are resolvable by plugins as well as those that are already
* in the registry.
*
* NOTE: this method should not be used in latency sensitive code paths.
* It may rely on "admin" API calls such as "list models", which may cause increased cold start latency.
*
* @returns All resolvable action metadata as a map of <key, action metadata>.
*/
async listResolvableActions() {
const resolvableActions = {};
await Promise.all(
Object.entries(this.pluginsByName).map(async ([pluginName, plugin]) => {
if (plugin.listActions) {
try {
(await plugin.listActions()).forEach((meta) => {
if (!meta.name) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `Invalid metadata when listing actions from ${pluginName} - name required`
});
}
if (!meta.actionType) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `Invalid metadata when listing actions from ${pluginName} - actionType required`
});
}
resolvableActions[`/${meta.actionType}/${meta.name}`] = meta;
});
} catch (e) {
logger.error(`Error listing actions for ${pluginName}
`, e);
}
}
})
);
for (const [key, action] of Object.entries(await this.listActions())) {
resolvableActions[key] = action.__action;
}
return {
...await this.parent?.listResolvableActions(),
...resolvableActions
};
}
/**
* Initializes all plugins in the registry.
*/
async initializeAllPlugins() {
if (this.allPluginsInitialized) {
return;
}
for (const pluginName of Object.keys(this.pluginsByName)) {
await this.initializePlugin(pluginName);
}
this.allPluginsInitialized = true;
}
/**
* Registers a plugin provider. This plugin must be initialized before it can be used by calling {@link initializePlugin} or {@link initializeAllPlugins}.
* @param name The name of the plugin to register.
* @param provider The plugin provider.
*/
registerPluginProvider(name, provider) {
if (this.pluginsByName[name]) {
throw new Error(`Plugin ${name} already registered`);
}
this.allPluginsInitialized = false;
let cached;
let isInitialized = false;
this.pluginsByName[name] = {
name: provider.name,
initializer: () => {
if (!isInitialized) {
cached = provider.initializer();
isInitialized = true;
}
return cached;
},
resolver: async (actionType, actionName) => {
if (provider.resolver) {
await provider.resolver(actionType, actionName);
}
},
listActions: async () => {
if (provider.listActions) {
return await provider.listActions();
}
return [];
}
};
}
/**
* Looks up a plugin.
* @param name The name of the plugin to lookup.
* @returns The plugin provider.
*/
lookupPlugin(name) {
return this.pluginsByName[name] || this.parent?.lookupPlugin(name);
}
/**
* Resolves a new Action dynamically by registering it.
* @param pluginName The name of the plugin
* @param actionType The type of the action
* @param actionName The name of the action
* @returns
*/
async resolvePluginAction(pluginName, actionType, actionName) {
const plugin = this.pluginsByName[pluginName];
if (plugin) {
return await runOutsideActionRuntimeContext(this, async () => {
if (plugin.resolver) {
await plugin.resolver(actionType, actionName);
}
});
}
}
/**
* Initializes a plugin already registered with {@link registerPluginProvider}.
* @param name The name of the plugin to initialize.
* @returns The plugin.
*/
async initializePlugin(name) {
if (this.pluginsByName[name]) {
return await runOutsideActionRuntimeContext(
this,
() => this.pluginsByName[name].initializer()
);
}
}
/**
* Registers a schema.
* @param name The name of the schema to register.
* @param data The schema to register (either a Zod schema or a JSON schema).
*/
registerSchema(name, data) {
if (this.schemasByName[name]) {
throw new Error(`Schema ${name} already registered`);
}
this.schemasByName[name] = data;
}
registerValue(type, name, value) {
if (!this.valueByTypeAndName[type]) {
this.valueByTypeAndName[type] = {};
}
this.valueByTypeAndName[type][name] = value;
}
async lookupValue(type, key) {
const pluginName = parsePluginName(key);
if (!this.valueByTypeAndName[type]?.[key] && pluginName) {
await this.initializePlugin(pluginName);
}
return this.valueByTypeAndName[type]?.[key] || this.parent?.lookupValue(type, key);
}
async listValues(type) {
await this.initializeAllPlugins();
return {
...await this.parent?.listValues(type) || {},
...this.valueByTypeAndName[type] || {}
};
}
/**
* Looks up a schema.
* @param name The name of the schema to lookup.
* @returns The schema.
*/
lookupSchema(name) {
return this.schemasByName[name] || this.parent?.lookupSchema(name);
}
}
class AsyncStore {
asls = {};
getStore(key) {
return this.asls[key]?.getStore();
}
run(key, store, callback) {
if (!this.asls[key]) {
this.asls[key] = new AsyncLocalStorage();
}
return this.asls[key].run(store, callback);
}
}
export {
AsyncStore,
Registry,
parseRegistryKey
};
//# sourceMappingURL=registry.mjs.map