Skip to content

Add compatibility w/ JinaAI Reranker API and Cohere Rerank API #797

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alvarobartt wants to merge 13 commits into
base: main
Choose a base branch
from
Open

Add compatibility w/ JinaAI Reranker API and Cohere Rerank API #797

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
b3e2578
Fix `disable_spans` docstring (add missing backslash)
alvarobartt Jan 7, 2026
d25c565
Run `pre-commit run --all-files` to fix `fmt`
alvarobartt Jan 7, 2026
4887145
Add `v1/rerank` w/ JinaAI Reranker API (missing errors)
alvarobartt Jan 7, 2026
d42cb5f
Add missing `prompt_tokens` in `JinaAIUsage`
alvarobartt Jan 9, 2026
4bbf38f
Add `v2/rerank` w/ Cohere Rerank API (missing errors)
alvarobartt Jan 9, 2026
0148651
Add comment on `max_tokens_per_doc` truncation
alvarobartt Jan 9, 2026
25e0950
Set `model` as not nullable and mandatory for JinaAI and Cohere
alvarobartt Jan 9, 2026
d2dcc54
Set `max_tokens_per_doc` optional
alvarobartt Jan 9, 2026
24fb02e
Truncate document if `max_tokens_per_doc` provided
alvarobartt Jan 9, 2026
53828eb
Add `CohereErrorResponse` and udpate in `v2/rerank`
alvarobartt Jan 9, 2026
ff0974f
Handle `JsonRejection` in `v2/rerank` to comply with Cohere API
alvarobartt Jan 9, 2026
1190ce2
Add `JinaAIErrorResponse` and handle validation errors
alvarobartt Jan 9, 2026
56a242d
Update `detail` for validation error
alvarobartt Jan 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

12 changes: 3 additions & 9 deletions backends/candle/tests/test_bert.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand All @@ -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);
Expand Down Expand Up @@ -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();
Expand Down
7 changes: 2 additions & 5 deletions backends/candle/tests/test_dense.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
4 changes: 1 addition & 3 deletions backends/candle/tests/test_gte.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
4 changes: 1 addition & 3 deletions backends/candle/tests/test_jina.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
8 changes: 2 additions & 6 deletions backends/candle/tests/test_modernbert.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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();
Expand Down
1 change: 1 addition & 0 deletions router/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
Loading
Loading