From 3dbc649b11d14955a3d3448ed6db373a0563dfe9 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <olivier@huggingface.co>
Date: Mon, 20 Nov 2023 10:33:44 +0100
Subject: [PATCH] fix: do not leak inputs on error (#1228)

Close #1225
---
 router/src/infer.rs      | 6 +++---
 router/src/validation.rs | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/router/src/infer.rs b/router/src/infer.rs
index cc34c466bbb..aa6dc66460d 100644
--- a/router/src/infer.rs
+++ b/router/src/infer.rs
@@ -82,7 +82,7 @@ impl Infer {
     }
 
     /// Add a new request to the queue and return a stream of InferStreamResponse
-    #[instrument(skip(self))]
+    #[instrument(skip_all)]
     pub(crate) async fn generate_stream(
         &self,
         request: GenerateRequest,
@@ -133,7 +133,7 @@ impl Infer {
     }
 
     /// Add a new request to the queue and return a InferResponse
-    #[instrument(skip(self))]
+    #[instrument(skip_all)]
     pub(crate) async fn generate(
         &self,
         request: GenerateRequest,
@@ -214,7 +214,7 @@ impl Infer {
     }
     /// Add best_of new requests to the queue and return a InferResponse of the sequence with
     /// the highest log probability per token
-    #[instrument(skip(self))]
+    #[instrument(skip(self, request))]
     pub(crate) async fn generate_best_of(
         &self,
         request: GenerateRequest,
diff --git a/router/src/validation.rs b/router/src/validation.rs
index 7a84640dcd7..1b47fc9749c 100644
--- a/router/src/validation.rs
+++ b/router/src/validation.rs
@@ -69,7 +69,7 @@ impl Validation {
         }
     }
 
-    #[instrument(skip_all)]
+    #[instrument(skip(self, inputs))]
     async fn validate_input(
         &self,
         inputs: String,