From 56786f9e326bf27340ff06dedb327d22d5ac58b4 Mon Sep 17 00:00:00 2001
From: Aiden Dai <daixb@amazon.com>
Date: Tue, 11 Jun 2024 10:53:56 +0800
Subject: [PATCH] Update api response

---
 src/api/models/bedrock.py | 3 ++-
 src/api/schema.py         | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/api/models/bedrock.py b/src/api/models/bedrock.py
index 96cdc05..6063ffb 100644
--- a/src/api/models/bedrock.py
+++ b/src/api/models/bedrock.py
@@ -463,6 +463,7 @@ class BedrockModel(BaseChatModel):
                 message = ChatResponseMessage(
                     tool_calls=[
                         ToolCall(
+                            index=0,
                             type="function",
                             id=delta["toolUse"]["toolUseId"],
                             function=ResponseFunction(
@@ -484,7 +485,7 @@ class BedrockModel(BaseChatModel):
                 message = ChatResponseMessage(
                     tool_calls=[
                         ToolCall(
-                            type="function",
+                            index=0,
                             function=ResponseFunction(
                                 arguments=delta["toolUse"]["input"],
                             )
diff --git a/src/api/schema.py b/src/api/schema.py
index 2b2f2fb..0ceb88b 100644
--- a/src/api/schema.py
+++ b/src/api/schema.py
@@ -22,6 +22,7 @@ class ResponseFunction(BaseModel):
 
 
 class ToolCall(BaseModel):
+    index: int | None = None
     id: str | None = None
     type: Literal["function"] = "function"
     function: ResponseFunction