From c728cb7015519338ab3d4ca2a2a04be6d3431c99 Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Thu, 18 Jul 2024 18:04:58 +0000
Subject: [PATCH] feat: add ToolChoice to docs

---
 docs/openapi.json    | 15 ++++++++++++++-
 router/src/lib.rs    |  2 +-
 router/src/server.rs |  3 ++-
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/docs/openapi.json b/docs/openapi.json
index 3e7050ab..7000c7b7 100644
--- a/docs/openapi.json
+++ b/docs/openapi.json
@@ -909,7 +909,7 @@
           "tool_choice": {
             "allOf": [
               {
-                "$ref": "#/components/schemas/ToolType"
+                "$ref": "#/components/schemas/ToolChoice"
               }
             ],
             "nullable": true
@@ -2035,6 +2035,14 @@
           }
         }
       },
+      "ToolChoice": {
+        "allOf": [
+          {
+            "$ref": "#/components/schemas/ToolType"
+          }
+        ],
+        "nullable": true
+      },
       "ToolType": {
         "oneOf": [
           {
@@ -2055,6 +2063,11 @@
                 "$ref": "#/components/schemas/FunctionName"
               }
             }
+          },
+          {
+            "type": "object",
+            "default": null,
+            "nullable": true
           }
         ]
       },
diff --git a/router/src/lib.rs b/router/src/lib.rs
index 43a6e056..3cf16077 100644
--- a/router/src/lib.rs
+++ b/router/src/lib.rs
@@ -854,7 +854,7 @@ pub struct FunctionName {
     pub name: String,
 }
 
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default, ToSchema)]
 #[serde(from = "ToolTypeDeserializer")]
 pub struct ToolChoice(pub Option<ToolType>);
 
diff --git a/router/src/server.rs b/router/src/server.rs
index 75fcce2b..c56c39a3 100644
--- a/router/src/server.rs
+++ b/router/src/server.rs
@@ -24,7 +24,7 @@ use crate::{
     CompletionRequest, CompletionType, DeltaToolCall, Function, Prompt, Tool, VertexRequest,
     VertexResponse,
 };
-use crate::{FunctionDefinition, HubPreprocessorConfig, ToolCall, ToolType};
+use crate::{FunctionDefinition, HubPreprocessorConfig, ToolCall, ToolChoice, ToolType};
 use async_stream::__private::AsyncStream;
 use axum::extract::Extension;
 use axum::http::{HeaderMap, Method, StatusCode};
@@ -1492,6 +1492,7 @@ pub async fn run(
     ToolCall,
     Function,
     FunctionDefinition,
+    ToolChoice,
     )
     ),
     tags(