@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
254 lines (250 loc) • 10.4 kB
text/typescript
import { Registry, HasRegistry } from '@genkit-ai/core/registry';
import { Message } from '../message.mjs';
import { V as ModelAction } from '../prompt-CckKZx8k.mjs';
import { ModelInfo, GenerateResponseChunkData, GenerateRequest, GenerateResponseData } from '../model-types.mjs';
import { Part } from '../parts.mjs';
import '@genkit-ai/core';
import 'handlebars';
import '../document-JsI88Acr.mjs';
import '../generate/chunk.mjs';
import '../session.mjs';
import '@genkit-ai/core/async';
import '../agent-types.mjs';
import '../generate/response.mjs';
import '../formats/types.mjs';
import '../resource.mjs';
/**
* Copyright 2026 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.
*/
/**
* A streamed chunk a mock model may emit. A bare string is shorthand for a
* single text part (`{ content: [{ text }] }`).
*/
type MockChunk = string | GenerateResponseChunkData;
/** Context passed to a {@link MockModelOptions.respond} callback. */
interface MockContext {
/**
* Emit a streamed chunk. A no-op unless the caller used `generateStream` /
* `{ onChunk }`. A string is shorthand for a text part.
*/
sendChunk: (chunk: MockChunk) => void;
}
/**
* What a {@link MockModelOptions.respond} callback may return. From lightest to
* fullest control:
* - a `string` — shorthand for a single text response;
* - an object with any of `text` / `toolRequests` / `content` — the common
* cases, assembled into a well-formed response for you;
* - a full {@link GenerateResponseData} (anything with a `message`) — used as-is.
*/
interface MockResponseObject {
/** Text content of the model message. */
text?: string;
/** Tool/function calls to emit, by tool name. */
toolRequests?: Array<{
name: string;
input?: unknown;
ref?: string;
}>;
/** Escape hatch: raw message parts, prepended before `text`/`toolRequests`. */
content?: Part[];
finishReason?: GenerateResponseData['finishReason'];
usage?: GenerateResponseData['usage'];
}
type MockResponse = string | MockResponseObject | GenerateResponseData;
/** A `respond` callback: given the request, returns the response to emit. */
type MockRespondFn = (request: GenerateRequest, context: MockContext) => MockResponse | Promise<MockResponse>;
/** Options for {@link mockModel}. */
interface MockModelOptions {
/** Registered model name. Defaults to `'mockModel'`. */
name?: string;
/**
* Model metadata (e.g. `supports`, `versions`).
*
* `supports.constrained` defaults to `'all'` (native constrained generation),
* so a `generate` with an `output.schema` reaches `respond` with that schema
* on `request.output.schema`. Pass `supports: { constrained: 'none' }` to
* instead exercise the framework's simulated path, where the schema is
* injected into the prompt (and thus visible in `lastRequestText`) and
* stripped from what `respond` sees.
*/
info?: ModelInfo;
/**
* What to respond with on each `generate` call. Defaults to empty text.
*
* - A **single response** (string / object) is returned on every call.
* - A **callback** `(request, { sendChunk }) => response` is invoked once per
* call, so you can branch on request history (for tool loops) and stream via
* `sendChunk`.
* - An **array** is a queue consumed one item per call, with the last item
* repeating once exhausted — handy for scripted multi-turn tests. A queued
* `Error` is thrown when reached, to inject a failure on a given turn.
* (Queued items are static, so streaming needs the callback form.)
*
* ```ts
* mockModel(ai, { respond: 'always this' });
* mockModel(ai, { respond: ['first', 'second'] });
* mockModel(ai, { respond: ['ok', new Error('rate limited')] });
* ```
*/
respond?: MockRespond;
}
/**
* Anything accepted as respond behavior: a single {@link MockResponse}, a
* per-call callback, or a queue of responses/errors. See
* {@link MockModelOptions.respond}.
*/
type MockRespond = MockRespondFn | MockResponse | Array<MockResponse | Error>;
/**
* A mock model with typed inspection of the calls it received. The extra
* members are read-only views over the recorded calls.
*/
type MockModel = ModelAction & {
/** The request from the most recent call, or `undefined` if never called. */
readonly lastRequest: GenerateRequest | undefined;
/**
* The final message of the most recent request, wrapped as a {@link Message}
* so you can read it the same way you read a response — `.text`, `.media`,
* `.toolRequests`, etc. `undefined` if the model was never called.
*
* ```ts
* assert.match(model.lastRequestMessage!.text, /Summarize: long text/);
* ```
*/
readonly lastRequestMessage: Message | undefined;
/**
* The full assembled conversation of the most recent request, flattened to a
* single string (system + every message, in order). Use it for prompt-
* assembly assertions on any mock — including ones returning structured
* output, where {@link echoModel} can't be used. `undefined` if never called.
*
* ```ts
* assert.match(model.lastRequestText!, /system: Be terse/);
* ```
*/
readonly lastRequestText: string | undefined;
/**
* The tool results fed back to the model in the most recent request, in order.
* Use it to assert which tools ran and what they returned, without digging
* through message content yourself. Empty if the model was never called or saw
* no tool results.
*
* ```ts
* assert.deepStrictEqual(model.toolResponses.map((t) => t.name), ['lookup']);
* ```
*/
readonly toolResponses: Array<{
name: string;
ref?: string;
output: unknown;
}>;
/** Every request this model received, oldest first. */
readonly requests: GenerateRequest[];
/** How many times this model was called. */
readonly requestCount: number;
/**
* Replaces the respond behavior for subsequent calls. Accepts everything
* {@link MockModelOptions.respond} does; an array re-arms as a fresh queue.
* Recorded history is untouched — use {@link MockModel.reset} for that.
*
* Together with `reset()` this supports the register-once idiom: define the
* mock once per test file, then give each test its own behavior.
*
* ```ts
* const model = mockModel(ai, { name: 'menuModel' });
* beforeEach(() => model.reset());
*
* test('...', async () => {
* model.respondWith({ text: 'scripted' });
* // ...
* });
* ```
*/
respondWith(respond: MockRespond): void;
/**
* Clears recorded history (`requests`, `requestCount`, …) and restores the
* respond behavior given at construction, re-arming a queued `respond` from
* its first item. Call it in `beforeEach` so tests sharing a mock stay
* order-independent.
*/
reset(): void;
};
/**
* Defines a programmable mock model for tests. Drive each call's response with
* `respond`, and inspect what the model was called with via `lastRequest` /
* `requests` / `requestCount`.
*
* ```ts
* const model = mockModel(ai, { respond: () => ({ text: 'a summary' }) });
* const out = (await ai.generate({ model, prompt: 'Summarize: ...' })).text;
* assert.match(model.lastRequest!.messages.at(-1)!.content[0].text!, /Summarize/);
* ```
*
* Streaming and tool calls are first-class:
*
* ```ts
* mockModel(ai, {
* respond: (req, { sendChunk }) => {
* sendChunk('hel');
* sendChunk('lo');
* return { text: 'hello' };
* },
* });
*
* mockModel(ai, {
* respond: () => ({ toolRequests: [{ name: 'lookup', input: { id: 1 } }] }),
* });
* ```
*
* For structured output, the mock defaults to native constrained generation
* (`supports.constrained: 'all'`), so `respond` sees `request.output.schema`
* and no schema blob is injected into the prompt. Override with
* `supports: { constrained: 'none' }` to test the simulated path.
*
* @param registry a `Genkit` instance (or anything holding a `Registry`).
* @param options model name, metadata, and the `respond` callback.
*/
declare function mockModel(registry: Registry | HasRegistry, options?: MockModelOptions): MockModel;
/** Options for {@link echoModel}. */
interface EchoModelOptions {
/** Registered model name. Defaults to `'echoModel'`. */
name?: string;
/** Model metadata. */
info?: ModelInfo;
}
/**
* A {@link mockModel} preset for *text* paths: a zero-config model that echoes
* the rendered request back as text, for asserting prompt and message assembly
* (what the model *would have seen*). Supports the same inspection members as
* {@link mockModel}.
*
* ```ts
* const model = echoModel(ai);
* const res = await ai.generate({ model, system: 'Be terse', prompt: 'hi' });
* assert.match(res.text, /system: Be terse/);
* ```
*
* Because it returns text, it can't satisfy a structured **output schema** —
* Genkit derives `output` by parsing the response text and validating it, which
* prose can't pass. If the request carries an output schema, `echoModel` throws
* an explanatory error. For structured-output paths, use {@link mockModel} with
* a conforming response and assert assembly via {@link MockModel.lastRequestText}
* / {@link MockModel.lastRequest} instead.
*
* @param registry a `Genkit` instance (or anything holding a `Registry`).
* @param options model name and metadata.
*/
declare function echoModel(registry: Registry | HasRegistry, options?: EchoModelOptions): MockModel;
export { type EchoModelOptions, type MockChunk, type MockContext, type MockModel, type MockModelOptions, type MockRespond, type MockRespondFn, type MockResponse, type MockResponseObject, echoModel, mockModel };