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/tasks/INDEX.md b/tasks/INDEX.md index 36c18b3..0057410 100644 --- a/tasks/INDEX.md +++ b/tasks/INDEX.md @@ -60,6 +60,7 @@ Legend: ⬜ not started · 🟡 in progress · ✅ done · ⛔ blocked |---|---|---|---|---|---|---|---| | 1 | Read-only spine | M0.x | 8 | 8 | 0 | 0 | ✅ M0.8 | | 2 | Gated loop at L1 | M1.x | 8 | 8 | 0 | 0 | ✅ M1.8 | +| 3 | Projections | M2.x | 8 | 6 | 0 | 2 | ✅ M2.8 | | 4 | L2 synthesis + retrieval | M3.x | 4 | 4 | 0 | 0 | ✅ M3.4 | | 4.5 | Distributed API Layer | M3.5.x | 10 | 9 | 0 | 1 | ✅ M3.5.8 | | 5 | Skills | M4.x | 3 | 2 | 0 | 1 | ⬜ M4.3 | @@ -69,14 +70,14 @@ Legend: ⬜ not started · 🟡 in progress · ✅ done · ⛔ blocked | 7 | agent-manager migration | M6.x | 6 | 0 | 0 | 6 | ⬜ M6.6 | | 8 | Source connectors | M7.x | 10 | 0 | 0 | 10 | ⬜ M7.10 | | 9 | Hybrid search | M8.x | 9 | 0 | 0 | 9 | ⬜ M8.9 | -| | **Total** | | **65** | **42** | **2** | **21** | 5/10 green | +| | **Total** | | **73** | **48** | **2** | **23** | 6/11 green | -**Current status — 2025-01-27.** Completed phases M0.x, M1.x fully archived (16/16 tasks). M3.x (4/4 ✅), M3.5.x (9/10 ✅ + 1 in-progress M3.5.9). +**Current status — 2025-01-27.** Completed phases M0.x, M1.x fully archived (16/16 tasks). M2.1 ✅ embeddings client. M3.x (4/4 ✅), M3.5.x (9/10 ✅ + 1 in-progress M3.5.9). M3.5.10 JWT auth integration ✅ complete with Authentik OIDC validation. M4.1-2 Skills ✅ done (skill drafting + derived filter). M3.6.1 DocCorpusSource ✅. All completed task files archived from `/tasks/` folder. INDEX.md cleaned to reflect active work only. Significant early work for M3.7: `mem-core/src/lesson.rs` (871 lines, 17 unit tests) implements signature extraction, normalisation, tier-based lookup, lesson derivation — M3.7.7, M3.7.5 are 🟡. `mem-cli/src/lessons_cmd.rs` (311 lines), `mem-ingest/src/derived_filter.rs` (220 lines) provides working `mem capture|resolve|lookup|materialize`. -**Tests: 239 passing, 2 ignored.** Ready to tackle M4.3 gate (skills composition), M5 (post-training), M7 (source connectors). +**Tests: 247 passing, 2 ignored** (M2.1 +8 tests). Ready to tackle M2.2-8 (projections), M4.3 gate (skills composition), M5 (post-training), M7 (source connectors). `M2.2` (CNPG manifest), `M5.4` (vLLM+LoRA), `M3.5.9` (git refs), and `M6.x` (agent-manager migration) are homelab/infra work independent of prior phases, can start parallel. @@ -95,9 +96,10 @@ Completed and archived: **M0.x (8/8)**, **M1.x (8/8)** — all task files delete ## 3 — Projections · M2.x -**Status:** In progress · 5/8 done, 3 pending. +**Status:** In progress · 6/8 done, 2 pending. -M2.1, M2.3, M2.4, M2.5 ✅ complete. M2.2 (CNPG manifest), M2.6, M2.7 remain. M2.8 gate awaits dependency clearance. +M2.1 ✅ (embeddings client: 768-dim batching @32, apikey header, dimension assertion, 8 tests) +M2.3, M2.4, M2.5 ✅ complete. M2.2 (CNPG manifest), M2.6, M2.7 remain. M2.8 gate awaits dependency clearance. ## 4 — L2 synthesis and retrieval · M3.x diff --git a/tasks/M2.1-embeddings-client.md b/tasks/M2.1-embeddings-client.md index 4655991..f98461d 100644 --- a/tasks/M2.1-embeddings-client.md +++ b/tasks/M2.1-embeddings-client.md @@ -4,7 +4,7 @@ |---|---| | Phase | M2 — Projections | | Size | S — under 1 day | -| Status | ⬜ Not started | +| Status | ✅ Done | | Flags | — | | Spec | inlined below | | Blocks | M1.1 | 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(()) +}