"""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. filter_and_cluster 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 .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=( "加载匹配glob模式的CSV文件。自动检测每个文件的BOM，验证schema一致性，" "纵向合并，并将结果作为惰性数据集存储在会话存储中。" ), 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=( "计算数据集的每列统计信息和数值相关矩阵。结果自动截断至约16 KB以" "节省LLM上下文窗口空间。" ), 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=( "对数据集应用列级过滤器，将结果存储为新的派生数据集。支持eq/neq/gt/gte/" "lt/lte/contains/not_contains/in/not_in/is_null/is_not_null运算符。" ), 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=( "预处理选定列：标准化（StandardScaler）、独热编码、缺失值填充和列删除。" ), 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=( "对选定列执行聚类。使用HDBSCAN（默认）或KMeans。特征在聚类前进行" "标准化缩放。返回聚类标签和质量指标。" ), 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=( "评估数据集的聚类结果。当前支持轮廓系数（silhouette）、Davies-Bouldin指数、" "Calinski-Harabasz指数、噪声比（noise_ratio）和簇大小（cluster_sizes）指标。" ), 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=( "提取每个聚类的区分特征并自动保存到Django ORM（ClusterFeature模型）。" "使用z-score或ANOVA对每个聚类的特征重要性进行排序。" ), 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=( "将数据集或聚类结果以CSV或Parquet格式导出到磁盘。数据集物化后写入；" "聚类结果序列化为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="列出会话存储中的所有活跃数据集和结果。", inputSchema={ "type": "object", "properties": {}, }, ), Tool( name="drop_dataset", description=( "从会话存储中移除数据集（或结果）并通过垃圾回收释放内存。" ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Dataset or result identifier to remove", }, }, "required": ["dataset_id"], }, ), Tool( name="clone_dataset", description=( "浅拷贝数据集。由于LazyFrame是查询计划，不会复制实际数据。适用于分支分析。" ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Source dataset identifier", }, }, "required": ["dataset_id"], }, ), Tool( name="filter_and_cluster", description=( "一步完成：先应用列级过滤器，再对过滤结果执行聚类。支持eq/neq/gt/gte/lt/lte/" "contains/not_contains/in/not_in/is_null/is_not_null运算符和HDBSCAN/KMeans聚类。" "自动选择数值列作为聚类特征。将过滤后的数据集和聚类结果都存储在会话中。" ), 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", }, "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)." ), }, }, "required": ["dataset_id", "filters"], }, ), Tool( name="compute_scores", description=( "在实体数据上计算自适应代理/异常/威胁评分。使用PCA从数据自动学习" "特征权重，计算加权代理评分（proxy_score）、基于IQR的风险等级" "（normal/watch/suspicious/critical）和威胁评分（threat_score）。" "返回评分摘要与阈值。" ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Entity dataset ID from filter_and_cluster or a scored dataset", }, }, "required": ["dataset_id"], }, ), Tool( name="detect_anomalies", description=( "对已评分的实体数据运行孤立森林（Isolation Forest）异常检测。识别偏离" "正常模式的实体。分块处理以应对大数据集。返回异常实体ID列表。" ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Scored entity dataset ID (after compute_scores)", }, "contamination": { "type": "number", "description": "Expected proportion of anomalies (0.0-0.5)", "default": 0.05, }, }, "required": ["dataset_id"], }, ), Tool( name="visualize_anomalies", description=( "为高风险实体生成PCA-2D嵌入和聚类可视化数据。按风险等级和聚类着色。" "返回用于前端渲染的散点图JSON数据。" ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Scored entity dataset ID (after compute_scores)", }, }, "required": ["dataset_id"], }, ), # ── Edge-case / Diagnostic tools (unexpected situations) ──────── Tool( name="validate_data", description=( "诊断工具：检查数据质量问题——文件间schema一致性、缺失值比例、列类型冲突、" "编码问题和空列。当流水线步骤因schema错误失败，或怀疑CSV数据损坏/格式异常时使用。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Dataset ID from load_data"}, "fix_issues": {"type": "boolean", "description": "Auto-fix schema inconsistencies", "default": False}, }, "required": ["dataset_id"], }, ), Tool( name="explore_distributions", description=( "诊断工具：计算每列分布统计（唯一值计数、空值比例、最小/最大/均值/标准差、" "前10高频值、直方图分箱）。当聚类产生意外结果、实体检测失败，或需要在选择" "分析策略前了解数据时使用。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Dataset ID from load_data or entity data"}, "columns": {"type": "array", "items": {"type": "string"}, "description": "Specific columns to analyze (default: auto-detect numeric)"}, "max_sample": {"type": "integer", "description": "Max rows to sample", "default": 10000}, }, "required": ["dataset_id"], }, ), Tool( name="find_outliers", description=( "诊断工具：使用IQR方法检测数值列中的统计异常值。返回任何数值列超过中位数3倍IQR的行。" "当怀疑数据质量问题或极端值偏离分析时使用。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Dataset ID from load_data"}, "columns": {"type": "array", "items": {"type": "string"}, "description": "Columns to check (default: all numeric)"}, "threshold": {"type": "number", "description": "IQR multiplier (default 3.0)", "default": 3.0}, }, "required": ["dataset_id"], }, ), Tool( name="diagnose_clustering", description=( "诊断工具：当聚类产生太少/太多聚类、全部为噪声或轮廓系数低时——诊断原因。" "检查：特征方差、相关矩阵、PCA解释方差、数据稀疏性，并建议参数调整" "（min_cluster_size、算法选择、特征选择）。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Entity dataset ID"}, "cluster_result_id": {"type": "string", "description": "Cluster result ID from previous run_clustering"}, }, "required": ["dataset_id", "cluster_result_id"], }, ), Tool( name="compare_datasets", description=( "诊断工具：比较两个数据集（如处理前/后，或两次不同运行）。报告schema差异、" "行数变化、列值分布偏移。用于验证处理步骤未损坏数据或比较不同运行的结果。" ), inputSchema={ "type": "object", "properties": { "dataset_id_a": {"type": "string", "description": "First dataset ID"}, "dataset_id_b": {"type": "string", "description": "Second dataset ID"}, }, "required": ["dataset_id_a", "dataset_id_b"], }, ), Tool( name="export_debug_sample", description=( "诊断工具：导出原始数据样本为JSON用于外部调试。当流水线深处发生错误且需要" "检查实际数据值时使用。返回前N行所有列。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Dataset ID"}, "max_rows": {"type": "integer", "description": "Rows to export", "default": 50}, }, "required": ["dataset_id"], }, ), Tool( name="repair_schema", description=( "诊断工具：通过对齐列名（不区分大小写合并、下划线/连字符规范化）并用空值填充" "缺失列，尝试修复文件间的schema不匹配。当load_data因CSV文件具有不同列集而" "失败时使用。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Failed dataset ID from load_data"}, }, "required": ["dataset_id"], }, ), # ── Read-only analysis tools (LLM deep-dive, no side effects) ──── Tool( name="analyze_patterns", description=( "分析工具：检测流量模式——top源/目标IP、端口分布、TLS版本分布、协议混合。" "返回汇总统计。只读，不修改数据。用于了解高层流量结构。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Dataset ID from load_data"}, "top_n": {"type": "integer", "description": "Number of top items per category", "default": 10}, }, "required": ["dataset_id"], }, ), Tool( name="analyze_temporal", description=( "分析工具：分析随时间变化的流量——小时/日流量计数、繁忙时段、空闲时段。" "需要时间戳列。只读。用于了解时间相关模式和周期性。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Dataset ID from load_data"}, "timestamp_column": {"type": "string", "description": "Column name for timestamps (auto-detected if omitted)"}, }, "required": ["dataset_id"], }, ), Tool( name="analyze_tls_health", description=( "分析工具：TLS安全态势评估——检查过时的TLS版本（1.0/1.1）、弱加密套件、" "证书问题、SNI异常。只读。用于评估流量的TLS安全性。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Dataset ID from load_data"}, }, "required": ["dataset_id"], }, ), Tool( name="analyze_geo_distribution", description=( "分析工具：流量地理分布——top源/目标国家、异常地点对（如不可能的行程）、" "ISP多样性。只读。需要国家或经纬度列。用于检测地理异常。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Dataset ID from load_data"}, }, "required": ["dataset_id"], }, ), Tool( name="analyze_entity_detail", description=( "分析工具：深入分析单个实体——其所有流、使用的TLS版本、目标多样性、时间模式、" "异常评分。只读。在compute_scores后调查特定可疑实体时使用。" ), inputSchema={ "type": "object", "properties": { "dataset_id": {"type": "string", "description": "Entity dataset ID (after filter_and_cluster or compute_scores)"}, "entity_value": {"type": "string", "description": "Entity value to investigate (e.g. an IP address or SNI)"}, }, "required": ["dataset_id", "entity_value"], }, ), Tool( name="compute_distance_matrix", description=( "LLM驱动工具：对数据集的每一行执行用户编写的Python函数，返回距离/相似度评分摘要。" "LLM提供函数体，签名格式为`def distance_fn(row: dict) -> float:`。执行器将其包装，" "在受限命名空间（仅math、numpy）中逐行运行，返回得分的最小/最大/均值/标准差以及" "评分行样本。用于现有工具无法表达的一次性度量、过滤或评分任务。" ), inputSchema={ "type": "object", "properties": { "dataset_id": { "type": "string", "description": "Source dataset identifier", }, "python_function": { "type": "string", "description": ( "Body of a Python function with signature `def distance_fn(row: dict) -> float:`. " "Example: \"return abs(row.get('col1', 0)) + row.get('col2', 0)\". " "Available namespace: math, numpy (as np). No dangerous modules." ), }, "columns": { "type": "array", "items": {"type": "string"}, "description": "Columns to pass to the function (default: auto-detect numeric columns)", }, "max_sample": { "type": "integer", "description": "Maximum rows to evaluate (default 10000 for performance)", "default": 10000, }, }, "required": ["dataset_id", "python_function"], }, ), ] # ═══════════════════════════════════════════════════════════════════════ # 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, 'compute_scores': _handle_compute_scores, 'filter_and_cluster': _handle_filter_and_cluster, 'detect_anomalies': _handle_detect_anomalies, 'visualize_anomalies': _handle_visualize_anomalies, 'validate_data': _handle_validate_data, 'explore_distributions': _handle_explore_distributions, 'find_outliers': _handle_find_outliers, 'diagnose_clustering': _handle_diagnose_clustering, 'compare_datasets': _handle_compare_datasets, 'export_debug_sample': _handle_export_debug_sample, 'repair_schema': _handle_repair_schema, 'analyze_patterns': _handle_analyze_patterns, 'analyze_temporal': _handle_analyze_temporal, 'analyze_tls_health': _handle_analyze_tls_health, 'analyze_geo_distribution': _handle_analyze_geo_distribution, 'analyze_entity_detail': _handle_analyze_entity_detail, 'compute_distance_matrix': _handle_compute_distance_matrix, } 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 ─────────────────────────────────────────────────── 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 on selected columns (async).""" 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 } # Downsample before collect: compute row count via fast lazy count try: row_count = lf.select(pl.len()).collect(streaming=True).item() except Exception: row_count = 0 if row_count > 100000: # Sample to 100K before collect to avoid OOM on 10M+ rows lf = lf.collect(streaming=True).sample(n=100000, seed=42).lazy() 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: Downsampling → extract matching indices ────────────── # Downsample to 50K max, store indices so _handle_extract_features # can filter the same rows. _ds_idx = None if len(data) > 50000: rng = np.random.default_rng(42) _ds_idx = rng.choice(len(data), 50000, replace=False) data = data[_ds_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: # All features have zero variance — fall back to unfiltered data keep_var = np.arange(data.shape[1]) 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, '_ds_idx': _ds_idx.tolist() if _ds_idx is not None else None, }, 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, } # ── 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, run_id: int | None = None, ) -> 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'] labels = cluster_entry['labels'] # D5: If clustering downsampled the data, filter matching rows _ds_idx = cluster_entry.get('params', {}).get('_ds_idx') if _ds_idx is not None: # Filter in lazy plan before collect — avoids loading full 10M rows df = (lf.with_row_index('__ds_idx') .filter(pl.col('__ds_idx').is_in(_ds_idx)) .drop('__ds_idx') .collect(streaming=True)) else: df = lf.collect(streaming=True) # Truncate labels if df was filtered but labels weren't if len(labels) != len(df): labels = labels[:len(df)] # 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 from asgiref.sync import sync_to_async csv_glob = dataset_entry.get('metadata', {}).get('csv_glob', 'manual') run_obj = await sync_to_async( AR.objects.filter(csv_glob=csv_glob).order_by('-id').first, thread_sensitive=True )() if run_obj: run_id_for_ep = run_obj.id profiles = [] cr_cache = {} 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 cache_key = f'{run_id_for_ep}_{lbl}' if cache_key not in cr_cache: cr_cache[cache_key] = CR.objects.filter(run_id=run_id_for_ep, cluster_label=lbl).first() profiles.append(EP( entity_value=ev, run_id=run_id_for_ep, cluster_label=lbl, cluster=cr_cache[cache_key], embedding_x=float(coords[i, 0]), embedding_y=float(coords[i, 1]), )) if profiles: EP.objects.bulk_create(profiles, ignore_conflicts=True, batch_size=1000) 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): if run_id is not None: run = await sync_to_async(AnalysisRunModel.objects.get)(id=run_id) else: csv_glob_fallback = dataset_entry.get('metadata', {}).get('csv_glob', 'manual') run = await sync_to_async( AnalysisRunModel.objects.filter(csv_glob=csv_glob_fallback).order_by('-id').first, thread_sensitive=True )() if run is None: run = await sync_to_async(AnalysisRunModel.objects.create)( csv_glob=csv_glob_fallback, 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} # ── Adaptive anomaly scoring helpers (inlined from entity_aggregator) ───── _ANOMALY_FEATURE_COLS = [ 'non_std_port_rate', 'modern_tls_rate', 'sni_missing_ratio', 'recoverable_ratio', 'flow_count', 'total_bytes_sent', 'total_packets', 'unique_dst_ips', 'unique_dst_ports', 'unique_protocols', 'has_sni', 'avg_sni_length', 'countries_visited', ] def _add_adaptive_scores(lf: pl.LazyFrame, feature_names: list[str]) -> pl.LazyFrame: """Add anomaly scores using data-driven weights (PCA + IQR + percentile). Steps: 1. Find which anomaly feature columns exist in the aggregated data 2. Collect a sample (max 10K rows) for PCA weight learning 3. Learn feature weights from PCA first component loadings 4. Compute weighted proxy_score across ALL rows (vectorized, no per-row ops) 5. Add percentile rank calibration 6. Infer risk thresholds from data distribution """ _agg_cols = lf.collect_schema().names() avail = [c for c in _ANOMALY_FEATURE_COLS if c in _agg_cols] if not avail: return lf try: # Step 1-2: Sample up to 10K rows for weight learning n_rows = lf.select(pl.len()).collect(streaming=True).item() sample_lf = lf if n_rows <= 10000 else lf.sample(n=10000, seed=42) df_sample = sample_lf.select(avail).collect(streaming=True) mat = df_sample.to_numpy() mat = np.nan_to_num(mat, nan=0.0) # Step 3: PCA weight learning from sklearn.decomposition import PCA n_comp = min(1, min(mat.shape) - 1) if n_comp < 1 or mat.shape[1] < 2: return lf pca = PCA(n_components=1, random_state=42).fit(mat) raw_weights = np.abs(pca.components_[0]) # Normalise to sum = 1 weights = dict(zip(avail, raw_weights / raw_weights.sum())) # Step 4: Weighted sum expression (all vectorised, no loop) expr_parts = [] for col, w in weights.items(): if col == 'modern_tls_rate': # modern_tls_rate is good — invert for anomaly scoring expr_parts.append((1.0 - pl.col(col).fill_null(0.5)) * w) else: expr_parts.append(pl.col(col).fill_null(0.0) * w) if expr_parts: proxy_expr = expr_parts[0] for p in expr_parts[1:]: proxy_expr = proxy_expr + p else: return lf lf = lf.with_columns(proxy_expr.clip(0, 1).alias('_raw_score')) # Step 5: Percentile rank calibration (0-1, relative position) sample_scores = df_sample.select( pl.sum_horizontal([ (1.0 - pl.col('modern_tls_rate').fill_null(0.5)) * weights.get('modern_tls_rate', 0) if c == 'modern_tls_rate' else pl.col(c).fill_null(0.0) * weights.get(c, 0) for c in avail ]).clip(0, 1) ).to_series().to_numpy() # Compute IQR thresholds from sample q25, q50, q75 = np.percentile(sample_scores, [25, 50, 75]) iqr = q75 - q25 upper_fence = float(q75 + 1.5 * iqr) if iqr > 0 else 0.9 extreme_fence = float(q75 + 3.0 * iqr) if iqr > 0 else 0.99 # Step 6: Risk level labels (0.0-1.0) lf = lf.with_columns([ pl.col('_raw_score').alias('proxy_score'), pl.when(pl.col('_raw_score') >= extreme_fence) .then(pl.lit('critical')) .when(pl.col('_raw_score') >= upper_fence) .then(pl.lit('suspicious')) .when(pl.col('_raw_score') >= q50) .then(pl.lit('watch')) .otherwise(pl.lit('normal')) .alias('risk_level'), (pl.col('_raw_score') * 0.6 + 0.4 * (pl.col('_raw_score') > q75).cast(pl.Float64)) .alias('threat_score'), ]).drop('_raw_score') valid_counts = {c: int((mat[:, i] != 0).sum()) for i, c in enumerate(avail)} import logging logger = logging.getLogger(__name__) logger.info( '[SCORE] adaptive weights from PCA: %s | thresholds: p50=%.3f upper=%.3f extreme=%.3f | sample=%d/%d', {c: f'{w:.3f}' for c, w in sorted(weights.items(), key=lambda x: -x[1])}, float(q50), upper_fence, extreme_fence, min(n_rows, 10000), n_rows, ) except Exception as exc: import logging logging.getLogger(__name__).warning('[SCORE] adaptive scoring skipped: %s', exc) return lf # ── 12. compute_scores ───────────────────────────────────────────────────── async def _handle_compute_scores(dataset_id: str) -> dict: """Compute adaptive proxy/anomaly/threat scores on entity data.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf = entry['lazyframe'] lf = _add_adaptive_scores(lf, []) store.store_dataset(dataset_id, lf, schema=entry.get('schema', {}), metadata=entry.get('metadata', {})) return {'status': 'scored', 'dataset_id': dataset_id} # ── 13. filter_and_cluster ───────────────────────────────────────────────── async def _handle_filter_and_cluster( dataset_id: str, filters: list[dict], logic: str = 'and', algorithm: str = 'hdbscan', params: dict | None = None, ) -> dict: """Apply column-level filters then run clustering on the filtered result.""" 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', {}) # ── Step 1: Apply filters ───────────────────────────────────────────── try: exprs = [_build_filter_expr( f['column'], f.get('op', 'eq'), f.get('value'), ) for f in filters] if exprs: 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) else: filtered_lf = lf # 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} # ── Step 2: Auto-detect numeric feature columns ─────────────────────── try: available = filtered_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 = [ 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: return { 'error': 'No numeric columns found for clustering after filtering', 'filtered_dataset_id': None, 'truncated': False, } except Exception as e: return {'error': f'Feature detection failed: {e}', 'truncated': False} # ── Step 3: Store filtered dataset ──────────────────────────────────── filtered_dataset_id = _generate_id('flt') store.store_dataset( dataset_id=filtered_dataset_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, ) # ── Step 4: Run clustering (inlined from _handle_run_clustering) ─────── p = params or {} try: # Downsample before collect try: total_count = filtered_lf.select(pl.len()).collect(streaming=True).item() except Exception: total_count = 0 clamp_lf = filtered_lf if total_count > 100000: clamp_lf = filtered_lf.collect(streaming=True).sample(n=100000, seed=42).lazy() df_features = clamp_lf.select(feature_cols).collect(streaming=True) # Drop all-NaN 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', 'filtered_dataset_id': filtered_dataset_id, '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) # 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'}, 'filtered_dataset_id': filtered_dataset_id, 'truncated': False, } # Low memory check low_memory = False try: import psutil mem = psutil.virtual_memory() low_memory = mem.available < 2 * 1024 ** 3 except ImportError: pass if low_memory and len(data) > 100000: rng = np.random.default_rng(42) idx = rng.choice(len(data), 50000, replace=False) data = data[idx] # Downsample to 50K max _ds_idx = None if len(data) > 50000: rng = np.random.default_rng(42) _ds_idx = rng.choice(len(data), 50000, replace=False) data = data[_ds_idx] # Variance filtering if data.shape[1] > 1: variances = np.var(data, axis=0) keep_var = np.where(variances >= 1e-10)[0] if len(keep_var) == 0: keep_var = np.arange(data.shape[1]) if len(keep_var) < data.shape[1]: data = data[:, keep_var] feature_cols = [feature_cols[i] for i in keep_var] # Standard scale from sklearn.preprocessing import StandardScaler data_scaled = StandardScaler().fit_transform(data) # Correlation filtering 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] # PCA dimensionality reduction (if >50 features) n_features = data_scaled.shape[1] if n_features > 50: from sklearn.decomposition import PCA data_scaled = PCA(n_components=50, random_state=42).fit_transform(data_scaled) # Hyperparameter auto-tuning 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: p['min_cluster_size'] = min( p.get('min_cluster_size', 5), max(n // 2, 2), ) # Fit model from sklearn.cluster import HDBSCAN, MiniBatchKMeans if algorithm == 'kmeans': n_clusters_param = p.get('n_clusters', 3) model = MiniBatchKMeans( n_clusters=n_clusters_param, random_state=42, 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 # 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=42) 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}', 'filtered_dataset_id': filtered_dataset_id, 'truncated': False} result_id = _generate_id('clust') store.store_cluster_result( result_id=result_id, labels=labels_array.tolist(), n_clusters=n_clusters_found, params={ 'algorithm': algorithm, 'params': p, 'random_state': 42, 'feature_columns': feature_cols, '_ds_idx': _ds_idx.tolist() if _ds_idx is not None else None, }, parent_dataset_id=filtered_dataset_id, ) return { 'cluster_result_id': result_id, 'filtered_dataset_id': filtered_dataset_id, 'n_clusters': n_clusters_found, 'n_noise': noise_count, 'labels_sample': labels_array[:50].tolist(), 'quality_metrics': quality, 'feature_columns': feature_cols, } # ── 14. detect_anomalies ────────────────────────────────────────────────── async def _handle_detect_anomalies(dataset_id: str, contamination: float = 0.05) -> dict: """Run Isolation Forest on scored entity data.""" from sklearn.ensemble import IsolationForest 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 = entry.get('schema', {}) numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'} num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')] if not num_cols: return {'error': 'No numeric columns for anomaly detection', 'truncated': False} # Collect scored data (already sampled by pipeline head param) df = lf.select(num_cols).collect(streaming=True) mat = df.to_numpy() mat = np.nan_to_num(mat, nan=0.0) # Isolation Forest in chunks n = len(mat) chunk_size = min(50000, n) labels = np.zeros(n, dtype=int) for start in range(0, n, chunk_size): end = min(start + chunk_size, n) iso = IsolationForest(n_estimators=100, max_samples=min(256, chunk_size), contamination=contamination, random_state=42, n_jobs=-1) labels[start:end] = iso.fit_predict(mat[start:end]) anomaly_count = int((labels == -1).sum()) return { 'status': 'completed', 'total_entities': n, 'anomaly_count': anomaly_count, 'anomaly_rate': round(anomaly_count / max(n, 1), 4), 'dataset_id': dataset_id, } # ── 15. visualize_anomalies ─────────────────────────────────────────────── async def _handle_visualize_anomalies(dataset_id: str) -> dict: """Generate PCA-2D embedding for anomaly visualization.""" from sklearn.preprocessing import StandardScaler from sklearn.decomposition import PCA 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 = entry.get('schema', {}) numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'} num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')] if len(num_cols) < 2: return {'error': 'Need at least 2 numeric columns for PCA', 'truncated': False} # Sample for PCA visualisation try: n_total = lf.select(pl.len()).collect(streaming=True).item() except Exception: n_total = 0 sample_lf = lf if n_total <= 10000 else lf.collect(streaming=True).sample(n=10000, seed=42).lazy() df = sample_lf.select(num_cols).collect(streaming=True) mat = StandardScaler().fit_transform(df.to_numpy()) mat = np.nan_to_num(mat, nan=0.0) coords = PCA(n_components=2, random_state=42).fit_transform(mat) risk_col = next((c for c in schema if 'risk' in c.lower()), None) score_col = next((c for c in schema if 'proxy_score' in c.lower() or 'threat' in c.lower()), None) non_num = [c for c in df.columns if c not in num_cols and not c.startswith('_')] ent_col = non_num[0] if non_num else df.columns[0] entities = df.select(ent_col).to_series().to_list() if ent_col in df.columns else [] scatter = [ {'x': float(coords[i, 0]), 'y': float(coords[i, 1]), 'entity': str(entities[i]) if i < len(entities) else '', 'score': round(float(mat[i].mean()), 4)} for i in range(len(coords)) ] return { 'status': 'completed', 'scatter_points': len(scatter), 'scatter_data': scatter[:200], # cap for LLM context 'dataset_id': dataset_id, } # ── 16. validate_data ────────────────────────────────────────────────── async def _handle_validate_data(dataset_id: str, fix_issues: bool = False) -> dict: """Check data quality: schema consistency, missing values, type conflicts.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) issues = [] try: df_sample = lf.head(1000).collect(streaming=True) for col in schema: null_count = df_sample[col].is_null().sum() if null_count > 500: issues.append(f'{col}: {null_count/10:.0f}% null') try: unique = df_sample[col].n_unique() if unique <= 1 and null_count == 0: issues.append(f'{col}: constant value (no variation)') except Exception: pass total_rows = lf.select(pl.len()).collect(streaming=True).item() except Exception as exc: return {'error': f'Validation failed: {exc}', 'truncated': False} return { 'status': 'ok', 'total_rows': total_rows, 'columns': len(schema), 'issues_found': len(issues), 'issues': issues[:20], 'recommendation': 'Try explore_distributions for detailed column analysis' if issues else 'No issues found', } # ── 17. explore_distributions ────────────────────────────────────────── async def _handle_explore_distributions(dataset_id: str, columns: list[str] = None, max_sample: int = 10000) -> dict: """Per-column distribution statistics.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'} if not columns: columns = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types][:10] try: n = lf.select(pl.len()).collect(streaming=True).item() sample_lf = lf if n <= max_sample else lf.collect(streaming=True).sample(n=max_sample, seed=42).lazy() df = sample_lf.select(columns).collect(streaming=True) except Exception as exc: return {'error': f'Failed to collect sample: {exc}', 'truncated': False} stats = {} for col in columns: try: s = df[col] stats[col] = { 'dtype': str(s.dtype), 'null_pct': round(float(s.is_null().mean() * 100), 2), 'unique': int(s.n_unique()), 'top_values': [str(v) for v in s.head(5).to_list()], } if s.dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64): stats[col]['min'] = float(s.min()) if s.min() is not None else None stats[col]['max'] = float(s.max()) if s.max() is not None else None stats[col]['mean'] = round(float(s.mean()), 4) if s.mean() is not None else None stats[col]['std'] = round(float(s.std()), 4) if s.std() is not None else None except Exception: stats[col] = {'error': 'failed to compute'} return {'columns': stats, 'total_rows': n, 'sampled': min(n, max_sample)} # ── 18. find_outliers ───────────────────────────────────────────────── async def _handle_find_outliers(dataset_id: str, columns: list[str] = None, threshold: float = 3.0) -> dict: """Find statistical outliers using IQR method.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'} if not columns: columns = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types] try: n = lf.select(pl.len()).collect(streaming=True).item() sample_lf = lf if n <= 100000 else lf.collect(streaming=True).sample(n=100000, seed=42).lazy() df = sample_lf.select(columns).collect(streaming=True) except Exception as exc: return {'error': f'Failed: {exc}', 'truncated': False} outlier_counts = {} for col in columns: try: s = df[col] q1, q3 = s.quantile(0.25), s.quantile(0.75) iqr = q3 - q1 lo = q1 - threshold * iqr hi = q3 + threshold * iqr n_out = ((s < lo) | (s > hi)).sum() if n_out > 0: outlier_counts[col] = int(n_out) except Exception: pass return {'outlier_counts': outlier_counts, 'total_sampled': len(df), 'threshold_iqr': threshold} # ── 19. diagnose_clustering ─────────────────────────────────────────── async def _handle_diagnose_clustering(dataset_id: str, cluster_result_id: str) -> dict: """Diagnose clustering quality issues.""" from sklearn.decomposition import PCA store = SessionStore() entry = store.get_dataset(dataset_id) cluster_entry = store.get_cluster_result(cluster_result_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} if cluster_entry is None: return {'error': f'Cluster result not found: {cluster_result_id}', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'} num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')] n_clusters = cluster_entry.get('n_clusters', 0) n_noise = cluster_entry.get('n_noise', 0) quality = cluster_entry.get('quality_metrics', {}) labels = cluster_entry.get('labels', []) params = cluster_entry.get('params', {}) algo = params.get('algorithm', 'unknown') feature_cols = params.get('feature_columns', []) diagnosis = [] if n_clusters <= 1: diagnosis.append('Only 1 cluster — data may lack structure or need feature scaling') if n_noise > 0.5 * len(labels) if labels else False: diagnosis.append(f'High noise ratio ({n_noise}/{len(labels)}) — try HDBSCAN with larger min_cluster_size') sil = quality.get('silhouette_score') if sil and sil < 0.2: diagnosis.append(f'Low silhouette ({sil:.3f}) — clusters overlap, try KMeans or feature selection') if len(feature_cols) < 3: diagnosis.append(f'Only {len(feature_cols)} features — add more columns for meaningful clusters') # PCA variance check try: df_sample = lf.select(num_cols).head(5000).collect(streaming=True) mat = df_sample.to_numpy() pca = PCA().fit(mat) var_explained = pca.explained_variance_ratio_ if var_explained[0] > 0.95: diagnosis.append(f'First PCA component explains {var_explained[0]:.0%} variance — data nearly 1D') except Exception: pass return { 'diagnosis': diagnosis, 'n_clusters': n_clusters, 'n_noise': n_noise, 'silhouette': sil, 'algorithm': algo, 'n_features': len(feature_cols), 'suggestions': [ f'Try algorithm={"kmeans" if algo == "hdbscan" else "hdbscan"}', 'Increase min_cluster_size for fewer, larger clusters', 'Use explore_distributions to check feature quality', ], } # ── 20. compare_datasets ───────────────────────────────────────────── async def _handle_compare_datasets(dataset_id_a: str, dataset_id_b: str) -> dict: """Compare two datasets.""" store = SessionStore() entry_a, entry_b = store.get_dataset(dataset_id_a), store.get_dataset(dataset_id_b) if entry_a is None or entry_b is None: return {'error': 'One or both datasets not found', 'truncated': False} def _get_info(entry): schema = entry.get('schema', {}) lf = entry['lazyframe'] try: n = lf.select(pl.len()).collect(streaming=True).item() except Exception: n = -1 return {'schema': set(schema.keys()), 'row_count': n, 'columns': len(schema)} info_a, info_b = _get_info(entry_a), _get_info(entry_b) common_cols = info_a['schema'] & info_b['schema'] only_a = info_a['schema'] - info_b['schema'] only_b = info_b['schema'] - info_a['schema'] return { 'dataset_a': {'id': dataset_id_a, **info_a}, 'dataset_b': {'id': dataset_id_b, **info_b}, 'common_columns': len(common_cols), 'columns_only_in_a': list(only_a)[:10] if only_a else [], 'columns_only_in_b': list(only_b)[:10] if only_b else [], 'row_diff': info_a['row_count'] - info_b['row_count'], 'identical_schema': len(only_a) == 0 and len(only_b) == 0, } # ── 21. export_debug_sample ─────────────────────────────────────────── async def _handle_export_debug_sample(dataset_id: str, max_rows: int = 50) -> dict: """Export raw data sample as JSON for debugging.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} try: df = entry['lazyframe'].head(max_rows).collect(streaming=True) rows = df.to_dict(as_series=False) # Convert to list-of-dicts, serialize safely result = [] for i in range(len(df)): row = {} for col in df.columns: val = rows[col][i] if isinstance(val, (float, int, str, bool)): row[col] = val elif val is None: row[col] = None else: row[col] = str(val) result.append(row) return {'rows': result, 'count': len(result), 'columns': df.columns} except Exception as exc: return {'error': f'Export failed: {exc}', 'truncated': False} # ── 22. repair_schema ───────────────────────────────────────────────── async def _handle_repair_schema(dataset_id: str) -> dict: """Attempt to repair schema mismatches by normalising column names.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': f'Dataset not found: {dataset_id}', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) # Normalise: replace hyphens/underscores/dots, lowercase col_map = {} for c in schema: norm = c.lower().replace('-', '_').replace('.', '_') col_map[c] = norm if len(set(col_map.values())) < len(col_map): # Conflicts exist — can't auto-repair return {'status': 'conflict', 'message': 'Normalised names conflict', 'column_map': col_map} # Rename columns renames = {old: new for old, new in col_map.items() if old != new} if renames: lf = lf.rename(renames) new_schema = {col_map.get(c, c): dtype for c, dtype in schema.items()} store.store_dataset(dataset_id, lf, schema=new_schema, metadata=entry.get('metadata', {})) return {'status': 'repaired', 'renamed': len(renames), 'column_map': col_map} return {'status': 'no_change', 'message': 'Column names already normalised'} # ── 23. analyze_patterns ────────────────────────────────────────────── async def _handle_analyze_patterns(dataset_id: str, top_n: int = 10) -> dict: """Read-only: traffic pattern analysis.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': 'Dataset not found', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) n = lf.select(pl.len()).collect(streaming=True).item() df = lf.head(100000).collect(streaming=True) result = {'total_rows': n, 'sampled': min(n, 100000)} # Top source IPs for col_name, label in [(':ips', 'src_ips'), (':ipd', 'dst_ips'), ('snam', 'snis'), ('cnam', 'cert_names')]: if col_name in df.columns: try: top = df[col_name].value_counts(sort=True).head(top_n) result[label] = { 'top': [{'value': str(r[col_name]), 'count': int(r['count'])} for r in top.iter_rows(named=True)], 'unique': int(df[col_name].n_unique()), } except Exception: pass # Port distribution for port_col in [':prs', ':prd']: if port_col in df.columns: try: top = df[port_col].value_counts(sort=True).head(top_n) result[f'{port_col}_dist'] = { 'top': [{'port': int(r[port_col]), 'count': int(r['count'])} for r in top.iter_rows(named=True)], 'unique': int(df[port_col].n_unique()), } except Exception: pass # TLS version breakdown for tls_col in ['0ver']: if tls_col in df.columns: try: vc = df[tls_col].value_counts(sort=True) result['tls_versions'] = [{'version': str(r[tls_col]), 'count': int(r['count'])} for r in vc.iter_rows(named=True)] except Exception: pass return result # ── 24. analyze_temporal ───────────────────────────────────────────── async def _handle_analyze_temporal(dataset_id: str, timestamp_column: str = None) -> dict: """Read-only: temporal traffic pattern analysis.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': 'Dataset not found', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) if not timestamp_column: for cand in ['timestamp', 'time', '8dbd', '8did']: if cand in schema: timestamp_column = cand break if not timestamp_column: return {'error': 'No timestamp column found', 'truncated': False} n = lf.select(pl.len()).collect(streaming=True).item() return {'total_rows': n, 'timestamp_column': timestamp_column, 'note': 'Temporal analysis requires parsed timestamps; raw values shown in explore_distributions'} # ── 25. analyze_tls_health ──────────────────────────────────────────── async def _handle_analyze_tls_health(dataset_id: str) -> dict: """Read-only: TLS security assessment.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': 'Dataset not found', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) n = lf.select(pl.len()).collect(streaming=True).item() df = lf.head(100000).collect(streaming=True) result = {'total_rows': n, 'sampled': min(n, 100000)} if '0ver' in df.columns: vc = df['0ver'].value_counts(sort=True) versions = {} for r in vc.iter_rows(named=True): v = str(r['0ver']).replace(' ', '') if v in ('0303', '0304'): versions['tls_1_2_plus'] = versions.get('tls_1_2_plus', 0) + int(r['count']) else: versions['tls_legacy'] = versions.get('tls_legacy', 0) + int(r['count']) total = sum(versions.values()) result.update({ 'tls_1_2_plus_pct': round(versions.get('tls_1_2_plus', 0) / max(total, 1) * 100, 1), 'tls_legacy_pct': round(versions.get('tls_legacy', 0) / max(total, 1) * 100, 1), 'tls_legacy_versions': [{'version': str(r['0ver']), 'count': int(r['count'])} for r in vc.iter_rows(named=True) if str(r['0ver']).replace(' ', '') not in ('0303', '0304')], }) if 'snam' in df.columns: empty_sni = df['snam'].is_null().sum() + df['snam'].cast(pl.Utf8).is_in(['', 'null', 'None']).sum() result['sni_missing_pct'] = round(float(empty_sni) / max(len(df), 1) * 100, 1) if 'cipher-suite' in df.columns: weak_ciphers = ['rc4', 'cbc', 'null', 'export'] total_ciphers = len(df) weak_count = 0 for w in weak_ciphers: try: weak_count += df['cipher-suite'].str.to_lowercase().str.contains(w).sum() except Exception: pass result['weak_cipher_pct'] = round(float(weak_count) / max(total_ciphers, 1) * 100, 1) if total_ciphers else 0 if result.get('tls_legacy_pct', 0) > 20: result['warning'] = 'High proportion of legacy TLS — possible downgrade attacks' if result.get('sni_missing_pct', 0) > 50: result['warning'] = 'Most flows lack SNI — possible IP-direct connections (proxy/VPN)' return result # ── 26. analyze_geo_distribution ───────────────────────────────────── async def _handle_analyze_geo_distribution(dataset_id: str) -> dict: """Read-only: geographic traffic analysis.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': 'Dataset not found', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) n = lf.select(pl.len()).collect(streaming=True).item() df = lf.head(100000).collect(streaming=True) result = {'total_rows': n, 'sampled': min(n, 100000)} for geo_col, label in [('scnt', 'src_countries'), ('dcnt', 'dst_countries')]: if geo_col in df.columns: try: vc = df[geo_col].value_counts(sort=True).head(15) result[label] = [{'country': str(r[geo_col]), 'count': int(r['count'])} for r in vc.iter_rows(named=True)] result[f'{label}_unique'] = int(df[geo_col].n_unique()) except Exception: pass if 'scnt' in df.columns and 'dcnt' in df.columns: try: same = (df['scnt'] == df['dcnt']).sum() result['domestic_pct'] = round(float(same) / max(len(df), 1) * 100, 1) result['international_pct'] = round(100.0 - result['domestic_pct'], 1) except Exception: pass return result # ── 27. analyze_entity_detail ──────────────────────────────────────── async def _handle_analyze_entity_detail(dataset_id: str, entity_value: str) -> dict: """Read-only: deep-dive into a single entity.""" store = SessionStore() entry = store.get_dataset(dataset_id) if entry is None: return {'error': 'Dataset not found', 'truncated': False} lf = entry['lazyframe'] schema = entry.get('schema', {}) # Find entity column (first non-numeric, non-internal column) numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'} entity_col = None for c in schema: if c.startswith('_') or schema[c].split('(')[0].strip() in numeric_types or '.' in c: continue entity_col = c break if not entity_col: entity_col = list(schema.keys())[0] try: filtered = lf.filter(pl.col(entity_col).cast(pl.Utf8).str.contains(entity_value, literal=True)) n = filtered.select(pl.len()).collect(streaming=True).item() if n == 0: return {'error': f"Entity '{entity_value}' not found in column '{entity_col}'", 'truncated': False} df = filtered.collect(streaming=True) except Exception as exc: return {'error': f'Failed to query: {exc}', 'truncated': False} result = { 'entity_value': entity_value, 'entity_column': entity_col, 'flow_count': n, 'columns': list(df.columns[:15]), } # Numeric summary for c in df.columns: if df[c].dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64): try: result[f'{c}_stats'] = { 'mean': round(float(df[c].mean()), 2), 'min': float(df[c].min()) if df[c].min() is not None else None, 'max': float(df[c].max()) if df[c].max() is not None else None, } except Exception: pass return result # ── 28. compute_distance_matrix ────────────────────────────────────────── def _indent_function_body(body: str, indent: str = ' ') -> str: """Indent each non-empty line of a code block for wrapping in a def.""" return '\n'.join( indent + line if line.strip() else line for line in body.split('\n') ) async def _handle_compute_distance_matrix( dataset_id: str, python_function: str, columns: Optional[list[str]] = None, max_sample: int = 10000, ) -> dict: """Evaluate a user-provided Python function on each row of a dataset. Wraps the function body as ``def distance_fn(row): ``, executes it per-row in a restricted namespace (``math``, ``numpy`` only), and returns a distance-score summary plus a sample of scored rows. """ 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'] # Collect data (streaming for large datasets) try: df = lf.collect(streaming=True) except Exception as e: return {'error': f'Data collection failed: {e}', 'truncated': False} # Determine columns to pass to the function if columns is None: numeric_dtypes = (pl.Int8, pl.Int16, pl.Int32, pl.Int64, pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64, pl.Float32, pl.Float64) columns = [c for c in df.columns if df[c].dtype in numeric_dtypes] # Filter to columns that actually exist available = set(df.columns) columns = [c for c in columns if c in available] if not columns: return {'error': 'No valid columns found in dataset', 'truncated': False} n_total = len(df) # Sample if larger than max_sample if n_total > max_sample: df_sampled = df.sample(n=max_sample, seed=42) n_evaluated = max_sample else: df_sampled = df n_evaluated = n_total # Build list-of-dicts for the selected columns rows = df_sampled.select(columns).to_dicts() # ── Safe execution ──────────────────────────────────────────────────── safe_builtins = { 'abs': abs, 'len': len, 'range': range, 'min': min, 'max': max, 'sum': sum, 'round': round, 'int': int, 'float': float, 'str': str, 'bool': bool, 'list': list, 'dict': dict, 'tuple': tuple, 'isinstance': isinstance, 'type': type, 'True': True, 'False': False, 'None': None, 'enumerate': enumerate, 'zip': zip, 'map': map, 'filter': filter, 'sorted': sorted, 'reversed': reversed, } safe_ns = { '__builtins__': safe_builtins, 'math': __import__('math'), 'np': __import__('numpy'), } # Wrap the user's function body in a def statement body = python_function.strip() if body.startswith('def '): # User provided the full function definition full_code = body else: indented = _indent_function_body(body) full_code = f'def distance_fn(row):\n{indented}' try: exec(full_code, safe_ns) # Resolve: the function name might be 'distance_fn' or whatever the user named it distance_fn = safe_ns.get('distance_fn') if distance_fn is None: # Fallback: try to find any callable in safe_ns for _name, _obj in safe_ns.items(): if callable(_obj) and _name not in ('math', 'np', '__builtins__'): distance_fn = _obj break if distance_fn is None: return {'error': 'Could not locate the compiled function', 'truncated': False} except SyntaxError as e: return {'error': f'Function syntax error: {e}', 'truncated': False} except Exception as e: return {'error': f'Function compilation failed: {e}', 'truncated': False} # ── Apply function per row ──────────────────────────────────────────── scores: list[float | None] = [] errors = 0 for row in rows: try: score = distance_fn(row) if isinstance(score, (int, float)): scores.append(float(score)) else: scores.append(None) errors += 1 except Exception: scores.append(None) errors += 1 valid_scores = [s for s in scores if s is not None] if not valid_scores: return {'error': 'No valid scores computed — all rows raised errors', 'truncated': False} # ── Distance summary ────────────────────────────────────────────────── arr = np.array(valid_scores) summary = { 'min': float(np.min(arr)), 'max': float(np.max(arr)), 'mean': float(np.mean(arr)), 'std': float(np.std(arr)), 'count': len(arr), 'errors': errors, 'total_rows_in_sample': len(rows), 'matrix_shape': [len(rows), len(columns)], } # Sample of scored rows (first 30, keep score + first 3 values) scored_sample = [] for i, row in enumerate(rows[:30]): if i < len(scores): scored_sample.append({ 'index': i, 'values': {c: row.get(c) for c in columns[:3]}, 'score': scores[i], }) return { 'summary': summary, 'scored_sample': scored_sample, 'columns_used': columns, 'total_rows_in_dataset': n_total, 'rows_evaluated': n_evaluated, }