From 08c4184eb2cd07637df5e79b849afec6bea0268e Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Mon, 4 Nov 2024 00:44:59 -0500
Subject: [PATCH] fix: add chat_tokenize endpoint to api docs (#2710)

---
 docs/openapi.json    | 56 ++++++++++++++++++++++++++++++++++++++++++++
 router/src/server.rs |  8 ++++++-
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/docs/openapi.json b/docs/openapi.json
index 903f7426..22b06720 100644
--- a/docs/openapi.json
+++ b/docs/openapi.json
@@ -101,6 +101,47 @@
         }
       }
     },
+    "/chat_tokenize": {
+      "post": {
+        "tags": [
+          "Text Generation Inference"
+        ],
+        "summary": "Template and tokenize ChatRequest",
+        "operationId": "get_chat_tokenize",
+        "requestBody": {
+          "content": {
+            "application/json": {
+              "schema": {
+                "$ref": "#/components/schemas/ChatRequest"
+              }
+            }
+          },
+          "required": true
+        },
+        "responses": {
+          "200": {
+            "description": "Templated and tokenized ChatRequest",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/ChatTokenizeResponse"
+                }
+              }
+            }
+          },
+          "404": {
+            "description": "Failed to tokenize ChatRequest",
+            "content": {
+              "application/json": {
+                "schema": {
+                  "$ref": "#/components/schemas/ErrorResponse"
+                }
+              }
+            }
+          }
+        }
+      }
+    },
     "/generate": {
       "post": {
         "tags": [
@@ -1092,6 +1133,21 @@
           }
         }
       },
+      "ChatTokenizeResponse": {
+        "type": "object",
+        "required": [
+          "tokenize_response",
+          "templated_text"
+        ],
+        "properties": {
+          "templated_text": {
+            "type": "string"
+          },
+          "tokenize_response": {
+            "$ref": "#/components/schemas/TokenizeResponse"
+          }
+        }
+      },
       "Chunk": {
         "type": "object",
         "required": [
diff --git a/router/src/server.rs b/router/src/server.rs
index 863607b1..7d8d518c 100644
--- a/router/src/server.rs
+++ b/router/src/server.rs
@@ -181,12 +181,16 @@ async fn openai_get_model_info(info: Extension<Info>) -> Json<ModelsInfo> {
     })
 }
 
+/// Template and tokenize ChatRequest
 #[utoipa::path(
     post,
     tag = "Text Generation Inference",
     path = "/chat_tokenize",
     request_body = ChatRequest,
-    responses((status = 200, description = "Templated and tokenized ChatRequest", body = ChatTokenizeResponse))
+    responses(
+    (status = 200, description = "Templated and tokenized ChatRequest", body = ChatTokenizeResponse),
+    (status = 404, description = "Failed to tokenize ChatRequest", body = ErrorResponse),
+    )
 )]
 async fn get_chat_tokenize(
     Extension(infer): Extension<Infer>,
@@ -1501,6 +1505,7 @@ tokenize,
 metrics,
 openai_get_model_info,
 sagemaker_compatibility,
+get_chat_tokenize,
 ),
 components(
 schemas(
@@ -1558,6 +1563,7 @@ Function,
 FunctionDefinition,
 ToolChoice,
 ModelInfo,
+ChatTokenizeResponse,
 )
 ),
 tags(