gmf_forge_ai_data.vector_stores

181class BaseVectorStore(ABC):
182    """
183    Abstract base class for vector store implementations.
184    
185    Vector stores provide storage and retrieval of document embeddings,
186    supporting various search strategies (vector similarity, keyword, hybrid).
187    
188    Implementations must provide:
189    - Document storage with embeddings
190    - Vector similarity search
191    - Optional keyword/hybrid search
192    - Document lifecycle management (add, update, delete)
193    """
194    
195    @abstractmethod
196    def add_documents(
197        self,
198        documents: List[Document],
199        generate_embeddings: bool = True
200    ) -> List[str]:
201        """
202        Add documents to the vector store.
203        
204        Args:
205            documents: List of Document objects to add
206            generate_embeddings: If True and documents lack embeddings,
207                                generate them automatically
208        
209        Returns:
210            List of document IDs that were added
211            
212        Raises:
213            ValueError: If documents are invalid or embeddings cannot be generated
214        """
215        pass
216    
217    @abstractmethod
218    def search(
219        self,
220        query: Optional[str] = None,
221        query_embedding: Optional[List[float]] = None,
222        top_k: int = 5,
223        filters: Optional[Dict[str, Any]] = None,
224        search_type: str = "vector"
225    ) -> List[SearchResult]:
226        """
227        Search the vector store.
228        
229        Args:
230            query: Text query (required for keyword/hybrid search)
231            query_embedding: Vector embedding of the query (required for vector search)
232            top_k: Number of results to return
233            filters: Metadata filters to apply (e.g., {"source": "docs.pdf"})
234            search_type: Type of search - "vector", "keyword", or "hybrid"
235        
236        Returns:
237            List of SearchResult objects, ordered by relevance
238            
239        Raises:
240            ValueError: If required parameters are missing for the search type
241        """
242        pass
243    
244    @abstractmethod
245    def delete_documents(self, document_ids: List[str]) -> int:
246        """
247        Delete documents from the vector store.
248        
249        Args:
250            document_ids: List of document IDs to delete
251        
252        Returns:
253            Number of documents successfully deleted
254        """
255        pass
256    
257    @abstractmethod
258    def update_document(
259        self,
260        document_id: str,
261        content: Optional[str] = None,
262        embedding: Optional[List[float]] = None,
263        metadata: Optional[Dict[str, Any]] = None
264    ) -> bool:
265        """
266        Update an existing document.
267        
268        Args:
269            document_id: ID of the document to update
270            content: New content (if provided)
271            embedding: New embedding (if provided)
272            metadata: New metadata (merged with existing)
273        
274        Returns:
275            True if update was successful, False if document not found
276        """
277        pass
278    
279    @abstractmethod
280    def get_document(self, document_id: str) -> Optional[Document]:
281        """
282        Retrieve a document by ID.
283        
284        Args:
285            document_id: ID of the document to retrieve
286        
287        Returns:
288            Document if found, None otherwise
289        """
290        pass
291    
292    @abstractmethod
293    def count(self) -> int:
294        """
295        Get the total number of documents in the vector store.
296        
297        Returns:
298            Total document count
299        """
300        pass
301    
302    @abstractmethod
303    def clear(self) -> None:
304        """
305        Remove all documents from the vector store.
306        
307        Warning: This operation is irreversible.
308        """
309        pass

 19@dataclass
 20class Document:
 21    """
 22    Represents a document or chunk with its content and metadata.
 23    
 24    This class is designed to be extensible - you can subclass it to add
 25    custom fields specific to your application domain.
 26    
 27    Base Attributes:
 28        id: Unique identifier for the document
 29        content: The text content of the document
 30        embedding: Optional vector embedding (list of floats)
 31        timestamp: Optional timestamp for time-weighted retrieval
 32        metadata: Flexible dictionary for custom properties
 33    
 34    Design Philosophy:
 35        The base Document keeps only truly universal fields. Domain-specific
 36        fields (source, page_number, category, etc.) should be added via
 37        inheritance to match your application's data model.
 38    
 39    Example - Using base Document with metadata:
 40        ```python
 41        doc = Document(
 42            id="doc_1",
 43            content="AI is transforming software development...",
 44            timestamp=datetime.now(),
 45            metadata={
 46                "source": "blog.pdf",
 47                "page_number": 5,
 48                "author": "John Doe",
 49                "category": "AI"
 50            }
 51        )
 52        ```
 53    
 54    Example - Custom Document schema via inheritance:
 55        ```python
 56        from dataclasses import dataclass
 57        from gmf_forge_ai_data.vector_stores import Document
 58        
 59        @dataclass
 60        class LegalDocument(Document):
 61            case_number: str
 62            court: str
 63            decision_date: datetime
 64            jurisdiction: str
 65            source: str = ""  # Add source as typed field
 66            page_number: Optional[int] = None
 67            
 68        legal_doc = LegalDocument(
 69            id="case_123",
 70            content="In the matter of...",
 71            case_number="2024-CV-12345",
 72            court="Supreme Court",
 73            decision_date=datetime(2024, 1, 15),
 74            jurisdiction="Federal",
 75            source="court_records.pdf",
 76            page_number=42
 77        )
 78        ```
 79    
 80    Example - E-commerce Product Document:
 81        ```python
 82        @dataclass
 83        class ProductDocument(Document):
 84            sku: str
 85            category: str
 86            price: float
 87            in_stock: bool
 88            brand: str
 89            # No source/page_number - not applicable to products
 90            
 91        product = ProductDocument(
 92            id="prod_456",
 93            content="High-performance laptop with...",
 94            sku="LAP-001",
 95            category="Electronics",
 96            price=1299.99,
 97            in_stock=True,
 98            brand="TechCorp"
 99        )
100        ```
101    """
102    id: str
103    content: str
104    embedding: Optional[List[float]] = None
105    timestamp: Optional[datetime] = None
106    metadata: Dict[str, Any] = field(default_factory=dict)
107    
108    def to_dict(self) -> Dict[str, Any]:
109        """
110        Convert Document to dictionary for serialization.
111        
112        Returns:
113            Dictionary representation of the document
114        """
115        doc_dict = asdict(self)
116        # Convert datetime to ISO string for serialization
117        if self.timestamp:
118            doc_dict['timestamp'] = self.timestamp.isoformat()
119        return doc_dict
120    
121    @classmethod
122    def from_dict(cls: Type[T], data: Dict[str, Any]) -> T:
123        """
124        Create Document from dictionary.
125        
126        Args:
127            data: Dictionary containing document fields
128        
129        Returns:
130            Document instance (or subclass instance)
131        """
132        # Handle datetime deserialization
133        if 'timestamp' in data and isinstance(data['timestamp'], str):
134            data['timestamp'] = datetime.fromisoformat(data['timestamp'])
135        
136        # Filter to only fields that exist in the dataclass
137        import inspect
138        valid_fields = set(inspect.signature(cls).parameters.keys())
139        filtered_data = {k: v for k, v in data.items() if k in valid_fields}
140        
141        return cls(**filtered_data)
142    
143    def update_metadata(self, **kwargs) -> None:
144        """
145        Update metadata fields.
146        
147        Args:
148            **kwargs: Key-value pairs to add/update in metadata
149        """
150        self.metadata.update(kwargs)
151    
152    def get_metadata(self, key: str, default: Any = None) -> Any:
153        """
154        Get a metadata value.
155        
156        Args:
157            key: Metadata key to retrieve
158            default: Default value if key not found
159        
160        Returns:
161            Metadata value or default
162        """
163        return self.metadata.get(key, default)

166@dataclass
167class SearchResult:
168    """
169    Represents a search result from a vector store.
170    
171    Attributes:
172        document: The retrieved document
173        score: Similarity score (0.0 to 1.0, higher is better)
174        rank: Position in the result list (0-indexed)
175    """
176    document: Document
177    score: float
178    rank: int

 17class InMemoryVectorStore(BaseVectorStore):
 18    """
 19    In-memory vector store using numpy for vector operations.
 20    
 21    Features:
 22    - Fast cosine similarity search using numpy
 23    - Keyword search using simple text matching
 24    - Hybrid search combining vector and keyword scores
 25    - Metadata filtering
 26    - No external dependencies (Azure, etc.)
 27    
 28    Ideal for:
 29    - Unit testing
 30    - Development and prototyping
 31    - Small datasets (< 10,000 documents)
 32    - CI/CD pipelines
 33    
 34    Note: All data is lost when the process terminates.
 35    """
 36    
 37    def __init__(self, embedding_dimension: int = 1536):
 38        """
 39        Initialize the in-memory vector store.
 40        
 41        Args:
 42            embedding_dimension: Expected dimension of embeddings (default 1536 for text-embedding-ada-002)
 43        """
 44        self.embedding_dimension = embedding_dimension
 45        self._documents: Dict[str, Document] = {}
 46        self._embeddings: Optional[np.ndarray] = None
 47        self._doc_ids: List[str] = []
 48    
 49    def add_documents(
 50        self,
 51        documents: List[Document],
 52        generate_embeddings: bool = True
 53    ) -> List[str]:
 54        """
 55        Add documents to the in-memory store.
 56        
 57        Args:
 58            documents: List of Document objects to add
 59            generate_embeddings: If True, validates embeddings are present
 60                                (actual generation should be done externally)
 61        
 62        Returns:
 63            List of document IDs that were added
 64        """
 65        if not documents:
 66            return []
 67        
 68        added_ids = []
 69        
 70        for doc in documents:
 71            # Validate document
 72            if not doc.id or not doc.content:
 73                raise ValueError(f"Document must have id and content: {doc}")
 74            
 75            if generate_embeddings and doc.embedding is None:
 76                raise ValueError(
 77                    f"Document {doc.id} lacks embedding. "
 78                    "Generate embeddings externally before adding to store."
 79                )
 80            
 81            if doc.embedding is not None and len(doc.embedding) != self.embedding_dimension:
 82                raise ValueError(
 83                    f"Document {doc.id} has embedding dimension {len(doc.embedding)}, "
 84                    f"expected {self.embedding_dimension}"
 85                )
 86            
 87            # Store document
 88            self._documents[doc.id] = deepcopy(doc)
 89            added_ids.append(doc.id)
 90        
 91        # Rebuild embedding matrix
 92        self._rebuild_embeddings()
 93        
 94        return added_ids
 95    
 96    def search(
 97        self,
 98        query: Optional[str] = None,
 99        query_embedding: Optional[List[float]] = None,
100        top_k: int = 5,
101        filters: Optional[Dict[str, Any]] = None,
102        search_type: str = "vector"
103    ) -> List[SearchResult]:
104        """
105        Search the in-memory vector store.
106        
107        Args:
108            query: Text query (required for keyword/hybrid search)
109            query_embedding: Vector embedding of the query (required for vector/hybrid search)
110            top_k: Number of results to return
111            filters: Metadata filters (e.g., {"source": "doc.pdf"})
112            search_type: "vector", "keyword", or "hybrid"
113        
114        Returns:
115            List of SearchResult objects, ordered by relevance
116        """
117        if search_type not in ["vector", "keyword", "hybrid"]:
118            raise ValueError(f"Invalid search_type: {search_type}. Must be 'vector', 'keyword', or 'hybrid'")
119        
120        if search_type in ["vector", "hybrid"] and query_embedding is None:
121            raise ValueError(f"{search_type} search requires query_embedding")
122        
123        if search_type in ["keyword", "hybrid"] and query is None:
124            raise ValueError(f"{search_type} search requires query text")
125        
126        if not self._documents:
127            return []
128        
129        # Apply metadata filters
130        filtered_docs = self._apply_filters(filters)
131        if not filtered_docs:
132            return []
133        
134        # Calculate scores based on search type
135        if search_type == "vector":
136            scores = self._vector_search(query_embedding, filtered_docs)
137        elif search_type == "keyword":
138            scores = self._keyword_search(query, filtered_docs)
139        else:  # hybrid
140            vector_scores = self._vector_search(query_embedding, filtered_docs)
141            keyword_scores = self._keyword_search(query, filtered_docs)
142            # Combine scores (50/50 weighting)
143            scores = {
144                doc_id: 0.5 * vector_scores.get(doc_id, 0.0) + 0.5 * keyword_scores.get(doc_id, 0.0)
145                for doc_id in filtered_docs
146            }
147        
148        # Sort by score and take top_k
149        sorted_results = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:top_k]
150        
151        # Build SearchResult objects
152        results = []
153        for rank, (doc_id, score) in enumerate(sorted_results):
154            results.append(SearchResult(
155                document=deepcopy(self._documents[doc_id]),
156                score=float(score),
157                rank=rank
158            ))
159        
160        return results
161    
162    def delete_documents(self, document_ids: List[str]) -> int:
163        """
164        Delete documents from the store.
165        
166        Args:
167            document_ids: List of document IDs to delete
168        
169        Returns:
170            Number of documents successfully deleted
171        """
172        deleted_count = 0
173        for doc_id in document_ids:
174            if doc_id in self._documents:
175                del self._documents[doc_id]
176                deleted_count += 1
177        
178        # Rebuild embedding matrix
179        self._rebuild_embeddings()
180        
181        return deleted_count
182    
183    def update_document(
184        self,
185        document_id: str,
186        content: Optional[str] = None,
187        embedding: Optional[List[float]] = None,
188        metadata: Optional[Dict[str, Any]] = None
189    ) -> bool:
190        """
191        Update an existing document.
192        
193        Args:
194            document_id: ID of the document to update
195            content: New content (if provided)
196            embedding: New embedding (if provided)
197            metadata: New metadata (merged with existing)
198        
199        Returns:
200            True if update was successful, False if document not found
201        """
202        if document_id not in self._documents:
203            return False
204        
205        doc = self._documents[document_id]
206        
207        if content is not None:
208            doc.content = content
209        
210        if embedding is not None:
211            if len(embedding) != self.embedding_dimension:
212                raise ValueError(
213                    f"Embedding dimension {len(embedding)} != expected {self.embedding_dimension}"
214                )
215            doc.embedding = embedding
216            self._rebuild_embeddings()
217        
218        if metadata is not None:
219            doc.metadata.update(metadata)
220        
221        return True
222    
223    def get_document(self, document_id: str) -> Optional[Document]:
224        """
225        Retrieve a document by ID.
226        
227        Args:
228            document_id: ID of the document to retrieve
229        
230        Returns:
231            Document if found, None otherwise
232        """
233        doc = self._documents.get(document_id)
234        return deepcopy(doc) if doc else None
235    
236    def count(self) -> int:
237        """Get the total number of documents in the store."""
238        return len(self._documents)
239    
240    def clear(self) -> None:
241        """Remove all documents from the store."""
242        self._documents.clear()
243        self._embeddings = None
244        self._doc_ids = []
245    
246    # Private helper methods
247    
248    def _rebuild_embeddings(self) -> None:
249        """Rebuild the embedding matrix from stored documents."""
250        docs_with_embeddings = [
251            (doc_id, doc) for doc_id, doc in self._documents.items()
252            if doc.embedding is not None
253        ]
254        
255        if not docs_with_embeddings:
256            self._embeddings = None
257            self._doc_ids = []
258            return
259        
260        self._doc_ids = [doc_id for doc_id, _ in docs_with_embeddings]
261        embeddings_list = [doc.embedding for _, doc in docs_with_embeddings]
262        self._embeddings = np.array(embeddings_list, dtype=np.float32)
263    
264    def _apply_filters(self, filters: Optional[Dict[str, Any]]) -> List[str]:
265        """
266        Apply metadata filters and return list of matching document IDs.
267        
268        Supports filtering on: source, page_number, timestamp, and metadata fields
269        Special operators: use tuples for range queries
270        - (">=", value) for greater than or equal
271        - ("<=", value) for less than or equal  
272        - ("range", min_val, max_val) for range queries
273        """
274        if filters is None:
275            return list(self._documents.keys())
276        
277        matching_ids = []
278        for doc_id, doc in self._documents.items():
279            matches = True
280            
281            for key, value in filters.items():
282                # Check document attributes (source, page_number, timestamp)
283                doc_value = getattr(doc, key, None)
284                
285                # If not a document attribute, check metadata
286                if doc_value is None and key in doc.metadata:
287                    doc_value = doc.metadata[key]
288                
289                # Skip if field doesn't exist
290                if doc_value is None:
291                    matches = False
292                    break
293                
294                # Handle tuple operators for range queries
295                if isinstance(value, tuple):
296                    operator = value[0]
297                    if operator == "range" and len(value) == 3:
298                        min_val, max_val = value[1], value[2]
299                        if not (min_val <= doc_value <= max_val):
300                            matches = False
301                            break
302                    elif operator in [">=", "gt"]:
303                        if not (doc_value >= value[1]):
304                            matches = False
305                            break
306                    elif operator in ["<=", "lt"]:
307                        if not (doc_value <= value[1]):
308                            matches = False
309                            break
310                # Simple equality check
311                elif doc_value != value:
312                    matches = False
313                    break
314            
315            if matches:
316                matching_ids.append(doc_id)
317        
318        return matching_ids
319    
320    def _vector_search(
321        self,
322        query_embedding: List[float],
323        candidate_ids: List[str]
324    ) -> Dict[str, float]:
325        """
326        Perform vector similarity search using cosine similarity.
327        
328        Returns:
329            Dict mapping document IDs to similarity scores (0.0 to 1.0)
330        """
331        if self._embeddings is None or len(self._doc_ids) == 0:
332            return {}
333        
334        query_vector = np.array(query_embedding, dtype=np.float32)
335        
336        # Filter embeddings to only candidate documents
337        candidate_indices = [
338            i for i, doc_id in enumerate(self._doc_ids)
339            if doc_id in candidate_ids
340        ]
341        
342        if not candidate_indices:
343            return {}
344        
345        candidate_embeddings = self._embeddings[candidate_indices]
346        candidate_doc_ids = [self._doc_ids[i] for i in candidate_indices]
347        
348        # Compute cosine similarity
349        query_norm = np.linalg.norm(query_vector)
350        if query_norm == 0:
351            return {doc_id: 0.0 for doc_id in candidate_doc_ids}
352        
353        doc_norms = np.linalg.norm(candidate_embeddings, axis=1)
354        dot_products = candidate_embeddings.dot(query_vector)
355        
356        # Avoid division by zero
357        similarities = np.zeros(len(candidate_doc_ids))
358        valid_mask = doc_norms > 0
359        similarities[valid_mask] = dot_products[valid_mask] / (doc_norms[valid_mask] * query_norm)
360        
361        # Convert from [-1, 1] to [0, 1]
362        similarities = (similarities + 1) / 2
363        
364        return {doc_id: float(sim) for doc_id, sim in zip(candidate_doc_ids, similarities)}
365    
366    def _keyword_search(
367        self,
368        query: str,
369        candidate_ids: List[str]
370    ) -> Dict[str, float]:
371        """
372        Perform simple keyword search using term matching.
373        
374        Returns:
375            Dict mapping document IDs to relevance scores (0.0 to 1.0)
376        """
377        query_lower = query.lower()
378        query_terms = set(query_lower.split())
379        
380        scores = {}
381        for doc_id in candidate_ids:
382            doc = self._documents[doc_id]
383            content_lower = doc.content.lower()
384            content_terms = set(content_lower.split())
385            
386            # Calculate Jaccard similarity as relevance score
387            if not query_terms:
388                scores[doc_id] = 0.0
389            else:
390                intersection = len(query_terms & content_terms)
391                union = len(query_terms | content_terms)
392                scores[doc_id] = intersection / union if union > 0 else 0.0
393        
394        return scores

 43class AzureAISearchVectorStore(BaseVectorStore):
 44    """
 45    Azure AI Search vector store for production RAG pipelines.
 46    
 47    Features:
 48    - Scalable vector search with HNSW (Hierarchical Navigable Small World)
 49    - Hybrid search (vector + keyword BM25)
 50    - Semantic ranking powered by Microsoft Bing
 51    - Automatic indexing of all document fields (universal + custom)
 52    - Efficient filtering on all indexed fields
 53    - High availability and disaster recovery
 54    - Managed service (no infrastructure management)
 55    
 56    Prerequisites:
 57    - Azure AI Search service (Basic tier or higher for vector search)
 58    - Service endpoint and API key **or** a managed-identity token provider
 59    - Index created with vector field configuration
 60    
 61    Usage:
 62    
 63    **Simple Mode** (base Document):
 64        ```python
 65        # API key auth
 66        store = AzureAISearchVectorStore(
 67            endpoint="https://my-search.search.windows.net",
 68            index_name="documents",
 69            api_key="...",
 70            embedding_dimension=1536
 71        )
 72
 73        # Managed identity / token provider auth
 74        from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 75        token_provider = get_bearer_token_provider(
 76            DefaultAzureCredential(),
 77            "https://search.azure.com/.default"  # Azure AI Search scope
 78        )
 79        store = AzureAISearchVectorStore(
 80            endpoint="https://my-search.search.windows.net",
 81            index_name="documents",
 82            token_provider=token_provider,
 83            embedding_dimension=1536
 84        )
 85        # Uses base Document with universal fields only
 86        ```
 87    
 88    **Custom Schema Mode** (domain-specific Document):
 89        ```python
 90        from dataclasses import dataclass
 91        from gmf_forge_ai_data.vector_stores import Document
 92        
 93        @dataclass
 94        class LegalDocument(Document):
 95            case_number: str = ""
 96            court: str = ""
 97            jurisdiction: str = ""
 98            decision_date: Optional[datetime] = None
 99        
100        # Pass document_type - all fields automatically indexed.
101        # If the index was provisioned externally (portal, Bicep, etc.) the
102        # vector and content field names will likely differ from the defaults
103        # ("embedding" / "content"). Always set vector_field_name and
104        # content_field_name to match your actual index schema.
105        store = AzureAISearchVectorStore(
106            endpoint="https://my-search.search.windows.net",
107            index_name="legal_documents",
108            api_key="...",
109            embedding_dimension=1536,
110            document_type=LegalDocument,
111            vector_field_name="contentVector",   # match your index schema
112            content_field_name="chunkContent",   # match your index schema
113        )
114        
115        # All fields indexed - filter on any field
116        results = store.search(
117            query_embedding=vector,
118            filters={
119                "jurisdiction": "Federal",
120                "court": "Supreme Court",
121                "decision_date": (">=", datetime(2024, 1, 1))
122            }
123        )
124        ```
125    """
126    
127    def __init__(
128        self,
129        endpoint: str,
130        index_name: str,
131        api_key: Optional[str] = None,
132        token_provider: Optional[Callable[[], str]] = None,
133        embedding_dimension: int = 1536,
134        document_type: type = Document,
135        vector_field_name: str = "embedding",
136        content_field_name: str = "content",
137    ):
138        """
139        Initialize Azure AI Search vector store.
140
141        The index must be pre-provisioned using ``AzureAISearchIndexBuilder``
142        before first use.  The constructor only establishes client connections
143        — it never creates or modifies indexes.
144
145        Exactly one of ``api_key`` or ``token_provider`` must be supplied.
146
147        Args:
148            endpoint: Azure AI Search service endpoint
149            index_name: Name of the search index
150            api_key: Azure AI Search API key. Use for local development or
151                when managed identity is not available.
152            token_provider: Zero-argument callable that returns a bearer token
153                string. Use for managed identity / workload identity scenarios.
154                The callable must request the **Azure AI Search** scope::
155
156                    from azure.identity import DefaultAzureCredential, get_bearer_token_provider
157                    token_provider = get_bearer_token_provider(
158                        DefaultAzureCredential(),
159                        "https://search.azure.com/.default"
160                    )
161
162                Note: this scope is different from Azure OpenAI / Cognitive Services
163                (``https://cognitiveservices.azure.com/.default``) — each service
164                requires its own token_provider.
165            embedding_dimension: Dimension of vector embeddings (default 1536)
166            document_type: Document class (base or custom subclass). All fields will be indexed.
167            vector_field_name: Name of the vector field in the Azure Search index (default "embedding").
168                Override this when the index was built externally with a different field name
169                (e.g. "contentVector", "chunkVector").
170            content_field_name: Name of the text content field in the Azure Search index (default "content").
171                Override when the index uses a different name (e.g. "chunkContent", "text").
172        """
173        if not api_key and not token_provider:
174            raise ValueError(
175                "Either api_key or token_provider must be supplied to AzureAISearchVectorStore."
176            )
177        self.endpoint = endpoint
178        self.index_name = index_name
179        self.embedding_dimension = embedding_dimension
180        self.document_type = document_type
181        self.vector_field_name = vector_field_name
182        self.content_field_name = content_field_name
183
184        if token_provider:
185            credential = _TokenProviderCredential(token_provider)
186        else:
187            credential = AzureKeyCredential(api_key)
188
189        # Initialize clients
190        self.search_client = SearchClient(
191            endpoint=endpoint,
192            index_name=index_name,
193            credential=credential
194        )
195    
196    @staticmethod
197    def _format_datetime_for_azure(dt: datetime) -> str:
198        """
199        Convert datetime to Azure Search DateTimeOffset format (ISO 8601 with timezone).
200        
201        Args:
202            dt: datetime object (timezone-aware or naive)
203            
204        Returns:
205            ISO 8601 string with timezone (e.g., '2024-01-15T10:30:00Z')
206        """
207        if dt.tzinfo is None:
208            # Timezone-naive datetime: assume UTC and append 'Z'
209            return dt.isoformat() + 'Z'
210        else:
211            # Timezone-aware datetime: use standard ISO format
212            return dt.isoformat()
213    
214    def add_documents(
215        self,
216        documents: List[Document],
217        generate_embeddings: bool = True
218    ) -> List[str]:
219        """
220        Add documents to Azure AI Search.
221        
222        Args:
223            documents: List of Document objects to add (can be base or custom subclasses)
224            generate_embeddings: If True, validates embeddings are present
225        
226        Returns:
227            List of document IDs that were added
228            
229        Note:
230            Custom document fields are automatically serialized and stored.
231            All fields remain accessible after retrieval.
232        """
233        if not documents:
234            return []
235        
236        # Helper function to format datetime for Azure Search DateTimeOffset
237        format_datetime_for_azure = self._format_datetime_for_azure
238        
239        # Helper function to serialize datetime objects in dictionaries
240        def serialize_for_json(obj):
241            """Recursively convert datetime objects to ISO format strings for JSON serialization."""
242            if isinstance(obj, datetime):
243                return format_datetime_for_azure(obj)
244            elif isinstance(obj, dict):
245                return {k: serialize_for_json(v) for k, v in obj.items()}
246            elif isinstance(obj, (list, tuple)):
247                return [serialize_for_json(item) for item in obj]
248            else:
249                return obj
250        
251        # Prepare documents for upload
252        search_documents = []
253        added_ids = []
254        
255        for doc in documents:
256            if not doc.id or not doc.content:
257                raise ValueError(f"Document must have id and content: {doc}")
258            
259            if generate_embeddings and doc.embedding is None:
260                raise ValueError(
261                    f"Document {doc.id} lacks embedding. "
262                    "Generate embeddings externally before adding."
263                )
264            
265            # Serialize entire document (including custom fields) to JSON
266            import json
267            import dataclasses
268            doc_dict = doc.to_dict()
269            
270            # Convert datetime objects for JSON serialization
271            doc_dict_serializable = serialize_for_json(doc_dict)
272            
273            # Convert to Azure Search document format
274            # Store base fields directly
275            search_doc = {
276                "id": doc.id,
277                "content": doc.content,
278                "embedding": doc.embedding if doc.embedding else [],
279                "timestamp": format_datetime_for_azure(doc.timestamp) if doc.timestamp else None,
280                "document_data": json.dumps(doc_dict_serializable)  # Keep for metadata and backwards compatibility
281            }
282            
283            # Add all custom fields as indexed fields
284            if dataclasses.is_dataclass(doc):
285                for field in dataclasses.fields(doc):
286                    # Skip base Document fields (already added above)
287                    if field.name in {'id', 'content', 'embedding', 'timestamp', 'metadata'}:
288                        continue
289                    
290                    field_value = getattr(doc, field.name, None)
291                    
292                    if field_value is not None:
293                        # Serialize datetime fields with Azure DateTimeOffset format
294                        if isinstance(field_value, datetime):
295                            search_doc[field.name] = format_datetime_for_azure(field_value)
296                        # Handle lists/dicts by converting to JSON string
297                        elif isinstance(field_value, (list, dict)):
298                            search_doc[field.name] = json.dumps(field_value)
299                        else:
300                            search_doc[field.name] = field_value
301            
302            search_documents.append(search_doc)
303            added_ids.append(doc.id)
304        
305        # Upload to Azure Search
306        result = self.search_client.upload_documents(documents=search_documents)
307        
308        # Check for failures
309        failed_count = sum(1 for r in result if not r.succeeded)
310        if failed_count > 0:
311            logger.warning(f"{failed_count} documents failed to upload")
312        
313        return added_ids
314    
315    def search(
316        self,
317        query: Optional[str] = None,
318        query_embedding: Optional[List[float]] = None,
319        top_k: int = 5,
320        filters: Optional[Dict[str, Any]] = None,
321        search_type: str = "vector"
322    ) -> List[SearchResult]:
323        """
324        Search Azure AI Search index.
325        
326        Args:
327            query: Text query (required for keyword/hybrid search)
328            query_embedding: Vector embedding (required for vector/hybrid search)
329            top_k: Number of results to return
330            filters: Metadata filters (e.g., {"source": "doc.pdf"})
331            search_type: "vector", "keyword", or "hybrid"
332        
333        Returns:
334            List of SearchResult objects, ordered by relevance
335        """
336        if search_type not in ["vector", "keyword", "hybrid"]:
337            raise ValueError(f"Invalid search_type: {search_type}")
338        
339        if search_type in ["vector", "hybrid"] and query_embedding is None:
340            raise ValueError(f"{search_type} search requires query_embedding")
341        
342        if search_type in ["keyword", "hybrid"] and query is None:
343            raise ValueError(f"{search_type} search requires query text")
344        
345        # Build filter expression
346        filter_expr = self._build_filter_expression(filters)
347        
348        # Execute search based on type
349        if search_type == "vector":
350            # Pure vector search
351            vector_query = VectorizedQuery(
352                vector=query_embedding,
353                k_nearest_neighbors=top_k,
354                fields=self.vector_field_name
355            )
356            
357            results = self.search_client.search(
358                search_text=None,
359                vector_queries=[vector_query],
360                filter=filter_expr,
361                top=top_k
362            )
363        
364        elif search_type == "keyword":
365            # Pure keyword search (BM25)
366            results = self.search_client.search(
367                search_text=query,
368                filter=filter_expr,
369                top=top_k,
370                include_total_count=False
371            )
372        
373        else:  # hybrid
374            # Hybrid search (vector + keyword)
375            vector_query = VectorizedQuery(
376                vector=query_embedding,
377                k_nearest_neighbors=top_k,
378                fields=self.vector_field_name
379            )
380            
381            results = self.search_client.search(
382                search_text=query,
383                vector_queries=[vector_query],
384                filter=filter_expr,
385                top=top_k
386            )
387        
388        # Convert to SearchResult objects
389        search_results = []
390        for rank, result in enumerate(results):
391            import json
392            
393            # Deserialize full document data
394            doc_data = json.loads(result.get("document_data", "{}") or "{}")
395            
396            if not doc_data:
397                # Index was built externally — fields are stored directly on the result.
398                # Collect all non-Azure-internal fields, excluding the vector blob.
399                doc_data = {
400                    k: v for k, v in result.items()
401                    if not k.startswith("@") and k != self.vector_field_name
402                }
403                # Remap content field when the index uses a non-standard name.
404                if self.content_field_name != "content" and self.content_field_name in doc_data:
405                    doc_data["content"] = doc_data.pop(self.content_field_name)
406            
407            # Reconstruct document using the configured document_type
408            doc = self.document_type.from_dict(doc_data)
409            
410            search_results.append(SearchResult(
411                document=doc,
412                score=result.get("@search.score", 0.0),
413                rank=rank
414            ))
415        
416        return search_results
417    
418    def delete_documents(self, document_ids: List[str]) -> int:
419        """
420        Delete documents from Azure AI Search.
421        
422        Args:
423            document_ids: List of document IDs to delete
424        
425        Returns:
426            Number of documents successfully deleted
427        """
428        if not document_ids:
429            return 0
430        
431        # Prepare delete operations
432        delete_docs = [{"id": doc_id} for doc_id in document_ids]
433        
434        # Execute delete
435        result = self.search_client.delete_documents(documents=delete_docs)
436        
437        # Count successful deletions
438        deleted_count = sum(1 for r in result if r.succeeded)
439        return deleted_count
440    
441    def update_document(
442        self,
443        document_id: str,
444        content: Optional[str] = None,
445        embedding: Optional[List[float]] = None,
446        metadata: Optional[Dict[str, Any]] = None,
447        timestamp: Optional[datetime] = None,
448        **custom_fields
449    ) -> bool:
450        """
451        Update an existing document in Azure AI Search.
452        
453        Args:
454            document_id: ID of the document to update
455            content: New content (if provided)
456            embedding: New embedding (if provided)
457            metadata: New metadata (merged with existing)
458            timestamp: New timestamp (if provided)
459            **custom_fields: Any custom fields to update
460        
461        Returns:
462            True if update was successful, False if document not found
463        """
464        try:
465            import json
466            
467            # Retrieve existing document
468            existing_doc = self.search_client.get_document(key=document_id)
469            doc_data = json.loads(existing_doc.get("document_data", "{}"))
470            
471            # Update fields in document data
472            if content is not None:
473                doc_data["content"] = content
474                existing_doc["content"] = content
475            
476            if embedding is not None:
477                doc_data["embedding"] = embedding
478                existing_doc["embedding"] = embedding
479            
480            if timestamp is not None:
481                doc_data["timestamp"] = self._format_datetime_for_azure(timestamp)
482                existing_doc["timestamp"] = self._format_datetime_for_azure(timestamp)
483            
484            if metadata is not None:
485                if "metadata" not in doc_data:
486                    doc_data["metadata"] = {}
487                doc_data["metadata"].update(metadata)
488            
489            # Update custom fields
490            for key, value in custom_fields.items():
491                doc_data[key] = value
492            
493            # Update document_data
494            existing_doc["document_data"] = json.dumps(doc_data)
495            
496            # Merge/upload updated document
497            result = self.search_client.merge_or_upload_documents(documents=[existing_doc])
498            
499            return result[0].succeeded
500        
501        except Exception as e:
502            logger.error(f"Failed to update document {document_id}: {e}")
503            return False
504    
505    def get_document(self, document_id: str) -> Optional[Document]:
506        """
507        Retrieve a document by ID from Azure AI Search.
508        
509        Args:
510            document_id: ID of the document to retrieve
511        
512        Returns:
513            Document if found, None otherwise (returns configured document_type)
514        """
515        try:
516            import json
517            result = self.search_client.get_document(key=document_id)
518            
519            # Deserialize full document data using configured document_type
520            doc_data = json.loads(result.get("document_data", "{}"))
521            return self.document_type.from_dict(doc_data)
522        
523        except Exception as e:
524            logger.debug(f"Document {document_id} not found: {e}")
525            return None
526
527    def get_document_by_parent(self, document_id: str) -> List[Document]:
528        """
529        Retrieve all chunks belonging to a parent document.
530
531        Args:
532            document_id: The value of the ``documentId`` field shared across
533                        all chunks of the same parent document.
534
535        Returns:
536            List of Document objects sorted by ``pageNumber``, or an empty
537            list if no matching chunks are found.
538        """
539        results = self.search(
540            query="*",
541            search_type="keyword",
542            filters={"documentId": document_id},
543            top_k=1000,
544        )
545        chunks = [r.document for r in results]
546        chunks.sort(key=lambda c: getattr(c, "pageNumber", 0))
547        return chunks
548
549    def count(self) -> int:
550        """
551        Get the total number of documents in the index.
552        
553        Returns:
554            Total document count
555        """
556        # Execute a search to get total count
557        results = self.search_client.search(
558            search_text="*",
559            include_total_count=True,
560            top=0
561        )
562        
563        # Get count from results
564        count = getattr(results, 'get_count', lambda: 0)()
565        return count if count is not None else 0
566    
567    def clear(self) -> None:
568        """
569        Remove all documents from the index.
570        
571        Warning: This operation is irreversible.
572        """
573        # Search for all document IDs
574        results = self.search_client.search(
575            search_text="*",
576            select=["id"],
577            top=1000  # Adjust batch size as needed
578        )
579        
580        # Collect all IDs
581        doc_ids = [r["id"] for r in results]
582        
583        # Delete in batches
584        if doc_ids:
585            self.delete_documents(doc_ids)
586            logger.info(f"Cleared {len(doc_ids)} documents from index '{self.index_name}'")
587    
588    def _build_filter_expression(self, filters: Optional[Dict[str, Any]]) -> Optional[str]:
589        """
590        Build OData filter expression from filter dictionary.
591        
592        Supports filtering on all indexed fields (universal + custom fields from document_type).
593        
594        Args:
595            filters: Dictionary of field: value pairs
596                    Special operators: use tuples for range queries
597                    - (">=", value) for greater than or equal
598                    - ("<=", value) for less than or equal
599                    - ("range", min_val, max_val) for range queries
600        
601        Returns:
602            OData filter string or None
603            
604        Example:
605            ```python
606            # Filter on universal field
607            filters = {"timestamp": (">=", datetime(2024, 1, 1))}
608            
609            # Filter on custom fields (if document_type has them)
610            filters = {
611                "court": "Supreme Court",
612                "case_number": "2024-001",
613                "decision_date": ("range", start_date, end_date)
614            }
615            ```
616        """
617        if not filters:
618            return None
619        
620        # Get indexed field names from document_type
621        import dataclasses
622        indexed_fields = {'id', 'content', 'timestamp'}  # Universal fields
623        
624        if dataclasses.is_dataclass(self.document_type):
625            for field in dataclasses.fields(self.document_type):
626                if field.name not in {'id', 'content', 'embedding', 'timestamp', 'metadata'}:
627                    indexed_fields.add(field.name)
628        
629        filter_parts = []
630        
631        for key, value in filters.items():
632            # Check if field is indexed
633            if key not in indexed_fields:
634                logger.warning(
635                    f"Filtering on '{key}' - field not found in document schema. "
636                    f"Indexed fields: {indexed_fields}"
637                )
638                continue
639            
640            # Handle tuple operators for range queries
641            if isinstance(value, tuple):
642                operator = value[0]
643                if operator == "range" and len(value) == 3:
644                    # Range query: field >= min and field <= max
645                    min_val = value[1]
646                    max_val = value[2]
647                    if isinstance(min_val, datetime):
648                        min_val = self._format_datetime_for_azure(min_val)
649                    if isinstance(max_val, datetime):
650                        max_val = self._format_datetime_for_azure(max_val)
651                    filter_parts.append(f"{key} ge {min_val} and {key} le {max_val}")
652                elif operator in [">=", "gt", ">="]:
653                    op_str = "ge"
654                    val = value[1]
655                    if isinstance(val, datetime):
656                        val = self._format_datetime_for_azure(val)
657                    filter_parts.append(f"{key} {op_str} {val}")
658                elif operator in ["<=", "lt", "<="]:
659                    op_str = "le"
660                    val = value[1]
661                    if isinstance(val, datetime):
662                        val = self._format_datetime_for_azure(val)
663                    filter_parts.append(f"{key} {op_str} {val}")
664            # Handle datetime equality
665            elif isinstance(value, datetime):
666                filter_parts.append(f"{key} eq {self._format_datetime_for_azure(value)}")
667            # Handle string equality (needs quotes in OData)
668            elif isinstance(value, str):
669                # Escape single quotes in the value
670                escaped_value = value.replace("'", "''")
671                filter_parts.append(f"{key} eq '{escaped_value}'")
672            # Handle boolean
673            elif isinstance(value, bool):
674                filter_parts.append(f"{key} eq {str(value).lower()}")
675            # Handle numeric types
676            elif isinstance(value, (int, float)):
677                filter_parts.append(f"{key} eq {value}")
678        
679        return " and ".join(filter_parts) if filter_parts else None