Merge pull request #72 from namtroi/phase5/dev

Phase5/dev

Author: namtroi

apps/ai-worker/requirements.txt

@@ -36,3 +36,6 @@ lxml>=5.1.0
 markdownify>=0.11.6
 ebooklib>=0.18
 openpyxl>=3.1.2
+
+# Phase 5: Hybrid Embeddings (dense + sparse)
+fastembed>=0.4.0

apps/ai-worker/src/embedder.py

@@ -0,0 +1,152 @@
+# apps/ai-worker/src/hybrid_embedder.py
+"""
+HybridEmbedder: Dense + Sparse vector generation.
+
+Phase 5: Uses fastembed for both dense (BGE) and sparse (BM25) embeddings.
+Replaces sentence-transformers for unified embedding approach.
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import structlog
+
+logger = structlog.get_logger()
+
+
+@dataclass
+class SparseVector:
+    """Sparse vector representation for BM25-style search."""
+
+    indices: List[int]
+    values: List[float]
+
+
+@dataclass
+class HybridVector:
+    """Combined dense + sparse vector for hybrid search."""
+
+    dense: List[float]  # 384 floats (BGE-small)
+    sparse: SparseVector  # Variable length
+
+
+class HybridEmbedder:
+    """
+    Generates both dense and sparse embeddings using fastembed.
+
+    - Dense: BAAI/bge-small-en-v1.5 (384 dimensions)
+    - Sparse: Qdrant/bm25 (BM25-based sparse vectors)
+    """
+
+    _instance: Optional["HybridEmbedder"] = None
+    _dense_model = None
+    _sparse_model = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        if self._dense_model is None:
+            self._load_models()
+
+    def _load_models(self):
+        """Load both embedding models."""
+        from fastembed import SparseTextEmbedding, TextEmbedding
+
+        logger.info("loading_hybrid_embedding_models")
+
+        # Dense model - same as before (BAAI/bge-small-en-v1.5)
+        logger.info("loading_dense_model", model="BAAI/bge-small-en-v1.5")
+        self._dense_model = TextEmbedding("BAAI/bge-small-en-v1.5")
+
+        # Sparse model - BM25 for keyword matching
+        logger.info("loading_sparse_model", model="Qdrant/bm25")
+        self._sparse_model = SparseTextEmbedding("Qdrant/bm25")
+
+        logger.info("hybrid_embedding_models_loaded")
+
+    def embed(self, texts: List[str]) -> List[HybridVector]:
+        """
+        Generate hybrid (dense + sparse) embeddings for texts.
+
+        Args:
+            texts: List of text strings to embed.
+
+        Returns:
+            List of HybridVector containing dense and sparse vectors.
+        """
+        if not texts:
+            return []
+
+        try:
+            # Generate both embedding types
+            dense_embeddings = list(self._dense_model.embed(texts))
+            sparse_embeddings = list(self._sparse_model.embed(texts))
+
+            # Combine into HybridVector
+            results = []
+            for dense, sparse in zip(dense_embeddings, sparse_embeddings):
+                results.append(
+                    HybridVector(
+                        dense=dense.tolist(),
+                        sparse=SparseVector(
+                            indices=sparse.indices.tolist(),
+                            values=sparse.values.tolist(),
+                        ),
+                    )
+                )
+
+            return results
+
+        except Exception as e:
+            logger.error("hybrid_embedding_failed", error=str(e))
+            raise
+
+    def embed_dense_only(self, texts: List[str]) -> List[List[float]]:
+        """
+        Generate only dense embeddings (backward compatibility).
+
+        Args:
+            texts: List of text strings to embed.
+
+        Returns:
+            List of dense vectors (384 floats each).
+        """
+        if not texts:
+            return []
+
+        embeddings = list(self._dense_model.embed(texts))
+        return [e.tolist() for e in embeddings]
+
+    def get_token_counts(self, texts: List[str]) -> List[int]:
+        """
+        Estimate token counts for texts.
+
+        Uses a simple heuristic based on word count.
+        For more accurate counts, use the dense model's tokenizer.
+        """
+        if not texts:
+            return []
+
+        # Simple estimation: ~0.75 tokens per word (typical for English)
+        # This is faster than loading tokenizer for each call
+        counts = []
+        for text in texts:
+            word_count = len(text.split())
+            # Cap at 512 (model max)
+            counts.append(min(int(word_count * 1.3), 512))
+        return counts
+
+
+# Singleton instance
+_hybrid_embedder: Optional[HybridEmbedder] = None
+
+
+def get_hybrid_embedder() -> HybridEmbedder:
+    """Get singleton HybridEmbedder instance."""
+    global _hybrid_embedder
+    if _hybrid_embedder is None:
+        _hybrid_embedder = HybridEmbedder()
+    return _hybrid_embedder

apps/ai-worker/src/hybrid_embedder.py

@@ -17,6 +17,7 @@
 
 from .callback import send_callback
 from .config import settings
+from .hybrid_embedder import HybridEmbedder
 from .logging_config import configure_logging, get_logger
 from .metrics import MetricsCollector
 from .models import (
@@ -70,21 +71,52 @@ async def readiness_check():
 
 @app.post("/embed", response_model=EmbedResponse)
 async def embed_texts(request: EmbedRequest):
-    """Generate embeddings for a list of texts."""
-    from .embedder import Embedder
+    """Generate dense-only embeddings for a list of texts (backward compatibility)."""
 
     if not request.texts:
         return EmbedResponse(embeddings=[])
 
     try:
-        embedder = Embedder()
-        embeddings = embedder.embed(request.texts)
+        embedder = HybridEmbedder()
+        embeddings = embedder.embed_dense_only(request.texts)
         return EmbedResponse(embeddings=embeddings)
     except Exception as e:
         logger.exception("embed_error", error=str(e))
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@app.post("/embed/query")
+async def embed_query(request: dict):
+    """
+    Generate hybrid embeddings for a search query.
+
+    Returns both dense (384d) and sparse (BM25) vectors for Qdrant hybrid search.
+    """
+    from .models import HybridEmbedResponse, SparseVectorModel
+
+    text = request.get("text", "")
+    if not text:
+        raise HTTPException(status_code=400, detail="text is required")
+
+    try:
+        embedder = HybridEmbedder()
+        vectors = embedder.embed([text])
+
+        if not vectors:
+            raise HTTPException(status_code=500, detail="Failed to generate embeddings")
+
+        return HybridEmbedResponse(
+            dense=vectors[0].dense,
+            sparse=SparseVectorModel(
+                indices=vectors[0].sparse.indices,
+                values=vectors[0].sparse.values,
+            ),
+        )
+    except Exception as e:
+        logger.exception("embed_query_error", error=str(e))
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @app.post("/process", response_model=ProcessResponse)
 async def process_document(request: ProcessRequest):
     """

apps/ai-worker/src/main.py

@@ -114,3 +114,33 @@ class EmbedResponse(BaseModel):
     """Response with generated embeddings."""
 
     embeddings: List[List[float]]
+
+
+# Phase 5: Hybrid Embedding Models for Query
+
+
+class SparseVectorModel(BaseModel):
+    """Sparse vector for BM25-style keyword matching."""
+
+    indices: List[int]
+    values: List[float]
+
+
+class HybridVectorModel(BaseModel):
+    """Combined dense + sparse vector for hybrid search."""
+
+    dense: List[float]
+    sparse: SparseVectorModel
+
+
+class HybridEmbedRequest(BaseModel):
+    """Request to generate hybrid embeddings for search query."""
+
+    text: str
+
+
+class HybridEmbedResponse(BaseModel):
+    """Response with hybrid embeddings for search."""
+
+    dense: List[float]
+    sparse: SparseVectorModel

apps/ai-worker/src/models.py

@@ -7,7 +7,7 @@
 from .chunkers.document_chunker import DocumentChunker
 from .chunkers.presentation_chunker import PresentationChunker
 from .chunkers.tabular_chunker import TabularChunker
-from .embedder import Embedder
+from .hybrid_embedder import HybridEmbedder
 from .logging_config import get_logger
 from .models import ProfileConfig
 from .quality.analyzer import QualityAnalyzer
@@ -48,7 +48,8 @@ def __init__(self, config: Optional[ProfileConfig] = None):
             penalty_per_flag=self.config.qualityPenaltyPerFlag,
         )
 
-        self.embedder = Embedder()
+        # Phase 5: Hybrid embedder (dense + sparse)
+        self.embedder = HybridEmbedder()
 
     def _strip_breadcrumb_prefix(self, content: str) -> str:
         """Remove breadcrumb prefix (> Chapter > Section) from content."""
@@ -291,15 +292,22 @@ def run(
             chunk["metadata"]["chunkType"] = category
             chunk["index"] = i
 
-        # 4. Generate embeddings and token counts (with timing)
+        # 4. Generate hybrid embeddings and token counts (with timing)
         texts = [c["content"] for c in chunks]
         embed_start = time.time()
-        embeddings = self.embedder.embed(texts)
+        hybrid_vectors = self.embedder.embed(texts)
         token_counts = self.embedder.get_token_counts(texts)
         embedding_time_ms = int((time.time() - embed_start) * 1000)
 
         for i, chunk in enumerate(chunks):
-            chunk["embedding"] = embeddings[i]
+            # Phase 5: Hybrid vector format for Qdrant
+            chunk["vector"] = {
+                "dense": hybrid_vectors[i].dense,
+                "sparse": {
+                    "indices": hybrid_vectors[i].sparse.indices,
+                    "values": hybrid_vectors[i].sparse.values,
+                },
+            }
             chunk["metadata"]["tokenCount"] = token_counts[i]
 
         logger.info(

apps/ai-worker/src/pipeline.py

@@ -96,13 +96,15 @@ class TestEmbeddingDimensions:
     """Regression tests for embedding dimensions."""
 
     def test_embedding_dimensions_unchanged(self):
-        """Embeddings should still be 384 dimensions."""
-        from src.embedder import Embedder
+        """HybridEmbedder should return 384 dense dimensions."""
+        from src.hybrid_embedder import HybridEmbedder
 
-        embedder = Embedder()
-        embedding = embedder.embed("Test text")
+        embedder = HybridEmbedder()
+        result = embedder.embed(["Test text"])
 
-        assert len(embedding) == 384
+        # HybridEmbedder returns list of HybridVector
+        assert len(result) == 1
+        assert len(result[0].dense) == 384
 
 
 class TestChunkStructure: