From 9ecbc3e5f0c9ed6624b9eb10d806975fe59b644a Mon Sep 17 00:00:00 2001
From: Jack Tysoe <jack.tysoe@konghq.com>
Date: Wed, 27 Nov 2024 01:49:17 +0000
Subject: [PATCH] fix(AG-178): fix gemini parsing multiple chained tool calls

(cherry picked from commit 581571b0c35035f0304282978610f2cbe4f2c8b4)
---
 kong/llm/drivers/gemini.lua | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/kong/llm/drivers/gemini.lua b/kong/llm/drivers/gemini.lua
index 1cbdd6095e5f..99ba1b4adbac 100644
--- a/kong/llm/drivers/gemini.lua
+++ b/kong/llm/drivers/gemini.lua
@@ -268,22 +268,23 @@ local function from_gemini_chat_openai(response, model_info, route_type)
       messages.model = model_info.name
 
     elseif is_tool_content(response) then
+      messages.choices[1] = {
+        index = 0,
+        message = {
+          role = "assistant",
+          tool_calls = {},
+        },
+      }
+
       local function_call_responses = response.candidates[1].content.parts
       for i, v in ipairs(function_call_responses) do
-        messages.choices[i] = {
-          index = 0,
-          message = {
-            role = "assistant",
-            tool_calls = {
-              {
-                ['function'] = {
-                  name = v.functionCall.name,
-                  arguments = cjson.encode(v.functionCall.args),
-                },
-              },
+        messages.choices[1].message.tool_calls[i] =
+          {
+            ['function'] = {
+              name = v.functionCall.name,
+              arguments = cjson.encode(v.functionCall.args),
             },
-          },
-        }
+          }
       end
     end