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}