Deploy Poimen Memory K8s cluster with ArgoCD tracking (M2.2, M3.5-M3.7)

crates/mem-llm/src/rerank.rs

@@ -0,0 +1,76 @@
+use anyhow::Result;
+use reqwest::Client;
+use serde_json::json;
+
+/// Rerank response item (bare array, not OpenAI envelope).
+#[derive(serde::Deserialize, Debug)]
+pub struct RerankScore {
+    pub index: usize,
+    pub score: f32,
+}
+
+/// Rerank client (BAAI/bge-reranker-base via TEI).
+pub struct RerankClient {
+    base_url: String,
+    api_key: String,
+    model: String,
+    timeout_secs: u64,
+}
+
+impl RerankClient {
+    /// Create rerank client.
+    pub fn new(base_url: &str, api_key: &str, model: &str) -> Result<Self> {
+        Ok(Self {
+            base_url: base_url.to_string(),
+            api_key: api_key.to_string(),
+            model: model.to_string(),
+            timeout_secs: 300,
+        })
+    }
+
+    /// Rerank query against texts, return scored items in score order.
+    /// Returns Vec<(index, score)> mapping back to input positions.
+    pub async fn rerank(&self, query: &str, texts: &[&str]) -> Result<Vec<(usize, f32)>> {
+        // Empty input returns empty without request
+        if texts.is_empty() {
+            return Ok(vec![]);
+        }
+
+        let url = format!("{}/rerank", self.base_url);
+
+        let client = Client::builder()
+            .timeout(std::time::Duration::from_secs(self.timeout_secs))
+            .build()?;
+
+        let payload = json!({
+            "query": query,
+            "texts": texts,
+        });
+
+        let response = client
+            .post(&url)
+            .header("apikey", &self.api_key)
+            .header("Content-Type", "application/json")
+            .json(&payload)
+            .send()
+            .await?;
+
+        if !response.status().is_success() {
+            return Err(anyhow::anyhow!("Rerank failed: {}", response.status()));
+        }
+
+        // Parse bare array (not OpenAI envelope)
+        let scores: Vec<RerankScore> = response.json().await?;
+
+        // Map back to input positions and scores
+        let mut results: Vec<(usize, f32)> = scores
+            .into_iter()
+            .map(|s| (s.index, s.score))
+            .collect();
+
+        // Sort by score descending (highest first)
+        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+
+        Ok(results)
+    }
+}