@genkit-ai/core
Version:
Genkit AI framework core libraries.
415 lines • 12.5 kB
JavaScript
import * as z from "zod";
import { getAsyncContext } from "./async-context.mjs";
import { Channel, lazy } from "./async.mjs";
import { getContext, runWithContext } from "./context.mjs";
import { parseSchema } from "./schema.mjs";
import {
SPAN_TYPE_ATTR,
runInNewSpan,
setCustomMetadataAttributes
} from "./tracing.mjs";
import { StatusCodes, StatusSchema } from "./statusTypes.mjs";
import { InMemoryStreamManager, StreamNotFoundError } from "./streaming.mjs";
const makeNoopAbortSignal = () => new AbortController().signal;
const ActionMetadataSchema = z.object({
key: z.string().optional(),
actionType: z.string().optional(),
name: z.string(),
description: z.string().optional(),
inputSchema: z.unknown().optional(),
inputJsonSchema: z.object({}).optional(),
outputSchema: z.unknown().optional(),
outputJsonSchema: z.object({}).optional(),
streamSchema: z.unknown().optional(),
initSchema: z.unknown().optional(),
initJsonSchema: z.object({}).optional(),
metadata: z.record(z.string(), z.any()).optional()
});
function actionWithMiddleware(action2, middleware) {
const wrapped = (async (req, options) => {
return (await wrapped.run(req, options)).result;
});
wrapped.__action = action2.__action;
wrapped.run = async (req, options) => {
let telemetry;
const dispatch = async (index, req2, opts) => {
if (index === middleware.length) {
const result = await action2.run(req2, opts);
telemetry = result.telemetry;
return result.result;
}
const currentMiddleware = middleware[index];
if (currentMiddleware.length === 3) {
return currentMiddleware(
req2,
opts,
async (modifiedReq, modifiedOptions) => dispatch(index + 1, modifiedReq || req2, modifiedOptions || opts)
);
} else if (currentMiddleware.length === 2) {
return currentMiddleware(
req2,
async (modifiedReq) => dispatch(index + 1, modifiedReq || req2, opts)
);
} else {
throw new Error("unspported middleware function shape");
}
};
wrapped.stream = action2.stream;
if (action2.streamBidi) {
wrapped.streamBidi = action2.streamBidi;
}
return { result: await dispatch(0, req, options), telemetry };
};
return wrapped;
}
function action(config, fn) {
const actionName = typeof config.name === "string" ? config.name : `${config.name.pluginId}/${config.name.actionId}`;
const actionMetadata = {
key: `/${config.actionType}/${actionName}`,
name: actionName,
description: config.description,
inputSchema: config.inputSchema,
inputJsonSchema: config.inputJsonSchema,
outputSchema: config.outputSchema,
outputJsonSchema: config.outputJsonSchema,
streamSchema: config.streamSchema,
initSchema: config.initSchema,
initJsonSchema: config.initJsonSchema,
metadata: config.metadata,
actionType: config.actionType
};
const actionFn = (async (input, options) => {
return (await actionFn.run(input, options)).result;
});
actionFn.__action = { ...actionMetadata };
actionFn.run = async (input, options) => {
if (config.inputSchema || config.inputJsonSchema) {
if (!options?.inputStream) {
input = parseSchema(input, {
schema: config.inputSchema,
jsonSchema: config.inputJsonSchema
});
} else {
const inputStream = options.inputStream;
options = {
...options,
inputStream: (async function* () {
for await (const item of inputStream) {
yield parseSchema(item, {
schema: config.inputSchema,
jsonSchema: config.inputJsonSchema
});
}
})()
};
}
}
if (config.initSchema || config.initJsonSchema) {
const validatedInit = parseSchema(options?.init, {
schema: config.initSchema,
jsonSchema: config.initJsonSchema
});
options = { ...options, init: validatedInit };
}
let traceId;
let spanId;
const genkitKey = actionFn.__action.key;
let output = await runInNewSpan(
{
metadata: {
name: actionName
},
labels: {
[SPAN_TYPE_ATTR]: "action",
"genkit:metadata:subtype": config.actionType,
...genkitKey ? { "genkit:key": genkitKey } : {},
...options?.telemetryLabels
}
},
async (metadata, span) => {
setCustomMetadataAttributes({
subtype: config.actionType
});
if (options?.context) {
setCustomMetadataAttributes({
context: JSON.stringify(options.context)
});
}
traceId = span.spanContext().traceId;
spanId = span.spanContext().spanId;
if (options?.onTraceStart) {
options.onTraceStart({ traceId, spanId });
}
metadata.name = actionName;
metadata.input = input;
if (options?.init !== void 0) {
metadata.init = options.init;
}
try {
const actFn = () => fn(input, {
...options,
// Context can either be explicitly set, or inherited from the parent action.
context: {
...actionFn.__registry?.context,
...options?.context ?? getContext()
},
streamingRequested: !!options?.onChunk && options.onChunk !== sentinelNoopStreamingCallback,
sendChunk: options?.onChunk ?? sentinelNoopStreamingCallback,
inputStream: options?.inputStream ?? asyncIterableFromArray([input]),
trace: {
traceId,
spanId
},
registry: actionFn.__registry,
abortSignal: options?.abortSignal ?? makeNoopAbortSignal()
});
const output2 = await runWithContext(options?.context, actFn);
metadata.output = JSON.stringify(output2);
return output2;
} catch (err) {
if (typeof err === "object") {
err.traceId = traceId;
}
throw err;
}
}
);
output = parseSchema(output, {
schema: config.outputSchema,
jsonSchema: config.outputJsonSchema
});
return {
result: output,
telemetry: {
traceId,
spanId
}
};
};
actionFn.stream = (input, opts) => {
let chunkStreamController;
const chunkStream = new ReadableStream({
start(controller) {
chunkStreamController = controller;
},
pull() {
},
cancel() {
}
});
const invocationPromise = actionFn.run(
!opts?.inputStream && config.inputSchema ? config.inputSchema.parse(input) : input,
{
onChunk: ((chunk) => {
chunkStreamController.enqueue(chunk);
}),
context: {
...actionFn.__registry?.context,
...opts?.context ?? getContext()
},
inputStream: opts?.inputStream,
abortSignal: opts?.abortSignal,
telemetryLabels: opts?.telemetryLabels,
init: opts?.init
}
).then((s) => s.result).finally(() => {
try {
chunkStreamController.close();
} catch (e) {
}
});
return {
output: invocationPromise,
stream: (async function* () {
const reader = chunkStream.getReader();
try {
while (true) {
const chunk = await reader.read();
if (chunk.value) {
yield chunk.value;
}
if (chunk.done) {
break;
}
}
return await invocationPromise;
} finally {
reader.cancel().catch(() => {
});
}
})()
};
};
if (config.use) {
return actionWithMiddleware(actionFn, config.use);
}
return actionFn;
}
function isAction(a) {
return typeof a === "function" && "__action" in a;
}
function defineAction(registry, config, fn) {
if (isInRuntimeContext()) {
throw new Error(
"Cannot define new actions at runtime.\nSee: https://github.com/genkit-ai/genkit/blob/main/docs/errors/no_new_actions_at_runtime.md"
);
}
const act = action(config, async (i, options) => {
await registry.initializeAllPlugins();
return await runInActionRuntimeContext(() => fn(i, options));
});
act.__action.actionType = config.actionType;
registry.registerAction(config.actionType, act);
return act;
}
function defineBidiAction(registry, config, fn) {
const act = bidiAction(config, (input) => {
return (async function* () {
await registry.initializeAllPlugins();
return yield* fn(input);
})();
});
registry.registerAction(config.actionType, act);
return act;
}
function bidiAction(config, fn) {
const meta = { ...config.metadata, bidi: true };
const act = action({ ...config, metadata: meta }, async (input, options) => {
const stream = options.inputStream;
const outputGen = fn({
...options,
inputStream: stream
});
const iter = outputGen[Symbol.asyncIterator]();
let result;
try {
while (true) {
const { value, done } = await iter.next();
if (done) {
result = value;
break;
}
options.sendChunk(value);
}
} finally {
if (iter.return) {
await iter.return(void 0);
}
}
return result;
});
act.streamBidi = (init, opts) => {
let channel;
let stream = opts?.inputStream;
if (!stream) {
channel = new Channel();
stream = channel;
}
const result = act.stream(void 0, {
...opts,
init,
inputStream: stream
});
return {
...result,
stream: (async function* () {
try {
for await (const chunk of result.stream) {
yield chunk;
}
} finally {
if (channel) {
channel.close();
}
}
})(),
send: (chunk) => {
if (!channel) {
throw new Error("Cannot send to a provided stream.");
}
channel.send(chunk);
},
close: () => {
if (!channel) {
throw new Error("Cannot close a provided stream.");
}
channel.close();
}
};
};
return act;
}
function defineActionAsync(registry, actionType, name, config, onInit) {
const actionName = typeof name === "string" ? name : `${name.pluginId}/${name.actionId}`;
const actionPromise = lazy(
() => config.then((resolvedConfig) => {
const act = action(
resolvedConfig,
async (i, options) => {
await registry.initializeAllPlugins();
return await runInActionRuntimeContext(
() => resolvedConfig.fn(i, options)
);
}
);
act.__action.actionType = actionType;
act.__action.key = `/${actionType}/${actionName}`;
onInit?.(act);
return act;
})
);
registry.registerActionAsync(actionType, actionName, actionPromise);
return actionPromise;
}
const streamingAlsKey = "core.action.streamingCallback";
const sentinelNoopStreamingCallback = () => null;
function runWithStreamingCallback(streamingCallback, fn) {
return getAsyncContext().run(
streamingAlsKey,
streamingCallback || sentinelNoopStreamingCallback,
fn
);
}
function getStreamingCallback() {
const cb = getAsyncContext().getStore(streamingAlsKey);
if (cb === sentinelNoopStreamingCallback) {
return void 0;
}
return cb;
}
const runtimeContextAslKey = "core.action.runtimeContext";
function isInRuntimeContext() {
return getAsyncContext().getStore(runtimeContextAslKey) === "runtime";
}
function runInActionRuntimeContext(fn) {
return getAsyncContext().run(runtimeContextAslKey, "runtime", fn);
}
function runOutsideActionRuntimeContext(fn) {
return getAsyncContext().run(runtimeContextAslKey, "outside", fn);
}
async function* asyncIterableFromArray(array) {
for (const item of array) {
yield item;
}
}
export {
ActionMetadataSchema,
InMemoryStreamManager,
StatusCodes,
StatusSchema,
StreamNotFoundError,
action,
actionWithMiddleware,
bidiAction,
defineAction,
defineActionAsync,
defineBidiAction,
getStreamingCallback,
isAction,
isInRuntimeContext,
runInActionRuntimeContext,
runOutsideActionRuntimeContext,
runWithStreamingCallback,
sentinelNoopStreamingCallback
};
//# sourceMappingURL=action.mjs.map