fix

predibase · Nov 8, 2024 · a4183eb · a4183eb
1 parent a2e14ce
commit a4183eb
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/router/src/health.rs b/router/src/health.rs
@@ -1,7 +1,10 @@
 use lorax_client::{
     Batch, NextTokenChooserParameters, Request, ShardInfo, ShardedClient,
-    StoppingCriteriaParameters,
+    StoppingCriteriaParameters, Cl
 };
+use crate::{
+  ClassifyRequest, EmbedRequest
+}
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
@@ -123,6 +126,8 @@ impl Health {
               self.generation_health.store(value, Ordering::SeqCst);
               return value
             }
+            // Return false - need to implement that shard type.
+            return false
+          }
         }
-    }
 }
diff --git a/router/src/server.rs b/router/src/server.rs
@@ -550,8 +550,7 @@ example = json ! ({"error": "unhealthy", "error_type": "healthcheck"})),
 )]
 /// Health check method
 async fn health(
-    infer: Extension<Infer>,
-    health: Extension<Health>,
+    mut health: Extension<Health>,
 ) -> Result<(), (StatusCode, Json<ErrorResponse>)> {
     match health.check().await {
       true => Ok(()),