@genkit-ai/core
Version:
Genkit AI framework core libraries.
1 lines • 6.84 kB
Source Map (JSON)
{"version":3,"sources":["../src/flow.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 type { z } from 'zod';\nimport { ActionFnArg, action, type Action } from './action.js';\nimport { Registry, type HasRegistry } from './registry.js';\nimport { SPAN_TYPE_ATTR, runInNewSpan } from './tracing.js';\n\n/**\n * Flow is an observable, streamable, (optionally) strongly typed function.\n */\nexport interface Flow<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n> extends Action<I, O, S> {}\n\n/**\n * Configuration for a streaming flow.\n */\nexport interface FlowConfig<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n> {\n /** Name of the flow. */\n name: string;\n /** Schema of the input to the flow. */\n inputSchema?: I;\n /** Schema of the output from the flow. */\n outputSchema?: O;\n /** Schema of the streaming chunks from the flow. */\n streamSchema?: S;\n /** Metadata of the flow used by tooling. */\n metadata?: Record<string, any>;\n}\n\n/**\n * Flow execution context for flow to access the streaming callback and\n * side-channel context data. The context itself is a function, a short-cut\n * for streaming callback.\n */\nexport interface FlowSideChannel<S> extends ActionFnArg<S> {\n (chunk: S): void;\n}\n\n/**\n * Function to be executed in the flow.\n */\nexport type FlowFn<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n> = (\n /** Input to the flow. */\n input: z.infer<I>,\n /** Callback for streaming functions only. */\n streamingCallback: FlowSideChannel<z.infer<S>>\n) => Promise<z.infer<O>> | z.infer<O>;\n\n/**\n * Defines a flow. This operates on the currently active registry.\n */\nexport function flow<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n>(config: FlowConfig<I, O, S> | string, fn: FlowFn<I, O, S>): Flow<I, O, S> {\n const resolvedConfig: FlowConfig<I, O, S> =\n typeof config === 'string' ? { name: config } : config;\n\n return flowAction(resolvedConfig, fn);\n}\n\n/**\n * Defines a non-streaming flow. This operates on the currently active registry.\n */\nexport function defineFlow<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n registry: Registry,\n config: FlowConfig<I, O, S> | string,\n fn: FlowFn<I, O, S>\n): Flow<I, O, S> {\n const f = flow(config, fn);\n\n registry.registerAction('flow', f);\n\n return f;\n}\n\n/**\n * Registers a flow as an action in the registry.\n */\nfunction flowAction<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n>(config: FlowConfig<I, O, S>, fn: FlowFn<I, O, S>): Flow<I, O, S> {\n return action(\n {\n actionType: 'flow',\n name: config.name,\n inputSchema: config.inputSchema,\n outputSchema: config.outputSchema,\n streamSchema: config.streamSchema,\n metadata: config.metadata,\n },\n async (\n input,\n { sendChunk, context, trace, abortSignal, streamingRequested }\n ) => {\n const ctx = sendChunk;\n (ctx as FlowSideChannel<z.infer<S>>).sendChunk = sendChunk;\n (ctx as FlowSideChannel<z.infer<S>>).context = context;\n (ctx as FlowSideChannel<z.infer<S>>).trace = trace;\n (ctx as FlowSideChannel<z.infer<S>>).abortSignal = abortSignal;\n (ctx as FlowSideChannel<z.infer<S>>).streamingRequested =\n streamingRequested;\n return fn(input, ctx as FlowSideChannel<z.infer<S>>);\n }\n );\n}\n\nexport function run<T>(\n name: string,\n func: () => Promise<T>,\n _?: Registry\n): Promise<T>;\n\nexport function run<T>(\n name: string,\n input: any,\n func: (input?: any) => Promise<T>,\n registry?: Registry\n): Promise<T>;\n\n/**\n * A flow step that executes the provided function. Each run step is recorded separately in the trace.\n */\nexport function run<T>(\n name: string,\n funcOrInput: () => Promise<T>,\n fnOrRegistry?: Registry | HasRegistry | ((input?: any) => Promise<T>),\n _?: Registry | HasRegistry\n): Promise<T> {\n let func;\n let input;\n let hasInput = false;\n if (typeof funcOrInput === 'function') {\n func = funcOrInput;\n } else {\n input = funcOrInput;\n hasInput = true;\n }\n if (typeof fnOrRegistry === 'function') {\n func = fnOrRegistry;\n }\n\n if (!func) {\n throw new Error('unable to resolve run function');\n }\n return runInNewSpan(\n {\n metadata: { name },\n labels: {\n [SPAN_TYPE_ATTR]: 'flowStep',\n },\n },\n async (meta) => {\n meta.input = input;\n const output = hasInput ? await func(input) : await func();\n meta.output = JSON.stringify(output);\n return output;\n }\n );\n}\n"],"mappings":"AAiBA,SAAsB,cAA2B;AAEjD,SAAS,gBAAgB,oBAAoB;AAyDtC,SAAS,KAId,QAAsC,IAAoC;AAC1E,QAAM,iBACJ,OAAO,WAAW,WAAW,EAAE,MAAM,OAAO,IAAI;AAElD,SAAO,WAAW,gBAAgB,EAAE;AACtC;AAKO,SAAS,WAKd,UACA,QACA,IACe;AACf,QAAM,IAAI,KAAK,QAAQ,EAAE;AAEzB,WAAS,eAAe,QAAQ,CAAC;AAEjC,SAAO;AACT;AAKA,SAAS,WAIP,QAA6B,IAAoC;AACjE,SAAO;AAAA,IACL;AAAA,MACE,YAAY;AAAA,MACZ,MAAM,OAAO;AAAA,MACb,aAAa,OAAO;AAAA,MACpB,cAAc,OAAO;AAAA,MACrB,cAAc,OAAO;AAAA,MACrB,UAAU,OAAO;AAAA,IACnB;AAAA,IACA,OACE,OACA,EAAE,WAAW,SAAS,OAAO,aAAa,mBAAmB,MAC1D;AACH,YAAM,MAAM;AACZ,MAAC,IAAoC,YAAY;AACjD,MAAC,IAAoC,UAAU;AAC/C,MAAC,IAAoC,QAAQ;AAC7C,MAAC,IAAoC,cAAc;AACnD,MAAC,IAAoC,qBACnC;AACF,aAAO,GAAG,OAAO,GAAkC;AAAA,IACrD;AAAA,EACF;AACF;AAkBO,SAAS,IACd,MACA,aACA,cACA,GACY;AACZ,MAAI;AACJ,MAAI;AACJ,MAAI,WAAW;AACf,MAAI,OAAO,gBAAgB,YAAY;AACrC,WAAO;AAAA,EACT,OAAO;AACL,YAAQ;AACR,eAAW;AAAA,EACb;AACA,MAAI,OAAO,iBAAiB,YAAY;AACtC,WAAO;AAAA,EACT;AAEA,MAAI,CAAC,MAAM;AACT,UAAM,IAAI,MAAM,gCAAgC;AAAA,EAClD;AACA,SAAO;AAAA,IACL;AAAA,MACE,UAAU,EAAE,KAAK;AAAA,MACjB,QAAQ;AAAA,QACN,CAAC,cAAc,GAAG;AAAA,MACpB;AAAA,IACF;AAAA,IACA,OAAO,SAAS;AACd,WAAK,QAAQ;AACb,YAAM,SAAS,WAAW,MAAM,KAAK,KAAK,IAAI,MAAM,KAAK;AACzD,WAAK,SAAS,KAAK,UAAU,MAAM;AACnC,aAAO;AAAA,IACT;AAAA,EACF;AACF;","names":[]}