@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
199 lines • 6.54 kB
JavaScript
import { Message } from "../message.mjs";
import {
defineModel
} from "../model.mjs";
function resolveRegistry(registry) {
return registry.registry ?? registry;
}
function cloneRequest(request) {
try {
return structuredClone(request);
} catch {
return {
...request,
messages: request.messages.map((m) => ({
...m,
content: m.content.map((c) => ({ ...c }))
}))
};
}
}
function toChunkData(chunk) {
return typeof chunk === "string" ? { content: [{ text: chunk }] } : chunk;
}
function renderPart(part) {
if (part.text !== void 0) return part.text;
if (part.media) {
const type = part.media.contentType ? ` ${part.media.contentType}` : "";
return `[media${type}: ${part.media.url}]`;
}
if (part.toolRequest) {
const { name, input } = part.toolRequest;
return `[toolRequest ${name}(${JSON.stringify(input)})]`;
}
if (part.toolResponse) {
const { name, output } = part.toolResponse;
return `[toolResponse ${name}: ${JSON.stringify(output)}]`;
}
if (part.reasoning !== void 0) return `[reasoning: ${part.reasoning}]`;
if (part.resource) return `[resource: ${part.resource.uri}]`;
if (part.data !== void 0) return `[data: ${JSON.stringify(part.data)}]`;
if (part.custom !== void 0)
return `[custom: ${JSON.stringify(part.custom)}]`;
return "";
}
function renderRequestText(request) {
return request.messages.map(
(m) => (m.role === "user" || m.role === "model" ? "" : `${m.role}: `) + m.content.map(renderPart).join("")
).join("\n");
}
function toRespondFn(respond) {
if (respond === void 0) {
return () => ({ text: "" });
}
if (typeof respond === "function") {
return respond;
}
const queue = Array.isArray(respond) ? respond : [respond];
if (queue.length === 0) {
return () => ({ text: "" });
}
let i = 0;
return () => {
const item = queue[Math.min(i, queue.length - 1)];
i++;
if (item instanceof Error) throw item;
return item;
};
}
function toResponseData(response) {
if (response === void 0 || response === null) {
return { message: { role: "model", content: [] }, finishReason: "stop" };
}
if (typeof response === "string") {
return {
message: { role: "model", content: [{ text: response }] },
finishReason: "stop"
};
}
if ("message" in response && response.message) {
const data = response;
return { ...data, finishReason: data.finishReason ?? "stop" };
}
const obj = response;
const content = [...obj.content ?? []];
if (obj.text !== void 0) {
content.push({ text: obj.text });
}
for (const tool of obj.toolRequests ?? []) {
content.push({
toolRequest: { name: tool.name, input: tool.input, ref: tool.ref }
});
}
return {
message: { role: "model", content },
finishReason: obj.finishReason ?? "stop",
usage: obj.usage
};
}
function mockModel(registry, options = {}) {
const requests = [];
let respond = toRespondFn(options.respond);
const model = defineModel(
resolveRegistry(registry),
{
apiVersion: "v2",
name: options.name ?? "mockModel",
// Forward only the metadata fields defineModel accepts; ModelInfo's
// `configSchema`/`stage` aren't part of DefineModelOptions.
versions: options.info?.versions,
label: options.info?.label,
// Default to native constrained generation (like modern provider models),
// so a structured-output request reaches `respond` with `output.schema`
// intact instead of the framework injecting a schema blob into the prompt
// and stripping it. Spread last so callers can opt out with
// `supports: { constrained: 'none' }` to exercise the simulated path.
supports: { constrained: "all", ...options.info?.supports }
},
async (request, { sendChunk }) => {
requests.push(cloneRequest(request));
const context = {
sendChunk: (chunk) => sendChunk?.(toChunkData(chunk))
};
return toResponseData(await respond(request, context));
}
);
Object.defineProperties(model, {
// Return clones so callers can't mutate recorded history through a view.
requests: { get: () => requests.map((r) => cloneRequest(r)) },
lastRequest: {
get: () => {
const last = requests[requests.length - 1];
return last ? cloneRequest(last) : void 0;
}
},
lastRequestMessage: {
get: () => {
const last = requests[requests.length - 1]?.messages.at(-1);
return last ? new Message(last) : void 0;
}
},
lastRequestText: {
get: () => {
const last = requests[requests.length - 1];
return last ? renderRequestText(last) : void 0;
}
},
toolResponses: {
get: () => (requests[requests.length - 1]?.messages ?? []).flatMap((m) => m.content).filter((p) => p.toolResponse).map((p) => ({
name: p.toolResponse.name,
ref: p.toolResponse.ref,
output: p.toolResponse.output
}))
},
requestCount: { get: () => requests.length },
respondWith: {
value: (next) => {
respond = toRespondFn(next);
}
},
reset: {
value: () => {
requests.length = 0;
respond = toRespondFn(options.respond);
}
}
});
return model;
}
function echoModel(registry, options = {}) {
return mockModel(registry, {
name: options.name ?? "echoModel",
// Declare native constrained support so the framework hands the output
// schema to the model directly (in `request.output.schema`) rather than
// injecting it as prompt text — that lets the guard below detect it
// reliably, and keeps the echo free of framework-injected schema blobs.
info: {
...options.info,
supports: { ...options.info?.supports, constrained: "all" }
},
respond: (request) => {
if (request.output?.schema) {
throw new Error(
"echoModel returns text and can't satisfy an output schema: this request asks for structured output. Either move `output: { schema }` to the generate()/flow call site so the prompt stays text-only, or use mockModel(...) with a conforming response and assert prompt assembly via model.lastRequestText / model.lastRequest."
);
}
return {
content: [
{ text: "Echo: " + renderRequestText(request) },
{ text: "; config: " + JSON.stringify(request.config) }
]
};
}
});
}
export {
echoModel,
mockModel
};
//# sourceMappingURL=mock-model.mjs.map