# Story 4.4: Vector Embedding and Indexing Status: done > [!NOTE] > **Schema Evolution (Post-Completion):** This story was originally completed with Excel-only support. The Milvus schema has since been updated to a unified structure supporting both Excel and Word documents. The schema now uses `document_type` (VARCHAR) and `metadata` (JSON) fields instead of individual `sheet_name`, `row_number`, and `has_symbols` fields. This document reflects the **current unified schema**. ## Story As a **system**, I want **records embedded as vectors and indexed in Milvus**, So that **semantic search finds conceptually similar content**. ## Acceptance Criteria 1. **AC4.4.1:** Convert records to text representation for embedding - Flatten content fields: `"Key1: Value1. Key2: Value2. ..."` - Include resolved fields if symbols resolved: `"Key_resolved: Meaning"` - Include source context: `"From: filename, Sheet: sheet_name, Row: row_number"` - Result format: `"Item Code: 2024-4_0019. Report: not applicable. Screen: Applicable/Yes. From: distribution_spec.xlsx, Sheet: Item Mapping, Row: 42"` 2. **AC4.4.2:** Batch embed records using Azure OpenAI text-embedding-3-large - Use Azure OpenAI `text-embedding-3-large` model (3072 dimensions) - Batch process 100 texts per API call for efficiency - Return list of embedding vectors (3072-dim floats) - Handle API errors and rate limits gracefully - Log embedding metrics (batch size, latency) 3. **AC4.4.3:** Create Milvus collection with proper schema - Collection name: `dsol_records` (from config) - Vector field: 3072 dimensions, FLOAT_VECTOR type - Scalar fields: record_id (VARCHAR), document_id (VARCHAR), filename (VARCHAR), document_type (VARCHAR), text_content (VARCHAR), metadata (JSON) - Index type: HNSW (Hierarchical Navigable Small World) for fast similarity search - Distance metric: Cosine similarity 4. **AC4.4.4:** Store vectors in Milvus with metadata - Each vector stored with: - `record_id`: UUID linking to PostgreSQL record - `document_id`: UUID for filtering by document - `filename`: For display in results - `document_type`: Document type ("excel" or "word") - `text_content`: For displaying context in results - `metadata`: JSON with type-specific fields (Excel: sheet_name/row_number/has_symbols, Word: page_number/contains_table) - Batch upsert 100 vectors per Milvus call - Use upsert (not insert) for idempotency 5. **AC4.4.5:** Performance target ≤ 30s per 1000 records - Embedding: ~2-3s per 100 records (10 batches for 1000 = ~25s) - Milvus upload: ~500ms per 100 vectors (10 batches = ~5s) - Total: ≤ 30s for 1000 records - Log timing metrics for monitoring 6. **AC4.4.6:** Error handling and retry logic - Retry on transient Azure OpenAI errors (rate limits, timeouts) - Retry on transient Milvus errors (connection failures) - Use exponential backoff (1s, 2s, 4s) - Log failed batches with context - Raise `IndexingError` on permanent failure 7. **AC4.4.7:** Integration with StructuredStore (Story 4.1) - Use record IDs from PostgreSQL as foreign keys - Ensure records are stored in PostgreSQL before embedding - Link vectors back to structured records via `record_id` - Support querying both stores simultaneously ## Tasks / Subtasks - [ ] **Task 1: Create Embeddings module** (AC: 4.4.1, 4.4.2) - [ ] Create `src/knowledge/embeddings.py` module - [ ] Implement `record_to_text(record)` function for text representation - [ ] Implement `EmbeddingService` class with Azure OpenAI client - [ ] Add `embed_batch(texts)` method (max 100 texts) - [ ] Use structlog for logging - [ ] **Task 2: Create Azure OpenAI client integration** (AC: 4.4.2) - [ ] Initialize OpenAI client with Azure credentials from config - [ ] Use `text-embedding-3-large` deployment name - [ ] Set dimensions=3072 explicitly - [ ] Handle API response parsing - [ ] Add retry logic for rate limits - [ ] **Task 3: Implement VectorStore class** (AC: 4.4.3, 4.4.4) - [ ] Create `src/knowledge/vector_store.py` (extends existing stub) - [ ] Implement `VectorStore` class with Milvus client - [ ] Add `create_collection()` method with schema definition - [ ] Add `index_records(records, document_id)` method - [ ] Add `_prepare_milvus_entities()` helper for data transformation - [ ] **Task 4: Implement Milvus collection creation** (AC: 4.4.3) - [ ] Define collection schema with all required fields - [ ] Create HNSW index on vector field - [ ] Set index parameters (M=16, efConstruction=200) - [ ] Handle collection already exists (idempotent) - [ ] Load collection after creation - [ ] **Task 5: Implement batch vector upload** (AC: 4.4.4, 4.4.5) - [ ] Split records into batches of 100 - [ ] For each batch: convert to text → embed → upload to Milvus - [ ] Use upsert operation for idempotency - [ ] Track timing metrics per batch - [ ] Log progress (batch X of Y) - [ ] **Task 6: Add error handling and retry logic** (AC: 4.4.6) - [ ] Wrap Azure OpenAI calls in try/except with retry - [ ] Wrap Milvus calls in try/except with retry - [ ] Implement exponential backoff - [ ] Log errors with structured context - [ ] Raise `IndexingError` from `src.core.exceptions` - [ ] **Task 7: Unit tests** (AC: All) - [ ] Create `tests/knowledge/test_embeddings.py` - [ ] Test `record_to_text()` with various record types - [ ] Test `embed_batch()` with mocked Azure OpenAI client - [ ] Test error handling and retry logic - [ ] Mock Azure OpenAI responses - [ ] **Task 8: Unit tests for VectorStore** (AC: All) - [ ] Create `tests/knowledge/test_vector_store.py` - [ ] Test collection creation with schema - [ ] Test entity preparation and data transformation - [ ] Test batch upload logic - [ ] Mock Milvus client - [ ] **Task 9: Integration tests** (AC: All) - [ ] Create `tests/knowledge/test_vector_store_integration.py` - [ ] Test embedding and indexing 100 real records - [ ] Verify vectors are stored in Milvus - [ ] Test similarity search retrieval - [ ] Measure performance (≤ 30s per 1000 records) - [ ] Use real Milvus instance (Docker) - [ ] **Task 10: Create usage examples** (AC: All) - [ ] Create `examples/vector_store_usage.py` - [ ] Demonstrate collection creation - [ ] Demonstrate batch embedding and indexing - [ ] Demonstrate similarity search - [ ] Show integration with StructuredStore ## Dev Notes ### Architecture Patterns and Constraints **From architecture.md:** - **Module location:** `src/knowledge/embeddings.py`, `src/knowledge/vector_store.py` (architecture.md:116-117) - **Vector store:** Milvus 1.16+ with HNSW index (architecture.md:20, 175-181) - **Embedding model:** Azure OpenAI text-embedding-3-large (3072 dims) (architecture.md:25) - **Distance metric:** Cosine similarity (architecture.md:180) - **Error handling:** Use `DSOLError` hierarchy from `src.core.exceptions` - **Logging:** Structured logging with structlog **From tech spec (knowledge-indexing-tech-spec.md):** - **Text representation format:** (lines 404-431) - Content fields: "Key1: Value1. Key2: Value2. ..." - Resolved fields: "Key_resolved: Meaning" (if symbols resolved) - Source context: "From: filename, Sheet: sheet_name, Row: row_number" - **Batch embedding:** (lines 433-451) - Azure OpenAI client with async API - Model: text-embedding-3-large - Dimensions: 3072 - Max batch size: 100 texts per call - Response parsing: `[item.embedding for item in response.data]` - **Milvus collection config:** (lines 453-466) - Collection name: "dsol_records" (from config) - Vector size: 3072 - Distance: Cosine - Indexing threshold: 10000 (auto-index after 10k vectors) - **Milvus entity structure:** (lines 468-482) - record_id (VARCHAR): Link to PostgreSQL - document_id (VARCHAR): For filtering - filename (VARCHAR): For display - document_type (VARCHAR): "excel" or "word" - text_content (VARCHAR): For result display - metadata (JSON): Type-specific fields - Excel: {sheet_name, row_number, has_symbols} - Word: {page_number, contains_table} - **Performance targets:** (lines 503-507) - Batch size: 100 records per embedding call - Embedding latency: ~2-3s per 100 records - Milvus upload: ~500ms per 100 vectors - Total: ≤ 30s per 1000 records **From PRD (prd.md):** - **FR20:** System supports semantic similarity search for conceptual queries - **FR21:** System maintains source location for all indexed content - **NFR3:** Query response time < 5 seconds (includes retrieval + answer generation) ### Milvus Schema Reference **Collection Schema:** ```python from pymilvus import ( CollectionSchema, FieldSchema, DataType, Collection, connections, ) # Define fields (Unified schema for Excel and Word documents) fields = [ FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True), FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=3072), FieldSchema(name="record_id", dtype=DataType.VARCHAR, max_length=36), FieldSchema(name="document_id", dtype=DataType.VARCHAR, max_length=36), FieldSchema(name="filename", dtype=DataType.VARCHAR, max_length=255), FieldSchema(name="document_type", dtype=DataType.VARCHAR, max_length=20), FieldSchema(name="text_content", dtype=DataType.VARCHAR, max_length=8000), FieldSchema(name="metadata", dtype=DataType.JSON), ] # Create schema schema = CollectionSchema(fields, description="DSOL document records") # Create collection collection = Collection(name="dsol_records", schema=schema) ``` **HNSW Index Configuration:** ```python from pymilvus import IndexType index_params = { "metric_type": "COSINE", "index_type": "HNSW", "params": { "M": 16, # Number of bi-directional links "efConstruction": 200 # Size of dynamic candidate list } } collection.create_index( field_name="vector", index_params=index_params ) ``` ### Text Representation Examples **Example 1: Simple Record** ```python record = ExtractedRecord( content={"Item Code": "2024-4_0019", "Report": "not applicable"}, headers=["Item Code", "Report"], _source={ "filename": "distribution_spec.xlsx", "sheet": "Item Mapping", "row": 42, } ) # Text representation: # "Item Code: 2024-4_0019. Report: not applicable. From: distribution_spec.xlsx, Sheet: Item Mapping, Row: 42" ``` **Example 2: Record with Symbol Resolution** ```python record = ExtractedRecord( content={"Item Code": "2024-4_0019", "Screen": "○"}, resolved_content={"Screen_resolved": "Applicable/Yes"}, headers=["Item Code", "Screen"], _source={ "filename": "distribution_spec.xlsx", "sheet": "Item Mapping", "row": 42, } ) # Text representation: # "Item Code: 2024-4_0019. Screen: ○. Screen_resolved: Applicable/Yes. From: distribution_spec.xlsx, Sheet: Item Mapping, Row: 42" ``` ### Azure OpenAI Integration **Client Setup:** ```python from openai import AsyncAzureOpenAI from src.core.config import settings client = AsyncAzureOpenAI( api_key=settings.azure_openai_api_key, api_version=settings.azure_openai_api_version, azure_endpoint=settings.azure_openai_endpoint, ) ``` **Embedding API Call:** ```python response = await client.embeddings.create( model=settings.azure_openai_embedding_deployment, # "text-embedding-3-large" input=texts[:100], # Max 100 texts per call dimensions=3072 ) embeddings = [item.embedding for item in response.data] # Returns: list[list[float]] with 3072 dimensions each ``` ### Integration with Story 4.1 **Workflow:** ```python from src.knowledge.structured_store import StructuredStore from src.knowledge.vector_store import VectorStore from src.extraction.models import ExtractedRecord # Step 1: Store records in PostgreSQL (Story 4.1) structured_store = StructuredStore(db_session) stored_ids = await structured_store.store_records(records, document_id) # Step 2: Embed and index in Milvus (Story 4.4) vector_store = VectorStore() await vector_store.index_records(records, document_id, stored_ids) # Both stores now contain the same records # - PostgreSQL: Structured data with exact lookups # - Milvus: Vector embeddings for semantic search ``` ### Prerequisites - **Story 4.1 COMPLETE:** StructuredStore must be implemented and storing records - **Database:** PostgreSQL with `extracted_records` table - **Milvus:** Running instance (Docker or cloud) - **Azure OpenAI:** Deployment with text-embedding-3-large model - **Environment:** AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, MILVUS_HOST configured ### Success Criteria - ✅ Records converted to text representation correctly - ✅ Azure OpenAI embedding works (3072 dimensions) - ✅ Milvus collection created with proper schema - ✅ Vectors stored with complete metadata - ✅ Performance: ≤ 30s per 1000 records - ✅ Error handling with retry logic working - ✅ Integration tests pass with real Milvus - ✅ Similarity search retrieves relevant results ### Out of Scope - **Query routing:** Determining exact vs semantic search (Epic 5, Story 5.2) - **Hybrid retrieval:** Combining PostgreSQL + Milvus results (Epic 5, Story 5.2) - **Answer generation:** LLM synthesis of answers (Epic 5, Story 5.3) - **Reranking:** Post-retrieval result reranking (future optimization) - **Incremental indexing:** Add/delete operations (Story 4.5) - **Collection partitioning:** Advanced Milvus optimization (future) ### Next Steps After Story 4.4: - **Story 4.5:** Incremental Indexing (add/delete/reindex operations) - **Epic 5:** Q&A & Answer Generation (uses both stores for retrieval) - Story 5.1: Query endpoint - Story 5.2: Query router (exact vs semantic) - Story 5.3: Answer generation with sources ## Technical Design ### Data Models **Text Representation:** ```python def record_to_text(record: ExtractedRecord) -> str: """Convert ExtractedRecord to text for embedding. Args: record: Extracted record with content, headers, and source metadata Returns: Text representation: "Key1: Value1. Key2: Value2. ... From: file, Sheet: sheet, Row: row" """ content_parts = [] # Add content fields for key, value in record.content.items(): if value is not None: content_parts.append(f"{key}: {value}") # Add resolved fields (symbol meanings) if record.resolved_content: for key, value in record.resolved_content.items(): if "_resolved" in key and value: content_parts.append(f"{key}: {value}") # Add source context source = record._source source_text = f"From: {source['filename']}, Sheet: {source['sheet']}, Row: {source['row']}" return f"{'. '.join(content_parts)}. {source_text}" ``` **Milvus Entity:** ```python from dataclasses import dataclass from uuid import UUID @dataclass class MilvusEntity: """Milvus entity structure for vector storage. Unified schema supporting both Excel and Word documents. """ vector: list[float] # 3072-dim embedding record_id: str # UUID linking to PostgreSQL document_id: str # UUID for filtering filename: str # For display document_type: str # "excel" or "word" text_content: str # For result display metadata: dict # Type-specific fields # Excel: {"sheet_name": str, "row_number": int, "has_symbols": bool} # Word: {"page_number": int, "contains_table": bool} ``` ### EmbeddingService Class ```python from openai import AsyncAzureOpenAI import structlog class EmbeddingService: """Service for converting text to embeddings using Azure OpenAI.""" def __init__(self): """Initialize with Azure OpenAI client.""" self.client = AsyncAzureOpenAI( api_key=settings.azure_openai_api_key, api_version=settings.azure_openai_api_version, azure_endpoint=settings.azure_openai_endpoint, ) self.logger = structlog.get_logger() async def embed_batch( self, texts: list[str], max_retries: int = 3 ) -> list[list[float]]: """Embed batch of texts using Azure OpenAI. Args: texts: List of text representations (max 100) max_retries: Number of retries for transient errors Returns: List of embedding vectors (3072 dimensions each) Raises: IndexingError: If embedding fails after retries """ if len(texts) > 100: raise ValueError("Batch size must be ≤ 100") for attempt in range(max_retries): try: response = await self.client.embeddings.create( model=settings.azure_openai_embedding_deployment, input=texts, dimensions=3072 ) embeddings = [item.embedding for item in response.data] self.logger.info( "embeddings_created", batch_size=len(texts), dimensions=3072, attempt=attempt + 1 ) return embeddings except Exception as e: if attempt == max_retries - 1: self.logger.error( "embedding_failed", batch_size=len(texts), error=str(e), attempt=attempt + 1 ) raise IndexingError(f"Embedding failed: {e}") # Exponential backoff await asyncio.sleep(2 ** attempt) ``` ### VectorStore Class ```python from pymilvus import Collection, connections import structlog class VectorStore: """Manages Milvus vector storage operations.""" def __init__(self): """Initialize with Milvus connection.""" self.logger = structlog.get_logger() self.embedding_service = EmbeddingService() async def create_collection(self): """Create Milvus collection with schema if not exists.""" # Implementation in Task 4 pass async def index_records( self, records: list[ExtractedRecord], document_id: UUID, record_ids: list[UUID] ) -> int: """Embed and index records in Milvus. Args: records: List of extracted records to index document_id: Document UUID for metadata record_ids: PostgreSQL record IDs from StructuredStore Returns: Number of vectors indexed Raises: IndexingError: If indexing fails """ start_time = time.time() total_indexed = 0 # Process in batches of 100 for i in range(0, len(records), 100): batch_records = records[i:i+100] batch_ids = record_ids[i:i+100] # Convert to text texts = [record_to_text(r) for r in batch_records] # Embed embeddings = await self.embedding_service.embed_batch(texts) # Prepare entities entities = self._prepare_milvus_entities( batch_records, embeddings, batch_ids, document_id ) # Upload to Milvus await self._upsert_vectors(entities) total_indexed += len(batch_records) self.logger.info( "batch_indexed", batch_num=i//100 + 1, batch_size=len(batch_records), total_indexed=total_indexed ) duration = time.time() - start_time self.logger.info( "indexing_complete", total_records=total_indexed, duration_seconds=duration, records_per_second=total_indexed / duration ) return total_indexed ``` ### Example Usage ```python from src.knowledge.structured_store import StructuredStore from src.knowledge.vector_store import VectorStore from src.db.session import get_db_session async def index_document_records(): """Example: Index document records in both stores.""" # Get extracted records from pipeline records = await extraction_pipeline.process_document(document_id) async with get_db_session() as session: # Step 1: Store in PostgreSQL structured_store = StructuredStore(session) record_ids = await structured_store.store_records(records, document_id) # Step 2: Embed and store in Milvus vector_store = VectorStore() vector_count = await vector_store.index_records( records, document_id, record_ids ) print(f"Indexed {vector_count} vectors in Milvus") ``` --- ## Learnings from Previous Stories **From Story 4.1: Structured Store (Status: done)** **Key Patterns:** - **Batch processing** - Process 500 records per transaction for performance - **Idempotent operations** - Use ON CONFLICT for safe retries - **Error handling** - Retry transient errors, raise IndexingError for permanent failures - **Performance logging** - Log metrics: records_per_second, duration_ms - **Async patterns** - Use SQLAlchemy async ORM throughout **From Story 4.2: Source Metadata Preservation (Status: done)** **Key Patterns:** - **JOIN queries** - Combine extracted_records with documents table for filename - **Dataclass pattern** - Use dataclasses for structured output (SourceCitation, ValidationReport) - **Validation** - Check completeness of source metadata before indexing - **Citation formatting** - Both human-readable and JSON output formats **From Story 4.3: Exact Lookup Indexing (Status: done)** **Key Patterns:** - **Index verification** - Query PostgreSQL system views to check index health - **Performance benchmarking** - Test queries against defined targets - **EXPLAIN ANALYZE** - Verify index usage in query plans - **Monitoring** - Track index usage statistics for production monitoring **Common Testing Patterns:** - **Async mock patterns** - Use AsyncMock for async methods, MagicMock for result objects - **Integration tests** - Test with real database/services when possible - **Performance tests** - Verify timing targets are met - **Example scripts** - Create both quickstart and comprehensive examples --- ## Definition of Done - [ ] All acceptance criteria (AC4.4.1 - AC4.4.7) implemented and tested - [ ] `EmbeddingService` class created in `src/knowledge/embeddings.py` - [ ] `VectorStore` class implemented in `src/knowledge/vector_store.py` - [ ] Text representation function working (`record_to_text()`) - [ ] Azure OpenAI embedding integration working (3072 dimensions) - [ ] Milvus collection creation with proper schema - [ ] Batch vector upload working (100 vectors per batch) - [ ] Unit tests pass (20+ tests across embeddings and vector_store) - [ ] Integration tests pass with real Milvus - [ ] Performance target met: ≤ 30s per 1000 records - [ ] Error handling with retry logic working - [ ] Example usage script created (`examples/vector_store_usage.py`) - [ ] Code follows project style (Black, Ruff, mypy strict) - [ ] Structured logging implemented - [ ] Documentation strings (docstrings) for all public methods - [ ] Sprint status updated to "done" - [ ] Git commit with proper commit message