Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion extensions/amazon-bedrock/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -264,19 +264,39 @@ describe("amazon-bedrock provider plugin", () => {
});
const result = callWrappedStream(provider, ANTHROPIC_MODEL, ANTHROPIC_MODEL_DESCRIPTOR);

// Anthropic models should get guardrailConfig
// Anthropic models get guardrailConfig + default anthropic_beta flags.
expect(result._capturedPayload).toEqual({
guardrailConfig: {
guardrailIdentifier: "guardrail-anthropic",
guardrailVersion: "2",
streamProcessingMode: "async",
trace: "disabled",
},
additionalModelRequestFields: {
anthropic_beta: [
"fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14",
],
},
});
// Anthropic models should NOT get cacheRetention: "none"
expect(result).not.toHaveProperty("cacheRetention", "none");
});

it("injects default anthropic_beta flags for Anthropic Bedrock models without guardrail", async () => {
const provider = await registerWithConfig(undefined);
const result = callWrappedStream(provider, ANTHROPIC_MODEL, ANTHROPIC_MODEL_DESCRIPTOR);

expect(result._capturedPayload).toEqual({
additionalModelRequestFields: {
anthropic_beta: [
"fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14",
],
},
});
});

it("injects guardrailConfig for non-Anthropic models with cacheRetention: none", async () => {
const provider = await registerWithConfig({
guardrail: {
Expand Down
20 changes: 19 additions & 1 deletion extensions/amazon-bedrock/register.sync.runtime.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,22 @@
} from "./api.js";
import { bedrockMemoryEmbeddingProviderAdapter } from "./memory-embedding-adapter.js";

const PI_AI_DEFAULT_ANTHROPIC_BETAS = [
"fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14",
] as const;

function createBedrockAnthropicBetasWrapper(baseStreamFn: StreamFn | undefined): StreamFn | undefined {
if (!baseStreamFn) return baseStreamFn;

Check failure on line 25 in extensions/amazon-bedrock/register.sync.runtime.ts

View workflow job for this annotation

GitHub Actions / check-lint

eslint(curly)

Expected { after 'if' condition.

Check failure on line 25 in extensions/amazon-bedrock/register.sync.runtime.ts

View workflow job for this annotation

GitHub Actions / check-additional-extension-bundled

eslint(curly)

Expected { after 'if' condition.
return (model, context, options) =>
streamWithPayloadPatch(baseStreamFn, model, context, options, (payload) => {
payload.additionalModelRequestFields = {
...(payload.additionalModelRequestFields as Record<string, unknown> | undefined),
anthropic_beta: [...PI_AI_DEFAULT_ANTHROPIC_BETAS],
};
});
}

type GuardrailConfig = {
guardrailIdentifier: string;
guardrailVersion: string;
Expand Down Expand Up @@ -82,7 +98,9 @@
api.registerMemoryEmbeddingProvider(bedrockMemoryEmbeddingProviderAdapter);

const baseWrapStreamFn = ({ modelId, streamFn }: { modelId: string; streamFn?: StreamFn }) =>
isAnthropicBedrockModel(modelId) ? streamFn : createBedrockNoCacheWrapper(streamFn);
isAnthropicBedrockModel(modelId)
? createBedrockAnthropicBetasWrapper(streamFn)
: createBedrockNoCacheWrapper(streamFn);

const cacheWrapStreamFn =
guardrail?.guardrailIdentifier && guardrail?.guardrailVersion
Expand Down
Loading