From 523ebf8cba952858bacb046fc1d4ceb965a58bde Mon Sep 17 00:00:00 2001
From: ochafik <ochafik@google.com>
Date: Fri, 27 Dec 2024 02:20:52 +0000
Subject: [PATCH] Simplify tool call grammars when there's only 1 tool

---
 common/tool-call.cpp | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/common/tool-call.cpp b/common/tool-call.cpp
index f6d509f4d..bc0de8ab2 100644
--- a/common/tool-call.cpp
+++ b/common/tool-call.cpp
@@ -448,7 +448,10 @@ llama_tool_call_handler llama_tool_call_handler_init(
                         {"properties", {
                             {"tool_calls", {
                                 {"type", "array"},
-                                {"items", json {{"anyOf", tool_call_schemas}}}
+                                {"items", tool_call_schemas.size() == 1 ? tool_call_schemas[0] : json {
+                                    {"anyOf", tool_call_schemas},
+                                }},
+                                {"minItems", 1},
                             }},
                         }},
                         {"required", json::array({"tool_calls"})},
@@ -456,7 +459,9 @@ llama_tool_call_handler llama_tool_call_handler_init(
                     : json {
                         {"type", "object"},
                         {"properties", {
-                            {"tool_call", json {{"anyOf", tool_call_schemas}}},
+                            {"tool_call", tool_call_schemas.size() == 1 ? tool_call_schemas[0] : json {
+                                {"anyOf", tool_call_schemas},
+                            }},
                         }},
                         {"required", json::array({"tool_call"})},
                     };
@@ -473,6 +478,7 @@ llama_tool_call_handler llama_tool_call_handler_init(
                                         : json_schema
                                     },
                                 }},
+                                {"required", json::array({"response"})},
                             },
                         })}
                     }
@@ -514,7 +520,7 @@ llama_tool_call_handler llama_tool_call_handler_init(
                 }
                 auto schema = json {
                     {"type", "array"},
-                    {"items", json {{"anyOf", schemas}}},
+                    {"items", schemas.size() == 1 ? schemas[0] : json {{"anyOf", schemas}}},
                     {"minItems", 1},
                 };
                 if (!parallel) {
@@ -548,7 +554,7 @@ llama_tool_call_handler llama_tool_call_handler_init(
                 }
                 auto schema = json {
                     {"type", "array"},
-                    {"items", json {{"anyOf", schemas}}},
+                    {"items", schemas.size() == 1 ? schemas[0] : json {{"anyOf", schemas}}},
                     {"minItems", 1},
                 };
                 if (!parallel) {