From cd3bd4d9e122652517335ee09b43fdf005d2196b Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Wed, 20 Mar 2024 16:30:52 +0000
Subject: [PATCH] fix: struct naming and min versions

---
 clients/python/pyproject.toml    | 2 +-
 integration-tests/pyproject.toml | 2 +-
 router/src/lib.rs                | 6 +++---
 router/src/server.rs             | 7 ++-----
 server/pyproject.toml            | 2 +-
 5 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/clients/python/pyproject.toml b/clients/python/pyproject.toml
index 4fe6e8b0..c7a885ef 100644
--- a/clients/python/pyproject.toml
+++ b/clients/python/pyproject.toml
@@ -12,7 +12,7 @@ repository = "https://github.com/huggingface/text-generation-inference"
 
 [tool.poetry.dependencies]
 python = "^3.7"
-pydantic = "> 1.10, < 3"
+pydantic = "> 2, < 3"
 aiohttp = "^3.8"
 huggingface-hub = ">= 0.12, < 1.0"
 
diff --git a/integration-tests/pyproject.toml b/integration-tests/pyproject.toml
index b0ef2435..f3599b29 100644
--- a/integration-tests/pyproject.toml
+++ b/integration-tests/pyproject.toml
@@ -5,7 +5,7 @@ description = "Text Generation Inference integration tests"
 authors = ["Nicolas Patry <nicolas@huggingface.co>"]
 
 [tool.poetry.dependencies]
-pydantic = "> 1.10, < 3"
+pydantic = "> 2, < 3"
 python = ">=3.9,<3.13"
 syrupy = "4.0.1"
 text-generation = "^0.6.0"
diff --git a/router/src/lib.rs b/router/src/lib.rs
index 7b195cb9..b231a31f 100644
--- a/router/src/lib.rs
+++ b/router/src/lib.rs
@@ -433,7 +433,7 @@ impl ChatCompletion {
         created: u64,
         details: Details,
         return_logprobs: bool,
-        tool_calls: Option<Vec<ChatCompletionMessageToolCall>>,
+        tool_calls: Option<Vec<ToolCall>>,
     ) -> Self {
         Self {
             id: String::new(),
@@ -764,7 +764,7 @@ pub(crate) struct ChatTemplateInputs<'a> {
 }
 
 #[derive(Clone, Deserialize, Serialize, ToSchema, Default, Debug)]
-pub(crate) struct ChatCompletionMessageToolCall {
+pub(crate) struct ToolCall {
     pub id: u32,
     pub r#type: String,
     pub function: FunctionDefinition,
@@ -781,7 +781,7 @@ pub(crate) struct Message {
     #[schema(example = "\"David\"")]
     pub name: Option<String>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub tool_calls: Option<Vec<ChatCompletionMessageToolCall>>,
+    pub tool_calls: Option<Vec<ToolCall>>,
 }
 
 #[derive(Clone, Debug, Deserialize, ToSchema)]
diff --git a/router/src/server.rs b/router/src/server.rs
index bcd54707..652bc699 100644
--- a/router/src/server.rs
+++ b/router/src/server.rs
@@ -14,10 +14,7 @@ use crate::{
     ChatRequest, CompatGenerateRequest, Completion, CompletionComplete, CompletionCompleteChunk,
     CompletionRequest, VertexRequest, VertexResponse,
 };
-use crate::{
-    ChatCompletionMessageToolCall, FunctionDefinition, FunctionRef, FunctionsMap, Properties,
-    ToolType, Tools,
-};
+use crate::{FunctionDefinition, FunctionRef, FunctionsMap, Properties, ToolCall, ToolType, Tools};
 use axum::extract::Extension;
 use axum::http::{HeaderMap, Method, StatusCode};
 use axum::response::sse::{Event, KeepAlive, Sse};
@@ -945,7 +942,7 @@ async fn chat_completions(
                     )
                 })?;
 
-            let tool_calls = vec![ChatCompletionMessageToolCall {
+            let tool_calls = vec![ToolCall {
                 id: 0,
                 r#type: "function".to_string(),
                 function: FunctionDefinition {
diff --git a/server/pyproject.toml b/server/pyproject.toml
index 75fb201a..764bdf95 100644
--- a/server/pyproject.toml
+++ b/server/pyproject.toml
@@ -34,7 +34,7 @@ peft = { version = "^0.9.0", optional = true }
 torch = { version = "^2.1.1", optional = true }
 scipy = "^1.11.1"
 pillow = "^10.0.0"
-outlines= { version = "0.0.36", optional = true }
+outlines= { version = "^0.0.36", optional = true }
 
 [tool.poetry.extras]
 torch = ["torch"]