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

crates/mem-llm/Cargo.toml

@@ -13,3 +13,4 @@ anyhow = { workspace = true }
 thiserror = { workspace = true }
 reqwest = { workspace = true }
 tracing = { workspace = true }
+chrono = { workspace = true }

crates/mem-llm/src/chat.rs

@@ -0,0 +1,253 @@
+use anyhow::{anyhow, Result};
+use reqwest::Client;
+use serde::{Deserialize, Serialize};
+use std::env;
+use std::time::Duration;
+use mem_core::gated_loop::LlmClient;
+
+/// Completion response from the model.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Completion {
+    pub text: String,
+    pub usage: Usage,
+    pub latency_ms: u64,
+}
+
+/// Token usage breakdown.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Usage {
+    pub prompt_tokens: u32,
+    pub completion_tokens: u32,
+    pub total_tokens: u32,
+}
+
+/// Chat client for the gateway.
+pub struct ChatClient {
+    base_url: String,
+    api_key: String,
+    model: String,
+    http: Client,
+    timeout: Duration,
+    max_retries: u32,
+}
+
+#[derive(Debug, Serialize)]
+struct Message {
+    role: String,
+    content: String,
+}
+
+#[derive(Debug, Serialize)]
+struct CompletionRequest {
+    model: String,
+    messages: Vec<Message>,
+    max_tokens: u32,
+}
+
+#[derive(Debug, Deserialize)]
+struct CompletionResponse {
+    choices: Vec<Choice>,
+    usage: ResponseUsage,
+}
+
+#[derive(Debug, Deserialize)]
+struct Choice {
+    message: MessageResponse,
+}
+
+#[derive(Debug, Deserialize)]
+struct MessageResponse {
+    role: String,
+    content: String,
+}
+
+#[derive(Debug, Deserialize)]
+struct ResponseUsage {
+    prompt_tokens: u32,
+    completion_tokens: u32,
+    total_tokens: u32,
+}
+
+impl LlmClient for ChatClient {
+    fn complete_blocking(&self, system: &str, user: &str, max_tokens: usize) -> Result<String> {
+        ChatClient::complete_blocking(self, system, user, max_tokens as u32)
+    }
+}
+
+impl ChatClient {
+    /// Create a new chat client.
+    ///
+    /// # Arguments
+    /// * `base_url` - Gateway base URL (e.g., `https://api.riotpiao.com/v1`)
+    /// * `api_key` - Authentication key
+    /// * `model` - Model identifier (e.g., `qwen2.5:3b-instruct`)
+    pub fn new(base_url: impl Into<String>, api_key: impl Into<String>, model: impl Into<String>) -> Result<Self> {
+        Ok(Self {
+            base_url: base_url.into(),
+            api_key: api_key.into(),
+            model: model.into(),
+            http: Client::new(),
+            timeout: Duration::from_secs(300),
+            max_retries: 3,
+        })
+    }
+
+    /// Set custom timeout.
+    pub fn with_timeout(mut self, timeout: Duration) -> Self {
+        self.timeout = timeout;
+        self
+    }
+
+    /// Set max retries for 5xx errors (default: 3).
+    pub fn with_max_retries(mut self, retries: u32) -> Self {
+        self.max_retries = retries;
+        self
+    }
+
+    /// Complete synchronously (blocks until response).
+    pub fn complete_blocking(&self, system: &str, user: &str, max_tokens: u32) -> Result<String> {
+        let rt = tokio::runtime::Runtime::new()?;
+        rt.block_on(async { 
+            let completion = self.complete(system, user, max_tokens).await?;
+            Ok(completion.text)
+        })
+    }
+
+    /// Complete a prompt.
+    pub async fn complete(&self, system: &str, user: &str, max_tokens: u32) -> Result<Completion> {
+        let url = format!("{}/qwen/chat/completions", self.base_url);
+
+        let request = CompletionRequest {
+            model: self.model.clone(),
+            messages: vec![
+                Message {
+                    role: "system".to_string(),
+                    content: system.to_string(),
+                },
+                Message {
+                    role: "user".to_string(),
+                    content: user.to_string(),
+                },
+            ],
+            max_tokens,
+        };
+
+        let body = serde_json::to_string(&request)?;
+
+        // Record request if MEM_LLM_RECORD is set
+        if let Ok(record_dir) = env::var("MEM_LLM_RECORD") {
+            let filename = format!("{}/request-{}.json", record_dir, chrono::Local::now().timestamp_millis());
+            let _ = std::fs::write(&filename, &body);
+        }
+
+        let start = std::time::Instant::now();
+        let mut last_error: Option<anyhow::Error> = None;
+
+        for attempt in 0..self.max_retries {
+            let response = self
+                .http
+                .post(&url)
+                .header("apikey", &self.api_key)
+                .header("Content-Type", "application/json")
+                .body(body.clone())
+                .timeout(self.timeout)
+                .send()
+                .await;
+
+            let response = match response {
+                Ok(r) => r,
+                Err(e) => {
+                    last_error = Some(anyhow!("Request failed: {}", e));
+                    if e.is_timeout() || e.is_status() {
+                        if attempt < self.max_retries - 1 {
+                            tokio::time::sleep(Duration::from_millis(100 * 2_u64.pow(attempt))).await;
+                            continue;
+                        }
+                    }
+                    return Err(last_error.unwrap());
+                }
+            };
+
+            let status = response.status();
+            let body_text = response.text().await.unwrap_or_default();
+
+            // Record response if MEM_LLM_RECORD is set
+            if let Ok(record_dir) = env::var("MEM_LLM_RECORD") {
+                let filename = format!(
+                    "{}/response-{}-{}.json",
+                    record_dir,
+                    chrono::Local::now().timestamp_millis(),
+                    status
+                );
+                let _ = std::fs::write(&filename, &body_text);
+            }
+
+            // Handle auth error
+            if status == 401 {
+                return Err(anyhow!(
+                    "Auth error (401): check apikey header format. Response: {}",
+                    body_text
+                ));
+            }
+
+            // 4xx errors should not be retried
+            if status.is_client_error() {
+                return Err(anyhow!("Client error ({}): {}", status, body_text));
+            }
+
+            // 5xx errors should be retried
+            if status.is_server_error() {
+                if attempt < self.max_retries - 1 {
+                    last_error = Some(anyhow!("Server error ({}): {}", status, body_text));
+                    tokio::time::sleep(Duration::from_millis(100 * 2_u64.pow(attempt))).await;
+                    continue;
+                } else {
+                    return Err(anyhow!("Server error ({}): {} (after {} retries)", status, body_text, self.max_retries));
+                }
+            }
+
+            // Parse success response
+            if status.is_success() {
+                let completion_response: CompletionResponse = serde_json::from_str(&body_text)?;
+
+                if completion_response.choices.is_empty() {
+                    return Err(anyhow!("No choices in response"));
+                }
+
+                let latency_ms = start.elapsed().as_millis() as u64;
+                let text = completion_response.choices[0].message.content.clone();
+                let usage = Usage {
+                    prompt_tokens: completion_response.usage.prompt_tokens,
+                    completion_tokens: completion_response.usage.completion_tokens,
+                    total_tokens: completion_response.usage.total_tokens,
+                };
+
+                return Ok(Completion {
+                    text,
+                    usage,
+                    latency_ms,
+                });
+            }
+
+            return Err(anyhow!("Unexpected status {}: {}", status, body_text));
+        }
+
+        Err(last_error.unwrap_or_else(|| anyhow!("Max retries exhausted")))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_completion_structs_serialize() {
+        let usage = Usage {
+            prompt_tokens: 100,
+            completion_tokens: 50,
+            total_tokens: 150,
+        };
+        let json = serde_json::to_string(&usage).unwrap();
+        assert!(json.contains("100"));
+    }
+}

crates/mem-llm/src/lib.rs

@@ -1 +1,5 @@
-pub mod placeholder {}
+pub mod chat;
+pub mod rerank;
+
+pub use chat::{ChatClient, Completion, Usage};
+pub use rerank::RerankClient;

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)
+    }
+}