Create artifact_registry.py

artifact_registry.py

@@ -0,0 +1,384 @@
+"""
+artifact_registry.py - Artifact Tracking & Management
+======================================================
+
+Tracks all files created during execution to ensure user access.
+Provides download manifests and file management.
+
+Author: AI Lab Team
+Last Updated: 2025-10-10
+Version: 1.0
+"""
+
+import json
+import os
+import shutil
+import uuid
+from datetime import datetime
+from typing import List, Dict, Any, Optional
+from logging_config import get_logger
+
+# Import config
+try:
+    import graph_config as cfg
+except ImportError:
+    class cfg:
+        USER_ARTIFACTS_DIR = "outputs/user_artifacts"
+        ARTIFACT_REGISTRY_FILE = "outputs/artifact_registry.json"
+        ARTIFACT_RETENTION_DAYS = 30
+        MAX_ARTIFACT_SIZE_MB = 100
+
+log = get_logger(__name__)
+
+
+class ArtifactRegistry:
+    """
+    Registry for tracking all created artifacts.
+    
+    Features:
+    - Tracks file creation
+    - Ensures user accessibility
+    - Provides download manifests
+    - Manages file lifecycle
+    """
+    
+    def __init__(self, registry_path: str = None):
+        """
+        Initialize artifact registry.
+        
+        Args:
+            registry_path: Path to registry JSON file
+        """
+        self.registry_path = registry_path or cfg.ARTIFACT_REGISTRY_FILE
+        self.artifacts = self._load_registry()
+        
+        # Ensure user artifacts directory exists
+        os.makedirs(cfg.USER_ARTIFACTS_DIR, exist_ok=True)
+        
+        log.info(f"Artifact Registry initialized: {len(self.artifacts)} artifacts")
+    
+    def register_artifact(self, filepath: str, metadata: Dict[str, Any]) -> tuple:
+        """
+        Register an artifact in the system.
+        
+        Args:
+            filepath: Original file path
+            metadata: Artifact metadata (type, session_id, user_request, etc.)
+        
+        Returns:
+            Tuple of (artifact_id, user_accessible_path)
+        """
+        # Generate unique ID
+        artifact_id = uuid.uuid4().hex[:8]
+        
+        # Ensure file is in user-accessible location
+        user_path = self._ensure_user_accessible(filepath)
+        
+        # Get file stats
+        try:
+            size_bytes = os.path.getsize(user_path)
+            size_mb = size_bytes / 1_000_000
+            
+            # Check size limit
+            if size_mb > cfg.MAX_ARTIFACT_SIZE_MB:
+                log.warning(
+                    f"⚠️ Artifact exceeds size limit: {size_mb:.2f}MB > "
+                    f"{cfg.MAX_ARTIFACT_SIZE_MB}MB"
+                )
+        except Exception:
+            size_bytes = 0
+            size_mb = 0.0
+        
+        # Create registry entry
+        self.artifacts[artifact_id] = {
+            "id": artifact_id,
+            "original_path": filepath,
+            "user_path": user_path,
+            "filename": os.path.basename(user_path),
+            "created_at": datetime.utcnow().isoformat(),
+            "size_bytes": size_bytes,
+            "size_mb": round(size_mb, 2),
+            "type": metadata.get("type", "unknown"),
+            "session_id": metadata.get("session_id"),
+            "user_request": metadata.get("user_request", "")[:200],
+            "language": metadata.get("language"),
+            "tier": metadata.get("tier"),
+            "metadata": metadata
+        }
+        
+        # Save registry
+        self._save_registry()
+        
+        log.info(f"✅ Artifact registered: {artifact_id} -> {os.path.basename(user_path)}")
+        
+        return artifact_id, user_path
+    
+    def _ensure_user_accessible(self, filepath: str) -> str:
+        """
+        Copy file to user-accessible location if needed.
+        
+        Args:
+            filepath: Original file path
+        
+        Returns:
+            Path in user-accessible directory
+        """
+        # If already in user directory, return as-is
+        if filepath.startswith(cfg.USER_ARTIFACTS_DIR):
+            return filepath
+        
+        # If file doesn't exist, return original path
+        if not os.path.exists(filepath):
+            log.warning(f"⚠️ File not found: {filepath}")
+            return filepath
+        
+        # Copy to user directory
+        filename = os.path.basename(filepath)
+        user_path = os.path.join(cfg.USER_ARTIFACTS_DIR, filename)
+        
+        # Handle name collisions
+        counter = 1
+        base, ext = os.path.splitext(filename)
+        while os.path.exists(user_path):
+            user_path = os.path.join(
+                cfg.USER_ARTIFACTS_DIR, 
+                f"{base}_{counter}{ext}"
+            )
+            counter += 1
+        
+        try:
+            shutil.copy2(filepath, user_path)
+            log.info(f"📋 Copied to user dir: {filename}")
+        except Exception as e:
+            log.error(f"Failed to copy {filepath}: {e}")
+            return filepath
+        
+        return user_path
+    
+    def get_session_artifacts(self, session_id: str) -> List[Dict[str, Any]]:
+        """
+        Get all artifacts for a session.
+        
+        Args:
+            session_id: Session identifier
+        
+        Returns:
+            List of artifact dicts
+        """
+        return [
+            artifact for artifact in self.artifacts.values()
+            if artifact.get("session_id") == session_id
+        ]
+    
+    def get_download_manifest(self, session_id: str) -> Dict[str, Any]:
+        """
+        Create download manifest for UI.
+        
+        Args:
+            session_id: Session identifier
+        
+        Returns:
+            Dict with download information
+        """
+        artifacts = self.get_session_artifacts(session_id)
+        
+        total_size_mb = sum(a.get("size_mb", 0) for a in artifacts)
+        
+        return {
+            "session_id": session_id,
+            "total_files": len(artifacts),
+            "total_size_mb": round(total_size_mb, 2),
+            "files": [
+                {
+                    "id": a["id"],
+                    "filename": a["filename"],
+                    "path": a["user_path"],
+                    "type": a["type"],
+                    "size_kb": round(a["size_bytes"] / 1000, 1),
+                    "size_mb": a.get("size_mb", 0),
+                    "created": a["created_at"],
+                    "language": a.get("language")
+                }
+                for a in artifacts
+            ]
+        }
+    
+    def get_artifact_by_id(self, artifact_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Get artifact by ID.
+        
+        Args:
+            artifact_id: Artifact identifier
+        
+        Returns:
+            Artifact dict or None
+        """
+        return self.artifacts.get(artifact_id)
+    
+    def delete_artifact(self, artifact_id: str) -> bool:
+        """
+        Delete an artifact.
+        
+        Args:
+            artifact_id: Artifact identifier
+        
+        Returns:
+            True if deleted successfully
+        """
+        artifact = self.artifacts.get(artifact_id)
+        if not artifact:
+            return False
+        
+        # Delete file
+        user_path = artifact.get("user_path")
+        if user_path and os.path.exists(user_path):
+            try:
+                os.remove(user_path)
+                log.info(f"🗑️ Deleted file: {os.path.basename(user_path)}")
+            except Exception as e:
+                log.error(f"Failed to delete {user_path}: {e}")
+                return False
+        
+        # Remove from registry
+        del self.artifacts[artifact_id]
+        self._save_registry()
+        
+        return True
+    
+    def cleanup_old_artifacts(self, days: int = None):
+        """
+        Delete artifacts older than specified days.
+        
+        Args:
+            days: Age threshold in days (uses config default if None)
+        """
+        days = days or cfg.ARTIFACT_RETENTION_DAYS
+        cutoff = datetime.utcnow().timestamp() - (days * 86400)
+        
+        deleted_count = 0
+        
+        for artifact_id, artifact in list(self.artifacts.items()):
+            created_str = artifact.get("created_at", "")
+            try:
+                created_dt = datetime.fromisoformat(created_str)
+                created_ts = created_dt.timestamp()
+                
+                if created_ts < cutoff:
+                    if self.delete_artifact(artifact_id):
+                        deleted_count += 1
+            except Exception as e:
+                log.debug(f"Cleanup skip {artifact_id}: {e}")
+        
+        if deleted_count > 0:
+            log.info(f"🧹 Cleaned up {deleted_count} old artifacts")
+    
+    def get_all_artifacts(self) -> List[Dict[str, Any]]:
+        """
+        Get all artifacts in registry.
+        
+        Returns:
+            List of all artifact dicts
+        """
+        return list(self.artifacts.values())
+    
+    def get_registry_stats(self) -> Dict[str, Any]:
+        """
+        Get registry statistics.
+        
+        Returns:
+            Dict with stats
+        """
+        artifacts = self.get_all_artifacts()
+        
+        total_size = sum(a.get("size_bytes", 0) for a in artifacts)
+        
+        # Count by type
+        type_counts = {}
+        for a in artifacts:
+            atype = a.get("type", "unknown")
+            type_counts[atype] = type_counts.get(atype, 0) + 1
+        
+        # Count by session
+        session_counts = {}
+        for a in artifacts:
+            session = a.get("session_id", "unknown")
+            session_counts[session] = session_counts.get(session, 0) + 1
+        
+        return {
+            "total_artifacts": len(artifacts),
+            "total_size_mb": round(total_size / 1_000_000, 2),
+            "by_type": type_counts,
+            "by_session": len(session_counts),
+            "unique_sessions": list(session_counts.keys())
+        }
+    
+    def _load_registry(self) -> Dict[str, Dict[str, Any]]:
+        """
+        Load registry from disk.
+        
+        Returns:
+            Dict of artifacts
+        """
+        if not os.path.exists(self.registry_path):
+            return {}
+        
+        try:
+            with open(self.registry_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+        except Exception as e:
+            log.error(f"Failed to load registry: {e}")
+            return {}
+    
+    def _save_registry(self):
+        """Save registry to disk."""
+        # Ensure directory exists
+        os.makedirs(os.path.dirname(self.registry_path), exist_ok=True)
+        
+        try:
+            with open(self.registry_path, 'w', encoding='utf-8') as f:
+                json.dump(self.artifacts, f, indent=2)
+        except Exception as e:
+            log.error(f"Failed to save registry: {e}")
+    
+    def export_manifest(self, output_path: str = None) -> str:
+        """
+        Export complete manifest to JSON file.
+        
+        Args:
+            output_path: Output file path
+        
+        Returns:
+            Path to exported manifest
+        """
+        output_path = output_path or os.path.join(
+            cfg.USER_ARTIFACTS_DIR, 
+            f"manifest_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.json"
+        )
+        
+        manifest = {
+            "exported_at": datetime.utcnow().isoformat(),
+            "stats": self.get_registry_stats(),
+            "artifacts": self.get_all_artifacts()
+        }
+        
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        
+        with open(output_path, 'w', encoding='utf-8') as f:
+            json.dump(manifest, f, indent=2)
+        
+        log.info(f"📊 Manifest exported: {output_path}")
+        return output_path
+
+
+# Global instance
+artifact_registry = ArtifactRegistry()
+
+
+# ============================================================================
+# EXPORTS
+# ============================================================================
+
+__all__ = [
+    'ArtifactRegistry',
+    'artifact_registry'
+]