@genkit-ai/flow
Version:
Genkit AI framework workflow APIs.
621 lines • 19.5 kB
JavaScript
import {
__async,
__asyncGenerator,
__await
} from "./chunk-7OAPEGJQ.mjs";
import {
FlowStateSchema,
defineAction,
getStreamingCallback,
config as globalConfig,
isDevEnv
} from "@genkit-ai/core";
import { logger } from "@genkit-ai/core/logging";
import { toJsonSchema } from "@genkit-ai/core/schema";
import {
SPAN_TYPE_ATTR,
newTrace,
setCustomMetadataAttribute,
setCustomMetadataAttributes
} from "@genkit-ai/core/tracing";
import { SpanStatusCode } from "@opentelemetry/api";
import * as bodyParser from "body-parser";
import { default as cors } from "cors";
import express from "express";
import { performance } from "node:perf_hooks";
import { Context } from "./context.js";
import {
FlowExecutionError,
FlowStillRunningError,
InterruptError,
getErrorMessage,
getErrorStack
} from "./errors.js";
import {
FlowActionInputSchema,
FlowInvokeEnvelopeMessageSchema
} from "./types.js";
import {
generateFlowId,
metadataPrefix,
runWithActiveContext
} from "./utils.js";
const streamDelimiter = "\n";
const CREATED_FLOWS = "genkit__CREATED_FLOWS";
function createdFlows() {
if (global[CREATED_FLOWS] === void 0) {
global[CREATED_FLOWS] = [];
}
return global[CREATED_FLOWS];
}
function defineFlow(config, steps) {
const f = new Flow(
{
name: config.name,
inputSchema: config.inputSchema,
outputSchema: config.outputSchema,
streamSchema: config.streamSchema,
experimentalDurable: !!config.experimentalDurable,
stateStore: globalConfig ? () => globalConfig.getFlowStateStore() : void 0,
authPolicy: config.authPolicy,
middleware: config.middleware,
// We always use local dispatcher in dev mode or when one is not provided.
invoker: (flow, msg, streamingCallback) => __async(this, null, function* () {
if (!isDevEnv() && config.invoker) {
return config.invoker(flow, msg, streamingCallback);
}
const state = yield flow.runEnvelope(msg, streamingCallback);
return state.operation;
}),
scheduler: (flow, msg, delay = 0) => __async(this, null, function* () {
if (!config.experimentalDurable) {
throw new Error(
"This flow is not durable, cannot use scheduling features."
);
}
if (!isDevEnv() && config.experimentalScheduler) {
return config.experimentalScheduler(flow, msg, delay);
}
setTimeout(() => flow.runEnvelope(msg), delay * 1e3);
})
},
steps
);
createdFlows().push(f);
wrapAsAction(f);
return f;
}
class Flow {
constructor(config, steps) {
this.steps = steps;
this.name = config.name;
this.inputSchema = config.inputSchema;
this.outputSchema = config.outputSchema;
this.streamSchema = config.streamSchema;
this.stateStore = config.stateStore;
this.invoker = config.invoker;
this.scheduler = config.scheduler;
this.experimentalDurable = config.experimentalDurable;
this.authPolicy = config.authPolicy;
this.middleware = config.middleware;
if (this.authPolicy && this.experimentalDurable) {
throw new Error("Durable flows can not define auth policies.");
}
}
/**
* Executes the flow with the input directly.
*
* This will either be called by runEnvelope when starting durable flows,
* or it will be called directly when starting non-durable flows.
*/
runDirectly(input, opts) {
return __async(this, null, function* () {
const flowId = generateFlowId();
const state = createNewState(flowId, this.name, input);
const ctx = new Context(this, flowId, state, opts.auth);
try {
yield this.executeSteps(
ctx,
this.steps,
"start",
opts.streamingCallback,
opts.labels
);
} finally {
if (isDevEnv() || this.experimentalDurable) {
yield ctx.saveState();
}
}
return state;
});
}
/**
* Executes the flow with the input in the envelope format.
*/
runEnvelope(req, streamingCallback, auth) {
return __async(this, null, function* () {
logger.debug(req, "runEnvelope");
if (req.start) {
return this.runDirectly(req.start.input, {
streamingCallback,
auth,
labels: req.start.labels
});
}
if (req.schedule) {
if (!this.experimentalDurable) {
throw new Error("Cannot schedule a non-durable flow");
}
if (!this.stateStore) {
throw new Error(
"Flow state store for durable flows must be configured"
);
}
const flowId = generateFlowId();
const state = createNewState(flowId, this.name, req.schedule.input);
try {
yield (yield this.stateStore()).save(flowId, state);
yield this.scheduler(
this,
{ runScheduled: { flowId } },
req.schedule.delay
);
} catch (e) {
state.operation.done = true;
state.operation.result = {
error: getErrorMessage(e),
stacktrace: getErrorStack(e)
};
yield (yield this.stateStore()).save(flowId, state);
}
return state;
}
if (req.state) {
if (!this.experimentalDurable) {
throw new Error("Cannot state check a non-durable flow");
}
if (!this.stateStore) {
throw new Error(
"Flow state store for durable flows must be configured"
);
}
const flowId = req.state.flowId;
const state = yield (yield this.stateStore()).load(flowId);
if (state === void 0) {
throw new Error(`Unable to find flow state for ${flowId}`);
}
return state;
}
if (req.runScheduled) {
if (!this.experimentalDurable) {
throw new Error("Cannot run scheduled non-durable flow");
}
if (!this.stateStore) {
throw new Error(
"Flow state store for durable flows must be configured"
);
}
const flowId = req.runScheduled.flowId;
const state = yield (yield this.stateStore()).load(flowId);
if (state === void 0) {
throw new Error(`Unable to find flow state for ${flowId}`);
}
const ctx = new Context(this, flowId, state);
try {
yield this.executeSteps(
ctx,
this.steps,
"runScheduled",
void 0,
void 0
);
} finally {
yield ctx.saveState();
}
return state;
}
if (req.resume) {
if (!this.experimentalDurable) {
throw new Error("Cannot resume a non-durable flow");
}
if (!this.stateStore) {
throw new Error(
"Flow state store for durable flows must be configured"
);
}
const flowId = req.resume.flowId;
const state = yield (yield this.stateStore()).load(flowId);
if (state === void 0) {
throw new Error(`Unable to find flow state for ${flowId}`);
}
if (!state.blockedOnStep) {
throw new Error(
"Unable to resume flow that's currently not interrupted"
);
}
state.eventsTriggered[state.blockedOnStep.name] = req.resume.payload;
const ctx = new Context(this, flowId, state);
try {
yield this.executeSteps(
ctx,
this.steps,
"resume",
void 0,
void 0
);
} finally {
yield ctx.saveState();
}
return state;
}
throw new Error(
"Unexpected envelope message case, must set one of: start, schedule, runScheduled, resume, retry, state"
);
});
}
// TODO: refactor me... this is a mess!
executeSteps(ctx, handler, dispatchType, streamingCallback, labels) {
return __async(this, null, function* () {
const startTimeMs = performance.now();
yield runWithActiveContext(ctx, () => __async(this, null, function* () {
let traceContext;
if (ctx.state.traceContext) {
traceContext = JSON.parse(ctx.state.traceContext);
}
let ctxLinks = traceContext ? [{ context: traceContext }] : [];
let errored = false;
const output = yield newTrace(
{
name: ctx.flow.name,
labels: {
[SPAN_TYPE_ATTR]: "flow"
},
links: ctxLinks
},
(metadata, rootSpan) => __async(this, null, function* () {
ctx.state.executions.push({
startTime: Date.now(),
traceIds: []
});
setCustomMetadataAttribute(
metadataPrefix(`execution`),
(ctx.state.executions.length - 1).toString()
);
if (labels) {
Object.keys(labels).forEach((label) => {
setCustomMetadataAttribute(
metadataPrefix(`label:${label}`),
labels[label]
);
});
}
setCustomMetadataAttributes({
[metadataPrefix("name")]: this.name,
[metadataPrefix("id")]: ctx.flowId
});
ctx.getCurrentExecution().traceIds.push(rootSpan.spanContext().traceId);
if (!traceContext) {
ctx.state.traceContext = JSON.stringify(rootSpan.spanContext());
}
setCustomMetadataAttribute(
metadataPrefix("dispatchType"),
dispatchType
);
try {
const input = this.inputSchema ? this.inputSchema.parse(ctx.state.input) : ctx.state.input;
metadata.input = input;
const output2 = yield handler(input, streamingCallback);
metadata.output = JSON.stringify(output2);
setCustomMetadataAttribute(metadataPrefix("state"), "done");
return output2;
} catch (e) {
if (e instanceof InterruptError) {
setCustomMetadataAttribute(
metadataPrefix("state"),
"interrupted"
);
} else {
metadata.state = "error";
rootSpan.setStatus({
code: SpanStatusCode.ERROR,
message: getErrorMessage(e)
});
if (e instanceof Error) {
rootSpan.recordException(e);
}
setCustomMetadataAttribute(metadataPrefix("state"), "error");
ctx.state.operation.done = true;
ctx.state.operation.result = {
error: getErrorMessage(e),
stacktrace: getErrorStack(e)
};
}
errored = true;
}
})
);
if (!errored) {
ctx.state.operation.done = true;
ctx.state.operation.result = { response: output };
}
}));
});
}
durableExpressHandler(req, res) {
return __async(this, null, function* () {
if (req.query.stream === "true") {
const respBody = {
error: {
status: "INVALID_ARGUMENT",
message: "Output from durable flows cannot be streamed"
}
};
res.status(400).send(respBody).end();
return;
}
let data = req.body;
if (req.body.data) {
data = req.body.data;
}
const envMsg = FlowInvokeEnvelopeMessageSchema.parse(data);
try {
const state = yield this.runEnvelope(envMsg);
res.status(200).send(state.operation).end();
} catch (e) {
const respBody = {
done: true,
result: {
error: getErrorMessage(e),
stacktrace: getErrorStack(e)
}
};
res.status(500).send(respBody).end();
}
});
}
nonDurableExpressHandler(req, res) {
return __async(this, null, function* () {
var _a, _b, _c, _d;
const { stream } = req.query;
const auth = req.auth;
let input = req.body.data;
try {
yield (_a = this.authPolicy) == null ? void 0 : _a.call(this, auth, input);
} catch (e) {
const respBody = {
error: {
status: "PERMISSION_DENIED",
message: e.message || "Permission denied to resource"
}
};
res.status(403).send(respBody).end();
return;
}
if (stream === "true") {
res.writeHead(200, {
"Content-Type": "text/plain",
"Transfer-Encoding": "chunked"
});
try {
const state = yield this.runDirectly(input, {
streamingCallback: (chunk) => {
res.write(JSON.stringify(chunk) + streamDelimiter);
},
auth
});
res.write(JSON.stringify(state.operation));
res.end();
} catch (e) {
const respBody = {
done: true,
result: {
error: getErrorMessage(e),
stacktrace: getErrorStack(e)
}
};
res.write(JSON.stringify(respBody));
res.end();
}
} else {
try {
const state = yield this.runDirectly(input, { auth });
if ((_b = state.operation.result) == null ? void 0 : _b.error) {
throw new Error((_c = state.operation.result) == null ? void 0 : _c.error);
}
res.status(200).send({
result: (_d = state.operation.result) == null ? void 0 : _d.response
}).end();
} catch (e) {
res.status(500).send({
error: {
status: "INTERNAL",
message: getErrorMessage(e),
details: getErrorStack(e)
}
}).end();
}
}
});
}
get expressHandler() {
return this.experimentalDurable ? this.durableExpressHandler.bind(this) : this.nonDurableExpressHandler.bind(this);
}
}
function runFlow(flow, payload, opts) {
return __async(this, null, function* () {
var _a, _b, _c, _d, _e;
if (!(flow instanceof Flow)) {
flow = flow.flow;
}
const input = flow.inputSchema ? flow.inputSchema.parse(payload) : payload;
yield (_a = flow.authPolicy) == null ? void 0 : _a.call(flow, opts == null ? void 0 : opts.withLocalAuthContext, payload);
if (flow.middleware) {
logger.warn(
`Flow (${flow.name}) middleware won't run when invoked with runFlow.`
);
}
const state = yield flow.runEnvelope(
{
start: {
input
}
},
void 0,
opts == null ? void 0 : opts.withLocalAuthContext
);
if (!state.operation.done) {
throw new FlowStillRunningError(
`flow ${state.name} did not finish execution`
);
}
if ((_b = state.operation.result) == null ? void 0 : _b.error) {
throw new FlowExecutionError(
state.operation.name,
(_c = state.operation.result) == null ? void 0 : _c.error,
(_d = state.operation.result) == null ? void 0 : _d.stacktrace
);
}
return (_e = state.operation.result) == null ? void 0 : _e.response;
});
}
function streamFlow(flowOrFlowWrapper, payload, opts) {
var _a, _b;
const flow = !(flowOrFlowWrapper instanceof Flow) ? flowOrFlowWrapper.flow : flowOrFlowWrapper;
let chunkStreamController;
const chunkStream = new ReadableStream({
start(controller) {
chunkStreamController = controller;
},
pull() {
},
cancel() {
}
});
const authPromise = (_b = (_a = flow.authPolicy) == null ? void 0 : _a.call(flow, opts == null ? void 0 : opts.withLocalAuthContext, payload)) != null ? _b : Promise.resolve();
const operationPromise = authPromise.then(
() => flow.runEnvelope(
{
start: {
input: flow.inputSchema ? flow.inputSchema.parse(payload) : payload
}
},
(c) => {
chunkStreamController.enqueue(c);
},
opts == null ? void 0 : opts.withLocalAuthContext
)
).then((s) => s.operation).finally(() => {
chunkStreamController.close();
});
return {
output() {
return operationPromise.then((op) => {
var _a2, _b2, _c2, _d;
if (!op.done) {
throw new FlowStillRunningError(
`flow ${op.name} did not finish execution`
);
}
if ((_a2 = op.result) == null ? void 0 : _a2.error) {
throw new FlowExecutionError(
op.name,
(_b2 = op.result) == null ? void 0 : _b2.error,
(_c2 = op.result) == null ? void 0 : _c2.stacktrace
);
}
return (_d = op.result) == null ? void 0 : _d.response;
});
},
stream() {
return __asyncGenerator(this, null, function* () {
const reader = chunkStream.getReader();
while (true) {
const chunk = yield new __await(reader.read());
if (chunk.value) {
yield chunk.value;
}
if (chunk.done) {
break;
}
}
return yield new __await(operationPromise);
});
}
};
}
function createNewState(flowId, name, input) {
return {
flowId,
name,
startTime: Date.now(),
input,
cache: {},
eventsTriggered: {},
blockedOnStep: null,
executions: [],
operation: {
name: flowId,
done: false
}
};
}
function wrapAsAction(flow) {
return defineAction(
{
actionType: "flow",
name: flow.name,
inputSchema: FlowActionInputSchema,
outputSchema: FlowStateSchema,
metadata: {
inputSchema: toJsonSchema({ schema: flow.inputSchema }),
outputSchema: toJsonSchema({ schema: flow.outputSchema }),
experimentalDurable: !!flow.experimentalDurable,
requiresAuth: !!flow.authPolicy
}
},
(envelope) => __async(this, null, function* () {
var _a, _b;
yield (_b = flow.authPolicy) == null ? void 0 : _b.call(
flow,
envelope.auth,
(_a = envelope.start) == null ? void 0 : _a.input
);
setCustomMetadataAttribute(metadataPrefix("wrapperAction"), "true");
return yield flow.runEnvelope(
envelope,
getStreamingCallback(),
envelope.auth
);
})
);
}
function startFlowsServer(params) {
var _a;
const port = (params == null ? void 0 : params.port) || (process.env.PORT ? parseInt(process.env.PORT) : 0) || 3400;
const pathPrefix = (_a = params == null ? void 0 : params.pathPrefix) != null ? _a : "";
const app = express();
app.use(bodyParser.json(params == null ? void 0 : params.jsonParserOptions));
app.use(cors(params == null ? void 0 : params.cors));
const flows = (params == null ? void 0 : params.flows) || createdFlows();
logger.info(`Starting flows server on port ${port}`);
flows.forEach((f) => {
var _a2;
const flowPath = `/${pathPrefix}${f.name}`;
logger.info(` - ${flowPath}`);
(_a2 = f.middleware) == null ? void 0 : _a2.forEach((m) => {
app.post(flowPath, m);
});
app.post(flowPath, f.expressHandler);
});
app.listen(port, () => {
console.log(`Flows server listening on port ${port}`);
});
}
export {
Flow,
defineFlow,
runFlow,
startFlowsServer,
streamFlow
};
//# sourceMappingURL=flow.mjs.map