@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
945 lines • 34.6 kB
JavaScript
;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
var agent_exports = {};
__export(agent_exports, {
AgentAbortRequestSchema: () => import_agent_types2.AgentAbortRequestSchema,
AgentAbortResponseSchema: () => import_agent_types2.AgentAbortResponseSchema,
AgentInitError: () => AgentInitError,
AgentInitSchema: () => import_agent_types2.AgentInitSchema,
AgentInputSchema: () => import_agent_types2.AgentInputSchema,
AgentOutputSchema: () => import_agent_types2.AgentOutputSchema,
AgentResultSchema: () => import_agent_types2.AgentResultSchema,
AgentStreamChunkSchema: () => import_agent_types2.AgentStreamChunkSchema,
GetSnapshotDataInputSchema: () => GetSnapshotDataInputSchema,
GetSnapshotRequestSchema: () => import_agent_types2.GetSnapshotRequestSchema,
JsonPatchOperationSchema: () => import_agent_types2.JsonPatchOperationSchema,
JsonPatchSchema: () => import_agent_types2.JsonPatchSchema,
SessionRunner: () => SessionRunner,
TurnEndSchema: () => import_agent_types2.TurnEndSchema,
defineAgent: () => defineAgent,
defineCustomAgent: () => defineCustomAgent,
definePromptAgent: () => definePromptAgent,
validateResumeAgainstHistory: () => validateResumeAgainstHistory
});
module.exports = __toCommonJS(agent_exports);
var import_core = require("@genkit-ai/core");
var import_async = require("@genkit-ai/core/async");
var import_agent_core = require("./agent-core.js");
var import_schema = require("@genkit-ai/core/schema");
var import_tracing = require("@genkit-ai/core/tracing");
var import_agent_types = require("./agent-types.js");
var import_generate = require("./generate.js");
var import_json_patch = require("./json-patch.js");
var import_prompt = require("./prompt.js");
var import_session_stores = require("./session-stores.js");
var import_session = require("./session.js");
var import_agent_types2 = require("./agent-types.js");
const DEFAULT_HEARTBEAT_INTERVAL_MS = 3e4;
const DEFAULT_HEARTBEAT_TIMEOUT_MS = 6e4;
function isHeartbeatExpired(snapshot, timeoutMs = DEFAULT_HEARTBEAT_TIMEOUT_MS) {
if (snapshot.status !== "pending" || !snapshot.heartbeatAt) {
return false;
}
const last = Date.parse(snapshot.heartbeatAt);
if (Number.isNaN(last)) {
return false;
}
return Date.now() - last > timeoutMs;
}
function toErrorDetails(e) {
return {
status: e?.status || "INTERNAL",
message: e?.message || "Internal failure",
// Only surface explicitly-provided structured details. Never fall back to
// the raw thrown value: it is serialized over the wire (AgentOutput.error)
// and persisted into snapshots, so leaking it could expose stack traces /
// internal state, or break JSON.stringify on circular error objects.
details: e?.detail ?? e?.details
};
}
function abortAwareMutator(input) {
return (current) => current?.status === "aborted" ? null : input;
}
function requireStore(store, operation, agentName) {
if (!store) {
throw new import_core.GenkitError({
status: "FAILED_PRECONDITION",
message: `${operation} requires a persistent store. Provide a 'store' when defining '${agentName}'.`
});
}
}
async function abortSnapshotInStore(store, snapshotId, options) {
let previousStatus;
await store.saveSnapshot(
snapshotId,
(current) => {
if (!current) return null;
previousStatus = current.status;
if (current.status === "completed" || current.status === "failed" || current.status === "aborted") {
return null;
}
return { ...current, status: "aborted" };
},
options
);
return previousStatus;
}
class SessionRunner {
session;
inputCh;
turnIndex = 0;
onEndTurn;
onDetach;
newSnapshotId;
/** The finish reason of the most recently completed turn. */
lastTurnFinishReason;
/**
* Error details of the most recent failed turn. Set when a turn throws and
* the runner resolves gracefully instead of propagating the exception.
*/
lastTurnError;
/**
* The state the most recently *successful* turn left behind. On a failed
* turn this is the state the failed turn started with - the last-good state
* returned to the caller (for client-managed agents).
*/
lastGoodState;
/**
* The snapshotId of the most recently *successful* (persisted, `done`) turn.
* On a failed turn this is the last-good snapshot the caller resumes from;
* `undefined` when no turn has succeeded yet (e.g. a first-turn failure).
*/
lastGoodSnapshotId;
lastSnapshot;
lastSnapshotVersion = 0;
store;
isDetached = false;
/**
* Aborts in-flight turns. When set and aborted, a turn that rejects out of
* `generate` is reported as `aborted` (not `failed`) and its failed snapshot
* write is skipped (the abort path already persisted the `aborted` status).
*/
abortSignal;
/**
* True until the first `customPatch` chunk of the current turn has been
* emitted. The first patch of every turn is a whole-document replace
* (re-basing clients that may not share the server's baseline); reset to
* `true` at the start of each turn.
*/
firstCustomPatchInTurn = true;
constructor(session, inputCh, options) {
this.session = session;
this.inputCh = inputCh;
this.lastSnapshot = options?.lastSnapshot;
this.store = options?.store;
this.abortSignal = options?.abortSignal;
this.onEndTurn = options?.onEndTurn;
this.onDetach = options?.onDetach;
this.lastGoodState = this.session.getState();
this.lastGoodSnapshotId = options?.lastSnapshot?.snapshotId;
}
// ── Session delegate methods ────────────────────────────────────────
// These forward to `this.session` so callers can write `sess.addMessages()`
// instead of the verbose `sess.session.addMessages()`.
/** Returns a deep copy of the current session state. */
getState() {
return this.session.getState();
}
/** Retrieves all messages associated with the session. */
getMessages() {
return this.session.getMessages();
}
/** Appends messages to the session. */
addMessages(messages) {
this.session.addMessages(messages);
}
/** Overwrites the session messages. */
setMessages(messages) {
this.session.setMessages(messages);
}
/** Retrieves the custom state of the session. */
getCustom() {
return this.session.getCustom();
}
/** Updates the custom state using a mutator function. */
updateCustom(fn) {
this.session.updateCustom(fn);
}
/** Retrieves the list of artifacts generated during the session. */
getArtifacts() {
return this.session.getArtifacts();
}
/** Adds artifacts to the session, deduplicating by name. */
addArtifacts(artifacts) {
this.session.addArtifacts(artifacts);
}
/** Invokes the end-of-turn callback, absorbing errors from a closed stream. */
notifyEndTurn(snapshotId, finishReason) {
try {
this.onEndTurn?.(snapshotId, finishReason);
} catch {
}
}
/**
* Executes the flow handler against incoming input messages sequentially.
*
* The handler receives the turn's {@link AgentInput} and a {@link TurnContext}
* whose `snapshotId` is *reserved up front* - it is the id the snapshot
* persisted at turn end will reuse. This lets a handler set up external,
* snapshot-correlated state (e.g. a git branch/worktree named after the
* snapshot) before generating, then commit it under that id.
*
* The handler may return a {@link TurnResult} carrying an explicit
* `finishReason` for the just-completed turn. When omitted, no per-turn
* reason is reported. Failures always report `failed`.
*/
async run(fn) {
for await (const input of this.inputCh) {
if (input.message) {
this.session.addMessages([input.message]);
}
this.firstCustomPatchInTurn = true;
const parentSnapshotId = this.lastSnapshot?.snapshotId;
if (this.store && !this.newSnapshotId) {
this.newSnapshotId = (0, import_session.reserveSnapshotId)();
}
const turnSnapshotId = this.newSnapshotId;
this.newSnapshotId = void 0;
const turnContext = {
snapshotId: turnSnapshotId,
parentSnapshotId,
turnIndex: this.turnIndex
};
try {
await (0, import_core.run)(`runTurn-${this.turnIndex + 1}`, input, async () => {
const turnResult = await fn(input, turnContext);
const finishReason = turnResult?.finishReason;
this.lastTurnFinishReason = finishReason;
this.lastTurnError = void 0;
const snapshotId = await this.maybeSnapshot(
"completed",
void 0,
turnSnapshotId,
finishReason
);
this.lastGoodState = this.session.getState();
this.lastGoodSnapshotId = snapshotId;
if (snapshotId) {
(0, import_tracing.setCustomMetadataAttribute)("agent:snapshotId", snapshotId);
}
this.notifyEndTurn(snapshotId, finishReason);
return { state: this.session.getState() };
});
this.turnIndex++;
} catch (e) {
if (this.abortSignal?.aborted) {
this.lastTurnFinishReason = "aborted";
this.lastTurnError = void 0;
this.notifyEndTurn(this.lastSnapshot?.snapshotId, "aborted");
break;
}
this.lastTurnFinishReason = "failed";
this.lastTurnError = toErrorDetails(e);
const snapshotId = await this.maybeSnapshot(
"failed",
this.lastTurnError,
turnSnapshotId,
"failed"
);
this.notifyEndTurn(snapshotId, "failed");
break;
}
}
}
/**
* Saves a snapshot of the current session state to the persistent store.
*
* When a store is configured every turn is persisted (snapshotting is no
* longer opt-out). Uses the mutator-based `saveSnapshot` to atomically check
* that the snapshot has not been concurrently aborted before writing -
* preventing a race where a "done" write could overwrite a concurrent
* "aborted" status.
*/
async maybeSnapshot(status, error, snapshotId, finishReason) {
if (!this.store || this.isDetached && snapshotId !== this.lastSnapshot?.snapshotId)
return this.lastSnapshot?.snapshotId;
const currentVersion = this.session.getVersion();
if (currentVersion === this.lastSnapshotVersion && !status) {
return this.lastSnapshot?.snapshotId;
}
const currentState = this.session.getState();
const snapshotInput = {
...snapshotId || this.newSnapshotId ? { snapshotId: snapshotId || this.newSnapshotId } : {},
// Stamp the session id onto every snapshot in the chain so callers can
// resolve a snapshot's session without reaching into its state.
sessionId: this.session.sessionId,
createdAt: (/* @__PURE__ */ new Date()).toISOString(),
updatedAt: (/* @__PURE__ */ new Date()).toISOString(),
state: currentState,
parentId: this.lastSnapshot?.snapshotId,
// Default to a resumable `completed` status. The only caller that omits a
// status is the post-invocation write (which fires when the handler
// mutates state after the last turn); persisting it as `completed` keeps
// it a valid resume target under the "only `completed` is resumable" rule.
status: status ?? "completed",
// Stamp an initial heartbeat on a `pending` (detached, in-flight)
// snapshot. A background heartbeat loop refreshes it; if it goes stale the
// snapshot is reported as `expired` on read (the worker is presumed dead).
...status === "pending" && { heartbeatAt: (/* @__PURE__ */ new Date()).toISOString() },
...finishReason && { finishReason },
error
};
const effectiveId = snapshotId || this.newSnapshotId;
const assignedId = await this.store.saveSnapshot(
effectiveId,
abortAwareMutator(snapshotInput),
{ context: (0, import_core.getContext)() }
);
if (assignedId === null) {
return effectiveId;
}
this.lastSnapshot = { ...snapshotInput, snapshotId: assignedId };
this.lastSnapshotVersion = currentVersion;
return assignedId;
}
}
const GetSnapshotDataInputSchema = import_core.z.object({
snapshotId: import_core.z.string().optional(),
sessionId: import_core.z.string().optional()
});
class AgentInitError extends import_core.GenkitError {
}
function assertInitMatchesStateManagement(config, init) {
if ((init?.snapshotId || init?.sessionId) && !config.store) {
throw new AgentInitError({
status: "FAILED_PRECONDITION",
message: `Cannot use '${init.snapshotId ? "snapshotId" : "sessionId"}' with agent '${config.name}': this agent has no store configured (client-managed state). Send 'state' instead.`
});
}
if (init?.state && config.store) {
throw new AgentInitError({
status: "FAILED_PRECONDITION",
message: `Cannot send 'state' to agent '${config.name}': this agent uses a server-managed store. Send 'snapshotId' or 'sessionId' instead.`
});
}
}
async function resolveSession(config, store, init, validateCustomState) {
if (init?.snapshotId) {
const snapshot = await store.getSnapshot({
snapshotId: init.snapshotId,
context: (0, import_core.getContext)()
});
if (!snapshot) {
throw new import_core.GenkitError({
status: "NOT_FOUND",
message: `Snapshot ${init.snapshotId} not found`
});
}
const snapshotSessionId = snapshot.sessionId ?? snapshot.state?.sessionId;
if (init.sessionId && snapshotSessionId !== init.sessionId) {
throw new AgentInitError({
status: "INVALID_ARGUMENT",
message: `Snapshot ${init.snapshotId} does not belong to session ${init.sessionId} (it belongs to ${snapshotSessionId ?? "an unknown session"}).`
});
}
if (snapshot.status !== "completed") {
throw new import_core.GenkitError({
status: "INVALID_ARGUMENT",
message: `Snapshot ${init.snapshotId} is not resumable (status: ${snapshot.status ?? "unknown"}). Only 'completed' snapshots can be resumed.`
});
}
validateCustomState(snapshot.state?.custom);
return {
snapshot,
session: new import_session.Session(snapshot.state)
};
}
if (init?.sessionId) {
let snapshot = await store.getSnapshot({
sessionId: init.sessionId,
context: (0, import_core.getContext)()
});
const visited = /* @__PURE__ */ new Set();
while (snapshot && snapshot.status !== "completed") {
if (visited.has(snapshot.snapshotId)) {
throw new import_core.GenkitError({
status: "FAILED_PRECONDITION",
message: `Session '${init.sessionId}' has a cyclic snapshot parent chain (snapshot '${snapshot.snapshotId}' was visited twice). Resume by snapshotId instead.`
});
}
visited.add(snapshot.snapshotId);
snapshot = snapshot.parentId ? await store.getSnapshot({
snapshotId: snapshot.parentId,
context: (0, import_core.getContext)()
}) : void 0;
}
if (snapshot) {
validateCustomState(snapshot.state?.custom);
return {
snapshot,
session: new import_session.Session(snapshot.state)
};
}
return {
session: new import_session.Session({
custom: void 0,
artifacts: [],
messages: [],
sessionId: init.sessionId
})
};
}
if (init?.state && !config.store) {
validateCustomState(init.state.custom);
return {
session: new import_session.Session(init.state)
};
}
return {
session: new import_session.Session({
custom: void 0,
artifacts: [],
messages: []
})
};
}
function pipeInputWithDetach(inputStream, target, getRunner, storeEnabled, rejectDetach) {
(async () => {
try {
for await (const input of inputStream) {
if (input.detach) {
if (!storeEnabled) {
rejectDetach(
new import_core.GenkitError({
status: "FAILED_PRECONDITION",
message: "Detach is only supported when a session store is provided."
})
);
} else {
const runner = getRunner();
const turnSnapshotId = runner.newSnapshotId || (0, import_session.reserveSnapshotId)();
runner.newSnapshotId = turnSnapshotId;
await runner.maybeSnapshot("pending", void 0, turnSnapshotId);
runner.isDetached = true;
if (runner.onDetach) {
runner.onDetach(turnSnapshotId);
}
}
const hasPayload = !!(input.message || input.resume?.restart?.length || input.resume?.respond?.length);
if (hasPayload) {
target.send(input);
}
} else {
target.send(input);
}
}
target.close();
} catch (e) {
target.error(e);
}
})();
}
function defineCustomAgent(registry, config, fn) {
const toClientState = (state) => {
if (config.clientTransform?.state) {
return config.clientTransform.state(state);
}
return state;
};
const stateJsonSchema = config.stateSchema ? (0, import_schema.toJsonSchema)({ schema: config.stateSchema }) : void 0;
const validateCustomState = (custom) => {
if (config.stateSchema && custom !== void 0) {
(0, import_schema.parseSchema)(custom, { schema: config.stateSchema });
}
};
const primaryAction = (0, import_core.defineBidiAction)(
registry,
{
name: config.name,
description: config.description,
actionType: "agent",
inputSchema: import_agent_types.AgentInputSchema,
outputSchema: import_agent_types.AgentOutputSchema,
streamSchema: import_agent_types.AgentStreamChunkSchema,
initSchema: import_agent_types.AgentInitSchema,
metadata: {
agent: {
stateManagement: config.store ? "server" : "client",
abortable: !!config.store?.onSnapshotStateChange,
...stateJsonSchema && { stateSchema: stateJsonSchema }
}
}
},
async function* (arg) {
const init = arg.init;
const store = config.store || new import_session_stores.InMemorySessionStore();
assertInitMatchesStateManagement(config, init);
let session;
let snapshot;
try {
({ session, snapshot } = await resolveSession(
config,
store,
init,
validateCustomState
));
} catch (e) {
if (e instanceof AgentInitError) {
throw e;
}
return {
finishReason: "failed",
error: toErrorDetails(e),
...!config.store && init?.state && { state: init.state }
};
}
(0, import_tracing.setCustomMetadataAttributes)({
"agent:sessionId": session.sessionId
});
let detachedSnapshotId;
let resolveDetach;
let rejectDetach;
const detachPromise = new Promise((resolve, reject) => {
resolveDetach = resolve;
rejectDetach = reject;
});
const abortController = new AbortController();
let unsubscribe = void 0;
let heartbeatTimer;
const stopHeartbeat = () => {
if (heartbeatTimer) {
clearInterval(heartbeatTimer);
heartbeatTimer = void 0;
}
};
let runner;
const emitChunk = (chunk) => {
try {
let toSend = chunk;
if (config.clientTransform?.chunk) {
toSend = config.clientTransform.chunk(chunk);
}
if (!toSend) return;
arg.sendChunk(toSend);
} catch {
}
};
const runnerInputChannel = new import_async.Channel();
pipeInputWithDetach(
arg.inputStream,
runnerInputChannel,
() => runner,
!!config.store,
(reason) => rejectDetach?.(reason)
);
runner = new SessionRunner(session, runnerInputChannel, {
store,
lastSnapshot: snapshot,
abortSignal: abortController.signal,
onDetach: (snapshotId) => {
detachedSnapshotId = snapshotId;
if (resolveDetach) {
resolveDetach();
}
const ctx = (0, import_core.getContext)();
heartbeatTimer = setInterval(() => {
void store.saveSnapshot(
snapshotId,
(current) => current?.status === "pending" ? { ...current, heartbeatAt: (/* @__PURE__ */ new Date()).toISOString() } : null,
{ context: ctx }
).catch(() => {
});
}, DEFAULT_HEARTBEAT_INTERVAL_MS);
heartbeatTimer.unref?.();
if (store.onSnapshotStateChange) {
unsubscribe = store.onSnapshotStateChange(
snapshotId,
(snap) => {
if (snap.status === "aborted") {
stopHeartbeat();
abortController.abort();
if (unsubscribe) unsubscribe();
}
},
{ context: (0, import_core.getContext)() }
);
}
},
onEndTurn: (snapshotId, finishReason2) => {
if (!runner.isDetached) {
emitChunk({
turnEnd: {
...config.store && { snapshotId },
...finishReason2 && { finishReason: finishReason2 }
}
});
}
}
});
const sendArtifactChunk = (a) => {
if (!runner.isDetached) {
emitChunk({ artifact: a });
}
};
session.on("artifactAdded", sendArtifactChunk);
session.on("artifactUpdated", sendArtifactChunk);
let lastSentCustom;
const sendCustomPatch = () => {
if (runner.isDetached) return;
const transformed = toClientState(session.getState())?.custom;
let patch;
if (runner.firstCustomPatchInTurn) {
patch = [
{ op: "replace", path: "", value: structuredClone(transformed) }
];
runner.firstCustomPatchInTurn = false;
} else {
patch = (0, import_json_patch.diff)(lastSentCustom, transformed);
}
lastSentCustom = structuredClone(transformed);
if (patch.length) {
emitChunk({ customPatch: patch });
}
};
session.on("customChanged", sendCustomPatch);
const sendChunk = (chunk) => {
if (!runner.isDetached) {
emitChunk(chunk);
}
};
const flowPromise = (async () => {
try {
const result2 = await (0, import_session.runWithSession)(
registry,
session,
() => fn(runner, {
sendChunk,
abortSignal: abortController.signal,
context: (0, import_core.getContext)()
})
);
const finalSnapshotId2 = await runner.maybeSnapshot();
return { result: result2, finalSnapshotId: finalSnapshotId2 };
} finally {
stopHeartbeat();
if (unsubscribe) unsubscribe();
session.off("artifactAdded", sendArtifactChunk);
session.off("artifactUpdated", sendArtifactChunk);
session.off("customChanged", sendCustomPatch);
}
})();
const outcome = await Promise.race([
flowPromise,
detachPromise.then(() => "detached")
]);
if (outcome === "detached") {
return {
sessionId: session.sessionId,
snapshotId: detachedSnapshotId,
finishReason: "detached",
...!config.store && { state: toClientState(session.getState()) }
};
}
const { result, finalSnapshotId } = outcome;
if (runner.lastTurnFinishReason === "failed" && runner.lastTurnError) {
const lastGood = runner.lastGoodState ?? session.getState();
const lastGoodMessages = lastGood.messages;
return {
sessionId: session.sessionId,
finishReason: "failed",
error: runner.lastTurnError,
...result.artifacts?.length && { artifacts: result.artifacts },
...lastGoodMessages?.length && {
message: lastGoodMessages[lastGoodMessages.length - 1]
},
// Server-managed: the last successful turn is already persisted (every
// turn is snapshotted), so point at its snapshot. The failed turn's
// own snapshot is persisted too but is not resumable - `sessionId`
// resume skips it back to this last-good `done` snapshot. Undefined on
// a first-turn failure (no successful turn yet; client holds the seed).
...config.store && { snapshotId: runner.lastGoodSnapshotId },
// Client-managed: return the last-good state directly.
...!config.store && { state: toClientState(lastGood) }
};
}
const finishReason = result.finishReason ?? runner.lastTurnFinishReason;
return {
sessionId: session.sessionId,
...result.artifacts?.length && { artifacts: result.artifacts },
...result.message && { message: result.message },
...finishReason && { finishReason },
...config.store && { snapshotId: finalSnapshotId },
...!config.store && { state: toClientState(session.getState()) }
};
}
);
const toClientSnapshot = (snapshot) => {
if (!config.clientTransform?.state || !snapshot.state) {
return snapshot;
}
return {
...snapshot,
state: config.clientTransform.state(snapshot.state)
};
};
const resolveSnapshot = async (lookup) => {
requireStore(config.store, "getSnapshotData", config.name);
const snapshot = await config.store.getSnapshot(lookup);
if (!snapshot) return void 0;
const effective = isHeartbeatExpired(snapshot) ? { ...snapshot, status: "expired" } : snapshot;
return toClientSnapshot(effective);
};
const runAbort = (snapshotId, options) => {
requireStore(config.store, "abort", config.name);
return abortSnapshotInStore(config.store, snapshotId, options);
};
const getSnapshotDataAction = (0, import_core.defineAction)(
registry,
{
name: config.name,
description: `Gets snapshot data for ${config.name} by snapshotId or sessionId`,
actionType: "agent-snapshot",
inputSchema: import_agent_types.GetSnapshotRequestSchema,
outputSchema: import_session.SessionSnapshotSchema.optional()
},
async (lookup) => resolveSnapshot({ ...lookup, context: (0, import_core.getContext)() })
);
const abortAgentAction = (0, import_core.defineAction)(
registry,
{
name: config.name,
description: `Aborts ${config.name} agent by snapshotId. Returns the snapshot id and its status after the abort attempt.`,
actionType: "agent-abort",
inputSchema: import_agent_types.AgentAbortRequestSchema,
outputSchema: import_agent_types.AgentAbortResponseSchema
},
async ({ snapshotId }) => {
const status = await runAbort(snapshotId, { context: (0, import_core.getContext)() });
return { snapshotId, status };
}
);
const composite = Object.assign(primaryAction, {
getSnapshotData: (opts) => resolveSnapshot(opts),
abort: (snapshotId, options) => runAbort(snapshotId, options),
getSnapshotDataAction,
abortAgentAction
});
const startBidi = (input, init, opts) => {
const bidi = primaryAction.streamBidi(init, {
abortSignal: opts.abortSignal
});
bidi.send(input);
bidi.close();
return bidi;
};
const transport = {
stateManagement: config.store ? "server" : "client",
runTurn(input, init, opts) {
const bidi = startBidi(input, init, opts);
return { stream: bidi.stream, output: bidi.output };
},
async getSnapshot(lookup) {
return composite.getSnapshotData(lookup);
},
abort(snapshotId) {
return composite.abort(snapshotId);
}
};
const agentApi = (0, import_agent_core.createAgentAPI)(transport);
Object.assign(composite, {
chat: agentApi.chat,
loadChat: agentApi.loadChat,
getSnapshot: agentApi.getSnapshot
});
return composite;
}
function definePromptAgent(registry, config) {
let cachedPromptAction;
const fn = async (sess, { sendChunk, abortSignal }) => {
await sess.run(async (input) => {
const promptInput = config.promptInput ?? {};
if (!cachedPromptAction) {
cachedPromptAction = await registry.lookupAction(
`/prompt/${config.promptName}`
);
if (!cachedPromptAction) {
throw new Error(
`Prompt '${config.promptName}' not found. Ensure it is defined before the agent is invoked.`
);
}
}
const historyTag = "_genkit_history";
const promptTag = "agentPreamble";
const history = (sess.getMessages() || []).map((m) => ({
...m,
metadata: { ...m.metadata, [historyTag]: true }
}));
const genOpts = await cachedPromptAction.__executablePrompt.render(
promptInput,
{ messages: history }
);
if (genOpts.messages) {
genOpts.messages = genOpts.messages.map((m) => {
if (m.metadata?.[historyTag]) {
const { [historyTag]: _, ...restMeta } = m.metadata;
return {
...m,
metadata: Object.keys(restMeta).length ? restMeta : void 0
};
}
return { ...m, metadata: { ...m.metadata, [promptTag]: true } };
});
}
if (input.resume) {
validateResumeAgainstHistory(input.resume, sess.getMessages());
genOpts.resume = {
...input.resume.restart?.length && {
restart: input.resume.restart
},
...input.resume.respond?.length && {
respond: input.resume.respond
}
};
}
const result = (0, import_generate.generateStream)(registry, { ...genOpts, abortSignal });
for await (const chunk of result.stream) {
sendChunk({ modelChunk: chunk });
}
const res = await result.response;
if (res.request?.messages) {
const msgs2 = res.request.messages.filter(
(m) => !m.metadata?.[promptTag]
);
if (res.message) {
msgs2.push(res.message);
}
sess.setMessages(msgs2);
} else if (res.message) {
sess.addMessages([res.message]);
}
if (res.finishReason === "interrupted") {
const parts = res.message?.content?.filter((p) => !!p.toolRequest) || [];
if (parts.length > 0) {
sendChunk({
modelChunk: {
role: "tool",
content: parts
}
});
}
}
return { finishReason: res.finishReason };
});
const msgs = sess.getMessages();
return {
artifacts: sess.getArtifacts(),
message: msgs.length > 0 ? msgs[msgs.length - 1] : void 0,
...sess.lastTurnFinishReason && {
finishReason: sess.lastTurnFinishReason
}
};
};
return defineCustomAgent(
registry,
{
name: config.promptName,
description: config.description,
stateSchema: config.stateSchema,
store: config.store,
clientTransform: config.clientTransform
},
fn
);
}
function validateResumeAgainstHistory(resume, history) {
const allToolRequests = [];
for (const msg of history) {
if (msg.role === "model") {
for (const part of msg.content) {
if (part.toolRequest) {
allToolRequests.push(part.toolRequest);
}
}
}
}
for (const restart of resume.restart || []) {
const { name, ref, input } = restart.toolRequest;
const match = allToolRequests.find(
(tr) => tr.name === name && tr.ref === ref
);
if (!match) {
throw new import_core.GenkitError({
status: "INVALID_ARGUMENT",
message: `resume.restart references tool '${name}'` + (ref ? ` (ref: ${ref})` : "") + ` which was not found in session history.`
});
}
if (!(0, import_core.deepEqual)(input, match.input)) {
throw new import_core.GenkitError({
status: "INVALID_ARGUMENT",
message: `resume.restart for tool '${name}'` + (ref ? ` (ref: ${ref})` : "") + ` has modified inputs that do not match the original tool request in session history. Restart inputs must exactly match the interrupted tool request.`
});
}
}
for (const respond of resume.respond || []) {
const { name, ref } = respond.toolResponse;
const match = allToolRequests.find(
(tr) => tr.name === name && tr.ref === ref
);
if (!match) {
throw new import_core.GenkitError({
status: "INVALID_ARGUMENT",
message: `resume.respond references tool '${name}'` + (ref ? ` (ref: ${ref})` : "") + ` which was not found in session history.`
});
}
}
}
function defineAgent(registry, config) {
const { stateSchema, store, clientTransform, promptInput, ...promptConfig } = config;
(0, import_prompt.definePrompt)(registry, promptConfig);
return definePromptAgent(registry, {
promptName: promptConfig.name,
description: promptConfig.description,
promptInput,
stateSchema,
store,
clientTransform
});
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
AgentAbortRequestSchema,
AgentAbortResponseSchema,
AgentInitError,
AgentInitSchema,
AgentInputSchema,
AgentOutputSchema,
AgentResultSchema,
AgentStreamChunkSchema,
GetSnapshotDataInputSchema,
GetSnapshotRequestSchema,
JsonPatchOperationSchema,
JsonPatchSchema,
SessionRunner,
TurnEndSchema,
defineAgent,
defineCustomAgent,
definePromptAgent,
validateResumeAgainstHistory
});
//# sourceMappingURL=agent.js.map