From 5d5456c1c5f71df23af2f7b147ccb667dd92c181 Mon Sep 17 00:00:00 2001
From: glay <glay.guo@gmail.com>
Date: Wed, 6 Nov 2024 17:23:53 +0800
Subject: [PATCH] =?UTF-8?q?=09=E4=BF=AE=E6=94=B9=EF=BC=9A=20=20=20=20=20ap?=
 =?UTF-8?q?p/api/bedrock.ts=20=09=E4=BF=AE=E6=94=B9=EF=BC=9A=20=20=20=20?=
 =?UTF-8?q?=20app/client/platforms/bedrock.ts?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app/api/bedrock.ts              | 25 +++++++++++++++++++++++++
 app/client/platforms/bedrock.ts |  9 ++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/app/api/bedrock.ts b/app/api/bedrock.ts
index 4a2a1fc5d..b2b04cc32 100644
--- a/app/api/bedrock.ts
+++ b/app/api/bedrock.ts
@@ -46,6 +46,7 @@ export interface ConverseRequest {
     description?: string;
     input_schema: any;
   }[];
+  stream?: boolean;
 }
 
 function supportsToolUse(modelId: string): boolean {
@@ -188,6 +189,30 @@ export async function handle(
       throw new Error("No stream in response");
     }
 
+    // If stream is false, accumulate the response and return as JSON
+    console.log("Body.stream==========" + body.stream);
+    if (body.stream === false) {
+      let fullResponse = {
+        content: "",
+      };
+
+      const responseStream =
+        response.stream as AsyncIterable<ConverseStreamOutput>;
+      for await (const event of responseStream) {
+        if (
+          "contentBlockDelta" in event &&
+          event.contentBlockDelta?.delta &&
+          "text" in event.contentBlockDelta.delta &&
+          event.contentBlockDelta.delta.text
+        ) {
+          fullResponse.content += event.contentBlockDelta.delta.text;
+        }
+      }
+
+      return NextResponse.json(fullResponse);
+    }
+
+    // Otherwise, return streaming response
     const stream = new ReadableStream({
       async start(controller) {
         try {
diff --git a/app/client/platforms/bedrock.ts b/app/client/platforms/bedrock.ts
index 6d236344e..8dca4eff5 100644
--- a/app/client/platforms/bedrock.ts
+++ b/app/client/platforms/bedrock.ts
@@ -38,9 +38,11 @@ export class BedrockApi implements LLMApi {
   }
 
   extractMessage(res: any) {
-    console.log("[Response] claude response: ", res);
-
-    return res?.content?.[0]?.text;
+    console.log("[Response] Bedrock not stream response: ", res);
+    if (res.error) {
+      return "```\n" + JSON.stringify(res, null, 4) + "\n```";
+    }
+    return res?.content ?? res;
   }
 
   async chat(options: ChatOptions): Promise<void> {
@@ -144,6 +146,7 @@ export class BedrockApi implements LLMApi {
         topP: modelConfig.top_p,
         stopSequences: [],
       },
+      stream: shouldStream,
     };
 
     const conversePath = `${ApiPath.Bedrock}/converse`;