gproxy_protocol/protocol/openai/
rerank.rs1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4
5use super::common::{Extra, OpenAiModelId, OpenAiWireModel};
6
7pub 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}