"""Tool definitions and handler functions for the TLS Analyzer MCP server. Each tool is defined by a :class:`Tool ` metadata object and an async handler function. Handlers interact with the :class:`SessionStore` singleton and the Django ORM. All 11 tools: 1. load_data 7. extract_features 2. profile_data 8. export_results 3. filter_data 9. list_datasets 4. preprocess_data 10. drop_dataset 5. run_clustering 11. clone_dataset 6. evaluate_clustering """ import json import os from pathlib import Path from typing import Any, Optional import numpy as np import polars as pl from mcp.types import Tool from .data_loader import load_csv_directory from .data_profiler import profile_dataset as _profile_dataset from .entity_detector import detect_entity_column from .entity_aggregator import aggregate_by_entity from .session_store import SessionStore, _generate_id # ═══════════════════════════════════════════════════════════════════════ # Tool metadata # ═══════════════════════════════════════════════════════════════════════ def _json_size(obj: Any) -> int: return len(json.dumps(obj, default=str, ensure_ascii=False)) _MAX_RESPONSE_KB = 64 _MAX_RESPONSE_BYTES = _MAX_RESPONSE_KB * 1024 def _truncate_response(result: dict, max_kb: int = _MAX_RESPONSE_KB) -> dict: """Truncate large lists/keys in *result* to stay under *max_kb*.""" if _json_size(result) <= max_kb * 1024: result.setdefault('truncated', False) return result # Drop verbose keys first for costly_key in ('column_stats', 'features', 'correlation_matrix', 'labels', 'cluster_labels'): if costly_key in result and isinstance(result[costly_key], dict): result[costly_key] = {'__truncated__': True, 'key_count': len(result[costly_key])} elif costly_key in result and isinstance(result[costly_key], list): lst = result[costly_key] if len(lst) > 10: result[costly_key] = lst[:10] + [{'__truncated__': True, 'omitted': len(lst) - 10}] result['truncated'] = True return result def get_tools_meta() -> list[Tool]: """Return the list of Tool metadata objects for MCP server registration. Call this from :meth:`TLSAnalyzerMCPServer._setup_tools`. """ return [ Tool( name="load_data", description=( "Load CSV files matching a glob pattern. Detects BOM per file, " "validates schema consistency, merges vertically, and stores the " "result as a lazy dataset in the session store." ), inputSchema={ "type": "object", "properties": { "csv_glob": { "type": "string", "description": "Glob pattern for CSV files (e.g. 'data/*.csv')", }, "config_path": { "type": "string", "description": "Optional path to a YAML config file", }, "encoding": { "type": "string", "description": "Fallback encoding when no BOM is found", "default": "utf-8", }, "delimiter": { "type": "string", "description": "Column delimiter", "default": ",", }, "schema_strict": { "type": "boolean", "description": "When True, raise on column mismatch; when False, merge leniently with nulls", "default": False, }, "recursive": { "type": "boolean", "description": "When True, use recursive glob for **/*.csv patterns", "default": False, }, }, "required": ["csv_glob"], }, ), Tool( name="profile_data", description=( "Compute per-column statistics and a numeric correlation matrix " "for a dataset. Result is automatically truncated to ~16 KB for " "LLM context-window efficiency." ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Dataset identifier from a previous load_data call", }, "sample_size": { "type": "integer", "description": "Maximum rows to profile (sampled if larger)", "default": 1000, }, "columns": { "type": "array", "items": {"type": "string"}, "description": "Columns to profile (default: all)", }, }, "required": ["dataset_id"], }, ), Tool( name="filter_data", description=( "Apply column-level filters to a dataset and store the result as " "a new derived dataset. Supports eq/neq/gt/gte/lt/lte/contains/" "not_contains/in/not_in/is_null/is_not_null operators." ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Source dataset identifier", }, "filters": { "type": "array", "items": { "type": "object", "properties": { "column": {"type": "string"}, "op": { "type": "string", "enum": [ "eq", "neq", "gt", "gte", "lt", "lte", "contains", "not_contains", "in", "not_in", "is_null", "is_not_null", ], }, "value": {"description": "Operand value"}, }, "required": ["column", "op"], }, "description": "List of filter expressions", }, "logic": { "type": "string", "enum": ["and", "or"], "description": "Combine filters with AND or OR logic", "default": "and", }, }, "required": ["dataset_id", "filters"], }, ), Tool( name="preprocess_data", description=( "Preprocess selected columns: scaling (standardize), one-hot encoding, " "missing-value imputation, and column dropping." ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Source dataset identifier", }, "columns": { "type": "array", "items": {"type": "string"}, "description": "Columns to preprocess", }, "config": { "type": "object", "properties": { "standardize": { "type": "boolean", "description": "Apply StandardScaler to numeric columns", "default": False, }, "onehot": { "type": "boolean", "description": "One-hot encode categorical columns", "default": False, }, "fillna": { "type": "string", "enum": ["mean", "median", "zero", "drop"], "description": "Missing value strategy", }, "drop": { "type": "array", "items": {"type": "string"}, "description": "Columns to drop entirely", }, }, }, }, "required": ["dataset_id", "columns"], }, ), Tool( name="run_clustering", description=( "Run clustering on selected columns. Uses HDBSCAN (default) or " "KMeans. Features are standard-scaled before clustering. Returns " "cluster labels and quality metrics." ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Dataset identifier", }, "cluster_columns": { "type": "array", "items": {"type": "string"}, "description": "Columns to use as clustering features", }, "algorithm": { "type": "string", "enum": ["hdbscan", "kmeans"], "description": "Clustering algorithm", "default": "hdbscan", }, "params": { "type": "object", "description": ( "Algorithm parameters. HDBSCAN: min_cluster_size (5), " "min_samples (None), metric ('euclidean'). " "KMeans: n_clusters (3), random_state (42)." ), }, "random_state": { "type": "integer", "description": "Random seed for reproducibility", "default": 42, }, }, "required": ["dataset_id", "cluster_columns"], }, ), Tool( name="evaluate_clustering", description=( "Evaluate a clustering result on a dataset. Currently supports " "silhouette, davies_bouldin, calinski_harabasz, noise_ratio, and " "cluster_sizes metrics." ), inputSchema={ "type": "object", "properties": { "cluster_result_id": { "type": "string", "description": "Result ID from run_clustering", }, "dataset_id": { "type": "string", "description": "Original dataset for metric computation", }, "metrics": { "type": "array", "items": {"type": "string"}, "description": ( "Metrics to compute. Options: silhouette, " "davies_bouldin, calinski_harabasz, noise_ratio, " "cluster_sizes (default: all)" ), }, }, "required": ["cluster_result_id", "dataset_id"], }, ), Tool( name="extract_features", description=( "Extract distinguishing features per cluster and auto-save to " "the Django ORM (ClusterFeature model). Uses z-score or ANOVA " "to rank feature importance per cluster." ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Dataset containing the feature columns", }, "cluster_result_id": { "type": "string", "description": "Result ID from run_clustering", }, "top_k": { "type": "integer", "description": "Top distinguishing features per cluster", "default": 10, }, "method": { "type": "string", "enum": ["zscore", "anova"], "description": "Method to rank feature importance", "default": "zscore", }, "save_to_db": { "type": "boolean", "description": "Save features to Django ORM if True", "default": True, }, }, "required": ["dataset_id", "cluster_result_id"], }, ), Tool( name="export_results", description=( "Export a dataset or cluster result to disk as CSV or Parquet. " "Datasets are materialised and written; cluster results are " "serialised as JSON." ), inputSchema={ "type": "object", "properties": { "result_id": { "type": "string", "description": "Dataset ID or cluster result ID to export", }, "output_path": { "type": "string", "description": "Output file path or directory", }, "format": { "type": "string", "enum": ["csv", "parquet", "json"], "description": "Output format", "default": "csv", }, "overwrite": { "type": "boolean", "description": "Overwrite existing file", "default": False, }, }, "required": ["result_id", "output_path"], }, ), Tool( name="list_datasets", description="List all active datasets and results in the session store.", inputSchema={ "type": "object", "properties": {}, }, ), Tool( name="drop_dataset", description=( "Remove a dataset (or result) from the session store and release " "memory via garbage collection." ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Dataset or result identifier to remove", }, }, "required": ["dataset_id"], }, ), Tool( name="clone_dataset", description=( "Shallow-clone a dataset. Since LazyFrames are query plans, " "no data is copied. Useful for branching analyses." ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Source dataset identifier", }, }, "required": ["dataset_id"], }, ), Tool( name="build_entity_profiles", description=( "Auto-detect or manually specify entity column(s) (e.g. src_ip, " "sni, user) and aggregate flow-level data into per-entity profiles. " "Computes flow stats, destination diversity, TLS features, protocol " "features, and time patterns. Results are persisted to the Django " "ORM (EntityProfile model) and stored as a new dataset in the session." ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Source dataset identifier from load_data", }, "entity_column": { "type": "string", "description": "Column to group by (optional; auto-detected if omitted)", }, "entity_columns": { "type": "array", "items": {"type": "string"}, "description": "Multiple columns to group by (optional; overrides entity_column)", }, "auto_detect": { "type": "boolean", "description": "Auto-detect entity column when entity_column is not given", "default": True, }, }, "required": ["dataset_id"], }, ), ] # ═══════════════════════════════════════════════════════════════════════ # Handler implementations # ═══════════════════════════════════════════════════════════════════════ async def handle_call(name: str, arguments: dict) -> dict: """Dispatch a tool call to the appropriate handler. Args: name: Tool name (must match one of the 11 tools). arguments: Tool-specific parameters. Returns: A JSON-serialisable dict. Raises: ValueError: If *name* is not a recognised tool. """ _handlers = { 'load_data': _handle_load_data, 'profile_data': _handle_profile_data, 'filter_data': _handle_filter_data, 'preprocess_data': _handle_preprocess_data, 'run_clustering': _handle_run_clustering, 'evaluate_clustering': _handle_evaluate_clustering, 'extract_features': _handle_extract_features, 'export_results': _handle_export_results, 'list_datasets': _handle_list_datasets, 'drop_dataset': _handle_drop_dataset, 'clone_dataset': _handle_clone_dataset, 'build_entity_profiles': _handle_build_entity_profiles, } handler = _handlers.get(name) if handler is None: raise ValueError(f"Unknown tool: '{name}'") result = await handler(**arguments) # Always ensure truncated key exists if 'truncated' not in result: result['truncated'] = False return _truncate_response(result) # ── 1. load_data ──────────────────────────────────────────────────────── async def _handle_load_data( csv_glob: str, config_path: Optional[str] = None, encoding: str = 'utf-8', delimiter: str = ',', schema_strict: bool = False, recursive: bool = False, ) -> dict: """Load, validate, merge CSVs, and store in session.""" try: lf, schema, row_count, file_count, memory_mb = load_csv_directory( glob_pattern=csv_glob, encoding=encoding, delimiter=delimiter, config_path=config_path, sample_rows=10000, schema_strict=schema_strict, recursive=recursive, ) except FileNotFoundError as e: return {'error': f'File not found: {e}', 'truncated': False} except ValueError as e: return {'error': f'Schema validation failed: {e}', 'truncated': False} store = SessionStore() dataset_id = _generate_id('ds') store.store_dataset( dataset_id=dataset_id, lazyframe=lf, schema=schema, metadata={ 'row_count': row_count, 'file_count': file_count, 'memory_mb': round(memory_mb, 2), 'csv_glob': csv_glob, 'encoding': encoding, }, ) return { 'dataset_id': dataset_id, 'schema': schema, 'row_count': row_count, 'file_count': file_count, 'memory_mb': round(memory_mb, 2), } # ── 2. profile_data ───────────────────────────────────────────────────── async def _handle_profile_data( dataset_id: str, sample_size: int = 1000, columns: Optional[list[str]] = None, ) -> dict: """Profile dataset columns and correlations.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf: pl.LazyFrame = entry['lazyframe'] try: profile = _profile_dataset(lf, sample_size=sample_size, columns=columns) except Exception as e: return {'error': f'Profiling failed: {e}', 'truncated': False} return profile # ── 3. filter_data ────────────────────────────────────────────────────── def _build_filter_expr(column: str, op: str, value: Any) -> pl.Expr: """Build a single Polars expression from a filter descriptor.""" col = pl.col(column) op_map = { 'eq': col == value, 'neq': col != value, 'gt': col > value, 'gte': col >= value, 'lt': col < value, 'lte': col <= value, 'contains': col.str.contains(value) if value is not None else col.cast(str).str.contains(''), 'not_contains': ~(col.str.contains(value) if value is not None else col.cast(str).str.contains('')), 'in': col.is_in(value) if isinstance(value, (list, tuple)) else col.is_in([value]), 'not_in': ~(col.is_in(value) if isinstance(value, (list, tuple)) else col.is_in([value])), 'is_null': col.is_null(), 'is_not_null': col.is_not_null(), } expr = op_map.get(op) if expr is None: raise ValueError(f"Unsupported filter operator: {op}") return expr async def _handle_filter_data( dataset_id: str, filters: list[dict], logic: str = 'and', ) -> dict: """Apply filters and store filtered dataset.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf: pl.LazyFrame = entry['lazyframe'] full_schema = entry.get('schema', {}) try: exprs = [_build_filter_expr( f['column'], f.get('op', 'eq'), f.get('value'), ) for f in filters] if not exprs: return {'error': 'No filter expressions could be built', 'truncated': False} if logic == 'and': combined = exprs[0] for e in exprs[1:]: combined = combined & e else: combined = exprs[0] for e in exprs[1:]: combined = combined | e filtered_lf = lf.filter(combined) # Estimate row count (fast streaming count) try: row_count = filtered_lf.select(pl.lit(1).count()).collect(streaming=True).item() except Exception: row_count = None except Exception as e: return {'error': f'Filter failed: {e}', 'truncated': False} new_id = _generate_id('flt') store.store_dataset( dataset_id=new_id, lazyframe=filtered_lf, schema=full_schema, metadata={ 'row_count': row_count, 'parent_dataset_id': dataset_id, 'filter_count': len(filters), 'logic': logic, }, parent_id=dataset_id, ) return { 'filtered_dataset_id': new_id, 'row_count': row_count, 'column_deltas': {'kept': list(full_schema.keys()), 'dropped': []}, } # ── 4. preprocess_data ────────────────────────────────────────────────── async def _handle_preprocess_data( dataset_id: str, columns: list[str], config: Optional[dict] = None, ) -> dict: """Preprocess selected columns (scale, encode, impute, drop).""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf: pl.LazyFrame = entry['lazyframe'] cfg = config or {} column_mapping: dict[str, str] = {} try: df = lf.collect(streaming=True) except Exception as e: return {'error': f'Data collection failed: {e}', 'truncated': False} # Resolve columns that actually exist available = set(df.columns) target_cols = [c for c in columns if c in available] if not target_cols: return {'error': 'None of the specified columns exist in the dataset', 'truncated': False} # Drop columns drop_cols = cfg.get('drop', []) if drop_cols: df = df.drop([c for c in drop_cols if c in df.columns]) for c in drop_cols: if c in available: column_mapping[c] = f'' # Fill missing values fill_strategy = cfg.get('fillna') if fill_strategy and fill_strategy != 'drop': for col in target_cols: if col not in df.columns: continue if df[col].null_count() == 0: continue if fill_strategy == 'mean': mean_val = df[col].drop_nulls().mean() if mean_val is not None: df = df.with_columns(pl.col(col).fill_null(mean_val)) elif fill_strategy == 'median': med_val = df[col].drop_nulls().median() if med_val is not None: df = df.with_columns(pl.col(col).fill_null(med_val)) elif fill_strategy == 'zero': df = df.with_columns(pl.col(col).fill_null(0)) elif fill_strategy == 'drop': df = df.drop_nulls(subset=target_cols) # Standard scaling (z-score) if cfg.get('standardize'): from sklearn.preprocessing import StandardScaler num_cols = [c for c in target_cols if c in df.columns and df[c].dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64)] if num_cols: scaler = StandardScaler() scaled = scaler.fit_transform(df.select(num_cols).to_numpy()) for i, col in enumerate(num_cols): col_scaled = f'{col}_scaled' df = df.with_columns(pl.Series(col_scaled, scaled[:, i])) column_mapping[col] = col_scaled # One-hot encoding if cfg.get('onehot'): str_cols = [c for c in target_cols if c in df.columns and df[c].dtype in (pl.Utf8, pl.String, pl.Categorical)] if str_cols: df = df.to_dummies(columns=str_cols, drop_first=False) for c in str_cols: column_mapping[c] = f'' new_id = _generate_id('pp') new_lf = df.lazy() new_schema = {name: str(dtype) for name, dtype in zip(new_lf.collect_schema().names(), new_lf.collect_schema().dtypes())} store.store_dataset( dataset_id=new_id, lazyframe=new_lf, schema=new_schema, metadata={ 'row_count': len(df), 'parent_dataset_id': dataset_id, 'preprocessing_config': cfg, }, parent_id=dataset_id, ) return { 'preprocessed_dataset_id': new_id, 'column_mapping': column_mapping, 'row_count': len(df), 'column_count': len(new_schema), } # ── 5. run_clustering ─────────────────────────────────────────────────── def _cluster_sync( dataset_id: str, cluster_columns: list[str], algorithm: str = 'hdbscan', params: Optional[dict] = None, random_state: int = 42, ) -> dict: """Run HDBSCAN or KMeans clustering on selected columns (synchronous core).""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf: pl.LazyFrame = entry['lazyframe'] p = params or {} try: # Collect cluster features - numeric only available = lf.collect_schema() available_names = available.names() available_dtypes = available.dtypes() numeric_dtypes = (pl.Int8, pl.Int16, pl.Int32, pl.Int64, pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64, pl.Float32, pl.Float64) feature_cols = [ c for c in cluster_columns if c in available_names and available_dtypes[available_names.index(c)] in numeric_dtypes ] if not feature_cols: # Fall back: auto-select all numeric columns (exclude internal ones) feature_cols = [ available_names[i] for i, dt in enumerate(available_dtypes) if dt in numeric_dtypes and not available_names[i].startswith('_') ][:10] if not feature_cols: all_numeric = [f"{available_names[i]}" for i, dt in enumerate(available_dtypes) if dt in numeric_dtypes] return { 'error': f'No numeric cluster columns found in dataset. Available: {all_numeric}', 'truncated': False } df_features = lf.select(feature_cols).collect(streaming=True) # Drop all-NaN columns (keep rows, drop useless columns) for col in df_features.columns: if df_features[col].is_null().all(): df_features = df_features.drop(col) feature_cols = [c for c in feature_cols if c != col] if df_features.width == 0: return {'error': 'No valid feature columns after dropping all-NaN columns', 'truncated': False} data = df_features.to_numpy() # Fill remaining NaN with column mean if data.dtype.kind == 'f': col_mean = np.nanmean(data, axis=0) col_mean = np.nan_to_num(col_mean, nan=0.0) data = np.where(np.isnan(data), col_mean, data) # Zero-sample / too-few-samples guard if len(data) < 3: return { 'cluster_result_id': _generate_id('clust'), 'n_clusters': 0, 'n_noise': 0, 'quality_metrics': {'note': 'too few samples'}, 'truncated': False, } # ── D14: Low memory check (safe fallback if psutil not installed) ── low_memory = False try: import psutil mem = psutil.virtual_memory() low_memory = mem.available < 2 * 1024**3 # 2 GB except ImportError: pass # ── D14: Auto-downsample for low memory (>100K rows → 50K) ──────── if low_memory and len(data) > 100000: rng = np.random.default_rng(42) idx = rng.choice(len(data), 50000, replace=False) data = data[idx] # ── D5: General downsampling for big data (>50K rows → 50K) ─────── if len(data) > 50000: rng = np.random.default_rng(42) idx = rng.choice(len(data), 50000, replace=False) data = data[idx] # ── H17: Variance filtering ─────────────────────────────────────── # Remove features with zero / near-zero variance (< 1e-10) if data.shape[1] > 1: variances = np.var(data, axis=0) keep_var = np.where(variances >= 1e-10)[0] if len(keep_var) == 0: return { 'cluster_result_id': _generate_id('clust'), 'n_clusters': 0, 'n_noise': 0, 'quality_metrics': {'note': 'all features have zero variance'}, 'truncated': False, } if len(keep_var) < data.shape[1]: data = data[:, keep_var] feature_cols = [feature_cols[i] for i in keep_var] # Log clustering input import logging _clust_logger = logging.getLogger(__name__) _clust_logger.info(f'[CLUSTER_INPUT] rows={len(data)} cols={data.shape[1] if len(data.shape) > 1 else 1} feature_cols={feature_cols}') # Standard scale from sklearn.preprocessing import StandardScaler data_scaled = StandardScaler().fit_transform(data) # ── H17: Correlation filtering ──────────────────────────────────── # Remove features with Pearson correlation > 0.95 (keep the first one) if data_scaled.shape[1] > 1: corr = np.corrcoef(data_scaled.T) upper = np.triu(np.abs(corr), 1) to_drop = [j for j in range(upper.shape[1]) if any(upper[:, j] > 0.95)] if to_drop: keep = [j for j in range(data_scaled.shape[1]) if j not in to_drop] data_scaled = data_scaled[:, keep] feature_cols = [feature_cols[j] for j in keep] # ── D14: PCA dimensionality reduction (if >50 features) ─────────── n_features = data_scaled.shape[1] if n_features > 50: from sklearn.decomposition import PCA pca = PCA(n_components=50, random_state=random_state) data_scaled = pca.fit_transform(data_scaled) # ── H17: HDBSCAN hyperparameter auto-tuning based on data size ──── n = len(data_scaled) if algorithm == 'hdbscan': min_cluster = p.get('min_cluster_size', max(3, min(50, n // 20))) p['min_cluster_size'] = min_cluster p['min_samples'] = p.get('min_samples', min_cluster) p['cluster_selection_epsilon'] = p.get('cluster_selection_epsilon', 0.0) else: # KMeans safety cap p['min_cluster_size'] = min( p.get('min_cluster_size', 5), max(n // 2, 2), ) # Clustering from sklearn.cluster import HDBSCAN, MiniBatchKMeans if algorithm == 'kmeans': n_clusters = p.get('n_clusters', 3) model = MiniBatchKMeans( n_clusters=n_clusters, random_state=random_state, batch_size=1024, n_init='auto', ) else: model = HDBSCAN( min_cluster_size=p['min_cluster_size'], min_samples=p.get('min_samples', None), metric=p.get('metric', 'euclidean'), cluster_selection_epsilon=p.get('cluster_selection_epsilon', 0.0), ) labels_array = model.fit_predict(data_scaled) n_clusters_found = int(len(set(labels_array)) - (1 if -1 in labels_array else 0)) noise_count = int((labels_array == -1).sum()) if algorithm == 'hdbscan' else 0 noise_ratio = round(noise_count / len(labels_array), 4) if algorithm == 'hdbscan' else 0.0 # Map back to original rows full_labels = labels_array # Quality metrics quality = { 'n_clusters': n_clusters_found, 'n_noise': noise_count, 'noise_ratio': noise_ratio, 'algorithm': algorithm, } if n_clusters_found > 1: try: from sklearn.metrics import silhouette_score score = silhouette_score(data_scaled, labels_array, random_state=random_state) quality['silhouette_score'] = round(float(score), 4) except Exception: quality['silhouette_score'] = None try: from sklearn.metrics import davies_bouldin_score dbs = davies_bouldin_score(data_scaled, labels_array) quality['davies_bouldin_score'] = round(float(dbs), 4) except Exception: quality['davies_bouldin_score'] = None try: from sklearn.metrics import calinski_harabasz_score chs = calinski_harabasz_score(data_scaled, labels_array) quality['calinski_harabasz_score'] = round(float(chs), 4) except Exception: quality['calinski_harabasz_score'] = None # Cluster sizes cluster_sizes = {} for label in np.unique(labels_array): cluster_sizes[int(label)] = int((labels_array == label).sum()) quality['cluster_sizes'] = cluster_sizes except Exception as e: return {'error': f'Clustering failed: {e}', 'truncated': False} result_id = _generate_id('clust') store.store_cluster_result( result_id=result_id, labels=full_labels.tolist(), n_clusters=n_clusters_found, params={ 'algorithm': algorithm, 'params': p, 'random_state': random_state, 'feature_columns': feature_cols, }, parent_dataset_id=dataset_id, ) return { 'cluster_result_id': result_id, 'n_clusters': n_clusters_found, 'n_noise': noise_count, 'labels_sample': full_labels[:50].tolist(), 'quality_metrics': quality, } async def _handle_run_clustering( dataset_id: str, cluster_columns: list[str], algorithm: str = 'hdbscan', params: Optional[dict] = None, random_state: int = 42, ) -> dict: """Run HDBSCAN or KMeans clustering (async wrapper — delegates to _cluster_sync).""" return _cluster_sync( dataset_id=dataset_id, cluster_columns=cluster_columns, algorithm=algorithm, params=params, random_state=random_state, ) # ── 6. evaluate_clustering ────────────────────────────────────────────── async def _handle_evaluate_clustering( cluster_result_id: str, dataset_id: str, metrics: Optional[list[str]] = None, ) -> dict: """Evaluate clustering result against the dataset.""" store = SessionStore() cluster_entry = store.get_cluster_result(cluster_result_id) if cluster_entry is None: return {'error': f'Cluster result not found: {cluster_result_id}', 'truncated': False} dataset_entry = store.get_dataset(dataset_id) if dataset_entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} labels = cluster_entry['labels'] default_metrics = ['silhouette', 'davies_bouldin', 'calinski_harabasz', 'noise_ratio', 'cluster_sizes'] requested = metrics or default_metrics scores: dict = {'n_clusters': cluster_entry['n_clusters']} try: lf: pl.LazyFrame = dataset_entry['lazyframe'] params = cluster_entry.get('params', {}) feature_cols = params.get('feature_columns', []) if not feature_cols and 'feature_columns' in params: pass elif feature_cols: df_features = lf.select(feature_cols).collect(streaming=True) data = df_features.to_numpy() # Align labels with data (handle missing rows) valid_mask = ~np.isnan(data).any(axis=1) if data.dtype.kind == 'f' else np.ones(len(data), dtype=bool) data_clean = data[valid_mask] labels_clean = np.array(labels)[:len(data_clean)] unique_labels = set(labels_clean) n_clusters_actual = len(unique_labels - {-1}) for m in requested: if m == 'silhouette' and n_clusters_actual >= 2: from sklearn.metrics import silhouette_score scores['silhouette'] = round(float( silhouette_score(data_clean, labels_clean, random_state=42)), 4) elif m == 'davies_bouldin' and n_clusters_actual >= 2: from sklearn.metrics import davies_bouldin_score scores['davies_bouldin'] = round(float( davies_bouldin_score(data_clean, labels_clean)), 4) elif m == 'calinski_harabasz' and n_clusters_actual >= 2: from sklearn.metrics import calinski_harabasz_score scores['calinski_harabasz'] = round(float( calinski_harabasz_score(data_clean, labels_clean)), 4) elif m == 'noise_ratio': noise = int((labels_clean == -1).sum()) scores['noise_ratio'] = round(noise / len(labels_clean), 4) elif m == 'cluster_sizes': sizes = {} for label in unique_labels: sizes[int(label)] = int((labels_clean == label).sum()) scores['cluster_sizes'] = sizes scores['n_clusters'] = cluster_entry['n_clusters'] except Exception as e: return {'error': f'Evaluation failed: {e}', 'truncated': False} return {'scores': scores} # ── 7. extract_features ───────────────────────────────────────────────── async def _handle_extract_features( dataset_id: str, cluster_result_id: str, top_k: int = 10, method: str = 'zscore', save_to_db: bool = True, ) -> dict: """Compute per-cluster distinguishing features and optionally persist.""" store = SessionStore() dataset_entry = store.get_dataset(dataset_id) if dataset_entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} cluster_entry = store.get_cluster_result(cluster_result_id) if cluster_entry is None: return {'error': f'Cluster result not found: {cluster_result_id}', 'truncated': False} try: lf: pl.LazyFrame = dataset_entry['lazyframe'] df = lf.collect(streaming=True) labels = cluster_entry['labels'] # Only use numeric columns for feature extraction numeric_cols = [c for c in df.columns if df[c].dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64, pl.UInt32, pl.UInt64)] if len(labels) != len(df): labels = labels[:len(df)] # Global statistics per column global_stats = {} for col in numeric_cols: s = df[col].drop_nulls() if len(s) > 0: global_stats[col] = { 'mean': float(s.mean()), 'std': float(s.std()) if s.std() is not None else 1.0, } unique_labels = sorted(set(labels)) all_features = [] for label in unique_labels: mask = np.array(labels) == label if mask.sum() < 2: continue cluster_df = df.filter(pl.Series('__mask__', mask)) cluster_features = [] for col in numeric_cols: c_series = cluster_df[col].drop_nulls() gs = global_stats.get(col) if gs is None or gs['std'] == 0 or len(c_series) == 0: continue c_mean = float(c_series.mean()) c_std = float(c_series.std()) if c_series.std() is not None else 0.0 if method == 'zscore': score = (c_mean - gs['mean']) / max(gs['std'], 1e-10) else: # ANOVA-like: F = (between-group variance) / (within-group variance) score = abs(c_mean - gs['mean']) / max(c_std + gs['std'], 1e-10) cluster_features.append({ 'feature_name': col, 'cluster_label': int(label), 'mean': round(c_mean, 4), 'std': round(c_std, 4), 'global_mean': round(gs['mean'], 4), 'global_std': round(gs['std'], 4), 'distinguishing_score': round(float(score), 4), 'distinguishing_method': method, }) # Sort by distinguishing score descending, take top_k cluster_features.sort(key=lambda x: abs(x['distinguishing_score']), reverse=True) all_features.extend(cluster_features[:top_k]) # ── Compute PCA-2D embedding ────────────────────────────────────── if len(numeric_cols) >= 2 and len(df) > 2: try: data_matrix = df.select(numeric_cols).to_numpy() from sklearn.preprocessing import StandardScaler from sklearn.decomposition import PCA data_scaled = StandardScaler().fit_transform(data_matrix) # Handle NaN/Inf values (np is imported at module level) data_scaled = np.nan_to_num(data_scaled, nan=0.0, posinf=0.0, neginf=0.0) pca_model = PCA(n_components=2, random_state=42) coords = pca_model.fit_transform(data_scaled) non_numeric = [c for c in df.columns if c not in numeric_cols and not c.startswith('_')] entity_col_pca = non_numeric[0] if non_numeric else df.columns[0] labels_list = cluster_entry.get('labels', []) # Save: create EntityProfile records with coordinates from analysis.models import EntityProfile as EP from analysis.models import ClusterResult as CR from analysis.models import AnalysisRun as AR csv_glob = dataset_entry.get('metadata', {}).get('csv_glob', 'manual') run_obj = AR.objects.filter(csv_glob=csv_glob).order_by('-id').first() if run_obj: run_id_for_ep = run_obj.id for i, row in enumerate(df.iter_rows(named=True)): if i >= len(coords): break ev = str(row.get(entity_col_pca, '')) if not ev: continue lbl = labels_list[i] if i < len(labels_list) else -1 cr = CR.objects.filter(run_id=run_id_for_ep, cluster_label=lbl).first() EP.objects.update_or_create(entity_value=ev, defaults={ 'run_id': run_id_for_ep, 'cluster_label': lbl, 'cluster': cr, 'embedding_x': float(coords[i,0]), 'embedding_y': float(coords[i,1])}) except Exception: import traceback, sys; print(f'[PCA_ERR] {traceback.format_exc()}', file=sys.stderr) # Persist to Django ORM db_saved = False if save_to_db: try: from asgiref.sync import sync_to_async from analysis.models import AnalysisRun as AnalysisRunModel ce = cluster_entry # local ref for closure async def _do_save(ce_inner): run, _ = await sync_to_async(AnalysisRunModel.objects.get_or_create, thread_sensitive=True)( csv_glob=dataset_entry.get('metadata', {}).get('csv_glob', 'manual'), defaults={'status': 'running', 'total_flows': 0} ) run.status = 'completed' run.cluster_count = ce_inner['n_clusters'] await sync_to_async(run.save)() return await sync_to_async(_save_features_to_db)(cluster_result_id, all_features, run_id=run.id) db_saved = await _do_save(ce) except Exception as e: import traceback, sys; print(f'[DB_SAVE_ERROR] {e}\n{traceback.format_exc()}', file=sys.stderr) db_saved = False except Exception as e: return {'error': f'Feature extraction failed: {e}', 'truncated': False} result_id = _generate_id('feat') store.store_feature_result( result_id=result_id, features=all_features, parent_cluster_result_id=cluster_result_id, ) return { 'feature_result_id': result_id, 'features': all_features, 'n_features': len(all_features), 'db_saved': db_saved, } def _save_features_to_db( cluster_result_id: str, features: list[dict], run_id: int | None = None, ) -> bool: """Save extracted features to the Django ORM ClusterFeature model. If run_id is provided, use it directly (skip csv_glob lookup). """ if run_id is not None: from analysis.models import AnalysisRun as AnalysisRun2 run = AnalysisRun2.objects.filter(id=run_id).first() if run is None: return False else: # Try to find the ClusterResult record by run association store = SessionStore() cluster_entry = store.get_cluster_result(cluster_result_id) if cluster_entry is None: return False parent_ds_id = cluster_entry.get('parent_dataset_id', '') from analysis.models import AnalysisRun ds_entry = store.get_dataset(parent_ds_id) if ds_entry is None: return False csv_glob = ds_entry.get('metadata', {}).get('csv_glob', '') # Compute n_clusters from features before using it unique_global = sorted(set(f['cluster_label'] for f in features)) n_clusters_global = len([l for l in unique_global if l >= 0]) try: run = AnalysisRun.objects.filter(csv_glob=csv_glob).order_by('-created_at').first() if run is None: run = AnalysisRun.objects.create( csv_glob=csv_glob, status='completed', cluster_count=n_clusters_global, ) except Exception: return False from analysis.models import ClusterResult as ClusterResultModel # Compute actual cluster sizes from features (not from cluster_entry — avoids scope issues) label_counts = {} for feat in features: lbl = feat['cluster_label'] label_counts[lbl] = label_counts.get(lbl, 0) + 1 unique_labels = sorted(set( f['cluster_label'] for f in features )) # Compute n_clusters from features (excluding noise -1) n_clusters = len([l for l in unique_labels if l >= 0]) for label in unique_labels: if label == -1: continue # skip noise cluster label_features = [f for f in features if f['cluster_label'] == label] if not label_features: continue cluster_db, _ = ClusterResultModel.objects.get_or_create( run=run, cluster_label=label, defaults={'size': label_counts.get(label, 1)}, ) # Update size if already exists if label_counts.get(label) is not None: ClusterResultModel.objects.filter( run=run, cluster_label=label ).update(size=label_counts[label]) from analysis.models import ClusterFeature as ClusterFeatureModel for feat in label_features: try: ClusterFeatureModel.objects.get_or_create( cluster=cluster_db, feature_name=feat['feature_name'], defaults={ 'mean': feat.get('mean'), 'std': feat.get('std'), 'median': None, 'p25': None, 'p75': None, 'missing_rate': None, 'distinguishing_score': feat.get('distinguishing_score'), 'distinguishing_method': feat.get('distinguishing_method', 'zscore'), } ) except Exception: continue # skip problematic features # Update run counts run.cluster_count = n_clusters run.save(update_fields=['cluster_count']) return True return True # ── 8. export_results ─────────────────────────────────────────────────── async def _handle_export_results( result_id: str, output_path: str, format: str = 'csv', overwrite: bool = False, ) -> dict: """Export a dataset (or result) to disk.""" store = SessionStore() out = Path(output_path) file_paths: list[str] = [] # Check if it's a dataset dataset_entry = store.get_dataset(result_id) if dataset_entry is not None: lf: pl.LazyFrame = dataset_entry['lazyframe'] try: df = lf.collect(streaming=True) if out.suffix == '': out = out / f"export_{result_id}.{format}" if out.exists() and not overwrite: return {'error': f'Output exists and overwrite=False: {out}', 'truncated': False} out.parent.mkdir(parents=True, exist_ok=True) if format == 'csv': df.write_csv(out) elif format == 'parquet': df.write_parquet(out) elif format == 'json': df.write_json(out) else: return {'error': f'Unsupported format: {format}', 'truncated': False} file_paths.append(str(out.resolve())) except Exception as e: return {'error': f'Export failed: {e}', 'truncated': False} return {'file_paths': file_paths, 'format': format, 'row_count': len(df)} # Check if it's a cluster result cluster_entry = store.get_cluster_result(result_id) if cluster_entry is not None: if out.suffix == '': out = out / f"cluster_{result_id}.json" out.parent.mkdir(parents=True, exist_ok=True) if out.exists() and not overwrite: return {'error': f'Output exists and overwrite=False: {out}', 'truncated': False} export_data = { 'result_id': result_id, 'n_clusters': cluster_entry['n_clusters'], 'params': cluster_entry['params'], 'labels_sample': cluster_entry['labels'][:1000], } with open(out, 'w', encoding='utf-8') as f: json.dump(export_data, f, default=str) file_paths.append(str(out.resolve())) return {'file_paths': file_paths, 'format': 'json'} # Check feature result feat_entry = store.get_feature_result(result_id) if feat_entry is not None: if out.suffix == '': out = out / f"features_{result_id}.json" out.parent.mkdir(parents=True, exist_ok=True) if out.exists() and not overwrite: return {'error': f'Output exists and overwrite=False: {out}', 'truncated': False} with open(out, 'w', encoding='utf-8') as f: json.dump(feat_entry['features'], f, default=str) file_paths.append(str(out.resolve())) return {'file_paths': file_paths, 'format': 'json'} return {'error': f'No dataset or result found for ID: {result_id}', 'truncated': False} # ── 9. list_datasets ──────────────────────────────────────────────────── async def _handle_list_datasets( _filter_type: Optional[str] = None, ) -> dict: """List all active datasets and results.""" store = SessionStore() datasets = store.list_datasets() return {'datasets': datasets, 'count': len(datasets)} # ── 10. drop_dataset ──────────────────────────────────────────────────── async def _handle_drop_dataset(dataset_id: str) -> dict: """Remove a dataset/result and release memory.""" store = SessionStore() dropped = (store.drop_dataset(dataset_id) or store.drop_cluster_result(dataset_id) or store.drop_feature_result(dataset_id)) if dropped: return {'dropped': True, 'dataset_id': dataset_id} return {'error': f'No dataset or result found: {dataset_id}', 'truncated': False} # ── 11. clone_dataset ─────────────────────────────────────────────────── async def _handle_clone_dataset(dataset_id: str) -> dict: """Shallow-clone a dataset.""" store = SessionStore() new_id = store.clone_dataset(dataset_id) if new_id is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} return {'original_dataset_id': dataset_id, 'cloned_dataset_id': new_id} # ── 12. build_entity_profiles ───────────────────────────────────────────── async def _handle_build_entity_profiles( dataset_id: str, entity_column: Optional[str] = None, entity_columns: Optional[list[str]] = None, auto_detect: bool = True, ) -> dict: """Auto-detect entity column(s), aggregate flows → entity profiles, persist. Steps: 1. Load dataset from session store. 2. Resolve entity column(s): *entity_columns* takes precedence, then *entity_column*, then auto-detection. 3. Run :func:`aggregate_by_entity` to build per-entity features. 4. Store aggregated result as a new dataset in the session. 5. Persist each entity's feature dict to the Django ORM (:class:`EntityProfile`). 6. Return summary (capped at ~8 KB). """ store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf: pl.LazyFrame = entry['lazyframe'] schema: dict[str, str] = entry.get('schema', {}) # ── Step 1: resolve entity column(s) ───────────────────────────────── used_columns: list[str] = [] detection_confidence: Optional[float] = None candidates: list[dict] = [] if entity_columns: # Multi-column support: use specified list directly missing = [c for c in entity_columns if c not in schema] if missing: return { 'error': f"Entity column(s) '{missing}' not found in schema. " f"Available: {list(schema.keys())[:20]}", 'truncated': False, } used_columns = entity_columns elif entity_column: if entity_column not in schema: return { 'error': f"Entity column '{entity_column}' not found in schema. " f"Available: {list(schema.keys())[:20]}", 'truncated': False, } used_columns = [entity_column] elif auto_detect: det_result = detect_entity_column(dataset_id) if det_result.get('error'): return {'error': f'Entity detection failed: {det_result["error"]}', 'truncated': False} candidates = det_result.get('candidates', []) recommended = det_result.get('recommended', '') used_columns = [recommended] if recommended else [] if candidates: detection_confidence = candidates[0]['score'] if not used_columns: return { 'error': 'No suitable entity column auto-detected. ' 'Provide entity_column manually.', 'truncated': False, } else: return { 'error': 'No entity_column provided and auto_detect is False.', 'truncated': False, } # ── Step 2: aggregate ──────────────────────────────────────────────── try: agg_lf, feature_cols = aggregate_by_entity(lf, used_columns, schema) except Exception as exc: return {'error': f'Aggregation failed: {exc}', 'truncated': False} # Collect to compute entity count and preview try: df_agg = agg_lf.collect(streaming=True) entity_count = len(df_agg) feature_columns = [c for c in df_agg.columns if not c.startswith('_')] except Exception as exc: return {'error': f'Failed to collect aggregated data: {exc}', 'truncated': False} # Build feature column schema agg_schema = {name: str(dtype) for name, dtype in zip(df_agg.columns, [df_agg[c].dtype for c in df_agg.columns])} # Primary entity column for metadata / DB persistence primary_entity_col = used_columns[0] # ── Step 3: store as new dataset in session ────────────────────────── entity_dataset_id = _generate_id('entity') store.store_dataset( dataset_id=entity_dataset_id, lazyframe=df_agg.lazy(), schema=agg_schema, metadata={ 'row_count': entity_count, 'entity_column': primary_entity_col, 'entity_columns': used_columns, 'parent_dataset_id': dataset_id, 'feature_columns': feature_columns, }, parent_id=dataset_id, ) # ── Step 4: persist to Django ORM (EntityProfile) ──────────────────── db_saved = False try: db_saved = _save_entity_profiles( dataset_id=dataset_id, entity_column=primary_entity_col, df_agg=df_agg, entity_count=entity_count, ) except Exception: db_saved = False # ── Step 5: build response (compact, ≤ 8 KB friendly) ──────────────── # Take a small sample of rows for the response sample_rows = df_agg.head(20).to_dict(as_series=False) if entity_count > 0 else {} result = { 'entity_dataset_id': entity_dataset_id, 'entity_count': entity_count, 'entity_columns_used': used_columns, 'feature_columns': feature_columns, 'detection_confidence': detection_confidence, 'detection_candidates': candidates[:5] if candidates else [], 'db_saved': db_saved, 'sample': sample_rows, } # Ensure truncation within ~8 KB (the global _truncate_response uses 64 KB # by default, but we cap sample here for this tool) import json body = json.dumps(result, default=str, ensure_ascii=False) if len(body) > 8 * 1024: # Drop sample and candidates, keep only counts result.pop('sample', None) result.pop('detection_candidates', None) result['truncated'] = True else: result['truncated'] = False return result def _save_entity_profiles( dataset_id: str, entity_column: str, df_agg: pl.DataFrame, entity_count: int, ) -> bool: """Persist aggregated entity profiles to the Django ORM. Creates/updates :class:`EntityProfile` records linked to the :class:`AnalysisRun` associated with *dataset_id*. """ import django import os os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'tianxuan.settings') django.setup() from analysis.models import AnalysisRun, EntityProfile store = SessionStore() ds_entry = store.get_dataset(dataset_id) if ds_entry is None: return False csv_glob = ds_entry.get('metadata', {}).get('csv_glob', '') try: run = AnalysisRun.objects.filter(csv_glob=csv_glob).order_by('-created_at').first() if run is None: run = AnalysisRun.objects.create( csv_glob=csv_glob, status='aggregating', entity_count=entity_count, entity_column=entity_column, ) except Exception: return False # Build EntityProfile records in bulk entity_values = df_agg[entity_column].to_list() feature_cols = [c for c in df_agg.columns if not c.startswith('_') and c != entity_column] profiles = [] for row_idx in range(len(df_agg)): entity_val = str(entity_values[row_idx]) feature_dict = { col: _serialize_value(df_agg[col][row_idx]) for col in feature_cols } profiles.append(EntityProfile( run=run, entity_value=entity_val, feature_json=feature_dict, )) try: EntityProfile.objects.bulk_create(profiles, ignore_conflicts=True) except Exception: return False # Update run metadata run.status = 'completed' run.entity_count = entity_count run.entity_column = entity_column run.save(update_fields=['status', 'entity_count', 'entity_column']) return True def _serialize_value(val) -> object: """Convert a Polars/Numpy scalar to a JSON-safe Python value.""" if val is None: return None if isinstance(val, (int, float, str, bool)): # Handle NaN / Inf if isinstance(val, float): import math if math.isnan(val) or math.isinf(val): return None return val if isinstance(val, (list, tuple)): return [_serialize_value(v) for v in val] return str(val)