stratifyai 0.1.0__py3-none-any.whl

stratifyai/utils/csv_extractor.py

@@ -0,0 +1,197 @@
+"""CSV/DataFrame schema extraction for intelligent file analysis.
+
+This module extracts compact schema information from CSV files to reduce
+token usage by 99%+ while preserving essential structure information.
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+import pandas as pd
+
+
+@dataclass
+class ColumnSchema:
+    """Schema information for a single column."""
+    name: str
+    dtype: str
+    null_count: int
+    null_percentage: float
+    unique_count: int
+    sample_values: List[Any]
+    numeric_stats: Optional[Dict[str, float]] = None
+
+
+@dataclass
+class CSVSchema:
+    """Complete schema information for a CSV file."""
+    file_path: str
+    row_count: int
+    column_count: int
+    columns: List[ColumnSchema]
+    memory_usage_mb: float
+    
+    def to_text(self) -> str:
+        """Convert schema to human-readable text representation.
+        
+        Returns:
+            Formatted schema description
+        """
+        lines = [
+            f"CSV File: {self.file_path}",
+            f"Dimensions: {self.row_count:,} rows × {self.column_count} columns",
+            f"Memory: {self.memory_usage_mb:.2f} MB",
+            "",
+            "Column Schema:"
+        ]
+        
+        for col in self.columns:
+            # Basic info
+            lines.append(f"\n  {col.name} ({col.dtype})")
+            lines.append(f"    - Null: {col.null_count:,} ({col.null_percentage:.1f}%)")
+            lines.append(f"    - Unique: {col.unique_count:,}")
+            
+            # Numeric stats if available
+            if col.numeric_stats:
+                stats = col.numeric_stats
+                lines.append(f"    - Range: {stats['min']:.2f} to {stats['max']:.2f}")
+                lines.append(f"    - Mean: {stats['mean']:.2f}, Median: {stats['median']:.2f}")
+                lines.append(f"    - Std: {stats['std']:.2f}")
+            
+            # Sample values
+            samples_str = ", ".join(str(v) for v in col.sample_values[:5])
+            lines.append(f"    - Samples: {samples_str}")
+        
+        return "\n".join(lines)
+
+
+def extract_csv_schema(
+    file_path: Path,
+    sample_size: int = 5,
+    max_rows: Optional[int] = None
+) -> CSVSchema:
+    """Extract schema information from a CSV file.
+    
+    Args:
+        file_path: Path to CSV file
+        sample_size: Number of sample values to extract per column
+        max_rows: Maximum number of rows to read (None = all)
+        
+    Returns:
+        CSVSchema object with extracted information
+        
+    Raises:
+        FileNotFoundError: If file doesn't exist
+        pd.errors.EmptyDataError: If CSV is empty
+        pd.errors.ParserError: If CSV is malformed
+    """
+    if not file_path.exists():
+        raise FileNotFoundError(f"CSV file not found: {file_path}")
+    
+    # Read CSV
+    df = pd.read_csv(file_path, nrows=max_rows)
+    
+    if df.empty:
+        raise pd.errors.EmptyDataError(f"CSV file is empty: {file_path}")
+    
+    # Extract column schemas
+    columns = []
+    for col_name in df.columns:
+        col_data = df[col_name]
+        
+        # Basic stats
+        null_count = col_data.isna().sum()
+        null_pct = (null_count / len(df)) * 100
+        unique_count = col_data.nunique()
+        
+        # Sample values (exclude nulls)
+        non_null_values = col_data.dropna()
+        if len(non_null_values) > 0:
+            sample_values = non_null_values.sample(
+                min(sample_size, len(non_null_values)),
+                random_state=42
+            ).tolist()
+        else:
+            sample_values = []
+        
+        # Numeric statistics if applicable
+        numeric_stats = None
+        if pd.api.types.is_numeric_dtype(col_data):
+            try:
+                numeric_stats = {
+                    'min': float(col_data.min()),
+                    'max': float(col_data.max()),
+                    'mean': float(col_data.mean()),
+                    'median': float(col_data.median()),
+                    'std': float(col_data.std())
+                }
+            except (ValueError, TypeError):
+                # Handle edge cases (e.g., all NaN)
+                pass
+        
+        columns.append(ColumnSchema(
+            name=col_name,
+            dtype=str(col_data.dtype),
+            null_count=int(null_count),
+            null_percentage=float(null_pct),
+            unique_count=int(unique_count),
+            sample_values=sample_values,
+            numeric_stats=numeric_stats
+        ))
+    
+    # Memory usage
+    memory_bytes = df.memory_usage(deep=True).sum()
+    memory_mb = memory_bytes / (1024 * 1024)
+    
+    return CSVSchema(
+        file_path=str(file_path),
+        row_count=len(df),
+        column_count=len(df.columns),
+        columns=columns,
+        memory_usage_mb=memory_mb
+    )
+
+
+def estimate_token_reduction(original_size: int, schema_size: int) -> float:
+    """Estimate token reduction percentage.
+    
+    Args:
+        original_size: Size of original CSV in characters
+        schema_size: Size of extracted schema in characters
+        
+    Returns:
+        Reduction percentage (0-100)
+    """
+    if original_size == 0:
+        return 0.0
+    
+    reduction = ((original_size - schema_size) / original_size) * 100
+    return max(0.0, min(100.0, reduction))
+
+
+def analyze_csv_file(file_path: Path) -> Dict[str, Any]:
+    """Analyze a CSV file and return comprehensive information.
+    
+    Args:
+        file_path: Path to CSV file
+        
+    Returns:
+        Dictionary with schema and metadata
+    """
+    schema = extract_csv_schema(file_path)
+    schema_text = schema.to_text()
+    
+    # Calculate original size
+    original_size = file_path.stat().st_size
+    schema_size = len(schema_text)
+    
+    reduction = estimate_token_reduction(original_size, schema_size)
+    
+    return {
+        'schema': schema,
+        'schema_text': schema_text,
+        'original_size_bytes': original_size,
+        'schema_size_bytes': schema_size,
+        'token_reduction_pct': reduction,
+        'recommended_action': 'Use schema for LLM analysis instead of full CSV'
+    }

stratifyai/utils/file_analyzer.py

@@ -0,0 +1,192 @@
+"""File analysis utilities for detecting file types and estimating token usage."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+from enum import Enum
+
+from .token_counter import estimate_tokens, check_token_limit
+
+
+class FileType(Enum):
+    """Supported file types for intelligent processing."""
+    CSV = "csv"
+    JSON = "json"
+    LOG = "log"
+    PYTHON = "python"
+    JAVASCRIPT = "javascript"
+    JAVA = "java"
+    GO = "go"
+    TEXT = "text"
+    MARKDOWN = "markdown"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class FileAnalysis:
+    """Results of file analysis."""
+    file_path: Path
+    file_type: FileType
+    file_size_bytes: int
+    file_size_mb: float
+    content_length: int
+    estimated_tokens: int
+    exceeds_threshold: bool
+    context_window: int
+    percentage_used: float
+    recommendation: str
+
+
+def detect_file_type(file_path: Path) -> FileType:
+    """
+    Detect the type of file based on extension and content.
+    
+    Args:
+        file_path: Path to the file
+        
+    Returns:
+        FileType enum value
+    """
+    suffix = file_path.suffix.lower()
+    
+    # Map file extensions to types
+    type_mapping = {
+        ".csv": FileType.CSV,
+        ".json": FileType.JSON,
+        ".log": FileType.LOG,
+        ".py": FileType.PYTHON,
+        ".js": FileType.JAVASCRIPT,
+        ".ts": FileType.JAVASCRIPT,
+        ".jsx": FileType.JAVASCRIPT,
+        ".tsx": FileType.JAVASCRIPT,
+        ".java": FileType.JAVA,
+        ".go": FileType.GO,
+        ".txt": FileType.TEXT,
+        ".md": FileType.MARKDOWN,
+        ".markdown": FileType.MARKDOWN,
+    }
+    
+    return type_mapping.get(suffix, FileType.UNKNOWN)
+
+
+def get_recommendation(
+    file_type: FileType,
+    estimated_tokens: int,
+    context_window: int,
+    percentage_used: float
+) -> str:
+    """
+    Get processing recommendation based on file analysis.
+    
+    Args:
+        file_type: Detected file type
+        estimated_tokens: Estimated token count
+        context_window: Model's context window
+        percentage_used: Percentage of context window used
+        
+    Returns:
+        Recommendation string
+    """
+    # File fits comfortably in context
+    if percentage_used < 0.5:
+        return "✓ File fits well in model context - direct upload recommended"
+    
+    # File approaching context limit
+    elif percentage_used < 0.8:
+        return "⚠ File uses >50% of context - consider chunking for better performance"
+    
+    # File exceeds safe threshold
+    else:
+        if file_type == FileType.CSV:
+            return "⚠ Large CSV detected - use schema extraction (--extract-mode schema) for 99% token reduction"
+        elif file_type == FileType.JSON:
+            return "⚠ Large JSON detected - use schema extraction (--extract-mode schema) for 95% token reduction"
+        elif file_type == FileType.LOG:
+            return "⚠ Large log file detected - use error extraction (--extract errors) for 90% token reduction"
+        elif file_type in [FileType.PYTHON, FileType.JAVASCRIPT, FileType.JAVA, FileType.GO]:
+            return "⚠ Large code file detected - use code extraction (--extract summary) for 80% token reduction"
+        elif estimated_tokens > context_window:
+            return "✗ File exceeds model context - chunking required (--chunked)"
+        else:
+            return "⚠ File near context limit - chunking recommended (--chunked) for 90% token reduction"
+
+
+def analyze_file(
+    file_path: Path,
+    provider: str = "openai",
+    model: str = "gpt-4o",
+    threshold: float = 0.8
+) -> FileAnalysis:
+    """
+    Analyze a file and provide recommendations for processing.
+    
+    Args:
+        file_path: Path to the file to analyze
+        provider: LLM provider for token estimation
+        model: LLM model for context window limits
+        threshold: Warning threshold (default 0.8 = 80%)
+        
+    Returns:
+        FileAnalysis object with complete analysis
+        
+    Raises:
+        FileNotFoundError: If file doesn't exist
+        
+    Examples:
+        >>> analysis = analyze_file(Path("data.csv"), "openai", "gpt-4o")
+        >>> print(f"Tokens: {analysis.estimated_tokens}")
+        >>> print(analysis.recommendation)
+    """
+    if not file_path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+    
+    # Get file info
+    file_size_bytes = file_path.stat().st_size
+    file_size_mb = file_size_bytes / (1024 * 1024)
+    
+    # Detect file type
+    file_type = detect_file_type(file_path)
+    
+    # Read file content (with size limit for very large files)
+    MAX_READ_SIZE = 10 * 1024 * 1024  # 10MB max for token estimation
+    try:
+        if file_size_bytes > MAX_READ_SIZE:
+            # For very large files, estimate based on sample
+            with open(file_path, 'r', encoding='utf-8') as f:
+                sample = f.read(MAX_READ_SIZE)
+            # Extrapolate token count
+            sample_tokens = estimate_tokens(sample, provider, model)
+            estimated_tokens = int(sample_tokens * (file_size_bytes / len(sample)))
+            content_length = file_size_bytes  # Approximate
+        else:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+            content_length = len(content)
+            estimated_tokens = estimate_tokens(content, provider, model)
+    except UnicodeDecodeError:
+        # Binary file - rough estimate based on size
+        content_length = file_size_bytes
+        estimated_tokens = int(file_size_bytes / 4)  # Very rough estimate
+    
+    # Check token limits
+    exceeds_threshold, context_window, percentage_used = check_token_limit(
+        estimated_tokens, provider, model, threshold
+    )
+    
+    # Get recommendation
+    recommendation = get_recommendation(
+        file_type, estimated_tokens, context_window, percentage_used
+    )
+    
+    return FileAnalysis(
+        file_path=file_path,
+        file_type=file_type,
+        file_size_bytes=file_size_bytes,
+        file_size_mb=file_size_mb,
+        content_length=content_length,
+        estimated_tokens=estimated_tokens,
+        exceeds_threshold=exceeds_threshold,
+        context_window=context_window,
+        percentage_used=percentage_used,
+        recommendation=recommendation
+    )

stratifyai/utils/json_extractor.py

@@ -0,0 +1,219 @@
+"""JSON schema extraction for intelligent file analysis.
+
+This module extracts compact schema information from JSON files to reduce
+token usage by 95%+ while preserving essential structure information.
+"""
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Any, Union
+
+
+@dataclass
+class JSONSchema:
+    """Schema information for a JSON structure."""
+    type: str  # object, array, string, number, boolean, null
+    keys: Optional[List[str]] = None  # For objects
+    value_schema: Optional['JSONSchema'] = None  # For arrays
+    nested_schemas: Optional[Dict[str, 'JSONSchema']] = None  # For objects
+    sample_values: Optional[List[Any]] = None  # For primitives and arrays
+    depth: int = 0
+    
+    def to_text(self, indent: int = 0) -> str:
+        """Convert schema to human-readable text representation.
+        
+        Args:
+            indent: Current indentation level
+            
+        Returns:
+            Formatted schema description
+        """
+        prefix = "  " * indent
+        lines = []
+        
+        if self.type == "object":
+            lines.append(f"{prefix}Object with {len(self.keys or [])} keys:")
+            if self.nested_schemas:
+                for key, schema in self.nested_schemas.items():
+                    lines.append(f"{prefix}  {key}:")
+                    lines.append(schema.to_text(indent + 2))
+        
+        elif self.type == "array":
+            lines.append(f"{prefix}Array:")
+            if self.value_schema:
+                lines.append(f"{prefix}  Elements:")
+                lines.append(self.value_schema.to_text(indent + 2))
+            if self.sample_values:
+                sample_str = ", ".join(str(v)[:50] for v in self.sample_values[:3])
+                lines.append(f"{prefix}  Sample: [{sample_str}]")
+        
+        else:
+            # Primitive type
+            lines.append(f"{prefix}{self.type}")
+            if self.sample_values:
+                sample_str = ", ".join(str(v)[:50] for v in self.sample_values[:5])
+                lines.append(f"{prefix}  Samples: {sample_str}")
+        
+        return "\n".join(lines)
+
+
+def infer_json_schema(
+    data: Any,
+    max_depth: int = 10,
+    current_depth: int = 0,
+    sample_size: int = 3
+) -> JSONSchema:
+    """Infer schema from JSON data structure.
+    
+    Args:
+        data: JSON data (dict, list, or primitive)
+        max_depth: Maximum nesting depth to analyze
+        current_depth: Current depth in recursion
+        sample_size: Number of sample values to collect
+        
+    Returns:
+        JSONSchema object describing the structure
+    """
+    if current_depth >= max_depth:
+        return JSONSchema(type="...", depth=current_depth)
+    
+    if data is None:
+        return JSONSchema(type="null", depth=current_depth)
+    
+    elif isinstance(data, bool):
+        return JSONSchema(
+            type="boolean",
+            sample_values=[data],
+            depth=current_depth
+        )
+    
+    elif isinstance(data, (int, float)):
+        return JSONSchema(
+            type="number",
+            sample_values=[data],
+            depth=current_depth
+        )
+    
+    elif isinstance(data, str):
+        return JSONSchema(
+            type="string",
+            sample_values=[data[:100]],  # Truncate long strings
+            depth=current_depth
+        )
+    
+    elif isinstance(data, list):
+        # Analyze array elements
+        sample_values = data[:sample_size] if data else []
+        
+        # Infer schema from first element (assuming homogeneous array)
+        value_schema = None
+        if data:
+            value_schema = infer_json_schema(
+                data[0],
+                max_depth,
+                current_depth + 1,
+                sample_size
+            )
+        
+        return JSONSchema(
+            type="array",
+            value_schema=value_schema,
+            sample_values=sample_values,
+            depth=current_depth
+        )
+    
+    elif isinstance(data, dict):
+        # Analyze object keys and values
+        keys = list(data.keys())
+        nested_schemas = {}
+        
+        for key in keys:
+            nested_schemas[key] = infer_json_schema(
+                data[key],
+                max_depth,
+                current_depth + 1,
+                sample_size
+            )
+        
+        return JSONSchema(
+            type="object",
+            keys=keys,
+            nested_schemas=nested_schemas,
+            depth=current_depth
+        )
+    
+    else:
+        return JSONSchema(type="unknown", depth=current_depth)
+
+
+def extract_json_schema(file_path: Path) -> Dict[str, Any]:
+    """Extract schema information from a JSON file.
+    
+    Args:
+        file_path: Path to JSON file
+        
+    Returns:
+        Dictionary with schema and metadata
+        
+    Raises:
+        FileNotFoundError: If file doesn't exist
+        json.JSONDecodeError: If JSON is malformed
+    """
+    if not file_path.exists():
+        raise FileNotFoundError(f"JSON file not found: {file_path}")
+    
+    # Read and parse JSON
+    with open(file_path, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+    
+    # Infer schema
+    schema = infer_json_schema(data)
+    schema_text = schema.to_text()
+    
+    # Calculate sizes
+    original_size = file_path.stat().st_size
+    schema_size = len(schema_text)
+    
+    # Determine structure type
+    if isinstance(data, dict):
+        structure = f"Object with {len(data)} keys"
+    elif isinstance(data, list):
+        structure = f"Array with {len(data)} elements"
+    else:
+        structure = f"Primitive: {type(data).__name__}"
+    
+    return {
+        'schema': schema,
+        'schema_text': schema_text,
+        'structure': structure,
+        'original_size_bytes': original_size,
+        'schema_size_bytes': schema_size,
+        'token_reduction_pct': ((original_size - schema_size) / original_size * 100) if original_size > 0 else 0.0,
+        'recommended_action': 'Use schema for LLM analysis instead of full JSON'
+    }
+
+
+def analyze_json_file(file_path: Path) -> str:
+    """Analyze a JSON file and return schema description.
+    
+    Args:
+        file_path: Path to JSON file
+        
+    Returns:
+        Human-readable schema description
+    """
+    result = extract_json_schema(file_path)
+    
+    lines = [
+        f"JSON File: {file_path}",
+        f"Structure: {result['structure']}",
+        f"Original size: {result['original_size_bytes']:,} bytes",
+        f"Schema size: {result['schema_size_bytes']:,} bytes",
+        f"Token reduction: {result['token_reduction_pct']:.1f}%",
+        "",
+        "Schema:",
+        result['schema_text']
+    ]
+    
+    return "\n".join(lines)