admin/Moltbot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

refactor: share stream payload patch helper

Browse Source
This commit is contained in:
Peter Steinberger
2026-03-13 19:26:29 +00:00
parent e08dc6f0af
commit 92dbb59b79
3 changed files with 30 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts
View File

@@ -7,6 +7,7 @@ import {
usesOpenAiStringModeAnthropicToolChoice,
} from "../provider-capabilities.js";
import { log } from "./logger.js";
import { streamWithPayloadPatch } from "./stream-payload-utils.js";
const ANTHROPIC_CONTEXT_1M_BETA = "context-1m-2025-08-07";
const ANTHROPIC_1M_MODEL_PREFIXES = ["claude-opus-4", "claude-sonnet-4"] as const;
@@ -341,18 +342,10 @@ export function createAnthropicFastModeWrapper(
return underlying(model, context, options);
}
const originalOnPayload = options?.onPayload;
return underlying(model, context, {
...options,
onPayload: (payload) => {
if (payload && typeof payload === "object") {
const payloadObj = payload as Record<string, unknown>;
if (payloadObj.service_tier === undefined) {
payloadObj.service_tier = serviceTier;
}
}
return originalOnPayload?.(payload, model);
},
return streamWithPayloadPatch(underlying, model, context, options, (payloadObj) => {
if (payloadObj.service_tier === undefined) {
payloadObj.service_tier = serviceTier;
}
});
};
}

17
src/agents/pi-embedded-runner/openai-stream-wrappers.ts
View File

@@ -2,6 +2,7 @@ import type { StreamFn } from "@mariozechner/pi-agent-core";
import type { SimpleStreamOptions } from "@mariozechner/pi-ai";
import { streamSimple } from "@mariozechner/pi-ai";
import { log } from "./logger.js";
import { streamWithPayloadPatch } from "./stream-payload-utils.js";
type OpenAIServiceTier = "auto" | "default" | "flex" | "priority";
type OpenAIReasoningEffort = "low" | "medium" | "high";
@@ -325,18 +326,10 @@ export function createOpenAIServiceTierWrapper(
) {
return underlying(model, context, options);
}
const originalOnPayload = options?.onPayload;
return underlying(model, context, {
...options,
onPayload: (payload) => {
if (payload && typeof payload === "object") {
const payloadObj = payload as Record<string, unknown>;
if (payloadObj.service_tier === undefined) {
payloadObj.service_tier = serviceTier;
}
}
return originalOnPayload?.(payload, model);
},
return streamWithPayloadPatch(underlying, model, context, options, (payloadObj) => {
if (payloadObj.service_tier === undefined) {
payloadObj.service_tier = serviceTier;
}
});
};
}

20
src/agents/pi-embedded-runner/stream-payload-utils.ts Normal file
View File

@@ -0,0 +1,20 @@
import type { StreamFn } from "@mariozechner/pi-agent-core";
export function streamWithPayloadPatch(
underlying: StreamFn,
model: Parameters<StreamFn>[0],
context: Parameters<StreamFn>[1],
options: Parameters<StreamFn>[2],
patchPayload: (payload: Record<string, unknown>) => void,
) {
const originalOnPayload = options?.onPayload;
return underlying(model, context, {
...options,
onPayload: (payload) => {
if (payload && typeof payload === "object") {
patchPayload(payload as Record<string, unknown>);
}
return originalOnPayload?.(payload, model);
},
});
}