@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
559 lines • 15.5 kB
JavaScript
import { applyPatch } from "./json-patch.mjs";
class AgentError extends Error {
status;
details;
state;
snapshotId;
response;
constructor(opts) {
super(opts.message);
this.name = "AgentError";
this.status = opts.status;
this.details = opts.details;
this.state = opts.state;
this.snapshotId = opts.snapshotId;
this.response = opts.response;
Object.setPrototypeOf(this, AgentError.prototype);
}
}
const TERMINAL_STATUSES = /* @__PURE__ */ new Set([
"completed",
"failed",
"aborted",
"expired"
]);
function sleep(ms) {
return new Promise((resolve) => setTimeout(resolve, ms));
}
function toAgentInput(input) {
if (typeof input === "string") {
return { message: { role: "user", content: [{ text: input }] } };
}
return input;
}
function partsText(parts) {
return (parts ?? []).map((p) => p.text ?? "").join("");
}
function partsReasoning(parts) {
return (parts ?? []).filter((p) => p.reasoning !== void 0).map((p) => p.reasoning ?? "").join("");
}
function firstMedia(parts) {
const p = (parts ?? []).find((p2) => !!p2.media);
return p?.media ? { url: p.media.url, contentType: p.media.contentType } : null;
}
function firstData(parts) {
const p = (parts ?? []).find((p2) => p2.data !== void 0);
return p?.data ?? null;
}
function toolRequestParts(parts) {
return (parts ?? []).filter((p) => !!p.toolRequest);
}
class AgentInterruptImpl {
name;
ref;
input;
constructor(part) {
this.name = part.toolRequest.name;
this.ref = part.toolRequest.ref;
this.input = part.toolRequest.input;
}
respond(output) {
return {
toolResponse: {
name: this.name,
...this.ref !== void 0 && { ref: this.ref },
output
}
};
}
restart() {
return {
toolRequest: {
name: this.name,
...this.ref !== void 0 && { ref: this.ref },
input: this.input
}
};
}
}
class AgentResponseImpl {
constructor(_raw, _messages, _fallbackState, _fallbackSessionId) {
this._raw = _raw;
this._messages = _messages;
this._fallbackState = _fallbackState;
this._fallbackSessionId = _fallbackSessionId;
}
get message() {
return this._raw.message;
}
get text() {
return partsText(this._raw.message?.content);
}
get reasoning() {
return partsReasoning(this._raw.message?.content);
}
get media() {
return firstMedia(this._raw.message?.content);
}
get data() {
return firstData(this._raw.message?.content);
}
get toolRequests() {
return toolRequestParts(this._raw.message?.content);
}
get interrupts() {
return this.toolRequests.filter((p) => !!p.metadata?.interrupt).map((p) => new AgentInterruptImpl(p));
}
get messages() {
return this._messages;
}
get finishReason() {
return this._raw.finishReason ?? "unknown";
}
get finishMessage() {
return this._raw.error?.message;
}
get raw() {
return this._raw;
}
get snapshotId() {
return this._raw.snapshotId;
}
get sessionId() {
return this._raw.sessionId ?? this._fallbackSessionId?.();
}
get state() {
const fromWire = this._raw.state?.custom;
return fromWire !== void 0 ? fromWire : this._fallbackState?.();
}
get artifacts() {
return this._raw.artifacts ?? this._raw.state?.artifacts ?? [];
}
assertValid() {
if (this.finishReason === "blocked") {
throw new Error(
`Generation blocked${this.finishMessage ? `: ${this.finishMessage}` : ""}.`
);
}
if (!this._raw.message) {
throw new Error("Agent response has no message.");
}
}
}
class AgentChunkImpl {
constructor(_raw, _previousText) {
this._raw = _raw;
this._previousText = _previousText;
}
/**
* The post-patch custom state, set by the {@link AgentChatImpl.sendStream}
* generator after it applies this chunk's `customPatch`. Left `undefined` on
* chunks that do not carry a custom-state update.
*/
_custom;
/** Internal: records the post-patch custom state this chunk reports. */
_setCustom(custom) {
this._custom = custom;
}
get _content() {
return this._raw.modelChunk?.content;
}
get text() {
return partsText(this._content);
}
get reasoning() {
return partsReasoning(this._content);
}
get accumulatedText() {
return this._previousText + this.text;
}
get toolRequests() {
return toolRequestParts(this._content);
}
get data() {
return firstData(this._content);
}
get media() {
return firstMedia(this._content);
}
get artifact() {
return this._raw.artifact;
}
get custom() {
return this._custom;
}
get raw() {
return this._raw;
}
}
class DetachedTaskImpl {
constructor(snapshotId, transport) {
this.snapshotId = snapshotId;
this.transport = transport;
}
async *poll(opts) {
const intervalMs = opts?.intervalMs ?? 1e3;
while (true) {
const snap = await this.transport.getSnapshot({
snapshotId: this.snapshotId
});
if (snap) {
yield snap;
if (snap.status && TERMINAL_STATUSES.has(snap.status)) {
return;
}
}
await sleep(intervalMs);
}
}
async wait(opts) {
let last;
for await (const snap of this.poll(opts)) {
last = snap;
}
if (!last) {
throw new Error(
`Detached task ${this.snapshotId} did not produce a snapshot.`
);
}
return last;
}
abort() {
return this.transport.abort(this.snapshotId);
}
}
class AgentChatImpl {
constructor(transport, connectInit) {
this.transport = transport;
this.connectInit = connectInit;
if (connectInit?.snapshotId) {
this.snapshotId = connectInit.snapshotId;
}
if (connectInit?.sessionId) {
this.sessionId = connectInit.sessionId;
}
if (connectInit?.state) {
this.hydrateFromState(connectInit.state);
}
}
snapshotId;
sessionId;
messages = [];
artifacts = [];
clientState;
get state() {
return this.clientState?.custom;
}
/**
* Replaces the tracked `clientState`/`messages`/`artifacts` aggregates with
* (copies of) those carried by a session state.
*/
hydrateFromState(state) {
this.clientState = state;
this.messages = state?.messages ? [...state.messages] : [];
this.artifacts = state?.artifacts ? [...state.artifacts] : [];
if (state?.sessionId) {
this.sessionId = state.sessionId;
}
}
/** Loads aggregates from a server snapshot (used by `loadChat`). */
_loadFromSnapshot(snapshot) {
this.snapshotId = snapshot.snapshotId;
this.hydrateFromState(snapshot.state);
}
/**
* Builds the init for the next turn from tracked aggregates. Always returns
* an object (never `undefined`) because the agent validates `init` against
* `AgentInitSchema` - an empty object is the valid "fresh session" init.
*/
buildInit() {
if (this.snapshotId) {
return { snapshotId: this.snapshotId };
}
if (this.clientState) {
return { state: this.clientState };
}
return this.connectInit ?? {};
}
/** Applies a completed turn's output to the running aggregates. */
applyOutput(raw) {
if (raw.snapshotId !== void 0) {
this.snapshotId = raw.snapshotId;
}
if (raw.sessionId !== void 0) {
this.sessionId = raw.sessionId;
}
if (raw.state !== void 0) {
this.clientState = raw.state;
}
if (this.transport.stateManagement === void 0) {
if (raw.snapshotId !== void 0) {
this.transport.stateManagement = "server";
} else if (raw.state !== void 0) {
this.transport.stateManagement = "client";
}
}
if (raw.state?.messages !== void 0) {
this.messages = [...raw.state.messages];
} else if (raw.message) {
this.messages.push(raw.message);
}
if (raw.state?.artifacts !== void 0) {
this.artifacts = [...raw.state.artifacts];
} else if (raw.artifacts?.length) {
for (const a of raw.artifacts) {
const idx = a.name ? this.artifacts.findIndex((x) => x.name === a.name) : -1;
if (idx >= 0) {
this.artifacts[idx] = a;
} else {
this.artifacts.push(a);
}
}
}
}
/**
* Wires an optional external abort signal to a fresh {@link AbortController}
* and returns it alongside a getter for whether the turn was aborted.
*/
setupAbort(opts) {
const controller = new AbortController();
let aborted = false;
if (opts?.abortSignal) {
if (opts.abortSignal.aborted) {
controller.abort();
aborted = true;
} else {
opts.abortSignal.addEventListener("abort", () => {
aborted = true;
controller.abort();
});
}
}
return {
controller,
isAborted: () => aborted || controller.signal.aborted
};
}
/**
* Resolves a turn's raw `output` into an {@link AgentResponse}, applying it to
* the running aggregates and throwing an {@link AgentError} on a failed turn.
* Aborted turns resolve to a synthetic `aborted` response. Shared by both
* {@link send} and {@link sendStream}.
*/
buildResponse(output, isAborted, messageCountBeforeTurn) {
return (async () => {
let raw;
try {
raw = await output;
} catch (e) {
if (isAborted()) {
raw = { finishReason: "aborted" };
} else {
throw this.toAgentError(e);
}
}
if ((raw.finishReason === "failed" || raw.finishReason === "aborted") && raw.state?.messages === void 0 && !raw.message) {
this.messages.length = messageCountBeforeTurn;
}
this.applyOutput(raw);
const response = new AgentResponseImpl(
raw,
[...this.messages],
() => this.state,
() => this.sessionId
);
if (raw.finishReason === "failed") {
throw new AgentError({
message: raw.error?.message ?? "Agent turn failed.",
status: raw.error?.status ?? "UNKNOWN",
details: raw.error?.details,
state: raw.state?.custom,
snapshotId: raw.snapshotId,
response
});
}
return response;
})();
}
async send(input, opts) {
const turn = this.sendStream(input, opts);
for await (const _ of turn.stream) {
}
return turn.response;
}
sendStream(input, opts) {
const agentInput = toAgentInput(input);
if (opts?.abortSignal?.aborted) {
const aborted = (async () => {
return new AgentResponseImpl(
{ finishReason: "aborted" },
[...this.messages],
() => this.state,
() => this.sessionId
);
})();
aborted.catch(() => {
});
return {
stream: (async function* () {
})(),
response: aborted,
abort() {
}
};
}
const messageCountBeforeTurn = this.messages.length;
if (agentInput.message) {
this.messages.push(agentInput.message);
}
const init = this.buildInit();
const { controller, isAborted } = this.setupAbort(opts);
const { stream: rawStream, output } = this.transport.runTurn(
agentInput,
init,
{ abortSignal: controller.signal }
);
const responsePromise = this.buildResponse(
output,
isAborted,
messageCountBeforeTurn
);
responsePromise.catch(() => {
});
const self = this;
const stream = (async function* () {
let previousText = "";
try {
for await (const raw of rawStream) {
const chunk = new AgentChunkImpl(raw, previousText);
previousText = chunk.accumulatedText;
if (raw.customPatch) {
self.applyCustomPatch(raw.customPatch);
chunk._setCustom(self.state);
}
yield chunk;
}
} catch (e) {
if (!isAborted()) {
await responsePromise;
throw self.toAgentError(e);
}
}
await responsePromise;
})();
return {
stream,
response: responsePromise,
abort() {
controller.abort();
}
};
}
resume(resume, opts) {
return this.send({ resume }, opts);
}
resumeStream(resume, opts) {
return this.sendStream({ resume }, opts);
}
/**
* Applies a streamed RFC 6902 JSON Patch to the locally tracked custom
* state, keeping {@link state} live as the turn streams. The first patch of
* a turn is a whole-document replace (rooted at `""`) that re-bases the
* client onto the server's current baseline.
*
* Transport requirement: `customPatch` chunks carry no sequence number and
* are applied positionally, so the transport MUST deliver them in order with
* no loss. The in-order SSE channel used today satisfies this; a lossy or
* reordering transport would yield silently-wrong state (each turn does begin
* with a whole-document replace, so corruption cannot persist past a turn
* boundary).
*/
applyCustomPatch(patch) {
const current = this.clientState?.custom;
const next = applyPatch(current, patch);
if (this.clientState) {
this.clientState = { ...this.clientState, custom: next };
} else {
this.clientState = { custom: next };
}
}
async detach(input) {
const agentInput = { ...toAgentInput(input), detach: true };
if (agentInput.message) {
this.messages.push(agentInput.message);
}
const init = this.buildInit();
const controller = new AbortController();
const { output } = this.transport.runTurn(agentInput, init, {
abortSignal: controller.signal
});
const raw = await output;
this.applyOutput(raw);
if (!raw.snapshotId) {
throw new Error("detach did not return a snapshotId.");
}
return new DetachedTaskImpl(raw.snapshotId, this.transport);
}
async abort() {
if (!this.snapshotId) {
return void 0;
}
return this.transport.abort(this.snapshotId);
}
toAgentError(e) {
if (e instanceof AgentError) {
return e;
}
const message = e instanceof Error ? e.message : String(e);
const match = /^([A-Z_]+):/.exec(message);
const status = match ? match[1] : "UNKNOWN";
const raw = {
finishReason: "failed",
error: { status, message }
};
const response = new AgentResponseImpl(
raw,
[...this.messages],
() => this.clientState?.custom,
() => this.sessionId
);
return new AgentError({
message,
status,
details: e,
state: this.clientState?.custom,
snapshotId: this.snapshotId,
response
});
}
}
function createAgentAPI(transport) {
return {
chat(init) {
return new AgentChatImpl(transport, init);
},
async loadChat(opts) {
const snapshot = await transport.getSnapshot(opts);
if (!snapshot) {
const id = "snapshotId" in opts ? opts.snapshotId : `session ${opts.sessionId}`;
throw new Error(`Snapshot ${id} not found.`);
}
const chat = new AgentChatImpl(transport);
chat._loadFromSnapshot(snapshot);
return chat;
},
getSnapshot(lookup) {
const normalized = typeof lookup === "string" ? { snapshotId: lookup } : lookup;
return transport.getSnapshot(normalized);
},
abort(snapshotId) {
return transport.abort(snapshotId);
}
};
}
export {
AgentChatImpl,
AgentError,
createAgentAPI
};
//# sourceMappingURL=agent-core.mjs.map