@genkit-ai/core
Version:
Genkit AI framework core libraries.
1 lines • 13.9 kB
Source Map (JSON)
{"version":3,"sources":["../src/background-action.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 { JSONSchema7 } from 'json-schema';\nimport * as z from 'zod';\nimport { Action, ActionMetadata, Middleware, action } from './action.js';\nimport { ActionContext } from './context.js';\nimport { GenkitError } from './error.js';\nimport { ActionType, Registry } from './registry.js';\nimport { toJsonSchema } from './schema.js';\n\n/**\n * Zod schema of an opration representing a background task.\n */\nexport const OperationSchema = z.object({\n action: z.string().optional(),\n id: z.string(),\n done: z.boolean().optional(),\n output: z.any().optional(),\n error: z.object({ message: z.string() }).passthrough().optional(),\n metadata: z.record(z.string(), z.any()).optional(),\n});\n\n/**\n * Background operation.\n */\nexport interface Operation<O = any> {\n action?: string;\n id: string;\n done?: boolean;\n output?: O;\n error?: { message: string; [key: string]: unknown };\n metadata?: Record<string, any>;\n}\n\n/**\n * Background action. Unlike regular action, background action can run for a long time in the background.\n * The returned operation can used to check the status of the background operation and retrieve the response.\n */\nexport interface BackgroundAction<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n RunOptions extends BackgroundActionRunOptions = BackgroundActionRunOptions,\n> {\n __action: ActionMetadata<I, O>;\n\n readonly startAction: Action<I, typeof OperationSchema>;\n readonly checkAction: Action<typeof OperationSchema, typeof OperationSchema>;\n readonly cancelAction?: Action<\n typeof OperationSchema,\n typeof OperationSchema\n >;\n readonly supportsCancel: boolean;\n\n start(\n input?: z.infer<I>,\n options?: RunOptions\n ): Promise<Operation<z.infer<O>>>;\n\n check(operation: Operation<z.infer<O>>): Promise<Operation<z.infer<O>>>;\n\n cancel(operation: Operation<z.infer<O>>): Promise<Operation<z.infer<O>>>;\n}\n\nexport async function lookupBackgroundAction<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n registry: Registry,\n key: string\n): Promise<BackgroundAction<I, O> | undefined> {\n const root: Action<I, typeof OperationSchema> = await registry.lookupAction<\n I,\n typeof OperationSchema,\n Action<I, typeof OperationSchema>\n >(key);\n if (!root) return undefined;\n const actionName = key.substring(key.indexOf('/', 1) + 1);\n return new BackgroundActionImpl(\n root,\n await registry.lookupAction<\n typeof OperationSchema,\n typeof OperationSchema,\n Action<typeof OperationSchema, typeof OperationSchema>\n >(`/check-operation/${actionName}/check`),\n await registry.lookupAction<\n typeof OperationSchema,\n typeof OperationSchema,\n Action<typeof OperationSchema, typeof OperationSchema>\n >(`/cancel-operation/${actionName}/cancel`)\n );\n}\n\nclass BackgroundActionImpl<\n I extends z.ZodTypeAny = z.ZodTypeAny,\n O extends z.ZodTypeAny = z.ZodTypeAny,\n RunOptions extends BackgroundActionRunOptions = BackgroundActionRunOptions,\n> implements BackgroundAction<I, O, RunOptions>\n{\n __action: ActionMetadata<I, O>;\n\n readonly startAction: Action<I, typeof OperationSchema>;\n readonly checkAction: Action<typeof OperationSchema, typeof OperationSchema>;\n readonly cancelAction?: Action<\n typeof OperationSchema,\n typeof OperationSchema\n >;\n\n constructor(\n startAction: Action<I, typeof OperationSchema>,\n checkAction: Action<typeof OperationSchema, typeof OperationSchema>,\n cancelAction:\n | Action<typeof OperationSchema, typeof OperationSchema>\n | undefined\n ) {\n this.__action = {\n name: startAction.__action.name,\n description: startAction.__action.description,\n inputSchema: startAction.__action.inputSchema,\n inputJsonSchema: startAction.__action.inputJsonSchema,\n metadata: startAction.__action.metadata,\n actionType: startAction.__action.actionType,\n };\n this.startAction = startAction;\n this.checkAction = checkAction;\n this.cancelAction = cancelAction;\n }\n\n async start(\n input?: z.infer<I>,\n options?: RunOptions\n ): Promise<Operation<z.infer<O>>> {\n return await this.startAction(input, options);\n }\n\n async check(\n operation: Operation<z.infer<O>>\n ): Promise<Operation<z.infer<O>>> {\n return await this.checkAction(operation);\n }\n\n get supportsCancel(): boolean {\n return !!this.cancelAction;\n }\n\n async cancel(\n operation: Operation<z.infer<O>>\n ): Promise<Operation<z.infer<O>>> {\n if (!this.cancelAction) {\n return operation;\n }\n return await this.cancelAction(operation);\n }\n}\n\n/**\n * Options (side channel) data to pass to the model.\n */\nexport interface BackgroundActionRunOptions {\n /**\n * Additional runtime context data (ex. auth context data).\n */\n context?: ActionContext;\n\n /**\n * Additional span attributes to apply to OT spans.\n */\n telemetryLabels?: Record<string, string>;\n}\n\n/**\n * Options (side channel) data to pass to the model.\n */\nexport interface BackgroundActionFnArg<S> {\n /**\n * Additional runtime context data (ex. auth context data).\n */\n context?: ActionContext;\n\n /**\n * Trace context containing trace and span IDs.\n */\n trace: {\n traceId: string;\n spanId: string;\n };\n}\n\n/**\n * Action factory params.\n */\nexport type BackgroundActionParams<\n I extends z.ZodTypeAny,\n O extends z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n> = {\n name: string;\n start: (\n input: z.infer<I>,\n options: BackgroundActionFnArg<z.infer<S>>\n ) => Promise<Operation<z.infer<O>>>;\n check: (input: Operation<z.infer<O>>) => Promise<Operation<z.infer<O>>>;\n cancel?: (input: Operation<z.infer<O>>) => Promise<Operation<z.infer<O>>>;\n actionType: ActionType;\n\n description?: string;\n inputSchema?: I;\n inputJsonSchema?: JSONSchema7;\n outputSchema?: O;\n outputJsonSchema?: JSONSchema7;\n metadata?: Record<string, any>;\n use?: Middleware<z.infer<I>, z.infer<O>>[];\n streamSchema?: S;\n};\n\n/**\n * Defines an action with the given config and registers it in the registry.\n */\nexport function defineBackgroundAction<\n I extends z.ZodTypeAny,\n O extends z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n registry: Registry,\n config: BackgroundActionParams<I, O, S>\n): BackgroundAction<I, O> {\n const act = backgroundAction(config);\n\n registerBackgroundAction(registry, act);\n\n return act;\n}\n\nexport function registerBackgroundAction(\n registry: Registry,\n act: BackgroundAction<any, any>,\n opts?: { namespace?: string }\n) {\n registry.registerAction(\n act.startAction.__action.actionType!,\n act.startAction,\n opts\n );\n registry.registerAction(\n act.checkAction.__action.actionType!,\n act.checkAction,\n opts\n );\n if (act.cancelAction) {\n registry.registerAction(\n act.cancelAction.__action.actionType!,\n act.cancelAction,\n opts\n );\n }\n}\n\n/**\n * Creates a background action with the given config.\n */\nexport function backgroundAction<\n I extends z.ZodTypeAny,\n O extends z.ZodTypeAny,\n S extends z.ZodTypeAny = z.ZodTypeAny,\n>(config: BackgroundActionParams<I, O, S>): BackgroundAction<I, O> {\n const startAction = action(\n {\n actionType: config.actionType,\n name: config.name,\n description: config.description,\n inputSchema: config.inputSchema,\n inputJsonSchema: config.inputJsonSchema,\n outputSchema: OperationSchema,\n metadata: {\n ...config.metadata,\n outputSchema: toJsonSchema({\n schema: config.outputSchema,\n jsonSchema: config.outputJsonSchema,\n }),\n },\n use: config.use,\n },\n async (input, options) => {\n const operation = await config.start(input, options);\n operation.action = `/${config.actionType}/${config.name}`;\n return operation;\n }\n );\n const checkAction = action(\n {\n actionType: 'check-operation',\n name: `${config.name}/check`,\n description: config.description,\n inputSchema: OperationSchema,\n inputJsonSchema: config.inputJsonSchema,\n outputSchema: OperationSchema,\n metadata: {\n ...config.metadata,\n outputSchema: toJsonSchema({\n schema: config.outputSchema,\n jsonSchema: config.outputJsonSchema,\n }),\n },\n },\n async (input) => {\n const operation = await config.check(input);\n operation.action = `/${config.actionType}/${config.name}`;\n return operation;\n }\n );\n let cancelAction:\n | Action<typeof OperationSchema, typeof OperationSchema>\n | undefined = undefined;\n if (config.cancel) {\n cancelAction = action(\n {\n actionType: 'cancel-operation',\n name: `${config.name}/cancel`,\n description: config.description,\n inputSchema: OperationSchema,\n inputJsonSchema: config.inputJsonSchema,\n outputSchema: OperationSchema,\n metadata: {\n ...config.metadata,\n outputSchema: toJsonSchema({\n schema: config.outputSchema,\n jsonSchema: config.outputJsonSchema,\n }),\n },\n },\n async (input) => {\n if (!config.cancel) {\n throw new GenkitError({\n status: 'UNAVAILABLE',\n message: `${config.name} does not support cancellation.`,\n });\n }\n const operation = await config.cancel(input);\n operation.action = `/${config.actionType}/${config.name}`;\n return operation;\n }\n );\n }\n\n return new BackgroundActionImpl(startAction, checkAction, cancelAction);\n}\n\nexport function isBackgroundAction(a: unknown): a is BackgroundAction {\n return a instanceof BackgroundActionImpl;\n}\n"],"mappings":"AAiBA,YAAY,OAAO;AACnB,SAA6C,cAAc;AAE3D,SAAS,mBAAmB;AAE5B,SAAS,oBAAoB;AAKtB,MAAM,kBAAkB,EAAE,OAAO;AAAA,EACtC,QAAQ,EAAE,OAAO,EAAE,SAAS;AAAA,EAC5B,IAAI,EAAE,OAAO;AAAA,EACb,MAAM,EAAE,QAAQ,EAAE,SAAS;AAAA,EAC3B,QAAQ,EAAE,IAAI,EAAE,SAAS;AAAA,EACzB,OAAO,EAAE,OAAO,EAAE,SAAS,EAAE,OAAO,EAAE,CAAC,EAAE,YAAY,EAAE,SAAS;AAAA,EAChE,UAAU,EAAE,OAAO,EAAE,OAAO,GAAG,EAAE,IAAI,CAAC,EAAE,SAAS;AACnD,CAAC;AA2CD,eAAsB,uBAIpB,UACA,KAC6C;AAC7C,QAAM,OAA0C,MAAM,SAAS,aAI7D,GAAG;AACL,MAAI,CAAC,KAAM,QAAO;AAClB,QAAM,aAAa,IAAI,UAAU,IAAI,QAAQ,KAAK,CAAC,IAAI,CAAC;AACxD,SAAO,IAAI;AAAA,IACT;AAAA,IACA,MAAM,SAAS,aAIb,oBAAoB,UAAU,QAAQ;AAAA,IACxC,MAAM,SAAS,aAIb,qBAAqB,UAAU,SAAS;AAAA,EAC5C;AACF;AAEA,MAAM,qBAKN;AAAA,EACE;AAAA,EAES;AAAA,EACA;AAAA,EACA;AAAA,EAKT,YACE,aACA,aACA,cAGA;AACA,SAAK,WAAW;AAAA,MACd,MAAM,YAAY,SAAS;AAAA,MAC3B,aAAa,YAAY,SAAS;AAAA,MAClC,aAAa,YAAY,SAAS;AAAA,MAClC,iBAAiB,YAAY,SAAS;AAAA,MACtC,UAAU,YAAY,SAAS;AAAA,MAC/B,YAAY,YAAY,SAAS;AAAA,IACnC;AACA,SAAK,cAAc;AACnB,SAAK,cAAc;AACnB,SAAK,eAAe;AAAA,EACtB;AAAA,EAEA,MAAM,MACJ,OACA,SACgC;AAChC,WAAO,MAAM,KAAK,YAAY,OAAO,OAAO;AAAA,EAC9C;AAAA,EAEA,MAAM,MACJ,WACgC;AAChC,WAAO,MAAM,KAAK,YAAY,SAAS;AAAA,EACzC;AAAA,EAEA,IAAI,iBAA0B;AAC5B,WAAO,CAAC,CAAC,KAAK;AAAA,EAChB;AAAA,EAEA,MAAM,OACJ,WACgC;AAChC,QAAI,CAAC,KAAK,cAAc;AACtB,aAAO;AAAA,IACT;AACA,WAAO,MAAM,KAAK,aAAa,SAAS;AAAA,EAC1C;AACF;AAiEO,SAAS,uBAKd,UACA,QACwB;AACxB,QAAM,MAAM,iBAAiB,MAAM;AAEnC,2BAAyB,UAAU,GAAG;AAEtC,SAAO;AACT;AAEO,SAAS,yBACd,UACA,KACA,MACA;AACA,WAAS;AAAA,IACP,IAAI,YAAY,SAAS;AAAA,IACzB,IAAI;AAAA,IACJ;AAAA,EACF;AACA,WAAS;AAAA,IACP,IAAI,YAAY,SAAS;AAAA,IACzB,IAAI;AAAA,IACJ;AAAA,EACF;AACA,MAAI,IAAI,cAAc;AACpB,aAAS;AAAA,MACP,IAAI,aAAa,SAAS;AAAA,MAC1B,IAAI;AAAA,MACJ;AAAA,IACF;AAAA,EACF;AACF;AAKO,SAAS,iBAId,QAAiE;AACjE,QAAM,cAAc;AAAA,IAClB;AAAA,MACE,YAAY,OAAO;AAAA,MACnB,MAAM,OAAO;AAAA,MACb,aAAa,OAAO;AAAA,MACpB,aAAa,OAAO;AAAA,MACpB,iBAAiB,OAAO;AAAA,MACxB,cAAc;AAAA,MACd,UAAU;AAAA,QACR,GAAG,OAAO;AAAA,QACV,cAAc,aAAa;AAAA,UACzB,QAAQ,OAAO;AAAA,UACf,YAAY,OAAO;AAAA,QACrB,CAAC;AAAA,MACH;AAAA,MACA,KAAK,OAAO;AAAA,IACd;AAAA,IACA,OAAO,OAAO,YAAY;AACxB,YAAM,YAAY,MAAM,OAAO,MAAM,OAAO,OAAO;AACnD,gBAAU,SAAS,IAAI,OAAO,UAAU,IAAI,OAAO,IAAI;AACvD,aAAO;AAAA,IACT;AAAA,EACF;AACA,QAAM,cAAc;AAAA,IAClB;AAAA,MACE,YAAY;AAAA,MACZ,MAAM,GAAG,OAAO,IAAI;AAAA,MACpB,aAAa,OAAO;AAAA,MACpB,aAAa;AAAA,MACb,iBAAiB,OAAO;AAAA,MACxB,cAAc;AAAA,MACd,UAAU;AAAA,QACR,GAAG,OAAO;AAAA,QACV,cAAc,aAAa;AAAA,UACzB,QAAQ,OAAO;AAAA,UACf,YAAY,OAAO;AAAA,QACrB,CAAC;AAAA,MACH;AAAA,IACF;AAAA,IACA,OAAO,UAAU;AACf,YAAM,YAAY,MAAM,OAAO,MAAM,KAAK;AAC1C,gBAAU,SAAS,IAAI,OAAO,UAAU,IAAI,OAAO,IAAI;AACvD,aAAO;AAAA,IACT;AAAA,EACF;AACA,MAAI,eAEY;AAChB,MAAI,OAAO,QAAQ;AACjB,mBAAe;AAAA,MACb;AAAA,QACE,YAAY;AAAA,QACZ,MAAM,GAAG,OAAO,IAAI;AAAA,QACpB,aAAa,OAAO;AAAA,QACpB,aAAa;AAAA,QACb,iBAAiB,OAAO;AAAA,QACxB,cAAc;AAAA,QACd,UAAU;AAAA,UACR,GAAG,OAAO;AAAA,UACV,cAAc,aAAa;AAAA,YACzB,QAAQ,OAAO;AAAA,YACf,YAAY,OAAO;AAAA,UACrB,CAAC;AAAA,QACH;AAAA,MACF;AAAA,MACA,OAAO,UAAU;AACf,YAAI,CAAC,OAAO,QAAQ;AAClB,gBAAM,IAAI,YAAY;AAAA,YACpB,QAAQ;AAAA,YACR,SAAS,GAAG,OAAO,IAAI;AAAA,UACzB,CAAC;AAAA,QACH;AACA,cAAM,YAAY,MAAM,OAAO,OAAO,KAAK;AAC3C,kBAAU,SAAS,IAAI,OAAO,UAAU,IAAI,OAAO,IAAI;AACvD,eAAO;AAAA,MACT;AAAA,IACF;AAAA,EACF;AAEA,SAAO,IAAI,qBAAqB,aAAa,aAAa,YAAY;AACxE;AAEO,SAAS,mBAAmB,GAAmC;AACpE,SAAO,aAAa;AACtB;","names":[]}