From 8ad20daf33617296aa3982f1db08bc8abaf40f83 Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Tue, 8 Oct 2024 12:35:48 -0400
Subject: [PATCH] CI (2599): Update ToolType input schema (#2601)

* Update ToolType input schema

* lint

* fix: run formatter

* fix: allow tool choide to be null

---------

Co-authored-by: Wauplin <lucainp@gmail.com>
---
 docs/openapi.json                | 23 +++++++++++++----------
 router/src/infer/tool_grammar.rs |  5 +----
 router/src/lib.rs                | 16 ++++++++++++----
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/docs/openapi.json b/docs/openapi.json
index 67394a14..d1b60f4d 100644
--- a/docs/openapi.json
+++ b/docs/openapi.json
@@ -2114,12 +2114,18 @@
       "ToolType": {
         "oneOf": [
           {
-            "type": "object",
-            "default": null,
-            "nullable": true
+            "type": "string",
+            "description": "Means the model can pick between generating a message or calling one or more tools.",
+            "enum": [
+              "auto"
+            ]
           },
           {
-            "type": "string"
+            "type": "string",
+            "description": "Means the model will not call any tool and instead generates a message.",
+            "enum": [
+              "none"
+            ]
           },
           {
             "type": "object",
@@ -2131,13 +2137,10 @@
                 "$ref": "#/components/schemas/FunctionName"
               }
             }
-          },
-          {
-            "type": "object",
-            "default": null,
-            "nullable": true
           }
-        ]
+        ],
+        "description": "Controls which (if any) tool is called by the model.",
+        "example": "auto"
       },
       "Url": {
         "type": "object",
diff --git a/router/src/infer/tool_grammar.rs b/router/src/infer/tool_grammar.rs
index 4fe15720..cc9bf31d 100644
--- a/router/src/infer/tool_grammar.rs
+++ b/router/src/infer/tool_grammar.rs
@@ -53,10 +53,7 @@ impl ToolGrammar {
 
         // if tools are provided and no tool_choice we default to the OneOf
         let tools_to_use = match tool_choice {
-            ToolType::FunctionName(name) => {
-                vec![Self::find_tool_by_name(&tools, &name)?]
-            }
-            ToolType::Function { function } => {
+            ToolType::Function(function) => {
                 vec![Self::find_tool_by_name(&tools, &function.name)?]
             }
             ToolType::OneOf => tools.clone(),
diff --git a/router/src/lib.rs b/router/src/lib.rs
index 0901bafa..b29c9395 100644
--- a/router/src/lib.rs
+++ b/router/src/lib.rs
@@ -957,12 +957,18 @@ pub fn default_tool_prompt() -> String {
 }
 
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize, ToSchema)]
-#[serde(untagged)]
+#[schema(example = "auto")]
+/// Controls which (if any) tool is called by the model.
 pub enum ToolType {
+    /// Means the model can pick between generating a message or calling one or more tools.
+    #[schema(rename = "auto")]
     OneOf,
-    FunctionName(String),
-    Function { function: FunctionName },
+    /// Means the model will not call any tool and instead generates a message.
+    #[schema(rename = "none")]
     NoTool,
+    /// Forces the model to call a specific tool.
+    #[schema(rename = "function")]
+    Function(FunctionName),
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
@@ -977,6 +983,7 @@ pub struct ToolChoice(pub Option<ToolType>);
 #[derive(Deserialize)]
 #[serde(untagged)]
 enum ToolTypeDeserializer {
+    Null,
     String(String),
     ToolType(ToolType),
 }
@@ -984,10 +991,11 @@ enum ToolTypeDeserializer {
 impl From<ToolTypeDeserializer> for ToolChoice {
     fn from(value: ToolTypeDeserializer) -> Self {
         match value {
+            ToolTypeDeserializer::Null => ToolChoice(None),
             ToolTypeDeserializer::String(s) => match s.as_str() {
                 "none" => ToolChoice(Some(ToolType::NoTool)),
                 "auto" => ToolChoice(Some(ToolType::OneOf)),
-                _ => ToolChoice(Some(ToolType::FunctionName(s))),
+                _ => ToolChoice(Some(ToolType::Function(FunctionName { name: s }))),
             },
             ToolTypeDeserializer::ToolType(tool_type) => ToolChoice(Some(tool_type)),
         }