1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

use serde::{Deserialize, Serialize};

#[derive(Serialize, Debug, Default)]
pub struct ReRankRequest<'input> {
    /// The search query.
    pub query: &'input str,
    /// A list of document strings to rerank.
    pub documents: &'input [String],
    /// The model to use.
    pub model: ReRankModel,
    /// The number of results to return, defaults to the length of the documents.
    pub top_n: Option<u64>,
    // The maximum number of chunks to derive from each document.
    pub max_chunks_per_doc: Option<u64>,
}

#[derive(strum_macros::Display, Serialize, Debug, Default)]
pub enum ReRankModel {
    #[strum(serialize = "rerank-english-v2.0")]
    #[serde(rename = "rerank-english-v2.0")]
    #[default]
    English,
    #[strum(serialize = "rerank-multilingual-v2.0")]
    #[serde(rename = "rerank-multilingual-v2.0")]
    Multilingual,
}

#[derive(Deserialize, Debug)]
pub(crate) struct ReRankResponse {
    /// List of ranked documents
    pub results: Vec<ReRankResult>,
}

#[derive(Deserialize, Debug, PartialEq)]
pub struct ReRankResult {
    /// The index of the input document
    pub index: u64,
    /// A relevance score assigned to the ranking
    pub relevance_score: f64,
}