zeroentropy_community/resources/models.rs
1use crate::client::Client;
2use crate::error::Result;
3use crate::types::{RerankDocument, RerankResponse};
4use serde::Serialize;
5
6/// Models resource for reranking operations
7pub struct Models<'a> {
8 client: &'a Client,
9}
10
11impl<'a> Models<'a> {
12 pub(crate) fn new(client: &'a Client) -> Self {
13 Self { client }
14 }
15
16 /// Rerank documents based on relevance to a query
17 ///
18 /// # Arguments
19 /// * `query` - The query to rank documents against
20 /// * `documents` - List of documents to rerank
21 /// * `model_id` - Optional model ID (defaults to best available model)
22 /// * `top_k` - Optional number of top results to return
23 ///
24 /// # Example
25 /// ```no_run
26 /// # use zeroentropy_community::{Client, RerankDocument};
27 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
28 /// let client = Client::from_env()?;
29 /// let documents = vec![
30 /// RerankDocument {
31 /// id: "doc1".to_string(),
32 /// text: "Rust is a systems programming language".to_string(),
33 /// },
34 /// RerankDocument {
35 /// id: "doc2".to_string(),
36 /// text: "Python is a high-level programming language".to_string(),
37 /// },
38 /// ];
39 ///
40 /// let response = client.models().rerank(
41 /// "systems programming",
42 /// documents,
43 /// None,
44 /// None,
45 /// ).await?;
46 /// # Ok(())
47 /// # }
48 /// ```
49 pub async fn rerank(
50 &self,
51 query: impl Into<String>,
52 documents: Vec<RerankDocument>,
53 model_id: Option<String>,
54 top_k: Option<u32>,
55 ) -> Result<RerankResponse> {
56 #[derive(Serialize)]
57 struct Request {
58 query: String,
59 documents: Vec<RerankDocument>,
60 #[serde(skip_serializing_if = "Option::is_none")]
61 model_id: Option<String>,
62 #[serde(skip_serializing_if = "Option::is_none")]
63 top_k: Option<u32>,
64 }
65
66 let body = Request {
67 query: query.into(),
68 documents,
69 model_id,
70 top_k,
71 };
72
73 self.client.post("/models/rerank", &body).await
74 }
75}