532 lines
18 KiB
TypeScript
532 lines
18 KiB
TypeScript
import { Effect, Schema } from "effect"
|
|
import { Route, type RouteModelInput } from "../route/client"
|
|
import { Endpoint } from "../route/endpoint"
|
|
import { Protocol } from "../route/protocol"
|
|
import {
|
|
Usage,
|
|
type CacheHint,
|
|
type FinishReason,
|
|
type LLMEvent,
|
|
type LLMRequest,
|
|
type ToolCallPart,
|
|
type ToolDefinition,
|
|
type ToolResultPart,
|
|
} from "../schema"
|
|
import { BedrockEventStream } from "./bedrock-event-stream"
|
|
import { JsonObject, optionalArray, ProviderShared } from "./shared"
|
|
import { BedrockAuth, type Credentials as BedrockCredentials } from "./utils/bedrock-auth"
|
|
import { BedrockCache } from "./utils/bedrock-cache"
|
|
import { BedrockMedia } from "./utils/bedrock-media"
|
|
import { ToolStream } from "./utils/tool-stream"
|
|
|
|
const ADAPTER = "bedrock-converse"
|
|
|
|
export type { Credentials as BedrockCredentials } from "./utils/bedrock-auth"
|
|
|
|
// =============================================================================
|
|
// Public Model Input
|
|
// =============================================================================
|
|
export type BedrockConverseModelInput = RouteModelInput & {
|
|
/**
|
|
* Bearer API key (Bedrock's newer API key auth). Sets the `Authorization`
|
|
* header and bypasses SigV4 signing. Mutually exclusive with `credentials`.
|
|
*/
|
|
readonly apiKey?: string
|
|
/**
|
|
* AWS credentials for SigV4 signing. The route signs each request at
|
|
* `toHttp` time using `aws4fetch`. Mutually exclusive with `apiKey`.
|
|
*/
|
|
readonly credentials?: BedrockCredentials
|
|
readonly headers?: Record<string, string>
|
|
}
|
|
|
|
// =============================================================================
|
|
// Request Body Schema
|
|
// =============================================================================
|
|
const BedrockTextBlock = Schema.Struct({
|
|
text: Schema.String,
|
|
})
|
|
type BedrockTextBlock = Schema.Schema.Type<typeof BedrockTextBlock>
|
|
|
|
const BedrockToolUseBlock = Schema.Struct({
|
|
toolUse: Schema.Struct({
|
|
toolUseId: Schema.String,
|
|
name: Schema.String,
|
|
input: Schema.Unknown,
|
|
}),
|
|
})
|
|
type BedrockToolUseBlock = Schema.Schema.Type<typeof BedrockToolUseBlock>
|
|
|
|
const BedrockToolResultContentItem = Schema.Union([
|
|
Schema.Struct({ text: Schema.String }),
|
|
Schema.Struct({ json: Schema.Unknown }),
|
|
])
|
|
|
|
const BedrockToolResultBlock = Schema.Struct({
|
|
toolResult: Schema.Struct({
|
|
toolUseId: Schema.String,
|
|
content: Schema.Array(BedrockToolResultContentItem),
|
|
status: Schema.optional(Schema.Literals(["success", "error"])),
|
|
}),
|
|
})
|
|
type BedrockToolResultBlock = Schema.Schema.Type<typeof BedrockToolResultBlock>
|
|
|
|
const BedrockReasoningBlock = Schema.Struct({
|
|
reasoningContent: Schema.Struct({
|
|
reasoningText: Schema.optional(
|
|
Schema.Struct({
|
|
text: Schema.String,
|
|
signature: Schema.optional(Schema.String),
|
|
}),
|
|
),
|
|
}),
|
|
})
|
|
|
|
const BedrockUserBlock = Schema.Union([
|
|
BedrockTextBlock,
|
|
BedrockMedia.ImageBlock,
|
|
BedrockMedia.DocumentBlock,
|
|
BedrockToolResultBlock,
|
|
BedrockCache.CachePointBlock,
|
|
])
|
|
type BedrockUserBlock = Schema.Schema.Type<typeof BedrockUserBlock>
|
|
|
|
const BedrockAssistantBlock = Schema.Union([
|
|
BedrockTextBlock,
|
|
BedrockReasoningBlock,
|
|
BedrockToolUseBlock,
|
|
BedrockCache.CachePointBlock,
|
|
])
|
|
type BedrockAssistantBlock = Schema.Schema.Type<typeof BedrockAssistantBlock>
|
|
|
|
const BedrockMessage = Schema.Union([
|
|
Schema.Struct({ role: Schema.Literal("user"), content: Schema.Array(BedrockUserBlock) }),
|
|
Schema.Struct({ role: Schema.Literal("assistant"), content: Schema.Array(BedrockAssistantBlock) }),
|
|
]).pipe(Schema.toTaggedUnion("role"))
|
|
type BedrockMessage = Schema.Schema.Type<typeof BedrockMessage>
|
|
|
|
const BedrockSystemBlock = Schema.Union([BedrockTextBlock, BedrockCache.CachePointBlock])
|
|
type BedrockSystemBlock = Schema.Schema.Type<typeof BedrockSystemBlock>
|
|
|
|
const BedrockTool = Schema.Struct({
|
|
toolSpec: Schema.Struct({
|
|
name: Schema.String,
|
|
description: Schema.String,
|
|
inputSchema: Schema.Struct({
|
|
json: JsonObject,
|
|
}),
|
|
}),
|
|
})
|
|
type BedrockTool = Schema.Schema.Type<typeof BedrockTool>
|
|
|
|
const BedrockToolChoice = Schema.Union([
|
|
Schema.Struct({ auto: Schema.Struct({}) }),
|
|
Schema.Struct({ any: Schema.Struct({}) }),
|
|
Schema.Struct({ tool: Schema.Struct({ name: Schema.String }) }),
|
|
])
|
|
|
|
const BedrockBodyFields = {
|
|
modelId: Schema.String,
|
|
messages: Schema.Array(BedrockMessage),
|
|
system: optionalArray(BedrockSystemBlock),
|
|
inferenceConfig: Schema.optional(
|
|
Schema.Struct({
|
|
maxTokens: Schema.optional(Schema.Number),
|
|
temperature: Schema.optional(Schema.Number),
|
|
topP: Schema.optional(Schema.Number),
|
|
stopSequences: optionalArray(Schema.String),
|
|
}),
|
|
),
|
|
toolConfig: Schema.optional(
|
|
Schema.Struct({
|
|
tools: Schema.Array(BedrockTool),
|
|
toolChoice: Schema.optional(BedrockToolChoice),
|
|
}),
|
|
),
|
|
additionalModelRequestFields: Schema.optional(JsonObject),
|
|
}
|
|
const BedrockConverseBody = Schema.Struct(BedrockBodyFields)
|
|
export type BedrockConverseBody = Schema.Schema.Type<typeof BedrockConverseBody>
|
|
|
|
const BedrockUsageSchema = Schema.Struct({
|
|
inputTokens: Schema.optional(Schema.Number),
|
|
outputTokens: Schema.optional(Schema.Number),
|
|
totalTokens: Schema.optional(Schema.Number),
|
|
cacheReadInputTokens: Schema.optional(Schema.Number),
|
|
cacheWriteInputTokens: Schema.optional(Schema.Number),
|
|
})
|
|
type BedrockUsageSchema = Schema.Schema.Type<typeof BedrockUsageSchema>
|
|
|
|
// Streaming event shape — the AWS event stream wraps each JSON payload by its
|
|
// `:event-type` header (e.g. `messageStart`, `contentBlockDelta`). We
|
|
// reconstruct that wrapping in `decodeFrames` below so the event schema can
|
|
// stay a plain discriminated record.
|
|
const BedrockEvent = Schema.Struct({
|
|
messageStart: Schema.optional(Schema.Struct({ role: Schema.String })),
|
|
contentBlockStart: Schema.optional(
|
|
Schema.Struct({
|
|
contentBlockIndex: Schema.Number,
|
|
start: Schema.optional(
|
|
Schema.Struct({
|
|
toolUse: Schema.optional(Schema.Struct({ toolUseId: Schema.String, name: Schema.String })),
|
|
}),
|
|
),
|
|
}),
|
|
),
|
|
contentBlockDelta: Schema.optional(
|
|
Schema.Struct({
|
|
contentBlockIndex: Schema.Number,
|
|
delta: Schema.optional(
|
|
Schema.Struct({
|
|
text: Schema.optional(Schema.String),
|
|
toolUse: Schema.optional(Schema.Struct({ input: Schema.String })),
|
|
reasoningContent: Schema.optional(
|
|
Schema.Struct({
|
|
text: Schema.optional(Schema.String),
|
|
signature: Schema.optional(Schema.String),
|
|
}),
|
|
),
|
|
}),
|
|
),
|
|
}),
|
|
),
|
|
contentBlockStop: Schema.optional(Schema.Struct({ contentBlockIndex: Schema.Number })),
|
|
messageStop: Schema.optional(
|
|
Schema.Struct({
|
|
stopReason: Schema.String,
|
|
additionalModelResponseFields: Schema.optional(Schema.Unknown),
|
|
}),
|
|
),
|
|
metadata: Schema.optional(
|
|
Schema.Struct({
|
|
usage: Schema.optional(BedrockUsageSchema),
|
|
metrics: Schema.optional(Schema.Unknown),
|
|
}),
|
|
),
|
|
internalServerException: Schema.optional(Schema.Struct({ message: Schema.String })),
|
|
modelStreamErrorException: Schema.optional(Schema.Struct({ message: Schema.String })),
|
|
validationException: Schema.optional(Schema.Struct({ message: Schema.String })),
|
|
throttlingException: Schema.optional(Schema.Struct({ message: Schema.String })),
|
|
serviceUnavailableException: Schema.optional(Schema.Struct({ message: Schema.String })),
|
|
})
|
|
type BedrockEvent = Schema.Schema.Type<typeof BedrockEvent>
|
|
|
|
// =============================================================================
|
|
// Request Lowering
|
|
// =============================================================================
|
|
const lowerTool = (tool: ToolDefinition): BedrockTool => ({
|
|
toolSpec: {
|
|
name: tool.name,
|
|
description: tool.description,
|
|
inputSchema: { json: tool.inputSchema },
|
|
},
|
|
})
|
|
|
|
const textWithCache = (
|
|
text: string,
|
|
cache: CacheHint | undefined,
|
|
): Array<BedrockTextBlock | BedrockCache.CachePointBlock> => {
|
|
const cachePoint = BedrockCache.block(cache)
|
|
return cachePoint ? [{ text }, cachePoint] : [{ text }]
|
|
}
|
|
|
|
const lowerToolChoice = (toolChoice: NonNullable<LLMRequest["toolChoice"]>) =>
|
|
ProviderShared.matchToolChoice("Bedrock Converse", toolChoice, {
|
|
auto: () => ({ auto: {} }) as const,
|
|
none: () => undefined,
|
|
required: () => ({ any: {} }) as const,
|
|
tool: (name) => ({ tool: { name } }) as const,
|
|
})
|
|
|
|
const lowerToolCall = (part: ToolCallPart): BedrockToolUseBlock => ({
|
|
toolUse: {
|
|
toolUseId: part.id,
|
|
name: part.name,
|
|
input: part.input,
|
|
},
|
|
})
|
|
|
|
const lowerToolResult = (part: ToolResultPart): BedrockToolResultBlock => ({
|
|
toolResult: {
|
|
toolUseId: part.id,
|
|
content:
|
|
part.result.type === "text" || part.result.type === "error"
|
|
? [{ text: ProviderShared.toolResultText(part) }]
|
|
: [{ json: part.result.value }],
|
|
status: part.result.type === "error" ? "error" : "success",
|
|
},
|
|
})
|
|
|
|
const lowerMessages = Effect.fn("BedrockConverse.lowerMessages")(function* (request: LLMRequest) {
|
|
const messages: BedrockMessage[] = []
|
|
|
|
for (const message of request.messages) {
|
|
if (message.role === "user") {
|
|
const content: BedrockUserBlock[] = []
|
|
for (const part of message.content) {
|
|
if (!ProviderShared.supportsContent(part, ["text", "media"]))
|
|
return yield* ProviderShared.unsupportedContent("Bedrock Converse", "user", ["text", "media"])
|
|
if (part.type === "text") {
|
|
content.push(...textWithCache(part.text, part.cache))
|
|
continue
|
|
}
|
|
if (part.type === "media") {
|
|
content.push(yield* BedrockMedia.lower(part))
|
|
continue
|
|
}
|
|
}
|
|
messages.push({ role: "user", content })
|
|
continue
|
|
}
|
|
|
|
if (message.role === "assistant") {
|
|
const content: BedrockAssistantBlock[] = []
|
|
for (const part of message.content) {
|
|
if (!ProviderShared.supportsContent(part, ["text", "reasoning", "tool-call"]))
|
|
return yield* ProviderShared.unsupportedContent("Bedrock Converse", "assistant", [
|
|
"text",
|
|
"reasoning",
|
|
"tool-call",
|
|
])
|
|
if (part.type === "text") {
|
|
content.push(...textWithCache(part.text, part.cache))
|
|
continue
|
|
}
|
|
if (part.type === "reasoning") {
|
|
content.push({
|
|
reasoningContent: {
|
|
reasoningText: { text: part.text, signature: part.encrypted },
|
|
},
|
|
})
|
|
continue
|
|
}
|
|
if (part.type === "tool-call") {
|
|
content.push(lowerToolCall(part))
|
|
continue
|
|
}
|
|
}
|
|
messages.push({ role: "assistant", content })
|
|
continue
|
|
}
|
|
|
|
const content: BedrockToolResultBlock[] = []
|
|
for (const part of message.content) {
|
|
if (!ProviderShared.supportsContent(part, ["tool-result"]))
|
|
return yield* ProviderShared.unsupportedContent("Bedrock Converse", "tool", ["tool-result"])
|
|
content.push(lowerToolResult(part))
|
|
}
|
|
messages.push({ role: "user", content })
|
|
}
|
|
|
|
return messages
|
|
})
|
|
|
|
// System prompts share the cache-point convention: emit the text block, then
|
|
// optionally a positional `cachePoint` marker.
|
|
const lowerSystem = (system: ReadonlyArray<LLMRequest["system"][number]>): BedrockSystemBlock[] =>
|
|
system.flatMap((part) => textWithCache(part.text, part.cache))
|
|
|
|
const fromRequest = Effect.fn("BedrockConverse.fromRequest")(function* (request: LLMRequest) {
|
|
const toolChoice = request.toolChoice ? yield* lowerToolChoice(request.toolChoice) : undefined
|
|
const generation = request.generation
|
|
return {
|
|
modelId: request.model.id,
|
|
messages: yield* lowerMessages(request),
|
|
system: request.system.length === 0 ? undefined : lowerSystem(request.system),
|
|
inferenceConfig:
|
|
generation?.maxTokens === undefined &&
|
|
generation?.temperature === undefined &&
|
|
generation?.topP === undefined &&
|
|
(generation?.stop === undefined || generation.stop.length === 0)
|
|
? undefined
|
|
: {
|
|
maxTokens: generation?.maxTokens,
|
|
temperature: generation?.temperature,
|
|
topP: generation?.topP,
|
|
stopSequences: generation?.stop,
|
|
},
|
|
toolConfig:
|
|
request.tools.length > 0 && request.toolChoice?.type !== "none"
|
|
? { tools: request.tools.map(lowerTool), toolChoice }
|
|
: undefined,
|
|
}
|
|
})
|
|
|
|
// =============================================================================
|
|
// Stream Parsing
|
|
// =============================================================================
|
|
const mapFinishReason = (reason: string): FinishReason => {
|
|
if (reason === "end_turn" || reason === "stop_sequence") return "stop"
|
|
if (reason === "max_tokens") return "length"
|
|
if (reason === "tool_use") return "tool-calls"
|
|
if (reason === "content_filtered" || reason === "guardrail_intervened") return "content-filter"
|
|
return "unknown"
|
|
}
|
|
|
|
const mapUsage = (usage: BedrockUsageSchema | undefined): Usage | undefined => {
|
|
if (!usage) return undefined
|
|
return new Usage({
|
|
inputTokens: usage.inputTokens,
|
|
outputTokens: usage.outputTokens,
|
|
totalTokens: ProviderShared.totalTokens(usage.inputTokens, usage.outputTokens, usage.totalTokens),
|
|
cacheReadInputTokens: usage.cacheReadInputTokens,
|
|
cacheWriteInputTokens: usage.cacheWriteInputTokens,
|
|
native: usage,
|
|
})
|
|
}
|
|
|
|
interface ParserState {
|
|
readonly tools: ToolStream.State<number>
|
|
// Bedrock splits the finish into `messageStop` (carries `stopReason`) and
|
|
// `metadata` (carries usage). Hold the terminal event in state so `onHalt`
|
|
// can emit exactly one finish after both chunks have had a chance to arrive.
|
|
readonly pendingFinish: { readonly reason: FinishReason; readonly usage?: Usage } | undefined
|
|
}
|
|
|
|
const step = (state: ParserState, event: BedrockEvent) =>
|
|
Effect.gen(function* () {
|
|
if (event.contentBlockStart?.start?.toolUse) {
|
|
const index = event.contentBlockStart.contentBlockIndex
|
|
return [
|
|
{
|
|
...state,
|
|
tools: ToolStream.start(state.tools, index, {
|
|
id: event.contentBlockStart.start.toolUse.toolUseId,
|
|
name: event.contentBlockStart.start.toolUse.name,
|
|
}),
|
|
},
|
|
[],
|
|
] as const
|
|
}
|
|
|
|
if (event.contentBlockDelta?.delta?.text) {
|
|
return [state, [{ type: "text-delta" as const, text: event.contentBlockDelta.delta.text }]] as const
|
|
}
|
|
|
|
if (event.contentBlockDelta?.delta?.reasoningContent?.text) {
|
|
return [
|
|
state,
|
|
[{ type: "reasoning-delta" as const, text: event.contentBlockDelta.delta.reasoningContent.text }],
|
|
] as const
|
|
}
|
|
|
|
if (event.contentBlockDelta?.delta?.toolUse) {
|
|
const index = event.contentBlockDelta.contentBlockIndex
|
|
const result = ToolStream.appendExisting(
|
|
ADAPTER,
|
|
state.tools,
|
|
index,
|
|
event.contentBlockDelta.delta.toolUse.input,
|
|
"Bedrock Converse tool delta is missing its tool call",
|
|
)
|
|
if (ToolStream.isError(result)) return yield* result
|
|
return [{ ...state, tools: result.tools }, result.event ? [result.event] : []] as const
|
|
}
|
|
|
|
if (event.contentBlockStop) {
|
|
const result = yield* ToolStream.finish(ADAPTER, state.tools, event.contentBlockStop.contentBlockIndex)
|
|
return [{ ...state, tools: result.tools }, result.event ? [result.event] : []] as const
|
|
}
|
|
|
|
if (event.messageStop) {
|
|
return [
|
|
{
|
|
...state,
|
|
pendingFinish: { reason: mapFinishReason(event.messageStop.stopReason), usage: state.pendingFinish?.usage },
|
|
},
|
|
[],
|
|
] as const
|
|
}
|
|
|
|
if (event.metadata) {
|
|
const usage = mapUsage(event.metadata.usage)
|
|
return [{ ...state, pendingFinish: { reason: state.pendingFinish?.reason ?? "stop", usage } }, []] as const
|
|
}
|
|
|
|
if (event.internalServerException || event.modelStreamErrorException || event.serviceUnavailableException) {
|
|
const message =
|
|
event.internalServerException?.message ??
|
|
event.modelStreamErrorException?.message ??
|
|
event.serviceUnavailableException?.message ??
|
|
"Bedrock Converse stream error"
|
|
return [state, [{ type: "provider-error" as const, message, retryable: true }]] as const
|
|
}
|
|
|
|
if (event.validationException || event.throttlingException) {
|
|
const message =
|
|
event.validationException?.message ?? event.throttlingException?.message ?? "Bedrock Converse error"
|
|
return [
|
|
state,
|
|
[{ type: "provider-error" as const, message, retryable: event.throttlingException !== undefined }],
|
|
] as const
|
|
}
|
|
|
|
return [state, []] as const
|
|
})
|
|
|
|
const framing = BedrockEventStream.framing(ADAPTER)
|
|
|
|
const onHalt = (state: ParserState): ReadonlyArray<LLMEvent> =>
|
|
state.pendingFinish
|
|
? [{ type: "request-finish", reason: state.pendingFinish.reason, usage: state.pendingFinish.usage }]
|
|
: []
|
|
|
|
// =============================================================================
|
|
// Protocol And Bedrock Route
|
|
// =============================================================================
|
|
/**
|
|
* The Bedrock Converse protocol — request body construction, body schema, and
|
|
* the streaming-event state machine.
|
|
*/
|
|
export const protocol = Protocol.make({
|
|
id: ADAPTER,
|
|
body: {
|
|
schema: BedrockConverseBody,
|
|
from: fromRequest,
|
|
},
|
|
stream: {
|
|
event: BedrockEvent,
|
|
initial: () => ({ tools: ToolStream.empty<number>(), pendingFinish: undefined }),
|
|
step,
|
|
onHalt,
|
|
},
|
|
})
|
|
|
|
export const route = Route.make({
|
|
id: ADAPTER,
|
|
protocol,
|
|
// Bedrock's URL embeds the region in the host (set on `model.baseURL` by
|
|
// the provider helper from credentials) and the validated modelId in the
|
|
// path. We read the validated body so the URL matches the body that gets
|
|
// signed.
|
|
endpoint: Endpoint.path<BedrockConverseBody>(
|
|
({ body }) => `/model/${encodeURIComponent(body.modelId)}/converse-stream`,
|
|
),
|
|
auth: BedrockAuth.auth,
|
|
framing,
|
|
})
|
|
|
|
export const nativeCredentials = BedrockAuth.nativeCredentials
|
|
|
|
const bedrockModel = Route.model(
|
|
route,
|
|
{
|
|
provider: "bedrock",
|
|
},
|
|
{
|
|
mapInput: (input: BedrockConverseModelInput) => {
|
|
const { credentials, ...rest } = input
|
|
const region = credentials?.region ?? "us-east-1"
|
|
return {
|
|
...rest,
|
|
baseURL: rest.baseURL ?? `https://bedrock-runtime.${region}.amazonaws.com`,
|
|
native: nativeCredentials(input.native, credentials),
|
|
}
|
|
},
|
|
},
|
|
)
|
|
|
|
export const model = bedrockModel
|
|
|
|
export * as BedrockConverse from "./bedrock-converse"
|