From 11cde0f8ba373ac79cb6fe94549a344c3fccef37 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Crist=C3=B3v=C3=A3o=20Neves=20Ferreira?=
 <andrecnf@gmail.com>
Date: Sat, 6 Jul 2024 03:49:08 +0200
Subject: [PATCH] fix: :bug: Fix Bedrock Converse's tool use blocks, when there
 are multiple consecutive function calls (#14386)

---
 .../llms/bedrock_converse/utils.py            | 49 ++++++++++++-------
 1 file changed, 30 insertions(+), 19 deletions(-)

diff --git a/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/utils.py b/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/utils.py
index 0762a9282f6f8..b61dd73532d21 100644
--- a/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/utils.py
+++ b/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/utils.py
@@ -120,28 +120,39 @@ def messages_to_converse_messages(
             status = message.additional_kwargs.get("status")
             if status:
                 content["toolResult"]["status"] = status
-            converse_message = {
-                "role": "user",
-                "content": [content],
-            }
-            converse_messages.append(converse_message)
+            converse_messages.append(
+                {
+                    "role": "user",
+                    "content": [content],
+                }
+            )
         else:
-            content = []
             if message.content:
                 # get the text of the message
-                content.append({"text": message.content})
-            # convert tool calls to the AWS Bedrock Converse format
-            tool_calls = message.additional_kwargs.get("tool_calls", [])
-            for tool_call in tool_calls:
-                assert "toolUseId" in tool_call, f"`toolUseId` not found in {tool_call}"
-                assert "input" in tool_call, f"`input` not found in {tool_call}"
-                assert "name" in tool_call, f"`name` not found in {tool_call}"
-                content.append({"toolUse": tool_call})
-            converse_message = {
-                "role": message.role.value,
-                "content": content,
-            }
-            converse_messages.append(converse_message)
+                converse_messages.append(
+                    {
+                        "role": message.role.value,
+                        "content": [{"text": message.content}],
+                    }
+                )
+        # convert tool calls to the AWS Bedrock Converse format
+        # NOTE tool calls might show up within any message,
+        # e.g. within assistant message or in consecutive tool calls,
+        # thus this tool call check is done for all messages
+        tool_calls = message.additional_kwargs.get("tool_calls", [])
+        content = []
+        for tool_call in tool_calls:
+            assert "toolUseId" in tool_call, f"`toolUseId` not found in {tool_call}"
+            assert "input" in tool_call, f"`input` not found in {tool_call}"
+            assert "name" in tool_call, f"`name` not found in {tool_call}"
+            content.append({"toolUse": tool_call})
+        if len(content) > 0:
+            converse_messages.append(
+                {
+                    "role": "assistant",  # tool calls are always from the assistant
+                    "content": content,
+                }
+            )
 
     return __merge_common_role_msgs(converse_messages), system_prompt.strip()