genkit
Version:
Genkit AI framework
489 lines • 14.1 kB
JavaScript
import {
GenkitAI,
defineHelper,
definePartial,
definePrompt,
defineTool,
generate,
loadPromptFolder,
modelRef,
prompt
} from "@genkit-ai/ai";
import {
defineEmbedder
} from "@genkit-ai/ai/embedder";
import {
defineEvaluator,
evaluate
} from "@genkit-ai/ai/evaluator";
import { configureFormats } from "@genkit-ai/ai/formats";
import {
defineBackgroundModel,
defineGenerateAction,
defineModel
} from "@genkit-ai/ai/model";
import {
defineReranker,
rerank
} from "@genkit-ai/ai/reranker";
import {
defineIndexer,
defineRetriever,
defineSimpleRetriever,
index,
retrieve
} from "@genkit-ai/ai/retriever";
import {
dynamicTool
} from "@genkit-ai/ai/tool";
import {
GenkitError,
ReflectionServer,
defineDynamicActionProvider,
defineFlow,
defineJsonSchema,
defineSchema,
isAction,
isBackgroundAction,
isDevEnv,
registerBackgroundAction,
setClientHeader
} from "@genkit-ai/core";
import { Channel } from "@genkit-ai/core/async";
import { logger } from "./logging.mjs";
import {
isPluginV2
} from "./plugin.mjs";
import { Registry } from "./registry.mjs";
import { SPAN_TYPE_ATTR, runInNewSpan } from "./tracing.mjs";
class Genkit extends GenkitAI {
/** Developer-configured options. */
options;
/** Reflection server for this registry. May be null if not started. */
reflectionServer = null;
/** List of flows that have been registered in this instance. */
flows = [];
get apiStability() {
return this.registry.apiStability;
}
constructor(options) {
const registry = new Registry();
super(registry);
this.options = options || {};
if (this.options.context) {
this.registry.context = this.options.context;
}
this.configure();
if (isDevEnv() && !disableReflectionApi) {
this.reflectionServer = new ReflectionServer(this.registry, {
configuredEnvs: ["dev"],
name: this.options.name
});
this.reflectionServer.start().catch((e) => logger.error);
}
if (options?.clientHeader) {
setClientHeader(options?.clientHeader);
}
}
/**
* Defines and registers a flow function.
*/
defineFlow(config, fn) {
const flow = defineFlow(this.registry, config, fn);
this.flows.push(flow);
return flow;
}
defineTool(config, fn) {
return defineTool(this.registry, config, fn);
}
/**
* Defines a dynamic tool. Dynamic tools are just like regular tools ({@link Genkit.defineTool}) but will not be registered in the
* Genkit registry and can be defined dynamically at runtime.
*/
dynamicTool(config, fn) {
return dynamicTool(config, fn);
}
/**
* Defines and registers a dynamic action provider (e.g. mcp host)
*/
defineDynamicActionProvider(config, fn) {
return defineDynamicActionProvider(this.registry, config, fn);
}
/**
* Defines and registers a schema from a Zod schema.
*
* Defined schemas can be referenced by `name` in prompts in place of inline schemas.
*/
defineSchema(name, schema) {
return defineSchema(this.registry, name, schema);
}
/**
* Defines and registers a schema from a JSON schema.
*
* Defined schemas can be referenced by `name` in prompts in place of inline schemas.
*/
defineJsonSchema(name, jsonSchema) {
return defineJsonSchema(this.registry, name, jsonSchema);
}
/**
* Defines a new model and adds it to the registry.
*/
defineModel(options, runner) {
return defineModel(this.registry, options, runner);
}
/**
* Defines a new background model and adds it to the registry.
*/
defineBackgroundModel(options) {
return defineBackgroundModel(this.registry, options);
}
/**
* Looks up a prompt by `name` (and optionally `variant`). Can be used to lookup
* .prompt files or prompts previously defined with {@link Genkit.definePrompt}
*/
prompt(name, options) {
return this.wrapExecutablePromptPromise(
`${name}${options?.variant ? `.${options?.variant}` : ""}`,
prompt(this.registry, name, {
...options,
dir: this.options.promptDir === void 0 ? "./prompts" : this.options.promptDir
})
);
}
wrapExecutablePromptPromise(name, promise) {
const executablePrompt = (async (input, opts) => {
return (await promise)(input, opts);
});
executablePrompt.ref = { name };
executablePrompt.render = async (input, opts) => {
return (await promise).render(input, opts);
};
executablePrompt.stream = (input, opts) => {
let channel = new Channel();
const generated = runInNewSpan(
this.registry,
{
metadata: {
name,
input
},
labels: {
[SPAN_TYPE_ATTR]: "dotprompt"
}
},
() => generate(
this.registry,
promise.then(
(action) => action.render(input, {
...opts,
onChunk: (chunk) => channel.send(chunk)
})
)
)
);
generated.then(
() => channel.close(),
(err) => channel.error(err)
);
return {
response: generated,
stream: channel
};
};
executablePrompt.asTool = async () => {
return (await promise).asTool();
};
return executablePrompt;
}
/**
* Defines and registers a prompt based on a function.
*
* This is an alternative to defining and importing a .prompt file, providing
* the most advanced control over how the final request to the model is made.
*
* @param options - Prompt metadata including model, model params,
* input/output schemas, etc
* @param fn - A function that returns a {@link GenerateRequest}. Any config
* parameters specified by the {@link GenerateRequest} will take precedence
* over any parameters specified by `options`.
*
* ```ts
* const hi = ai.definePrompt(
* {
* name: 'hi',
* input: {
* schema: z.object({
* name: z.string(),
* }),
* },
* config: {
* temperature: 1,
* },
* },
* async (input) => {
* return {
* messages: [ { role: 'user', content: [{ text: `hi ${input.name}` }] } ],
* };
* }
* );
* const { text } = await hi({ name: 'Genkit' });
* ```
*/
definePrompt(options, templateOrFn) {
if (templateOrFn) {
if (options.messages) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: "Cannot specify template/function argument and `options.messages` at the same time"
});
}
if (typeof templateOrFn === "string") {
return definePrompt(this.registry, {
...options,
messages: templateOrFn
});
} else {
return definePrompt(this.registry, {
...options,
messages: async (input) => {
const response = await templateOrFn(input);
return response.messages;
}
});
}
}
return definePrompt(this.registry, options);
}
/**
* Creates a retriever action for the provided {@link RetrieverFn} implementation.
*/
defineRetriever(options, runner) {
return defineRetriever(this.registry, options, runner);
}
/**
* 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.
*/
defineSimpleRetriever(options, handler) {
return defineSimpleRetriever(this.registry, options, handler);
}
/**
* Creates an indexer action for the provided {@link IndexerFn} implementation.
*/
defineIndexer(options, runner) {
return defineIndexer(this.registry, options, runner);
}
/**
* Creates evaluator action for the provided {@link EvaluatorFn} implementation.
*/
defineEvaluator(options, runner) {
return defineEvaluator(this.registry, options, runner);
}
/**
* Creates embedder model for the provided {@link EmbedderFn} model implementation.
*/
defineEmbedder(options, runner) {
return defineEmbedder(this.registry, options, runner);
}
/**
* create a handlebars helper (https://handlebarsjs.com/guide/block-helpers.html) to be used in dotprompt templates.
*/
defineHelper(name, fn) {
defineHelper(this.registry, name, fn);
}
/**
* Creates a handlebars partial (https://handlebarsjs.com/guide/partials.html) to be used in dotprompt templates.
*/
definePartial(name, source) {
definePartial(this.registry, name, source);
}
/**
* Creates a reranker action for the provided {@link RerankerFn} implementation.
*/
defineReranker(options, runner) {
return defineReranker(this.registry, options, runner);
}
/**
* Evaluates the given `dataset` using the specified `evaluator`.
*/
evaluate(params) {
return evaluate(this.registry, params);
}
/**
* Reranks documents from a {@link RerankerArgument} based on the provided query.
*/
rerank(params) {
return rerank(this.registry, params);
}
/**
* Indexes `documents` using the provided `indexer`.
*/
index(params) {
return index(this.registry, params);
}
/**
* Retrieves documents from the `retriever` based on the provided `query`.
*/
retrieve(params) {
return retrieve(this.registry, params);
}
/**
* Configures the Genkit instance.
*/
configure() {
const activeRegistry = this.registry;
defineGenerateAction(activeRegistry);
configureFormats(activeRegistry);
const plugins = [...this.options.plugins ?? []];
if (this.options.model) {
this.registry.registerValue(
"defaultModel",
"defaultModel",
toModelRef(this.options.model)
);
}
if (this.options.promptDir !== null) {
loadPromptFolder(
this.registry,
this.options.promptDir ?? "./prompts",
""
);
}
plugins.forEach((plugin) => {
if (isPluginV2(plugin)) {
logger.debug(`Registering v2 plugin ${plugin.name}...`);
plugin.middleware?.()?.forEach((middleware) => {
activeRegistry.registerValue(
"middleware",
middleware.name,
middleware
);
});
activeRegistry.registerPluginProvider(plugin.name, {
name: plugin.name,
async initializer() {
logger.debug(`Initializing plugin ${plugin.name}:`);
if (!plugin.init) return;
const resolvedActions = await plugin.init();
resolvedActions?.forEach((resolvedAction) => {
registerActionV2(activeRegistry, resolvedAction, plugin);
});
},
async resolver(action, target) {
if (!plugin.resolve) return;
const resolvedAction = await plugin.resolve(action, target);
if (resolvedAction) {
registerActionV2(activeRegistry, resolvedAction, plugin);
}
},
async listActions() {
if (typeof plugin.list === "function") {
return (await plugin.list()).map((a) => {
if (a.name.startsWith(`${plugin.name}/`)) {
return a;
}
return {
...a,
// Apply namespace for v2 plugins.
name: `${plugin.name}/${a.name}`
};
});
}
return [];
}
});
} else {
const loadedPlugin = plugin(this);
logger.debug(`Registering plugin ${loadedPlugin.name}...`);
activeRegistry.registerPluginProvider(loadedPlugin.name, {
name: loadedPlugin.name,
async initializer() {
logger.debug(`Initializing plugin ${loadedPlugin.name}:`);
await loadedPlugin.initializer();
},
async resolver(action, target) {
if (loadedPlugin.resolver) {
await loadedPlugin.resolver(action, target);
}
},
async listActions() {
if (loadedPlugin.listActions) {
return await loadedPlugin.listActions();
}
return [];
}
});
}
});
}
/**
* Stops all servers.
*/
async stopServers() {
await this.reflectionServer?.stop();
this.reflectionServer = null;
}
}
function registerActionV2(registry, resolvedAction, plugin) {
if (isBackgroundAction(resolvedAction)) {
registerBackgroundAction(registry, resolvedAction, {
namespace: plugin.name
});
} else if (isAction(resolvedAction)) {
if (!resolvedAction.__action.actionType) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: "Action type is missing for " + resolvedAction.__action.name
});
}
registry.registerAction(
resolvedAction.__action.actionType,
resolvedAction,
{ namespace: plugin.name }
);
} else {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: "Unknown action type returned from plugin " + plugin.name
});
}
}
function genkit(options) {
return new Genkit(options);
}
const shutdown = async () => {
logger.debug("Shutting down all Genkit servers...");
await ReflectionServer.stopAll();
process.exit(0);
};
process.on("SIGTERM", shutdown);
process.on("SIGINT", shutdown);
let disableReflectionApi = false;
function __disableReflectionApi() {
disableReflectionApi = true;
}
function toModelRef(modelArg) {
if (modelArg === void 0) {
return void 0;
}
if (typeof modelArg === "string") {
return modelRef({ name: modelArg });
}
if (modelArg.name) {
return modelArg;
}
const modelAction = modelArg;
return modelRef({
name: modelAction.__action.name
});
}
export {
Genkit,
__disableReflectionApi,
genkit
};
//# sourceMappingURL=genkit.mjs.map