vv-llm 0.1.1

Rust implementation surface for vv-llm
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

use crate::{RerankResponse, RerankResult, VvLlmError};
use async_trait::async_trait;
use serde_json::{json, Value};

use super::RerankClient;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RerankMapping {
    pub method: String,
    pub path: String,
}

impl RerankMapping {
    pub fn default_siliconflow() -> Self {
        Self {
            method: "POST".to_string(),
            path: "/rerank".to_string(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct CustomJsonHttpRerankClient {
    model: String,
    api_base: String,
    api_key: String,
    mapping: RerankMapping,
    http: reqwest::Client,
}

impl CustomJsonHttpRerankClient {
    pub fn new(
        model: impl Into<String>,
        api_base: impl Into<String>,
        api_key: impl Into<String>,
        mapping: RerankMapping,
    ) -> Self {
        Self {
            model: model.into(),
            api_base: api_base.into(),
            api_key: api_key.into(),
            mapping,
            http: reqwest::Client::new(),
        }
    }

    pub fn build_request_body(&self, query: &str, documents: &[&str]) -> Result<Value, VvLlmError> {
        self.build_request_body_with_top_n(query, documents, None)
    }

    pub fn build_request_body_with_top_n(
        &self,
        query: &str,
        documents: &[&str],
        top_n: Option<usize>,
    ) -> Result<Value, VvLlmError> {
        let mut body = json!({
            "model": self.model,
            "query": query,
            "documents": documents,
        });
        if let Some(top_n) = top_n {
            body["top_n"] = json!(top_n);
        }
        Ok(body)
    }

    pub fn endpoint_url(&self) -> String {
        format!(
            "{}{}",
            self.api_base.trim_end_matches('/'),
            self.mapping.path.as_str()
        )
    }
}

#[async_trait]
impl RerankClient for CustomJsonHttpRerankClient {
    fn provider_name(&self) -> &'static str {
        "custom-json-http"
    }

    async fn rerank(&self, query: &str, documents: &[&str]) -> Result<RerankResponse, VvLlmError> {
        let response = self
            .http
            .post(self.endpoint_url())
            .bearer_auth(&self.api_key)
            .json(&self.build_request_body(query, documents)?)
            .send()
            .await
            .map_err(|error| VvLlmError::Http(error.to_string()))?
            .error_for_status()
            .map_err(|error| VvLlmError::Http(error.to_string()))?
            .json::<Value>()
            .await
            .map_err(|error| VvLlmError::Http(error.to_string()))?;

        let results = response
            .get("results")
            .and_then(Value::as_array)
            .map(|items| {
                items
                    .iter()
                    .map(|item| RerankResult {
                        index: item.get("index").and_then(Value::as_u64).unwrap_or(0) as usize,
                        relevance_score: item
                            .get("relevance_score")
                            .and_then(Value::as_f64)
                            .unwrap_or(0.0) as f32,
                    })
                    .collect()
            })
            .unwrap_or_default();

        Ok(RerankResponse { results })
    }
}