ds-agent-cli 0.1.0

package/src/utils/schema_extraction.py

@@ -0,0 +1,132 @@
+"""
+Local Schema Extraction (No LLM)
+Fast, cheap extraction of dataset metadata without sending to LLM.
+"""
+
+import polars as pl
+from pathlib import Path
+from typing import Dict, Any, Optional
+
+
+def extract_schema_local(file_path: str, sample_rows: int = 5) -> Dict[str, Any]:
+    """
+    Extract dataset schema and basic stats locally without LLM.
+    
+    Returns:
+        - column names and types
+        - row/column counts
+        - missing value counts
+        - small sample for reference
+        - memory usage
+    """
+    try:
+        # Read with Polars (faster than pandas)
+        if file_path.endswith('.csv'):
+            # 🔥 FIX: Use infer_schema_length and ignore_errors to handle mixed-type columns
+            # This prevents failures like: could not parse `835.159865` as dtype `i64`
+            try:
+                df = pl.read_csv(file_path, infer_schema_length=10000, ignore_errors=True)
+            except Exception:
+                # Final fallback: read everything as strings, then let Polars infer
+                try:
+                    import pandas as pd
+                    pdf = pd.read_csv(file_path, low_memory=False)
+                    df = pl.from_pandas(pdf)
+                except Exception as e2:
+                    return {
+                        'error': f"Failed to read CSV: {str(e2)}",
+                        'file_path': file_path
+                    }
+        elif file_path.endswith('.parquet'):
+            df = pl.read_parquet(file_path)
+        else:
+            # Fallback to pandas
+            import pandas as pd
+            pdf = pd.read_csv(file_path, low_memory=False)
+            df = pl.from_pandas(pdf)
+        
+        # Basic metadata
+        schema_info = {
+            'file_path': file_path,
+            'file_size_mb': round(Path(file_path).stat().st_size / (1024 * 1024), 2),
+            'num_rows': df.shape[0],
+            'num_columns': df.shape[1],
+            'columns': {}
+        }
+        
+        # Per-column metadata
+        for col in df.columns:
+            col_series = df[col]
+            dtype_str = str(col_series.dtype)
+            
+            col_info = {
+                'dtype': dtype_str,
+                'missing_count': col_series.null_count(),
+                'missing_pct': round(col_series.null_count() / len(col_series) * 100, 2),
+                'unique_count': col_series.n_unique() if len(col_series) < 100000 else None  # Skip for huge datasets
+            }
+            
+            # Type-specific stats (lightweight)
+            if dtype_str in ['Int64', 'Float64', 'Int32', 'Float32']:
+                try:
+                    col_info['min'] = float(col_series.min())
+                    col_info['max'] = float(col_series.max())
+                    col_info['mean'] = float(col_series.mean())
+                except:
+                    pass
+            
+            schema_info['columns'][col] = col_info
+        
+        # Small sample for LLM context (only first few rows)
+        sample_data = df.head(sample_rows).to_dicts()
+        schema_info['sample_rows'] = sample_data
+        
+        # Categorize columns
+        schema_info['numeric_columns'] = [
+            col for col, info in schema_info['columns'].items()
+            if 'Int' in info['dtype'] or 'Float' in info['dtype']
+        ]
+        schema_info['categorical_columns'] = [
+            col for col, info in schema_info['columns'].items()
+            if info['dtype'] in ['Utf8', 'String'] or (
+                info.get('unique_count') is not None and 
+                info.get('unique_count') < 50 and 
+                col not in schema_info['numeric_columns']
+            )
+        ]
+        schema_info['datetime_columns'] = [
+            col for col, info in schema_info['columns'].items()
+            if 'Date' in info['dtype'] or 'Time' in info['dtype']
+        ]
+        
+        return schema_info
+        
+    except Exception as e:
+        return {
+            'error': f"Failed to extract schema: {str(e)}",
+            'file_path': file_path
+        }
+
+
+def infer_task_type(target_column: str, schema_info: Dict[str, Any]) -> Optional[str]:
+    """
+    Infer ML task type from target column without LLM.
+    """
+    if not target_column or target_column not in schema_info.get('columns', {}):
+        return None
+    
+    target_info = schema_info['columns'][target_column]
+    
+    # Numeric with many unique values → regression
+    if target_info['dtype'] in ['Int64', 'Float64', 'Int32', 'Float32']:
+        unique_count = target_info.get('unique_count')
+        if unique_count and unique_count > 20:
+            return 'regression'
+        elif unique_count and unique_count <= 10:
+            return 'classification'
+    
+    # Categorical or low cardinality → classification
+    if target_info['dtype'] in ['Utf8', 'String'] or target_info.get('unique_count', 0) <= 20:
+        return 'classification'
+    
+    return None

package/src/utils/semantic_layer.py

@@ -0,0 +1,392 @@
+"""
+Semantic Layer using SBERT for Column Understanding and Agent Routing
+
+Provides semantic understanding of dataset columns and agent intent matching
+using sentence-transformers embeddings.
+"""
+
+import numpy as np
+from typing import Dict, Any, List, Optional, Tuple
+import polars as pl
+from pathlib import Path
+import json
+
+# SBERT for semantic embeddings
+try:
+    from sentence_transformers import SentenceTransformer
+    import torch
+    SBERT_AVAILABLE = True
+except ImportError:
+    SBERT_AVAILABLE = False
+    print("⚠️ sentence-transformers not available. Install with: pip install sentence-transformers")
+
+# Sklearn for similarity
+try:
+    from sklearn.metrics.pairwise import cosine_similarity
+    SKLEARN_AVAILABLE = True
+except ImportError:
+    SKLEARN_AVAILABLE = False
+
+
+class SemanticLayer:
+    """
+    Semantic understanding layer using SBERT embeddings.
+    
+    Features:
+    - Column semantic embedding (name + sample values + dtype)
+    - Semantic column matching (find similar columns)
+    - Agent intent routing (semantic task → agent mapping)
+    - Target column inference (semantic similarity to "target")
+    """
+    
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        """
+        Initialize semantic layer with SBERT model.
+        
+        Args:
+            model_name: Sentence-transformer model name
+                - all-MiniLM-L6-v2: Fast, 384 dims (recommended)
+                - all-mpnet-base-v2: Better quality, 768 dims, slower
+                - paraphrase-MiniLM-L6-v2: Good for short texts
+        """
+        self.model_name = model_name
+        self.model = None
+        self.enabled = SBERT_AVAILABLE and SKLEARN_AVAILABLE
+        
+        if self.enabled:
+            try:
+                print(f"🧠 Loading SBERT model: {model_name}...")
+                # Try loading with trust_remote_code for better compatibility
+                self.model = SentenceTransformer(model_name, trust_remote_code=True)
+                # Use GPU if available
+                if torch.cuda.is_available():
+                    self.model = self.model.to('cuda')
+                    print("✅ SBERT loaded on GPU")
+                else:
+                    print("✅ SBERT loaded on CPU")
+            except Exception as e:
+                print(f"⚠️ Failed to load SBERT model: {e}")
+                print(f"   Falling back to keyword-based routing (semantic features disabled)")
+                self.enabled = False
+        else:
+            print("⚠️ SBERT semantic layer disabled (missing dependencies)")
+    
+    def encode_column(self, column_name: str, dtype: str, 
+                      sample_values: Optional[List[Any]] = None,
+                      stats: Optional[Dict[str, Any]] = None) -> np.ndarray:
+        """
+        Create semantic embedding for a column.
+        
+        Combines column name, data type, sample values, and stats into
+        a text description that captures the column's semantic meaning.
+        
+        Args:
+            column_name: Name of the column
+            dtype: Data type (Int64, Float64, Utf8, etc.)
+            sample_values: Sample values from the column
+            stats: Optional statistics (mean, min, max, etc.)
+        
+        Returns:
+            Embedding vector (numpy array)
+        
+        Example:
+            >>> encode_column("annual_salary", "Float64", [50000, 75000], {"mean": 65000})
+            >>> # Returns embedding for "annual_salary (Float64 numeric): values like 50000, 75000, mean 65000"
+        """
+        if not self.enabled:
+            return np.zeros(384)  # Dummy embedding
+        
+        # Build semantic description
+        description_parts = [f"Column name: {column_name}"]
+        
+        # Add type information
+        type_desc = self._interpret_dtype(dtype)
+        description_parts.append(f"Type: {type_desc}")
+        
+        # Add sample values
+        if sample_values:
+            # Format samples nicely
+            samples_str = ", ".join([str(v)[:50] for v in sample_values[:5] if v is not None])
+            description_parts.append(f"Example values: {samples_str}")
+        
+        # Add statistics
+        if stats:
+            if 'mean' in stats and stats['mean'] is not None:
+                description_parts.append(f"Mean: {stats['mean']:.2f}")
+            if 'unique_count' in stats and stats['unique_count'] is not None:
+                description_parts.append(f"Unique values: {stats['unique_count']}")
+            if 'null_percentage' in stats and stats['null_percentage'] is not None:
+                description_parts.append(f"Missing: {stats['null_percentage']:.1f}%")
+        
+        # Combine into single text
+        text = ". ".join(description_parts)
+        
+        # Generate embedding
+        try:
+            embedding = self.model.encode(text, convert_to_numpy=True, show_progress_bar=False)
+            return embedding
+        except Exception as e:
+            print(f"⚠️ Error encoding column {column_name}: {e}")
+            return np.zeros(self.model.get_sentence_embedding_dimension())
+    
+    def _interpret_dtype(self, dtype: str) -> str:
+        """Convert polars dtype to human-readable description."""
+        dtype_lower = str(dtype).lower()
+        
+        if 'int' in dtype_lower or 'float' in dtype_lower:
+            return "numeric continuous or count data"
+        elif 'bool' in dtype_lower:
+            return "boolean flag"
+        elif 'utf8' in dtype_lower or 'str' in dtype_lower:
+            return "text or categorical label"
+        elif 'date' in dtype_lower or 'time' in dtype_lower:
+            return "temporal timestamp"
+        else:
+            return "data values"
+    
+    def find_similar_columns(self, query_column: str, column_embeddings: Dict[str, np.ndarray],
+                            top_k: int = 3, threshold: float = 0.6) -> List[Tuple[str, float]]:
+        """
+        Find columns semantically similar to query column.
+        
+        Use case: Detect duplicates or related columns
+        Example: "Salary" → finds ["Annual_Income", "Compensation", "Pay"]
+        
+        Args:
+            query_column: Column name to search for
+            column_embeddings: Dict mapping column names to their embeddings
+            top_k: Number of similar columns to return
+            threshold: Minimum similarity score (0-1)
+        
+        Returns:
+            List of (column_name, similarity_score) tuples
+        """
+        if not self.enabled or query_column not in column_embeddings:
+            return []
+        
+        query_emb = column_embeddings[query_column].reshape(1, -1)
+        
+        similarities = []
+        for col_name, col_emb in column_embeddings.items():
+            if col_name == query_column:
+                continue
+            
+            sim = cosine_similarity(query_emb, col_emb.reshape(1, -1))[0][0]
+            if sim >= threshold:
+                similarities.append((col_name, float(sim)))
+        
+        # Sort by similarity descending
+        similarities.sort(key=lambda x: x[1], reverse=True)
+        return similarities[:top_k]
+    
+    def infer_target_column(self, column_embeddings: Dict[str, np.ndarray],
+                           task_description: str) -> Optional[Tuple[str, float]]:
+        """
+        Infer which column is likely the target/label for prediction.
+        
+        Uses semantic similarity between column descriptions and task description.
+        
+        Args:
+            column_embeddings: Dict mapping column names to embeddings
+            task_description: User's task description
+        
+        Returns:
+            (column_name, confidence_score) or None
+        
+        Example:
+            >>> infer_target_column(embeddings, "predict house prices")
+            >>> ("Price", 0.85)  # High confidence "Price" is target
+        """
+        if not self.enabled:
+            return None
+        
+        # Encode task description
+        task_emb = self.model.encode(task_description, convert_to_numpy=True, show_progress_bar=False)
+        task_emb = task_emb.reshape(1, -1)
+        
+        # Find column with highest similarity to task
+        best_col = None
+        best_score = 0.0
+        
+        for col_name, col_emb in column_embeddings.items():
+            sim = cosine_similarity(task_emb, col_emb.reshape(1, -1))[0][0]
+            if sim > best_score:
+                best_score = sim
+                best_col = col_name
+        
+        # Only return if confidence is reasonable
+        if best_score >= 0.4:  # Threshold for target inference
+            return (best_col, float(best_score))
+        
+        return None
+    
+    def route_to_agent(self, task_description: str, 
+                       agent_descriptions: Dict[str, str]) -> Tuple[str, float]:
+        """
+        Route task to appropriate specialist agent using semantic similarity.
+        
+        Replaces keyword-based routing with semantic understanding.
+        
+        Args:
+            task_description: User's task description
+            agent_descriptions: Dict mapping agent_key → agent description
+        
+        Returns:
+            (agent_key, confidence_score)
+        
+        Example:
+            >>> route_to_agent("build a predictive model", {
+            ...     "modeling_agent": "Expert in ML training and models",
+            ...     "viz_agent": "Expert in visualizations"
+            ... })
+            >>> ("modeling_agent", 0.92)
+        """
+        if not self.enabled:
+            # Fallback to first agent
+            return list(agent_descriptions.keys())[0], 0.5
+        
+        # Encode task
+        task_emb = self.model.encode(task_description, convert_to_numpy=True, show_progress_bar=False)
+        task_emb = task_emb.reshape(1, -1)
+        
+        # Encode agent descriptions
+        best_agent = None
+        best_score = 0.0
+        
+        for agent_key, agent_desc in agent_descriptions.items():
+            agent_emb = self.model.encode(agent_desc, convert_to_numpy=True, show_progress_bar=False)
+            agent_emb = agent_emb.reshape(1, -1)
+            
+            sim = cosine_similarity(task_emb, agent_emb)[0][0]
+            if sim > best_score:
+                best_score = sim
+                best_agent = agent_key
+        
+        return best_agent, float(best_score)
+    
+    def semantic_column_match(self, target_name: str, available_columns: List[str],
+                             threshold: float = 0.6) -> Optional[Tuple[str, float]]:
+        """
+        Find best matching column for a target name using fuzzy semantic matching.
+        
+        Better than string fuzzy matching because it understands synonyms:
+        - "salary" matches "annual_income", "compensation", "pay"
+        - "target" matches "label", "class", "outcome"
+        
+        Args:
+            target_name: Column name to find (might not exist exactly)
+            available_columns: List of actual column names in dataset
+            threshold: Minimum similarity to consider a match
+        
+        Returns:
+            (matched_column, confidence) or None
+        
+        Example:
+            >>> semantic_column_match("salary", ["Annual_Income", "Name", "Age"])
+            >>> ("Annual_Income", 0.78)
+        """
+        if not self.enabled:
+            # Fallback to exact match
+            if target_name in available_columns:
+                return (target_name, 1.0)
+            return None
+        
+        # Encode target
+        target_emb = self.model.encode(target_name, convert_to_numpy=True, show_progress_bar=False)
+        target_emb = target_emb.reshape(1, -1)
+        
+        # Find best match
+        best_col = None
+        best_score = 0.0
+        
+        for col in available_columns:
+            col_emb = self.model.encode(col, convert_to_numpy=True, show_progress_bar=False)
+            col_emb = col_emb.reshape(1, -1)
+            
+            sim = cosine_similarity(target_emb, col_emb)[0][0]
+            if sim > best_score:
+                best_score = sim
+                best_col = col
+        
+        if best_score >= threshold:
+            return (best_col, float(best_score))
+        
+        return None
+    
+    def enrich_dataset_info(self, dataset_info: Dict[str, Any], 
+                           file_path: str, sample_size: int = 100) -> Dict[str, Any]:
+        """
+        Enrich dataset_info with semantic column embeddings.
+        
+        Adds 'column_embeddings' and 'semantic_insights' to dataset_info.
+        
+        Args:
+            dataset_info: Dataset info from schema_extraction
+            file_path: Path to CSV file
+            sample_size: Number of rows to sample for encoding
+        
+        Returns:
+            Enhanced dataset_info with semantic layer
+        """
+        if not self.enabled:
+            return dataset_info
+        
+        try:
+            # Load dataset
+            df = pl.read_csv(file_path, n_rows=sample_size)
+            
+            column_embeddings = {}
+            
+            for col_name, col_info in dataset_info['columns'].items():
+                # Get sample values
+                sample_values = df[col_name].head(5).to_list()
+                
+                # Create embedding
+                embedding = self.encode_column(
+                    column_name=col_name,
+                    dtype=col_info['dtype'],
+                    sample_values=sample_values,
+                    stats={
+                        'unique_count': col_info.get('unique_count'),
+                        'missing_pct': col_info.get('missing_pct'),
+                        'mean': col_info.get('mean')
+                    }
+                )
+                
+                column_embeddings[col_name] = embedding
+            
+            # Add to dataset_info
+            dataset_info['column_embeddings'] = column_embeddings
+            
+            # Detect similar columns (potential duplicates)
+            similar_pairs = []
+            cols = list(column_embeddings.keys())
+            for i, col1 in enumerate(cols):
+                similar = self.find_similar_columns(col1, column_embeddings, top_k=1, threshold=0.75)
+                if similar:
+                    similar_pairs.append((col1, similar[0][0], similar[0][1]))
+            
+            dataset_info['semantic_insights'] = {
+                'similar_columns': similar_pairs,
+                'total_columns_embedded': len(column_embeddings)
+            }
+            
+            print(f"🧠 Semantic layer: Embedded {len(column_embeddings)} columns")
+            if similar_pairs:
+                print(f"   Found {len(similar_pairs)} similar column pairs (potential duplicates)")
+            
+        except Exception as e:
+            print(f"⚠️ Error enriching dataset with semantic layer: {e}")
+        
+        return dataset_info
+
+
+# Global semantic layer instance (lazy loaded)
+_semantic_layer = None
+
+def get_semantic_layer() -> SemanticLayer:
+    """Get or create global semantic layer instance."""
+    global _semantic_layer
+    if _semantic_layer is None:
+        _semantic_layer = SemanticLayer()
+    return _semantic_layer