From 6dc080aeaf41a4da40eaa558d997bedaddf73611 Mon Sep 17 00:00:00 2001
From: adavyas <adavyasharma@gmail.com>
Date: Sat, 14 Mar 2026 17:01:20 -0700
Subject: [PATCH] Use Gemini function names for tool responses

---
 src/utils/clients.py        | 12 +++++++++++-
 tests/utils/test_clients.py |  2 +-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/utils/clients.py b/src/utils/clients.py
index 0591fc85..b2d1054f 100644
--- a/src/utils/clients.py
+++ b/src/utils/clients.py
@@ -130,6 +130,7 @@ def _build_gemini_contents_from_messages(
     """Convert generic chat messages into Gemini contents and system instruction."""
     system_instruction_parts: list[str] = []
     gemini_contents: list[dict[str, Any]] = []
+    tool_use_names: dict[str, str] = {}
 
     for msg in messages:
         role = msg.get("role", "user")
@@ -168,6 +169,9 @@ def _build_gemini_contents_from_messages(
                     continue
 
                 if block_type == "tool_use" and isinstance(block.get("name"), str):
+                    tool_use_id = block.get("id")
+                    if isinstance(tool_use_id, str):
+                        tool_use_names[tool_use_id] = block["name"]
                     tool_args = block.get("input")
                     parts.append(
                         {
@@ -183,10 +187,16 @@ def _build_gemini_contents_from_messages(
                     tool_result = block.get("content")
                     if not isinstance(tool_result, str):
                         tool_result = json.dumps(tool_result)
+                    tool_use_id = block.get("tool_use_id")
+                    function_name = (
+                        tool_use_names.get(tool_use_id)
+                        if isinstance(tool_use_id, str)
+                        else None
+                    )
                     parts.append(
                         {
                             "function_response": {
-                                "name": str(block.get("tool_use_id", "tool_result")),
+                                "name": function_name or "tool_result",
                                 "response": {
                                     "result": tool_result,
                                     "is_error": bool(block.get("is_error", False)),
diff --git a/tests/utils/test_clients.py b/tests/utils/test_clients.py
index b0103f8c..3eb57311 100644
--- a/tests/utils/test_clients.py
+++ b/tests/utils/test_clients.py
@@ -982,7 +982,7 @@ class TestGoogleClient:
                 "parts": [
                     {
                         "function_response": {
-                            "name": "toolu_1",
+                            "name": "search_memory",
                             "response": {
                                 "result": "Alice likes tea",
                                 "is_error": False,