"""Milvus vector store operations for document embeddings.""" import asyncio import json import time from dataclasses import dataclass from uuid import UUID import structlog from pymilvus import ( Collection, CollectionSchema, DataType, FieldSchema, connections, utility, ) from src.core.config import settings from src.core.exceptions import IndexingError from src.extraction.models import ExtractedRecord from src.knowledge.embeddings import EmbeddingService, record_to_text logger = structlog.get_logger(__name__) # Milvus JSON field limit is 65536 bytes; use 60KB threshold for safety margin MILVUS_JSON_MAX_BYTES = 60000 def truncate_metadata_to_fit(metadata: dict, max_bytes: int = MILVUS_JSON_MAX_BYTES) -> dict: """Truncate metadata dict to fit within Milvus JSON field size limit. Strategy: 1. Remove large string values (>1000 chars) first 2. If still too large, remove list/dict values 3. Final fallback: keep only essential fields Args: metadata: Original metadata dict max_bytes: Maximum allowed size in bytes (default: 60KB) Returns: Truncated metadata dict that fits within the limit """ def get_json_size(d: dict) -> int: return len(json.dumps(d, ensure_ascii=False, default=str).encode('utf-8')) # Check if already within limit original_size = get_json_size(metadata) if original_size <= max_bytes: return metadata result = metadata.copy() logger.warning( "metadata_exceeds_milvus_limit", original_size_bytes=original_size, max_bytes=max_bytes, message="Truncating metadata to fit Milvus JSON field limit" ) # Phase 1: Truncate large string values for key, value in list(result.items()): if isinstance(value, str) and len(value) > 1000: result[key] = value[:997] + "..." if get_json_size(result) <= max_bytes: return result # Phase 2: Remove list/dict values that are large for key, value in list(result.items()): if isinstance(value, (list, dict)): value_size = len(json.dumps(value, ensure_ascii=False, default=str).encode('utf-8')) if value_size > 5000: del result[key] if get_json_size(result) <= max_bytes: return result # Phase 3: Remove remaining large values progressively while get_json_size(result) > max_bytes and result: # Find largest value and remove it largest_key = None largest_size = 0 for key, value in result.items(): value_size = len(json.dumps(value, ensure_ascii=False, default=str).encode('utf-8')) if value_size > largest_size: largest_size = value_size largest_key = key if largest_key: del result[largest_key] else: break # Mark that metadata was truncated if result != metadata: result['_metadata_truncated'] = True return result @dataclass class MilvusEntity: """Milvus entity structure for vector storage. Unified schema supporting both Excel and Word documents. Attributes: vector: 3072-dimensional embedding vector record_id: UUID linking to PostgreSQL extracted_records table document_id: UUID for filtering by document tenant_id: Tenant UUID for data isolation (data-isolation-spec Section 3) node_id: Node UUID for data isolation (data-isolation-spec Section 3) filename: For display in search results document_type: Type of document ("excel" or "word") text_content: Original text for result display metadata: JSON dict with type-specific fields (flexible schema) Base fields for Excel: {"sheet_name": str, "row_number": int, "has_symbols": bool, "element_type": str} Additional fields from DETAILED pipeline: - Images: "image_number", "position_key", "image_format", "width", "height", "cell_reference", "from_cell", "to_cell", "mime_type" - Flowcharts: "chunk_type", "shape_count", "connector_count", "connection_count", "bounding_box" - Shapes: Shape-specific annotations and positions - Note: image_base64 is NOT stored to avoid bloating Milvus Base fields for Word: {"chunk_index": int, "contains_table": bool, "contains_image": bool} domain_category: Domain category (finance, hr, engineering, etc.) domain_terms: List of domain-specific terms for pre-filtering tags: Semantic tags for filtering (e.g., req:functional, risk:security) """ vector: list[float] record_id: str document_id: str tenant_id: str node_id: str filename: str document_type: str text_content: str metadata: dict domain_category: str = "general" domain_terms: list[str] = None tags: list[str] = None def __post_init__(self): """Initialize domain_terms and tags to empty lists if None.""" if self.domain_terms is None: self.domain_terms = [] if self.tags is None: self.tags = [] def _truncate_to_byte_length(s: str, max_bytes: int) -> str: """Truncate a string so its UTF-8 byte length does not exceed max_bytes.""" encoded = s.encode("utf-8") if len(encoded) <= max_bytes: return s # Truncate bytes and decode safely (ignore partial multibyte chars) return encoded[:max_bytes].decode("utf-8", errors="ignore") def normalize_tags(raw_tags: list, max_length: int = 50) -> list[str]: """Convert TagAssignment objects or dicts to simple tag name strings. The SemanticTagger returns TagAssignment objects with structure: {"tag_id": "...", "name": "req:functional", "confidence": 0.95, ...} Milvus expects simple string arrays: ["req:functional", "topic:auth"] Args: raw_tags: List of tags (TagAssignment dicts, strings, or objects) max_length: Max byte length per tag string (Milvus VARCHAR limit is in bytes) Returns: List of tag name strings suitable for Milvus storage """ if not raw_tags: return [] normalized = [] for tag in raw_tags: if isinstance(tag, str): normalized.append(_truncate_to_byte_length(tag, max_length)) elif isinstance(tag, dict): tag_name = tag.get("name") if tag_name: normalized.append(_truncate_to_byte_length(tag_name, max_length)) elif hasattr(tag, "name"): normalized.append(_truncate_to_byte_length(tag.name, max_length)) else: logger.warning("unknown_tag_format", tag_type=type(tag).__name__) return normalized class VectorStore: """Manages Milvus vector storage operations for document records. Handles embedding generation, collection creation, and batch vector upload with retry logic for robust indexing. """ def __init__(self) -> None: """Initialize VectorStore with Milvus connection and EmbeddingService.""" self.logger = structlog.get_logger() self.embedding_service = EmbeddingService() self.collection_name = settings.milvus_collection self.vector_dim = settings.embedding_dimensions self._connected = False def _connect(self) -> None: """Connect to Milvus server if not already connected.""" if not self._connected: connections.connect( alias="default", host=settings.milvus_host, port=settings.milvus_port, ) self._connected = True self.logger.info( "milvus_connected", host=settings.milvus_host, port=settings.milvus_port, ) def _disconnect(self) -> None: """Disconnect from Milvus server.""" if self._connected: connections.disconnect("default") self._connected = False self.logger.info("milvus_disconnected") async def create_collection(self) -> None: """Create Milvus collection with schema if it doesn't exist. Creates the collection with: - 3072-dimensional vector field - Scalar fields for metadata (record_id, document_id, filename, etc.) - HNSW index for fast similarity search - Cosine distance metric Idempotent: Does nothing if collection already exists. Raises: IndexingError: If collection creation fails """ self._connect() try: # Check if collection already exists if utility.has_collection(self.collection_name): self.logger.info( "collection_already_exists", collection_name=self.collection_name, ) return # Define unified schema fields for both 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=self.vector_dim, ), 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=65535, ), FieldSchema( name="metadata", dtype=DataType.JSON, ), # Domain-aware RAG fields FieldSchema( name="domain_category", dtype=DataType.VARCHAR, max_length=64, default_value="general", ), FieldSchema( name="domain_terms", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=50, max_length=100, ), # Semantic tagging fields FieldSchema( name="tags", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=20, max_length=100, # Increased to support longer multilingual tags (e.g., Japanese) ), # Tenant + node isolation fields (data-isolation-spec Section 3.2) FieldSchema( name="tenant_id", dtype=DataType.VARCHAR, max_length=36, ), FieldSchema( name="node_id", dtype=DataType.VARCHAR, max_length=36, ), ] # Create schema schema = CollectionSchema( fields, description="DSOL document records with embeddings", ) # Create collection collection = Collection( name=self.collection_name, schema=schema, ) # Create HNSW index for vector similarity search 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, ) # Create INVERTED indexes for domain fields (for efficient filtering) collection.create_index( field_name="domain_category", index_params={"index_type": "INVERTED"}, ) collection.create_index( field_name="domain_terms", index_params={"index_type": "INVERTED"}, ) collection.create_index( field_name="tags", index_params={"index_type": "INVERTED"}, ) # Tenant isolation indexes (data-isolation-spec Section 3.2) collection.create_index( field_name="tenant_id", index_params={"index_type": "INVERTED"}, ) collection.create_index( field_name="node_id", index_params={"index_type": "INVERTED"}, ) # Load collection into memory for searching collection.load() self.logger.info( "collection_created", collection_name=self.collection_name, vector_dim=self.vector_dim, index_type="HNSW", ) except Exception as e: self.logger.error( "collection_creation_failed", collection_name=self.collection_name, error=str(e), ) raise IndexingError( f"Failed to create Milvus collection: {e}", details={"collection_name": self.collection_name, "error": str(e)}, ) from e async def index_records( self, records: list[ExtractedRecord] | list[dict], document_id: UUID, record_ids: list[UUID], document_type: str = "excel", filename: str = "", tenant_id: str = "", node_id: str = "", ) -> int: """Embed and index records in Milvus. Supports both Excel ExtractedRecords and Word chunk dictionaries. Processes records in batches of 100: 1. Convert records to text representation 2. Generate embeddings via Azure OpenAI 3. Upload vectors to Milvus with metadata Args: records: List of ExtractedRecord objects (Excel) or chunk dicts (Word) document_id: Document UUID for metadata record_ids: PostgreSQL record IDs from StructuredStore document_type: "excel" or "word" (default: "excel") filename: Document filename (required for Word documents) tenant_id: Tenant UUID string (required for data isolation) node_id: Node UUID string (required for data isolation) Returns: Total number of vectors indexed Raises: IndexingError: If indexing fails ValueError: If records and record_ids lengths don't match """ if len(records) != len(record_ids): raise ValueError( f"Records count ({len(records)}) must match " f"record_ids count ({len(record_ids)})" ) if not records: self.logger.info("no_records_to_index") return 0 # Ensure collection exists before indexing await self.create_collection() self._connect() start_time = time.perf_counter() total_indexed = 0 batch_size = 100 try: # Process in batches of 100 for i in range(0, len(records), batch_size): batch_start = time.perf_counter() batch_records = records[i : i + batch_size] batch_ids = record_ids[i : i + batch_size] # Convert records to text based on document type if document_type == "excel": texts = [record_to_text(record) for record in batch_records] else: # Word chunks texts = [record.get("text", "") for record in batch_records] # Generate embeddings embeddings = await self.embedding_service.embed_batch(texts) # Prepare Milvus entities entities = self._prepare_milvus_entities( batch_records, embeddings, batch_ids, document_id, document_type, filename, tenant_id=tenant_id, node_id=node_id, ) # Upload to Milvus with retry await self._upsert_vectors(entities) total_indexed += len(batch_records) batch_duration_ms = (time.perf_counter() - batch_start) * 1000 self.logger.info( "batch_indexed", batch_num=i // batch_size + 1, batch_size=len(batch_records), total_indexed=total_indexed, total_records=len(records), batch_duration_ms=round(batch_duration_ms, 2), ) # Calculate overall metrics total_duration_s = time.perf_counter() - start_time self.logger.info( "indexing_complete", total_records=total_indexed, duration_seconds=round(total_duration_s, 2), records_per_second=round(total_indexed / total_duration_s, 2), ) return total_indexed except Exception as e: self.logger.error( "indexing_failed", total_indexed=total_indexed, total_records=len(records), error=str(e), ) raise IndexingError( f"Indexing failed after {total_indexed}/{len(records)} records: {e}", details={ "total_indexed": total_indexed, "total_records": len(records), "error": str(e), }, ) from e async def index_with_vectors( self, records: list[ExtractedRecord] | list[dict], vectors: list[list[float]], document_id: UUID, record_ids: list[UUID], document_type: str = "excel", filename: str = "", domain_data: list[dict] | None = None, tenant_id: str = "", node_id: str = "", ) -> int: """Index records with pre-computed embedding vectors in Milvus. Use this method when embeddings have already been generated (e.g., by GenerateEmbeddingsOperator in Airflow). Skips the embedding generation step for efficiency. Args: records: List of ExtractedRecord objects (Excel) or chunk dicts (Word) vectors: Pre-computed embedding vectors (3072 dims each) document_id: Document UUID for metadata record_ids: PostgreSQL record IDs from StructuredStore document_type: "excel" or "word" (default: "excel") filename: Document filename (required for Word documents) domain_data: Optional list of dicts with domain_category and domain_terms tenant_id: Tenant UUID string (required for data isolation) node_id: Node UUID string (required for data isolation) Returns: Total number of vectors indexed Raises: IndexingError: If indexing fails ValueError: If records, vectors, and record_ids lengths don't match """ if len(records) != len(vectors) or len(records) != len(record_ids): raise ValueError( f"Records ({len(records)}), vectors ({len(vectors)}), and " f"record_ids ({len(record_ids)}) counts must match" ) if not records: self.logger.info("no_records_to_index") return 0 # Ensure collection exists before indexing await self.create_collection() self._connect() start_time = time.perf_counter() try: # Prepare Milvus entities with pre-computed vectors entities = self._prepare_milvus_entities( records, vectors, record_ids, document_id, document_type, filename, domain_data, tenant_id=tenant_id, node_id=node_id, ) # Upload to Milvus with retry await self._upsert_vectors(entities) total_indexed = len(entities) total_duration_s = time.perf_counter() - start_time self.logger.info( "indexing_with_vectors_complete", total_records=total_indexed, duration_seconds=round(total_duration_s, 2), records_per_second=round(total_indexed / total_duration_s, 2) if total_duration_s > 0 else 0, ) return total_indexed except Exception as e: self.logger.error( "indexing_with_vectors_failed", total_records=len(records), error=str(e), ) raise IndexingError( f"Indexing with pre-computed vectors failed: {e}", details={ "total_records": len(records), "error": str(e), }, ) from e def _prepare_milvus_entities( self, records: list[ExtractedRecord] | list[dict], embeddings: list[list[float]], record_ids: list[UUID], document_id: UUID, document_type: str = "excel", filename: str = "", domain_data: list[dict] | None = None, tenant_id: str = "", node_id: str = "", ) -> list[MilvusEntity]: """Prepare Milvus entities from records and embeddings. Supports both Excel ExtractedRecords and Word chunk dictionaries. Args: records: List of ExtractedRecord objects (Excel) or chunk dicts (Word) embeddings: List of embedding vectors (3072 dims each) record_ids: PostgreSQL record UUIDs document_id: Document UUID document_type: "excel" or "word" filename: Document filename (required for Word docs) domain_data: Optional list of dicts with domain_category and domain_terms tenant_id: Tenant UUID string for data isolation node_id: Node UUID string for data isolation Returns: List of MilvusEntity objects ready for upload """ entities = [] for idx, (record, embedding, record_id) in enumerate(zip(records, embeddings, record_ids)): # Extract domain data for this record (if provided) domain_category = "general" domain_terms = [] tags = [] if domain_data and idx < len(domain_data): domain_category = domain_data[idx].get("domain_category", "general") domain_terms = [ _truncate_to_byte_length(t, 100) for t in domain_data[idx].get("domain_terms", []) ] if document_type == "excel": # Excel record - ExtractedRecord object has_symbols = hasattr(record, "resolved_content") and bool( record.resolved_content ) text_content = record_to_text(record) # Milvus VARCHAR limit is 65535 bytes # For Excel cells, use 8000 char limit (8000 × 3 = 24000 bytes worst case) if len(text_content) > 8000: text_content = text_content[:7997] + "..." # Start with base metadata metadata = { "sheet_name": record._source["sheet"], "row_number": record._source["row"], "has_symbols": has_symbols, } # Add cell colors if available if "cell_colors" in record._source: metadata["cell_colors"] = record._source["cell_colors"] # Add record_json for standard table records (LARGE TABLE pipeline only) # Only add for records WITHOUT _legacy_metadata (excludes DETAILED pipeline: images, flowcharts, shapes) if not (hasattr(record, '_legacy_metadata') and record._legacy_metadata): # Validate JSON serialization and size try: record_json_str = json.dumps(record.content, ensure_ascii=False, default=str) json_size_kb = len(record_json_str.encode('utf-8')) / 1024 # Milvus JSON field limit: 65KB (leave margin for safety) if json_size_kb > 50: logger.warning( f"record_json_too_large", sheet=record._source["sheet"], row=record._source["row"], size_kb=round(json_size_kb, 1), message="Truncating to first 20 fields" ) # Store first 20 fields as fallback truncated = dict(list(record.content.items())[:20]) metadata['record_json'] = truncated metadata['record_json_truncated'] = True else: metadata['record_json'] = record.content except (TypeError, ValueError) as e: logger.warning( f"record_json_serialization_error", sheet=record._source["sheet"], row=record._source["row"], error=str(e), message="Storing as string" ) metadata['record_json'] = str(record.content) # Preserve ALL detailed metadata from DETAILED pipeline if available if hasattr(record, '_legacy_metadata') and record._legacy_metadata: legacy_meta = record._legacy_metadata element_type = legacy_meta.get('element_type', 'Text') metadata['element_type'] = element_type # Remove row_number for ALL DETAILED pipeline records # (DETAILED chunks are not table rows, they're images/flowcharts/shapes/text-by-sheet) metadata.pop('row_number', None) metadata.pop('has_symbols', None) # Get full metadata from DETAILED pipeline full_metadata = legacy_meta.get('full_metadata', {}) if full_metadata: for key, value in full_metadata.items(): # Skip fields already added, internal fields, or image_base64 (too large for Milvus) if key not in ['sheet_name', 'sheet_number', 'image_base64'] and not key.startswith('_'): metadata[key] = value file_name = record._source["filename"] # Extract tags from record if available (normalize to simple strings) raw_tags = getattr(record, "tags", []) or [] tags = normalize_tags(raw_tags) else: # Word record - dict from SemanticChunker text_content = record.get("text", "") # Milvus VARCHAR limit is 65535 bytes (not characters) # Use 1.5x factor for mixed ASCII/multibyte text # 65535 / 1.5 = 43690, use 43000 for safety margin if len(text_content) > 43000: text_content = text_content[:42997] + "..." # Extract Word-specific metadata word_metadata = record.get("metadata", {}) metadata = { "page_number": word_metadata.get("page_number"), "contains_table": word_metadata.get("contains_table", False), } file_name = filename # Extract tags from record if available (normalize to simple strings) raw_tags = record.get("tags", []) or [] tags = normalize_tags(raw_tags) # Ensure metadata fits within Milvus JSON field limit (65536 bytes) metadata = truncate_metadata_to_fit(metadata) entity = MilvusEntity( vector=embedding, record_id=str(record_id), document_id=str(document_id), tenant_id=tenant_id, node_id=node_id, filename=file_name, document_type=document_type, text_content=text_content, metadata=metadata, domain_category=domain_category, domain_terms=domain_terms, tags=tags, ) entities.append(entity) return entities async def _upsert_vectors( self, entities: list[MilvusEntity], max_retries: int = 3, batch_size: int = 500, ) -> None: """Upload vectors to Milvus with batching and retry logic. Inserts in batches to avoid exceeding the gRPC 64MB message limit. Args: entities: List of MilvusEntity objects to upload max_retries: Number of retries for transient errors batch_size: Number of entities per insert batch Raises: IndexingError: If upload fails after all retries """ collection = Collection(self.collection_name) total_batches = (len(entities) - 1) // batch_size + 1 self.logger.info( "inserting_vectors", total=len(entities), batch_size=batch_size, batches=total_batches, ) for batch_num in range(total_batches): start = batch_num * batch_size end = min(start + batch_size, len(entities)) batch = entities[start:end] # Prepare data in columnar format for Milvus with unified schema # Note: Exclude 'id' field as it's auto_id=True # Order must match collection schema field order data = [ [e.vector for e in batch], [e.record_id for e in batch], [e.document_id for e in batch], [e.filename for e in batch], [e.document_type for e in batch], [e.text_content for e in batch], [e.metadata for e in batch], [e.domain_category for e in batch], [e.domain_terms for e in batch], [e.tags for e in batch], [e.tenant_id for e in batch], [e.node_id for e in batch], ] for attempt in range(max_retries): try: collection.insert(data) self.logger.debug( "batch_inserted", batch=batch_num + 1, total_batches=total_batches, count=len(batch), attempt=attempt + 1, ) break except Exception as e: is_last_attempt = attempt == max_retries - 1 if is_last_attempt: self.logger.error( "insert_failed_permanently", batch=batch_num + 1, count=len(batch), error=str(e), attempts=max_retries, ) raise IndexingError( f"Failed to insert batch {batch_num + 1} after {max_retries} attempts: {e}", details={ "batch": batch_num + 1, "count": len(batch), "attempts": max_retries, "error": str(e), }, ) from e backoff_seconds = 2**attempt self.logger.warning( "insert_attempt_failed_retrying", batch=batch_num + 1, count=len(batch), error=str(e), attempt=attempt + 1, max_retries=max_retries, backoff_seconds=backoff_seconds, ) await asyncio.sleep(backoff_seconds) # Flush once after all batches to make data available for queries collection.flush() self.logger.info( "all_vectors_inserted", total=len(entities), batches=total_batches, ) def __enter__(self): """Context manager entry.""" self._connect() return self def __exit__(self, exc_type, exc_val, exc_tb): """Context manager exit.""" self._disconnect() async def delete_vectors_by_document_id( document_id: UUID, tenant_id: str = "", node_id: str = "", ) -> int: """Delete all vectors for a document from Milvus. Args: document_id: Document UUID to filter vectors. tenant_id: Tenant UUID for scoped deletion (defense-in-depth). node_id: Node UUID for scoped deletion (defense-in-depth). Returns: Number of vectors deleted. Raises: Exception: If Milvus deletion fails. """ try: # Connect to Milvus connections.connect( alias="default", host=settings.milvus_host, port=settings.milvus_port, ) collection = Collection(settings.milvus_collection) # Delete by document_id, scoped by tenant + node if provided expr = f'document_id == "{str(document_id)}"' if tenant_id: expr += f' and tenant_id == "{tenant_id}"' if node_id: expr += f' and node_id == "{node_id}"' result = collection.delete(expr) logger.info( "milvus_vectors_deleted", document_id=str(document_id), tenant_id=tenant_id, node_id=node_id, delete_count=result.delete_count, ) # Disconnect connections.disconnect("default") return result.delete_count except Exception as e: logger.error( "milvus_deletion_failed", document_id=str(document_id), error=str(e), ) raise