From f17367e883550e7e04c60ff2bcf2d918885e6988 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Fri, 18 Apr 2025 12:48:07 +0200
Subject: [PATCH] Fixing CI

---
 router/src/sagemaker.rs | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/router/src/sagemaker.rs b/router/src/sagemaker.rs
index 750ef222b..520831c88 100644
--- a/router/src/sagemaker.rs
+++ b/router/src/sagemaker.rs
@@ -67,16 +67,26 @@ pub(crate) async fn sagemaker_compatibility(
     default_return_full_text: Extension<bool>,
     infer: Extension<Infer>,
     compute_type: Extension<ComputeType>,
+    context: Extension<Option<opentelemetry::Context>>,
     info: Extension<Info>,
     Json(req): Json<SagemakerRequest>,
 ) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
     match req {
         SagemakerRequest::Generate(req) => {
-            compat_generate(default_return_full_text, infer, compute_type, Json(req)).await
+            compat_generate(
+                default_return_full_text,
+                infer,
+                compute_type,
+                context,
+                Json(req),
+            )
+            .await
+        }
+        SagemakerRequest::Chat(req) => {
+            chat_completions(infer, compute_type, info, context, Json(req)).await
         }
-        SagemakerRequest::Chat(req) => chat_completions(infer, compute_type, info, Json(req)).await,
         SagemakerRequest::Completion(req) => {
-            completions(infer, compute_type, info, Json(req)).await
+            completions(infer, compute_type, info, context, Json(req)).await
         }
     }
 }