Skip to content

Commit d8458a1

Browse files
vincentkocvincentkoc
committed
refactor(providers): share transport stream helpers
1 parent fcec417 commit d8458a1

5 files changed

Lines changed: 234 additions & 103 deletions

‎src/agents/anthropic-transport-stream.ts‎

Lines changed: 35 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ import Anthropic from "@anthropic-ai/sdk";
22
import type { StreamFn } from "@mariozechner/pi-agent-core";
33
import {
44
calculateCost,
5-
createAssistantMessageEventStream,
65
getEnvApiKey,
76
parseStreamingJson,
87
type AnthropicOptions,
@@ -12,9 +11,16 @@ import {
1211
type ThinkingLevel,
1312
} from "@mariozechner/pi-ai";
1413
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./copilot-dynamic-headers.js";
15-
import { sanitizeTransportPayloadText } from "./openai-transport-stream.js";
1614
import { buildGuardedModelFetch } from "./provider-transport-fetch.js";
1715
import { transformTransportMessages } from "./transport-message-transform.js";
16+
import {
17+
createEmptyTransportUsage,
18+
createWritableTransportEventStream,
19+
failTransportStream,
20+
finalizeTransportStream,
21+
mergeTransportHeaders,
22+
sanitizeTransportPayloadText,
23+
} from "./transport-stream-shared.js";
1824

1925
const CLAUDE_CODE_VERSION = "2.1.75";
2026
const CLAUDE_CODE_TOOLS = [
@@ -86,10 +92,6 @@ type MutableAssistantOutput = {
8692
errorMessage?: string;
8793
};
8894

89-
function sanitizeAnthropicText(text: string): string {
90-
return sanitizeTransportPayloadText(text);
91-
}
92-
9395
function supportsAdaptiveThinking(modelId: string): boolean {
9496
return (
9597
modelId.includes("opus-4-6") ||
@@ -143,18 +145,6 @@ function adjustMaxTokensForThinking(params: {
143145
return { maxTokens, thinkingBudget };
144146
}
145147

146-
function mergeHeaders(
147-
...headerSources: Array<Record<string, string> | undefined>
148-
): Record<string, string> | undefined {
149-
const merged: Record<string, string> = {};
150-
for (const headers of headerSources) {
151-
if (headers) {
152-
Object.assign(merged, headers);
153-
}
154-
}
155-
return Object.keys(merged).length > 0 ? merged : undefined;
156-
}
157-
158148
function isAnthropicOAuthToken(apiKey: string): boolean {
159149
return apiKey.includes("sk-ant-oat");
160150
}
@@ -197,15 +187,15 @@ function convertContentBlocks(
197187
) {
198188
const hasImages = content.some((item) => item.type === "image");
199189
if (!hasImages) {
200-
return sanitizeAnthropicText(
190+
return sanitizeTransportPayloadText(
201191
content.map((item) => ("text" in item ? item.text : "")).join("\n"),
202192
);
203193
}
204194
const blocks = content.map((block) => {
205195
if (block.type === "text") {
206196
return {
207197
type: "text",
208-
text: sanitizeAnthropicText(block.text),
198+
text: sanitizeTransportPayloadText(block.text),
209199
};
210200
}
211201
return {
@@ -245,7 +235,7 @@ function convertAnthropicMessages(
245235
if (msg.content.trim().length > 0) {
246236
params.push({
247237
role: "user",
248-
content: sanitizeAnthropicText(msg.content),
238+
content: sanitizeTransportPayloadText(msg.content),
249239
});
250240
}
251241
continue;
@@ -260,7 +250,7 @@ function convertAnthropicMessages(
260250
item.type === "text"
261251
? {
262252
type: "text",
263-
text: sanitizeAnthropicText(item.text),
253+
text: sanitizeTransportPayloadText(item.text),
264254
}
265255
: {
266256
type: "image",
@@ -293,7 +283,7 @@ function convertAnthropicMessages(
293283
if (block.text.trim().length > 0) {
294284
blocks.push({
295285
type: "text",
296-
text: sanitizeAnthropicText(block.text),
286+
text: sanitizeTransportPayloadText(block.text),
297287
});
298288
}
299289
continue;
@@ -312,12 +302,12 @@ function convertAnthropicMessages(
312302
if (!block.thinkingSignature || block.thinkingSignature.trim().length === 0) {
313303
blocks.push({
314304
type: "text",
315-
text: sanitizeAnthropicText(block.thinking),
305+
text: sanitizeTransportPayloadText(block.thinking),
316306
});
317307
} else {
318308
blocks.push({
319309
type: "thinking",
320-
thinking: sanitizeAnthropicText(block.thinking),
310+
thinking: sanitizeTransportPayloadText(block.thinking),
321311
signature: block.thinkingSignature,
322312
});
323313
}
@@ -454,7 +444,7 @@ function createAnthropicTransportClient(params: {
454444
authToken: apiKey,
455445
baseURL: model.baseUrl,
456446
dangerouslyAllowBrowser: true,
457-
defaultHeaders: mergeHeaders(
447+
defaultHeaders: mergeTransportHeaders(
458448
{
459449
accept: "application/json",
460450
"anthropic-dangerous-direct-browser-access": "true",
@@ -483,7 +473,7 @@ function createAnthropicTransportClient(params: {
483473
authToken: apiKey,
484474
baseURL: model.baseUrl,
485475
dangerouslyAllowBrowser: true,
486-
defaultHeaders: mergeHeaders(
476+
defaultHeaders: mergeTransportHeaders(
487477
{
488478
accept: "application/json",
489479
"anthropic-dangerous-direct-browser-access": "true",
@@ -504,7 +494,7 @@ function createAnthropicTransportClient(params: {
504494
apiKey,
505495
baseURL: model.baseUrl,
506496
dangerouslyAllowBrowser: true,
507-
defaultHeaders: mergeHeaders(
497+
defaultHeaders: mergeTransportHeaders(
508498
{
509499
accept: "application/json",
510500
"anthropic-dangerous-direct-browser-access": "true",
@@ -544,7 +534,7 @@ function buildAnthropicParams(
544534
? [
545535
{
546536
type: "text",
547-
text: sanitizeAnthropicText(context.systemPrompt),
537+
text: sanitizeTransportPayloadText(context.systemPrompt),
548538
...(cacheControl ? { cache_control: cacheControl } : {}),
549539
},
550540
]
@@ -554,7 +544,7 @@ function buildAnthropicParams(
554544
params.system = [
555545
{
556546
type: "text",
557-
text: sanitizeAnthropicText(context.systemPrompt),
547+
text: sanitizeTransportPayloadText(context.systemPrompt),
558548
...(cacheControl ? { cache_control: cacheControl } : {}),
559549
},
560550
];
@@ -639,22 +629,15 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
639629
return (rawModel, context, rawOptions) => {
640630
const model = rawModel as AnthropicTransportModel;
641631
const options = rawOptions as AnthropicTransportOptions | undefined;
642-
const stream = createAssistantMessageEventStream();
632+
const { eventStream, stream } = createWritableTransportEventStream();
643633
void (async () => {
644634
const output: MutableAssistantOutput = {
645635
role: "assistant",
646636
content: [],
647637
api: "anthropic-messages",
648638
provider: model.provider,
649639
model: model.id,
650-
usage: {
651-
input: 0,
652-
output: 0,
653-
cacheRead: 0,
654-
cacheWrite: 0,
655-
totalTokens: 0,
656-
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
657-
},
640+
usage: createEmptyTransportUsage(),
658641
stopReason: "stop",
659642
timestamp: Date.now(),
660643
};
@@ -895,24 +878,21 @@ export function createAnthropicMessagesTransportStreamFn(): StreamFn {
895878
calculateCost(model, output.usage);
896879
}
897880
}
898-
if (transportOptions.signal?.aborted) {
899-
throw new Error("Request was aborted");
900-
}
901-
if (output.stopReason === "aborted" || output.stopReason === "error") {
902-
throw new Error("An unknown error occurred");
903-
}
904-
stream.push({ type: "done", reason: output.stopReason as never, message: output as never });
905-
stream.end();
881+
finalizeTransportStream({ stream, output, signal: transportOptions.signal });
906882
} catch (error) {
907-
for (const block of output.content) {
908-
delete block.index;
909-
}
910-
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
911-
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
912-
stream.push({ type: "error", reason: output.stopReason as never, error: output as never });
913-
stream.end();
883+
failTransportStream({
884+
stream,
885+
output,
886+
signal: options?.signal,
887+
error,
888+
cleanup: () => {
889+
for (const block of output.content) {
890+
delete block.index;
891+
}
892+
},
893+
});
914894
}
915895
})();
916-
return stream as ReturnType<StreamFn>;
896+
return eventStream as ReturnType<StreamFn>;
917897
};
918898
}

0 commit comments

Comments
 (0)