openclaw/extensions/amazon-bedrock/index.ts

100 lines
3.2 KiB
TypeScript

import type { StreamFn } from "@mariozechner/pi-agent-core";
import { definePluginEntry } from "openclaw/plugin-sdk/plugin-entry";
import {
createBedrockNoCacheWrapper,
isAnthropicBedrockModel,
streamWithPayloadPatch,
} from "openclaw/plugin-sdk/provider-stream";
import {
mergeImplicitBedrockProvider,
resolveBedrockConfigApiKey,
resolveImplicitBedrockProvider,
} from "./api.js";
type GuardrailConfig = {
guardrailIdentifier: string;
guardrailVersion: string;
streamProcessingMode?: "sync" | "async";
trace?: "enabled" | "disabled" | "enabled_full";
};
function createGuardrailWrapStreamFn(
innerWrapStreamFn: (ctx: { modelId: string; streamFn?: StreamFn }) => StreamFn | null | undefined,
guardrailConfig: GuardrailConfig,
): (ctx: { modelId: string; streamFn?: StreamFn }) => StreamFn | null | undefined {
return (ctx) => {
const inner = innerWrapStreamFn(ctx);
if (!inner) return inner;
return (model, context, options) => {
return streamWithPayloadPatch(inner, model, context, options, (payload) => {
const gc: Record<string, unknown> = {
guardrailIdentifier: guardrailConfig.guardrailIdentifier,
guardrailVersion: guardrailConfig.guardrailVersion,
};
if (guardrailConfig.streamProcessingMode) {
gc.streamProcessingMode = guardrailConfig.streamProcessingMode;
}
if (guardrailConfig.trace) {
gc.trace = guardrailConfig.trace;
}
payload.guardrailConfig = gc;
});
};
};
}
const PROVIDER_ID = "amazon-bedrock";
const CLAUDE_46_MODEL_RE = /claude-(?:opus|sonnet)-4(?:\.|-)6(?:$|[-.])/i;
export default definePluginEntry({
id: PROVIDER_ID,
name: "Amazon Bedrock Provider",
description: "Bundled Amazon Bedrock provider policy plugin",
register(api) {
const guardrail = (api.pluginConfig as Record<string, unknown> | undefined)?.guardrail as
| GuardrailConfig
| undefined;
const baseWrapStreamFn = ({ modelId, streamFn }: { modelId: string; streamFn?: StreamFn }) =>
isAnthropicBedrockModel(modelId) ? streamFn : createBedrockNoCacheWrapper(streamFn);
const wrapStreamFn =
guardrail?.guardrailIdentifier && guardrail?.guardrailVersion
? createGuardrailWrapStreamFn(baseWrapStreamFn, guardrail)
: baseWrapStreamFn;
api.registerProvider({
id: PROVIDER_ID,
label: "Amazon Bedrock",
docsPath: "/providers/models",
auth: [],
catalog: {
order: "simple",
run: async (ctx) => {
const implicit = await resolveImplicitBedrockProvider({
config: ctx.config,
env: ctx.env,
});
if (!implicit) {
return null;
}
return {
provider: mergeImplicitBedrockProvider({
existing: ctx.config.models?.providers?.[PROVIDER_ID],
implicit,
}),
};
},
},
resolveConfigApiKey: ({ env }) => resolveBedrockConfigApiKey(env),
capabilities: {
providerFamily: "anthropic",
dropThinkingBlockModelHints: ["claude"],
},
wrapStreamFn,
resolveDefaultThinkingLevel: ({ modelId }) =>
CLAUDE_46_MODEL_RE.test(modelId.trim()) ? "adaptive" : undefined,
});
},
});