chore: import upstream snapshot with attribution
CI / Migrate Dev DB (push) Has been skipped
CI / Detect Version (push) Has been cancelled
CI / Migrate DB (push) Has been cancelled
CI / Build Dev ECR (./docker/app.Dockerfile, ECR_APP) (push) Has been cancelled
CI / Build Dev ECR (./docker/db.Dockerfile, ECR_MIGRATIONS) (push) Has been cancelled
CI / Build Dev ECR (./docker/pii.Dockerfile, ECR_PII) (push) Has been cancelled
CI / Build Dev ECR (./docker/realtime.Dockerfile, ECR_REALTIME) (push) Has been cancelled
CI / Deploy Trigger.dev (Dev) (push) Has been cancelled
CI / Build AMD64 (./docker/app.Dockerfile, ECR_APP, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build AMD64 (./docker/db.Dockerfile, ECR_MIGRATIONS, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build AMD64 (./docker/pii.Dockerfile, ECR_PII, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build AMD64 (./docker/realtime.Dockerfile, ECR_REALTIME, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/app.Dockerfile, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/db.Dockerfile, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/pii.Dockerfile, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/realtime.Dockerfile, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Check Docs Changes (push) Has been cancelled
CI / Process Docs (push) Has been cancelled
CI / Create GitHub Release (push) Has been cancelled
CI / Test and Build (push) Has been cancelled
Publish CLI Package / publish-npm (push) Has been cancelled
Publish Python SDK / publish-pypi (push) Has been cancelled
Publish TypeScript SDK / publish-npm (push) Has been cancelled
CI / Migrate Dev DB (push) Has been skipped
CI / Detect Version (push) Has been cancelled
CI / Migrate DB (push) Has been cancelled
CI / Build Dev ECR (./docker/app.Dockerfile, ECR_APP) (push) Has been cancelled
CI / Build Dev ECR (./docker/db.Dockerfile, ECR_MIGRATIONS) (push) Has been cancelled
CI / Build Dev ECR (./docker/pii.Dockerfile, ECR_PII) (push) Has been cancelled
CI / Build Dev ECR (./docker/realtime.Dockerfile, ECR_REALTIME) (push) Has been cancelled
CI / Deploy Trigger.dev (Dev) (push) Has been cancelled
CI / Build AMD64 (./docker/app.Dockerfile, ECR_APP, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build AMD64 (./docker/db.Dockerfile, ECR_MIGRATIONS, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build AMD64 (./docker/pii.Dockerfile, ECR_PII, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build AMD64 (./docker/realtime.Dockerfile, ECR_REALTIME, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/app.Dockerfile, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/db.Dockerfile, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/pii.Dockerfile, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/realtime.Dockerfile, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Check Docs Changes (push) Has been cancelled
CI / Process Docs (push) Has been cancelled
CI / Create GitHub Release (push) Has been cancelled
CI / Test and Build (push) Has been cancelled
Publish CLI Package / publish-npm (push) Has been cancelled
Publish Python SDK / publish-pypi (push) Has been cancelled
Publish TypeScript SDK / publish-npm (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,123 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
const mockSend = vi.fn()
|
||||
|
||||
vi.mock('@aws-sdk/client-bedrock-runtime', () => ({
|
||||
BedrockRuntimeClient: vi.fn().mockImplementation(
|
||||
class {
|
||||
send = mockSend
|
||||
}
|
||||
),
|
||||
ConverseCommand: vi.fn(),
|
||||
ConverseStreamCommand: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/providers/bedrock/utils', () => ({
|
||||
getBedrockInferenceProfileId: vi
|
||||
.fn()
|
||||
.mockReturnValue('us.anthropic.claude-3-5-sonnet-20241022-v2:0'),
|
||||
checkForForcedToolUsage: vi.fn(),
|
||||
createReadableStreamFromBedrockStream: vi.fn(),
|
||||
generateToolUseId: vi.fn().mockReturnValue('tool-1'),
|
||||
}))
|
||||
|
||||
vi.mock('@/providers/models', () => ({
|
||||
getProviderFileAttachment: vi
|
||||
.fn()
|
||||
.mockReturnValue({ maxBytes: 10 * 1024 * 1024, strategy: 'inline' }),
|
||||
INLINE_ATTACHMENT_MAX_BYTES: 10 * 1024 * 1024,
|
||||
getProviderModels: vi.fn().mockReturnValue([]),
|
||||
getProviderDefaultModel: vi.fn().mockReturnValue('us.anthropic.claude-3-5-sonnet-20241022-v2:0'),
|
||||
}))
|
||||
|
||||
vi.mock('@/providers/utils', () => ({
|
||||
calculateCost: vi.fn().mockReturnValue({ input: 0, output: 0, total: 0, pricing: null }),
|
||||
prepareToolExecution: vi.fn(),
|
||||
prepareToolsWithUsageControl: vi.fn().mockReturnValue({
|
||||
tools: [],
|
||||
toolChoice: 'auto',
|
||||
forcedTools: [],
|
||||
}),
|
||||
sumToolCosts: vi.fn().mockReturnValue(0),
|
||||
}))
|
||||
|
||||
vi.mock('@/tools', () => ({
|
||||
executeTool: vi.fn(),
|
||||
}))
|
||||
|
||||
import { BedrockRuntimeClient } from '@aws-sdk/client-bedrock-runtime'
|
||||
import { bedrockProvider } from '@/providers/bedrock/index'
|
||||
import { clearProviderClientCacheForTests } from '@/providers/client-cache'
|
||||
|
||||
describe('bedrockProvider credential handling', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
clearProviderClientCacheForTests()
|
||||
mockSend.mockResolvedValue({
|
||||
output: { message: { content: [{ text: 'response' }] } },
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
})
|
||||
})
|
||||
|
||||
const baseRequest = {
|
||||
model: 'us.anthropic.claude-3-5-sonnet-20241022-v2:0',
|
||||
systemPrompt: 'You are helpful.',
|
||||
messages: [{ role: 'user' as const, content: 'Hello' }],
|
||||
}
|
||||
|
||||
it('throws when only bedrockAccessKeyId is provided', async () => {
|
||||
await expect(
|
||||
bedrockProvider.executeRequest({
|
||||
...baseRequest,
|
||||
bedrockAccessKeyId: 'AKIAIOSFODNN7EXAMPLE',
|
||||
})
|
||||
).rejects.toThrow('Both bedrockAccessKeyId and bedrockSecretKey must be provided together')
|
||||
})
|
||||
|
||||
it('throws when only bedrockSecretKey is provided', async () => {
|
||||
await expect(
|
||||
bedrockProvider.executeRequest({
|
||||
...baseRequest,
|
||||
bedrockSecretKey: 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
|
||||
})
|
||||
).rejects.toThrow('Both bedrockAccessKeyId and bedrockSecretKey must be provided together')
|
||||
})
|
||||
|
||||
it('creates client with explicit credentials when both are provided', async () => {
|
||||
await bedrockProvider.executeRequest({
|
||||
...baseRequest,
|
||||
bedrockAccessKeyId: 'AKIAIOSFODNN7EXAMPLE',
|
||||
bedrockSecretKey: 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
|
||||
})
|
||||
|
||||
expect(BedrockRuntimeClient).toHaveBeenCalledWith({
|
||||
region: 'us-east-1',
|
||||
credentials: {
|
||||
accessKeyId: 'AKIAIOSFODNN7EXAMPLE',
|
||||
secretAccessKey: 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
|
||||
},
|
||||
})
|
||||
})
|
||||
|
||||
it('creates client without credentials when neither is provided', async () => {
|
||||
await bedrockProvider.executeRequest(baseRequest)
|
||||
|
||||
expect(BedrockRuntimeClient).toHaveBeenCalledWith({
|
||||
region: 'us-east-1',
|
||||
})
|
||||
})
|
||||
|
||||
it('uses custom region when provided', async () => {
|
||||
await bedrockProvider.executeRequest({
|
||||
...baseRequest,
|
||||
bedrockRegion: 'eu-west-1',
|
||||
})
|
||||
|
||||
expect(BedrockRuntimeClient).toHaveBeenCalledWith({
|
||||
region: 'eu-west-1',
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,956 @@
|
||||
import {
|
||||
type Message as BedrockMessage,
|
||||
BedrockRuntimeClient,
|
||||
type BedrockRuntimeClientConfig,
|
||||
type ContentBlock,
|
||||
type ConversationRole,
|
||||
ConverseCommand,
|
||||
type ConverseResponse,
|
||||
ConverseStreamCommand,
|
||||
type SystemContentBlock,
|
||||
type Tool,
|
||||
type ToolConfiguration,
|
||||
type ToolResultBlock,
|
||||
type ToolUseBlock,
|
||||
} from '@aws-sdk/client-bedrock-runtime'
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { getErrorMessage, toError } from '@sim/utils/errors'
|
||||
import type { IterationToolCall, StreamingExecution } from '@/executor/types'
|
||||
import { MAX_TOOL_ITERATIONS } from '@/providers'
|
||||
import { buildBedrockMessageContent } from '@/providers/attachments'
|
||||
import {
|
||||
checkForForcedToolUsage,
|
||||
createReadableStreamFromBedrockStream,
|
||||
generateToolUseId,
|
||||
getBedrockInferenceProfileId,
|
||||
} from '@/providers/bedrock/utils'
|
||||
import { getCachedProviderClient } from '@/providers/client-cache'
|
||||
import { getProviderDefaultModel, getProviderModels } from '@/providers/models'
|
||||
import { createStreamingExecution } from '@/providers/streaming-execution'
|
||||
import { enrichLastModelSegment } from '@/providers/trace-enrichment'
|
||||
import type {
|
||||
FunctionCallResponse,
|
||||
ProviderConfig,
|
||||
ProviderRequest,
|
||||
ProviderResponse,
|
||||
TimeSegment,
|
||||
} from '@/providers/types'
|
||||
import { ProviderError } from '@/providers/types'
|
||||
import {
|
||||
calculateCost,
|
||||
prepareToolExecution,
|
||||
prepareToolsWithUsageControl,
|
||||
sumToolCosts,
|
||||
} from '@/providers/utils'
|
||||
import { executeTool } from '@/tools'
|
||||
|
||||
const logger = createLogger('BedrockProvider')
|
||||
|
||||
function enrichLastModelSegmentFromBedrockResponse(
|
||||
timeSegments: TimeSegment[],
|
||||
response: ConverseResponse,
|
||||
extras: { model: string }
|
||||
): void {
|
||||
const blocks: ContentBlock[] = response.output?.message?.content ?? []
|
||||
|
||||
const assistantText = blocks
|
||||
.filter((b): b is ContentBlock & { text: string } => 'text' in b && typeof b.text === 'string')
|
||||
.map((b) => b.text)
|
||||
.join('\n')
|
||||
const assistantContent = assistantText.length > 0 ? assistantText : undefined
|
||||
|
||||
const toolCalls: IterationToolCall[] = blocks
|
||||
.filter((b): b is ContentBlock & { toolUse: ToolUseBlock } => 'toolUse' in b && !!b.toolUse)
|
||||
.map((b) => {
|
||||
const input = b.toolUse.input
|
||||
return {
|
||||
id: b.toolUse.toolUseId ?? '',
|
||||
name: b.toolUse.name ?? '',
|
||||
arguments:
|
||||
input && typeof input === 'object' && !Array.isArray(input)
|
||||
? (input as Record<string, unknown>)
|
||||
: {},
|
||||
}
|
||||
})
|
||||
|
||||
const inputTokens = response.usage?.inputTokens
|
||||
const outputTokens = response.usage?.outputTokens
|
||||
|
||||
let cost: { input: number; output: number; total: number } | undefined
|
||||
if (typeof inputTokens === 'number' && typeof outputTokens === 'number') {
|
||||
const full = calculateCost(extras.model, inputTokens, outputTokens)
|
||||
cost = { input: full.input, output: full.output, total: full.total }
|
||||
}
|
||||
|
||||
enrichLastModelSegment(timeSegments, {
|
||||
assistantContent,
|
||||
toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
|
||||
finishReason: response.stopReason ?? undefined,
|
||||
tokens:
|
||||
inputTokens !== undefined || outputTokens !== undefined
|
||||
? {
|
||||
input: inputTokens,
|
||||
output: outputTokens,
|
||||
total:
|
||||
typeof inputTokens === 'number' && typeof outputTokens === 'number'
|
||||
? inputTokens + outputTokens
|
||||
: undefined,
|
||||
}
|
||||
: undefined,
|
||||
cost,
|
||||
provider: 'aws.bedrock',
|
||||
})
|
||||
}
|
||||
|
||||
export const bedrockProvider: ProviderConfig = {
|
||||
id: 'bedrock',
|
||||
name: 'AWS Bedrock',
|
||||
description: 'AWS Bedrock foundation models',
|
||||
version: '1.0.0',
|
||||
models: getProviderModels('bedrock'),
|
||||
defaultModel: getProviderDefaultModel('bedrock'),
|
||||
|
||||
executeRequest: async (
|
||||
request: ProviderRequest
|
||||
): Promise<ProviderResponse | StreamingExecution> => {
|
||||
const region = request.bedrockRegion || 'us-east-1'
|
||||
const bedrockModelId = getBedrockInferenceProfileId(request.model, region)
|
||||
|
||||
logger.info('Bedrock request', {
|
||||
requestModel: request.model,
|
||||
inferenceProfileId: bedrockModelId,
|
||||
region,
|
||||
})
|
||||
|
||||
const hasAccessKey = Boolean(request.bedrockAccessKeyId)
|
||||
const hasSecretKey = Boolean(request.bedrockSecretKey)
|
||||
if (hasAccessKey !== hasSecretKey) {
|
||||
throw new Error(
|
||||
'Both bedrockAccessKeyId and bedrockSecretKey must be provided together. ' +
|
||||
'Provide both for explicit credentials, or omit both to use the AWS default credential chain.'
|
||||
)
|
||||
}
|
||||
|
||||
const clientConfig: BedrockRuntimeClientConfig = { region }
|
||||
if (request.bedrockAccessKeyId && request.bedrockSecretKey) {
|
||||
clientConfig.credentials = {
|
||||
accessKeyId: request.bedrockAccessKeyId,
|
||||
secretAccessKey: request.bedrockSecretKey,
|
||||
}
|
||||
}
|
||||
|
||||
// Key on the full credential (access key id + secret) so a corrected secret
|
||||
// under the same access key id yields a fresh client rather than a stale one.
|
||||
const credentialKey =
|
||||
request.bedrockAccessKeyId && request.bedrockSecretKey
|
||||
? `${request.bedrockAccessKeyId}:${request.bedrockSecretKey}`
|
||||
: 'default-chain'
|
||||
const client = getCachedProviderClient(
|
||||
`bedrock::${region}::${credentialKey}`,
|
||||
() => new BedrockRuntimeClient(clientConfig)
|
||||
)
|
||||
|
||||
const messages: BedrockMessage[] = []
|
||||
const systemContent: SystemContentBlock[] = []
|
||||
|
||||
if (request.systemPrompt) {
|
||||
systemContent.push({ text: request.systemPrompt })
|
||||
}
|
||||
|
||||
if (request.context) {
|
||||
messages.push({
|
||||
role: 'user' as ConversationRole,
|
||||
content: [{ text: request.context }],
|
||||
})
|
||||
}
|
||||
|
||||
if (request.messages) {
|
||||
for (const msg of request.messages) {
|
||||
if (msg.role === 'function' || msg.role === 'tool') {
|
||||
const toolResultBlock: ToolResultBlock = {
|
||||
toolUseId: msg.tool_call_id || msg.name || generateToolUseId('tool'),
|
||||
content: [{ text: msg.content || '' }],
|
||||
}
|
||||
messages.push({
|
||||
role: 'user' as ConversationRole,
|
||||
content: [{ toolResult: toolResultBlock }],
|
||||
})
|
||||
} else if (msg.function_call || msg.tool_calls) {
|
||||
const toolCall = msg.function_call || msg.tool_calls?.[0]?.function
|
||||
if (toolCall) {
|
||||
const toolUseBlock: ToolUseBlock = {
|
||||
toolUseId: msg.tool_calls?.[0]?.id || generateToolUseId(toolCall.name),
|
||||
name: toolCall.name,
|
||||
input: JSON.parse(toolCall.arguments),
|
||||
}
|
||||
messages.push({
|
||||
role: 'assistant' as ConversationRole,
|
||||
content: [{ toolUse: toolUseBlock }],
|
||||
})
|
||||
}
|
||||
} else {
|
||||
const role: ConversationRole = msg.role === 'assistant' ? 'assistant' : 'user'
|
||||
const content = buildBedrockMessageContent(msg.content, msg.files, 'bedrock')
|
||||
messages.push({
|
||||
role,
|
||||
// double-cast-allowed: shared attachment builder emits Bedrock Converse content blocks while keeping provider-neutral attachment types
|
||||
content: content as unknown as ContentBlock[],
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (messages.length === 0) {
|
||||
messages.push({
|
||||
role: 'user' as ConversationRole,
|
||||
content: [{ text: request.systemPrompt || 'Hello' }],
|
||||
})
|
||||
systemContent.length = 0
|
||||
}
|
||||
|
||||
let structuredOutputTool: Tool | undefined
|
||||
const structuredOutputToolName = 'structured_output'
|
||||
|
||||
if (request.responseFormat) {
|
||||
const schema = request.responseFormat.schema || request.responseFormat
|
||||
const schemaName = request.responseFormat.name || 'response'
|
||||
|
||||
structuredOutputTool = {
|
||||
toolSpec: {
|
||||
name: structuredOutputToolName,
|
||||
description: `Output the response as structured JSON matching the ${schemaName} schema. You MUST call this tool to provide your final response.`,
|
||||
inputSchema: {
|
||||
json: schema,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
logger.info(`Using Tool Use approach for structured outputs: ${schemaName}`)
|
||||
}
|
||||
|
||||
let bedrockTools: Tool[] | undefined
|
||||
let toolChoice: any = { auto: {} }
|
||||
let preparedTools: ReturnType<typeof prepareToolsWithUsageControl> | null = null
|
||||
|
||||
if (request.tools?.length) {
|
||||
bedrockTools = request.tools.map((tool) => ({
|
||||
toolSpec: {
|
||||
name: tool.id,
|
||||
description: tool.description,
|
||||
inputSchema: {
|
||||
json: {
|
||||
type: 'object',
|
||||
properties: tool.parameters.properties,
|
||||
required: tool.parameters.required,
|
||||
},
|
||||
},
|
||||
},
|
||||
}))
|
||||
|
||||
try {
|
||||
preparedTools = prepareToolsWithUsageControl(
|
||||
bedrockTools.map((t) => ({
|
||||
name: t.toolSpec?.name || '',
|
||||
description: t.toolSpec?.description || '',
|
||||
input_schema: t.toolSpec?.inputSchema?.json,
|
||||
})),
|
||||
request.tools,
|
||||
logger,
|
||||
'bedrock'
|
||||
)
|
||||
|
||||
const { tools: filteredTools, toolChoice: tc } = preparedTools
|
||||
|
||||
if (filteredTools?.length) {
|
||||
bedrockTools = filteredTools.map((t: any) => ({
|
||||
toolSpec: {
|
||||
name: t.name,
|
||||
description: t.description,
|
||||
inputSchema: { json: t.input_schema },
|
||||
},
|
||||
}))
|
||||
|
||||
if (typeof tc === 'object' && tc !== null) {
|
||||
if (tc.type === 'tool' && tc.name) {
|
||||
toolChoice = { tool: { name: tc.name } }
|
||||
logger.info(`Using Bedrock tool_choice format: force tool "${tc.name}"`)
|
||||
} else if (tc.type === 'function' && tc.function?.name) {
|
||||
toolChoice = { tool: { name: tc.function.name } }
|
||||
logger.info(`Using Bedrock tool_choice format: force tool "${tc.function.name}"`)
|
||||
} else if (tc.type === 'any') {
|
||||
toolChoice = { any: {} }
|
||||
logger.info('Using Bedrock tool_choice format: any tool')
|
||||
} else {
|
||||
toolChoice = { auto: {} }
|
||||
}
|
||||
} else if (tc === 'none') {
|
||||
toolChoice = undefined
|
||||
bedrockTools = undefined
|
||||
} else {
|
||||
toolChoice = { auto: {} }
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error('Error in prepareToolsWithUsageControl:', { error })
|
||||
toolChoice = { auto: {} }
|
||||
}
|
||||
} else if (structuredOutputTool) {
|
||||
bedrockTools = [structuredOutputTool]
|
||||
toolChoice = { tool: { name: structuredOutputToolName } }
|
||||
logger.info('Using structured_output tool as only tool (forced)')
|
||||
}
|
||||
|
||||
const hasToolContentInMessages = messages.some((msg) =>
|
||||
msg.content?.some(
|
||||
(block) =>
|
||||
('toolUse' in block && block.toolUse) || ('toolResult' in block && block.toolResult)
|
||||
)
|
||||
)
|
||||
|
||||
const toolConfig: ToolConfiguration | undefined = bedrockTools?.length
|
||||
? {
|
||||
tools: bedrockTools,
|
||||
toolChoice,
|
||||
}
|
||||
: hasToolContentInMessages && request.tools?.length
|
||||
? {
|
||||
tools: request.tools.map((tool) => ({
|
||||
toolSpec: {
|
||||
name: tool.id,
|
||||
description: tool.description,
|
||||
inputSchema: {
|
||||
json: {
|
||||
type: 'object',
|
||||
properties: tool.parameters.properties,
|
||||
required: tool.parameters.required,
|
||||
},
|
||||
},
|
||||
},
|
||||
})),
|
||||
toolChoice: { auto: {} },
|
||||
}
|
||||
: undefined
|
||||
|
||||
if (hasToolContentInMessages && !toolConfig) {
|
||||
throw new Error(
|
||||
'Messages contain tool use/result blocks but no tools were provided. ' +
|
||||
'Bedrock requires toolConfig when processing messages with tool content.'
|
||||
)
|
||||
}
|
||||
|
||||
const systemPromptWithSchema = systemContent
|
||||
|
||||
const inferenceConfig: { temperature: number; maxTokens?: number } = {
|
||||
temperature: Number.parseFloat(String(request.temperature ?? 0.7)),
|
||||
}
|
||||
if (request.maxTokens != null) {
|
||||
inferenceConfig.maxTokens = Number.parseInt(String(request.maxTokens))
|
||||
}
|
||||
|
||||
const shouldStreamToolCalls = request.streamToolCalls ?? false
|
||||
|
||||
if (request.stream && (!bedrockTools || bedrockTools.length === 0)) {
|
||||
logger.info('Using streaming response for Bedrock request (no tools)')
|
||||
|
||||
const providerStartTime = Date.now()
|
||||
const providerStartTimeISO = new Date(providerStartTime).toISOString()
|
||||
|
||||
const command = new ConverseStreamCommand({
|
||||
modelId: bedrockModelId,
|
||||
messages,
|
||||
system: systemPromptWithSchema.length > 0 ? systemPromptWithSchema : undefined,
|
||||
inferenceConfig,
|
||||
})
|
||||
|
||||
const streamResponse = await client.send(
|
||||
command,
|
||||
request.abortSignal ? { abortSignal: request.abortSignal } : undefined
|
||||
)
|
||||
|
||||
if (!streamResponse.stream) {
|
||||
throw new Error('No stream returned from Bedrock')
|
||||
}
|
||||
|
||||
const bedrockStream = streamResponse.stream
|
||||
const streamingResult = createStreamingExecution({
|
||||
model: request.model,
|
||||
providerStartTime,
|
||||
providerStartTimeISO,
|
||||
timing: { kind: 'simple', segmentName: request.model },
|
||||
initialTokens: { input: 0, output: 0, total: 0 },
|
||||
initialCost: { total: 0.0, input: 0.0, output: 0.0 },
|
||||
isStreaming: true,
|
||||
createStream: ({ output, finalizeTiming }) =>
|
||||
createReadableStreamFromBedrockStream(bedrockStream, (content, usage) => {
|
||||
output.content = content
|
||||
output.tokens = {
|
||||
input: usage.inputTokens,
|
||||
output: usage.outputTokens,
|
||||
total: usage.inputTokens + usage.outputTokens,
|
||||
}
|
||||
|
||||
const costResult = calculateCost(request.model, usage.inputTokens, usage.outputTokens)
|
||||
output.cost = {
|
||||
input: costResult.input,
|
||||
output: costResult.output,
|
||||
total: costResult.total,
|
||||
}
|
||||
|
||||
finalizeTiming()
|
||||
}),
|
||||
})
|
||||
|
||||
return streamingResult
|
||||
}
|
||||
|
||||
const providerStartTime = Date.now()
|
||||
const providerStartTimeISO = new Date(providerStartTime).toISOString()
|
||||
|
||||
try {
|
||||
const initialCallTime = Date.now()
|
||||
const originalToolChoice = toolChoice
|
||||
const forcedTools = preparedTools?.forcedTools || []
|
||||
let usedForcedTools: string[] = []
|
||||
|
||||
const command = new ConverseCommand({
|
||||
modelId: bedrockModelId,
|
||||
messages,
|
||||
system: systemPromptWithSchema.length > 0 ? systemPromptWithSchema : undefined,
|
||||
inferenceConfig,
|
||||
toolConfig,
|
||||
})
|
||||
|
||||
let currentResponse = await client.send(
|
||||
command,
|
||||
request.abortSignal ? { abortSignal: request.abortSignal } : undefined
|
||||
)
|
||||
const firstResponseTime = Date.now() - initialCallTime
|
||||
|
||||
let content = ''
|
||||
let hasExtractedStructuredOutput = false
|
||||
if (currentResponse.output?.message?.content) {
|
||||
const structuredOutputCall = currentResponse.output.message.content.find(
|
||||
(block): block is ContentBlock & { toolUse: ToolUseBlock } =>
|
||||
'toolUse' in block && block.toolUse?.name === structuredOutputToolName
|
||||
)
|
||||
|
||||
if (structuredOutputCall && structuredOutputTool) {
|
||||
content = JSON.stringify(structuredOutputCall.toolUse.input, null, 2)
|
||||
hasExtractedStructuredOutput = true
|
||||
logger.info('Extracted structured output from tool call')
|
||||
} else {
|
||||
const textBlocks = currentResponse.output.message.content.filter(
|
||||
(block): block is ContentBlock & { text: string } => 'text' in block
|
||||
)
|
||||
content = textBlocks.map((block) => block.text).join('\n')
|
||||
}
|
||||
}
|
||||
|
||||
const tokens = {
|
||||
input: currentResponse.usage?.inputTokens || 0,
|
||||
output: currentResponse.usage?.outputTokens || 0,
|
||||
total:
|
||||
(currentResponse.usage?.inputTokens || 0) + (currentResponse.usage?.outputTokens || 0),
|
||||
}
|
||||
|
||||
const initialCost = calculateCost(
|
||||
request.model,
|
||||
currentResponse.usage?.inputTokens || 0,
|
||||
currentResponse.usage?.outputTokens || 0
|
||||
)
|
||||
const cost = {
|
||||
input: initialCost.input,
|
||||
output: initialCost.output,
|
||||
total: initialCost.total,
|
||||
pricing: initialCost.pricing,
|
||||
}
|
||||
|
||||
const toolCalls: FunctionCallResponse[] = []
|
||||
const toolResults: Record<string, unknown>[] = []
|
||||
const currentMessages = [...messages]
|
||||
let iterationCount = 0
|
||||
let hasUsedForcedTool = false
|
||||
let modelTime = firstResponseTime
|
||||
let toolsTime = 0
|
||||
|
||||
const timeSegments: TimeSegment[] = [
|
||||
{
|
||||
type: 'model',
|
||||
name: request.model,
|
||||
startTime: initialCallTime,
|
||||
endTime: initialCallTime + firstResponseTime,
|
||||
duration: firstResponseTime,
|
||||
},
|
||||
]
|
||||
|
||||
enrichLastModelSegmentFromBedrockResponse(timeSegments, currentResponse, {
|
||||
model: request.model,
|
||||
})
|
||||
|
||||
const initialToolUseContentBlocks = (currentResponse.output?.message?.content || []).filter(
|
||||
(block): block is ContentBlock & { toolUse: ToolUseBlock } => 'toolUse' in block
|
||||
)
|
||||
const toolUseBlocks = initialToolUseContentBlocks.map((block) => ({
|
||||
name: block.toolUse.name || '',
|
||||
}))
|
||||
|
||||
const firstCheckResult = checkForForcedToolUsage(
|
||||
toolUseBlocks,
|
||||
originalToolChoice,
|
||||
forcedTools,
|
||||
usedForcedTools
|
||||
)
|
||||
if (firstCheckResult) {
|
||||
hasUsedForcedTool = firstCheckResult.hasUsedForcedTool
|
||||
usedForcedTools = firstCheckResult.usedForcedTools
|
||||
}
|
||||
|
||||
while (iterationCount < MAX_TOOL_ITERATIONS) {
|
||||
const textContentBlocks = (currentResponse.output?.message?.content || []).filter(
|
||||
(block): block is ContentBlock & { text: string } => 'text' in block
|
||||
)
|
||||
const textContent = textContentBlocks.map((block) => block.text).join('\n')
|
||||
|
||||
if (textContent) {
|
||||
content = textContent
|
||||
}
|
||||
|
||||
const toolUseContentBlocks = (currentResponse.output?.message?.content || []).filter(
|
||||
(block): block is ContentBlock & { toolUse: ToolUseBlock } => 'toolUse' in block
|
||||
)
|
||||
const currentToolUses = toolUseContentBlocks.map((block) => block.toolUse)
|
||||
|
||||
if (!currentToolUses || currentToolUses.length === 0) {
|
||||
break
|
||||
}
|
||||
|
||||
const toolsStartTime = Date.now()
|
||||
|
||||
const toolExecutionPromises = currentToolUses.map(async (toolUse: ToolUseBlock) => {
|
||||
const toolCallStartTime = Date.now()
|
||||
const toolName = toolUse.name || ''
|
||||
const toolArgs = (toolUse.input as Record<string, any>) || {}
|
||||
const toolUseId = toolUse.toolUseId || generateToolUseId(toolName)
|
||||
|
||||
try {
|
||||
const tool = request.tools?.find((t) => t.id === toolName)
|
||||
if (!tool) return null
|
||||
|
||||
const { toolParams, executionParams } = prepareToolExecution(tool, toolArgs, request)
|
||||
const result = await executeTool(toolName, executionParams, {
|
||||
signal: request.abortSignal,
|
||||
})
|
||||
const toolCallEndTime = Date.now()
|
||||
|
||||
return {
|
||||
toolUseId,
|
||||
toolName,
|
||||
toolArgs,
|
||||
toolParams,
|
||||
result,
|
||||
startTime: toolCallStartTime,
|
||||
endTime: toolCallEndTime,
|
||||
duration: toolCallEndTime - toolCallStartTime,
|
||||
}
|
||||
} catch (error) {
|
||||
const toolCallEndTime = Date.now()
|
||||
logger.error('Error processing tool call:', { error, toolName })
|
||||
|
||||
return {
|
||||
toolUseId,
|
||||
toolName,
|
||||
toolArgs,
|
||||
toolParams: {},
|
||||
result: {
|
||||
success: false,
|
||||
output: undefined,
|
||||
error: getErrorMessage(error, 'Tool execution failed'),
|
||||
},
|
||||
startTime: toolCallStartTime,
|
||||
endTime: toolCallEndTime,
|
||||
duration: toolCallEndTime - toolCallStartTime,
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
const executionResults = await Promise.allSettled(toolExecutionPromises)
|
||||
|
||||
const assistantContent: ContentBlock[] = currentToolUses.map((toolUse: ToolUseBlock) => ({
|
||||
toolUse: {
|
||||
toolUseId: toolUse.toolUseId,
|
||||
name: toolUse.name,
|
||||
input: toolUse.input,
|
||||
},
|
||||
}))
|
||||
currentMessages.push({
|
||||
role: 'assistant' as ConversationRole,
|
||||
content: assistantContent,
|
||||
})
|
||||
|
||||
const toolResultContent: ContentBlock[] = []
|
||||
|
||||
for (const settledResult of executionResults) {
|
||||
if (settledResult.status === 'rejected' || !settledResult.value) continue
|
||||
|
||||
const {
|
||||
toolUseId,
|
||||
toolName,
|
||||
toolArgs,
|
||||
toolParams,
|
||||
result,
|
||||
startTime,
|
||||
endTime,
|
||||
duration,
|
||||
} = settledResult.value
|
||||
|
||||
timeSegments.push({
|
||||
type: 'tool',
|
||||
name: toolName,
|
||||
startTime,
|
||||
endTime,
|
||||
duration,
|
||||
})
|
||||
|
||||
let resultContent: any
|
||||
if (result.success) {
|
||||
toolResults.push(result.output!)
|
||||
resultContent = result.output
|
||||
} else {
|
||||
resultContent = {
|
||||
error: true,
|
||||
message: result.error || 'Tool execution failed',
|
||||
tool: toolName,
|
||||
}
|
||||
}
|
||||
|
||||
toolCalls.push({
|
||||
name: toolName,
|
||||
arguments: toolParams,
|
||||
startTime: new Date(startTime).toISOString(),
|
||||
endTime: new Date(endTime).toISOString(),
|
||||
duration,
|
||||
result: resultContent,
|
||||
success: result.success,
|
||||
})
|
||||
|
||||
const toolResultBlock: ToolResultBlock = {
|
||||
toolUseId,
|
||||
content: [{ text: JSON.stringify(resultContent) }],
|
||||
}
|
||||
toolResultContent.push({ toolResult: toolResultBlock })
|
||||
}
|
||||
|
||||
if (toolResultContent.length > 0) {
|
||||
currentMessages.push({
|
||||
role: 'user' as ConversationRole,
|
||||
content: toolResultContent,
|
||||
})
|
||||
}
|
||||
|
||||
const thisToolsTime = Date.now() - toolsStartTime
|
||||
toolsTime += thisToolsTime
|
||||
|
||||
let nextToolChoice = toolChoice
|
||||
if (typeof originalToolChoice === 'object' && hasUsedForcedTool && forcedTools.length > 0) {
|
||||
const remainingTools = forcedTools.filter((tool) => !usedForcedTools.includes(tool))
|
||||
|
||||
if (remainingTools.length > 0) {
|
||||
nextToolChoice = { tool: { name: remainingTools[0] } }
|
||||
logger.info(`Forcing next tool: ${remainingTools[0]}`)
|
||||
} else {
|
||||
nextToolChoice = { auto: {} }
|
||||
logger.info('All forced tools have been used, switching to auto')
|
||||
}
|
||||
} else if (hasUsedForcedTool && typeof originalToolChoice === 'object') {
|
||||
nextToolChoice = { auto: {} }
|
||||
logger.info('Switching to auto tool choice after forced tool was used')
|
||||
}
|
||||
|
||||
const nextModelStartTime = Date.now()
|
||||
|
||||
const nextCommand = new ConverseCommand({
|
||||
modelId: bedrockModelId,
|
||||
messages: currentMessages,
|
||||
system: systemPromptWithSchema.length > 0 ? systemPromptWithSchema : undefined,
|
||||
inferenceConfig,
|
||||
toolConfig: bedrockTools?.length
|
||||
? { tools: bedrockTools, toolChoice: nextToolChoice }
|
||||
: undefined,
|
||||
})
|
||||
|
||||
currentResponse = await client.send(
|
||||
nextCommand,
|
||||
request.abortSignal ? { abortSignal: request.abortSignal } : undefined
|
||||
)
|
||||
|
||||
const nextToolUseContentBlocks = (currentResponse.output?.message?.content || []).filter(
|
||||
(block): block is ContentBlock & { toolUse: ToolUseBlock } => 'toolUse' in block
|
||||
)
|
||||
const nextToolUseBlocks = nextToolUseContentBlocks.map((block) => ({
|
||||
name: block.toolUse.name || '',
|
||||
}))
|
||||
|
||||
const nextCheckResult = checkForForcedToolUsage(
|
||||
nextToolUseBlocks,
|
||||
nextToolChoice,
|
||||
forcedTools,
|
||||
usedForcedTools
|
||||
)
|
||||
if (nextCheckResult) {
|
||||
hasUsedForcedTool = nextCheckResult.hasUsedForcedTool
|
||||
usedForcedTools = nextCheckResult.usedForcedTools
|
||||
}
|
||||
|
||||
const nextModelEndTime = Date.now()
|
||||
const thisModelTime = nextModelEndTime - nextModelStartTime
|
||||
|
||||
timeSegments.push({
|
||||
type: 'model',
|
||||
name: request.model,
|
||||
startTime: nextModelStartTime,
|
||||
endTime: nextModelEndTime,
|
||||
duration: thisModelTime,
|
||||
})
|
||||
|
||||
enrichLastModelSegmentFromBedrockResponse(timeSegments, currentResponse, {
|
||||
model: request.model,
|
||||
})
|
||||
|
||||
modelTime += thisModelTime
|
||||
|
||||
if (currentResponse.usage) {
|
||||
tokens.input += currentResponse.usage.inputTokens || 0
|
||||
tokens.output += currentResponse.usage.outputTokens || 0
|
||||
tokens.total +=
|
||||
(currentResponse.usage.inputTokens || 0) + (currentResponse.usage.outputTokens || 0)
|
||||
|
||||
const iterationCost = calculateCost(
|
||||
request.model,
|
||||
currentResponse.usage.inputTokens || 0,
|
||||
currentResponse.usage.outputTokens || 0
|
||||
)
|
||||
cost.input += iterationCost.input
|
||||
cost.output += iterationCost.output
|
||||
cost.total += iterationCost.total
|
||||
}
|
||||
|
||||
iterationCount++
|
||||
}
|
||||
|
||||
if (structuredOutputTool && request.tools?.length) {
|
||||
logger.info('Making final call with forced structured_output tool')
|
||||
|
||||
const structuredOutputStartTime = Date.now()
|
||||
|
||||
const structuredOutputCommand = new ConverseCommand({
|
||||
modelId: bedrockModelId,
|
||||
messages: currentMessages,
|
||||
system: systemPromptWithSchema.length > 0 ? systemPromptWithSchema : undefined,
|
||||
inferenceConfig,
|
||||
toolConfig: {
|
||||
tools: [structuredOutputTool],
|
||||
toolChoice: { tool: { name: structuredOutputToolName } },
|
||||
},
|
||||
})
|
||||
|
||||
const structuredResponse = await client.send(
|
||||
structuredOutputCommand,
|
||||
request.abortSignal ? { abortSignal: request.abortSignal } : undefined
|
||||
)
|
||||
const structuredOutputEndTime = Date.now()
|
||||
|
||||
timeSegments.push({
|
||||
type: 'model',
|
||||
name: 'Structured output extraction',
|
||||
startTime: structuredOutputStartTime,
|
||||
endTime: structuredOutputEndTime,
|
||||
duration: structuredOutputEndTime - structuredOutputStartTime,
|
||||
})
|
||||
|
||||
enrichLastModelSegmentFromBedrockResponse(timeSegments, structuredResponse, {
|
||||
model: request.model,
|
||||
})
|
||||
|
||||
modelTime += structuredOutputEndTime - structuredOutputStartTime
|
||||
|
||||
const structuredOutputCall = structuredResponse.output?.message?.content?.find(
|
||||
(block): block is ContentBlock & { toolUse: ToolUseBlock } =>
|
||||
'toolUse' in block && block.toolUse?.name === structuredOutputToolName
|
||||
)
|
||||
|
||||
if (structuredOutputCall) {
|
||||
content = JSON.stringify(structuredOutputCall.toolUse.input, null, 2)
|
||||
hasExtractedStructuredOutput = true
|
||||
logger.info('Extracted structured output from forced tool call')
|
||||
} else {
|
||||
logger.warn('Structured output tool was forced but no tool call found in response')
|
||||
}
|
||||
|
||||
if (structuredResponse.usage) {
|
||||
tokens.input += structuredResponse.usage.inputTokens || 0
|
||||
tokens.output += structuredResponse.usage.outputTokens || 0
|
||||
tokens.total +=
|
||||
(structuredResponse.usage.inputTokens || 0) +
|
||||
(structuredResponse.usage.outputTokens || 0)
|
||||
|
||||
const structuredCost = calculateCost(
|
||||
request.model,
|
||||
structuredResponse.usage.inputTokens || 0,
|
||||
structuredResponse.usage.outputTokens || 0
|
||||
)
|
||||
cost.input += structuredCost.input
|
||||
cost.output += structuredCost.output
|
||||
cost.total += structuredCost.total
|
||||
}
|
||||
}
|
||||
|
||||
const providerEndTime = Date.now()
|
||||
const providerEndTimeISO = new Date(providerEndTime).toISOString()
|
||||
const totalDuration = providerEndTime - providerStartTime
|
||||
|
||||
if (request.stream && !shouldStreamToolCalls && !hasExtractedStructuredOutput) {
|
||||
logger.info('Using streaming for final Bedrock response after tool processing')
|
||||
|
||||
const messagesHaveToolContent = currentMessages.some((msg) =>
|
||||
msg.content?.some(
|
||||
(block) =>
|
||||
('toolUse' in block && block.toolUse) || ('toolResult' in block && block.toolResult)
|
||||
)
|
||||
)
|
||||
|
||||
const streamToolConfig: ToolConfiguration | undefined =
|
||||
messagesHaveToolContent && request.tools?.length
|
||||
? {
|
||||
tools: request.tools.map((tool) => ({
|
||||
toolSpec: {
|
||||
name: tool.id,
|
||||
description: tool.description,
|
||||
inputSchema: {
|
||||
json: {
|
||||
type: 'object',
|
||||
properties: tool.parameters.properties,
|
||||
required: tool.parameters.required,
|
||||
},
|
||||
},
|
||||
},
|
||||
})),
|
||||
toolChoice: { auto: {} },
|
||||
}
|
||||
: undefined
|
||||
|
||||
const streamCommand = new ConverseStreamCommand({
|
||||
modelId: bedrockModelId,
|
||||
messages: currentMessages,
|
||||
system: systemPromptWithSchema.length > 0 ? systemPromptWithSchema : undefined,
|
||||
inferenceConfig,
|
||||
toolConfig: streamToolConfig,
|
||||
})
|
||||
|
||||
const streamResponse = await client.send(
|
||||
streamCommand,
|
||||
request.abortSignal ? { abortSignal: request.abortSignal } : undefined
|
||||
)
|
||||
|
||||
if (!streamResponse.stream) {
|
||||
throw new Error('No stream returned from Bedrock')
|
||||
}
|
||||
|
||||
const bedrockStream = streamResponse.stream
|
||||
const streamingResult = createStreamingExecution({
|
||||
model: request.model,
|
||||
providerStartTime,
|
||||
providerStartTimeISO,
|
||||
timing: {
|
||||
kind: 'accumulated',
|
||||
modelTime,
|
||||
toolsTime,
|
||||
firstResponseTime,
|
||||
iterations: iterationCount + 1,
|
||||
timeSegments,
|
||||
},
|
||||
initialTokens: { input: tokens.input, output: tokens.output, total: tokens.total },
|
||||
initialCost: {
|
||||
input: cost.input,
|
||||
output: cost.output,
|
||||
toolCost: undefined as number | undefined,
|
||||
total: cost.total,
|
||||
},
|
||||
toolCalls:
|
||||
toolCalls.length > 0 ? { list: toolCalls, count: toolCalls.length } : undefined,
|
||||
isStreaming: true,
|
||||
createStream: ({ output, finalizeTiming }) =>
|
||||
createReadableStreamFromBedrockStream(bedrockStream, (streamContent, usage) => {
|
||||
output.content = streamContent
|
||||
output.tokens = {
|
||||
input: tokens.input + usage.inputTokens,
|
||||
output: tokens.output + usage.outputTokens,
|
||||
total: tokens.total + usage.inputTokens + usage.outputTokens,
|
||||
}
|
||||
|
||||
const streamCost = calculateCost(request.model, usage.inputTokens, usage.outputTokens)
|
||||
const tc = sumToolCosts(toolResults)
|
||||
output.cost = {
|
||||
input: cost.input + streamCost.input,
|
||||
output: cost.output + streamCost.output,
|
||||
toolCost: tc || undefined,
|
||||
total: cost.total + streamCost.total + tc,
|
||||
}
|
||||
|
||||
finalizeTiming()
|
||||
}),
|
||||
})
|
||||
|
||||
return streamingResult
|
||||
}
|
||||
|
||||
return {
|
||||
content,
|
||||
model: request.model,
|
||||
tokens,
|
||||
cost: {
|
||||
input: cost.input,
|
||||
output: cost.output,
|
||||
total: cost.total,
|
||||
pricing: cost.pricing,
|
||||
},
|
||||
toolCalls:
|
||||
toolCalls.length > 0
|
||||
? toolCalls.map((tc) => ({
|
||||
name: tc.name,
|
||||
arguments: tc.arguments as Record<string, any>,
|
||||
startTime: tc.startTime,
|
||||
endTime: tc.endTime,
|
||||
duration: tc.duration,
|
||||
result: tc.result,
|
||||
}))
|
||||
: undefined,
|
||||
toolResults: toolResults.length > 0 ? toolResults : undefined,
|
||||
timing: {
|
||||
startTime: providerStartTimeISO,
|
||||
endTime: providerEndTimeISO,
|
||||
duration: totalDuration,
|
||||
modelTime,
|
||||
toolsTime,
|
||||
firstResponseTime,
|
||||
iterations: iterationCount + 1,
|
||||
timeSegments,
|
||||
},
|
||||
}
|
||||
} catch (error) {
|
||||
const providerEndTime = Date.now()
|
||||
const providerEndTimeISO = new Date(providerEndTime).toISOString()
|
||||
const totalDuration = providerEndTime - providerStartTime
|
||||
|
||||
logger.error('Error in Bedrock request:', {
|
||||
error,
|
||||
duration: totalDuration,
|
||||
})
|
||||
|
||||
throw new ProviderError(toError(error).message, {
|
||||
startTime: providerStartTimeISO,
|
||||
endTime: providerEndTimeISO,
|
||||
duration: totalDuration,
|
||||
})
|
||||
}
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
/**
|
||||
* @vitest-environment node
|
||||
*/
|
||||
import { describe, expect, it } from 'vitest'
|
||||
import { getBedrockInferenceProfileId } from '@/providers/bedrock/utils'
|
||||
|
||||
describe('getBedrockInferenceProfileId', () => {
|
||||
it.concurrent('prefixes geo inference profile for models that require it', () => {
|
||||
expect(
|
||||
getBedrockInferenceProfileId('bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0', 'us-east-1')
|
||||
).toBe('us.anthropic.claude-sonnet-4-5-20250929-v1:0')
|
||||
expect(getBedrockInferenceProfileId('bedrock/amazon.nova-pro-v1:0', 'eu-west-1')).toBe(
|
||||
'eu.amazon.nova-pro-v1:0'
|
||||
)
|
||||
expect(
|
||||
getBedrockInferenceProfileId('bedrock/meta.llama4-scout-17b-instruct-v1:0', 'us-west-2')
|
||||
).toBe('us.meta.llama4-scout-17b-instruct-v1:0')
|
||||
})
|
||||
|
||||
it.concurrent('returns already-prefixed inference profile IDs unchanged', () => {
|
||||
expect(
|
||||
getBedrockInferenceProfileId('us.anthropic.claude-sonnet-4-5-20250929-v1:0', 'us-east-1')
|
||||
).toBe('us.anthropic.claude-sonnet-4-5-20250929-v1:0')
|
||||
expect(getBedrockInferenceProfileId('global.amazon.nova-2-lite-v1:0', 'us-east-1')).toBe(
|
||||
'global.amazon.nova-2-lite-v1:0'
|
||||
)
|
||||
})
|
||||
|
||||
it.concurrent('returns the bare model ID for models without geo profile support', () => {
|
||||
expect(
|
||||
getBedrockInferenceProfileId('bedrock/mistral.mistral-large-3-675b-instruct', 'us-east-1')
|
||||
).toBe('mistral.mistral-large-3-675b-instruct')
|
||||
expect(
|
||||
getBedrockInferenceProfileId('bedrock/mistral.ministral-3-8b-instruct', 'eu-west-1')
|
||||
).toBe('mistral.ministral-3-8b-instruct')
|
||||
expect(getBedrockInferenceProfileId('bedrock/cohere.command-r-plus-v1:0', 'us-east-1')).toBe(
|
||||
'cohere.command-r-plus-v1:0'
|
||||
)
|
||||
expect(
|
||||
getBedrockInferenceProfileId('bedrock/mistral.mixtral-8x7b-instruct-v0:1', 'ap-southeast-1')
|
||||
).toBe('mistral.mixtral-8x7b-instruct-v0:1')
|
||||
expect(
|
||||
getBedrockInferenceProfileId('bedrock/amazon.titan-text-premier-v1:0', 'us-east-1')
|
||||
).toBe('amazon.titan-text-premier-v1:0')
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,143 @@
|
||||
import type { ConverseStreamOutput } from '@aws-sdk/client-bedrock-runtime'
|
||||
import { createLogger } from '@sim/logger'
|
||||
import { randomFloat } from '@sim/utils/random'
|
||||
import { trackForcedToolUsage } from '@/providers/utils'
|
||||
|
||||
const logger = createLogger('BedrockUtils')
|
||||
|
||||
export interface BedrockStreamUsage {
|
||||
inputTokens: number
|
||||
outputTokens: number
|
||||
}
|
||||
|
||||
export function createReadableStreamFromBedrockStream(
|
||||
bedrockStream: AsyncIterable<ConverseStreamOutput>,
|
||||
onComplete?: (content: string, usage: BedrockStreamUsage) => void
|
||||
): ReadableStream<Uint8Array> {
|
||||
let fullContent = ''
|
||||
let inputTokens = 0
|
||||
let outputTokens = 0
|
||||
|
||||
return new ReadableStream({
|
||||
async start(controller) {
|
||||
try {
|
||||
for await (const event of bedrockStream) {
|
||||
if (event.contentBlockDelta?.delta?.text) {
|
||||
const text = event.contentBlockDelta.delta.text
|
||||
fullContent += text
|
||||
controller.enqueue(new TextEncoder().encode(text))
|
||||
} else if (event.metadata?.usage) {
|
||||
inputTokens = event.metadata.usage.inputTokens ?? 0
|
||||
outputTokens = event.metadata.usage.outputTokens ?? 0
|
||||
}
|
||||
}
|
||||
|
||||
if (onComplete) {
|
||||
onComplete(fullContent, { inputTokens, outputTokens })
|
||||
}
|
||||
|
||||
controller.close()
|
||||
} catch (err) {
|
||||
controller.error(err)
|
||||
}
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
export function checkForForcedToolUsage(
|
||||
toolUseBlocks: Array<{ name: string }>,
|
||||
toolChoice: any,
|
||||
forcedTools: string[],
|
||||
usedForcedTools: string[]
|
||||
): { hasUsedForcedTool: boolean; usedForcedTools: string[] } | null {
|
||||
if (typeof toolChoice === 'object' && toolChoice !== null && toolUseBlocks.length > 0) {
|
||||
const adaptedToolCalls = toolUseBlocks.map((tool) => ({ name: tool.name }))
|
||||
const adaptedToolChoice = toolChoice.tool
|
||||
? { function: { name: toolChoice.tool.name } }
|
||||
: toolChoice
|
||||
|
||||
return trackForcedToolUsage(
|
||||
adaptedToolCalls,
|
||||
adaptedToolChoice,
|
||||
logger,
|
||||
'bedrock',
|
||||
forcedTools,
|
||||
usedForcedTools
|
||||
)
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a unique tool use ID for Bedrock.
|
||||
* AWS Bedrock requires toolUseId to be 1-64 characters, pattern [a-zA-Z0-9_-]+
|
||||
*/
|
||||
export function generateToolUseId(toolName: string): string {
|
||||
const timestamp = Date.now().toString(36) // Base36 timestamp (9 chars)
|
||||
const random = randomFloat().toString(36).substring(2, 7) // 5 random chars
|
||||
const suffix = `-${timestamp}-${random}` // ~15 chars
|
||||
const maxNameLength = 64 - suffix.length
|
||||
const truncatedName = toolName.substring(0, maxNameLength).replace(/[^a-zA-Z0-9_-]/g, '_')
|
||||
return `${truncatedName}${suffix}`
|
||||
}
|
||||
|
||||
/**
|
||||
* Models whose AWS model cards state geo/cross-region inference profiles are
|
||||
* not supported ("Geo inference ID: Not supported"). These must be invoked
|
||||
* with the bare in-region model ID — prefixing them with a geo profile
|
||||
* (e.g. us.mistral...) produces an invalid model identifier.
|
||||
*/
|
||||
const GEO_PROFILE_UNSUPPORTED_MODEL_IDS = new Set([
|
||||
'mistral.mistral-large-3-675b-instruct',
|
||||
'mistral.mistral-large-2411-v1:0',
|
||||
'mistral.mistral-large-2407-v1:0',
|
||||
'mistral.magistral-small-2509',
|
||||
'mistral.ministral-3-14b-instruct',
|
||||
'mistral.ministral-3-8b-instruct',
|
||||
'mistral.ministral-3-3b-instruct',
|
||||
'mistral.mixtral-8x7b-instruct-v0:1',
|
||||
'amazon.titan-text-premier-v1:0',
|
||||
'cohere.command-r-v1:0',
|
||||
'cohere.command-r-plus-v1:0',
|
||||
])
|
||||
|
||||
/**
|
||||
* Converts a model ID to the Bedrock inference profile format.
|
||||
* AWS Bedrock requires inference profile IDs (e.g., us.anthropic.claude-...)
|
||||
* for on-demand invocation of newer models, while some models only accept
|
||||
* the bare in-region model ID.
|
||||
*
|
||||
* @param modelId - The model ID (e.g., "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0")
|
||||
* @param region - The AWS region (e.g., "us-east-1")
|
||||
* @returns The inference profile ID (e.g., "us.anthropic.claude-sonnet-4-5-20250929-v1:0")
|
||||
*/
|
||||
export function getBedrockInferenceProfileId(modelId: string, region: string): string {
|
||||
const baseModelId = modelId.startsWith('bedrock/') ? modelId.slice(8) : modelId
|
||||
|
||||
if (/^(us-gov|us|eu|apac|au|ca|jp|global)\./.test(baseModelId)) {
|
||||
return baseModelId
|
||||
}
|
||||
|
||||
if (GEO_PROFILE_UNSUPPORTED_MODEL_IDS.has(baseModelId)) {
|
||||
return baseModelId
|
||||
}
|
||||
|
||||
let inferencePrefix: string
|
||||
if (region.startsWith('us-gov-')) {
|
||||
inferencePrefix = 'us-gov'
|
||||
} else if (region.startsWith('us-') || region.startsWith('ca-')) {
|
||||
inferencePrefix = 'us'
|
||||
} else if (region.startsWith('eu-') || region === 'il-central-1') {
|
||||
inferencePrefix = 'eu'
|
||||
} else if (region.startsWith('ap-') || region.startsWith('me-')) {
|
||||
inferencePrefix = 'apac'
|
||||
} else if (region.startsWith('sa-')) {
|
||||
inferencePrefix = 'us'
|
||||
} else if (region.startsWith('af-')) {
|
||||
inferencePrefix = 'eu'
|
||||
} else {
|
||||
inferencePrefix = 'us'
|
||||
}
|
||||
|
||||
return `${inferencePrefix}.${baseModelId}`
|
||||
}
|
||||
Reference in New Issue
Block a user