diff --git a/crates/mem-llm/src/embeddings.rs b/crates/mem-llm/src/embeddings.rs index 9826968..a7330f5 100644 --- a/crates/mem-llm/src/embeddings.rs +++ b/crates/mem-llm/src/embeddings.rs @@ -3,13 +3,20 @@ use pgvector::Vector; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::env; +use std::time::Duration; -/// Embeddings client for Ollama +/// Embeddings dimensionality — must match schema and HNSW index +const EMBEDDINGS_DIM: usize = 768; +/// Maximum batch size for embeddings API (gateway limit: 32) +const BATCH_SIZE: usize = 32; + +/// Embeddings client for TEI/Ollama via api.riotpiao.com gateway +/// Batches requests at ≤32 texts per call, preserves input order #[derive(Clone)] pub struct EmbeddingsClient { base_url: String, model: String, - #[allow(dead_code)] + api_key: String, http: Client, } @@ -20,45 +27,132 @@ struct EmbeddingRequest { } #[derive(Debug, Deserialize)] -struct EmbeddingResponse { - embeddings: Vec>, - model: String, +#[serde(untagged)] +enum EmbeddingResponse { + Success { + #[serde(default)] + object: String, + data: Vec, + #[serde(default)] + usage: serde_json::Value, + }, + Error { + error: serde_json::Value, + }, +} + +#[derive(Debug, Deserialize)] +struct EmbeddingData { + embedding: Vec, + #[serde(default)] + index: usize, } impl EmbeddingsClient { /// Create from environment - /// Uses api.riotpiao.com gateway (nomic-ai/nomic-embed-text-v2-moe model) + /// Uses api.riotpiao.com gateway (nomic-ai/nomic-embed-text-v2-moe model, 768-dim) pub fn from_env() -> Result { - let base_url = env::var("LLM_API_BASE").unwrap_or_else(|_| "https://api.riotpiao.com".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(); + let api_key = env::var("LLM_API_KEY") + .unwrap_or_else(|_| String::new()); + + let http = Client::builder() + .timeout(Duration::from_secs(30)) + .build()?; Ok(Self { base_url, model, - http: Client::new(), + api_key, + http, }) } - /// Embed a single text string - pub async fn embed(&self, text: &str) -> Result { - let embeddings = self.embed_batch(&[text.to_string()]).await?; - Ok(embeddings.into_iter().next().ok_or_else(|| anyhow::anyhow!("empty embedding response"))?) + /// Embed a single text string, returning a 768-dim vector + pub async fn embed_one(&self, text: &str) -> Result { + let embeddings = self.embed(&[text.to_string()]).await?; + Ok(embeddings + .into_iter() + .next() + .ok_or_else(|| anyhow!("empty embedding response"))?) } - /// Embed multiple texts in a batch using api.riotpiao.com gateway - pub async fn embed_batch(&self, texts: &[String]) -> Result> { + /// Embed multiple texts, batched at ≤32 per request, preserving input order + /// Returns exactly N vectors for N input texts, each 768-dim + /// + /// **Batching:** Splits input into chunks of ≤32, processes each via POST /v1/embeddings + /// **Order:** Preserves input order across batch boundaries + /// **Assertion:** Every returned vector must be exactly 768-dim, else errors loudly with model name + /// **Headers:** Sends apikey even though route currently doesn't require auth (future-proofing) + pub async fn embed(&self, texts: &[String]) -> Result> { + if texts.is_empty() { + return Ok(Vec::new()); + } + + let mut all_vectors = Vec::new(); + + // Split into batches of ≤32 + for batch in texts.chunks(BATCH_SIZE) { + let batch_vecs = self.embed_batch_internal(batch).await?; + all_vectors.extend(batch_vecs); + } + + Ok(all_vectors) + } + + /// Internal: embed a single batch of ≤32 texts + async fn embed_batch_internal(&self, texts: &[String]) -> Result> { let req = EmbeddingRequest { model: self.model.clone(), input: texts.to_vec(), }; let url = format!("{}/v1/embeddings", self.base_url); - let resp: EmbeddingResponse = self.http.post(&url).json(&req).send().await?.json().await?; + let mut builder = self.http.post(&url); - Ok(resp - .embeddings - .into_iter() - .map(Vector::from) - .collect()) + // Send apikey header even though route currently doesn't require auth + // This future-proofs for when the route's auth plugin gets enabled + if !self.api_key.is_empty() { + builder = builder.header("apikey", &self.api_key); + } + + let resp = builder.json(&req).send().await?; + let _status = resp.status(); + let body: EmbeddingResponse = resp.json().await?; + + match body { + EmbeddingResponse::Error { error } => { + Err(anyhow!("embeddings API error: {}", error)) + } + EmbeddingResponse::Success { data, .. } => { + let mut vectors: Vec = Vec::new(); + for item in data { + // Assert exactly 768 dimensions + if item.embedding.len() != EMBEDDINGS_DIM { + return Err(anyhow!( + "model {} returned {}-dim vector, expected {}", + self.model, + item.embedding.len(), + EMBEDDINGS_DIM + )); + } + vectors.push(Vector::from(item.embedding)); + } + Ok(vectors) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_batch_size_constant() { + assert_eq!(BATCH_SIZE, 32); + assert_eq!(EMBEDDINGS_DIM, 768); } } diff --git a/tests/it_embeddings.rs b/tests/it_embeddings.rs new file mode 100644 index 0000000..92b6697 --- /dev/null +++ b/tests/it_embeddings.rs @@ -0,0 +1,262 @@ +use anyhow::Result; +use mem_llm::EmbeddingsClient; +use serde_json::json; +use wiremock::matchers::{header, method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +#[tokio::test] +async fn a1_batches_at_32() -> Result<()> { + // Test: 100 inputs produce exactly 4 requests (32+32+32+4) + let server = MockServer::start().await; + + // Mock POST /v1/embeddings to count calls and return 768-dim vectors + let mut call_count = 0; + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "object": "list", + "data": (0..32).map(|i| json!({ + "embedding": vec![0.1; 768], + "index": i + })).collect::>(), + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }))) + .mount(&server) + .await; + + let client = EmbeddingsClient::from_env(); + let mut client = client?; + client.base_url = server.uri(); + + // 100 inputs + let texts: Vec = (0..100).map(|i| format!("text {}", i)).collect(); + let result = client.embed(&texts).await?; + + assert_eq!(result.len(), 100, "Should return 100 vectors for 100 inputs"); + Ok(()) +} + +#[tokio::test] +async fn a2_order_preserved() -> Result<()> { + // Test: order is preserved across batch boundaries + // Mock returns vectors with distinct values based on input index + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(|req: &wiremock::Request| { + // Parse request body to extract input texts + let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap_or_default(); + let input = body["input"].as_array().unwrap_or(&vec![]); + + let data: Vec<_> = input + .iter() + .enumerate() + .map(|(idx, text)| { + let text_str = text.as_str().unwrap_or(""); + // Extract the number from "text N" to create distinguishable vectors + let marker = text_str + .split_whitespace() + .last() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0.0); + + json!({ + "embedding": vec![marker; 768], // Distinctive marker value + "index": idx + }) + }) + .collect(); + + ResponseTemplate::new(200).set_body_json(json!({ + "object": "list", + "data": data, + "usage": {} + })) + }) + .mount(&server) + .await; + + let mut client = EmbeddingsClient::from_env()?; + client.base_url = server.uri(); + + // 70 inputs to cross batch boundary (32 + 32 + 6) + let texts: Vec = (0..70).map(|i| format!("text {}", i)).collect(); + let result = client.embed(&texts).await?; + + // Verify order: each vector's first element should match input index + for (i, vec) in result.iter().enumerate() { + let first_val = vec.as_ref()[0]; + let expected = i as f32; + assert!( + (first_val - expected).abs() < 0.01, + "Vector {} has marker {}, expected {}", + i, + first_val, + expected + ); + } + + Ok(()) +} + +#[tokio::test] +async fn a3_dimension_asserted() -> Result<()> { + // Test: 512-dim response triggers error naming the model + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "object": "list", + "data": [{ + "embedding": vec![0.1; 512], // Wrong dimension! + "index": 0 + }], + "usage": {} + }))) + .mount(&server) + .await; + + let mut client = EmbeddingsClient::from_env()?; + client.base_url = server.uri(); + + let result = client.embed(&["test".to_string()]).await; + + assert!(result.is_err(), "Should error on dimension mismatch"); + let err_msg = format!("{:?}", result.unwrap_err()); + assert!( + err_msg.contains("768") && err_msg.contains("512"), + "Error should name both dimensions: {}", + err_msg + ); + assert!( + err_msg.contains("nomic-ai/nomic-embed-text-v2-moe"), + "Error should name the model: {}", + err_msg + ); + + Ok(()) +} + +#[tokio::test] +async fn a4_apikey_sent() -> Result<()> { + // Test: apikey header is present in request + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .and(header("apikey", wiremock::matchers::Matcher::regex(".*"))) // Match any apikey value + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "object": "list", + "data": [{ + "embedding": vec![0.1; 768], + "index": 0 + }], + "usage": {} + }))) + .mount(&server) + .await; + + let mut client = EmbeddingsClient::from_env()?; + client.base_url = server.uri(); + client.api_key = "test-key-12345".to_string(); + + let result = client.embed(&["test".to_string()]).await?; + + assert_eq!(result.len(), 1, "Should return 1 vector"); + Ok(()) +} + +#[tokio::test] +#[ignore] // Live test against real gateway +async fn a5_live_dims() -> Result<()> { + // Test: Real gateway returns 768-dim vectors + // Run with: cargo test a5_live_dims -- --ignored + let client = EmbeddingsClient::from_env()?; + + let result = client.embed(&["hello world".to_string()]).await?; + + assert_eq!(result.len(), 1, "Should return 1 vector"); + assert_eq!(result[0].as_ref().len(), 768, "Should be 768-dim"); + + Ok(()) +} + +#[tokio::test] +async fn test_empty_input() -> Result<()> { + let client = EmbeddingsClient::from_env()?; + let result = client.embed(&[]).await?; + assert_eq!(result.len(), 0, "Empty input should return empty output"); + Ok(()) +} + +#[tokio::test] +async fn test_batch_boundary_32() -> Result<()> { + // Exact boundary: 32 inputs should fit in 1 batch + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "object": "list", + "data": (0..32).map(|i| json!({ + "embedding": vec![0.1; 768], + "index": i + })).collect::>(), + "usage": {} + }))) + .mount(&server) + .await; + + let mut client = EmbeddingsClient::from_env()?; + client.base_url = server.uri(); + + let texts: Vec = (0..32).map(|i| format!("text {}", i)).collect(); + let result = client.embed(&texts).await?; + + assert_eq!(result.len(), 32, "32 inputs = 1 batch"); + Ok(()) +} + +#[tokio::test] +async fn test_batch_boundary_33() -> Result<()> { + // Over boundary: 33 inputs should need 2 batches (32+1) + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "object": "list", + "data": (0..32).map(|i| json!({ + "embedding": vec![0.1; 768], + "index": i + })).collect::>(), + "usage": {} + }))) + .mount(&server) + .await; + + // Add a separate mock for the 1-element batch + Mock::given(method("POST")) + .and(path("/v1/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "object": "list", + "data": [{ + "embedding": vec![0.1; 768], + "index": 0 + }], + "usage": {} + }))) + .mount(&server) + .await; + + let mut client = EmbeddingsClient::from_env()?; + client.base_url = server.uri(); + + let texts: Vec = (0..33).map(|i| format!("text {}", i)).collect(); + let result = client.embed(&texts).await?; + + assert_eq!(result.len(), 33, "33 inputs = 2 batches"); + Ok(()) +}