vyctor 0.1.0

A fast CLI tool for semantic file search using vector embeddings
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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

//! Voyage AI reranker implementation
//!
//! Uses the Voyage AI rerank API to reorder documents by relevance.
//! API docs: https://docs.voyageai.com/reference/reranker-api

use super::provider::{DocumentToRerank, RerankResult, Reranker};
use anyhow::{Context, Result};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};

/// Voyage AI reranker
#[derive(Debug)]
pub struct VoyageReranker {
    api_key: String,
    model: String,
    base_url: String,
    client: reqwest::Client,
}

impl VoyageReranker {
    /// Create a new Voyage reranker
    pub fn new(api_key: &str, model: &str, base_url: &str) -> Self {
        Self {
            api_key: api_key.to_string(),
            model: model.to_string(),
            base_url: base_url.trim_end_matches('/').to_string(),
            client: reqwest::Client::new(),
        }
    }
}

#[async_trait]
impl Reranker for VoyageReranker {
    fn model_name(&self) -> &str {
        &self.model
    }

    async fn rerank(
        &self,
        query: &str,
        documents: Vec<DocumentToRerank>,
    ) -> Result<Vec<RerankResult>> {
        if documents.is_empty() {
            return Ok(vec![]);
        }

        // Build the document list for the API
        let doc_contents: Vec<&str> = documents.iter().map(|d| d.content.as_str()).collect();

        let request = VoyageRerankRequest {
            model: &self.model,
            query,
            documents: &doc_contents,
            return_documents: false,
        };

        let url = format!("{}/rerank", self.base_url);

        let response = self
            .client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .context("Failed to send rerank request to Voyage AI")?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_default();
            anyhow::bail!("Voyage rerank API error ({}): {}", status, error_text);
        }

        let response: VoyageRerankResponse = response
            .json()
            .await
            .context("Failed to parse Voyage rerank response")?;

        // Map the response back to our result format, preserving original IDs
        let results = response
            .data
            .into_iter()
            .map(|d| {
                // The index in the response corresponds to the original document order
                let original_id = documents[d.index].id;
                RerankResult {
                    id: original_id,
                    relevance_score: d.relevance_score,
                }
            })
            .collect();

        Ok(results)
    }
}

/// Request payload for Voyage rerank API
#[derive(Debug, Serialize)]
struct VoyageRerankRequest<'a> {
    model: &'a str,
    query: &'a str,
    documents: &'a [&'a str],
    #[serde(skip_serializing_if = "std::ops::Not::not")]
    return_documents: bool,
}

/// Response from Voyage rerank API
#[derive(Debug, Deserialize)]
struct VoyageRerankResponse {
    data: Vec<VoyageRerankData>,
    #[allow(dead_code)]
    usage: Option<VoyageRerankUsage>,
}

/// Individual rerank result from Voyage
#[derive(Debug, Deserialize)]
struct VoyageRerankData {
    /// Index of the document in the original input
    index: usize,
    /// Relevance score (higher = more relevant)
    relevance_score: f64,
}

/// Usage statistics from Voyage rerank API
#[derive(Debug, Deserialize)]
struct VoyageRerankUsage {
    #[allow(dead_code)]
    total_tokens: usize,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_reranker_creation() {
        let reranker = VoyageReranker::new("test-key", "rerank-2.5", "https://api.voyageai.com/v1");
        assert_eq!(reranker.model_name(), "rerank-2.5");
    }
}