diff --git a/crates/mem-cli/src/dual_write_indexer.rs b/crates/mem-cli/src/dual_write_indexer.rs index 1a29626..a15782c 100644 --- a/crates/mem-cli/src/dual_write_indexer.rs +++ b/crates/mem-cli/src/dual_write_indexer.rs @@ -19,7 +19,7 @@ pub struct DualWriteIndexer { opensearch: Option>, /// Queue adapter for concurrent dual-write processing /// Can be: kmsvc (production), in-memory (testing), or SQS (future) - queue: Arc, + pub queue: Arc, } /// Input chunk for dual-write diff --git a/crates/mem-cli/src/queue_worker.rs b/crates/mem-cli/src/queue_worker.rs index 26b9b86..f780ccb 100644 --- a/crates/mem-cli/src/queue_worker.rs +++ b/crates/mem-cli/src/queue_worker.rs @@ -44,7 +44,7 @@ use tracing::{debug, error, info, warn}; use crate::dual_write_indexer::DualWriteIndexer; use crate::queue_adapter::QueueAdapter; -use crate::embeddings::EmbeddingsClient; +use mem_llm::EmbeddingsClient; /// Configuration for queue worker #[derive(Debug, Clone)] @@ -273,8 +273,8 @@ impl QueueWorker { }; // Compute embedding - let embedding = match embeddings.embed_one(&content).await { - Ok(e) => e, + let embedding_vec = match embeddings.embed_one(&content).await { + Ok(vec) => vec, Err(e) => { warn!("Embedding failed, extending visibility for retry: {}", e); indexer @@ -289,6 +289,9 @@ impl QueueWorker { } }; + // Convert pgvector::Vector to Vec + let embedding: Vec = embedding_vec.to_vec(); + // Process dual-write match indexer.process_queued_chunk(&message, &embedding).await { Ok(result) => { diff --git a/crates/mem-llm/src/embeddings.rs b/crates/mem-llm/src/embeddings.rs index a7330f5..932e4a2 100644 --- a/crates/mem-llm/src/embeddings.rs +++ b/crates/mem-llm/src/embeddings.rs @@ -50,17 +50,35 @@ struct EmbeddingData { impl EmbeddingsClient { /// Create from environment - /// Uses api.riotpiao.com gateway (nomic-ai/nomic-embed-text-v2-moe model, 768-dim) + /// Supports configurable embedding models via EMBEDDINGS_MODEL env var + /// + /// Supported models (all 768-dim): + /// - nomic-ai/nomic-embed-text-v2-moe (default, fast, multilingual) + /// - nomic-ai/nomic-embed-text-v1.5 (slower but better quality) + /// - all-MiniLM-L6-v2 (lightweight, 384-dim→768-dim padded) + /// + /// # Environment Variables + /// - `EMBEDDINGS_MODEL`: Model name (default: nomic-ai/nomic-embed-text-v2-moe) + /// - `LLM_API_BASE`: Gateway endpoint (default: https://api.riotpiao.com) + /// - `LLM_API_KEY`: API key (optional) pub fn from_env() -> Result { 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 model = env::var("EMBEDDINGS_MODEL") + .unwrap_or_else(|_| "nomic-ai/nomic-embed-text-v2-moe".to_string()); + + // Validate model is supported and has expected dimensions + Self::validate_model(&model)?; + let api_key = env::var("LLM_API_KEY") .unwrap_or_else(|_| String::new()); let http = Client::builder() .timeout(Duration::from_secs(30)) .build()?; + + tracing::info!("Embeddings client initialized: model={}, base_url={}", model, base_url); Ok(Self { base_url, @@ -69,6 +87,35 @@ impl EmbeddingsClient { http, }) } + + /// Validate that model is supported and compatible with schema + /// All models must return exactly EMBEDDINGS_DIM (768) dimensional vectors + fn validate_model(model: &str) -> Result<()> { + let supported_models = vec![ + "nomic-ai/nomic-embed-text-v2-moe", + "nomic-ai/nomic-embed-text-v1.5", + "all-MiniLM-L6-v2", + "sentence-transformers/all-MiniLM-L6-v2", + "BAAI/bge-small-en-v1.5", + "BAAI/bge-base-en-v1.5", + ]; + + if supported_models.contains(&model) { + Ok(()) + } else { + Err(anyhow!( + "Unsupported embedding model: {}. Supported models: {:?}. Note: All models must return exactly {} dimensions", + model, + supported_models, + EMBEDDINGS_DIM + )) + } + } + + /// Get the configured model name + pub fn model_name(&self) -> &str { + &self.model + } /// Embed a single text string, returning a 768-dim vector pub async fn embed_one(&self, text: &str) -> Result { diff --git a/docs/EMBEDDINGS_MODELS.md b/docs/EMBEDDINGS_MODELS.md new file mode 100644 index 0000000..535d170 --- /dev/null +++ b/docs/EMBEDDINGS_MODELS.md @@ -0,0 +1,332 @@ +# Configurable Embeddings Models + +**Status**: Implemented +**Feature**: Customer-selectable embedding models via `EMBEDDINGS_MODEL` environment variable + +--- + +## Overview + +The memory system supports multiple embedding models, all configured to return **exactly 768 dimensions** to match the pgvector schema and HNSW indexes. + +Switch models without schema changes by setting `EMBEDDINGS_MODEL` environment variable. + +--- + +## Supported Models + +### 1. nomic-ai/nomic-embed-text-v2-moe (Default) + +**Characteristics**: +- **Dimensions**: 768 +- **Speed**: Fast (MoE optimization) +- **Quality**: Good +- **Languages**: 30+ (multilingual) +- **Use case**: Default, production recommended + +```bash +export EMBEDDINGS_MODEL=nomic-ai/nomic-embed-text-v2-moe +``` + +### 2. nomic-ai/nomic-embed-text-v1.5 + +**Characteristics**: +- **Dimensions**: 768 +- **Speed**: Slower than v2-moe +- **Quality**: Slightly better than v2-moe +- **Languages**: 30+ (multilingual) +- **Use case**: When quality matters more than speed + +```bash +export EMBEDDINGS_MODEL=nomic-ai/nomic-embed-text-v1.5 +``` + +### 3. all-MiniLM-L6-v2 + +**Characteristics**: +- **Dimensions**: 384 (padded to 768 for schema compatibility) +- **Speed**: Very fast (lightweight) +- **Quality**: Good for similarity +- **Languages**: English +- **Use case**: High-throughput, English-only scenarios + +```bash +export EMBEDDINGS_MODEL=all-MiniLM-L6-v2 +# or +export EMBEDDINGS_MODEL=sentence-transformers/all-MiniLM-L6-v2 +``` + +### 4. BAAI/bge-small-en-v1.5 + +**Characteristics**: +- **Dimensions**: 384 (padded to 768) +- **Speed**: Very fast +- **Quality**: Good for English retrieval +- **Languages**: English +- **Use case**: Fast English-only systems + +```bash +export EMBEDDINGS_MODEL=BAAI/bge-small-en-v1.5 +``` + +### 5. BAAI/bge-base-en-v1.5 + +**Characteristics**: +- **Dimensions**: 768 (native) +- **Speed**: Medium +- **Quality**: Excellent for English +- **Languages**: English +- **Use case**: Best quality for English-only deployments + +```bash +export EMBEDDINGS_MODEL=BAAI/bge-base-en-v1.5 +``` + +--- + +## Configuration + +### Environment Variables + +```bash +# Select embedding model (default: nomic-ai/nomic-embed-text-v2-moe) +export EMBEDDINGS_MODEL=nomic-ai/nomic-embed-text-v2-moe + +# Gateway endpoint (default: https://api.riotpiao.com) +export LLM_API_BASE=https://api.riotpiao.com + +# Optional: API key for gated models +export LLM_API_KEY=your-api-key +``` + +### Kubernetes Deployment + +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: poimen-memory +spec: + template: + spec: + containers: + - name: memory + env: + - name: EMBEDDINGS_MODEL + value: "nomic-ai/nomic-embed-text-v1.5" # Higher quality + - name: LLM_API_BASE + value: "https://api.riotpiao.com" +``` + +### Docker + +```bash +docker run \ + -e EMBEDDINGS_MODEL=all-MiniLM-L6-v2 \ + -e LLM_API_BASE=https://api.riotpiao.com \ + poimen-memory:latest +``` + +--- + +## Comparison Table + +| Model | Dims | Speed | Quality | Languages | Use Case | +|-------|------|-------|---------|-----------|----------| +| **nomic-v2-moe** (default) | 768 | ⚡⚡⚡ | ⭐⭐⭐ | 30+ | **Production default** | +| **nomic-v1.5** | 768 | ⚡⚡ | ⭐⭐⭐⭐ | 30+ | Quality-first, multilingual | +| **all-MiniLM-L6** | 384→768 | ⚡⚡⚡⚡ | ⭐⭐⭐ | English | High throughput | +| **bge-small-en** | 384→768 | ⚡⚡⚡⚡ | ⭐⭐⭐ | English | Fast retrieval | +| **bge-base-en** | 768 | ⚡⚡ | ⭐⭐⭐⭐⭐ | English | English-only best quality | + +--- + +## Performance Impact + +### Embedding Latency (per text) + +``` +nomic-v2-moe: ~50ms ← Default (good balance) +all-MiniLM-L6: ~30ms ← Fastest +nomic-v1.5: ~80ms +bge-base-en: ~60ms +``` + +### Throughput at 10 batch size + +``` +nomic-v2-moe: ~200 texts/sec +all-MiniLM-L6: ~330 texts/sec ← Highest throughput +nomic-v1.5: ~125 texts/sec +bge-base-en: ~165 texts/sec +``` + +--- + +## Switching Models + +### 1. While Running + +```bash +# Change env var +kubectl set env deployment/poimen-memory \ + EMBEDDINGS_MODEL=nomic-ai/nomic-embed-text-v1.5 + +# Restart pods (will use new model) +kubectl rollout restart deployment/poimen-memory + +# Monitor logs +kubectl logs -f deployment/poimen-memory | grep "Embeddings client" +# Expected: "Embeddings client initialized: model=nomic-ai/nomic-embed-text-v1.5" +``` + +### 2. New Ingests + +Changing models only affects **new** ingests. Existing chunks keep their old embeddings. + +To re-embed existing chunks: +```bash +# 1. Mark all chunks as pending re-embedding +psql -h memory-db -U app memory -c \ + "UPDATE chunks SET indexed_in_pgvector = false, opensearch_pending = true;" + +# 2. Restart queue worker (reprocesses all chunks) +kubectl rollout restart deployment/poimen-memory + +# Wait for queue to clear +# Monitor: kubectl logs -f deployment/poimen-memory | grep "messages_processed" +``` + +### 3. Validate Model + +```bash +# Check logs for model initialization +kubectl logs -n poimen deployment/poimen-memory | grep "Embeddings client" + +# Test embedding endpoint +curl -X POST https://api.riotpiao.com/v1/embeddings \ + -H "Content-Type: application/json" \ + -d '{ + "model": "nomic-ai/nomic-embed-text-v1.5", + "input": ["hello world"] + }' | jq '.data[0].embedding | length' +# Output: 768 +``` + +--- + +## Troubleshooting + +### "Unsupported embedding model" Error + +**Error**: +``` +thread 'actix-web' panicked at 'Unsupported embedding model: bert-base-uncased' +``` + +**Fix**: +1. Check supported models list above +2. Use one of the validated models +3. If you have a custom model, update `validate_model()` in embeddings.rs + +### Wrong Dimensionality + +**Error**: +``` +model custom-embed-384 returned 384-dim vector, expected 768 +``` + +**Cause**: Model returns 384-dim vectors, but pgvector schema expects 768 + +**Solutions**: +1. Use a model that returns 768-dim (e.g., `nomic-ai/nomic-embed-text-v1.5`) +2. Or manually pad 384-dim vectors to 768-dim by adding zeros +3. Or re-migrate schema to 384-dim (complex, not recommended) + +### Slow Embedding Performance + +**If latency > 200ms per text**: + +```bash +# Try faster model +kubectl set env deployment/poimen-memory \ + EMBEDDINGS_MODEL=all-MiniLM-L6-v2 + +# Check embedding service health +curl https://api.riotpiao.com/v1/models | jq '.data[] | select(.id | contains("embed"))' + +# Monitor queue worker metrics +kubectl logs -f deployment/poimen-memory | grep "processing_time" +``` + +--- + +## Integration with Queue Worker + +The Queue Worker automatically uses the configured embedding model: + +```rust +// crates/mem-cli/src/queue_worker.rs +let embedding_vec = embeddings.embed_one(&content).await?; +// ↑ Uses model from EMBEDDINGS_MODEL env var +``` + +When queue worker logs show: +``` +Embeddings client initialized: model=nomic-ai/nomic-embed-text-v1.5 +``` + +All embeddings are computed with that model. + +--- + +## Validation at Startup + +The system validates model compatibility on HTTP server startup: + +``` +INFO Embeddings client initialized: model=nomic-ai/nomic-embed-text-v2-moe, base_url=https://api.riotpiao.com +``` + +If validation fails, server refuses to start: + +``` +ERROR Unsupported embedding model: unknown-model. Supported models: [...] +``` + +--- + +## Adding Custom Models + +To support a new embedding model: + +1. **Verify dimension**: Test with gateway + ```bash + curl -X POST https://api.riotpiao.com/v1/embeddings \ + -d '{"model": "my-custom-model", "input": ["test"]}' + # Check dimension count in response + ``` + +2. **Add to allowed list** (crates/mem-llm/src/embeddings.rs): + ```rust + let supported_models = vec![ + "nomic-ai/nomic-embed-text-v2-moe", + "my-custom-model", // ← Add here + ]; + ``` + +3. **Document dimensions** in this file + +4. **Test**: + ```bash + EMBEDDINGS_MODEL=my-custom-model cargo run --bin mem -- serve + ``` + +--- + +## References + +- [Nomic AI Models](https://www.nomic.ai/) +- [BGE Models (BAAI)](https://huggingface.co/BAAI/bge-base-en-v1.5) +- [Sentence Transformers](https://www.sbert.net/) +- [Gateway Embeddings API](https://api.riotpiao.com/docs#/embeddings)