@genkit-ai/core/lib/flow.mjs.map

Genkit AI framework core libraries.

{"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 { AsyncLocalStorage } from 'node:async_hooks';\nimport type { z } from 'zod';\nimport { ActionFnArg, defineAction, 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 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 resolvedConfig: FlowConfig<I, O, S> =\n    typeof config === 'string' ? { name: config } : config;\n\n  return defineFlowAction(registry, resolvedConfig, fn);\n}\n\n/**\n * Registers a flow as an action in the registry.\n */\nfunction defineFlowAction<\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>,\n  fn: FlowFn<I, O, S>\n): Flow<I, O, S> {\n  return defineAction(\n    registry,\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      return await legacyRegistryAls.run(registry, () => {\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}\n\nconst legacyRegistryAls = new AsyncLocalStorage<Registry>();\n\nexport function run<T>(\n  name: string,\n  func: () => Promise<T>,\n  registry?: Registry\n): Promise<T>;\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  maybeRegistry?: Registry | HasRegistry\n): Promise<T> {\n  let func;\n  let input;\n  let registry: Registry | undefined;\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  } else if (\n    fnOrRegistry instanceof Registry ||\n    (fnOrRegistry as HasRegistry)?.registry\n  ) {\n    registry = (fnOrRegistry as HasRegistry)?.registry\n      ? (fnOrRegistry as HasRegistry)?.registry\n      : (fnOrRegistry as Registry);\n  }\n  if (maybeRegistry) {\n    registry = (maybeRegistry as HasRegistry).registry\n      ? (maybeRegistry as HasRegistry).registry\n      : (maybeRegistry as Registry);\n  }\n\n  if (!registry) {\n    registry = legacyRegistryAls.getStore();\n  }\n  if (!registry) {\n    throw new Error(\n      'Unable to resolve registry. Consider explicitly passing Genkit instance.'\n    );\n  }\n\n  if (!func) {\n    throw new Error('unable to resolve run function');\n  }\n  return runInNewSpan(\n    registry,\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":"AAgBA,SAAS,yBAAyB;AAElC,SAAsB,oBAAiC;AACvD,SAAS,gBAAkC;AAC3C,SAAS,gBAAgB,oBAAoB;AAyDtC,SAAS,WAKd,UACA,QACA,IACe;AACf,QAAM,iBACJ,OAAO,WAAW,WAAW,EAAE,MAAM,OAAO,IAAI;AAElD,SAAO,iBAAiB,UAAU,gBAAgB,EAAE;AACtD;AAKA,SAAS,iBAKP,UACA,QACA,IACe;AACf,SAAO;AAAA,IACL;AAAA,IACA;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,aAAO,MAAM,kBAAkB,IAAI,UAAU,MAAM;AACjD,cAAM,MAAM;AACZ,QAAC,IAAoC,YAAY;AACjD,QAAC,IAAoC,UAAU;AAC/C,QAAC,IAAoC,QAAQ;AAC7C,QAAC,IAAoC,cAAc;AACnD,QAAC,IAAoC,qBACnC;AACF,eAAO,GAAG,OAAO,GAAkC;AAAA,MACrD,CAAC;AAAA,IACH;AAAA,EACF;AACF;AAEA,MAAM,oBAAoB,IAAI,kBAA4B;AAiBnD,SAAS,IACd,MACA,aACA,cACA,eACY;AACZ,MAAI;AACJ,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,WACE,wBAAwB,YACvB,cAA8B,UAC/B;AACA,eAAY,cAA8B,WACrC,cAA8B,WAC9B;AAAA,EACP;AACA,MAAI,eAAe;AACjB,eAAY,cAA8B,WACrC,cAA8B,WAC9B;AAAA,EACP;AAEA,MAAI,CAAC,UAAU;AACb,eAAW,kBAAkB,SAAS;AAAA,EACxC;AACA,MAAI,CAAC,UAAU;AACb,UAAM,IAAI;AAAA,MACR;AAAA,IACF;AAAA,EACF;AAEA,MAAI,CAAC,MAAM;AACT,UAAM,IAAI,MAAM,gCAAgC;AAAA,EAClD;AACA,SAAO;AAAA,IACL;AAAA,IACA;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":[]}