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

Genkit AI framework workflow APIs.

{"version":3,"sources":["../src/context.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 { FlowState, FlowStateExecution, Operation } from '@genkit-ai/core';\nimport { toJsonSchema } from '@genkit-ai/core/schema';\nimport {\n  SPAN_TYPE_ATTR,\n  runInNewSpan,\n  setCustomMetadataAttribute,\n  setCustomMetadataAttributes,\n} from '@genkit-ai/core/tracing';\nimport { logger } from 'firebase-functions/v1';\nimport { z } from 'zod';\nimport { InterruptError } from './errors.js';\nimport { Flow, RunStepConfig } from './flow.js';\nimport { metadataPrefix } from './utils.js';\n\n/**\n * Context object encapsulates flow execution state at runtime.\n */\nexport class Context<\n  I extends z.ZodTypeAny,\n  O extends z.ZodTypeAny,\n  S extends z.ZodTypeAny,\n> {\n  private seenSteps: Record<string, number> = {};\n\n  constructor(\n    readonly flow: Flow<I, O, S>,\n    readonly flowId: string,\n    readonly state: FlowState,\n    readonly auth?: unknown\n  ) {}\n\n  private isCached(stepName: string): boolean {\n    return this.state.cache.hasOwnProperty(stepName);\n  }\n  private getCached<T>(stepName: string): T {\n    return this.state.cache[stepName].value;\n  }\n  private updateCachedValue(stepName: string, value: any) {\n    this.state.cache[stepName] = value ? { value } : { empty: true };\n  }\n\n  private async memoize<T>(\n    stepName: string,\n    func: () => Promise<T>\n  ): Promise<[T, boolean]> {\n    if (this.isCached(stepName)) {\n      return [this.getCached(stepName), true];\n    }\n    const value = await func();\n    this.updateCachedValue(stepName, value);\n    return [value, false];\n  }\n\n  async saveState() {\n    if (this.flow.stateStore) {\n      await (await this.flow.stateStore()).save(this.flowId, this.state);\n    }\n  }\n\n  // Runs provided function in the current context. The config can specify retry and other behaviors.\n  async run<T>(\n    config: RunStepConfig,\n    input: any | undefined,\n    func: () => Promise<T>\n  ): Promise<T> {\n    return await runInNewSpan(\n      {\n        metadata: {\n          name: config.name,\n        },\n        labels: {\n          [SPAN_TYPE_ATTR]: 'flowStep',\n        },\n      },\n      async (metadata, _, isRoot) => {\n        const stepName = this.resolveStepName(config.name);\n        setCustomMetadataAttributes({\n          [metadataPrefix('stepType')]: 'run',\n          [metadataPrefix('stepName')]: config.name,\n          [metadataPrefix('resolvedStepName')]: stepName,\n        });\n        if (input !== undefined) {\n          metadata.input = input;\n        }\n        const [value, wasCached] = isRoot\n          ? await this.memoize(stepName, func)\n          : [await func(), false];\n        if (wasCached) {\n          setCustomMetadataAttribute(metadataPrefix('state'), 'cached');\n        } else {\n          setCustomMetadataAttribute(metadataPrefix('state'), 'run');\n          if (value !== undefined) {\n            metadata.output = JSON.stringify(value);\n          }\n        }\n        return value;\n      }\n    );\n  }\n\n  private resolveStepName(name: string) {\n    if (this.seenSteps[name] !== undefined) {\n      this.seenSteps[name]++;\n      name += `-${this.seenSteps[name]}`;\n    } else {\n      this.seenSteps[name] = 0;\n    }\n    return name;\n  }\n\n  // Executes interrupt step in the current context.\n  async interrupt<I extends z.ZodTypeAny, O>(\n    stepName: string,\n    func: (payload: I) => Promise<O>,\n    responseSchema: I | null,\n    skipCache?: boolean\n  ): Promise<O> {\n    return await runInNewSpan(\n      {\n        metadata: {\n          name: stepName,\n        },\n        labels: {\n          [SPAN_TYPE_ATTR]: 'flowStep',\n        },\n      },\n      async (metadata) => {\n        const resolvedStepName = this.resolveStepName(stepName);\n        setCustomMetadataAttributes({\n          [metadataPrefix('stepType')]: 'interrupt',\n          [metadataPrefix('stepName')]: stepName,\n          [metadataPrefix('resolvedStepName')]: resolvedStepName,\n        });\n        if (!skipCache && this.isCached(resolvedStepName)) {\n          setCustomMetadataAttribute(metadataPrefix('state'), 'skipped');\n          return this.getCached(resolvedStepName);\n        }\n        // TODO: refactor this.\n        if (this.state.eventsTriggered.hasOwnProperty(resolvedStepName)) {\n          let value;\n          try {\n            value = await func(\n              this.state.eventsTriggered[resolvedStepName] as I\n            );\n          } catch (e) {\n            if (e instanceof InterruptError) {\n              setCustomMetadataAttribute(metadataPrefix('state'), 'interrupt');\n            } else {\n              setCustomMetadataAttribute(metadataPrefix('state'), 'error');\n            }\n            throw e;\n          }\n          this.state.blockedOnStep = null;\n          if (!skipCache) {\n            this.updateCachedValue(resolvedStepName, value);\n          }\n          setCustomMetadataAttribute(metadataPrefix('state'), 'dispatch');\n          if (value !== undefined) {\n            metadata.output = JSON.stringify(value);\n          }\n          return value;\n        }\n        logger.debug('blockedOnStep', resolvedStepName);\n        this.state.blockedOnStep = { name: resolvedStepName };\n        if (responseSchema) {\n          this.state.blockedOnStep.schema = JSON.stringify(\n            toJsonSchema({ schema: responseSchema })\n          );\n        }\n        setCustomMetadataAttribute(metadataPrefix('state'), 'interrupted');\n        throw new InterruptError();\n      }\n    );\n  }\n\n  // Sleep for the specified number of seconds.\n  async sleep<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(\n    stepName: string,\n    seconds: number\n  ): Promise<O> {\n    const resolvedStepName = this.resolveStepName(stepName);\n    if (this.isCached(resolvedStepName)) {\n      setCustomMetadataAttribute(metadataPrefix('state'), 'skipped');\n      return this.getCached(resolvedStepName);\n    }\n\n    await this.flow.scheduler(\n      this.flow,\n      {\n        runScheduled: {\n          flowId: this.flowId,\n        },\n      },\n      seconds\n    );\n    this.updateCachedValue(resolvedStepName, undefined);\n    return this.interrupt(\n      stepName,\n      (input: z.infer<I>): z.infer<O> => input,\n      null\n    );\n  }\n\n  /**\n   * Wait for the provided flow to complete execution. This will do a poll.\n   * Poll will be done with an exponential backoff (configurable).\n   */\n  async waitFor(opts: {\n    flow: Flow<z.ZodTypeAny, z.ZodTypeAny, z.ZodTypeAny>;\n    stepName: string;\n    flowIds: string[];\n    pollingConfig?: PollingConfig;\n  }): Promise<Operation[]> {\n    const resolvedStepName = this.resolveStepName(opts.stepName);\n    if (this.isCached(resolvedStepName)) {\n      return this.getCached(resolvedStepName);\n    }\n    const states = await this.getFlowsOperations(opts.flow, opts.flowIds);\n    if (states.includes(undefined)) {\n      throw new Error(\n        'Unable to resolve flow state for ' +\n          opts.flowIds[states.indexOf(undefined)]\n      );\n    }\n    const ops = states.map((s) => s!.operation);\n    if (ops.map((op) => op.done).reduce((a, b) => a && b)) {\n      // all done.\n      this.updateCachedValue(resolvedStepName, states);\n      return ops;\n    }\n    await this.flow.scheduler(\n      this.flow,\n      {\n        runScheduled: {\n          flowId: this.flowId,\n        },\n      },\n      opts.pollingConfig?.interval || 5\n    );\n    throw new InterruptError();\n  }\n\n  private async getFlowsOperations(\n    flow: Flow<z.ZodTypeAny, z.ZodTypeAny, z.ZodTypeAny>,\n    flowIds: string[]\n  ): Promise<(FlowState | undefined)[]> {\n    return await Promise.all(\n      flowIds.map(async (id) => {\n        if (!flow.stateStore) {\n          throw new Error('Flow state store must be configured');\n        }\n        return (await flow.stateStore()).load(id);\n      })\n    );\n  }\n\n  /**\n   * Returns current active execution state.\n   */\n  getCurrentExecution(): FlowStateExecution {\n    return this.state.executions[this.state.executions.length - 1];\n  }\n}\n\nexport interface PollingConfig {\n  // TODO: add more options\n  interval: number;\n}\n"],"mappings":";;;AAiBA,SAAS,oBAAoB;AAC7B;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,EACA;AAAA,OACK;AACP,SAAS,cAAc;AAEvB,SAAS,sBAAsB;AAE/B,SAAS,sBAAsB;AAKxB,MAAM,QAIX;AAAA,EAGA,YACW,MACA,QACA,OACA,MACT;AAJS;AACA;AACA;AACA;AANX,SAAQ,YAAoC,CAAC;AAAA,EAO1C;AAAA,EAEK,SAAS,UAA2B;AAC1C,WAAO,KAAK,MAAM,MAAM,eAAe,QAAQ;AAAA,EACjD;AAAA,EACQ,UAAa,UAAqB;AACxC,WAAO,KAAK,MAAM,MAAM,QAAQ,EAAE;AAAA,EACpC;AAAA,EACQ,kBAAkB,UAAkB,OAAY;AACtD,SAAK,MAAM,MAAM,QAAQ,IAAI,QAAQ,EAAE,MAAM,IAAI,EAAE,OAAO,KAAK;AAAA,EACjE;AAAA,EAEc,QACZ,UACA,MACuB;AAAA;AACvB,UAAI,KAAK,SAAS,QAAQ,GAAG;AAC3B,eAAO,CAAC,KAAK,UAAU,QAAQ,GAAG,IAAI;AAAA,MACxC;AACA,YAAM,QAAQ,MAAM,KAAK;AACzB,WAAK,kBAAkB,UAAU,KAAK;AACtC,aAAO,CAAC,OAAO,KAAK;AAAA,IACtB;AAAA;AAAA,EAEM,YAAY;AAAA;AAChB,UAAI,KAAK,KAAK,YAAY;AACxB,eAAO,MAAM,KAAK,KAAK,WAAW,GAAG,KAAK,KAAK,QAAQ,KAAK,KAAK;AAAA,MACnE;AAAA,IACF;AAAA;AAAA;AAAA,EAGM,IACJ,QACA,OACA,MACY;AAAA;AACZ,aAAO,MAAM;AAAA,QACX;AAAA,UACE,UAAU;AAAA,YACR,MAAM,OAAO;AAAA,UACf;AAAA,UACA,QAAQ;AAAA,YACN,CAAC,cAAc,GAAG;AAAA,UACpB;AAAA,QACF;AAAA,QACA,CAAO,UAAU,GAAG,WAAW;AAC7B,gBAAM,WAAW,KAAK,gBAAgB,OAAO,IAAI;AACjD,sCAA4B;AAAA,YAC1B,CAAC,eAAe,UAAU,CAAC,GAAG;AAAA,YAC9B,CAAC,eAAe,UAAU,CAAC,GAAG,OAAO;AAAA,YACrC,CAAC,eAAe,kBAAkB,CAAC,GAAG;AAAA,UACxC,CAAC;AACD,cAAI,UAAU,QAAW;AACvB,qBAAS,QAAQ;AAAA,UACnB;AACA,gBAAM,CAAC,OAAO,SAAS,IAAI,SACvB,MAAM,KAAK,QAAQ,UAAU,IAAI,IACjC,CAAC,MAAM,KAAK,GAAG,KAAK;AACxB,cAAI,WAAW;AACb,uCAA2B,eAAe,OAAO,GAAG,QAAQ;AAAA,UAC9D,OAAO;AACL,uCAA2B,eAAe,OAAO,GAAG,KAAK;AACzD,gBAAI,UAAU,QAAW;AACvB,uBAAS,SAAS,KAAK,UAAU,KAAK;AAAA,YACxC;AAAA,UACF;AACA,iBAAO;AAAA,QACT;AAAA,MACF;AAAA,IACF;AAAA;AAAA,EAEQ,gBAAgB,MAAc;AACpC,QAAI,KAAK,UAAU,IAAI,MAAM,QAAW;AACtC,WAAK,UAAU,IAAI;AACnB,cAAQ,IAAI,KAAK,UAAU,IAAI,CAAC;AAAA,IAClC,OAAO;AACL,WAAK,UAAU,IAAI,IAAI;AAAA,IACzB;AACA,WAAO;AAAA,EACT;AAAA;AAAA,EAGM,UACJ,UACA,MACA,gBACA,WACY;AAAA;AACZ,aAAO,MAAM;AAAA,QACX;AAAA,UACE,UAAU;AAAA,YACR,MAAM;AAAA,UACR;AAAA,UACA,QAAQ;AAAA,YACN,CAAC,cAAc,GAAG;AAAA,UACpB;AAAA,QACF;AAAA,QACA,CAAO,aAAa;AAClB,gBAAM,mBAAmB,KAAK,gBAAgB,QAAQ;AACtD,sCAA4B;AAAA,YAC1B,CAAC,eAAe,UAAU,CAAC,GAAG;AAAA,YAC9B,CAAC,eAAe,UAAU,CAAC,GAAG;AAAA,YAC9B,CAAC,eAAe,kBAAkB,CAAC,GAAG;AAAA,UACxC,CAAC;AACD,cAAI,CAAC,aAAa,KAAK,SAAS,gBAAgB,GAAG;AACjD,uCAA2B,eAAe,OAAO,GAAG,SAAS;AAC7D,mBAAO,KAAK,UAAU,gBAAgB;AAAA,UACxC;AAEA,cAAI,KAAK,MAAM,gBAAgB,eAAe,gBAAgB,GAAG;AAC/D,gBAAI;AACJ,gBAAI;AACF,sBAAQ,MAAM;AAAA,gBACZ,KAAK,MAAM,gBAAgB,gBAAgB;AAAA,cAC7C;AAAA,YACF,SAAS,GAAG;AACV,kBAAI,aAAa,gBAAgB;AAC/B,2CAA2B,eAAe,OAAO,GAAG,WAAW;AAAA,cACjE,OAAO;AACL,2CAA2B,eAAe,OAAO,GAAG,OAAO;AAAA,cAC7D;AACA,oBAAM;AAAA,YACR;AACA,iBAAK,MAAM,gBAAgB;AAC3B,gBAAI,CAAC,WAAW;AACd,mBAAK,kBAAkB,kBAAkB,KAAK;AAAA,YAChD;AACA,uCAA2B,eAAe,OAAO,GAAG,UAAU;AAC9D,gBAAI,UAAU,QAAW;AACvB,uBAAS,SAAS,KAAK,UAAU,KAAK;AAAA,YACxC;AACA,mBAAO;AAAA,UACT;AACA,iBAAO,MAAM,iBAAiB,gBAAgB;AAC9C,eAAK,MAAM,gBAAgB,EAAE,MAAM,iBAAiB;AACpD,cAAI,gBAAgB;AAClB,iBAAK,MAAM,cAAc,SAAS,KAAK;AAAA,cACrC,aAAa,EAAE,QAAQ,eAAe,CAAC;AAAA,YACzC;AAAA,UACF;AACA,qCAA2B,eAAe,OAAO,GAAG,aAAa;AACjE,gBAAM,IAAI,eAAe;AAAA,QAC3B;AAAA,MACF;AAAA,IACF;AAAA;AAAA;AAAA,EAGM,MACJ,UACA,SACY;AAAA;AACZ,YAAM,mBAAmB,KAAK,gBAAgB,QAAQ;AACtD,UAAI,KAAK,SAAS,gBAAgB,GAAG;AACnC,mCAA2B,eAAe,OAAO,GAAG,SAAS;AAC7D,eAAO,KAAK,UAAU,gBAAgB;AAAA,MACxC;AAEA,YAAM,KAAK,KAAK;AAAA,QACd,KAAK;AAAA,QACL;AAAA,UACE,cAAc;AAAA,YACZ,QAAQ,KAAK;AAAA,UACf;AAAA,QACF;AAAA,QACA;AAAA,MACF;AACA,WAAK,kBAAkB,kBAAkB,MAAS;AAClD,aAAO,KAAK;AAAA,QACV;AAAA,QACA,CAAC,UAAkC;AAAA,QACnC;AAAA,MACF;AAAA,IACF;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAMM,QAAQ,MAKW;AAAA;AApO3B;AAqOI,YAAM,mBAAmB,KAAK,gBAAgB,KAAK,QAAQ;AAC3D,UAAI,KAAK,SAAS,gBAAgB,GAAG;AACnC,eAAO,KAAK,UAAU,gBAAgB;AAAA,MACxC;AACA,YAAM,SAAS,MAAM,KAAK,mBAAmB,KAAK,MAAM,KAAK,OAAO;AACpE,UAAI,OAAO,SAAS,MAAS,GAAG;AAC9B,cAAM,IAAI;AAAA,UACR,sCACE,KAAK,QAAQ,OAAO,QAAQ,MAAS,CAAC;AAAA,QAC1C;AAAA,MACF;AACA,YAAM,MAAM,OAAO,IAAI,CAAC,MAAM,EAAG,SAAS;AAC1C,UAAI,IAAI,IAAI,CAAC,OAAO,GAAG,IAAI,EAAE,OAAO,CAAC,GAAG,MAAM,KAAK,CAAC,GAAG;AAErD,aAAK,kBAAkB,kBAAkB,MAAM;AAC/C,eAAO;AAAA,MACT;AACA,YAAM,KAAK,KAAK;AAAA,QACd,KAAK;AAAA,QACL;AAAA,UACE,cAAc;AAAA,YACZ,QAAQ,KAAK;AAAA,UACf;AAAA,QACF;AAAA,UACA,UAAK,kBAAL,mBAAoB,aAAY;AAAA,MAClC;AACA,YAAM,IAAI,eAAe;AAAA,IAC3B;AAAA;AAAA,EAEc,mBACZ,MACA,SACoC;AAAA;AACpC,aAAO,MAAM,QAAQ;AAAA,QACnB,QAAQ,IAAI,CAAO,OAAO;AACxB,cAAI,CAAC,KAAK,YAAY;AACpB,kBAAM,IAAI,MAAM,qCAAqC;AAAA,UACvD;AACA,kBAAQ,MAAM,KAAK,WAAW,GAAG,KAAK,EAAE;AAAA,QAC1C,EAAC;AAAA,MACH;AAAA,IACF;AAAA;AAAA;AAAA;AAAA;AAAA,EAKA,sBAA0C;AACxC,WAAO,KAAK,MAAM,WAAW,KAAK,MAAM,WAAW,SAAS,CAAC;AAAA,EAC/D;AACF;","names":[]}