Jesse_Chen
275 lines
11 KiB
TypeScript
275 lines
11 KiB
TypeScript
import type { ConsultationRuntimeState } from "../mastra/consultation-tools.ts";
|
|
import {
|
|
agentExecutionReceiptSchema,
|
|
consultationAgentPublicEventSchema,
|
|
publicActivityPhaseSchema,
|
|
type AgentExecutionReceipt,
|
|
type ConsultationAgentPublicEvent,
|
|
} from "./consultation-agent-events.ts";
|
|
import { createVisibleTextTransformer } from "./stream-text-response.ts";
|
|
|
|
type Chunk = { type?: string; payload?: Record<string, unknown>; data?: unknown };
|
|
type ChunkStream = AsyncIterable<unknown> | ReadableStream<unknown>;
|
|
|
|
async function* readChunks(stream: ChunkStream): AsyncIterable<Chunk> {
|
|
const values = Symbol.asyncIterator in stream
|
|
? stream as AsyncIterable<unknown>
|
|
: (async function* () {
|
|
const reader = (stream as ReadableStream<unknown>).getReader();
|
|
try {
|
|
while (true) {
|
|
const { done, value } = await reader.read();
|
|
if (done) return;
|
|
yield value;
|
|
}
|
|
} finally {
|
|
reader.releaseLock();
|
|
}
|
|
})();
|
|
for await (const value of values) {
|
|
if (value && typeof value === "object") yield value as Chunk;
|
|
}
|
|
}
|
|
type Status = "ready" | "degraded" | "blocked";
|
|
|
|
export const ENSURE_FINAL_RESPONSE_FALLBACK =
|
|
"本次计算已完成，但暂时没有生成可展示的回答。请换一个角度提问，我会基于已完成的计算继续说明。";
|
|
|
|
export function ensureFinalResponseText(output: string, contractIsReady: boolean) {
|
|
if (!contractIsReady || /\S/.test(output)) return null;
|
|
return ENSURE_FINAL_RESPONSE_FALLBACK;
|
|
}
|
|
|
|
type EventOptions = {
|
|
runId: string;
|
|
requestId: string;
|
|
toolStatus: () => Status;
|
|
receipt: () => AgentExecutionReceipt;
|
|
};
|
|
|
|
function activity(value: unknown): ConsultationAgentPublicEvent | null {
|
|
if (!value || typeof value !== "object") return null;
|
|
const data = value as { phase?: unknown; label?: unknown };
|
|
const phase = publicActivityPhaseSchema.safeParse(data.phase);
|
|
if (!phase.success || typeof data.label !== "string") return null;
|
|
return { type: "activity", phase: phase.data, label: data.label.slice(0, 120) };
|
|
}
|
|
|
|
function safeToolError(error: unknown) {
|
|
if (error instanceof DOMException && error.name === "AbortError") return "cancelled" as const;
|
|
if (error instanceof DOMException && error.name === "TimeoutError") return "timeout" as const;
|
|
return "calculation_failed" as const;
|
|
}
|
|
|
|
function mapChunk(
|
|
chunk: Chunk,
|
|
options: EventOptions,
|
|
startedAt: Map<string, number>,
|
|
jyotishSkillCallIds: Set<string>,
|
|
): ConsultationAgentPublicEvent[] {
|
|
const payload = chunk.payload ?? {};
|
|
if (chunk.type === "data-jyotish-activity") {
|
|
const event = activity(chunk.data);
|
|
return event ? [event] : [];
|
|
}
|
|
if (chunk.type === "tool-call") {
|
|
const toolName = payload.toolName;
|
|
const callId = typeof payload.toolCallId === "string" ? payload.toolCallId : "tool";
|
|
if (toolName === "skill" && (payload.args as { name?: unknown } | undefined)?.name === "jyotish-vedic-astrology") {
|
|
jyotishSkillCallIds.add(callId);
|
|
return [{ type: "skill.started", name: "jyotish-vedic-astrology" }];
|
|
}
|
|
if (toolName === "run-jyotish-consultation") {
|
|
startedAt.set(callId, Date.now());
|
|
return [{ type: "tool.started", callId, tool: "run-jyotish-consultation", label: "正在计算个人星盘" }];
|
|
}
|
|
}
|
|
if (chunk.type === "tool-result") {
|
|
const toolName = payload.toolName;
|
|
const callId = typeof payload.toolCallId === "string" ? payload.toolCallId : "tool";
|
|
if (toolName === "skill" && jyotishSkillCallIds.delete(callId)) {
|
|
return [{ type: "skill.completed", name: "jyotish-vedic-astrology" }];
|
|
}
|
|
if (toolName === "run-jyotish-consultation") {
|
|
return [{
|
|
type: "tool.completed", callId, tool: "run-jyotish-consultation", status: options.toolStatus(),
|
|
durationMs: Math.max(0, Date.now() - (startedAt.get(callId) ?? Date.now())),
|
|
}];
|
|
}
|
|
}
|
|
if (chunk.type === "tool-error") {
|
|
const toolName = payload.toolName;
|
|
if (toolName === "run-jyotish-consultation") {
|
|
return [{
|
|
type: "tool.failed",
|
|
callId: typeof payload.toolCallId === "string" ? payload.toolCallId : "tool",
|
|
tool: "run-jyotish-consultation",
|
|
code: safeToolError(payload.error),
|
|
}];
|
|
}
|
|
}
|
|
return [];
|
|
}
|
|
|
|
export async function collectAgentPublicEvents(stream: ChunkStream | Iterable<Chunk>, options: EventOptions) {
|
|
const events: ConsultationAgentPublicEvent[] = [{ type: "run.started", runId: options.runId, requestId: options.requestId }];
|
|
const startedAt = new Map<string, number>();
|
|
const jyotishSkillCallIds = new Set<string>();
|
|
for await (const chunk of stream instanceof ReadableStream || Symbol.asyncIterator in stream ? readChunks(stream as ChunkStream) : stream) {
|
|
events.push(...mapChunk(chunk, options, startedAt, jyotishSkillCallIds));
|
|
if (chunk.type === "text-delta" && typeof chunk.payload?.text === "string") {
|
|
events.push({ type: "answer.delta", text: chunk.payload.text });
|
|
}
|
|
}
|
|
events.push({ type: "run.completed", receipt: agentExecutionReceiptSchema.parse(options.receipt()) });
|
|
return events.map((event) => consultationAgentPublicEventSchema.parse(event));
|
|
}
|
|
|
|
type StreamAgentResponseOptions = EventOptions & {
|
|
state: ConsultationRuntimeState;
|
|
stream: ChunkStream;
|
|
transformText?: (text: string) => string;
|
|
requireTool: boolean;
|
|
retry?: () => Promise<ChunkStream>;
|
|
continueAfterDisconnect?: boolean;
|
|
headers?: HeadersInit;
|
|
onFirstActivity?: () => void | Promise<void>;
|
|
onFirstOutput?: () => void | Promise<void>;
|
|
onComplete?: (output: string, receipt: AgentExecutionReceipt) => void | Promise<void>;
|
|
onError?: (error: unknown, emitted: boolean, output: string) => void | Promise<void>;
|
|
onCancel?: (emitted: boolean) => void | Promise<void>;
|
|
};
|
|
|
|
function contractReady(options: StreamAgentResponseOptions) {
|
|
return options.state.jyotishSkillLoaded
|
|
&& (!options.requireTool || (options.state.consultationToolCompleted && options.state.consultationToolCallCount === 1));
|
|
}
|
|
|
|
export function streamAgentResponse(options: StreamAgentResponseOptions) {
|
|
const encoder = new TextEncoder();
|
|
let disconnected = false;
|
|
let settled = false;
|
|
let settling = false;
|
|
let emitted = false;
|
|
let firstActivity = false;
|
|
let firstOutput = false;
|
|
let fullOutput = "";
|
|
const startedAt = new Map<string, number>();
|
|
const jyotishSkillCallIds = new Set<string>();
|
|
const send = (controller: ReadableStreamDefaultController<Uint8Array> | undefined, event: ConsultationAgentPublicEvent) => {
|
|
if (!firstActivity && (event.type === "skill.started" || event.type === "tool.started" || event.type === "activity")) {
|
|
firstActivity = true;
|
|
void Promise.resolve(options.onFirstActivity?.()).catch(() => {});
|
|
}
|
|
if (!disconnected && controller) controller.enqueue(encoder.encode(`${JSON.stringify(consultationAgentPublicEventSchema.parse(event))}\n`));
|
|
};
|
|
|
|
async function consumeAttempt(controller: ReadableStreamDefaultController<Uint8Array> | undefined, stream: ChunkStream) {
|
|
const visible = createVisibleTextTransformer(options.transformText ?? ((value) => value));
|
|
let held = "";
|
|
let attemptOutput = "";
|
|
let composingSent = false;
|
|
const outputText = async (text: string) => {
|
|
attemptOutput += text;
|
|
held += text;
|
|
if (!held || !contractReady(options)) return;
|
|
if (!composingSent) {
|
|
composingSent = true;
|
|
send(controller, { type: "activity", phase: "answer-composition", label: "正在组织回答" });
|
|
}
|
|
if (!firstOutput && /\S/.test(held)) {
|
|
firstOutput = true;
|
|
await options.onFirstOutput?.();
|
|
}
|
|
send(controller, { type: "answer.delta", text: held });
|
|
fullOutput += held;
|
|
if (/\S/.test(held)) emitted = true;
|
|
held = "";
|
|
};
|
|
for await (const chunk of readChunks(stream)) {
|
|
for (const event of mapChunk(chunk, options, startedAt, jyotishSkillCallIds)) send(controller, event);
|
|
if (chunk.type === "text-delta" && typeof chunk.payload?.text === "string") {
|
|
await outputText(visible.push(chunk.payload.text));
|
|
}
|
|
}
|
|
await outputText(visible.finish(""));
|
|
return { held, attemptOutput };
|
|
}
|
|
|
|
const body = new ReadableStream<Uint8Array>({
|
|
start(controller) {
|
|
void (async () => {
|
|
send(controller, { type: "run.started", runId: options.runId, requestId: options.requestId });
|
|
try {
|
|
const first = await consumeAttempt(controller, options.stream);
|
|
if (!contractReady(options) && options.retry) {
|
|
if (options.state.steps.length < 32) {
|
|
options.state.steps.push({ sequence: options.state.steps.length + 1, kind: "validation", name: "runtime-contract-retry", status: "completed" });
|
|
}
|
|
send(controller, { type: "activity", phase: "loading-method", label: "正在补齐方法与计算步骤" });
|
|
await consumeAttempt(controller, await options.retry());
|
|
}
|
|
if (!contractReady(options)) throw new Error("runtime_contract_incomplete");
|
|
const ensuredFinalResponse = ensureFinalResponseText(fullOutput, contractReady(options));
|
|
if (ensuredFinalResponse) {
|
|
if (options.state.steps.length < 32) {
|
|
options.state.steps.push({ sequence: options.state.steps.length + 1, kind: "validation", name: "ensure-final-response", status: "completed" });
|
|
}
|
|
if (!firstOutput) {
|
|
firstOutput = true;
|
|
await options.onFirstOutput?.();
|
|
}
|
|
send(controller, { type: "answer.delta", text: ensuredFinalResponse });
|
|
fullOutput = ensuredFinalResponse;
|
|
emitted = true;
|
|
}
|
|
if (!/\S/.test(fullOutput)) {
|
|
if (/\S/.test(first.held) || /\S/.test(first.attemptOutput)) throw new Error("runtime_contract_incomplete");
|
|
throw new Error("empty_answer");
|
|
}
|
|
settling = true;
|
|
const receipt = agentExecutionReceiptSchema.parse(options.receipt());
|
|
await options.onComplete?.(fullOutput, receipt);
|
|
settled = true;
|
|
settling = false;
|
|
send(controller, { type: "run.completed", receipt });
|
|
if (!disconnected) controller.close();
|
|
} catch (error) {
|
|
if (settled) return;
|
|
settled = true;
|
|
settling = false;
|
|
try {
|
|
await options.onError?.(error, emitted, fullOutput);
|
|
} catch {}
|
|
const code = error instanceof Error && error.message === "runtime_contract_incomplete"
|
|
? "runtime_contract_incomplete" as const
|
|
: error instanceof Error && error.message === "empty_answer"
|
|
? "empty_answer" as const
|
|
: "calculation_failed" as const;
|
|
send(controller, { type: "run.failed", code, message: code === "runtime_contract_incomplete" ? "Agent 未完成必要的方法与计算步骤，本次不会扣点。" : "咨询暂时无法完成，本次不会扣点。" });
|
|
if (!disconnected) controller.close();
|
|
}
|
|
})();
|
|
},
|
|
async cancel() {
|
|
if (settled) return;
|
|
if (settling || options.continueAfterDisconnect) {
|
|
disconnected = true;
|
|
return;
|
|
}
|
|
settled = true;
|
|
await options.onCancel?.(emitted);
|
|
},
|
|
});
|
|
return new Response(body, {
|
|
headers: {
|
|
"cache-control": "no-cache, no-transform",
|
|
"content-type": "application/x-ndjson; charset=utf-8",
|
|
"x-accel-buffering": "no",
|
|
"x-ayanam-mode": "mastra-agentic",
|
|
"x-ayanam-request-id": options.requestId,
|
|
...options.headers,
|
|
},
|
|
});
|
|
}
|