ds-agent-cli 0.1.0

package/src/storage/user_files_service.py

@@ -0,0 +1,288 @@
+"""
+User Files Service - Manages file metadata in Supabase
+
+This service:
+1. Tracks all user files (plots, CSVs, reports, models) in Supabase
+2. Provides file listing for the Assets panel
+3. Handles file expiration and cleanup coordination
+4. Works with R2StorageService for actual file storage
+"""
+
+import os
+from datetime import datetime, timedelta
+from typing import Optional, Dict, Any, List
+from dataclasses import dataclass
+from enum import Enum
+
+# Supabase client import
+try:
+    from supabase import create_client, Client
+except ImportError:
+    print("Warning: supabase package not installed. Run: pip install supabase")
+    Client = None
+
+SUPABASE_URL = os.getenv("SUPABASE_URL", "")
+SUPABASE_SERVICE_KEY = os.getenv("SUPABASE_SERVICE_KEY", "")  # Use service key for backend
+
+
+class FileType(Enum):
+    PLOT = "plot"
+    CSV = "csv"
+    REPORT = "report"
+    MODEL = "model"
+
+
+@dataclass
+class UserFile:
+    """Represents a user file record."""
+    id: str
+    user_id: str
+    session_id: Optional[str]
+    file_type: FileType
+    file_name: str
+    r2_key: str
+    size_bytes: int
+    mime_type: str
+    metadata: Dict[str, Any]
+    created_at: datetime
+    expires_at: datetime
+    download_url: Optional[str] = None
+
+
+class UserFilesService:
+    """Service for managing user file metadata in Supabase."""
+    
+    def __init__(self):
+        """Initialize Supabase client."""
+        if not SUPABASE_URL or not SUPABASE_SERVICE_KEY:
+            raise ValueError("SUPABASE_URL and SUPABASE_SERVICE_KEY must be set")
+        
+        self.client: Client = create_client(SUPABASE_URL, SUPABASE_SERVICE_KEY)
+        self.table = "user_files"
+    
+    # ==================== CREATE ====================
+    
+    def create_file_record(
+        self,
+        user_id: str,
+        file_type: FileType,
+        file_name: str,
+        r2_key: str,
+        size_bytes: int,
+        session_id: Optional[str] = None,
+        mime_type: str = "application/octet-stream",
+        metadata: Optional[Dict[str, Any]] = None,
+        expires_in_days: int = 7
+    ) -> UserFile:
+        """
+        Create a file record in Supabase.
+        
+        Args:
+            user_id: User ID
+            file_type: Type of file
+            file_name: Display name
+            r2_key: R2 storage key
+            size_bytes: File size
+            session_id: Optional chat session ID
+            mime_type: MIME type
+            metadata: Additional metadata (plot type, metrics, etc.)
+            expires_in_days: Days until file expires
+            
+        Returns:
+            Created UserFile record
+        """
+        expires_at = datetime.utcnow() + timedelta(days=expires_in_days)
+        
+        data = {
+            "user_id": user_id,
+            "session_id": session_id,
+            "file_type": file_type.value,
+            "file_name": file_name,
+            "r2_key": r2_key,
+            "size_bytes": size_bytes,
+            "mime_type": mime_type,
+            "metadata": metadata or {},
+            "expires_at": expires_at.isoformat()
+        }
+        
+        result = self.client.table(self.table).insert(data).execute()
+        
+        if result.data:
+            return self._to_user_file(result.data[0])
+        raise Exception("Failed to create file record")
+    
+    # ==================== READ ====================
+    
+    def get_user_files(
+        self,
+        user_id: str,
+        file_type: Optional[FileType] = None,
+        session_id: Optional[str] = None,
+        include_expired: bool = False
+    ) -> List[UserFile]:
+        """
+        Get all files for a user.
+        
+        Args:
+            user_id: User ID
+            file_type: Optional filter by type
+            session_id: Optional filter by session
+            include_expired: Include expired files
+            
+        Returns:
+            List of UserFile records
+        """
+        query = self.client.table(self.table)\
+            .select("*")\
+            .eq("user_id", user_id)\
+            .eq("is_deleted", False)
+        
+        if file_type:
+            query = query.eq("file_type", file_type.value)
+        
+        if session_id:
+            query = query.eq("session_id", session_id)
+        
+        if not include_expired:
+            query = query.gt("expires_at", datetime.utcnow().isoformat())
+        
+        query = query.order("created_at", desc=True)
+        
+        result = query.execute()
+        
+        return [self._to_user_file(row) for row in (result.data or [])]
+    
+    def get_file_by_id(self, file_id: str) -> Optional[UserFile]:
+        """Get a specific file by ID."""
+        result = self.client.table(self.table)\
+            .select("*")\
+            .eq("id", file_id)\
+            .single()\
+            .execute()
+        
+        if result.data:
+            return self._to_user_file(result.data)
+        return None
+    
+    def get_file_by_r2_key(self, r2_key: str) -> Optional[UserFile]:
+        """Get a file by R2 key."""
+        result = self.client.table(self.table)\
+            .select("*")\
+            .eq("r2_key", r2_key)\
+            .single()\
+            .execute()
+        
+        if result.data:
+            return self._to_user_file(result.data)
+        return None
+    
+    def get_session_files(self, session_id: str) -> List[UserFile]:
+        """Get all files for a chat session."""
+        result = self.client.table(self.table)\
+            .select("*")\
+            .eq("session_id", session_id)\
+            .eq("is_deleted", False)\
+            .order("created_at", desc=True)\
+            .execute()
+        
+        return [self._to_user_file(row) for row in (result.data or [])]
+    
+    # ==================== UPDATE ====================
+    
+    def extend_expiration(self, file_id: str, additional_days: int = 7) -> bool:
+        """Extend file expiration date."""
+        file = self.get_file_by_id(file_id)
+        if not file:
+            return False
+        
+        new_expires = datetime.utcnow() + timedelta(days=additional_days)
+        
+        result = self.client.table(self.table)\
+            .update({"expires_at": new_expires.isoformat()})\
+            .eq("id", file_id)\
+            .execute()
+        
+        return bool(result.data)
+    
+    # ==================== DELETE ====================
+    
+    def soft_delete_file(self, file_id: str) -> bool:
+        """Soft delete a file (mark as deleted)."""
+        result = self.client.table(self.table)\
+            .update({"is_deleted": True})\
+            .eq("id", file_id)\
+            .execute()
+        
+        return bool(result.data)
+    
+    def hard_delete_file(self, file_id: str) -> bool:
+        """Permanently delete a file record."""
+        result = self.client.table(self.table)\
+            .delete()\
+            .eq("id", file_id)\
+            .execute()
+        
+        return bool(result.data)
+    
+    def get_expired_files(self) -> List[UserFile]:
+        """Get all expired files for cleanup."""
+        result = self.client.table(self.table)\
+            .select("*")\
+            .lt("expires_at", datetime.utcnow().isoformat())\
+            .eq("is_deleted", False)\
+            .execute()
+        
+        return [self._to_user_file(row) for row in (result.data or [])]
+    
+    # ==================== STATS ====================
+    
+    def get_user_storage_stats(self, user_id: str) -> Dict[str, Any]:
+        """Get storage statistics for a user."""
+        files = self.get_user_files(user_id, include_expired=False)
+        
+        stats = {
+            "total_files": len(files),
+            "total_size_bytes": sum(f.size_bytes for f in files),
+            "by_type": {}
+        }
+        
+        for file_type in FileType:
+            type_files = [f for f in files if f.file_type == file_type]
+            stats["by_type"][file_type.value] = {
+                "count": len(type_files),
+                "size_bytes": sum(f.size_bytes for f in type_files)
+            }
+        
+        stats["total_size_mb"] = round(stats["total_size_bytes"] / (1024 * 1024), 2)
+        
+        return stats
+    
+    # ==================== HELPERS ====================
+    
+    def _to_user_file(self, row: Dict[str, Any]) -> UserFile:
+        """Convert database row to UserFile object."""
+        return UserFile(
+            id=row["id"],
+            user_id=row["user_id"],
+            session_id=row.get("session_id"),
+            file_type=FileType(row["file_type"]),
+            file_name=row["file_name"],
+            r2_key=row["r2_key"],
+            size_bytes=row.get("size_bytes", 0),
+            mime_type=row.get("mime_type", "application/octet-stream"),
+            metadata=row.get("metadata", {}),
+            created_at=datetime.fromisoformat(row["created_at"].replace("Z", "+00:00")),
+            expires_at=datetime.fromisoformat(row["expires_at"].replace("Z", "+00:00"))
+        )
+
+
+# ==================== SINGLETON ====================
+
+_files_service: Optional[UserFilesService] = None
+
+def get_files_service() -> UserFilesService:
+    """Get or create UserFilesService singleton."""
+    global _files_service
+    if _files_service is None:
+        _files_service = UserFilesService()
+    return _files_service

package/src/tools/__init__.py

@@ -0,0 +1,335 @@
+"""Tools module initialization - All 44 tools."""
+
+# Basic Tools (10)
+from .data_profiling import (
+    profile_dataset,
+    detect_data_quality_issues,
+    analyze_correlations,
+    get_smart_summary,  # NEW: Enhanced data summary
+    detect_label_errors  # NEW: cleanlab label error detection
+)
+
+from .data_cleaning import (
+    clean_missing_values,
+    handle_outliers,
+    fix_data_types
+)
+
+from .data_type_conversion import (
+    force_numeric_conversion,
+    smart_type_inference
+)
+
+# Data Wrangling Tools (3) - NEW
+from .data_wrangling import (
+    merge_datasets,
+    concat_datasets,
+    reshape_dataset
+)
+
+from .feature_engineering import (
+    create_time_features,
+    encode_categorical
+)
+
+from .model_training import (
+    train_baseline_models,
+    generate_model_report
+)
+
+# AutoGluon-Powered Training (9) - Classification, Regression, Time Series, Optimization
+from .autogluon_training import (
+    train_with_autogluon,
+    predict_with_autogluon,
+    forecast_with_autogluon,
+    optimize_autogluon_model,
+    analyze_autogluon_model,
+    extend_autogluon_training,
+    train_multilabel_autogluon,
+    backtest_timeseries,
+    analyze_timeseries_model
+)
+
+# Advanced Analysis Tools (5)
+from .advanced_analysis import (
+    perform_eda_analysis,
+    detect_model_issues,
+    detect_anomalies,
+    detect_and_handle_multicollinearity,
+    perform_statistical_tests
+)
+
+# Advanced Feature Engineering Tools (4)
+from .advanced_feature_engineering import (
+    create_interaction_features,
+    create_aggregation_features,
+    engineer_text_features,
+    auto_feature_engineering
+)
+
+# Advanced Preprocessing Tools (3)
+from .advanced_preprocessing import (
+    handle_imbalanced_data,
+    perform_feature_scaling,
+    split_data_strategically
+)
+
+# Advanced Training Tools (3)
+from .advanced_training import (
+    hyperparameter_tuning,
+    train_ensemble_models,
+    perform_cross_validation
+)
+
+# Business Intelligence Tools (4)
+from .business_intelligence import (
+    perform_cohort_analysis,
+    perform_rfm_analysis,
+    detect_causal_relationships,
+    generate_business_insights
+)
+
+# Computer Vision Tools (3)
+from .computer_vision import (
+    extract_image_features,
+    perform_image_clustering,
+    analyze_tabular_image_hybrid
+)
+
+# NLP/Text Analytics Tools (4)
+from .nlp_text_analytics import (
+    perform_topic_modeling,
+    perform_named_entity_recognition,
+    analyze_sentiment_advanced,
+    perform_text_similarity
+)
+
+# Production/MLOps Tools (5 + 2 new)
+from .production_mlops import (
+    monitor_model_drift,
+    explain_predictions,
+    generate_model_card,
+    perform_ab_test_analysis,
+    detect_feature_leakage,
+    monitor_drift_evidently,  # NEW: Evidently drift reports
+    explain_with_dtreeviz  # NEW: Decision tree visualization
+)
+
+# Time Series Tools (3)
+from .time_series import (
+    forecast_time_series,
+    detect_seasonality_trends,
+    create_time_series_features
+)
+
+# Advanced Insights Tools (6) - NEW
+from .advanced_insights import (
+    analyze_root_cause,
+    detect_trends_and_seasonality,
+    detect_anomalies_advanced,
+    perform_hypothesis_testing,
+    analyze_distribution,
+    perform_segment_analysis
+)
+
+# Automated Pipeline Tools (2) - NEW
+from .auto_pipeline import (
+    auto_ml_pipeline,
+    auto_feature_selection
+)
+
+# Visualization Tools (5) - NEW
+from .visualization_engine import (
+    generate_all_plots,
+    generate_data_quality_plots,
+    generate_eda_plots,
+    generate_model_performance_plots,
+    generate_feature_importance_plot
+)
+
+# Interactive Plotly Visualizations (6) - NEW PHASE 2
+from .plotly_visualizations import (
+    generate_interactive_scatter,
+    generate_interactive_histogram,
+    generate_interactive_correlation_heatmap,
+    generate_interactive_box_plots,
+    generate_interactive_time_series,
+    generate_plotly_dashboard
+)
+
+# EDA Report Generation (2) - NEW PHASE 2
+from .eda_reports import (
+    generate_ydata_profiling_report,
+    generate_sweetviz_report  # NEW: Sweetviz EDA with comparison
+)
+
+# Code Interpreter (2) - NEW PHASE 2 - CRITICAL for True AI Agent
+from .code_interpreter import (
+    execute_python_code,
+    execute_code_from_file
+)
+
+# Cloud Data Sources (4) - NEW: BigQuery Integration
+from .cloud_data_sources import (
+    load_bigquery_table,
+    write_bigquery_table,
+    profile_bigquery_table,
+    query_bigquery
+)
+
+from .tools_registry import TOOLS, get_tool_by_name, get_all_tool_names
+
+from .enhanced_feature_engineering import (
+    create_ratio_features,
+    create_statistical_features,
+    create_log_features,
+    create_binned_features,
+)
+
+__all__ = [
+    # Basic Data Profiling (4 + 1 new) - UPDATED
+    "profile_dataset",
+    "detect_data_quality_issues",
+    "analyze_correlations",
+    "get_smart_summary",  # NEW
+    "detect_label_errors",  # NEW: cleanlab
+    
+    # Basic Data Cleaning (3)
+    "clean_missing_values",
+    "handle_outliers",
+    "fix_data_types",
+    
+    # Data Type Conversion (2)
+    "force_numeric_conversion",
+    "smart_type_inference",
+    
+    # Data Wrangling (3) - NEW
+    "merge_datasets",
+    "concat_datasets",
+    "reshape_dataset",
+    
+    # Basic Feature Engineering (2)
+    "create_time_features",
+    "encode_categorical",
+    
+    # Basic Model Training (2)
+    "train_baseline_models",
+    "generate_model_report",
+    
+    # AutoGluon Training (9) - NEW
+    "train_with_autogluon",
+    "predict_with_autogluon",
+    "forecast_with_autogluon",
+    "optimize_autogluon_model",
+    "analyze_autogluon_model",
+    "extend_autogluon_training",
+    "train_multilabel_autogluon",
+    "backtest_timeseries",
+    "analyze_timeseries_model",
+    
+    # Advanced Analysis (5)
+    "perform_eda_analysis",
+    "detect_model_issues",
+    "detect_anomalies",
+    "detect_and_handle_multicollinearity",
+    "perform_statistical_tests",
+    
+    # Advanced Feature Engineering (4)
+    "create_interaction_features",
+    "create_aggregation_features",
+    "engineer_text_features",
+    "auto_feature_engineering",
+    
+    # Advanced Preprocessing (3)
+    "handle_imbalanced_data",
+    "perform_feature_scaling",
+    "split_data_strategically",
+    
+    # Advanced Training (3)
+    "hyperparameter_tuning",
+    "train_ensemble_models",
+    "perform_cross_validation",
+    
+    # Business Intelligence (4)
+    "perform_cohort_analysis",
+    "perform_rfm_analysis",
+    "detect_causal_relationships",
+    "generate_business_insights",
+    
+    # Computer Vision (3)
+    "extract_image_features",
+    "perform_image_clustering",
+    "analyze_tabular_image_hybrid",
+    
+    # NLP/Text Analytics (4)
+    "perform_topic_modeling",
+    "perform_named_entity_recognition",
+    "analyze_sentiment_advanced",
+    "perform_text_similarity",
+    
+    # Production/MLOps (5 + 2 new)
+    "monitor_model_drift",
+    "explain_predictions",
+    "generate_model_card",
+    "perform_ab_test_analysis",
+    "detect_feature_leakage",
+    "monitor_drift_evidently",  # NEW: Evidently
+    "explain_with_dtreeviz",  # NEW: dtreeviz
+    
+    # Time Series (3)
+    "forecast_time_series",
+    "detect_seasonality_trends",
+    "create_time_series_features",
+    
+    # Advanced Insights (6) - NEW
+    "analyze_root_cause",
+    "detect_trends_and_seasonality",
+    "detect_anomalies_advanced",
+    "perform_hypothesis_testing",
+    "analyze_distribution",
+    "perform_segment_analysis",
+    
+    # Automated Pipeline (2) - NEW
+    "auto_ml_pipeline",
+    "auto_feature_selection",
+    
+    # Visualization (5) - NEW
+    "generate_all_plots",
+    "generate_data_quality_plots",
+    "generate_eda_plots",
+    "generate_model_performance_plots",
+    "generate_feature_importance_plot",
+    
+    # Interactive Plotly Visualizations (6) - NEW PHASE 2
+    "generate_interactive_scatter",
+    "generate_interactive_histogram",
+    "generate_interactive_correlation_heatmap",
+    "generate_interactive_box_plots",
+    "generate_interactive_time_series",
+    "generate_plotly_dashboard",
+    
+    # EDA Report Generation (2) - NEW PHASE 2
+    "generate_ydata_profiling_report",
+    "generate_sweetviz_report",  # NEW: Sweetviz
+    
+    # Code Interpreter (2) - NEW PHASE 2 - CRITICAL for True AI Agent
+    "execute_python_code",
+    "execute_code_from_file",
+    
+    # Cloud Data Sources (4) - NEW: BigQuery Integration
+    "load_bigquery_table",
+    "write_bigquery_table",
+    "profile_bigquery_table",
+    "query_bigquery",
+    
+    # Enhanced Feature Engineering (4) - NEW
+    "create_ratio_features",
+    "create_statistical_features",
+    "create_log_features",
+    "create_binned_features",
+    
+    # Registry
+    "TOOLS",
+    "get_tool_by_name",
+    "get_all_tool_names",
+]