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

common/environment_manager.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""
+Environment Manager - Smart Environment Detection and Activation
+
+This module automatically detects and uses the best available Python environment:
+1. UV environment (.venv) if available and working
+2. Current active environment if it has required packages
+3. System Python as fallback
+
+Developers can override by setting FORCE_ENV environment variable or by
+activating their preferred environment before running scripts.
+"""
+
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class EnvironmentManager:
+    """Smart environment detection and management."""
+
+    def __init__(self):
+        self.project_root = Path(__file__).parent.parent
+        self.uv_env_path = self.project_root / ".venv"
+        self.current_env = os.environ.get("VIRTUAL_ENV")
+        self.force_env = os.environ.get("FORCE_ENV")  # Allow developer override
+
+    def get_best_python_executable(self) -> str:
+        """
+        Get the best Python executable to use.
+
+        Returns:
+            Path to Python executable
+        """
+        # Developer override
+        if self.force_env:
+            if self.force_env == "system":
+                return sys.executable
+            elif self.force_env == "uv" and self.uv_env_path.exists():
+                return str(self.uv_env_path / "bin" / "python")
+            else:
+                logger.warning(
+                    f"FORCE_ENV={self.force_env} not valid, using auto-detection"
+                )
+
+        # If already in a virtual environment and it has required packages, use it
+        if self.current_env and self._check_environment_has_iris():
+            logger.debug(f"Using current virtual environment: {self.current_env}")
+            return sys.executable
+
+        # Check UV environment if available
+        if self.uv_env_path.exists():
+            uv_python = self.uv_env_path / "bin" / "python"
+            if uv_python.exists() and self._check_environment_has_iris(str(uv_python)):
+                logger.debug(f"Using UV environment: {self.uv_env_path}")
+                return str(uv_python)
+
+        # Fallback to current Python
+        logger.debug("Using current Python executable")
+        return sys.executable
+
+    def _check_environment_has_iris(self, python_exe: Optional[str] = None) -> bool:
+        """
+        Check if an environment has the required IRIS packages.
+
+        Args:
+            python_exe: Python executable to check (None for current)
+
+        Returns:
+            True if environment has required packages
+        """
+        if python_exe is None:
+            python_exe = sys.executable
+
+        try:
+            # Quick check for intersystems_irispython package
+            result = subprocess.run(
+                [
+                    python_exe,
+                    "-c",
+                    "try: import iris; print(hasattr(iris, 'connect')); except ImportError: import iris; print(hasattr(iris, 'connect'))",
+                ],
+                capture_output=True,
+                text=True,
+                timeout=5,
+            )
+
+            return result.returncode == 0 and "True" in result.stdout
+
+        except Exception as e:
+            logger.debug(f"Environment check failed for {python_exe}: {e}")
+            return False
+
+    def get_environment_info(self) -> Dict[str, Any]:
+        """
+        Get information about the current environment setup.
+
+        Returns:
+            Dictionary with environment information
+        """
+        best_python = self.get_best_python_executable()
+
+        info = {
+            "best_python": best_python,
+            "current_python": sys.executable,
+            "current_venv": self.current_env,
+            "uv_env_exists": self.uv_env_path.exists(),
+            "force_env": self.force_env,
+            "iris_available": self._check_environment_has_iris(),
+            "best_iris_available": self._check_environment_has_iris(best_python),
+        }
+
+        return info
+
+    def ensure_iris_available(self) -> bool:
+        """
+        Ensure IRIS packages are available in the best environment.
+
+        Returns:
+            True if IRIS is available, False otherwise
+        """
+        best_python = self.get_best_python_executable()
+        return self._check_environment_has_iris(best_python)
+
+    def run_in_best_environment(
+        self, script_path: str, args: list = None
+    ) -> subprocess.CompletedProcess:
+        """
+        Run a script in the best available environment.
+
+        Args:
+            script_path: Path to Python script
+            args: Additional arguments
+
+        Returns:
+            CompletedProcess result
+        """
+        best_python = self.get_best_python_executable()
+        cmd = [best_python, script_path]
+        if args:
+            cmd.extend(args)
+
+        logger.debug(f"Running: {' '.join(cmd)}")
+        return subprocess.run(cmd, cwd=self.project_root)
+
+
+def get_iris_connection_in_best_env():
+    """
+    Get an IRIS connection using the best available environment.
+    This function automatically handles environment switching if needed.
+    """
+    env_manager = EnvironmentManager()
+
+    # If we're already in the best environment, use direct import
+    if env_manager.get_best_python_executable() == sys.executable:
+        try:
+            import iris
+
+            if hasattr(iris, "connect"):
+                return iris.connect(
+                    hostname=os.environ.get("IRIS_HOST", "localhost"),
+                    port=int(os.environ.get("IRIS_PORT", "1974")),
+                    namespace=os.environ.get("IRIS_NAMESPACE", "USER"),
+                    username=os.environ.get("IRIS_USERNAME", "SuperUser"),
+                    password=os.environ.get("IRIS_PASSWORD", "SYS"),
+                )
+        except Exception as e:
+            logger.warning(f"Direct IRIS import failed: {e}")
+
+    # Otherwise, we need to spawn a subprocess in the correct environment
+    logger.info("Switching to optimal environment for IRIS connection")
+    raise RuntimeError(
+        "IRIS connection requires environment switch. "
+        "Please use the connection manager or run in the correct environment."
+    )
+
+
+def check_environment_status():
+    """Check and report environment status."""
+    env_manager = EnvironmentManager()
+    info = env_manager.get_environment_info()
+
+    print("=== Environment Status ===")
+    print(f"Current Python: {info['current_python']}")
+    print(f"Best Python: {info['best_python']}")
+    print(f"Current VENV: {info['current_venv'] or 'None'}")
+    print(f"UV Environment: {'Available' if info['uv_env_exists'] else 'Not found'}")
+    print(f"Force Environment: {info['force_env'] or 'Auto-detect'}")
+    print(f"IRIS in Current: {'✅' if info['iris_available'] else '❌'}")
+    print(f"IRIS in Best: {'✅' if info['best_iris_available'] else '❌'}")
+
+    if not info["best_iris_available"]:
+        print("\n⚠️  Warning: IRIS packages not available in any environment")
+        print("   Consider running: uv sync  # if using UV")
+        print("   Or: pip install intersystems-irispython")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Environment Manager")
+    parser.add_argument("--check", action="store_true", help="Check environment status")
+    parser.add_argument("--run", help="Run script in best environment")
+    parser.add_argument("args", nargs="*", help="Arguments for script")
+
+    args = parser.parse_args()
+
+    if args.check:
+        check_environment_status()
+    elif args.run:
+        env_manager = EnvironmentManager()
+        result = env_manager.run_in_best_environment(args.run, args.args)
+        sys.exit(result.returncode)
+    else:
+        parser.print_help()

common/huggingface_utils.py

@@ -0,0 +1,206 @@
+"""
+HuggingFace model download utilities with rate limiting and retry logic.
+"""
+
+import logging
+import random
+import time
+from functools import wraps
+from typing import Any, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+def retry_with_exponential_backoff(
+    max_retries: int = 5,
+    base_delay: float = 1.0,
+    max_delay: float = 60.0,
+    exponential_base: float = 2.0,
+    jitter: bool = True,
+):
+    """
+    Decorator for retrying functions with exponential backoff.
+
+    Args:
+        max_retries: Maximum number of retry attempts
+        base_delay: Initial delay in seconds
+        max_delay: Maximum delay in seconds
+        exponential_base: Base for exponential backoff
+        jitter: Whether to add random jitter to delays
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            last_exception = None
+
+            for attempt in range(max_retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except Exception as e:
+                    last_exception = e
+
+                    # Check if this is a rate limiting error
+                    error_str = str(e).lower()
+                    is_rate_limit = any(
+                        indicator in error_str
+                        for indicator in [
+                            "rate limit",
+                            "too many requests",
+                            "429",
+                            "quota exceeded",
+                            "service unavailable",
+                            "503",
+                            "timeout",
+                        ]
+                    )
+
+                    if attempt == max_retries:
+                        logger.error(f"Failed after {max_retries} retries: {e}")
+                        raise last_exception
+
+                    if is_rate_limit:
+                        # Calculate delay with exponential backoff
+                        delay = min(base_delay * (exponential_base**attempt), max_delay)
+
+                        # Add jitter to prevent thundering herd
+                        if jitter:
+                            delay *= 0.5 + random.random() * 0.5
+
+                        logger.warning(
+                            f"Rate limit detected (attempt {attempt + 1}/{max_retries + 1}). "
+                            f"Retrying in {delay:.2f} seconds: {e}"
+                        )
+                        time.sleep(delay)
+                    else:
+                        # For non-rate-limit errors, fail immediately
+                        logger.error(f"Non-rate-limit error encountered: {e}")
+                        raise e
+
+            raise last_exception
+
+        return wrapper
+
+    return decorator
+
+
+@retry_with_exponential_backoff(max_retries=5, base_delay=2.0, max_delay=120.0)
+def download_huggingface_model(
+    model_name: str, trust_remote_code: bool = False, **kwargs
+) -> Tuple[Any, Any]:
+    """
+    Download HuggingFace model and tokenizer with retry logic.
+
+    Args:
+        model_name: Name of the model to download
+        trust_remote_code: Whether to trust remote code
+        **kwargs: Additional arguments for model loading
+
+    Returns:
+        Tuple of (tokenizer, model)
+    """
+    try:
+        from transformers import AutoModel, AutoTokenizer
+
+        logger.info(f"Downloading HuggingFace model: {model_name}")
+
+        # Download tokenizer first
+        logger.debug(f"Loading tokenizer for {model_name}")
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name, trust_remote_code=trust_remote_code, **kwargs
+        )
+
+        # Download model
+        logger.debug(f"Loading model for {model_name}")
+        model = AutoModel.from_pretrained(
+            model_name, trust_remote_code=trust_remote_code, **kwargs
+        )
+
+        logger.info(f"Successfully downloaded HuggingFace model: {model_name}")
+        return tokenizer, model
+
+    except ImportError as e:
+        logger.error(f"transformers library not available: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Failed to download model {model_name}: {e}")
+        raise
+
+
+def get_cached_model(
+    model_name: str, cache_dict: dict, trust_remote_code: bool = False, **kwargs
+) -> Tuple[Any, Any]:
+    """
+    Get model from cache or download if not cached.
+
+    Args:
+        model_name: Name of the model
+        cache_dict: Dictionary to use for caching
+        trust_remote_code: Whether to trust remote code
+        **kwargs: Additional arguments for model loading
+
+    Returns:
+        Tuple of (tokenizer, model)
+    """
+    if model_name not in cache_dict:
+        logger.info(f"Model {model_name} not in cache, downloading...")
+        tokenizer, model = download_huggingface_model(
+            model_name, trust_remote_code=trust_remote_code, **kwargs
+        )
+        cache_dict[model_name] = (tokenizer, model)
+        logger.info(f"Cached model {model_name}")
+    else:
+        logger.info(f"Using cached model {model_name}")
+        tokenizer, model = cache_dict[model_name]
+
+    return tokenizer, model
+
+
+def clear_model_cache(cache_dict: dict, model_name: Optional[str] = None):
+    """
+    Clear model cache.
+
+    Args:
+        cache_dict: Dictionary containing cached models
+        model_name: Specific model to clear, or None to clear all
+    """
+    if model_name:
+        if model_name in cache_dict:
+            del cache_dict[model_name]
+            logger.info(f"Cleared cache for model {model_name}")
+    else:
+        cache_dict.clear()
+        logger.info("Cleared all model cache")
+
+
+# Global cache for models
+_global_model_cache = {}
+
+
+def get_global_cached_model(
+    model_name: str, trust_remote_code: bool = False, **kwargs
+) -> Tuple[Any, Any]:
+    """
+    Get model from global cache or download if not cached.
+
+    Args:
+        model_name: Name of the model
+        trust_remote_code: Whether to trust remote code
+        **kwargs: Additional arguments for model loading
+
+    Returns:
+        Tuple of (tokenizer, model)
+    """
+    return get_cached_model(
+        model_name, _global_model_cache, trust_remote_code, **kwargs
+    )
+
+
+def clear_global_cache(model_name: Optional[str] = None):
+    """
+    Clear global model cache.
+
+    Args:
+        model_name: Specific model to clear, or None to clear all
+    """
+    clear_model_cache(_global_model_cache, model_name)

common/iris_client.py

@@ -0,0 +1,92 @@
+"""
+Simple IRIS client for database operations
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .iris_connection_manager import IRISConnectionManager
+
+logger = logging.getLogger(__name__)
+
+
+class IRISClient:
+    """Simple IRIS database client"""
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """Initialize IRIS client with configuration"""
+        self.config = config or {}
+        self.connection_manager = IRISConnectionManager()
+        self._connection = None
+
+    def connect(self) -> bool:
+        """Establish database connection"""
+        try:
+            self._connection = self.connection_manager.get_connection(self.config)
+            return True
+        except Exception as e:
+            logger.error(f"Failed to connect to IRIS: {e}")
+            return False
+
+    def query(self, sql: str, params: Optional[tuple] = None) -> List[Dict[str, Any]]:
+        """Execute query and return results"""
+        if not self._connection:
+            if not self.connect():
+                raise ConnectionError("No database connection available")
+
+        try:
+            cursor = self._connection.cursor()
+            cursor.execute(sql, params or ())
+
+            # Get column names
+            columns = (
+                [desc[0] for desc in cursor.description] if cursor.description else []
+            )
+
+            # Fetch results
+            rows = cursor.fetchall()
+
+            # Convert to list of dictionaries
+            results = []
+            for row in rows:
+                results.append(dict(zip(columns, row)))
+
+            return results
+
+        except Exception as e:
+            logger.error(f"Query failed: {e}")
+            raise
+
+    def execute(self, sql: str, params: Optional[tuple] = None) -> int:
+        """Execute non-query SQL statement"""
+        if not self._connection:
+            if not self.connect():
+                raise ConnectionError("No database connection available")
+
+        try:
+            cursor = self._connection.cursor()
+            cursor.execute(sql, params or ())
+            self._connection.commit()
+            return cursor.rowcount
+        except Exception as e:
+            logger.error(f"Execute failed: {e}")
+            self._connection.rollback()
+            raise
+
+    def close(self):
+        """Close database connection"""
+        if self._connection:
+            try:
+                self._connection.close()
+            except:
+                pass
+            self._connection = None
+
+    def __enter__(self):
+        """Context manager entry"""
+        self.connect()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit"""
+        self.close()