Enhance data_types.py with vector database optimized pydantic models

[manavgup]

📋 Overview

Enhance the pydantic data models in vectordbs/data_types.py to provide vector database optimized structures that eliminate manual parsing and improve type safety across all vector database implementations.

🎯 Goals

• Eliminate manual dict parsing in vector DB implementations
• Provide type-safe request/response models
• Enable better error handling with structured responses
• Standardize serialization/deserialization patterns
• Improve developer experience with IDE completion and validation

🔧 Technical Specifications

New Pydantic Models to Add

1. EmbeddedChunk Model

class EmbeddedChunk(DocumentChunk):
    \"\"\"Chunk guaranteed to have embeddings for vector DB storage\"\"\"
    embeddings: List[float]  # Required, not optional
    
    @classmethod
    def from_chunk(cls, chunk: DocumentChunk) -> 'EmbeddedChunk':
        \"\"\"Convert DocumentChunk to EmbeddedChunk with validation\"\"\"
        
    def to_vector_metadata(self) -> Dict[str, Any]:
        \"\"\"Serialize metadata for vector DB storage\"\"\"
        
    def to_vector_db(self) -> Dict[str, Any]:
        \"\"\"Complete serialization for vector DB insertion\"\"\"

2. Request/Response Models

class DocumentIngestionRequest(BaseModel):
    \"\"\"Request for adding documents to vector DB\"\"\"
    collection_name: str
    documents: List[Document] 
    batch_size: Optional[int] = 100
    
    def get_embedded_chunks(self) -> List[EmbeddedChunk]:
        \"\"\"Extract all chunks that have embeddings\"\"\"

class VectorSearchRequest(BaseModel):
    \"\"\"Standardized search request\"\"\"
    collection_name: str
    query: Union[str, List[float]]  # Text or embedding vector
    filters: Optional[DocumentMetadataFilter] = None
    limit: int = 10
    include_metadata: bool = True
    
class CollectionConfig(BaseModel):
    \"\"\"Vector DB collection configuration\"\"\"
    name: str
    dimension: int
    metric: VectorMetric = VectorMetric.COSINE
    cloud_provider: Optional[str] = None
    region: Optional[str] = None
    
    def validate_for_vector_db(self, db_type: str) -> None:
        \"\"\"Validate configuration for specific vector DB type\"\"\"

3. Generic Response Models

class VectorDBResponse(BaseModel, Generic[T]):
    \"\"\"Standardized response wrapper for all vector DB operations\"\"\"
    success: bool
    data: Optional[T] = None
    error: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    
    @classmethod
    def success(cls, data: T, metadata: Optional[Dict] = None) -> 'VectorDBResponse[T]':
        \"\"\"Create success response\"\"\"
        
    @classmethod 
    def error(cls, message: str, metadata: Optional[Dict] = None) -> 'VectorDBResponse[T]':
        \"\"\"Create error response\"\"\"

# Type aliases for common responses
DocumentIngestionResponse = VectorDBResponse[List[str]]
CollectionResponse = VectorDBResponse[str] 
SearchResponse = VectorDBResponse[List[QueryResult]]
HealthCheckResponse = VectorDBResponse[Dict[str, Any]]

✅ Acceptance Criteria

Functional Requirements

• All new pydantic models validate correctly with proper type hints
• EmbeddedChunk enforces non-null embeddings at creation time
• Request models provide convenient methods for data extraction
• Response models support both success and error states
• Serialization methods produce vector DB compatible formats
• Validation methods catch configuration errors early

Technical Requirements

• All models inherit from BaseModel with proper type annotations
• Use pydantic v2 features (Field, computed fields, validators)
• Include comprehensive docstrings with examples
• Support JSON serialization/deserialization
• Maintain backward compatibility with existing Document/DocumentChunk usage
• Add proper repr methods for debugging

Testing Requirements

• Unit tests for all new models and methods
• Validation testing for edge cases and error conditions
• Serialization/deserialization round-trip tests
• Performance tests for large document batches
• Integration tests with existing codebase

🔄 Implementation Details

File Changes

• vectordbs/data_types.py - Primary implementation
• tests/unit/test_data_types.py - Comprehensive test coverage
• tests/integration/test_vector_models.py - Integration testing

Dependencies

• Pydantic v2
• Python 3.12+ type hints
• No breaking changes to existing models

Migration Strategy

• All new models are additive - existing code continues to work
• Provide utility functions to convert between old and new patterns
• Gradual adoption in downstream code

🧪 Testing Strategy

Unit Tests

def test_embedded_chunk_validation():
    \"\"\"Test EmbeddedChunk requires embeddings\"\"\"
    
def test_ingestion_request_chunk_filtering():
    \"\"\"Test DocumentIngestionRequest filters embedded chunks correctly\"\"\"
    
def test_response_model_serialization():
    \"\"\"Test VectorDBResponse serializes properly\"\"\"
    
def test_collection_config_validation():
    \"\"\"Test CollectionConfig validates for different vector DBs\"\"\"

📊 Success Metrics

• Zero manual dict parsing in vector DB implementations after adoption
• 100% type coverage with mypy
• <100ms serialization time for 1000 document chunks
• Backward compatibility maintained for all existing usages

🔗 Related Issues

• Depends on: None (foundational work)
• Blocks: Enhanced VectorStore Base Class (TBD)
• Blocks: Vector DB Implementation Refactoring (TBD)

Priority: High
Estimated Effort: Medium (3-5 days)
Risk Level: Low (additive changes only)

[manavgup]

@claude please reivew this issue and the codebase and build a detailed implementation plan. See also #212 and #213 and figure out the best path for implementation.

[claude]

Claude finished @manavgup's task —— View job

---

📋 Implementation Plan for Enhanced Data Types

I've reviewed the codebase and identified the current pain points and opportunities for improvement. Here's a comprehensive implementation plan:

✅ Completed Analysis:

• Review current vectordbs/data_types.py implementation
• Examine vector_store.py base class to understand usage patterns
• Review existing vector DB implementations for manual parsing patterns
• Analyze dependencies and relationships
• Build detailed implementation plan with proper ordering
• Provide recommendations for implementation approach

🔍 Current State Analysis

Manual Parsing Issues Identified:

1. MilvusStore: Lines 172-194 manually extract fields from documents/chunks for insertion

   • Manual creation of lists for each field
   • Error-prone field extraction with nested conditions
   • No standardized serialization pattern

2. ChromaDBStore: Lines 87-98 manually build metadata dictionaries

   • Direct dict construction with conditional logic
   • Inconsistent handling of optional fields
   • Manual type conversions

3. Search Result Processing: Both stores manually construct QueryResult objects

   • Direct attribute access with getattr()
   • Manual score extraction and conversion
   • Redundant code patterns across implementations

📐 Detailed Implementation Plan

Phase 1: Core Model Enhancements (Priority: High)

1.1 EmbeddedChunk Model

# Location: backend/vectordbs/data_types.py

class EmbeddedChunk(DocumentChunk):
    """Chunk guaranteed to have embeddings for vector DB storage"""
    
    # Override to make embeddings required
    embeddings: Embeddings  # Not optional
    
    @field_validator('embeddings')
    def validate_embeddings(cls, v):
        """Ensure embeddings are not empty"""
        if not v or len(v) == 0:
            raise ValueError("EmbeddedChunk requires non-empty embeddings")
        return v
    
    @classmethod
    def from_chunk(cls, chunk: DocumentChunk) -> 'EmbeddedChunk':
        """Convert DocumentChunk to EmbeddedChunk with validation"""
        if not chunk.embeddings:
            raise ValueError("Cannot create EmbeddedChunk without embeddings")
        return cls(**chunk.model_dump())
    
    def to_vector_metadata(self) -> dict[str, Any]:
        """Serialize metadata for vector DB storage"""
        metadata = {}
        if self.metadata:
            metadata.update({
                'source': str(self.metadata.source),
                'document_id': self.metadata.document_id or '',
                'page_number': self.metadata.page_number or 0,
                'chunk_number': self.metadata.chunk_number or 0,
            })
        return metadata
    
    def to_vector_db(self) -> dict[str, Any]:
        """Complete serialization for vector DB insertion"""
        return {
            'chunk_id': self.chunk_id,
            'text': self.text,
            'embeddings': self.embeddings,
            'metadata': self.to_vector_metadata(),
            'document_id': self.document_id
        }

1.2 Request Models

class DocumentIngestionRequest(BaseModel):
    """Request for adding documents to vector DB"""
    collection_name: str
    documents: list[Document]
    batch_size: int = 100
    
    def get_embedded_chunks(self) -> list[EmbeddedChunk]:
        """Extract all chunks that have embeddings"""
        embedded = []
        for doc in self.documents:
            for chunk in doc.chunks:
                if chunk.embeddings:
                    embedded.append(EmbeddedChunk.from_chunk(chunk))
        return embedded
    
    def batch_chunks(self) -> Iterator[list[EmbeddedChunk]]:
        """Yield batches of embedded chunks"""
        chunks = self.get_embedded_chunks()
        for i in range(0, len(chunks), self.batch_size):
            yield chunks[i:i + self.batch_size]

class VectorSearchRequest(BaseModel):
    """Standardized search request"""
    collection_name: str
    query: str | list[float]
    filters: DocumentMetadataFilter | None = None
    limit: int = 10
    include_metadata: bool = True
    
    @property
    def is_text_query(self) -> bool:
        return isinstance(self.query, str)
    
    @property
    def is_vector_query(self) -> bool:
        return isinstance(self.query, list)

Phase 2: Configuration Models (Priority: High)

2.1 Collection Configuration

class VectorMetric(str, Enum):
    """Supported vector similarity metrics"""
    COSINE = "cosine"
    L2 = "l2"
    IP = "ip"  # Inner product

class CollectionConfig(BaseModel):
    """Vector DB collection configuration"""
    name: str
    dimension: int
    metric: VectorMetric = VectorMetric.COSINE
    cloud_provider: str | None = None
    region: str | None = None
    index_type: str | None = None
    
    def validate_for_milvus(self) -> None:
        """Validate configuration for Milvus"""
        if self.dimension < 1 or self.dimension > 32768:
            raise ValueError("Milvus dimension must be between 1 and 32768")
        if self.index_type and self.index_type not in ["IVF_FLAT", "IVF_SQ8", "HNSW"]:
            raise ValueError(f"Unsupported index type for Milvus: {self.index_type}")
    
    def validate_for_chroma(self) -> None:
        """Validate configuration for ChromaDB"""
        if self.dimension < 1:
            raise ValueError("ChromaDB dimension must be positive")
    
    def validate_for_vector_db(self, db_type: str) -> None:
        """Validate configuration for specific vector DB type"""
        validators = {
            'milvus': self.validate_for_milvus,
            'chroma': self.validate_for_chroma,
            # Add more as needed
        }
        if validator := validators.get(db_type.lower()):
            validator()

Phase 3: Response Models (Priority: Medium)

3.1 Generic Response Wrapper

from typing import Generic, TypeVar

T = TypeVar('T')

class VectorDBResponse(BaseModel, Generic[T]):
    """Standardized response wrapper for all vector DB operations"""
    success: bool
    data: T | None = None
    error: str | None = None
    metadata: dict[str, Any] | None = None
    
    @classmethod
    def success_response(cls, data: T, metadata: dict | None = None) -> 'VectorDBResponse[T]':
        """Create success response"""
        return cls(success=True, data=data, metadata=metadata)
    
    @classmethod
    def error_response(cls, message: str, metadata: dict | None = None) -> 'VectorDBResponse[T]':
        """Create error response"""
        return cls(success=False, error=message, metadata=metadata)
    
    def unwrap(self) -> T:
        """Unwrap data or raise exception"""
        if not self.success or self.data is None:
            raise ValueError(f"Failed operation: {self.error}")
        return self.data

# Type aliases
DocumentIngestionResponse = VectorDBResponse[list[str]]
CollectionResponse = VectorDBResponse[str]
SearchResponse = VectorDBResponse[list[QueryResult]]
HealthCheckResponse = VectorDBResponse[dict[str, Any]]

Phase 4: Vector Store Integration (Priority: High)

4.1 Update MilvusStore to Use New Models

# backend/vectordbs/milvus_store.py

def add_documents(self, collection_name: str, documents: list[Document]) -> list[str]:
    """Simplified version using new models"""
    request = DocumentIngestionRequest(
        collection_name=collection_name,
        documents=documents
    )
    
    collection = self._get_collection(collection_name)
    document_ids = []
    
    for batch in request.batch_chunks():
        # Use standardized serialization
        data = defaultdict(list)
        for chunk in batch:
            serialized = chunk.to_vector_db()
            for key, value in serialized.items():
                data[key].append(value)
        
        collection.insert(list(data.values()))
        document_ids.extend([c.document_id for c in batch])
    
    collection.flush()
    return document_ids

📝 Implementation Sequence

1. Week 1: Core Models

   • Day 1-2: Implement EmbeddedChunk and validation logic
   • Day 3: Implement Request/Response models
   • Day 4: Write comprehensive unit tests
   • Day 5: Integration testing with existing code

2. Week 2: Vector Store Updates

   • Day 1-2: Update MilvusStore to use new models
   • Day 3: Update ChromaDBStore
   • Day 4: Update other vector stores
   • Day 5: End-to-end testing

3. Week 3: Documentation & Migration

   • Day 1: Write migration guide
   • Day 2: Update API documentation
   • Day 3-5: Support gradual migration in services

🎯 Recommendations

1. Start with EmbeddedChunk: This is the foundation that eliminates most manual parsing
2. Use Builder Pattern: For complex models, consider builder methods for cleaner construction
3. Maintain Backward Compatibility: Keep existing models, add conversion methods
4. Incremental Adoption: Update one vector store at a time to minimize risk
5. Performance Testing: Benchmark serialization performance with large datasets

⚠️ Risk Mitigation

• All new models are additive - no breaking changes
• Provide conversion utilities between old and new patterns
• Comprehensive test coverage before any integration
• Feature flag for gradual rollout if needed

📊 Success Metrics

• ✅ Zero manual dict construction in vector stores
• ✅ 100% type coverage with mypy
• ✅ <100ms serialization for 1000 chunks
• ✅ Reduced LOC in vector store implementations by 30%

---

Note: Since issues #212 and #213 couldn't be accessed directly, this plan focuses on the requirements from issue #211. The implementation is designed to be flexible enough to accommodate related enhancements from those issues when they're available.

Summary: This implementation plan provides a structured approach to enhance the pydantic models in vectordbs/data_types.py, eliminating manual parsing and improving type safety across all vector database implementations. The plan is designed for incremental adoption with zero breaking changes.