From b3917ff6957b53f2b52af0d8ef8fb3c3bf92557d Mon Sep 17 00:00:00 2001
From: David Holtz <david.richard.holtz@gmail.com>
Date: Tue, 15 Oct 2024 15:14:04 +0000
Subject: [PATCH] fix: add limit to internal stream function too

---
 router/src/server.rs | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/router/src/server.rs b/router/src/server.rs
index e82716fc..458ef554 100644
--- a/router/src/server.rs
+++ b/router/src/server.rs
@@ -478,7 +478,7 @@ async fn generate_stream(
 async fn generate_stream_internal(
     infer: Infer,
     ComputeType(compute_type): ComputeType,
-    Json(req): Json<GenerateRequest>,
+    Json(mut req): Json<GenerateRequest>,
     span: tracing::Span,
 ) -> (
     HeaderMap,
@@ -487,7 +487,11 @@ async fn generate_stream_internal(
     let start_time = Instant::now();
     metrics::counter!("tgi_request_count").increment(1);
 
-    tracing::debug!("Input: {}", req.inputs);
+    // Do not long ultra long inputs, like image payloads.
+    tracing::debug!(
+        "Input: {}",
+        &req.inputs.chars().take(1000).collect::<String>()
+    );
 
     let compute_characters = req.inputs.chars().count();
 
@@ -508,6 +512,10 @@ async fn generate_stream_internal(
         if req.parameters.return_full_text.unwrap_or(false) {
             add_prompt = Some(req.inputs.clone());
         }
+        if req.parameters.max_new_tokens.is_none() {
+            req.parameters.max_new_tokens = Some(100);
+        }
+
         let details = req.parameters.details;
 
         let best_of = req.parameters.best_of.unwrap_or(1);