feat(ai): round-trip Bedrock redacted reasoning (#38623)
This commit is contained in:
parent
65c5c7e3f6
commit
ea010ab3a4
2 changed files with 99 additions and 13 deletions
|
|
@ -66,14 +66,15 @@ const BedrockToolResultBlock = Schema.Struct({
|
||||||
type BedrockToolResultBlock = Schema.Schema.Type<typeof BedrockToolResultBlock>
|
type BedrockToolResultBlock = Schema.Schema.Type<typeof BedrockToolResultBlock>
|
||||||
|
|
||||||
const BedrockReasoningBlock = Schema.Struct({
|
const BedrockReasoningBlock = Schema.Struct({
|
||||||
reasoningContent: Schema.Struct({
|
reasoningContent: Schema.Union([
|
||||||
reasoningText: Schema.optional(
|
Schema.Struct({
|
||||||
Schema.Struct({
|
reasoningText: Schema.Struct({
|
||||||
text: Schema.String,
|
text: Schema.String,
|
||||||
signature: Schema.optional(Schema.String),
|
signature: Schema.optional(Schema.String),
|
||||||
}),
|
}),
|
||||||
),
|
}),
|
||||||
}),
|
Schema.Struct({ redactedContent: Schema.String }),
|
||||||
|
]),
|
||||||
})
|
})
|
||||||
|
|
||||||
const BedrockUserBlock = Schema.Union([
|
const BedrockUserBlock = Schema.Union([
|
||||||
|
|
@ -181,6 +182,8 @@ const BedrockEvent = Schema.Struct({
|
||||||
Schema.Struct({
|
Schema.Struct({
|
||||||
text: Schema.optional(Schema.String),
|
text: Schema.optional(Schema.String),
|
||||||
signature: Schema.optional(Schema.String),
|
signature: Schema.optional(Schema.String),
|
||||||
|
// Blob fields in Bedrock's JSON event stream are base64 strings.
|
||||||
|
redactedContent: Schema.optional(Schema.String),
|
||||||
}),
|
}),
|
||||||
),
|
),
|
||||||
}),
|
}),
|
||||||
|
|
@ -260,6 +263,13 @@ const reasoningSignature = (part: ReasoningPart) => {
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const reasoningRedactedData = (part: ReasoningPart) => {
|
||||||
|
const bedrock = part.providerMetadata?.bedrock
|
||||||
|
return ProviderShared.isRecord(bedrock) && typeof bedrock.redactedData === "string"
|
||||||
|
? bedrock.redactedData
|
||||||
|
: undefined
|
||||||
|
}
|
||||||
|
|
||||||
const lowerToolCall = (part: ToolCallPart): BedrockToolUseBlock => ({
|
const lowerToolCall = (part: ToolCallPart): BedrockToolUseBlock => ({
|
||||||
toolUse: {
|
toolUse: {
|
||||||
toolUseId: part.id,
|
toolUseId: part.id,
|
||||||
|
|
@ -349,11 +359,13 @@ const lowerMessages = Effect.fn("BedrockConverse.lowerMessages")(function* (
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if (part.type === "reasoning") {
|
if (part.type === "reasoning") {
|
||||||
content.push({
|
const signature = reasoningSignature(part)
|
||||||
reasoningContent: {
|
const redactedData = reasoningRedactedData(part)
|
||||||
reasoningText: { text: part.text, signature: reasoningSignature(part) },
|
if (signature === undefined && redactedData !== undefined) {
|
||||||
},
|
content.push({ reasoningContent: { redactedContent: redactedData } })
|
||||||
})
|
continue
|
||||||
|
}
|
||||||
|
content.push({ reasoningContent: { reasoningText: { text: part.text, signature } } })
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if (part.type === "tool-call") {
|
if (part.type === "tool-call") {
|
||||||
|
|
@ -519,12 +531,20 @@ const step = (state: ParserState, event: BedrockEvent) =>
|
||||||
const index = event.contentBlockDelta.contentBlockIndex
|
const index = event.contentBlockDelta.contentBlockIndex
|
||||||
const reasoning = event.contentBlockDelta.delta.reasoningContent
|
const reasoning = event.contentBlockDelta.delta.reasoningContent
|
||||||
const events: LLMEvent[] = []
|
const events: LLMEvent[] = []
|
||||||
|
const lifecycle = reasoning.text
|
||||||
|
? Lifecycle.reasoningDelta(state.lifecycle, events, `reasoning-${index}`, reasoning.text)
|
||||||
|
: reasoning.redactedContent !== undefined
|
||||||
|
? Lifecycle.reasoningStart(
|
||||||
|
state.lifecycle,
|
||||||
|
events,
|
||||||
|
`reasoning-${index}`,
|
||||||
|
bedrockMetadata({ redactedData: reasoning.redactedContent }),
|
||||||
|
)
|
||||||
|
: state.lifecycle
|
||||||
return [
|
return [
|
||||||
{
|
{
|
||||||
...state,
|
...state,
|
||||||
lifecycle: reasoning.text
|
lifecycle,
|
||||||
? Lifecycle.reasoningDelta(state.lifecycle, events, `reasoning-${index}`, reasoning.text)
|
|
||||||
: state.lifecycle,
|
|
||||||
reasoningSignatures: reasoning.signature
|
reasoningSignatures: reasoning.signature
|
||||||
? { ...state.reasoningSignatures, [index]: reasoning.signature }
|
? { ...state.reasoningSignatures, [index]: reasoning.signature }
|
||||||
: state.reasoningSignatures,
|
: state.reasoningSignatures,
|
||||||
|
|
|
||||||
|
|
@ -467,6 +467,72 @@ describe("Bedrock Converse route", () => {
|
||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
it.effect("round-trips streamed redacted reasoning with tool use into a continuation request", () =>
|
||||||
|
Effect.gen(function* () {
|
||||||
|
// Bedrock represents redactedContent blobs as base64 strings on its JSON
|
||||||
|
// wire. The provider owns the payload and requires byte-exact replay.
|
||||||
|
const redactedData = "cmVkYWN0ZWQtdGhpbmtpbmc="
|
||||||
|
const response = yield* LLMClient.generate(
|
||||||
|
LLM.updateRequest(baseRequest, {
|
||||||
|
tools: [{ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } }],
|
||||||
|
}),
|
||||||
|
).pipe(
|
||||||
|
Effect.provide(
|
||||||
|
fixedBytes(
|
||||||
|
eventStreamBody(
|
||||||
|
["messageStart", { role: "assistant" }],
|
||||||
|
[
|
||||||
|
"contentBlockDelta",
|
||||||
|
{ contentBlockIndex: 0, delta: { reasoningContent: { redactedContent: redactedData } } },
|
||||||
|
],
|
||||||
|
["contentBlockStop", { contentBlockIndex: 0 }],
|
||||||
|
[
|
||||||
|
"contentBlockStart",
|
||||||
|
{
|
||||||
|
contentBlockIndex: 1,
|
||||||
|
start: { toolUse: { toolUseId: "tool_1", name: "lookup" } },
|
||||||
|
},
|
||||||
|
],
|
||||||
|
[
|
||||||
|
"contentBlockDelta",
|
||||||
|
{ contentBlockIndex: 1, delta: { toolUse: { input: '{"query":"weather"}' } } },
|
||||||
|
],
|
||||||
|
["contentBlockStop", { contentBlockIndex: 1 }],
|
||||||
|
["messageStop", { stopReason: "tool_use" }],
|
||||||
|
),
|
||||||
|
),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
const prepared = yield* LLMClient.prepare<BedrockConverse.BedrockConverseBody>(
|
||||||
|
LLM.request({
|
||||||
|
model,
|
||||||
|
messages: [
|
||||||
|
Message.user("Say hello."),
|
||||||
|
response.message,
|
||||||
|
Message.tool({ id: "tool_1", name: "lookup", result: "sunny", resultType: "text" }),
|
||||||
|
],
|
||||||
|
tools: [{ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } }],
|
||||||
|
cache: "none",
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
|
||||||
|
expect(prepared.body.messages).toEqual([
|
||||||
|
{ role: "user", content: [{ text: "Say hello." }] },
|
||||||
|
{
|
||||||
|
role: "assistant",
|
||||||
|
content: [
|
||||||
|
{ reasoningContent: { redactedContent: redactedData } },
|
||||||
|
{ toolUse: { toolUseId: "tool_1", name: "lookup", input: { query: "weather" } } },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
{
|
||||||
|
role: "user",
|
||||||
|
content: [{ toolResult: { toolUseId: "tool_1", content: [{ text: "sunny" }], status: "success" } }],
|
||||||
|
},
|
||||||
|
])
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
|
||||||
it.effect("classifies throttlingException as a rate limit", () =>
|
it.effect("classifies throttlingException as a rate limit", () =>
|
||||||
Effect.gen(function* () {
|
Effect.gen(function* () {
|
||||||
const body = eventStreamBody(
|
const body = eventStreamBody(
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue