ferrin_spec/
reranking_model.rs

1//! Reranking model interface.
2
3use std::future::Future;
4
5use serde::Deserialize;
6use serde::Serialize;
7use tokio_util::sync::CancellationToken;
8
9use crate::error::ProviderError;
10use crate::json::JsonObject;
11use crate::language_model::ResponseMetadata;
12use crate::shared::Headers;
13use crate::shared::ModelId;
14use crate::shared::ProviderId;
15use crate::shared::ProviderMetadata;
16use crate::shared::ProviderOptions;
17use crate::shared::Warning;
18
19/// A model that orders documents by relevance to a query.
20pub trait RerankingModel: Send + Sync + 'static {
21    /// Provider identifier.
22    fn provider(&self) -> &ProviderId;
23
24    /// Model identifier.
25    fn model_id(&self) -> &ModelId;
26
27    /// Ranks `options.documents` against `options.query`.
28    fn do_rerank(
29        &self,
30        options: RerankOptions,
31    ) -> impl Future<Output = Result<RerankResult, ProviderError>> + Send;
32}
33
34/// Documents to rank: all strings or all JSON objects.
35#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
36#[serde(tag = "type", rename_all = "lowercase")]
37#[non_exhaustive]
38pub enum RerankDocuments {
39    /// Plain text documents.
40    Text {
41        /// The documents.
42        values: Vec<String>,
43    },
44    /// Structured documents.
45    Object {
46        /// The documents.
47        values: Vec<JsonObject>,
48    },
49}
50
51impl RerankDocuments {
52    /// Number of documents.
53    #[must_use]
54    pub fn len(&self) -> usize {
55        match self {
56            Self::Text { values } => values.len(),
57            Self::Object { values } => values.len(),
58        }
59    }
60
61    /// Returns `true` when there are no documents.
62    #[must_use]
63    pub fn is_empty(&self) -> bool {
64        self.len() == 0
65    }
66}
67
68/// Options for a rerank call.
69#[derive(Debug, Clone)]
70pub struct RerankOptions {
71    /// The query.
72    pub query: String,
73    /// Documents to rank.
74    pub documents: RerankDocuments,
75    /// Return only the top `n` results.
76    pub top_n: Option<usize>,
77    /// Provider-specific options keyed by provider name.
78    pub provider_options: ProviderOptions,
79    /// Additional request headers.
80    pub headers: Headers,
81    /// Cancellation token.
82    pub cancellation: CancellationToken,
83}
84
85impl RerankOptions {
86    /// Creates options for `query` over `documents`.
87    #[must_use]
88    pub fn new(query: impl Into<String>, documents: RerankDocuments) -> Self {
89        Self {
90            query: query.into(),
91            documents,
92            top_n: None,
93            provider_options: ProviderOptions::new(),
94            headers: Headers::new(),
95            cancellation: CancellationToken::new(),
96        }
97    }
98}
99
100/// One entry of a ranking.
101#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
102pub struct RankedDocument {
103    /// Index of the document in the input list.
104    pub index: usize,
105    /// Relevance score (higher is more relevant).
106    pub relevance_score: f64,
107}
108
109/// Result of a rerank call.
110#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
111pub struct RerankResult {
112    /// Ranking, most relevant first.
113    pub ranking: Vec<RankedDocument>,
114    /// Provider-specific metadata.
115    #[serde(default, skip_serializing_if = "Option::is_none")]
116    pub provider_metadata: Option<ProviderMetadata>,
117    /// Warnings.
118    #[serde(default)]
119    pub warnings: Vec<Warning>,
120    /// Response metadata.
121    #[serde(default)]
122    pub response: ResponseMetadata,
123}
ferrin_spec/reranking_model.rs

ferrin_spec/
reranking_model.rs
