From b3e25785eb5cd9f3341f4644676c2dabd940beb3 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Wed, 7 Jan 2026 15:37:39 +0100
Subject: [PATCH 01/13] Fix `disable_spans` docstring (add missing backslash)

---
 router/src/main.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/router/src/main.rs b/router/src/main.rs
index afee836e6..4f88831b4 100644
--- a/router/src/main.rs
+++ b/router/src/main.rs
@@ -166,7 +166,7 @@ struct Args {
     #[clap(long, env)]
     json_output: bool,
 
-    // Whether or not to include the log trace through spans
+    /// Whether or not to include the log trace through spans
     #[clap(long, env)]
     disable_spans: bool,
 

From d25c565d502720eb95767883cb9ee56743ecab57 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Wed, 7 Jan 2026 15:38:10 +0100
Subject: [PATCH 02/13] Run `pre-commit run --all-files` to fix `fmt`

---
 backends/candle/tests/test_bert.rs       | 12 +++---------
 backends/candle/tests/test_dense.rs      |  7 ++-----
 backends/candle/tests/test_gte.rs        |  4 +---
 backends/candle/tests/test_jina.rs       |  4 +---
 backends/candle/tests/test_modernbert.rs |  8 ++------
 5 files changed, 9 insertions(+), 26 deletions(-)

diff --git a/backends/candle/tests/test_bert.rs b/backends/candle/tests/test_bert.rs
index 2497e5a34..60285da28 100644
--- a/backends/candle/tests/test_bert.rs
+++ b/backends/candle/tests/test_bert.rs
@@ -165,9 +165,7 @@ fn test_emotions() -> Result<()> {
 
     let matcher = relative_matcher();
 
-    let predictions: Vec<Vec<f32>> = backend
-        .predict(input_batch)?.into_values()
-        .collect();
+    let predictions: Vec<Vec<f32>> = backend.predict(input_batch)?.into_values().collect();
     let predictions_batch = SnapshotScores::from(predictions);
     insta::assert_yaml_snapshot!("emotions_batch", predictions_batch, &matcher);
 
@@ -177,9 +175,7 @@ fn test_emotions() -> Result<()> {
         vec![],
     );
 
-    let predictions: Vec<Vec<f32>> = backend
-        .predict(input_single)?.into_values()
-        .collect();
+    let predictions: Vec<Vec<f32>> = backend.predict(input_single)?.into_values().collect();
     let predictions_single = SnapshotScores::from(predictions);
 
     insta::assert_yaml_snapshot!("emotions_single", predictions_single, &matcher);
@@ -217,9 +213,7 @@ fn test_bert_classification() -> Result<()> {
         vec![],
     );
 
-    let predictions: Vec<Vec<f32>> = backend
-        .predict(input_single)?.into_values()
-        .collect();
+    let predictions: Vec<Vec<f32>> = backend.predict(input_single)?.into_values().collect();
     let predictions_single = SnapshotScores::from(predictions);
 
     let matcher = relative_matcher();
diff --git a/backends/candle/tests/test_dense.rs b/backends/candle/tests/test_dense.rs
index 35ff2e318..6adc397ab 100644
--- a/backends/candle/tests/test_dense.rs
+++ b/backends/candle/tests/test_dense.rs
@@ -62,11 +62,8 @@ fn test_stella_en_400m_v5_default_dense() -> Result<()> {
 #[test]
 #[serial_test::serial]
 fn test_stella_en_400m_v5_dense_768() -> Result<()> {
-    let (model_root, dense_paths) = download_artifacts(
-        "dunzhang/stella_en_400M_v5",
-        None,
-        Some("2_Dense_768"),
-    )?;
+    let (model_root, dense_paths) =
+        download_artifacts("dunzhang/stella_en_400M_v5", None, Some("2_Dense_768"))?;
     let tokenizer = load_tokenizer(&model_root)?;
 
     let backend = CandleBackend::new(
diff --git a/backends/candle/tests/test_gte.rs b/backends/candle/tests/test_gte.rs
index 66429afef..3ae06aed3 100644
--- a/backends/candle/tests/test_gte.rs
+++ b/backends/candle/tests/test_gte.rs
@@ -164,9 +164,7 @@ fn test_gte_classification() -> Result<()> {
         vec![],
     );
 
-    let predictions: Vec<Vec<f32>> = backend
-        .predict(input_single)?.into_values()
-        .collect();
+    let predictions: Vec<Vec<f32>> = backend.predict(input_single)?.into_values().collect();
     let predictions_single = SnapshotScores::from(predictions);
 
     let matcher = relative_matcher();
diff --git a/backends/candle/tests/test_jina.rs b/backends/candle/tests/test_jina.rs
index af63041e0..9ebc39457 100644
--- a/backends/candle/tests/test_jina.rs
+++ b/backends/candle/tests/test_jina.rs
@@ -70,9 +70,7 @@ fn test_jina_rerank() -> Result<()> {
         vec![],
     );
 
-    let predictions: Vec<Vec<f32>> = backend
-        .predict(input_single)?.into_values()
-        .collect();
+    let predictions: Vec<Vec<f32>> = backend.predict(input_single)?.into_values().collect();
 
     let predictions = SnapshotScores::from(predictions);
     insta::assert_yaml_snapshot!("jinabert_reranker_single", predictions, &relative_matcher());
diff --git a/backends/candle/tests/test_modernbert.rs b/backends/candle/tests/test_modernbert.rs
index 14c65f37d..83d4a2057 100644
--- a/backends/candle/tests/test_modernbert.rs
+++ b/backends/candle/tests/test_modernbert.rs
@@ -194,9 +194,7 @@ fn test_modernbert_classification() -> Result<()> {
         vec![],
     );
 
-    let predictions: Vec<Vec<f32>> = backend
-        .predict(input_single)?.into_values()
-        .collect();
+    let predictions: Vec<Vec<f32>> = backend.predict(input_single)?.into_values().collect();
     let predictions_single = SnapshotScores::from(predictions);
 
     let matcher = relative_matcher();
@@ -231,9 +229,7 @@ fn test_modernbert_classification_mean_pooling() -> Result<()> {
         vec![],
     );
 
-    let predictions: Vec<Vec<f32>> = backend
-        .predict(input_single)?.into_values()
-        .collect();
+    let predictions: Vec<Vec<f32>> = backend.predict(input_single)?.into_values().collect();
     let predictions_single = SnapshotScores::from(predictions);
 
     let matcher = relative_matcher();

From 488714567bd5b2d8a945ce9513538f3085176153 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Wed, 7 Jan 2026 15:39:54 +0100
Subject: [PATCH 03/13] Add `v1/rerank` w/ JinaAI Reranker API (missing errors)

---
 router/src/http/server.rs | 213 +++++++++++++++++++++++++++++++++++++-
 router/src/http/types.rs  |  49 +++++++++
 2 files changed, 261 insertions(+), 1 deletion(-)

diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index 10bdef4c6..b9f964524 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -2,6 +2,7 @@
 use crate::http::types::{
     DecodeRequest, DecodeResponse, EmbedAllRequest, EmbedAllResponse, EmbedRequest, EmbedResponse,
     EmbedSparseRequest, EmbedSparseResponse, Embedding, EncodingFormat, Input, InputIds, InputType,
+    JinaAIDocument, JinaAIRerankRequest, JinaAIRerankResponse, JinaAIResult, JinaAIUsage,
     OpenAICompatEmbedding, OpenAICompatErrorResponse, OpenAICompatRequest, OpenAICompatResponse,
     OpenAICompatUsage, PredictInput, PredictRequest, PredictResponse, Prediction, Rank,
     RerankRequest, RerankResponse, Sequence, SimilarityInput, SimilarityParameters,
@@ -473,6 +474,209 @@ async fn rerank(
     Ok((headers, Json(response)))
 }
 
+#[utoipa::path(
+post,
+tag = "Text Embeddings Inference",
+path = "/v1/rerank",
+request_body = JinaAIRerankRequest,
+responses(
+(status = 200, description = "Ranks", body = JinaAIRerankResponse),
+(status = 424, description = "Rerank Error", body = ErrorResponse,
+example = json ! ({"error": "Inference failed", "error_type": "backend"})),
+(status = 429, description = "Model is overloaded", body = ErrorResponse,
+example = json ! ({"error": "Model is overloaded", "error_type": "overloaded"})),
+(status = 422, description = "Tokenization error", body = ErrorResponse,
+example = json ! ({"error": "Tokenization error", "error_type": "tokenizer"})),
+(status = 400, description = "Batch is empty", body = ErrorResponse,
+example = json ! ({"error": "Batch is empty", "error_type": "empty"})),
+(status = 413, description = "Batch size error", body = ErrorResponse,
+example = json ! ({"error": "Batch size error", "error_type": "validation"})),
+)
+)]
+#[instrument(
+    skip_all,
+    fields(total_time, tokenization_time, queue_time, inference_time,)
+)]
+async fn jinaai_rerank(
+    infer: Extension<Infer>,
+    info: Extension<Info>,
+    Extension(context): Extension<Option<opentelemetry::Context>>,
+    Json(req): Json<JinaAIRerankRequest>,
+) -> Result<(HeaderMap, Json<JinaAIRerankResponse>), (StatusCode, Json<ErrorResponse>)> {
+    let span = tracing::Span::current();
+    if let Some(context) = context {
+        span.set_parent(context);
+    }
+
+    let start_time = Instant::now();
+
+    if req.documents.is_empty() {
+        let message = "`documents` cannot be empty".to_string();
+        tracing::error!("{message}");
+        let err = ErrorResponse {
+            error: message,
+            error_type: ErrorType::Empty,
+        };
+        let counter = metrics::counter!("te_request_failure", "err" => "validation");
+        counter.increment(1);
+        Err(err)?;
+    }
+
+    match &info.model_type {
+        ModelType::Reranker(_) => Ok(()),
+        ModelType::Classifier(_) | ModelType::Embedding(_) => {
+            let counter = metrics::counter!("te_request_failure", "err" => "model_type");
+            counter.increment(1);
+            let message = "model is not a re-ranker model".to_string();
+            Err(TextEmbeddingsError::Backend(BackendError::Inference(
+                message,
+            )))
+        }
+    }
+    .map_err(|err| {
+        tracing::error!("{err}");
+        ErrorResponse::from(err)
+    })?;
+
+    let rerank_inner = move |query: String, document: String, infer: Infer| async move {
+        let permit = infer.try_acquire_permit().map_err(ErrorResponse::from)?;
+
+        let response = infer
+            .predict(
+                (query, document),
+                false,
+                TruncationDirection::Right.into(),
+                false,
+                permit,
+            )
+            .await
+            .map_err(ErrorResponse::from)?;
+
+        let score = response.results[0];
+
+        Ok::<(usize, Duration, Duration, Duration, f32), ErrorResponse>((
+            response.metadata.prompt_tokens,
+            response.metadata.tokenization,
+            response.metadata.queue,
+            response.metadata.inference,
+            score,
+        ))
+    };
+
+    let (response, metadata) = {
+        let counter = metrics::counter!("te_request_count", "method" => "batch");
+        counter.increment(1);
+
+        let batch_size = req.documents.len();
+        if batch_size > info.max_client_batch_size {
+            let message = format!(
+                "batch size {batch_size} > maximum allowed batch size {}",
+                info.max_client_batch_size
+            );
+            tracing::error!("{message}");
+            let err = ErrorResponse {
+                error: message,
+                error_type: ErrorType::Validation,
+            };
+            let counter = metrics::counter!("te_request_failure", "err" => "batch_size");
+            counter.increment(1);
+            Err(err)?;
+        }
+
+        let mut futures = Vec::with_capacity(batch_size);
+        let query_chars = req.query.chars().count();
+        let mut compute_chars = query_chars * batch_size;
+
+        for document in &req.documents {
+            compute_chars += document.chars().count();
+            let local_infer = infer.clone();
+            futures.push(rerank_inner(
+                req.query.clone(),
+                document.clone(),
+                local_infer.0,
+            ))
+        }
+        let results = join_all(futures)
+            .await
+            .into_iter()
+            .collect::<Result<Vec<(usize, Duration, Duration, Duration, f32)>, ErrorResponse>>()?;
+
+        let mut ranks = Vec::with_capacity(batch_size);
+        let mut total_tokenization_time = 0;
+        let mut total_queue_time = 0;
+        let mut total_inference_time = 0;
+        let mut total_compute_tokens = 0;
+
+        for (index, r) in results.into_iter().enumerate() {
+            total_compute_tokens += r.0;
+            total_tokenization_time += r.1.as_nanos() as u64;
+            total_queue_time += r.2.as_nanos() as u64;
+            total_inference_time += r.3.as_nanos() as u64;
+            let document = if req.return_documents {
+                Some(JinaAIDocument {
+                    text: req.documents[index].clone(),
+                })
+            } else {
+                None
+            };
+
+            let score = r.4;
+            // Check that s is not NaN or the partial_cmp below will panic
+            if score.is_nan() {
+                Err(ErrorResponse {
+                    error: "score is NaN".to_string(),
+                    error_type: ErrorType::Backend,
+                })?;
+            }
+
+            ranks.push(JinaAIResult {
+                index,
+                document,
+                relevance_score: score,
+            });
+        }
+
+        ranks.sort_by(|x, y| x.relevance_score.partial_cmp(&y.relevance_score).unwrap());
+        ranks.reverse();
+        ranks.truncate(req.top_n.map_or(ranks.len(), |x| x.min(ranks.len())));
+
+        let batch_size = batch_size as u64;
+
+        let counter = metrics::counter!("te_request_success", "method" => "batch");
+        counter.increment(1);
+
+        (
+            ranks,
+            ResponseMetadata::new(
+                compute_chars,
+                total_compute_tokens,
+                start_time,
+                Duration::from_nanos(total_tokenization_time / batch_size),
+                Duration::from_nanos(total_queue_time / batch_size),
+                Duration::from_nanos(total_inference_time / batch_size),
+            ),
+        )
+    };
+
+    metadata.record_span(&span);
+    metadata.record_metrics();
+
+    let response = JinaAIRerankResponse {
+        model: info.model_id.clone(),
+        object: "list",
+        results: response,
+        usage: JinaAIUsage {
+            total_tokens: metadata.compute_tokens,
+        },
+    };
+
+    let headers = HeaderMap::from(metadata);
+
+    tracing::info!("Success");
+
+    Ok((headers, Json(response)))
+}
+
 /// Get Sentence Similarity. Returns a 424 status code if the model is not an embedding model.
 #[utoipa::path(
 post,
@@ -1757,9 +1961,16 @@ pub async fn run(
         .route("/similarity", post(similarity))
         .route("/tokenize", post(tokenize))
         .route("/decode", post(decode))
-        // OpenAI compat route
+        // OpenAI Embeddings API compatible routes
+        // Reference: https://platform.openai.com/docs/api-reference/embeddings
         .route("/embeddings", post(openai_embed))
         .route("/v1/embeddings", post(openai_embed))
+        // JinaAI Reranker API compatible route
+        // Reference: https://jina.ai/reranker/#apiform
+        .route("/v1/rerank", post(jinaai_rerank))
+        // // Cohere Rerank API compatible route
+        // // Reference: https://docs.cohere.com/reference/rerank
+        // .route("/v2/rerank", post(cohere_rerank))
         // Vertex compat route
         .route("/vertex", post(vertex_compatibility));
 
diff --git a/router/src/http/types.rs b/router/src/http/types.rs
index a2a01c478..2e4e195a3 100644
--- a/router/src/http/types.rs
+++ b/router/src/http/types.rs
@@ -273,6 +273,55 @@ pub(crate) struct Rank {
 #[derive(Serialize, ToSchema)]
 pub(crate) struct RerankResponse(pub Vec<Rank>);
 
+#[derive(Deserialize, ToSchema)]
+pub(crate) struct JinaAIRerankRequest {
+    #[allow(dead_code)]
+    #[schema(nullable = true, example = "null")]
+    pub model: Option<String>,
+    #[schema(example = "What is Deep Learning?")]
+    pub query: String,
+    #[schema(example = json!(["Deep Learning is ..."]))]
+    pub documents: Vec<String>,
+    #[schema(example = "3", nullable = true)]
+    pub top_n: Option<usize>,
+    #[serde(default)]
+    #[schema(default = "false", example = "false")]
+    pub return_documents: bool,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct JinaAIDocument {
+    #[schema(example = "Deep Learning is ...")]
+    pub text: String,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct JinaAIResult {
+    #[schema(example = "0")]
+    pub index: usize,
+    #[schema(nullable = true, default = "null")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub document: Option<JinaAIDocument>,
+    #[schema(example = "1.0")]
+    pub relevance_score: f32,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct JinaAIUsage {
+    #[schema(example = "512")]
+    pub total_tokens: usize,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct JinaAIRerankResponse {
+    #[schema(example = "thenlper/gte-base")]
+    pub model: String,
+    #[schema(example = "list")]
+    pub object: &'static str,
+    pub usage: JinaAIUsage,
+    pub results: Vec<JinaAIResult>,
+}
+
 #[derive(Deserialize, ToSchema, Debug)]
 #[serde(untagged)]
 pub(crate) enum InputType {

From d42cb5f4b049751b48b80e6f45fb67281a8d8412 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 13:28:06 +0100
Subject: [PATCH 04/13] Add missing `prompt_tokens` in `JinaAIUsage`

---
 router/src/http/server.rs | 1 +
 router/src/http/types.rs  | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index b9f964524..4f80f8a1f 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -666,6 +666,7 @@ async fn jinaai_rerank(
         object: "list",
         results: response,
         usage: JinaAIUsage {
+            prompt_tokens: metadata.compute_tokens,
             total_tokens: metadata.compute_tokens,
         },
     };
diff --git a/router/src/http/types.rs b/router/src/http/types.rs
index 2e4e195a3..04cb7f131 100644
--- a/router/src/http/types.rs
+++ b/router/src/http/types.rs
@@ -308,6 +308,8 @@ pub(crate) struct JinaAIResult {
 
 #[derive(Serialize, ToSchema)]
 pub(crate) struct JinaAIUsage {
+    #[schema(example = "512")]
+    pub prompt_tokens: usize,
     #[schema(example = "512")]
     pub total_tokens: usize,
 }

From 4bbf38f910c96c7bad4b9c4e495ee9a5d4fd8bbc Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 14:16:15 +0100
Subject: [PATCH 05/13] Add `v2/rerank` w/ Cohere Rerank API (missing errors)

---
 Cargo.lock                |   7 +-
 router/Cargo.toml         |   1 +
 router/src/http/server.rs | 228 ++++++++++++++++++++++++++++++++++++--
 router/src/http/types.rs  |  75 +++++++++++++
 4 files changed, 299 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 58bffaba1..56c20ea95 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4639,6 +4639,7 @@ dependencies = [
  "tracing-subscriber",
  "utoipa",
  "utoipa-swagger-ui",
+ "uuid",
  "veil",
  "vergen",
 ]
@@ -5472,11 +5473,13 @@ checksum = "e2eebbbfe4093922c2b6734d7c679ebfebd704a0d7e56dfcb0d05818ce28977d"
 
 [[package]]
 name = "uuid"
-version = "1.16.0"
+version = "1.19.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "458f7a779bf54acc9f347480ac654f68407d3aab21269a6e3c9f922acd9e2da9"
+checksum = "e2e054861b4bd027cd373e18e8d8d8e6548085000e41290d95ce0c373a654b4a"
 dependencies = [
  "getrandom 0.3.2",
+ "js-sys",
+ "wasm-bindgen",
 ]
 
 [[package]]
diff --git a/router/Cargo.toml b/router/Cargo.toml
index 381d611c0..1f7a3b102 100644
--- a/router/Cargo.toml
+++ b/router/Cargo.toml
@@ -60,6 +60,7 @@ tokio-stream = { version = "0.1.14", optional = true }
 # Optional
 cudarc = { workspace = true, optional = true }
 intel-mkl-src = { workspace = true, optional = true }
+uuid = { version = "1.19.0", features = ["v4"] }
 
 # Malloc trim hack for linux
 [target.'cfg(target_os = "linux")'.dependencies]
diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index 4f80f8a1f..59ad7469a 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -1,14 +1,15 @@
 /// HTTP Server logic
 use crate::http::types::{
-    DecodeRequest, DecodeResponse, EmbedAllRequest, EmbedAllResponse, EmbedRequest, EmbedResponse,
-    EmbedSparseRequest, EmbedSparseResponse, Embedding, EncodingFormat, Input, InputIds, InputType,
-    JinaAIDocument, JinaAIRerankRequest, JinaAIRerankResponse, JinaAIResult, JinaAIUsage,
-    OpenAICompatEmbedding, OpenAICompatErrorResponse, OpenAICompatRequest, OpenAICompatResponse,
-    OpenAICompatUsage, PredictInput, PredictRequest, PredictResponse, Prediction, Rank,
-    RerankRequest, RerankResponse, Sequence, SimilarityInput, SimilarityParameters,
-    SimilarityRequest, SimilarityResponse, SimpleToken, SparseValue, TokenizeInput,
-    TokenizeRequest, TokenizeResponse, TruncationDirection, VertexPrediction, VertexRequest,
-    VertexResponse,
+    CohereApiVersion, CohereMeta, CohereMetaTokens, CohereRerankRequest, CohereRerankResponse,
+    CohereResult, DecodeRequest, DecodeResponse, EmbedAllRequest, EmbedAllResponse, EmbedRequest,
+    EmbedResponse, EmbedSparseRequest, EmbedSparseResponse, Embedding, EncodingFormat, Input,
+    InputIds, InputType, JinaAIDocument, JinaAIRerankRequest, JinaAIRerankResponse, JinaAIResult,
+    JinaAIUsage, OpenAICompatEmbedding, OpenAICompatErrorResponse, OpenAICompatRequest,
+    OpenAICompatResponse, OpenAICompatUsage, PredictInput, PredictRequest, PredictResponse,
+    Prediction, Rank, RerankRequest, RerankResponse, Sequence, SimilarityInput,
+    SimilarityParameters, SimilarityRequest, SimilarityResponse, SimpleToken, SparseValue,
+    TokenizeInput, TokenizeRequest, TokenizeResponse, TruncationDirection, VertexPrediction,
+    VertexRequest, VertexResponse,
 };
 use crate::{
     logging, shutdown, ClassifierModel, EmbeddingModel, ErrorResponse, ErrorType, Info, ModelType,
@@ -43,6 +44,7 @@ use tracing::instrument;
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use utoipa::OpenApi;
 use utoipa_swagger_ui::SwaggerUi;
+use uuid::Uuid;
 
 ///Text Embeddings Inference endpoint info
 #[utoipa::path(
@@ -678,6 +680,211 @@ async fn jinaai_rerank(
     Ok((headers, Json(response)))
 }
 
+#[utoipa::path(
+post,
+tag = "Text Embeddings Inference",
+path = "/v2/rerank",
+request_body = CohereRerankRequest,
+responses(
+(status = 200, description = "Ranks", body = CohereRerankResponse),
+(status = 424, description = "Rerank Error", body = ErrorResponse,
+example = json ! ({"error": "Inference failed", "error_type": "backend"})),
+(status = 429, description = "Model is overloaded", body = ErrorResponse,
+example = json ! ({"error": "Model is overloaded", "error_type": "overloaded"})),
+(status = 422, description = "Tokenization error", body = ErrorResponse,
+example = json ! ({"error": "Tokenization error", "error_type": "tokenizer"})),
+(status = 400, description = "Batch is empty", body = ErrorResponse,
+example = json ! ({"error": "Batch is empty", "error_type": "empty"})),
+(status = 413, description = "Batch size error", body = ErrorResponse,
+example = json ! ({"error": "Batch size error", "error_type": "validation"})),
+)
+)]
+#[instrument(
+    skip_all,
+    fields(total_time, tokenization_time, queue_time, inference_time,)
+)]
+async fn cohere_rerank(
+    infer: Extension<Infer>,
+    info: Extension<Info>,
+    Extension(context): Extension<Option<opentelemetry::Context>>,
+    Json(req): Json<CohereRerankRequest>,
+) -> Result<(HeaderMap, Json<CohereRerankResponse>), (StatusCode, Json<ErrorResponse>)> {
+    let span = tracing::Span::current();
+    if let Some(context) = context {
+        span.set_parent(context);
+    }
+
+    let start_time = Instant::now();
+
+    if req.documents.is_empty() {
+        let message = "`documents` cannot be empty".to_string();
+        tracing::error!("{message}");
+        let err = ErrorResponse {
+            error: message,
+            error_type: ErrorType::Empty,
+        };
+        let counter = metrics::counter!("te_request_failure", "err" => "validation");
+        counter.increment(1);
+        Err(err)?;
+    }
+
+    match &info.model_type {
+        ModelType::Reranker(_) => Ok(()),
+        ModelType::Classifier(_) | ModelType::Embedding(_) => {
+            let counter = metrics::counter!("te_request_failure", "err" => "model_type");
+            counter.increment(1);
+            let message = "model is not a re-ranker model".to_string();
+            Err(TextEmbeddingsError::Backend(BackendError::Inference(
+                message,
+            )))
+        }
+    }
+    .map_err(|err| {
+        tracing::error!("{err}");
+        ErrorResponse::from(err)
+    })?;
+
+    let rerank_inner = move |query: String, document: String, infer: Infer| async move {
+        let permit = infer.try_acquire_permit().map_err(ErrorResponse::from)?;
+
+        let response = infer
+            .predict(
+                (query, document),
+                false,
+                TruncationDirection::Right.into(),
+                false,
+                permit,
+            )
+            .await
+            .map_err(ErrorResponse::from)?;
+
+        let score = response.results[0];
+
+        Ok::<(usize, Duration, Duration, Duration, f32), ErrorResponse>((
+            response.metadata.prompt_tokens,
+            response.metadata.tokenization,
+            response.metadata.queue,
+            response.metadata.inference,
+            score,
+        ))
+    };
+
+    let (response, metadata) = {
+        let counter = metrics::counter!("te_request_count", "method" => "batch");
+        counter.increment(1);
+
+        let batch_size = req.documents.len();
+        if batch_size > info.max_client_batch_size {
+            let message = format!(
+                "batch size {batch_size} > maximum allowed batch size {}",
+                info.max_client_batch_size
+            );
+            tracing::error!("{message}");
+            let err = ErrorResponse {
+                error: message,
+                error_type: ErrorType::Validation,
+            };
+            let counter = metrics::counter!("te_request_failure", "err" => "batch_size");
+            counter.increment(1);
+            Err(err)?;
+        }
+
+        let mut futures = Vec::with_capacity(batch_size);
+        let query_chars = req.query.chars().count();
+        let mut compute_chars = query_chars * batch_size;
+
+        for document in &req.documents {
+            compute_chars += document.chars().count();
+            let local_infer = infer.clone();
+            futures.push(rerank_inner(
+                req.query.clone(),
+                document.clone(),
+                local_infer.0,
+            ))
+        }
+        let results = join_all(futures)
+            .await
+            .into_iter()
+            .collect::<Result<Vec<(usize, Duration, Duration, Duration, f32)>, ErrorResponse>>()?;
+
+        let mut ranks = Vec::with_capacity(batch_size);
+        let mut total_tokenization_time = 0;
+        let mut total_queue_time = 0;
+        let mut total_inference_time = 0;
+        let mut total_compute_tokens = 0;
+
+        for (index, r) in results.into_iter().enumerate() {
+            total_compute_tokens += r.0;
+            total_tokenization_time += r.1.as_nanos() as u64;
+            total_queue_time += r.2.as_nanos() as u64;
+            total_inference_time += r.3.as_nanos() as u64;
+
+            let score = r.4;
+            // Check that s is not NaN or the partial_cmp below will panic
+            if score.is_nan() {
+                Err(ErrorResponse {
+                    error: "score is NaN".to_string(),
+                    error_type: ErrorType::Backend,
+                })?;
+            }
+
+            ranks.push(CohereResult {
+                index,
+                relevance_score: score as f64,
+            });
+        }
+
+        ranks.sort_by(|x, y| x.relevance_score.partial_cmp(&y.relevance_score).unwrap());
+        ranks.reverse();
+        ranks.truncate(req.top_n.map_or(ranks.len(), |x| x.min(ranks.len())));
+
+        let batch_size = batch_size as u64;
+
+        let counter = metrics::counter!("te_request_success", "method" => "batch");
+        counter.increment(1);
+
+        (
+            ranks,
+            ResponseMetadata::new(
+                compute_chars,
+                total_compute_tokens,
+                start_time,
+                Duration::from_nanos(total_tokenization_time / batch_size),
+                Duration::from_nanos(total_queue_time / batch_size),
+                Duration::from_nanos(total_inference_time / batch_size),
+            ),
+        )
+    };
+
+    metadata.record_span(&span);
+    metadata.record_metrics();
+
+    let response = CohereRerankResponse {
+        id: Some(Uuid::new_v4().to_string()),
+        results: response,
+        meta: Some(CohereMeta {
+            api_version: Some(CohereApiVersion {
+                version: env!("CARGO_PKG_VERSION").to_string(),
+                is_deprecated: None,
+                is_experimental: None,
+            }),
+            billed_units: None,
+            tokens: Some(CohereMetaTokens {
+                input_tokens: Some(metadata.compute_tokens),
+                output_tokens: Some(metadata.compute_tokens),
+            }),
+            cached_tokens: None,
+            warnings: None,
+        }),
+    };
+
+    let headers = HeaderMap::from(metadata);
+
+    tracing::info!("Success");
+
+    Ok((headers, Json(response)))
+}
+
 /// Get Sentence Similarity. Returns a 424 status code if the model is not an embedding model.
 #[utoipa::path(
 post,
@@ -1968,10 +2175,11 @@ pub async fn run(
         .route("/v1/embeddings", post(openai_embed))
         // JinaAI Reranker API compatible route
         // Reference: https://jina.ai/reranker/#apiform
+        // Swagger Docs: https://api.jina.ai/docs#/rerank/rank_v1_rerank_post
         .route("/v1/rerank", post(jinaai_rerank))
         // // Cohere Rerank API compatible route
         // // Reference: https://docs.cohere.com/reference/rerank
-        // .route("/v2/rerank", post(cohere_rerank))
+        .route("/v2/rerank", post(cohere_rerank))
         // Vertex compat route
         .route("/vertex", post(vertex_compatibility));
 
diff --git a/router/src/http/types.rs b/router/src/http/types.rs
index 04cb7f131..6ed6bb1b2 100644
--- a/router/src/http/types.rs
+++ b/router/src/http/types.rs
@@ -324,6 +324,81 @@ pub(crate) struct JinaAIRerankResponse {
     pub results: Vec<JinaAIResult>,
 }
 
+#[derive(Deserialize, ToSchema)]
+pub(crate) struct CohereRerankRequest {
+    #[allow(dead_code)]
+    #[schema(nullable = true, example = "null")]
+    pub model: Option<String>,
+    #[schema(example = "What is Deep Learning?")]
+    pub query: String,
+    #[schema(example = json!(["Deep Learning is ..."]))]
+    pub documents: Vec<String>,
+    #[schema(example = "3", nullable = true)]
+    pub top_n: Option<usize>,
+    #[allow(dead_code)]
+    #[serde(default)]
+    #[schema(default = "4096", example = "2048")]
+    pub max_tokens_per_doc: usize,
+    // Set `skip_serializing` given that the `priority` field is used internally in the Cohere API,
+    // but in Text Embeddings Inference there's not a priority queue defined by a field when the
+    // request is sent to the server.
+    #[allow(dead_code)]
+    #[serde(default, skip_serializing)]
+    #[schema(default = "0", example = "0")]
+    pub priority: usize,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct CohereResult {
+    #[schema(example = "0")]
+    pub index: usize,
+    #[schema(example = "1.0")]
+    pub relevance_score: f64,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct CohereMetaTokens {
+    #[schema(example = "128", nullable = true)]
+    pub input_tokens: Option<usize>,
+    #[schema(example = "128", nullable = true)]
+    pub output_tokens: Option<usize>,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct CohereApiVersion {
+    pub version: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[schema(default = "false", example = "false", nullable = true)]
+    pub is_deprecated: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[schema(default = "false", example = "false", nullable = true)]
+    pub is_experimental: Option<bool>,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct CohereMeta {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub api_version: Option<CohereApiVersion>,
+    #[allow(dead_code)]
+    #[serde(skip_serializing)]
+    pub billed_units: Option<serde_json::Value>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tokens: Option<CohereMetaTokens>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub cached_tokens: Option<usize>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub warnings: Option<Vec<String>>,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct CohereRerankResponse {
+    #[schema(example = "07734bd2-2473-4f07-94e1-0d9f0e6843cf", nullable = true)]
+    pub id: Option<String>,
+    #[schema(nullable = true)]
+    pub meta: Option<CohereMeta>,
+    pub results: Vec<CohereResult>,
+}
+
 #[derive(Deserialize, ToSchema, Debug)]
 #[serde(untagged)]
 pub(crate) enum InputType {

From 01486514443954c543d9d52eb841b28802012704 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 16:59:59 +0100
Subject: [PATCH 06/13] Add comment on `max_tokens_per_doc` truncation

---
 router/src/http/server.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index 59ad7469a..6ad5052a7 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -750,7 +750,9 @@ async fn cohere_rerank(
         let response = infer
             .predict(
                 (query, document),
-                false,
+                // Cohere Rerank API expects the `max_tokens_per_doc` parameter that always
+                // truncates to 4096 by default, despite the model's max length
+                true,
                 TruncationDirection::Right.into(),
                 false,
                 permit,

From 25e095034cd8d350adcac4a064c2119a556dd601 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 18:01:39 +0100
Subject: [PATCH 07/13] Set `model` as not nullable and mandatory for JinaAI
 and Cohere

---
 router/src/http/types.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/router/src/http/types.rs b/router/src/http/types.rs
index 6ed6bb1b2..b509e86eb 100644
--- a/router/src/http/types.rs
+++ b/router/src/http/types.rs
@@ -276,8 +276,8 @@ pub(crate) struct RerankResponse(pub Vec<Rank>);
 #[derive(Deserialize, ToSchema)]
 pub(crate) struct JinaAIRerankRequest {
     #[allow(dead_code)]
-    #[schema(nullable = true, example = "null")]
-    pub model: Option<String>,
+    #[schema(example = "cross-encoder/ms-marco-MiniLM-L6-v2")]
+    pub model: String,
     #[schema(example = "What is Deep Learning?")]
     pub query: String,
     #[schema(example = json!(["Deep Learning is ..."]))]
@@ -316,7 +316,7 @@ pub(crate) struct JinaAIUsage {
 
 #[derive(Serialize, ToSchema)]
 pub(crate) struct JinaAIRerankResponse {
-    #[schema(example = "thenlper/gte-base")]
+    #[schema(example = "cross-encoder/ms-marco-MiniLM-L6-v2")]
     pub model: String,
     #[schema(example = "list")]
     pub object: &'static str,
@@ -327,8 +327,8 @@ pub(crate) struct JinaAIRerankResponse {
 #[derive(Deserialize, ToSchema)]
 pub(crate) struct CohereRerankRequest {
     #[allow(dead_code)]
-    #[schema(nullable = true, example = "null")]
-    pub model: Option<String>,
+    #[schema(example = "cross-encoder/ms-marco-MiniLM-L6-v2")]
+    pub model: String,
     #[schema(example = "What is Deep Learning?")]
     pub query: String,
     #[schema(example = json!(["Deep Learning is ..."]))]

From d2dcc54eab12c92b6bc83dbf6f8687bfd910601a Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 18:01:53 +0100
Subject: [PATCH 08/13] Set `max_tokens_per_doc` optional

---
 router/src/http/types.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/router/src/http/types.rs b/router/src/http/types.rs
index b509e86eb..f8db997b9 100644
--- a/router/src/http/types.rs
+++ b/router/src/http/types.rs
@@ -338,7 +338,7 @@ pub(crate) struct CohereRerankRequest {
     #[allow(dead_code)]
     #[serde(default)]
     #[schema(default = "4096", example = "2048")]
-    pub max_tokens_per_doc: usize,
+    pub max_tokens_per_doc: Option<usize>,
     // Set `skip_serializing` given that the `priority` field is used internally in the Cohere API,
     // but in Text Embeddings Inference there's not a priority queue defined by a field when the
     // request is sent to the server.

From 24fb02edc07b91c27b8a04d8d3898c527671be72 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 18:02:14 +0100
Subject: [PATCH 09/13] Truncate document if `max_tokens_per_doc` provided

---
 router/src/http/server.rs | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index 6ad5052a7..8cfcaf1f3 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -744,9 +744,36 @@ async fn cohere_rerank(
         ErrorResponse::from(err)
     })?;
 
-    let rerank_inner = move |query: String, document: String, infer: Infer| async move {
+    let rerank_inner = move |query: String,
+                             document: String,
+                             max_tokens_per_doc: Option<usize>,
+                             infer: Infer| async move {
         let permit = infer.try_acquire_permit().map_err(ErrorResponse::from)?;
 
+        // TODO: This is most likely not ideal given that we're tokenizing each document, then
+        // truncating it if applicable and decoding the input IDs back into a string; so as to then
+        // call predict with each query and truncated document (if applicable) combination, which
+        // might not be ideal, but seems the most straightforward approach. In any case, this
+        // should be revisited.
+        let document = if let Some(max_tokens) = max_tokens_per_doc {
+            let (_, encoding) = infer
+                .tokenize(document.clone(), false, None)
+                .await
+                .map_err(ErrorResponse::from)?;
+
+            let token_ids = encoding.get_ids();
+            if token_ids.len() > max_tokens {
+                infer
+                    .decode(token_ids[..max_tokens].to_vec(), false)
+                    .await
+                    .map_err(ErrorResponse::from)?
+            } else {
+                document
+            }
+        } else {
+            document
+        };
+
         let response = infer
             .predict(
                 (query, document),
@@ -801,6 +828,7 @@ async fn cohere_rerank(
             futures.push(rerank_inner(
                 req.query.clone(),
                 document.clone(),
+                req.max_tokens_per_doc,
                 local_infer.0,
             ))
         }

From 53828eb8818f1645bb356aa7cf1cf17a7cb0adf3 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 18:22:53 +0100
Subject: [PATCH 10/13] Add `CohereErrorResponse` and udpate in `v2/rerank`

---
 router/src/http/server.rs | 57 ++++++++++++++++++++++++---------------
 router/src/http/types.rs  |  6 +++++
 2 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index 8cfcaf1f3..e3445be13 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -1,15 +1,15 @@
 /// HTTP Server logic
 use crate::http::types::{
-    CohereApiVersion, CohereMeta, CohereMetaTokens, CohereRerankRequest, CohereRerankResponse,
-    CohereResult, DecodeRequest, DecodeResponse, EmbedAllRequest, EmbedAllResponse, EmbedRequest,
-    EmbedResponse, EmbedSparseRequest, EmbedSparseResponse, Embedding, EncodingFormat, Input,
-    InputIds, InputType, JinaAIDocument, JinaAIRerankRequest, JinaAIRerankResponse, JinaAIResult,
-    JinaAIUsage, OpenAICompatEmbedding, OpenAICompatErrorResponse, OpenAICompatRequest,
-    OpenAICompatResponse, OpenAICompatUsage, PredictInput, PredictRequest, PredictResponse,
-    Prediction, Rank, RerankRequest, RerankResponse, Sequence, SimilarityInput,
-    SimilarityParameters, SimilarityRequest, SimilarityResponse, SimpleToken, SparseValue,
-    TokenizeInput, TokenizeRequest, TokenizeResponse, TruncationDirection, VertexPrediction,
-    VertexRequest, VertexResponse,
+    CohereApiVersion, CohereErrorResponse, CohereMeta, CohereMetaTokens, CohereRerankRequest,
+    CohereRerankResponse, CohereResult, DecodeRequest, DecodeResponse, EmbedAllRequest,
+    EmbedAllResponse, EmbedRequest, EmbedResponse, EmbedSparseRequest, EmbedSparseResponse,
+    Embedding, EncodingFormat, Input, InputIds, InputType, JinaAIDocument, JinaAIRerankRequest,
+    JinaAIRerankResponse, JinaAIResult, JinaAIUsage, OpenAICompatEmbedding,
+    OpenAICompatErrorResponse, OpenAICompatRequest, OpenAICompatResponse, OpenAICompatUsage,
+    PredictInput, PredictRequest, PredictResponse, Prediction, Rank, RerankRequest, RerankResponse,
+    Sequence, SimilarityInput, SimilarityParameters, SimilarityRequest, SimilarityResponse,
+    SimpleToken, SparseValue, TokenizeInput, TokenizeRequest, TokenizeResponse,
+    TruncationDirection, VertexPrediction, VertexRequest, VertexResponse,
 };
 use crate::{
     logging, shutdown, ClassifierModel, EmbeddingModel, ErrorResponse, ErrorType, Info, ModelType,
@@ -687,16 +687,16 @@ path = "/v2/rerank",
 request_body = CohereRerankRequest,
 responses(
 (status = 200, description = "Ranks", body = CohereRerankResponse),
-(status = 424, description = "Rerank Error", body = ErrorResponse,
-example = json ! ({"error": "Inference failed", "error_type": "backend"})),
-(status = 429, description = "Model is overloaded", body = ErrorResponse,
-example = json ! ({"error": "Model is overloaded", "error_type": "overloaded"})),
-(status = 422, description = "Tokenization error", body = ErrorResponse,
-example = json ! ({"error": "Tokenization error", "error_type": "tokenizer"})),
-(status = 400, description = "Batch is empty", body = ErrorResponse,
-example = json ! ({"error": "Batch is empty", "error_type": "empty"})),
-(status = 413, description = "Batch size error", body = ErrorResponse,
-example = json ! ({"error": "Batch size error", "error_type": "validation"})),
+(status = 424, description = "Rerank Error", body = CohereErrorResponse,
+example = json ! ({"id": "12345678-1234-1234-1234-123456789abc", "message": "Inference failed"})),
+(status = 429, description = "Model is overloaded", body = CohereErrorResponse,
+example = json ! ({"id": "12345678-1234-1234-1234-123456789abc", "message": "Model is overloaded"})),
+(status = 422, description = "Tokenization error", body = CohereErrorResponse,
+example = json ! ({"id": "12345678-1234-1234-1234-123456789abc", "message": "Tokenization error"})),
+(status = 400, description = "Batch is empty", body = CohereErrorResponse,
+example = json ! ({"id": "12345678-1234-1234-1234-123456789abc", "message": "Batch is empty"})),
+(status = 413, description = "Batch size error", body = CohereErrorResponse,
+example = json ! ({"id": "12345678-1234-1234-1234-123456789abc", "message": "Batch size error"})),
 )
 )]
 #[instrument(
@@ -708,7 +708,7 @@ async fn cohere_rerank(
     info: Extension<Info>,
     Extension(context): Extension<Option<opentelemetry::Context>>,
     Json(req): Json<CohereRerankRequest>,
-) -> Result<(HeaderMap, Json<CohereRerankResponse>), (StatusCode, Json<ErrorResponse>)> {
+) -> Result<(HeaderMap, Json<CohereRerankResponse>), (StatusCode, Json<CohereErrorResponse>)> {
     let span = tracing::Span::current();
     if let Some(context) = context {
         span.set_parent(context);
@@ -2362,6 +2362,21 @@ impl From<ErrorResponse> for (StatusCode, Json<OpenAICompatErrorResponse>) {
     }
 }
 
+impl From<ErrorResponse> for CohereErrorResponse {
+    fn from(value: ErrorResponse) -> Self {
+        CohereErrorResponse {
+            id: Some(uuid::Uuid::new_v4().to_string()),
+            message: Some(value.error),
+        }
+    }
+}
+
+impl From<ErrorResponse> for (StatusCode, Json<CohereErrorResponse>) {
+    fn from(err: ErrorResponse) -> Self {
+        (StatusCode::from(&err.error_type), Json(err.into()))
+    }
+}
+
 impl From<serde_json::Error> for ErrorResponse {
     fn from(err: serde_json::Error) -> Self {
         ErrorResponse {
diff --git a/router/src/http/types.rs b/router/src/http/types.rs
index f8db997b9..358549a5f 100644
--- a/router/src/http/types.rs
+++ b/router/src/http/types.rs
@@ -399,6 +399,12 @@ pub(crate) struct CohereRerankResponse {
     pub results: Vec<CohereResult>,
 }
 
+#[derive(Serialize, ToSchema)]
+pub(crate) struct CohereErrorResponse {
+    pub id: Option<String>,
+    pub message: Option<String>,
+}
+
 #[derive(Deserialize, ToSchema, Debug)]
 #[serde(untagged)]
 pub(crate) enum InputType {

From ff0974f0ebaaf1c2b043d2e6306265ba16f6801d Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 18:47:07 +0100
Subject: [PATCH 11/13] Handle `JsonRejection` in `v2/rerank` to comply with
 Cohere API

---
 router/src/http/server.rs | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index e3445be13..0e70fdc03 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -17,7 +17,7 @@ use crate::{
 };
 use ::http::HeaderMap;
 use anyhow::Context;
-use axum::extract::{DefaultBodyLimit, Extension};
+use axum::extract::{rejection::JsonRejection, DefaultBodyLimit, Extension};
 use axum::http::HeaderValue;
 use axum::http::{Method, StatusCode};
 use axum::routing::{get, post};
@@ -707,13 +707,29 @@ async fn cohere_rerank(
     infer: Extension<Infer>,
     info: Extension<Info>,
     Extension(context): Extension<Option<opentelemetry::Context>>,
-    Json(req): Json<CohereRerankRequest>,
+    payload: Result<Json<CohereRerankRequest>, JsonRejection>,
 ) -> Result<(HeaderMap, Json<CohereRerankResponse>), (StatusCode, Json<CohereErrorResponse>)> {
     let span = tracing::Span::current();
     if let Some(context) = context {
         span.set_parent(context);
     }
 
+    // NOTE: Required to capture the JSON parsing errors, so that those follow the specification
+    // for the Cohere Rerank API, otherwise those are returned as strings
+    let req = match payload {
+        Ok(req) => req,
+        Err(rejection) => {
+            let message = format!("Invalid JSON format: {}", rejection.body_text());
+            tracing::error!(message);
+
+            let response = CohereErrorResponse {
+                id: Some(uuid::Uuid::new_v4().to_string()),
+                message: Some(message),
+            };
+            return Err((StatusCode::BAD_REQUEST, Json(response)));
+        }
+    };
+
     let start_time = Instant::now();
 
     if req.documents.is_empty() {

From 1190ce2e3d6808110d55f49087cd4ad028b8748b Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Fri, 9 Jan 2026 19:11:50 +0100
Subject: [PATCH 12/13] Add `JinaAIErrorResponse` and handle validation errors

---
 router/src/http/server.rs | 71 ++++++++++++++++++++++++++++-----------
 router/src/http/types.rs  | 15 +++++++++
 2 files changed, 66 insertions(+), 20 deletions(-)

diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index 0e70fdc03..68fd7d111 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -3,8 +3,8 @@ use crate::http::types::{
     CohereApiVersion, CohereErrorResponse, CohereMeta, CohereMetaTokens, CohereRerankRequest,
     CohereRerankResponse, CohereResult, DecodeRequest, DecodeResponse, EmbedAllRequest,
     EmbedAllResponse, EmbedRequest, EmbedResponse, EmbedSparseRequest, EmbedSparseResponse,
-    Embedding, EncodingFormat, Input, InputIds, InputType, JinaAIDocument, JinaAIRerankRequest,
-    JinaAIRerankResponse, JinaAIResult, JinaAIUsage, OpenAICompatEmbedding,
+    Embedding, EncodingFormat, Input, InputIds, InputType, JinaAIDocument, JinaAIErrorResponse,
+    JinaAIRerankRequest, JinaAIRerankResponse, JinaAIResult, JinaAIUsage, OpenAICompatEmbedding,
     OpenAICompatErrorResponse, OpenAICompatRequest, OpenAICompatResponse, OpenAICompatUsage,
     PredictInput, PredictRequest, PredictResponse, Prediction, Rank, RerankRequest, RerankResponse,
     Sequence, SimilarityInput, SimilarityParameters, SimilarityRequest, SimilarityResponse,
@@ -483,16 +483,16 @@ path = "/v1/rerank",
 request_body = JinaAIRerankRequest,
 responses(
 (status = 200, description = "Ranks", body = JinaAIRerankResponse),
-(status = 424, description = "Rerank Error", body = ErrorResponse,
-example = json ! ({"error": "Inference failed", "error_type": "backend"})),
-(status = 429, description = "Model is overloaded", body = ErrorResponse,
-example = json ! ({"error": "Model is overloaded", "error_type": "overloaded"})),
-(status = 422, description = "Tokenization error", body = ErrorResponse,
-example = json ! ({"error": "Tokenization error", "error_type": "tokenizer"})),
-(status = 400, description = "Batch is empty", body = ErrorResponse,
-example = json ! ({"error": "Batch is empty", "error_type": "empty"})),
-(status = 413, description = "Batch size error", body = ErrorResponse,
-example = json ! ({"error": "Batch size error", "error_type": "validation"})),
+(status = 424, description = "Rerank Error", body = JinaAIErrorResponse,
+example = json ! ({"detail": "[RID: 12345678-1234-1234-1234-123456789abc] Inference failed"})),
+(status = 429, description = "Model is overloaded", body = JinaAIErrorResponse,
+example = json ! ({"detail": "[RID: 12345678-1234-1234-1234-123456789abc] Model is overloaded"})),
+(status = 422, description = "Tokenization error", body = JinaAIErrorResponse,
+example = json ! ({"detail": "[RID: 12345678-1234-1234-1234-123456789abc] Tokenization error"})),
+(status = 400, description = "Batch is empty", body = JinaAIErrorResponse,
+example = json ! ({"detail": "[RID: 12345678-1234-1234-1234-123456789abc] Batch is empty"})),
+(status = 413, description = "Batch size error", body = JinaAIErrorResponse,
+example = json ! ({"detail": "[RID: 12345678-1234-1234-1234-123456789abc] Batch size error"})),
 )
 )]
 #[instrument(
@@ -503,13 +503,33 @@ async fn jinaai_rerank(
     infer: Extension<Infer>,
     info: Extension<Info>,
     Extension(context): Extension<Option<opentelemetry::Context>>,
-    Json(req): Json<JinaAIRerankRequest>,
-) -> Result<(HeaderMap, Json<JinaAIRerankResponse>), (StatusCode, Json<ErrorResponse>)> {
+    payload: Result<Json<JinaAIRerankRequest>, JsonRejection>,
+) -> Result<(HeaderMap, Json<JinaAIRerankResponse>), (StatusCode, Json<JinaAIErrorResponse>)> {
     let span = tracing::Span::current();
     if let Some(context) = context {
         span.set_parent(context);
     }
 
+    // NOTE: Required to capture the JSON parsing errors, so that those follow the specification
+    // for the JinaAI Rerank API, otherwise those are returned as strings
+    let req = match payload {
+        Ok(req) => req,
+        Err(rejection) => {
+            let message = format!("Invalid JSON format: {}", rejection.body_text());
+            tracing::error!(message);
+
+            let response = JinaAIErrorResponse {
+                detail: Some(format!(
+                    "[RID: {}] Validation error: {}",
+                    uuid::Uuid::new_v4(),
+                    message
+                )),
+                errors: None,
+            };
+            return Err((StatusCode::BAD_REQUEST, Json(response)));
+        }
+    };
+
     let start_time = Instant::now();
 
     if req.documents.is_empty() {
@@ -673,11 +693,9 @@ async fn jinaai_rerank(
         },
     };
 
-    let headers = HeaderMap::from(metadata);
-
     tracing::info!("Success");
 
-    Ok((headers, Json(response)))
+    Ok((HeaderMap::from(metadata), Json(response)))
 }
 
 #[utoipa::path(
@@ -924,11 +942,9 @@ async fn cohere_rerank(
         }),
     };
 
-    let headers = HeaderMap::from(metadata);
-
     tracing::info!("Success");
 
-    Ok((headers, Json(response)))
+    Ok((HeaderMap::from(metadata), Json(response)))
 }
 
 /// Get Sentence Similarity. Returns a 424 status code if the model is not an embedding model.
@@ -2393,6 +2409,21 @@ impl From<ErrorResponse> for (StatusCode, Json<CohereErrorResponse>) {
     }
 }
 
+impl From<ErrorResponse> for JinaAIErrorResponse {
+    fn from(value: ErrorResponse) -> Self {
+        JinaAIErrorResponse {
+            detail: Some(format!("[RID: {}] {}", uuid::Uuid::new_v4(), value.error)),
+            errors: None,
+        }
+    }
+}
+
+impl From<ErrorResponse> for (StatusCode, Json<JinaAIErrorResponse>) {
+    fn from(err: ErrorResponse) -> Self {
+        (StatusCode::from(&err.error_type), Json(err.into()))
+    }
+}
+
 impl From<serde_json::Error> for ErrorResponse {
     fn from(err: serde_json::Error) -> Self {
         ErrorResponse {
diff --git a/router/src/http/types.rs b/router/src/http/types.rs
index 358549a5f..e912ae403 100644
--- a/router/src/http/types.rs
+++ b/router/src/http/types.rs
@@ -324,6 +324,21 @@ pub(crate) struct JinaAIRerankResponse {
     pub results: Vec<JinaAIResult>,
 }
 
+#[derive(Serialize, ToSchema)]
+pub(crate) struct JinaAIError {
+    pub field: Option<String>,
+    pub message: Option<Vec<JinaAIError>>,
+    #[serde(rename(serialize = "type"))]
+    pub error_t: Option<Vec<JinaAIError>>,
+    pub input: Option<Vec<JinaAIError>>,
+}
+
+#[derive(Serialize, ToSchema)]
+pub(crate) struct JinaAIErrorResponse {
+    pub detail: Option<String>,
+    pub errors: Option<Vec<JinaAIError>>,
+}
+
 #[derive(Deserialize, ToSchema)]
 pub(crate) struct CohereRerankRequest {
     #[allow(dead_code)]

From 56a242d663446d56f3e97a05b198b1bd56c3e698 Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Sat, 10 Jan 2026 19:06:54 +0100
Subject: [PATCH 13/13] Update `detail` for validation error

---
 router/src/http/server.rs | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/router/src/http/server.rs b/router/src/http/server.rs
index 68fd7d111..b405f0307 100644
--- a/router/src/http/server.rs
+++ b/router/src/http/server.rs
@@ -519,11 +519,7 @@ async fn jinaai_rerank(
             tracing::error!(message);
 
             let response = JinaAIErrorResponse {
-                detail: Some(format!(
-                    "[RID: {}] Validation error: {}",
-                    uuid::Uuid::new_v4(),
-                    message
-                )),
+                detail: Some(format!("[RID: {}] {}", uuid::Uuid::new_v4(), message)),
                 errors: None,
             };
             return Err((StatusCode::BAD_REQUEST, Json(response)));