openclaw/extensions/openrouter/stream.ts

53 lines
1.6 KiB
TypeScript

import type { StreamFn } from "@mariozechner/pi-agent-core";
import { buildProviderStreamFamilyHooks } from "openclaw/plugin-sdk/provider-stream";
import type { ProviderWrapStreamFnContext } from "openclaw/plugin-sdk/plugin-entry";
const OPENROUTER_THINKING_STREAM_HOOKS = buildProviderStreamFamilyHooks("openrouter-thinking");
function injectOpenRouterRouting(
baseStreamFn: StreamFn | undefined,
providerRouting?: Record<string, unknown>,
): StreamFn | undefined {
if (!providerRouting) {
return baseStreamFn;
}
return (model, context, options) =>
(
baseStreamFn ??
((nextModel) => {
throw new Error(
`OpenRouter routing wrapper requires an underlying streamFn for ${String(nextModel.id)}.`,
);
})
)(
{
...model,
compat: { ...model.compat, openRouterRouting: providerRouting },
} as typeof model,
context,
options,
);
}
export function wrapOpenRouterProviderStream(
ctx: ProviderWrapStreamFnContext,
): StreamFn | null | undefined {
const providerRouting =
ctx.extraParams?.provider != null && typeof ctx.extraParams.provider === "object"
? (ctx.extraParams.provider as Record<string, unknown>)
: undefined;
const routedStreamFn = providerRouting
? injectOpenRouterRouting(ctx.streamFn, providerRouting)
: ctx.streamFn;
const wrapStreamFn = OPENROUTER_THINKING_STREAM_HOOKS.wrapStreamFn ?? undefined;
if (!wrapStreamFn) {
return routedStreamFn;
}
return (
wrapStreamFn({
...ctx,
streamFn: routedStreamFn,
}) ?? undefined
);
}