UNPKG

@genkit-ai/ai

Version:

Genkit AI framework generative AI APIs.

263 lines (239 loc) 6.87 kB
/** * Copyright 2025 Google LLC * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ import { Action, ActionContext, detachedAction, DetachedAction, GenkitError, isAction, isDetachedAction, z, } from '@genkit-ai/core'; import { Registry } from '@genkit-ai/core/registry'; import uriTemplate from 'uri-templates'; import { PartSchema } from './model-types.js'; /** * Options for defining a resource. */ export interface ResourceOptions { /** * Resource name. If not specified, uri or template will be used as name. */ name?: string; /** * The URI of the resource. Can contain template variables. */ uri?: string; /** * The URI template (ex. `my://resource/{id}`). See RFC6570 for specification. */ template?: string; /** * A description of the resource. */ description?: string; /** * Resource metadata. */ metadata?: Record<string, any>; } export const ResourceInputSchema = z.object({ uri: z.string(), }); export type ResourceInput = z.infer<typeof ResourceInputSchema>; export const ResourceOutputSchema = z.object({ content: z.array(PartSchema), }); export type ResourceOutput = z.infer<typeof ResourceOutputSchema>; /** * A function that returns parts for a given resource. */ export type ResourceFn = ( input: ResourceInput, ctx: ActionContext ) => ResourceOutput | Promise<ResourceOutput>; /** * A resource action. */ export interface ResourceAction extends Action<typeof ResourceInputSchema, typeof ResourceOutputSchema> { matches(input: ResourceInput): boolean; } /** * Defines a resource. * * @param registry The registry to register the resource with. * @param opts The resource options. * @param fn The resource function. * @returns The resource action. */ export function defineResource( registry: Registry, opts: ResourceOptions, fn: ResourceFn ): ResourceAction { const action = dynamicResource(opts, fn).attach(registry) as ResourceAction; // when defined via defineResource it's not dynamic. delete action.__action.metadata?.dynamic; action.matches = createMatcher(opts.uri, opts.template); registry.registerAction('resource', action); return action; } /** * A dynamic action with a `resource` type. Dynamic resources are detached actions -- not associated with any registry. */ export type DynamicResourceAction = DetachedAction< typeof ResourceInputSchema, typeof ResourceOutputSchema > & { __action: { metadata: { type: 'resource'; }; }; matches(input: ResourceInput): boolean; }; /** * Finds a matching resource in the registry. If not found returns undefined. */ export async function findMatchingResource( registry: Registry, input: ResourceInput ): Promise<ResourceAction | undefined> { for (const actKeys of Object.keys(await registry.listResolvableActions())) { if (actKeys.startsWith('/resource/')) { const resource = (await registry.lookupAction(actKeys)) as ResourceAction; if (resource.matches(input)) { return resource; } } } return undefined; } /** Checks whether provided object is a dynamic resource. */ export function isDynamicResourceAction( t: unknown ): t is DynamicResourceAction { return ( (isDetachedAction(t) || isAction(t)) && t.__action.metadata?.type === 'resource' && !!t.__action.metadata?.dynamic ); } /** * Defines a dynamic resource. Dynamic resources are just like regular resources but will not be * registered in the Genkit registry and can be defined dynamically at runtime. */ export function dynamicResource( opts: ResourceOptions, fn: ResourceFn ): DynamicResourceAction { const uri = opts.uri ?? opts.template; if (!uri) { throw new GenkitError({ status: 'INVALID_ARGUMENT', message: `must specify either url or template options`, }); } const matcher = createMatcher(opts.uri, opts.template); const action = detachedAction( { actionType: 'resource', name: opts.name ?? uri, description: opts.description, inputSchema: ResourceInputSchema, outputSchema: ResourceOutputSchema, metadata: { resource: { uri: opts.uri, template: opts.template, }, ...opts.metadata, type: 'resource', dynamic: true, }, }, async (input, ctx) => { const templateMatch = matcher(input); if (!templateMatch) { throw new GenkitError({ status: 'INVALID_ARGUMENT', message: `input ${input} did not match template ${uri}`, }); } const parts = await fn(input, ctx); parts.content.map((p) => { if (!p.metadata) { p.metadata = {}; } if (p.metadata?.resource) { if (!(p.metadata as any).resource.parent) { (p.metadata as any).resource.parent = { uri: input.uri, }; if (opts.template) { (p.metadata as any).resource.parent.template = opts.template; } } } else { (p.metadata as any).resource = { uri: input.uri, }; if (opts.template) { (p.metadata as any).resource.template = opts.template; } } return p; }); return parts; } ) as DynamicResourceAction; action.matches = matcher; return { __action: { ...action.__action, metadata: { ...action.__action.metadata, type: 'resource', }, }, attach(registry) { const bound = action.attach(registry) as ResourceAction; bound.matches = matcher; return bound; }, matches: matcher, } as DynamicResourceAction; } function createMatcher( uriOpt: string | undefined, templateOpt: string | undefined ): (input: ResourceInput) => boolean { // TODO: normalize resource URI during comparisons // foo://bar?baz=1&qux=2 and foo://bar?qux=2&baz=1 are equivalent URIs but would not match. if (uriOpt) { return (input: ResourceInput) => input.uri === uriOpt; } if (templateOpt) { const template = uriTemplate(templateOpt); return (input: ResourceInput) => template!.fromUri(input.uri) !== undefined; } throw new GenkitError({ status: 'INVALID_ARGUMENT', message: 'must specify either url or template options', }); }