fix streaming for tooluse in ts

awslabs · Sep 25, 2024 · e00478f · e00478f
1 parent cd7764f
commit e00478f
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/typescript/src/agents/bedrockLLMAgent.ts b/typescript/src/agents/bedrockLLMAgent.ts
@@ -193,7 +193,12 @@ export class BedrockLLMAgent extends Agent {
 
       while (continueWithTools && maxRecursions > 0){
         // send the conversation to Amazon Bedrock
-        const bedrockResponse = await this.handleSingleResponse(converseCmd);
+        let bedrockResponse;
+        if (this.streaming) {
+          bedrockResponse =  this.handleStreamingResponse(converseCmd);
+        } else {
+          bedrockResponse = await this.handleSingleResponse(converseCmd);
+        }
 
         // Append the model's response to the ongoing conversation
         conversation.push(bedrockResponse);
@@ -218,7 +223,7 @@ export class BedrockLLMAgent extends Agent {
       if (this.streaming) {
         return this.handleStreamingResponse(converseCmd);
       } else {
-        return this.handleSingleResponse(converseCmd);
+        return await this.handleSingleResponse(converseCmd);
       }
     }
   }