contextsynapse 1.0.0__py3-none-any.whl

contextsynapse/extraction/ray_normalization.py

@@ -0,0 +1,233 @@
+"""
+Ray-Based Parallel Normalization
+
+Uses Ray for distributed parallel processing of page normalization.
+Consistent with chunking implementation using Ray.
+"""
+
+import logging
+from typing import Dict, List, Any, Optional
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Check if Ray is available
+RAY_AVAILABLE = False
+try:
+    import ray
+    RAY_AVAILABLE = True
+except ImportError:
+    RAY_AVAILABLE = False
+    logger.warning("Ray not available. Install with: pip install ray")
+
+
+def _init_ray_if_needed() -> bool:
+    """Initialize Ray if not already initialized."""
+    if not RAY_AVAILABLE:
+        return False
+    
+    try:
+        if not ray.is_initialized():
+            ray.init(ignore_reinit_error=True, num_cpus=None)  # Use all available CPUs
+        return True
+    except Exception as e:
+        logger.warning(f"Ray initialization failed: {e}")
+        return False
+
+
+if RAY_AVAILABLE:
+    @ray.remote(num_cpus=1, memory=512 * 1024 * 1024)  # 1 CPU, 512MB per worker
+    def remote_build_page(
+        page_data: Dict[str, Any],
+        page_no: int
+    ) -> Dict[str, Any]:
+        """
+        Remote function to build normalized page structure.
+        
+        This is called by Ray workers - processes one page at a time.
+        No dependencies between pages - perfect for parallelization!
+        
+        Args:
+            page_data: Raw page data from extraction
+            page_no: Page number (1-indexed)
+            
+        Returns:
+            Normalized page dictionary
+        """
+        try:
+            from .normalizer import Normalizer
+            
+            # Create minimal normalizer instance
+            # Note: file_manager not needed for page building
+            normalizer = Normalizer(file_manager=None)
+            return normalizer._build_page(page_data, page_no)
+        except Exception as e:
+            logger.error(f"Error normalizing page {page_no}: {e}")
+            # Return empty page structure on error
+            return {
+                "page_no": page_no,
+                "flat_text": "",
+                "raw_blocks": [],
+                "normalized_blocks": [],
+                "markdown": "",
+                "hierarchy": {},
+            }
+    
+    @ray.remote(num_cpus=1, memory=256 * 1024 * 1024)  # 1 CPU, 256MB per worker
+    def remote_build_table(
+        table_data: Dict[str, Any],
+        table_index: int
+    ) -> Dict[str, Any]:
+        """Remote function to build normalized table structure."""
+        try:
+            from .normalizer import Normalizer
+            normalizer = Normalizer(file_manager=None)
+            return normalizer._build_table(table_data, table_index)
+        except Exception as e:
+            logger.error(f"Error normalizing table {table_index}: {e}")
+            return {}
+    
+    @ray.remote(num_cpus=1, memory=256 * 1024 * 1024)  # 1 CPU, 256MB per worker
+    def remote_build_image(
+        image_data: Dict[str, Any],
+        image_index: int
+    ) -> Dict[str, Any]:
+        """Remote function to build normalized image structure."""
+        try:
+            from .normalizer import Normalizer
+            normalizer = Normalizer(file_manager=None)
+            return normalizer._build_image(image_data, image_index)
+        except Exception as e:
+            logger.error(f"Error normalizing image {image_index}: {e}")
+            return {}
+
+
+def normalize_pages_parallel_ray(
+    extraction_result: "ExtractionResult",
+    max_concurrent: Optional[int] = None,
+    use_streaming: bool = True
+) -> List[Dict[str, Any]]:
+    """
+    Normalize pages in parallel using Ray.
+    
+    Args:
+        extraction_result: Extraction result with pages
+        max_concurrent: Maximum concurrent Ray tasks (None = auto)
+        use_streaming: If True, process results as they complete (better memory)
+        
+    Returns:
+        List of normalized pages (in order)
+    """
+    if not _init_ray_if_needed():
+        logger.warning("Ray not available, falling back to sequential processing")
+        from .normalizer import Normalizer
+        normalizer = Normalizer(file_manager=None)
+        return [
+            normalizer._build_page(page_data, i+1)
+            for i, page_data in enumerate(extraction_result.pages)
+        ]
+    
+    # Submit all pages to Ray
+    futures = [
+        remote_build_page.remote(page_data, i+1)
+        for i, page_data in enumerate(extraction_result.pages)
+    ]
+    
+    pages = []
+    
+    if use_streaming and max_concurrent:
+        # Process results as they complete (streaming, better memory)
+        remaining = list(futures)
+        while remaining:
+            ready, remaining = ray.wait(
+                remaining,
+                num_returns=min(max_concurrent, len(remaining)),
+                timeout=1.0
+            )
+            for future in ready:
+                try:
+                    page = ray.get(future)
+                    pages.append(page)
+                except Exception as e:
+                    logger.error(f"Error getting page result: {e}")
+                    # Add empty page to maintain order
+                    pages.append({
+                        "page_no": len(pages) + 1,
+                        "flat_text": "",
+                        "raw_blocks": [],
+                        "normalized_blocks": [],
+                        "markdown": "",
+                        "hierarchy": {},
+                    })
+    else:
+        # Get all results at once
+        pages = ray.get(futures)
+    
+    # Sort by page number to ensure order (Ray may complete out of order)
+    pages.sort(key=lambda p: p.get("page_no", 0))
+    
+    return pages
+
+
+def normalize_content_parallel_ray(
+    extraction_result: "ExtractionResult",
+    max_concurrent: Optional[int] = None,
+    use_streaming: bool = True
+) -> Dict[str, Any]:
+    """
+    Normalize all content (pages, tables, images) in parallel using Ray.
+    
+    Args:
+        extraction_result: Extraction result
+        max_concurrent: Maximum concurrent Ray tasks
+        use_streaming: Process results as they complete
+        
+    Returns:
+        Content dictionary with normalized pages, tables, images
+    """
+    if not _init_ray_if_needed():
+        logger.warning("Ray not available, falling back to sequential processing")
+        from .normalizer import Normalizer
+        normalizer = Normalizer(file_manager=None)
+        return normalizer._build_content(extraction_result, None)
+    
+    # Normalize pages in parallel
+    pages = normalize_pages_parallel_ray(
+        extraction_result,
+        max_concurrent=max_concurrent,
+        use_streaming=use_streaming
+    )
+    
+    # Normalize tables in parallel (if any)
+    if extraction_result.tables:
+        table_futures = [
+            remote_build_table.remote(table_data, i+1)
+            for i, table_data in enumerate(extraction_result.tables)
+        ]
+        tables = ray.get(table_futures) if table_futures else []
+    else:
+        tables = []
+    
+    # Normalize images in parallel (if any)
+    if extraction_result.images:
+        image_futures = [
+            remote_build_image.remote(image_data, i+1)
+            for i, image_data in enumerate(extraction_result.images)
+        ]
+        images = ray.get(image_futures) if image_futures else []
+    else:
+        images = []
+    
+    return {
+        "pages": pages,
+        "tables": tables,
+        "images": images,
+        "web": [],  # Can be parallelized similarly
+        "audio": [],
+        "video": [],
+        "extensions": {},
+    }
+
+
+
+

contextsynapse/extraction/ray_runner.py

@@ -0,0 +1,207 @@
+"""
+Ray Runner
+
+Ray-based scaling utilities for parallel extraction.
+"""
+
+from typing import List, Optional, Dict, Any
+import logging
+import os
+
+try:
+    import ray
+    RAY_AVAILABLE = True
+except ImportError:
+    RAY_AVAILABLE = False
+    # Create a dummy ray object to avoid AttributeError
+    class DummyRay:
+        @staticmethod
+        def remote(func):
+            return func
+        @staticmethod
+        def is_initialized():
+            return False
+        @staticmethod
+        def init(*args, **kwargs):
+            pass
+        @staticmethod
+        def get(futures):
+            return [f() if callable(f) else f for f in futures]
+        @staticmethod
+        def shutdown():
+            pass
+    ray = DummyRay()
+
+from .config import ExtractionConfig
+from .extract_engine import run_extraction
+
+logger = logging.getLogger(__name__)
+
+
+def _init_ray_if_needed() -> bool:
+    """
+    Initialize Ray if not already initialized.
+    
+    Returns:
+        True if Ray is available and initialized, False otherwise
+    """
+    if not RAY_AVAILABLE:
+        logger.warning("Ray is not available. Install with: pip install ray")
+        return False
+    
+    if not ray.is_initialized():
+        try:
+            ray.init(ignore_reinit_error=True)
+            logger.info("Ray initialized successfully")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to initialize Ray: {e}")
+            return False
+    
+    return True
+
+
+@ray.remote
+def remote_extract_document(config_dict: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Remote Ray task for document extraction.
+    
+    Args:
+        config_dict: ExtractionConfig as dictionary
+        
+    Returns:
+        Dictionary with document_id and status
+    """
+    try:
+        config = ExtractionConfig.from_dict(config_dict)
+        document_id = run_extraction(config)
+        return {
+            "success": True,
+            "document_id": document_id,
+            "error": None
+        }
+    except Exception as e:
+        logger.error(f"Extraction failed: {e}")
+        return {
+            "success": False,
+            "document_id": None,
+            "error": str(e)
+        }
+
+
+def extract_many_documents_with_ray(
+    configs: List[ExtractionConfig],
+    num_cpus: Optional[int] = None,
+    num_gpus: Optional[int] = None
+) -> List[Dict[str, Any]]:
+    """
+    Extract multiple documents in parallel using Ray.
+    
+    Args:
+        configs: List of extraction configurations
+        num_cpus: Number of CPUs to use (optional)
+        num_gpus: Number of GPUs to use (optional)
+        
+    Returns:
+        List of results, each containing document_id and status
+    """
+    if not _init_ray_if_needed():
+        logger.warning("Ray not available, falling back to sequential extraction")
+        results = []
+        for config in configs:
+            try:
+                document_id = run_extraction(config)
+                results.append({
+                    "success": True,
+                    "document_id": document_id,
+                    "error": None
+                })
+            except Exception as e:
+                results.append({
+                    "success": False,
+                    "document_id": None,
+                    "error": str(e)
+                })
+        return results
+    
+    # Convert configs to dictionaries for Ray serialization
+    config_dicts = [config.to_dict() for config in configs]
+    
+    # Submit tasks
+    logger.info(f"Submitting {len(config_dicts)} extraction tasks to Ray")
+    futures = [remote_extract_document.remote(cfg_dict) for cfg_dict in config_dicts]
+    
+    # Get results
+    results = ray.get(futures)
+    
+    # Log summary
+    successful = sum(1 for r in results if r["success"])
+    failed = len(results) - successful
+    logger.info(f"Extraction complete: {successful} successful, {failed} failed")
+    
+    return results
+
+
+def extract_folder_with_ray(
+    folder_path: str,
+    reader: str = "AUTO",
+    detect: Optional[List[str]] = None,
+    output_dir: str = "contextcore_data",
+    num_cpus: Optional[int] = None
+) -> List[Dict[str, Any]]:
+    """
+    Extract all supported files from a folder using Ray.
+    
+    Args:
+        folder_path: Path to folder containing files
+        reader: Extractor to use (default: "AUTO")
+        detect: List of modalities to detect (default: ["TEXT"])
+        output_dir: Output directory
+        num_cpus: Number of CPUs to use
+        
+    Returns:
+        List of extraction results
+    """
+    from pathlib import Path
+    
+    folder = Path(folder_path)
+    if not folder.exists():
+        raise FileNotFoundError(f"Folder not found: {folder_path}")
+    
+    # Find all supported files
+    supported_extensions = {
+        ".pdf", ".docx", ".pptx", ".csv", ".txt", ".html", ".htm",
+        ".png", ".jpg", ".jpeg", ".mp3", ".wav", ".mp4", ".avi"
+    }
+    
+    files = [
+        f for f in folder.iterdir()
+        if f.is_file() and f.suffix.lower() in supported_extensions
+    ]
+    
+    if not files:
+        logger.warning(f"No supported files found in {folder_path}")
+        return []
+    
+    logger.info(f"Found {len(files)} files to extract")
+    
+    # Create configs
+    configs = [
+        ExtractionConfig(
+            file_path=str(f),
+            reader=reader,
+            detect=detect or ["TEXT"],
+            output_dir=output_dir
+        )
+        for f in files
+    ]
+    
+    # Run with Ray
+    return extract_many_documents_with_ray(configs, num_cpus=num_cpus)
+
+
+def shutdown_ray() -> None:
+    """Shutdown Ray cluster."""
+    if RAY_AVAILABLE and ray.is_initialized():
+        ray.shutdown()
+        logger.info("Ray shutdown complete")

contextsynapse/extraction/redis_schema_registry.py

@@ -0,0 +1,222 @@
+"""
+Redis-backed Schema Registry — custom schemas stored in Redis instead of SQLite.
+
+Builtin schemas still read from disk (config/schemas/*.yaml).
+Custom schemas stored as Redis hashes: schema:custom:{name} → fields.
+
+Same public API as SchemaRegistry — drop-in replacement.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+_BUILTIN_DIR = Path(__file__).resolve().parent.parent / "config" / "schemas"
+
+INDUSTRY_SCHEMAS = {
+    "technology": "sdlc",
+    "healthcare": "healthcare",
+    "finance": "finance",
+    "legal": "legal",
+    "education": "education",
+    "retail": "retail",
+    "research": "knowledge_base",
+    "general": "knowledge_base",
+}
+
+
+class RedisSchemaRegistry:
+    """Manages predefined and custom extraction schemas with Redis backend."""
+
+    KEY_PREFIX = "schema:custom:"
+    NAMES_KEY = "schema:custom:names"
+
+    def __init__(self, redis_url: str = None, builtin_dir: str = None):
+        import redis
+        self._redis_url = redis_url or os.environ.get("CONTEXTSYNAPSE_REDIS_URL") or os.environ.get("AICONTEXTDB_REDIS_URL", "redis://localhost:6379/0")
+        self._r = redis.from_url(self._redis_url, decode_responses=True)
+        self._builtin_dir = Path(builtin_dir) if builtin_dir else _BUILTIN_DIR
+
+        # Migrate from SQLite if Redis is empty
+        self._migrate_from_sqlite()
+
+    def list_schemas(self) -> List[Dict[str, Any]]:
+        schemas = []
+
+        # Builtin schemas (from disk)
+        if self._builtin_dir.exists():
+            for f in sorted(self._builtin_dir.glob("*.yaml")):
+                try:
+                    raw = yaml.safe_load(f.read_text(encoding="utf-8"))
+                    name = raw.get("name", f.stem)
+                    node_types = list(raw.get("node_types", {}).keys())
+                    edge_types = list(raw.get("edge_types", {}).keys())
+                    schemas.append({
+                        "name": name,
+                        "source": "builtin",
+                        "file": f.name,
+                        "node_types": node_types,
+                        "edge_types": edge_types,
+                        "node_count": len(node_types),
+                        "edge_count": len(edge_types),
+                        "description": raw.get("description", ""),
+                    })
+                except Exception as e:
+                    logger.debug("Failed to load schema %s: %s", f.name, e)
+
+        # Custom schemas (from Redis)
+        names = self._r.smembers(self.NAMES_KEY)
+        if names:
+            pipe = self._r.pipeline()
+            for name in sorted(names):
+                pipe.hgetall(f"{self.KEY_PREFIX}{name}")
+            results = pipe.execute()
+
+            for meta in results:
+                if not meta:
+                    continue
+                try:
+                    raw = yaml.safe_load(meta.get("yaml_content", ""))
+                    node_types = list(raw.get("node_types", {}).keys()) if raw else []
+                    edge_types = list(raw.get("edge_types", {}).keys()) if raw else []
+                    schemas.append({
+                        "name": meta["name"],
+                        "source": "custom",
+                        "description": meta.get("description", ""),
+                        "industry": meta.get("industry", ""),
+                        "tags": json.loads(meta.get("tags", "[]")),
+                        "node_types": node_types,
+                        "edge_types": edge_types,
+                        "node_count": len(node_types),
+                        "edge_count": len(edge_types),
+                        "created_at": meta.get("created_at", ""),
+                    })
+                except Exception:
+                    pass
+
+        return schemas
+
+    def get_schema(self, name: str) -> Optional[str]:
+        # Builtin first
+        if self._builtin_dir.exists():
+            for f in self._builtin_dir.glob("*.yaml"):
+                try:
+                    raw = yaml.safe_load(f.read_text(encoding="utf-8"))
+                    if raw.get("name") == name or f.stem == name:
+                        return f.read_text(encoding="utf-8")
+                except Exception:
+                    pass
+
+        # Custom from Redis
+        content = self._r.hget(f"{self.KEY_PREFIX}{name}", "yaml_content")
+        if content:
+            return content
+
+        # Industry mapping
+        mapped = INDUSTRY_SCHEMAS.get(name)
+        if mapped and mapped != name:
+            return self.get_schema(mapped)
+
+        return None
+
+    def get_schema_object(self, name: str):
+        yaml_content = self.get_schema(name)
+        if not yaml_content:
+            return None
+        from .schema_loader import load_schema_from_yaml_string
+        return load_schema_from_yaml_string(yaml_content, name=name)
+
+    def save_custom(
+        self,
+        name: str,
+        yaml_content: str,
+        description: str = "",
+        industry: str = "",
+        tags: Optional[List[str]] = None,
+    ) -> Dict[str, Any]:
+        raw = yaml.safe_load(yaml_content)
+        if not isinstance(raw, dict) or "node_types" not in raw:
+            raise ValueError("Schema must be a YAML mapping with 'node_types' key")
+
+        now = datetime.now(timezone.utc).isoformat()
+
+        # Preserve original created_at
+        existing = self._r.hget(f"{self.KEY_PREFIX}{name}", "created_at")
+        created_at = existing or now
+
+        pipe = self._r.pipeline()
+        pipe.hset(f"{self.KEY_PREFIX}{name}", mapping={
+            "name": name,
+            "yaml_content": yaml_content,
+            "description": description,
+            "industry": industry,
+            "tags": json.dumps(tags or []),
+            "created_at": created_at,
+            "updated_at": now,
+        })
+        pipe.sadd(self.NAMES_KEY, name)
+        pipe.execute()
+
+        node_types = list(raw.get("node_types", {}).keys())
+        return {"name": name, "node_types": node_types, "status": "saved"}
+
+    def delete_custom(self, name: str) -> bool:
+        pipe = self._r.pipeline()
+        pipe.delete(f"{self.KEY_PREFIX}{name}")
+        pipe.srem(self.NAMES_KEY, name)
+        results = pipe.execute()
+        return results[0] > 0  # True if the hash was deleted
+
+    def get_for_industry(self, industry: str) -> Optional[str]:
+        mapped = INDUSTRY_SCHEMAS.get(industry.lower())
+        if mapped:
+            return self.get_schema(mapped)
+        return self.get_schema("knowledge_base")
+
+    # ------------------------------------------------------------------
+    # Migration from SQLite
+    # ------------------------------------------------------------------
+
+    def _migrate_from_sqlite(self):
+        if self._r.scard(self.NAMES_KEY) > 0:
+            return
+
+        db_path = "contextcore_data/schemas.db"
+        if not os.path.exists(db_path):
+            return
+
+        try:
+            import sqlite3
+            conn = sqlite3.connect(db_path)
+            conn.row_factory = sqlite3.Row
+            rows = conn.execute("SELECT * FROM custom_schemas").fetchall()
+            conn.close()
+
+            if not rows:
+                return
+
+            pipe = self._r.pipeline()
+            for r in rows:
+                pipe.hset(f"{self.KEY_PREFIX}{r['name']}", mapping={
+                    "name": r["name"],
+                    "yaml_content": r["yaml_content"],
+                    "description": r["description"] or "",
+                    "industry": r["industry"] or "",
+                    "tags": r["tags"] or "[]",
+                    "created_at": r["created_at"],
+                    "updated_at": r["updated_at"],
+                })
+                pipe.sadd(self.NAMES_KEY, r["name"])
+            pipe.execute()
+            logger.info("[REDIS-SCHEMA] Migrated %d custom schemas from SQLite to Redis", len(rows))
+        except Exception as e:
+            logger.warning("[REDIS-SCHEMA] Migration from SQLite failed: %s", e)