@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
288 lines • 8.33 kB
JavaScript
import {
GenkitError,
OperationSchema,
action,
backgroundAction,
defineAction,
registerBackgroundAction,
z
} from "@genkit-ai/core";
import { logger } from "@genkit-ai/core/logging";
import { toJsonSchema } from "@genkit-ai/core/schema";
import { performance } from "node:perf_hooks";
import {
CustomPartSchema,
DataPartSchema,
MediaPartSchema,
TextPartSchema,
ToolRequestPartSchema,
ToolResponsePartSchema
} from "./document.js";
import {
GenerateRequestSchema,
GenerateResponseSchema
} from "./model-types.js";
import {
augmentWithContext,
simulateConstrainedGeneration
} from "./model/middleware.js";
import { defineGenerateAction } from "./generate/action.js";
export * from "./model-types.js";
function model(options, runner) {
const act = action(modelActionOptions(options), (input, ctx) => {
const startTimeMs = performance.now();
return runner(input, ctx).then((response) => {
const timedResponse = {
...response,
latencyMs: performance.now() - startTimeMs
};
return timedResponse;
});
});
Object.assign(act, {
__configSchema: options.configSchema || z.unknown()
});
return act;
}
function modelActionOptions(options) {
const label = options.label || options.name;
const middleware = getModelMiddleware(options);
return {
actionType: "model",
name: options.name,
description: label,
inputSchema: GenerateRequestSchema,
outputSchema: GenerateResponseSchema,
metadata: {
model: {
label,
customOptions: options.configSchema ? toJsonSchema({ schema: options.configSchema }) : void 0,
versions: options.versions,
supports: options.supports
}
},
use: middleware
};
}
function defineModel(registry, options, runner) {
const act = defineAction(
registry,
modelActionOptions(options),
(input, ctx) => {
const startTimeMs = performance.now();
const secondParam = options.apiVersion === "v2" ? ctx : ctx.streamingRequested ? ctx.sendChunk : void 0;
return runner(input, secondParam).then((response) => {
const timedResponse = {
...response,
latencyMs: performance.now() - startTimeMs
};
return timedResponse;
});
}
);
Object.assign(act, {
__configSchema: options.configSchema || z.unknown()
});
return act;
}
function defineBackgroundModel(registry, options) {
const act = backgroundModel(options);
registerBackgroundAction(registry, act);
return act;
}
function backgroundModel(options) {
const label = options.label || options.name;
const middleware = getModelMiddleware(options);
const act = backgroundAction({
actionType: "background-model",
name: options.name,
description: label,
inputSchema: GenerateRequestSchema,
outputSchema: GenerateResponseSchema,
metadata: {
model: {
label,
customOptions: options.configSchema ? toJsonSchema({ schema: options.configSchema }) : void 0,
versions: options.versions,
supports: options.supports
}
},
use: middleware,
async start(request) {
const startTimeMs = performance.now();
const response = await options.start(request);
Object.assign(response, {
latencyMs: performance.now() - startTimeMs
});
return response;
},
async check(op) {
return options.check(op);
},
cancel: options.cancel ? async (op) => {
if (!options.cancel) {
throw new GenkitError({
status: "UNIMPLEMENTED",
message: "cancel not implemented"
});
}
return options.cancel(op);
} : void 0
});
Object.assign(act, {
__configSchema: options.configSchema || z.unknown()
});
return act;
}
function getModelMiddleware(options) {
const middleware = options.use || [];
if (!options?.supports?.context) middleware.push(augmentWithContext());
const constratedSimulator = simulateConstrainedGeneration();
middleware.push((req, next) => {
if (!options?.supports?.constrained || options?.supports?.constrained === "none" || options?.supports?.constrained === "no-tools" && (req.tools?.length ?? 0) > 0) {
return constratedSimulator(req, next);
}
return next(req);
});
return middleware;
}
function modelActionMetadata({
name,
info,
configSchema,
background
}) {
return {
actionType: background ? "background-model" : "model",
name,
inputJsonSchema: toJsonSchema({ schema: GenerateRequestSchema }),
outputJsonSchema: background ? toJsonSchema({ schema: OperationSchema }) : toJsonSchema({ schema: GenerateResponseSchema }),
metadata: {
model: {
...info,
customOptions: configSchema ? toJsonSchema({ schema: configSchema }) : void 0
}
}
};
}
function modelRef(options) {
let name = options.name;
if (options.namespace && !name.startsWith(options.namespace + "/")) {
name = `${options.namespace}/${name}`;
}
const ref = {
...options,
name
};
ref.withConfig = (cfg) => {
return modelRef({
...options,
name,
config: cfg
});
};
ref.withVersion = (version) => {
return modelRef({
...options,
name,
version
});
};
return ref;
}
function getBasicUsageStats(input, response) {
const inputCounts = getPartCounts(input.flatMap((md) => md.content));
const outputCounts = getPartCounts(
Array.isArray(response) ? response.flatMap((c) => c.message.content) : response.content
);
return {
inputCharacters: inputCounts.characters,
inputImages: inputCounts.images,
inputVideos: inputCounts.videos,
inputAudioFiles: inputCounts.audio,
outputCharacters: outputCounts.characters,
outputImages: outputCounts.images,
outputVideos: outputCounts.videos,
outputAudioFiles: outputCounts.audio
};
}
function getPartCounts(parts) {
return parts.reduce(
(counts, part) => {
const isImage = part.media?.contentType?.startsWith("image") || part.media?.url?.startsWith("data:image");
const isVideo = part.media?.contentType?.startsWith("video") || part.media?.url?.startsWith("data:video");
const isAudio = part.media?.contentType?.startsWith("audio") || part.media?.url?.startsWith("data:audio");
return {
characters: counts.characters + (part.text?.length || 0),
images: counts.images + (isImage ? 1 : 0),
videos: counts.videos + (isVideo ? 1 : 0),
audio: counts.audio + (isAudio ? 1 : 0)
};
},
{ characters: 0, images: 0, videos: 0, audio: 0 }
);
}
async function resolveModel(registry, model2, options) {
let out;
let modelId;
if (!model2) {
model2 = await registry.lookupValue("defaultModel", "defaultModel");
}
if (!model2) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: "Must supply a `model` to `generate()` calls."
});
}
if (typeof model2 === "string") {
modelId = model2;
out = { modelAction: await lookupModel(registry, model2) };
} else if (model2.hasOwnProperty("__action")) {
modelId = model2.__action.name;
out = { modelAction: model2 };
} else {
const ref = model2;
modelId = ref.name;
out = {
modelAction: await lookupModel(registry, ref.name),
config: {
...ref.config
},
version: ref.version
};
}
if (!out.modelAction) {
throw new GenkitError({
status: "NOT_FOUND",
message: `Model '${modelId}' not found`
});
}
if (options?.warnDeprecated && out.modelAction.__action.metadata?.model?.stage === "deprecated") {
logger.warn(
`Model '${out.modelAction.__action.name}' is deprecated and may be removed in a future release.`
);
}
return out;
}
async function lookupModel(registry, model2) {
return await registry.lookupAction(`/model/${model2}`) || await registry.lookupAction(`/background-model/${model2}`);
}
export {
CustomPartSchema,
DataPartSchema,
MediaPartSchema,
TextPartSchema,
ToolRequestPartSchema,
ToolResponsePartSchema,
backgroundModel,
defineBackgroundModel,
defineGenerateAction,
defineModel,
getBasicUsageStats,
model,
modelActionMetadata,
modelRef,
resolveModel,
simulateConstrainedGeneration
};
//# sourceMappingURL=model.mjs.map