From 70e3f7c9a574fe6b6018596d317c6e4420b45d88 Mon Sep 17 00:00:00 2001 From: Story Crater Bot <19826264+Riotpiaole@users.noreply.github.com> Date: Sun, 23 Aug 2026 18:34:24 -0700 Subject: [PATCH] fix: wire all LLM clients to api.riotpiao.com gateway - EmbeddingsClient: use /v1/embeddings from gateway (nomic-ai/nomic-embed-text-v2-moe) - RerankClient: use /v1/rerank from gateway (BAAI/bge-reranker-base) - ChatClient: support gateway without auth headers (future: Bearer token) - All clients use LLM_API_BASE env var (default: https://api.riotpiao.com) - Ready for Kubernetes deployment with proper API routing --- crates/mem-cli/src/http_server.rs | 6 ++-- crates/mem-llm/src/chat.rs | 11 +++--- crates/mem-llm/src/embeddings.rs | 11 +++--- crates/mem-llm/src/rerank.rs | 56 +++++++++++++++++++++---------- 4 files changed, 53 insertions(+), 31 deletions(-) diff --git a/crates/mem-cli/src/http_server.rs b/crates/mem-cli/src/http_server.rs index 9ba2877..4aaad22 100644 --- a/crates/mem-cli/src/http_server.rs +++ b/crates/mem-cli/src/http_server.rs @@ -1,6 +1,6 @@ use actix_web::{web, App, HttpServer, HttpResponse, HttpRequest, middleware::Logger}; use anyhow::Result; -use mem_llm::{ChatClient, EmbeddingsClient, RerankClient}; +use mem_llm::{EmbeddingsClient, RerankClient}; use mem_store::{init_schema, VectorStore}; use serde_json::json; use sqlx::PgPool; @@ -49,9 +49,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res let vector_store = Arc::new(VectorStore::new(pool.clone())); let embeddings = Arc::new(EmbeddingsClient::from_env()?); let ingest_worker = Arc::new(IngestWorker::new(pool.clone(), (*embeddings).clone())); - - // Create a placeholder reranker (TODO: implement from_env) - let reranker = RerankClient::new("http://localhost:8000", "test", "cross-encoder")?; + let reranker = RerankClient::from_env()?; let query_worker = Arc::new(QueryWorker::new(VectorStore::new(pool.clone()), (*embeddings).clone(), reranker)); let state = web::Data::new(AppState { diff --git a/crates/mem-llm/src/chat.rs b/crates/mem-llm/src/chat.rs index 54bcb56..c72529c 100644 --- a/crates/mem-llm/src/chat.rs +++ b/crates/mem-llm/src/chat.rs @@ -144,10 +144,13 @@ impl ChatClient { let mut last_error: Option = None; for attempt in 0..self.max_retries { - let response = self - .http - .post(&url) - .header("apikey", &self.api_key) + let mut req = self.http.post(&url); + // Only add apikey header if it's not empty (for backward compatibility) + if !self.api_key.is_empty() && !self.api_key.starts_with("http") { + req = req.header("apikey", &self.api_key); + } + + let response = req .header("Content-Type", "application/json") .body(body.clone()) .timeout(self.timeout) diff --git a/crates/mem-llm/src/embeddings.rs b/crates/mem-llm/src/embeddings.rs index 8781d60..9826968 100644 --- a/crates/mem-llm/src/embeddings.rs +++ b/crates/mem-llm/src/embeddings.rs @@ -26,10 +26,11 @@ struct EmbeddingResponse { } impl EmbeddingsClient { - /// Create from environment (OLLAMA_BASE_URL, EMBEDDINGS_MODEL) + /// Create from environment + /// Uses api.riotpiao.com gateway (nomic-ai/nomic-embed-text-v2-moe model) pub fn from_env() -> Result { - let base_url = env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://ollama:11434".to_string()); - let model = env::var("EMBEDDINGS_MODEL").unwrap_or_else(|_| "nomic-embed-text-v2-moe".to_string()); + let base_url = env::var("LLM_API_BASE").unwrap_or_else(|_| "https://api.riotpiao.com".to_string()); + let model = "nomic-ai/nomic-embed-text-v2-moe".to_string(); Ok(Self { base_url, @@ -44,14 +45,14 @@ impl EmbeddingsClient { Ok(embeddings.into_iter().next().ok_or_else(|| anyhow::anyhow!("empty embedding response"))?) } - /// Embed multiple texts in a batch + /// Embed multiple texts in a batch using api.riotpiao.com gateway pub async fn embed_batch(&self, texts: &[String]) -> Result> { let req = EmbeddingRequest { model: self.model.clone(), input: texts.to_vec(), }; - let url = format!("{}/api/embed", self.base_url); + let url = format!("{}/v1/embeddings", self.base_url); let resp: EmbeddingResponse = self.http.post(&url).json(&req).send().await?.json().await?; Ok(resp diff --git a/crates/mem-llm/src/rerank.rs b/crates/mem-llm/src/rerank.rs index aec7852..0f4dcd3 100644 --- a/crates/mem-llm/src/rerank.rs +++ b/crates/mem-llm/src/rerank.rs @@ -1,33 +1,51 @@ use anyhow::Result; use reqwest::Client; -use serde_json::json; +use serde::{Deserialize, Serialize}; +use std::time::Duration; -/// Rerank response item (bare array, not OpenAI envelope). -#[derive(serde::Deserialize, Debug)] +/// Rerank score result +#[derive(Serialize, Deserialize, Debug, Clone)] pub struct RerankScore { pub index: usize, pub score: f32, } -/// Rerank client (BAAI/bge-reranker-base via TEI). +/// Rerank response from gateway +#[derive(Deserialize)] +struct RerankResponse { + results: Vec, +} + +/// Rerank client using api.riotpiao.com gateway (BAAI/bge-reranker-base model) 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 { + /// Create rerank client pointing to gateway + pub fn new(base_url: &str, _api_key: &str, model: &str) -> Result { Ok(Self { base_url: base_url.to_string(), - api_key: api_key.to_string(), model: model.to_string(), timeout_secs: 300, }) } + /// Create from environment (uses api.riotpiao.com) + pub fn from_env() -> Result { + let base_url = std::env::var("LLM_API_BASE") + .unwrap_or_else(|_| "https://api.riotpiao.com".to_string()); + let model = "BAAI/bge-reranker-base".to_string(); + + Ok(Self { + base_url, + model, + 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> { @@ -36,39 +54,41 @@ impl RerankClient { return Ok(vec![]); } - let url = format!("{}/rerank", self.base_url); + let url = format!("{}/v1/rerank", self.base_url); let client = Client::builder() - .timeout(std::time::Duration::from_secs(self.timeout_secs)) + .timeout(Duration::from_secs(self.timeout_secs)) .build()?; - let payload = json!({ + let payload = serde_json::json!({ + "model": self.model, "query": query, "texts": texts, + "top_k": texts.len(), }); 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())); + let error_text = response.text().await.unwrap_or_default(); + return Err(anyhow::anyhow!("Rerank failed: {}", error_text)); } - // Parse bare array (not OpenAI envelope) - let scores: Vec = response.json().await?; + // Parse gateway response (OpenAI format with results field) + let resp: RerankResponse = response.json().await?; - // Map back to input positions and scores - let mut results: Vec<(usize, f32)> = scores + // Map to (index, score) and sort by score descending + let mut results: Vec<(usize, f32)> = resp + .results .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)