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/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> = backend - .predict(input_batch)?.into_values() - .collect(); + let predictions: Vec> = 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> = backend - .predict(input_single)?.into_values() - .collect(); + let predictions: Vec> = 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> = backend - .predict(input_single)?.into_values() - .collect(); + let predictions: Vec> = 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> = backend - .predict(input_single)?.into_values() - .collect(); + let predictions: Vec> = 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> = backend - .predict(input_single)?.into_values() - .collect(); + let predictions: Vec> = 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> = backend - .predict(input_single)?.into_values() - .collect(); + let predictions: Vec> = 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> = backend - .predict(input_single)?.into_values() - .collect(); + let predictions: Vec> = backend.predict(input_single)?.into_values().collect(); let predictions_single = SnapshotScores::from(predictions); let matcher = relative_matcher(); 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 10bdef4c6..b405f0307 100644 --- a/router/src/http/server.rs +++ b/router/src/http/server.rs @@ -1,13 +1,15 @@ /// HTTP Server logic use crate::http::types::{ - DecodeRequest, DecodeResponse, EmbedAllRequest, EmbedAllResponse, EmbedRequest, EmbedResponse, - EmbedSparseRequest, EmbedSparseResponse, Embedding, EncodingFormat, Input, InputIds, InputType, - 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, JinaAIErrorResponse, + 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, @@ -15,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}; @@ -42,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( @@ -473,6 +476,473 @@ 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 = 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( + skip_all, + fields(total_time, tokenization_time, queue_time, inference_time,) +)] +async fn jinaai_rerank( + infer: Extension, + info: Extension, + Extension(context): Extension>, + payload: Result, JsonRejection>, +) -> Result<(HeaderMap, Json), (StatusCode, Json)> { + 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: {}] {}", uuid::Uuid::new_v4(), message)), + errors: None, + }; + return Err((StatusCode::BAD_REQUEST, Json(response))); + } + }; + + 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::, 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 { + prompt_tokens: metadata.compute_tokens, + total_tokens: metadata.compute_tokens, + }, + }; + + tracing::info!("Success"); + + Ok((HeaderMap::from(metadata), 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 = 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( + skip_all, + fields(total_time, tokenization_time, queue_time, inference_time,) +)] +async fn cohere_rerank( + infer: Extension, + info: Extension, + Extension(context): Extension>, + payload: Result, JsonRejection>, +) -> Result<(HeaderMap, Json), (StatusCode, Json)> { + 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() { + 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, + max_tokens_per_doc: Option, + 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), + // 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, + ) + .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(), + req.max_tokens_per_doc, + local_infer.0, + )) + } + let results = join_all(futures) + .await + .into_iter() + .collect::, 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, + }), + }; + + tracing::info!("Success"); + + Ok((HeaderMap::from(metadata), Json(response))) +} + /// Get Sentence Similarity. Returns a 424 status code if the model is not an embedding model. #[utoipa::path( post, @@ -1757,9 +2227,17 @@ 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 + // 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)) // Vertex compat route .route("/vertex", post(vertex_compatibility)); @@ -1912,6 +2390,36 @@ impl From for (StatusCode, Json) { } } +impl From for CohereErrorResponse { + fn from(value: ErrorResponse) -> Self { + CohereErrorResponse { + id: Some(uuid::Uuid::new_v4().to_string()), + message: Some(value.error), + } + } +} + +impl From for (StatusCode, Json) { + fn from(err: ErrorResponse) -> Self { + (StatusCode::from(&err.error_type), Json(err.into())) + } +} + +impl From for JinaAIErrorResponse { + fn from(value: ErrorResponse) -> Self { + JinaAIErrorResponse { + detail: Some(format!("[RID: {}] {}", uuid::Uuid::new_v4(), value.error)), + errors: None, + } + } +} + +impl From for (StatusCode, Json) { + fn from(err: ErrorResponse) -> Self { + (StatusCode::from(&err.error_type), Json(err.into())) + } +} + impl From 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 a2a01c478..e912ae403 100644 --- a/router/src/http/types.rs +++ b/router/src/http/types.rs @@ -273,6 +273,153 @@ pub(crate) struct Rank { #[derive(Serialize, ToSchema)] pub(crate) struct RerankResponse(pub Vec); +#[derive(Deserialize, ToSchema)] +pub(crate) struct JinaAIRerankRequest { + #[allow(dead_code)] + #[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 ..."]))] + pub documents: Vec, + #[schema(example = "3", nullable = true)] + pub top_n: Option, + #[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, + #[schema(example = "1.0")] + pub relevance_score: f32, +} + +#[derive(Serialize, ToSchema)] +pub(crate) struct JinaAIUsage { + #[schema(example = "512")] + pub prompt_tokens: usize, + #[schema(example = "512")] + pub total_tokens: usize, +} + +#[derive(Serialize, ToSchema)] +pub(crate) struct JinaAIRerankResponse { + #[schema(example = "cross-encoder/ms-marco-MiniLM-L6-v2")] + pub model: String, + #[schema(example = "list")] + pub object: &'static str, + pub usage: JinaAIUsage, + pub results: Vec, +} + +#[derive(Serialize, ToSchema)] +pub(crate) struct JinaAIError { + pub field: Option, + pub message: Option>, + #[serde(rename(serialize = "type"))] + pub error_t: Option>, + pub input: Option>, +} + +#[derive(Serialize, ToSchema)] +pub(crate) struct JinaAIErrorResponse { + pub detail: Option, + pub errors: Option>, +} + +#[derive(Deserialize, ToSchema)] +pub(crate) struct CohereRerankRequest { + #[allow(dead_code)] + #[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 ..."]))] + pub documents: Vec, + #[schema(example = "3", nullable = true)] + pub top_n: Option, + #[allow(dead_code)] + #[serde(default)] + #[schema(default = "4096", example = "2048")] + pub max_tokens_per_doc: Option, + // 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, + #[schema(example = "128", nullable = true)] + pub output_tokens: Option, +} + +#[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, + #[serde(skip_serializing_if = "Option::is_none")] + #[schema(default = "false", example = "false", nullable = true)] + pub is_experimental: Option, +} + +#[derive(Serialize, ToSchema)] +pub(crate) struct CohereMeta { + #[serde(skip_serializing_if = "Option::is_none")] + pub api_version: Option, + #[allow(dead_code)] + #[serde(skip_serializing)] + pub billed_units: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub cached_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub warnings: Option>, +} + +#[derive(Serialize, ToSchema)] +pub(crate) struct CohereRerankResponse { + #[schema(example = "07734bd2-2473-4f07-94e1-0d9f0e6843cf", nullable = true)] + pub id: Option, + #[schema(nullable = true)] + pub meta: Option, + pub results: Vec, +} + +#[derive(Serialize, ToSchema)] +pub(crate) struct CohereErrorResponse { + pub id: Option, + pub message: Option, +} + #[derive(Deserialize, ToSchema, Debug)] #[serde(untagged)] pub(crate) enum InputType { 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,