Skip to main content

gproxy_protocol/protocol/openai/
rerank.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4
5use super::common::{Extra, OpenAiModelId, OpenAiWireModel};
6
7/// OpenAI-shaped rerank wire model used by compatible providers such as
8/// OpenRouter and DashScope. OpenAI itself does not expose this endpoint.
9pub type RerankWireModel = OpenAiWireModel<RerankRequest, RerankResponse>;
10
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
12#[non_exhaustive]
13pub struct RerankRequest {
14    pub model: OpenAiModelId,
15    pub query: String,
16    pub documents: Vec<RerankDocument>,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub top_n: Option<u32>,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub instruct: Option<String>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub return_documents: Option<bool>,
23    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
24    pub extra: Extra,
25}
26
27#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
28#[serde(untagged)]
29#[non_exhaustive]
30pub enum RerankDocument {
31    Text(String),
32    Structured(RerankDocumentContent),
33}
34
35#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
36#[non_exhaustive]
37pub struct RerankDocumentContent {
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub text: Option<String>,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub image: Option<String>,
42    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
43    pub extra: Extra,
44}
45
46#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
47#[non_exhaustive]
48pub struct RerankResponse {
49    pub model: OpenAiModelId,
50    pub results: Vec<RerankResult>,
51    #[serde(skip_serializing_if = "Option::is_none")]
52    pub id: Option<String>,
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub object: Option<String>,
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub provider: Option<String>,
57    #[serde(skip_serializing_if = "Option::is_none")]
58    pub usage: Option<RerankUsage>,
59    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
60    pub extra: Extra,
61}
62
63#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
64#[non_exhaustive]
65pub struct RerankResult {
66    pub index: u32,
67    pub relevance_score: f64,
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub document: Option<RerankDocumentContent>,
70    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
71    pub extra: Extra,
72}
73
74#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, gproxy_protocol_macros::WireBuilder)]
75#[non_exhaustive]
76pub struct RerankUsage {
77    #[serde(skip_serializing_if = "Option::is_none")]
78    pub total_tokens: Option<u64>,
79    #[serde(skip_serializing_if = "Option::is_none")]
80    pub search_units: Option<u64>,
81    #[serde(skip_serializing_if = "Option::is_none")]
82    pub cost: Option<f64>,
83    #[serde(default, flatten, skip_serializing_if = "BTreeMap::is_empty")]
84    pub extra: Extra,
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn accepts_text_and_multimodal_documents() {
93        let request: RerankRequest = serde_json::from_value(serde_json::json!({
94            "model": "qwen3-vl-rerank",
95            "query": "the red car",
96            "documents": ["plain text", {"text": "a car", "image": "https://x/car.png"}],
97            "top_n": 1,
98            "provider": {"order": ["Alibaba"]}
99        }))
100        .unwrap();
101        assert!(matches!(request.documents[0], RerankDocument::Text(_)));
102        assert!(matches!(
103            request.documents[1],
104            RerankDocument::Structured(_)
105        ));
106        assert!(request.extra.contains_key("provider"));
107    }
108}