Skip to content

Commit 0ea823d

Browse files
authored
Merge pull request #180 from berkmancenter/cj/bedrock-v2
Migrate Bedrock integration to V2
2 parents 3b4757e + c83b069 commit 0ea823d

5 files changed

Lines changed: 234 additions & 65 deletions

File tree

docs/pages/installing/index.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@ Note that this will work for any OpenAI compatible LLM provider.
3131
1. Configure `BEDROCK_API_KEY` and `BEDROCK_BASE_URL` in your `.env` file.
3232
2. When creating a Conversation with an Agent, specify `llmPlatform` to be `bedrock` and `llmModel` to be an available Bedrock model.
3333

34+
Set `BEDROCK_BASE_URL` to the endpoint up to (but not including) the `/model` path segment. LLM Engine appends the standard Amazon Bedrock InvokeModel path, so each request goes to `{BEDROCK_BASE_URL}/model/{llmModel}/invoke`. Point it at the Amazon Bedrock runtime host or any gateway that fronts it.
35+
36+
Authentication sends `BEDROCK_API_KEY` as an `x-api-key` header, which assumes a gateway that signs the upstream AWS request for you. To call Amazon Bedrock directly (AWS requires SigV4-signed requests), replace the transport in `src/agents/helpers/bedrockGateway.ts`, or use LangChain's BedrockChat with AWS credentials instead.
37+
3438
### Google Generative AI (including Gemini)
3539

3640
1. Configure `GOOGLE_API_KEY` and `GOOGLE_BASE_URL` in your `.env` file.
Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import logger from '../../config/logger.js'
2+
import config from '../../config/config.js'
3+
import { transformPayloadForClaude } from './claudeHandler.js'
4+
5+
/**
6+
* Bedrock gateway transport.
7+
*
8+
* This module owns how we reach a specific Amazon Bedrock endpoint: the request URL and the
9+
* auth header. It is the one place to change when pointing the app at a different Bedrock
10+
* provider.
11+
*
12+
* The default targets an API-key gateway. It sends the native Bedrock InvokeModel request to
13+
* `{BEDROCK_BASE_URL}/model/{modelId}/invoke` with an `x-api-key` header, and assumes the
14+
* gateway signs the upstream AWS request (SigV4) on its side.
15+
*
16+
* If you call Amazon Bedrock directly, AWS requires SigV4-signed requests rather than an api
17+
* key. In that case, set `BEDROCK_BASE_URL` to your endpoint and replace the auth below, or
18+
* drop this custom fetch and use LangChain's BedrockChat with AWS credentials instead.
19+
*/
20+
21+
// Allowed Bedrock model id characters. Reject anything else (/, ?, #, ...) so a model id
22+
// cannot alter the URL path; the colon in dated ids is valid in a path segment and stays.
23+
const BEDROCK_MODEL_ID_PATTERN = /^[A-Za-z0-9._:-]+$/
24+
25+
/**
26+
* Builds the Bedrock InvokeModel URL for a model: `{baseUrl}/model/{modelId}/invoke`.
27+
*
28+
* `baseUrl` is everything up to the `/model` segment for your endpoint. For Amazon Bedrock
29+
* that is the bedrock-runtime host; for a gateway it is the gateway base path. The model id
30+
* goes straight into the path and can come from user-set config, so this validates it against
31+
* the known Bedrock character set first and rejects anything else. The colon in dated ids
32+
* (for example ...-v1:0) is valid in a URL path and is kept raw.
33+
*
34+
* @throws if the model id contains characters that are not valid in a Bedrock model id.
35+
*/
36+
export function buildBedrockInvokeUrl(baseUrl: string, modelId: string): string {
37+
if (!BEDROCK_MODEL_ID_PATTERN.test(modelId)) {
38+
throw new Error(`Invalid Bedrock model id: "${modelId}"`)
39+
}
40+
// Drop any trailing slash on the base url so the path does not end up with a double slash.
41+
const normalizedBaseUrl = baseUrl.replace(/\/+$/, '')
42+
return `${normalizedBaseUrl}/model/${modelId}/invoke`
43+
}
44+
45+
// Create a custom fetch function for Bedrock Claude or legacy LLM
46+
export function createClaudeFetchFn(defaultLLMModel: string, defaultLLMPlatform: string) {
47+
return async function fetchFn(url: string, init: Parameters<typeof fetch>[1]) {
48+
const fetchImpl = async () => {
49+
try {
50+
let bodyContent: unknown = {}
51+
if (init?.body) {
52+
if (
53+
typeof init.body === 'string' &&
54+
(init.body.trim().toLowerCase().startsWith('<!doctype') || init.body.trim().toLowerCase().startsWith('<html'))
55+
) {
56+
logger.error('init.body appears to be HTML, not JSON:', init.body)
57+
throw new Error('init.body is HTML, not JSON')
58+
}
59+
try {
60+
bodyContent = JSON.parse(init.body as string)
61+
} catch (err) {
62+
logger.error('init.body is not valid JSON:', init.body)
63+
throw err
64+
}
65+
}
66+
67+
// Transform payload for Claude if needed
68+
bodyContent = transformPayloadForClaude(bodyContent, defaultLLMModel, defaultLLMPlatform)
69+
70+
// Send the native Bedrock InvokeModel request body directly.
71+
const bodyString = JSON.stringify(bodyContent)
72+
const modifiedInit = {
73+
...(init || {}),
74+
body: bodyString,
75+
headers: {
76+
'Content-Type': 'application/json',
77+
// Auth for an api-key gateway. Direct AWS Bedrock would use SigV4 signing instead.
78+
'x-api-key': config.llms.bedrock.key
79+
}
80+
}
81+
// Bedrock routes by model id in the path; buildBedrockInvokeUrl validates the id first.
82+
const invokeUrl = buildBedrockInvokeUrl(config.llms.bedrock.baseUrl, defaultLLMModel)
83+
const response = await fetch(invokeUrl, modifiedInit)
84+
const contentType = response.headers.get('content-type') || ''
85+
if (!response.ok) {
86+
const responseText = await response.text()
87+
logger.error('Error response Text from proxy:', responseText)
88+
throw new Error(`Bedrock proxy error: ${response.status} ${response.statusText}\n${responseText}`)
89+
}
90+
if (contentType.includes('application/json')) {
91+
return response
92+
}
93+
const responseText = await response.text()
94+
throw new Error(`Expected JSON from Bedrock proxy, got ${contentType}: ${responseText}`)
95+
} catch (err) {
96+
logger.error('Error in fetchFn:', err)
97+
throw err
98+
}
99+
}
100+
return fetchImpl()
101+
}
102+
}

src/agents/helpers/claudeHandler.ts

Lines changed: 0 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,3 @@
1-
import logger from '../../config/logger.js'
2-
import config from '../../config/config.js'
3-
41
// Helper to determine if Bedrock Claude format should be used
52
export function shouldUseClaudeFormat(model: string | undefined, platform: string | undefined): boolean {
63
if (platform && typeof platform === 'string') {
@@ -129,64 +126,3 @@ export function transformPayloadForClaude(bodyContent: unknown, defaultLLMModel:
129126
tools
130127
})
131128
}
132-
133-
// Create a custom fetch function for Bedrock Claude or legacy LLM
134-
export function createClaudeFetchFn(defaultLLMModel: string, defaultLLMPlatform: string) {
135-
return async function fetchFn(url: string, init: Parameters<typeof fetch>[1]) {
136-
const fetchImpl = async () => {
137-
try {
138-
let bodyContent: unknown = {}
139-
if (init?.body) {
140-
if (
141-
typeof init.body === 'string' &&
142-
(init.body.trim().toLowerCase().startsWith('<!doctype') || init.body.trim().toLowerCase().startsWith('<html'))
143-
) {
144-
logger.error('init.body appears to be HTML, not JSON:', init.body)
145-
throw new Error('init.body is HTML, not JSON')
146-
}
147-
try {
148-
bodyContent = JSON.parse(init.body as string)
149-
} catch (err) {
150-
logger.error('init.body is not valid JSON:', init.body)
151-
throw err
152-
}
153-
}
154-
155-
// Transform payload for Claude if needed
156-
bodyContent = transformPayloadForClaude(bodyContent, defaultLLMModel, defaultLLMPlatform)
157-
158-
const body = {
159-
body: bodyContent,
160-
modelId: defaultLLMModel,
161-
contentType: 'application/json',
162-
accept: 'application/json'
163-
}
164-
const bodyString = JSON.stringify(body)
165-
const modifiedInit = {
166-
...(init || {}),
167-
body: bodyString,
168-
headers: {
169-
'Content-Type': 'application/json',
170-
'x-api-key': config.llms.bedrock.key
171-
}
172-
}
173-
const response = await fetch(config.llms.bedrock.baseUrl, modifiedInit)
174-
const contentType = response.headers.get('content-type') || ''
175-
if (!response.ok) {
176-
const responseText = await response.text()
177-
logger.error('Error response Text from proxy:', responseText)
178-
throw new Error(`Bedrock proxy error: ${response.status} ${response.statusText}\n${responseText}`)
179-
}
180-
if (contentType.includes('application/json')) {
181-
return response
182-
}
183-
const responseText = await response.text()
184-
throw new Error(`Expected JSON from Bedrock proxy, got ${contentType}: ${responseText}`)
185-
} catch (err) {
186-
logger.error('Error in fetchFn:', err)
187-
throw err
188-
}
189-
}
190-
return fetchImpl()
191-
}
192-
}

src/agents/helpers/getModelChat.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { google } from 'googleapis'
55
import { ChatGoogleGenerativeAI } from '@langchain/google-genai'
66
import { GoogleGenerativeAI } from '@google/generative-ai'
77
import config from '../../config/config.js'
8-
import { createClaudeFetchFn } from './claudeHandler.js'
8+
import { createClaudeFetchFn } from './bedrockGateway.js'
99
import { LlmPlatforms, LlmPlatformDetails, LlmModelDetails } from '../../types/index.types.js'
1010

1111
const PERSPECTIVE_API_URL = 'https://commentanalyzer.googleapis.com/$discovery/rest?version=v1alpha1'
Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
import { createClaudeFetchFn, buildBedrockInvokeUrl } from '../../../src/agents/helpers/bedrockGateway.js'
2+
import config from '../../../src/config/config.js'
3+
4+
describe('createClaudeFetchFn (Bedrock gateway transport)', () => {
5+
// The transport posts the native Bedrock InvokeModel request to {baseUrl}/model/{modelId}/invoke
6+
// with an x-api-key header. These tests pin that contract. The model id is colon-free here so the
7+
// URL assertion does not depend on how the colon in dated ids is handled.
8+
const BASE_URL = 'https://bedrock.example.com'
9+
const API_KEY = 'test-bedrock-key'
10+
const MODEL_ID = 'us.anthropic.claude-sonnet-4-6'
11+
12+
// Stash the real values once; each test swaps in known values and afterEach restores them.
13+
const originalFetch = globalThis.fetch
14+
const originalBaseUrl = config.llms.bedrock.baseUrl
15+
const originalKey = config.llms.bedrock.key
16+
17+
let capturedUrl: string | undefined
18+
let capturedInit: { body: string; headers: Record<string, string> } | undefined
19+
20+
beforeEach(() => {
21+
config.llms.bedrock.baseUrl = BASE_URL
22+
config.llms.bedrock.key = API_KEY
23+
capturedUrl = undefined
24+
capturedInit = undefined
25+
26+
globalThis.fetch = ((url: string, init: { body: string; headers: Record<string, string> }) => {
27+
capturedUrl = url
28+
capturedInit = init
29+
return Promise.resolve({
30+
ok: true,
31+
status: 200,
32+
statusText: 'OK',
33+
headers: { get: () => 'application/json' },
34+
json: async () => ({ content: [{ type: 'text', text: 'ok' }] }),
35+
text: async () => '{}'
36+
})
37+
}) as unknown as typeof globalThis.fetch
38+
})
39+
40+
afterEach(() => {
41+
globalThis.fetch = originalFetch
42+
config.llms.bedrock.baseUrl = originalBaseUrl
43+
config.llms.bedrock.key = originalKey
44+
})
45+
46+
const callFetchFn = () => {
47+
const fetchFn = createClaudeFetchFn(MODEL_ID, 'bedrock')
48+
return fetchFn('https://ignored.example', {
49+
body: JSON.stringify({
50+
system: 'You are a helpful assistant',
51+
messages: [{ role: 'user', content: 'Hello world' }],
52+
max_tokens: 100
53+
})
54+
})
55+
}
56+
57+
it('POSTs to the /model/{modelId}/invoke endpoint', async () => {
58+
await callFetchFn()
59+
expect(capturedUrl).toBe(`${BASE_URL}/model/${MODEL_ID}/invoke`)
60+
})
61+
62+
it('sends the native Bedrock request body, not a wrapper envelope', async () => {
63+
await callFetchFn()
64+
const sentBody = JSON.parse(capturedInit!.body)
65+
66+
// Native Bedrock fields sit at the top level of the request.
67+
expect(sentBody.anthropic_version).toBe('bedrock-2023-05-31')
68+
expect(sentBody.max_tokens).toBe(100)
69+
expect(Array.isArray(sentBody.messages)).toBe(true)
70+
71+
// No wrapper keys around the request.
72+
expect(sentBody).not.toHaveProperty('body')
73+
expect(sentBody).not.toHaveProperty('modelId')
74+
expect(sentBody).not.toHaveProperty('accept')
75+
expect(sentBody).not.toHaveProperty('contentType')
76+
})
77+
78+
it('keeps the x-api-key auth header from config', async () => {
79+
await callFetchFn()
80+
expect(capturedInit!.headers['x-api-key']).toBe(API_KEY)
81+
})
82+
83+
it('refuses an unsafe model id before calling fetch', async () => {
84+
const fetchFn = createClaudeFetchFn('evil/../path', 'bedrock')
85+
await expect(
86+
fetchFn('https://ignored.example', {
87+
body: JSON.stringify({
88+
system: 'You are a helpful assistant',
89+
messages: [{ role: 'user', content: 'Hello world' }],
90+
max_tokens: 100
91+
})
92+
})
93+
).rejects.toThrow()
94+
expect(capturedUrl).toBeUndefined()
95+
})
96+
})
97+
98+
describe('buildBedrockInvokeUrl', () => {
99+
const BASE = 'https://bedrock.example.com'
100+
101+
it('builds the /model/{modelId}/invoke path', () => {
102+
expect(buildBedrockInvokeUrl(BASE, 'us.anthropic.claude-sonnet-4-6')).toBe(
103+
`${BASE}/model/us.anthropic.claude-sonnet-4-6/invoke`
104+
)
105+
})
106+
107+
it('keeps the raw colon in dated model ids', () => {
108+
expect(buildBedrockInvokeUrl(BASE, 'us.anthropic.claude-haiku-4-5-20251001-v1:0')).toBe(
109+
`${BASE}/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke`
110+
)
111+
})
112+
113+
it('strips a trailing slash from the base url so the path has no double slash', () => {
114+
expect(buildBedrockInvokeUrl(`${BASE}/`, 'us.anthropic.claude-sonnet-4-6')).toBe(
115+
`${BASE}/model/us.anthropic.claude-sonnet-4-6/invoke`
116+
)
117+
})
118+
119+
it('rejects model ids that contain URL-structural characters', () => {
120+
expect(() => buildBedrockInvokeUrl(BASE, 'foo/bar')).toThrow()
121+
expect(() => buildBedrockInvokeUrl(BASE, 'foo?x=1')).toThrow()
122+
expect(() => buildBedrockInvokeUrl(BASE, 'foo#frag')).toThrow()
123+
expect(() => buildBedrockInvokeUrl(BASE, 'foo%2e')).toThrow()
124+
expect(() => buildBedrockInvokeUrl(BASE, 'has space')).toThrow()
125+
expect(() => buildBedrockInvokeUrl(BASE, '')).toThrow()
126+
})
127+
})

0 commit comments

Comments
 (0)