{"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.mjs';\nimport { Registry, type HasRegistry } from './registry.mjs';\nimport { SPAN_TYPE_ATTR, runInNewSpan } from './tracing.mjs';\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":[]}