From 010508cec8e60e3995f4d0a0f98f30c499c28b59 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <olivier@huggingface.co>
Date: Wed, 21 Feb 2024 15:30:45 +0100
Subject: [PATCH] fix: fix openapi schema (#1586)

---
 docs/openapi.json    | 117 +++++++++++++++++++++++++++++++++++++++----
 router/src/lib.rs    |   2 +-
 router/src/server.rs |  12 +++--
 3 files changed, 118 insertions(+), 13 deletions(-)

diff --git a/docs/openapi.json b/docs/openapi.json
index 89ffa0d7a..f8e52a8db 100644
--- a/docs/openapi.json
+++ b/docs/openapi.json
@@ -637,6 +637,35 @@
           }
         }
       },
+      "ChatCompletionComplete": {
+        "type": "object",
+        "required": [
+          "index",
+          "message",
+          "finish_reason"
+        ],
+        "properties": {
+          "finish_reason": {
+            "type": "string"
+          },
+          "index": {
+            "type": "integer",
+            "format": "int32",
+            "minimum": 0
+          },
+          "logprobs": {
+            "allOf": [
+              {
+                "$ref": "#/components/schemas/ChatCompletionLogprobs"
+              }
+            ],
+            "nullable": true
+          },
+          "message": {
+            "$ref": "#/components/schemas/Message"
+          }
+        }
+      },
       "ChatCompletionDelta": {
         "type": "object",
         "required": [
@@ -654,6 +683,59 @@
           }
         }
       },
+      "ChatCompletionLogprob": {
+        "type": "object",
+        "required": [
+          "token",
+          "logprob",
+          "top_logprobs"
+        ],
+        "properties": {
+          "logprob": {
+            "type": "number",
+            "format": "float"
+          },
+          "token": {
+            "type": "string"
+          },
+          "top_logprobs": {
+            "type": "array",
+            "items": {
+              "$ref": "#/components/schemas/ChatCompletionTopLogprob"
+            }
+          }
+        }
+      },
+      "ChatCompletionLogprobs": {
+        "type": "object",
+        "required": [
+          "content"
+        ],
+        "properties": {
+          "content": {
+            "type": "array",
+            "items": {
+              "$ref": "#/components/schemas/ChatCompletionLogprob"
+            }
+          }
+        }
+      },
+      "ChatCompletionTopLogprob": {
+        "type": "object",
+        "required": [
+          "token",
+          "logprob"
+        ],
+        "properties": {
+          "logprob": {
+            "type": "number",
+            "format": "float"
+          },
+          "token": {
+            "type": "string"
+          }
+        }
+      },
       "ChatRequest": {
         "type": "object",
         "required": [
@@ -1038,15 +1120,7 @@
                 ]
               },
               "value": {
-                "type": "string",
-                "description": "A string that represents a [JSON Schema](https://json-schema.org/).\n\nJSON Schema is a declarative language that allows to annotate JSON documents\nwith types and descriptions.",
-                "example": {
-                  "properties": {
-                    "location": {
-                      "type": "string"
-                    }
-                  }
-                }
+                "description": "A string that represents a [JSON Schema](https://json-schema.org/).\n\nJSON Schema is a declarative language that allows to annotate JSON documents\nwith types and descriptions."
               }
             }
           },
@@ -1362,6 +1436,31 @@
         "items": {
           "$ref": "#/components/schemas/SimpleToken"
         }
+      },
+      "Usage": {
+        "type": "object",
+        "required": [
+          "prompt_tokens",
+          "completion_tokens",
+          "total_tokens"
+        ],
+        "properties": {
+          "completion_tokens": {
+            "type": "integer",
+            "format": "int32",
+            "minimum": 0
+          },
+          "prompt_tokens": {
+            "type": "integer",
+            "format": "int32",
+            "minimum": 0
+          },
+          "total_tokens": {
+            "type": "integer",
+            "format": "int32",
+            "minimum": 0
+          }
+        }
       }
     }
   },
diff --git a/router/src/lib.rs b/router/src/lib.rs
index c6928a5a9..113e26420 100644
--- a/router/src/lib.rs
+++ b/router/src/lib.rs
@@ -347,7 +347,7 @@ pub(crate) struct ChatCompletionTopLogprob {
     logprob: f32,
 }
 
-#[derive(Clone, Deserialize, Serialize)]
+#[derive(Clone, Deserialize, Serialize, ToSchema)]
 pub(crate) struct Usage {
     pub prompt_tokens: u32,
     pub completion_tokens: u32,
diff --git a/router/src/server.rs b/router/src/server.rs
index ebde7133f..9fdd66cc3 100644
--- a/router/src/server.rs
+++ b/router/src/server.rs
@@ -3,11 +3,12 @@ use crate::health::Health;
 use crate::infer::{InferError, InferResponse, InferStreamResponse};
 use crate::validation::ValidationError;
 use crate::{
-    BestOfSequence, ChatCompletion, ChatCompletionChoice, ChatCompletionChunk, ChatCompletionDelta,
-    ChatCompletionLogprobs, ChatRequest, CompatGenerateRequest, Details, ErrorResponse,
+    BestOfSequence, ChatCompletion, ChatCompletionChoice, ChatCompletionChunk,
+    ChatCompletionComplete, ChatCompletionDelta, ChatCompletionLogprob, ChatCompletionLogprobs,
+    ChatCompletionTopLogprob, ChatRequest, CompatGenerateRequest, Details, ErrorResponse,
     FinishReason, GenerateParameters, GenerateRequest, GenerateResponse, GrammarType, HubModelInfo,
     HubTokenizerConfig, Infer, Info, Message, PrefillToken, SimpleToken, StreamDetails,
-    StreamResponse, Token, TokenizeResponse, Validation, VertexRequest, VertexResponse,
+    StreamResponse, Token, TokenizeResponse, Usage, Validation, VertexRequest, VertexResponse,
 };
 use axum::extract::Extension;
 use axum::http::{HeaderMap, Method, StatusCode};
@@ -896,9 +897,13 @@ pub async fn run(
     GrammarType,
     ChatRequest,
     Message,
+    ChatCompletionComplete,
     ChatCompletionChoice,
     ChatCompletionDelta,
     ChatCompletionChunk,
+    ChatCompletionLogprob,
+    ChatCompletionLogprobs,
+    ChatCompletionTopLogprob,
     ChatCompletion,
     GenerateParameters,
     PrefillToken,
@@ -913,6 +918,7 @@ pub async fn run(
     StreamDetails,
     ErrorResponse,
     GrammarType,
+    Usage,
     )
     ),
     tags(