iris-vector-rag 0.2.1__py3-none-any.whl

common/batch_utils.py

@@ -0,0 +1,294 @@
+"""
+Batch processing utilities for entity extraction.
+
+This module provides utilities for batching documents, retry logic with exponential
+backoff, and batch processing metrics tracking.
+
+Implementation:
+- BatchQueue: FIFO queue with token-aware batching (FR-006)
+- Retry logic: Exponential backoff with batch splitting (FR-005)
+- BatchMetricsTracker: Global singleton for metrics tracking (FR-007)
+"""
+
+from collections import deque
+from typing import List, Optional, Callable
+import time
+import logging
+
+from iris_rag.core.models import (
+    Document,
+    BatchExtractionResult,
+    ProcessingMetrics,
+    BatchStatus,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class BatchQueue:
+    """
+    FIFO queue for batching documents with token budget enforcement.
+
+    This class supports FR-006 (dynamic batch sizing based on token budget) and
+    provides O(1) queue operations using collections.deque.
+
+    Attributes:
+        token_budget: Default token budget for batches (default: 8192)
+        _queue: Internal deque storing (document, token_count) tuples
+    """
+
+    def __init__(self, token_budget: int = 8192):
+        """
+        Initialize batch queue.
+
+        Args:
+            token_budget: Default token budget for batches (default: 8192 per FR-006)
+        """
+        self.token_budget = token_budget
+        self._queue: deque = deque()
+
+    def add_document(self, document: Document, token_count: int) -> None:
+        """
+        Add a document to the queue.
+
+        Args:
+            document: Document to add
+            token_count: Token count for the document
+
+        Raises:
+            ValueError: If token_count is negative
+        """
+        if token_count < 0:
+            raise ValueError("Token count must be non-negative")
+
+        self._queue.append((document, token_count))
+
+    def get_next_batch(self, token_budget: Optional[int] = None) -> Optional[List[Document]]:
+        """
+        Get next batch of documents within token budget.
+
+        This implements first-fit batching strategy: documents are added to the batch
+        in FIFO order until adding the next document would exceed the token budget.
+
+        Special cases (per FR-006):
+        - If queue is empty, returns None
+        - If first document exceeds budget, returns it anyway (can't skip)
+        - If token_budget is 0 or negative, returns first document
+
+        Args:
+            token_budget: Token budget for this batch (overrides default if provided)
+
+        Returns:
+            List of documents for the batch, or None if queue is empty
+        """
+        if not self._queue:
+            return None
+
+        # Use provided budget or default
+        budget = token_budget if token_budget is not None else self.token_budget
+
+        batch = []
+        cumulative_tokens = 0
+
+        # Build batch until we exceed budget
+        while self._queue:
+            doc, tokens = self._queue[0]  # Peek at first document
+
+            # First document always goes in batch (even if exceeds budget)
+            if not batch:
+                batch.append(doc)
+                cumulative_tokens += tokens
+                self._queue.popleft()  # Remove from queue
+                continue
+
+            # Check if adding this document would exceed budget
+            if cumulative_tokens + tokens > budget:
+                break  # Don't add this document
+
+            # Add document to batch
+            batch.append(doc)
+            cumulative_tokens += tokens
+            self._queue.popleft()  # Remove from queue
+
+        return batch if batch else None
+
+
+def extract_batch_with_retry(
+    documents: List[Document],
+    extract_fn: Callable[[List[Document]], BatchExtractionResult],
+    max_retries: int = 3,
+    retry_delays: List[float] = None,
+) -> BatchExtractionResult:
+    """
+    Extract batch with exponential backoff retry logic (FR-005).
+
+    This function implements:
+    - Exponential backoff: 2s, 4s, 8s delays (configurable)
+    - Batch splitting: If all retries fail, split batch in half and retry
+    - Retry tracking: Tracks retry attempts in result
+
+    Args:
+        documents: List of documents to extract entities from
+        extract_fn: Function to call for extraction (takes List[Document], returns BatchExtractionResult)
+        max_retries: Maximum retry attempts (default: 3 per FR-005)
+        retry_delays: Retry delays in seconds (default: [2, 4, 8] per FR-005)
+
+    Returns:
+        BatchExtractionResult with extraction results or error
+
+    Examples:
+        >>> def my_extractor(docs):
+        ...     return BatchExtractionResult(batch_id="test", success_status=True)
+        >>> result = extract_batch_with_retry(documents, my_extractor)
+    """
+    if retry_delays is None:
+        retry_delays = [2.0, 4.0, 8.0]  # FR-005 exponential backoff
+
+    retry_count = 0
+
+    for attempt in range(max_retries + 1):  # +1 for initial attempt
+        try:
+            # Attempt extraction
+            result = extract_fn(documents)
+
+            # If successful, return result with retry count
+            if result.success_status:
+                result.retry_count = retry_count
+                return result
+
+            # If extraction failed, retry
+            logger.warning(
+                f"Batch extraction failed (attempt {attempt + 1}/{max_retries + 1}): {result.error_message}"
+            )
+
+        except Exception as e:
+            logger.error(
+                f"Batch extraction exception (attempt {attempt + 1}/{max_retries + 1}): {e}"
+            )
+
+        # If not last attempt, wait and retry
+        if attempt < max_retries:
+            retry_count += 1
+            delay = retry_delays[min(attempt, len(retry_delays) - 1)]
+            logger.info(f"Retrying batch extraction after {delay}s delay...")
+            time.sleep(delay)
+        else:
+            # All retries exhausted
+            logger.error(
+                f"All {max_retries} retry attempts exhausted for batch of {len(documents)} documents"
+            )
+
+            # Split batch and retry if batch has multiple documents
+            if len(documents) > 1:
+                logger.info(
+                    f"Splitting batch of {len(documents)} documents into 2 sub-batches"
+                )
+                mid = len(documents) // 2
+                batch1 = documents[:mid]
+                batch2 = documents[mid:]
+
+                # Recursively retry each sub-batch (with retry_count tracking)
+                result1 = extract_batch_with_retry(
+                    batch1, extract_fn, max_retries, retry_delays
+                )
+                result2 = extract_batch_with_retry(
+                    batch2, extract_fn, max_retries, retry_delays
+                )
+
+                # Merge results
+                merged_result = BatchExtractionResult(
+                    batch_id=result1.batch_id,  # Use first sub-batch ID
+                    per_document_entities={
+                        **result1.per_document_entities,
+                        **result2.per_document_entities,
+                    },
+                    per_document_relationships={
+                        **result1.per_document_relationships,
+                        **result2.per_document_relationships,
+                    },
+                    processing_time=result1.processing_time + result2.processing_time,
+                    success_status=result1.success_status and result2.success_status,
+                    retry_count=retry_count + result1.retry_count + result2.retry_count,
+                    error_message=None
+                    if (result1.success_status and result2.success_status)
+                    else f"Sub-batch errors: {result1.error_message}; {result2.error_message}",
+                )
+
+                return merged_result
+
+            # Single document batch - return failure
+            return BatchExtractionResult(
+                batch_id="failed",
+                success_status=False,
+                retry_count=retry_count,
+                error_message=f"Failed after {max_retries} retries (single document)",
+            )
+
+
+class BatchMetricsTracker:
+    """
+    Global singleton for tracking batch processing metrics (FR-007).
+
+    This class provides a global singleton instance for tracking metrics across
+    all batch processing operations. It implements FR-007 statistics requirements.
+
+    Usage:
+        >>> tracker = BatchMetricsTracker.get_instance()
+        >>> tracker.update_with_batch(batch_result, batch_size=10)
+        >>> metrics = tracker.get_statistics()
+    """
+
+    _instance: Optional["BatchMetricsTracker"] = None
+
+    def __init__(self):
+        """Initialize metrics tracker with default values."""
+        self._metrics = ProcessingMetrics()
+
+    @classmethod
+    def get_instance(cls) -> "BatchMetricsTracker":
+        """
+        Get global singleton instance of BatchMetricsTracker.
+
+        Returns:
+            Global BatchMetricsTracker instance
+        """
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def reset_instance(cls) -> None:
+        """
+        Reset global singleton instance (useful for testing).
+
+        This creates a fresh instance with zeroed metrics.
+        """
+        cls._instance = None
+
+    def update_with_batch(
+        self,
+        batch_result: BatchExtractionResult,
+        batch_size: int,
+        single_doc_baseline_time: float = 7.2,
+    ) -> None:
+        """
+        Update metrics with a new batch result.
+
+        Args:
+            batch_result: Result from batch processing
+            batch_size: Number of documents in the batch
+            single_doc_baseline_time: Baseline time per document (default: 7.2s from spec)
+        """
+        self._metrics.update_with_batch(
+            batch_result, batch_size, single_doc_baseline_time
+        )
+
+    def get_statistics(self) -> ProcessingMetrics:
+        """
+        Get current processing statistics (FR-007).
+
+        Returns:
+            ProcessingMetrics with current statistics
+        """
+        return self._metrics

common/config.py

@@ -0,0 +1,58 @@
+"""
+Common configuration module for IRIS RAG system
+Provides unified access to configuration settings
+"""
+
+import os
+import sys
+from pathlib import Path
+
+# Add config path to sys.path
+config_dir = Path(__file__).parent.parent / "config"
+sys.path.insert(0, str(config_dir))
+
+from loader import get_config
+
+
+def get_iris_config():
+    """Get IRIS database configuration"""
+    config = get_config()
+    if not config:
+        raise ValueError("Failed to load configuration")
+
+    db_config = config.get("database", {})
+    return {
+        "host": db_config.get("db_host", "localhost"),
+        "port": db_config.get("db_port", 1974),
+        "username": db_config.get("db_user", "SuperUser"),
+        "password": db_config.get("db_password", "SYS"),
+        "namespace": db_config.get("db_namespace", "USER"),
+        "connection_type": "dbapi",  # Default to DBAPI
+    }
+
+
+def get_embedding_config():
+    """Get embedding model configuration"""
+    config = get_config()
+    if not config:
+        raise ValueError("Failed to load configuration")
+
+    return config.get("embedding_model", {})
+
+
+def get_storage_config():
+    """Get storage configuration"""
+    config = get_config()
+    if not config:
+        raise ValueError("Failed to load configuration")
+
+    return config.get("storage", {})
+
+
+def get_pipeline_config():
+    """Get pipeline configuration"""
+    config = get_config()
+    if not config:
+        raise ValueError("Failed to load configuration")
+
+    return config.get("pipelines", {})

common/connection_manager.py

@@ -0,0 +1,242 @@
+"""
+Centralized connection manager for RAG pipelines
+Handles JDBC, ODBC, and future dbapi connections with a unified interface
+"""
+
+import logging
+import os
+from contextlib import contextmanager
+from typing import Any, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectionManager:
+    """Manages database connections for RAG pipelines"""
+
+    def __init__(self, connection_type: str = "odbc"):
+        """
+        Initialize connection manager
+
+        Args:
+            connection_type: Either "jdbc", "odbc", or "dbapi" (default: "odbc" for stability)
+        """
+        self.connection_type = connection_type.lower()
+        self._connection = None
+        self._cursor = None
+
+        if self.connection_type not in ["jdbc", "odbc", "dbapi"]:
+            raise ValueError(
+                f"Invalid connection type: {connection_type}. Must be 'jdbc', 'odbc', or 'dbapi'"
+            )
+
+    def connect(self):
+        """Establish database connection based on configured type"""
+        if self._connection:
+            return self._connection
+
+        if self.connection_type == "jdbc":
+            try:
+                from common.iris_connection_manager import get_iris_jdbc_connection
+
+                self._connection = get_iris_jdbc_connection()
+                logger.info("Established JDBC connection")
+            except Exception as e:
+                logger.warning(f"JDBC connection failed: {e}, falling back to ODBC")
+                self.connection_type = "odbc"
+                return self.connect()
+
+        elif self.connection_type == "dbapi":
+            try:
+                from common.iris_dbapi_connector import get_iris_dbapi_connection
+
+                self._connection = get_iris_dbapi_connection()
+                if not self._connection:  # If get_iris_dbapi_connection returns None
+                    raise ConnectionError("DBAPI connection returned None")
+                logger.info("Established DBAPI connection")
+            except Exception as e:
+                logger.warning(f"DBAPI connection failed: {e}, falling back to ODBC")
+                self.connection_type = "odbc"  # Fallback
+                # Clear any partial connection state before retrying
+                if self._connection and hasattr(self._connection, "close"):
+                    try:
+                        self._connection.close()
+                    except Exception as close_err:
+                        logger.error(
+                            f"Error closing partial DBAPI connection during fallback: {close_err}"
+                        )
+                self._connection = None
+                return self.connect()  # Retry with ODBC
+
+        else:  # odbc (or fallback from dbapi/jdbc)
+            from common.iris_connector import get_iris_connection
+
+            self._connection = get_iris_connection()
+            logger.info("Established ODBC connection")
+
+        return self._connection
+
+    def execute(self, query: str, params: Optional[List[Any]] = None) -> List[Any]:
+        """
+        Execute a query with optional parameters
+
+        Args:
+            query: SQL query to execute
+            params: Optional list of parameters for parameterized queries
+
+        Returns:
+            List of result rows
+        """
+        if not self._connection:
+            self.connect()
+
+        if self.connection_type == "jdbc":
+            # JDBC connector has execute method
+            return self._connection.execute(query, params or [])
+        else:  # Covers ODBC and DBAPI
+            # ODBC and DBAPI use a similar cursor-based approach
+            cursor = self._connection.cursor()
+            try:
+                if params:
+                    cursor.execute(query, params)
+                else:
+                    cursor.execute(query)
+
+                # For SELECT queries, fetchall. For others (INSERT, UPDATE, DELETE),
+                # fetchall might not be appropriate or available, or might return empty.
+                # Standard DBAPI cursor.description is None for non-SELECT.
+                if cursor.description:
+                    return cursor.fetchall()
+                else:
+                    # For non-SELECT, perhaps return rowcount or an empty list
+                    # For now, returning empty list to maintain consistency with existing behavior
+                    # if fetchall() was called on a non-SELECT for ODBC.
+                    return []
+            finally:
+                cursor.close()
+
+    def execute_many(self, query: str, params_list: List[List[Any]]) -> None:
+        """
+        Execute a query multiple times with different parameters
+
+        Args:
+            query: SQL query to execute
+            params_list: List of parameter lists
+        """
+        if not self._connection:
+            self.connect()
+
+        if self.connection_type == "jdbc":
+            # Execute in a loop for JDBC
+            for params in params_list:
+                self._connection.execute(query, params)
+        else:  # Covers ODBC and DBAPI
+            # ODBC and DBAPI support executemany
+            cursor = self._connection.cursor()
+            try:
+                cursor.executemany(query, params_list)
+                self._connection.commit()  # Ensure commit for DBAPI as well
+            finally:
+                cursor.close()
+
+    @contextmanager
+    def cursor(self):
+        """Get a cursor for direct database operations"""
+        if not self._connection:
+            self.connect()
+
+        if self.connection_type == "jdbc":
+            # For JDBC, we'll create a wrapper that mimics cursor behavior
+            class JDBCCursorWrapper:
+                def __init__(self, connection):
+                    self.connection = connection
+                    self._results = None
+
+                def execute(self, query, params=None):
+                    self._results = self.connection.execute(query, params or [])
+
+                def fetchall(self):
+                    return self._results or []
+
+                def fetchone(self):
+                    if self._results and len(self._results) > 0:
+                        return self._results[0]
+                    return None
+
+                def close(self):
+                    pass  # No-op for JDBC
+
+            yield JDBCCursorWrapper(self._connection)
+        else:  # Covers ODBC and DBAPI
+            # ODBC and DBAPI native cursor
+            cursor = self._connection.cursor()
+            try:
+                yield cursor
+            finally:
+                cursor.close()
+
+    def close(self):
+        """Close the database connection"""
+        if self._connection:
+            if self.connection_type == "jdbc":
+                # JDBC connection might not have close method
+                if hasattr(self._connection, "close"):
+                    self._connection.close()
+            else:  # Covers ODBC and DBAPI
+                if hasattr(
+                    self._connection, "close"
+                ):  # Good practice to check for DBAPI too
+                    self._connection.close()
+            self._connection = None
+            logger.info(f"Closed {self.connection_type.upper()} connection")
+
+    def __enter__(self):
+        """Context manager entry"""
+        self.connect()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit"""
+        self.close()
+        return False
+
+
+# Global connection manager instance (can be configured at startup)
+_global_connection_manager = None
+
+
+def get_connection_manager(connection_type: Optional[str] = None) -> ConnectionManager:
+    """
+    Get or create a global connection manager
+
+    Args:
+        connection_type: Override connection type (default uses environment or "odbc" for stability)
+
+    Returns:
+        ConnectionManager instance
+    """
+    global _global_connection_manager
+
+    if connection_type:
+        # Create new manager with specified type
+        return ConnectionManager(connection_type)
+
+    if _global_connection_manager is None:
+        # Create default manager based on environment
+        # Default to ODBC for stability until JDBC issues are resolved
+        default_type = os.environ.get("RAG_CONNECTION_TYPE", "odbc")
+        _global_connection_manager = ConnectionManager(default_type)
+
+    return _global_connection_manager
+
+
+def set_global_connection_type(connection_type: str):
+    """
+    Set the global connection type for all pipelines
+
+    Args:
+        connection_type: Either "jdbc", "odbc", or "dbapi"
+    """
+    global _global_connection_manager
+    _global_connection_manager = ConnectionManager(connection_type)
+    logger.info(f"Set global connection type to: {connection_type}")