Procházet zdrojové kódy

feat(ai): round-trip Bedrock redacted reasoning (#38623)

Aiden Cline před 2 týdny
rodič
revize
ea010ab3a4

+ 33 - 13
packages/ai/src/protocols/bedrock-converse.ts

@@ -66,14 +66,15 @@ const BedrockToolResultBlock = Schema.Struct({
 type BedrockToolResultBlock = Schema.Schema.Type<typeof BedrockToolResultBlock>
 
 const BedrockReasoningBlock = Schema.Struct({
-  reasoningContent: Schema.Struct({
-    reasoningText: Schema.optional(
-      Schema.Struct({
+  reasoningContent: Schema.Union([
+    Schema.Struct({
+      reasoningText: Schema.Struct({
         text: Schema.String,
         signature: Schema.optional(Schema.String),
       }),
-    ),
-  }),
+    }),
+    Schema.Struct({ redactedContent: Schema.String }),
+  ]),
 })
 
 const BedrockUserBlock = Schema.Union([
@@ -181,6 +182,8 @@ const BedrockEvent = Schema.Struct({
             Schema.Struct({
               text: 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 => ({
   toolUse: {
     toolUseId: part.id,
@@ -349,11 +359,13 @@ const lowerMessages = Effect.fn("BedrockConverse.lowerMessages")(function* (
           continue
         }
         if (part.type === "reasoning") {
-          content.push({
-            reasoningContent: {
-              reasoningText: { text: part.text, signature: reasoningSignature(part) },
-            },
-          })
+          const signature = reasoningSignature(part)
+          const redactedData = reasoningRedactedData(part)
+          if (signature === undefined && redactedData !== undefined) {
+            content.push({ reasoningContent: { redactedContent: redactedData } })
+            continue
+          }
+          content.push({ reasoningContent: { reasoningText: { text: part.text, signature } } })
           continue
         }
         if (part.type === "tool-call") {
@@ -519,12 +531,20 @@ const step = (state: ParserState, event: BedrockEvent) =>
       const index = event.contentBlockDelta.contentBlockIndex
       const reasoning = event.contentBlockDelta.delta.reasoningContent
       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 [
         {
           ...state,
-          lifecycle: reasoning.text
-            ? Lifecycle.reasoningDelta(state.lifecycle, events, `reasoning-${index}`, reasoning.text)
-            : state.lifecycle,
+          lifecycle,
           reasoningSignatures: reasoning.signature
             ? { ...state.reasoningSignatures, [index]: reasoning.signature }
             : state.reasoningSignatures,

+ 66 - 0
packages/ai/test/provider/bedrock-converse.test.ts

@@ -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", () =>
     Effect.gen(function* () {
       const body = eventStreamBody(