openclaw/src/agents/transport-stream-shared.ts

110 lines
3.0 KiB
TypeScript

import { createAssistantMessageEventStream } from "@mariozechner/pi-ai";
export type TransportUsage = {
input: number;
output: number;
cacheRead: number;
cacheWrite: number;
totalTokens: number;
cost: { input: number; output: number; cacheRead: number; cacheWrite: number; total: number };
};
export type WritableTransportStream = {
push(event: unknown): void;
end(): void;
};
type TransportOutputShape = {
stopReason: string;
errorMessage?: string;
};
export function sanitizeTransportPayloadText(text: string): string {
return text.replace(
/[\uD800-\uDBFF](?![\uDC00-\uDFFF])|(?<![\uD800-\uDBFF])[\uDC00-\uDFFF]/g,
"",
);
}
export function mergeTransportHeaders(
...headerSources: Array<Record<string, string> | undefined>
): Record<string, string> | undefined {
const merged: Record<string, string> = {};
for (const headers of headerSources) {
if (headers) {
Object.assign(merged, headers);
}
}
return Object.keys(merged).length > 0 ? merged : undefined;
}
export function mergeTransportMetadata<T extends Record<string, unknown>>(
payload: T,
metadata?: Record<string, string>,
): T {
if (!metadata || Object.keys(metadata).length === 0) {
return payload;
}
const existingMetadata =
payload.metadata && typeof payload.metadata === "object" && !Array.isArray(payload.metadata)
? (payload.metadata as Record<string, string>)
: undefined;
return {
...payload,
metadata: {
...existingMetadata,
...metadata,
},
};
}
export function createEmptyTransportUsage(): TransportUsage {
return {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
}
export function createWritableTransportEventStream() {
const eventStream = createAssistantMessageEventStream();
return {
eventStream,
stream: eventStream as unknown as WritableTransportStream,
};
}
export function finalizeTransportStream(params: {
stream: WritableTransportStream;
output: TransportOutputShape;
signal?: AbortSignal;
}): void {
const { stream, output, signal } = params;
if (signal?.aborted) {
throw new Error("Request was aborted");
}
if (output.stopReason === "aborted" || output.stopReason === "error") {
throw new Error("An unknown error occurred");
}
stream.push({ type: "done", reason: output.stopReason as never, message: output as never });
stream.end();
}
export function failTransportStream(params: {
stream: WritableTransportStream;
output: TransportOutputShape;
signal?: AbortSignal;
error: unknown;
cleanup?: () => void;
}): void {
const { stream, output, signal, error, cleanup } = params;
cleanup?.();
output.stopReason = signal?.aborted ? "aborted" : "error";
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
stream.push({ type: "error", reason: output.stopReason as never, error: output as never });
stream.end();
}