From ad9ee0391ffbc860ca6e56cb36a906d4c12474e1 Mon Sep 17 00:00:00 2001
From: Azriel Hoh <azriel@healthpoint.co.nz>
Date: Thu, 30 Oct 2025 10:36:08 +1300
Subject: [PATCH] Return `tool_calls` when `finish_reason` is `"max_tokens"`.

---
 src/api/models/bedrock.py | 43 +++++++++++++++++++++++++--------------
 1 file changed, 28 insertions(+), 15 deletions(-)

diff --git a/src/api/models/bedrock.py b/src/api/models/bedrock.py
index 4253e1ff..ea56ba4a 100644
--- a/src/api/models/bedrock.py
+++ b/src/api/models/bedrock.py
@@ -858,6 +858,29 @@ def _estimate_reasoning_tokens(self, content: list[dict]) -> int:
             return len(ENCODER.encode(reasoning_text))
         return 0
 
+    def _append_tool_use(
+        self,
+        message: ChatResponseMessage,
+        tool: dict,
+    ):
+        # https://docs.aws.amazon.com/bedrock/latest/userguide/tool-use.html#tool-use-examples
+        tool_calls = []
+        tool_calls.append(
+            ToolCall(
+                id=tool["toolUseId"],
+                type="function",
+                function=ResponseFunction(
+                    name=tool["name"],
+                    arguments=json.dumps(tool["input"]),
+                ),
+            )
+        )
+        if message.tool_calls is None:
+            message.tool_calls = tool_calls
+        else:
+            message.tool_calls += tool_calls
+        message.content = None
+
     def _create_response(
         self,
         model: str,
@@ -874,23 +897,10 @@ def _create_response(
             role="assistant",
         )
         if finish_reason == "tool_use":
-            # https://docs.aws.amazon.com/bedrock/latest/userguide/tool-use.html#tool-use-examples
-            tool_calls = []
             for part in content:
                 if "toolUse" in part:
                     tool = part["toolUse"]
-                    tool_calls.append(
-                        ToolCall(
-                            id=tool["toolUseId"],
-                            type="function",
-                            function=ResponseFunction(
-                                name=tool["name"],
-                                arguments=json.dumps(tool["input"]),
-                            ),
-                        )
-                    )
-            message.tool_calls = tool_calls
-            message.content = None
+                    self._append_tool_use(message, tool)
         else:
             message.content = ""
             for c in content:
@@ -900,9 +910,12 @@ def _create_response(
                     ].get("text", "")
                 elif "text" in c:
                     message.content = c["text"]
+                elif "toolUse" in c:
+                    tool = c["toolUse"]
+                    self._append_tool_use(message, tool)
                 else:
                     logger.warning(
-                        "Unknown tag in message content " + ",".join(c.keys())
+                        "Unknown tag in message content " + ",".join(c.keys()) + ". finish_reason is: " + finish_reason
                     )
             if message.reasoning_content:
                 message.content = f"<think>{message.reasoning_content}</think>{message.content}"