refactor: split tool_registry.py (3368 lines) into analysis/tools/ package
- 15 sub-modules: _helpers, _registry, _dispatch, load_data, profile, filter, preprocess, clustering, evaluate, features, export, entities, anomalies, diagnostics, analysis, data_mgmt, distance_matrix - Extracted shared helpers: _resolve_dataset(), _count_rows(), _get_numeric_columns(), NUMERIC_DTYPES, NUMERIC_TYPE_NAMES - Extracted _cluster_core() for shared clustering pipeline - Thin tool_registry.py wrapper preserves backward compatibility - Old file preserved as tool_registry_backup.py - manage.py check passes (0 issues)
This commit is contained in:
PM-pinou
+4
-3366
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,66 @@
|
||||
"""Tool registry package — split from analysis/tool_registry.py.
|
||||
|
||||
Re-exports all tool metadata, dispatcher, and handlers.
|
||||
"""
|
||||
|
||||
# ── Helpers ────────────────────────────────────────────────────────────────
|
||||
from ._helpers import (
|
||||
_json_size,
|
||||
_truncate_response,
|
||||
_build_filter_expr,
|
||||
_resolve_dataset,
|
||||
_count_rows,
|
||||
_get_numeric_columns,
|
||||
NUMERIC_DTYPES,
|
||||
NUMERIC_TYPE_NAMES,
|
||||
)
|
||||
|
||||
# ── Registry + dispatch ────────────────────────────────────────────────────
|
||||
from ._registry import get_tools_meta
|
||||
from ._dispatch import handle_call
|
||||
|
||||
# ── Handlers (tool implementations) ────────────────────────────────────────
|
||||
from .load_data import _handle_load_data
|
||||
from .profile import _handle_profile_data
|
||||
from .filter import _handle_filter_data
|
||||
from .preprocess import _handle_preprocess_data
|
||||
from .clustering import _handle_run_clustering, _handle_filter_and_cluster
|
||||
from .evaluate import _handle_evaluate_clustering
|
||||
from .features import _handle_extract_features, _save_features_to_db
|
||||
from .export import _handle_export_results
|
||||
from .entities import (
|
||||
_handle_build_entity_profiles,
|
||||
_handle_compute_scores,
|
||||
_add_adaptive_scores,
|
||||
_ENTITY_KEYWORDS,
|
||||
_NUMERIC_TYPES_FOR_AGG,
|
||||
_ANOMALY_FEATURE_COLS,
|
||||
)
|
||||
from .anomalies import _handle_detect_anomalies, _handle_visualize_anomalies
|
||||
from .diagnostics import (
|
||||
_handle_validate_data,
|
||||
_handle_explore_distributions,
|
||||
_handle_find_outliers,
|
||||
_handle_diagnose_clustering,
|
||||
_handle_compare_datasets,
|
||||
_handle_export_debug_sample,
|
||||
_handle_repair_schema,
|
||||
)
|
||||
from .analysis import (
|
||||
_handle_analyze_patterns,
|
||||
_handle_analyze_temporal,
|
||||
_handle_analyze_fft,
|
||||
_handle_analyze_tls_health,
|
||||
_handle_analyze_geo_distribution,
|
||||
_handle_analyze_entity_detail,
|
||||
)
|
||||
from .data_mgmt import (
|
||||
_handle_list_datasets,
|
||||
_handle_drop_dataset,
|
||||
_handle_clone_dataset,
|
||||
)
|
||||
from .distance_matrix import _handle_compute_distance_matrix
|
||||
|
||||
# Auto-profile all public functions in this module
|
||||
from analysis.profile_util import auto_profile_module # noqa: E402
|
||||
auto_profile_module(__name__)
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Tool call dispatcher — maps tool names to handler functions."""
|
||||
from ._helpers import _truncate_response
|
||||
from .load_data import _handle_load_data
|
||||
from .profile import _handle_profile_data
|
||||
from .filter import _handle_filter_data
|
||||
from .preprocess import _handle_preprocess_data
|
||||
from .clustering import _handle_run_clustering, _handle_filter_and_cluster
|
||||
from .evaluate import _handle_evaluate_clustering
|
||||
from .features import _handle_extract_features
|
||||
from .export import _handle_export_results
|
||||
from .entities import _handle_build_entity_profiles, _handle_compute_scores
|
||||
from .anomalies import _handle_detect_anomalies, _handle_visualize_anomalies
|
||||
from .diagnostics import (
|
||||
_handle_validate_data, _handle_explore_distributions, _handle_find_outliers,
|
||||
_handle_diagnose_clustering, _handle_compare_datasets,
|
||||
_handle_export_debug_sample, _handle_repair_schema,
|
||||
)
|
||||
from .analysis import (
|
||||
_handle_analyze_patterns, _handle_analyze_temporal, _handle_analyze_fft,
|
||||
_handle_analyze_tls_health, _handle_analyze_geo_distribution,
|
||||
_handle_analyze_entity_detail,
|
||||
)
|
||||
from .data_mgmt import _handle_list_datasets, _handle_drop_dataset, _handle_clone_dataset
|
||||
from .distance_matrix import _handle_compute_distance_matrix
|
||||
|
||||
|
||||
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 30 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,
|
||||
'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_fft': _handle_analyze_fft,
|
||||
'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)
|
||||
@@ -0,0 +1,134 @@
|
||||
"""Shared utilities for tool handlers.
|
||||
|
||||
Provides:
|
||||
- NUMERIC_DTYPES / NUMERIC_TYPE_NAMES: shared constants for dtype filtering
|
||||
- _json_size, _truncate_response: LLM context-window helpers
|
||||
- _build_filter_expr: Polars filter expression builder
|
||||
- _resolve_dataset: one-liner SessionStore lookup (replaces 30+ repeated blocks)
|
||||
- _count_rows: fast streaming row count
|
||||
- _get_numeric_columns: extract numeric column names from schema
|
||||
"""
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Shared constants
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
NUMERIC_DTYPES = (
|
||||
pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
||||
pl.Float32, pl.Float64,
|
||||
)
|
||||
|
||||
NUMERIC_TYPE_NAMES = {
|
||||
'Int8', 'Int16', 'Int32', 'Int64',
|
||||
'UInt8', 'UInt16', 'UInt32', 'UInt64',
|
||||
'Float32', 'Float64',
|
||||
}
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# LLM response helpers
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
_MAX_RESPONSE_KB = 64
|
||||
|
||||
|
||||
def _json_size(obj: Any) -> int:
|
||||
"""Estimate JSON byte size of *obj*."""
|
||||
return len(json.dumps(obj, default=str, ensure_ascii=False))
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Filter expression builder
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
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
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# SessionStore helpers
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
def _resolve_dataset(dataset_id: str) -> tuple:
|
||||
"""Resolve a dataset from SessionStore.
|
||||
|
||||
Returns:
|
||||
(entry, None) on success, where *entry* is the SessionStore entry dict.
|
||||
(None, error_dict) if the dataset is not found.
|
||||
"""
|
||||
store = SessionStore()
|
||||
entry = store.get_dataset(dataset_id)
|
||||
if entry is None:
|
||||
return None, {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||||
return entry, None
|
||||
|
||||
|
||||
def _count_rows(lf: pl.LazyFrame) -> int:
|
||||
"""Count rows in a LazyFrame via fast streaming count."""
|
||||
try:
|
||||
return lf.select(pl.len()).collect(streaming=True).item()
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
def _get_numeric_columns(lf: pl.LazyFrame, schema: dict, max_cols: int = 10) -> list[str]:
|
||||
"""Get numeric column names from a dataset.
|
||||
|
||||
Prefers *schema* (string-keyed); falls back to LazyFrame collect_schema().
|
||||
"""
|
||||
if schema:
|
||||
return [c for c in schema
|
||||
if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES
|
||||
and not c.startswith('_')][:max_cols]
|
||||
# Fallback to LazyFrame schema
|
||||
available = lf.collect_schema()
|
||||
return [available.names()[i] for i, dt in enumerate(available.dtypes())
|
||||
if dt in NUMERIC_DTYPES and not available.names()[i].startswith('_')][:max_cols]
|
||||
@@ -0,0 +1,754 @@
|
||||
"""Tool metadata definitions for the MCP server.
|
||||
|
||||
All 30 Tool(...) descriptor objects and the get_tools_meta() function.
|
||||
"""
|
||||
from mcp.types import Tool
|
||||
|
||||
|
||||
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": ["agglomerative", "hdbscan", "kmeans"],
|
||||
"description": "Clustering algorithm (default: agglomerative/Ward)",
|
||||
"default": "agglomerative",
|
||||
},
|
||||
"params": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Algorithm parameters. Agglomerative: n_clusters (5), metric ('euclidean'), "
|
||||
"linkage ('ward'). 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": ["agglomerative", "hdbscan", "kmeans"],
|
||||
"description": "Clustering algorithm (default: agglomerative/Ward)",
|
||||
"default": "agglomerative",
|
||||
},
|
||||
"params": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Algorithm parameters. Agglomerative: n_clusters (5), metric ('euclidean'), "
|
||||
"linkage ('ward'). HDBSCAN: min_cluster_size (5), "
|
||||
"min_samples (None), metric ('euclidean'). "
|
||||
"KMeans: n_clusters (3), random_state (42)."
|
||||
),
|
||||
},
|
||||
},
|
||||
"required": ["dataset_id", "filters"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="build_entity_profiles",
|
||||
description=(
|
||||
"按实体列（源IP、SNI等）分组原始流数据。自动检测实体列（如 src_ip, "
|
||||
"dst_ip, sni, host, mac, domain），按列分组并计算每个实体的聚合特征："
|
||||
"流数量、唯一端口数、TLS版本分布、字节总量、唯一目标数。"
|
||||
"生成的数据集每个实体一行。结果存储在会话中，可供后续 compute_scores 使用。"
|
||||
),
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"dataset_id": {
|
||||
"type": "string",
|
||||
"description": "Dataset ID from load_data or profile_data",
|
||||
},
|
||||
"entity_column": {
|
||||
"type": "string",
|
||||
"description": "Specific entity column to group by (e.g. 'src_ip')",
|
||||
},
|
||||
"entity_columns": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Multiple entity columns to group by",
|
||||
},
|
||||
"auto_detect": {
|
||||
"type": "boolean",
|
||||
"description": "Auto-detect entity columns (IP, SNI, host, MAC, domain)",
|
||||
"default": True,
|
||||
},
|
||||
},
|
||||
"required": ["dataset_id"],
|
||||
},
|
||||
),
|
||||
Tool(
|
||||
name="compute_scores",
|
||||
description=(
|
||||
"在实体数据上计算自适应代理/异常/威胁评分。从数据自动学习"
|
||||
"特征权重，计算加权代理评分（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=(
|
||||
"为高风险实体生成UMAP-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=(
|
||||
"诊断工具：当聚类产生太少/太多聚类、全部为噪声或轮廓系数低时——诊断原因。"
|
||||
"检查：特征方差、相关矩阵、特征降维分析、数据稀疏性，并建议参数调整"
|
||||
"（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_fft",
|
||||
description=(
|
||||
"FFT频谱分析——从时间戳列提取周期模式和频率特征。适用于检测时间周期性"
|
||||
"（如每日/每周流量模式）。返回频率、幅度、top周期、频谱质心、周期性评分。"
|
||||
"只读，无副作用。"
|
||||
),
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"dataset_id": {
|
||||
"type": "string",
|
||||
"description": "Dataset ID from load_data",
|
||||
},
|
||||
"timestamp_column": {
|
||||
"type": "string",
|
||||
"description": "Column name containing timestamps or numeric time values",
|
||||
},
|
||||
"top_n": {
|
||||
"type": "integer",
|
||||
"description": "Number of top frequency peaks to return",
|
||||
"default": 5,
|
||||
},
|
||||
},
|
||||
"required": ["dataset_id", "timestamp_column"],
|
||||
},
|
||||
),
|
||||
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"],
|
||||
},
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,326 @@
|
||||
"""Read-only analysis handlers — patterns, temporal, FFT, TLS health, geo, entity detail."""
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore
|
||||
from ._helpers import _resolve_dataset, _count_rows, NUMERIC_TYPE_NAMES
|
||||
|
||||
|
||||
async def _handle_analyze_patterns(dataset_id: str, top_n: int = 10) -> dict:
|
||||
"""Read-only: traffic pattern analysis."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
lf = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
n = _count_rows(lf)
|
||||
df = lf.head(100000).collect(streaming=True)
|
||||
|
||||
result = {'total_rows': n, 'sampled': min(n, 100000)}
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
|
||||
async def _handle_analyze_temporal(dataset_id: str, timestamp_column: str = None) -> dict:
|
||||
"""Read-only: temporal traffic pattern analysis."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
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 = _count_rows(lf)
|
||||
return {'total_rows': n, 'timestamp_column': timestamp_column,
|
||||
'note': 'Temporal analysis requires parsed timestamps; raw values shown in explore_distributions'}
|
||||
|
||||
|
||||
async def _handle_analyze_fft(
|
||||
dataset_id: str,
|
||||
timestamp_column: str,
|
||||
top_n: int = 5,
|
||||
) -> dict:
|
||||
"""FFT spectrum analysis on a timestamp column to detect periodic patterns."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
|
||||
col = None
|
||||
if timestamp_column in schema:
|
||||
col = timestamp_column
|
||||
else:
|
||||
col_lower = timestamp_column.lower()
|
||||
for c in schema:
|
||||
if c.lower() == col_lower:
|
||||
col = c
|
||||
break
|
||||
if col is None:
|
||||
return {'error': f"Timestamp column '{timestamp_column}' not found in dataset. Available: {list(schema.keys())[:20]}",
|
||||
'truncated': False}
|
||||
|
||||
df = lf.select(pl.col(col).alias('__ts__')).collect(streaming=True)
|
||||
n_raw = len(df)
|
||||
if n_raw < 4:
|
||||
return {'error': f'Need at least 4 samples for FFT; got {n_raw}', 'truncated': False}
|
||||
|
||||
ts_series = df['__ts__']
|
||||
dtype = ts_series.dtype
|
||||
str_dtypes = (pl.Utf8, pl.String)
|
||||
|
||||
if dtype in str_dtypes:
|
||||
try:
|
||||
ts_parsed = ts_series.str.strptime(pl.Datetime, format=None, strict=False)
|
||||
ts_numeric = ts_parsed.cast(pl.Int64) / 1_000_000_000.0
|
||||
except Exception:
|
||||
try:
|
||||
ts_numeric = ts_series.cast(pl.Float64)
|
||||
except Exception as e:
|
||||
return {'error': f'Cannot convert string column to numeric or datetime: {e}',
|
||||
'truncated': False}
|
||||
elif dtype in (pl.Datetime, pl.Date):
|
||||
ts_numeric = ts_series.cast(pl.Int64) / 1_000_000_000.0
|
||||
else:
|
||||
try:
|
||||
ts_numeric = ts_series.cast(pl.Float64)
|
||||
except Exception as e:
|
||||
return {'error': f'Cannot cast column to float: {e}', 'truncated': False}
|
||||
|
||||
df_sorted = pl.DataFrame({'__val__': ts_numeric}).drop_nulls().sort('__val__')
|
||||
n = len(df_sorted)
|
||||
if n < 4:
|
||||
return {'error': f'After removing nulls, only {n} samples remain (need ≥4)', 'truncated': False}
|
||||
|
||||
values = df_sorted['__val__'].to_numpy().astype(np.float64)
|
||||
|
||||
values_centered = values - np.mean(values)
|
||||
fft_complex = np.fft.fft(values_centered)
|
||||
fft_magnitudes = np.abs(fft_complex)
|
||||
freqs = np.fft.fftfreq(n)
|
||||
|
||||
pos_mask = freqs > 0
|
||||
freqs_pos = freqs[pos_mask]
|
||||
mags_pos = fft_magnitudes[pos_mask]
|
||||
|
||||
if len(freqs_pos) == 0:
|
||||
return {'error': 'Insufficient positive frequencies for FFT analysis', 'truncated': False}
|
||||
|
||||
peak_indices = np.argsort(mags_pos)[::-1][:min(top_n, len(mags_pos))]
|
||||
top_periods = []
|
||||
for idx in peak_indices:
|
||||
freq = float(freqs_pos[idx])
|
||||
mag = float(mags_pos[idx])
|
||||
period = float(1.0 / freq) if freq > 1e-10 else float('inf')
|
||||
if n > 1:
|
||||
time_span = float(values[-1] - values[0])
|
||||
sample_period = time_span / max(n - 1, 1)
|
||||
time_period = period * sample_period
|
||||
else:
|
||||
time_period = float('inf')
|
||||
top_periods.append({
|
||||
'period': round(period, 2),
|
||||
'magnitude': round(mag, 4),
|
||||
'frequency': round(freq, 6),
|
||||
'time_domain_period_seconds': round(time_period, 2) if time_period != float('inf') else None,
|
||||
})
|
||||
|
||||
mag_sum = float(np.sum(mags_pos))
|
||||
if mag_sum > 0:
|
||||
spectral_centroid = float(np.sum(freqs_pos * mags_pos) / mag_sum)
|
||||
else:
|
||||
spectral_centroid = 0.0
|
||||
|
||||
mean_mag = float(np.mean(mags_pos)) if len(mags_pos) > 0 else 1.0
|
||||
max_mag = float(mags_pos[peak_indices[0]]) if len(peak_indices) > 0 else 0.0
|
||||
periodicity_score = round(max_mag / max(mean_mag, 1e-10), 4)
|
||||
|
||||
decimate_step = max(1, len(freqs_pos) // 200)
|
||||
freqs_decimated = freqs_pos[::decimate_step].tolist()
|
||||
mags_decimated = mags_pos[::decimate_step].tolist()
|
||||
|
||||
return {
|
||||
'frequencies': [round(float(f), 6) for f in freqs_decimated],
|
||||
'magnitudes': [round(float(m), 4) for m in mags_decimated],
|
||||
'top_periods': top_periods,
|
||||
'spectral_centroid': round(spectral_centroid, 6),
|
||||
'periodicity_score': periodicity_score,
|
||||
'n_samples': n,
|
||||
'n_raw': n_raw,
|
||||
'nan_dropped': n_raw - n,
|
||||
'timestamp_column': col,
|
||||
}
|
||||
|
||||
|
||||
async def _handle_analyze_tls_health(dataset_id: str) -> dict:
|
||||
"""Read-only: TLS security assessment."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
lf = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
n = _count_rows(lf)
|
||||
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
|
||||
|
||||
|
||||
async def _handle_analyze_geo_distribution(dataset_id: str) -> dict:
|
||||
"""Read-only: geographic traffic analysis."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
lf = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
n = _count_rows(lf)
|
||||
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
|
||||
|
||||
|
||||
async def _handle_analyze_entity_detail(dataset_id: str, entity_value: str) -> dict:
|
||||
"""Read-only: deep-dive into a single entity."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
lf = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
|
||||
entity_col = None
|
||||
for c in schema:
|
||||
if c.startswith('_') or schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES 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 = _count_rows(filtered)
|
||||
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:
|
||||
raise
|
||||
|
||||
result = {
|
||||
'entity_value': entity_value,
|
||||
'entity_column': entity_col,
|
||||
'flow_count': n,
|
||||
'columns': list(df.columns[:15]),
|
||||
}
|
||||
|
||||
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
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Anomaly detection + visualization handlers."""
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore
|
||||
from ._helpers import _resolve_dataset, NUMERIC_TYPE_NAMES
|
||||
|
||||
|
||||
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
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
num_cols = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES and not c.startswith('_')]
|
||||
if not num_cols:
|
||||
return {'error': 'No numeric columns for anomaly detection', 'truncated': False}
|
||||
|
||||
df = lf.select(num_cols).collect(streaming=True)
|
||||
mat = df.to_numpy()
|
||||
mat = np.nan_to_num(mat, nan=0.0)
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
async def _handle_visualize_anomalies(dataset_id: str) -> dict:
|
||||
"""Generate UMAP-2D embedding for anomaly visualization."""
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
import umap
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
num_cols = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES and not c.startswith('_')]
|
||||
if len(num_cols) < 2:
|
||||
return {'error': 'Need at least 2 numeric columns for UMAP', 'truncated': False}
|
||||
|
||||
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)
|
||||
reducer = umap.UMAP(n_components=2, random_state=42,
|
||||
n_neighbors=15, min_dist=0.1, metric='euclidean')
|
||||
coords = reducer.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],
|
||||
'dataset_id': dataset_id,
|
||||
}
|
||||
@@ -0,0 +1,462 @@
|
||||
"""Clustering handlers — HDBSCAN, KMeans, AgglomerativeClustering.
|
||||
|
||||
Contains:
|
||||
- _handle_run_clustering: cluster selected columns
|
||||
- _handle_filter_and_cluster: filter then auto-cluster
|
||||
- _cluster_core: shared clustering pipeline extracted from both handlers
|
||||
"""
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore, _generate_id
|
||||
from ._helpers import (
|
||||
NUMERIC_DTYPES,
|
||||
_build_filter_expr,
|
||||
_resolve_dataset,
|
||||
_count_rows,
|
||||
)
|
||||
|
||||
_clust_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _cluster_core(
|
||||
data: np.ndarray,
|
||||
feature_cols: list[str],
|
||||
algorithm: str,
|
||||
params: dict,
|
||||
random_state: int = 42,
|
||||
):
|
||||
"""Shared clustering pipeline: scale, filter, reduce, fit, score.
|
||||
|
||||
Args:
|
||||
data: Numpy array (n_samples, n_features).
|
||||
feature_cols: Column names matching data columns.
|
||||
algorithm: 'hdbscan', 'kmeans', or 'agglomerative'.
|
||||
params: Algorithm-specific parameters dict.
|
||||
random_state: Random seed.
|
||||
|
||||
Returns:
|
||||
(labels_array, n_clusters_found, noise_count, quality_metrics, _ds_idx)
|
||||
"""
|
||||
p = dict(params)
|
||||
|
||||
# ── Low memory check ───────────────────────────────────────────────
|
||||
low_memory = False
|
||||
try:
|
||||
import psutil
|
||||
mem = psutil.virtual_memory()
|
||||
low_memory = mem.available < 2 * 1024 ** 3 # 2 GB
|
||||
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]
|
||||
|
||||
# ── Downsampling 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]
|
||||
|
||||
# Log clustering input
|
||||
_clust_logger.info(
|
||||
'[CLUSTER_INPUT] rows=%d cols=%d feature_cols=%s',
|
||||
len(data), data.shape[1] if len(data.shape) > 1 else 1, feature_cols,
|
||||
)
|
||||
|
||||
# Standard scale
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
_clust_logger.info('[CLUSTER_SCALE] Starting StandardScaler...')
|
||||
data_scaled = StandardScaler().fit_transform(data)
|
||||
_clust_logger.info('[CLUSTER_SCALE] Done. shape=%s', data_scaled.shape)
|
||||
|
||||
# ── 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]
|
||||
|
||||
# ── SVD dimensionality reduction (if >50 features) ─────────────────
|
||||
n_features = data_scaled.shape[1]
|
||||
if n_features > 50:
|
||||
from sklearn.decomposition import TruncatedSVD
|
||||
svd = TruncatedSVD(n_components=50, random_state=random_state)
|
||||
data_scaled = svd.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),
|
||||
)
|
||||
|
||||
# ── Model fit ──────────────────────────────────────────────────────
|
||||
from sklearn.cluster import HDBSCAN, MiniBatchKMeans, AgglomerativeClustering
|
||||
if algorithm == 'kmeans':
|
||||
n_clusters_param = p.get('n_clusters', 3)
|
||||
model = MiniBatchKMeans(
|
||||
n_clusters=n_clusters_param, random_state=random_state,
|
||||
batch_size=1024, n_init='auto',
|
||||
)
|
||||
elif algorithm in ('agglomerative', 'ward'):
|
||||
n_clusters_param = p.get('n_clusters', 5)
|
||||
model = AgglomerativeClustering(
|
||||
n_clusters=n_clusters_param,
|
||||
metric=p.get('metric', 'euclidean'),
|
||||
linkage=p.get('linkage', 'ward'),
|
||||
distance_threshold=None,
|
||||
)
|
||||
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=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
|
||||
|
||||
return labels_array, n_clusters_found, noise_count, quality, _ds_idx
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# _handle_run_clustering
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
async def _handle_run_clustering(
|
||||
dataset_id: str,
|
||||
cluster_columns: list[str],
|
||||
algorithm: str = 'agglomerative',
|
||||
params: Optional[dict] = None,
|
||||
random_state: int = 42,
|
||||
) -> dict:
|
||||
"""Run HDBSCAN, KMeans, or Agglomerative clustering on selected columns."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
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()
|
||||
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:
|
||||
# Fallback: try coercing Utf8 columns to Float64
|
||||
utf8_cols = [
|
||||
available_names[i] for i, dt in enumerate(available_dtypes)
|
||||
if dt == pl.Utf8 and not available_names[i].startswith('_')
|
||||
][:10]
|
||||
if utf8_cols:
|
||||
for col in utf8_cols:
|
||||
try:
|
||||
casted = lf.select(
|
||||
pl.col(col).cast(pl.Float64, strict=False).alias(col)
|
||||
).collect(streaming=True)
|
||||
null_ratio = casted[col].is_null().sum() / max(len(casted), 1)
|
||||
if null_ratio < 0.9:
|
||||
lf = lf.with_columns(
|
||||
pl.col(col).cast(pl.Float64, strict=False)
|
||||
)
|
||||
feature_cols.append(col)
|
||||
except Exception:
|
||||
pass
|
||||
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
|
||||
row_count = _count_rows(lf)
|
||||
if row_count > 100000:
|
||||
lf = lf.collect(streaming=True).sample(n=100000, seed=42).lazy()
|
||||
|
||||
df_features = 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',
|
||||
'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 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,
|
||||
}
|
||||
|
||||
labels_array, n_found, noise_count, quality, _ds_idx = _cluster_core(
|
||||
data, feature_cols, algorithm, p, random_state,
|
||||
)
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
result_id = _generate_id('clust')
|
||||
store = SessionStore()
|
||||
store.store_cluster_result(
|
||||
result_id=result_id,
|
||||
labels=labels_array.tolist(),
|
||||
n_clusters=n_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_found,
|
||||
'n_noise': noise_count,
|
||||
'labels_sample': labels_array[:50].tolist(),
|
||||
'quality_metrics': quality,
|
||||
}
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# _handle_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."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
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
|
||||
|
||||
try:
|
||||
row_count = _count_rows(filtered_lf)
|
||||
except Exception:
|
||||
row_count = None
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
# ── Step 2: Auto-detect numeric feature columns ───────────────────────
|
||||
try:
|
||||
available = filtered_lf.collect_schema()
|
||||
available_names = available.names()
|
||||
available_dtypes = available.dtypes()
|
||||
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:
|
||||
raise
|
||||
|
||||
# ── Step 3: Store filtered dataset ────────────────────────────────────
|
||||
filtered_dataset_id = _generate_id('flt')
|
||||
store = SessionStore()
|
||||
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 ────────────────────────────────────────────
|
||||
p = params or {}
|
||||
try:
|
||||
total_count = _count_rows(filtered_lf)
|
||||
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)
|
||||
|
||||
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()
|
||||
|
||||
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)
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
labels_array, n_found, noise_count, quality, _ds_idx = _cluster_core(
|
||||
data, feature_cols, algorithm, p, 42,
|
||||
)
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
result_id = _generate_id('clust')
|
||||
store.store_cluster_result(
|
||||
result_id=result_id,
|
||||
labels=labels_array.tolist(),
|
||||
n_clusters=n_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_found,
|
||||
'n_noise': noise_count,
|
||||
'labels_sample': labels_array[:50].tolist(),
|
||||
'quality_metrics': quality,
|
||||
'feature_columns': feature_cols,
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Data management handlers — list, drop, clone datasets."""
|
||||
from typing import Optional
|
||||
|
||||
from ..session_store import SessionStore
|
||||
|
||||
|
||||
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)}
|
||||
|
||||
|
||||
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}
|
||||
|
||||
|
||||
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}
|
||||
@@ -0,0 +1,253 @@
|
||||
"""Diagnostic tool handlers — validate, explore, diagnose, compare, repair."""
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore
|
||||
from ._helpers import _resolve_dataset, _count_rows, NUMERIC_TYPE_NAMES
|
||||
|
||||
|
||||
async def _handle_validate_data(dataset_id: str, fix_issues: bool = False) -> dict:
|
||||
"""Check data quality: schema consistency, missing values, type conflicts."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
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
|
||||
except Exception as exc:
|
||||
raise
|
||||
total_rows = _count_rows(lf)
|
||||
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',
|
||||
}
|
||||
|
||||
|
||||
async def _handle_explore_distributions(dataset_id: str, columns: list[str] = None,
|
||||
max_sample: int = 10000) -> dict:
|
||||
"""Per-column distribution statistics."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
lf = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
|
||||
if not columns:
|
||||
columns = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES][:10]
|
||||
if len(columns) > 20:
|
||||
columns = columns[:20]
|
||||
|
||||
try:
|
||||
n = _count_rows(lf)
|
||||
if n > max_sample:
|
||||
sample_lf = lf.select(columns).sample(n=max_sample, seed=42)
|
||||
else:
|
||||
sample_lf = lf.select(columns)
|
||||
df = sample_lf.collect(streaming=True)
|
||||
except Exception as exc:
|
||||
return {'error': f'explore_distributions failed: {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)}
|
||||
|
||||
|
||||
async def _handle_find_outliers(dataset_id: str, columns: list[str] = None,
|
||||
threshold: float = 3.0) -> dict:
|
||||
"""Find statistical outliers using IQR method."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
lf = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
if not columns:
|
||||
columns = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES]
|
||||
try:
|
||||
n = _count_rows(lf)
|
||||
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:
|
||||
raise
|
||||
|
||||
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}
|
||||
|
||||
|
||||
async def _handle_diagnose_clustering(dataset_id: str, cluster_result_id: str) -> dict:
|
||||
"""Diagnose clustering quality issues."""
|
||||
from sklearn.decomposition import TruncatedSVD
|
||||
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', {})
|
||||
num_cols = [c for c in schema if schema[c].split('(')[0].strip() in NUMERIC_TYPE_NAMES 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')
|
||||
|
||||
try:
|
||||
df_sample = lf.select(num_cols).head(5000).collect(streaming=True)
|
||||
mat = df_sample.to_numpy()
|
||||
n_comp = min(mat.shape)
|
||||
svd = TruncatedSVD(n_components=n_comp, random_state=42).fit(mat)
|
||||
var_explained = svd.explained_variance_ratio_
|
||||
if var_explained[0] > 0.95:
|
||||
diagnosis.append(f'First SVD 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',
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
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 = _count_rows(lf)
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
async def _handle_export_debug_sample(dataset_id: str, max_rows: int = 50) -> dict:
|
||||
"""Export raw data sample as JSON for debugging."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
try:
|
||||
df = entry['lazyframe'].head(max_rows).collect(streaming=True)
|
||||
rows = df.to_dict(as_series=False)
|
||||
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:
|
||||
raise
|
||||
|
||||
|
||||
async def _handle_repair_schema(dataset_id: str) -> dict:
|
||||
"""Attempt to repair schema mismatches by normalising column names."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
lf = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
col_map = {}
|
||||
for c in schema:
|
||||
norm = c.lower().replace('-', '_').replace('.', '_')
|
||||
col_map[c] = norm
|
||||
if len(set(col_map.values())) < len(col_map):
|
||||
return {'status': 'conflict', 'message': 'Normalised names conflict', 'column_map': col_map}
|
||||
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 = SessionStore()
|
||||
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'}
|
||||
@@ -0,0 +1,141 @@
|
||||
"""compute_distance_matrix handler — LLM-driven per-row evaluation."""
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore
|
||||
from ._helpers import _resolve_dataset, NUMERIC_DTYPES
|
||||
|
||||
|
||||
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."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
df = lf.collect(streaming=True)
|
||||
|
||||
if columns is None:
|
||||
columns = [c for c in df.columns
|
||||
if df[c].dtype in NUMERIC_DTYPES]
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
rows = df_sampled.select(columns).to_dicts()
|
||||
|
||||
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'),
|
||||
}
|
||||
|
||||
body = python_function.strip()
|
||||
if body.startswith('def '):
|
||||
full_code = body
|
||||
else:
|
||||
indented = _indent_function_body(body)
|
||||
full_code = f'def distance_fn(row):\n{indented}'
|
||||
|
||||
try:
|
||||
exec(full_code, safe_ns)
|
||||
distance_fn = safe_ns.get('distance_fn')
|
||||
if distance_fn is None:
|
||||
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}
|
||||
|
||||
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}
|
||||
|
||||
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)],
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
"""Entity profile building + adaptive scoring handlers."""
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore, _generate_id
|
||||
from ..profile_util import profile
|
||||
from ._helpers import _resolve_dataset, _count_rows, NUMERIC_DTYPES
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Constants
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
_ENTITY_KEYWORDS = [
|
||||
'src_ip', 'dst_ip', 'source_ip', 'destination_ip', 'src_addr', 'dst_addr',
|
||||
'sip', 'dip', 'ip_src', 'ip_dst', 'saddr', 'daddr',
|
||||
'sni', 'host', 'domain', 'uri', 'url', 'fqdn', 'server_name',
|
||||
'mac', 'mac_addr', 'src_mac', 'dst_mac',
|
||||
'user_agent', 'ua', 'asn', 'organization', 'org',
|
||||
]
|
||||
|
||||
_NUMERIC_TYPES_FOR_AGG = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
||||
pl.Float32, pl.Float64)
|
||||
|
||||
_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',
|
||||
]
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# _handle_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:
|
||||
"""Group raw flow data by entity column(s) and compute per-entity aggregate features."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
schema = entry.get('schema', {})
|
||||
col_names = list(schema.keys())
|
||||
|
||||
# ── Resolve entity columns ──────────────────────────────────────────
|
||||
if entity_column and entity_column in col_names:
|
||||
group_cols = [entity_column]
|
||||
elif entity_columns:
|
||||
group_cols = [c for c in entity_columns if c in col_names]
|
||||
if not group_cols:
|
||||
return {'error': f'None of entity_columns {entity_columns} exist in dataset. Available: {col_names[:20]}',
|
||||
'truncated': False}
|
||||
elif auto_detect:
|
||||
matched = []
|
||||
for kw in _ENTITY_KEYWORDS:
|
||||
found = [c for c in col_names if kw.lower() in c.lower()]
|
||||
matched.extend(found)
|
||||
seen = set()
|
||||
group_cols = [c for c in matched if not (c in seen or seen.add(c))]
|
||||
if not group_cols:
|
||||
for c in col_names:
|
||||
if c.startswith('_'):
|
||||
continue
|
||||
group_cols = [c]
|
||||
break
|
||||
group_cols = group_cols[:3]
|
||||
else:
|
||||
return {'error': 'No entity column specified and auto_detect is disabled',
|
||||
'truncated': False}
|
||||
|
||||
if not group_cols:
|
||||
return {'error': 'Could not determine entity column(s) for grouping',
|
||||
'truncated': False}
|
||||
|
||||
# ── Build aggregation expressions ──────────────────────────────────
|
||||
agg_exprs: list[pl.Expr] = [
|
||||
pl.len().alias('flow_count'),
|
||||
pl.col(group_cols[0]).n_unique().alias(f'unique_{group_cols[0]}'),
|
||||
]
|
||||
|
||||
numeric_cols = [c for c in col_names if c not in group_cols
|
||||
and schema[c].startswith(('Int', 'UInt', 'Float'))]
|
||||
for c in numeric_cols[:20]:
|
||||
agg_exprs.append(pl.sum(c).alias(f'total_{c}'))
|
||||
agg_exprs.append(pl.mean(c).alias(f'avg_{c}'))
|
||||
|
||||
str_cols = [c for c in col_names if c not in group_cols
|
||||
and schema[c].startswith(('Utf8', 'String', 'Cat'))]
|
||||
for c in str_cols[:10]:
|
||||
agg_exprs.append(pl.col(c).n_unique().alias(f'unique_{c}'))
|
||||
|
||||
for kw in ['dst_ip', 'dip', 'dst_addr', 'daddr', 'ip_dst', 'dest_ip',
|
||||
'dst_port', 'dport', 'destination_port', 'dest_port',
|
||||
'port', 'protocol', 'proto', 'sni', 'server_name',
|
||||
'cipher_suite', 'cipher', 'tls_version', 'version', '0ver',
|
||||
'country', 'scnt', 'dcnt', 'asn', 'organization']:
|
||||
found = [c for c in col_names if c not in group_cols and kw.lower() in c.lower()]
|
||||
for c in found:
|
||||
alias_name = f'unique_{c}'
|
||||
if alias_name not in [str(a.meta.output_name()) if hasattr(a, 'meta') else '' for a in agg_exprs]:
|
||||
agg_exprs.append(pl.col(c).n_unique().alias(alias_name))
|
||||
|
||||
# ── Execute grouping ───────────────────────────────────────────────
|
||||
try:
|
||||
entity_lf = lf.group_by(group_cols).agg(agg_exprs)
|
||||
n_entities = _count_rows(entity_lf)
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
# ── Store result ───────────────────────────────────────────────────
|
||||
new_schema = {name: str(dtype) for name, dtype in
|
||||
zip(entity_lf.collect_schema().names(), entity_lf.collect_schema().dtypes())}
|
||||
new_id = f'entity_{_generate_id("ds")}'
|
||||
store = SessionStore()
|
||||
store.store_dataset(
|
||||
dataset_id=new_id,
|
||||
lazyframe=entity_lf,
|
||||
schema=new_schema,
|
||||
metadata={
|
||||
'row_count': n_entities,
|
||||
'entity_columns': group_cols,
|
||||
'parent_dataset_id': dataset_id,
|
||||
'n_agg_features': len(agg_exprs),
|
||||
},
|
||||
parent_id=dataset_id,
|
||||
)
|
||||
|
||||
return {
|
||||
'dataset_id': new_id,
|
||||
'entity_columns': group_cols,
|
||||
'n_entities': n_entities,
|
||||
'n_features': len(agg_exprs),
|
||||
'aggregated_columns': list(new_schema.keys())[:20],
|
||||
}
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Adaptive scoring helpers
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
@profile
|
||||
def _add_adaptive_scores(lf: pl.LazyFrame, feature_names: list[str]) -> pl.LazyFrame:
|
||||
"""Add anomaly scores using data-driven weights (SVD + IQR + percentile)."""
|
||||
_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:
|
||||
n_rows = _count_rows(lf)
|
||||
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)
|
||||
|
||||
from sklearn.decomposition import TruncatedSVD
|
||||
n_comp = min(1, min(mat.shape) - 1)
|
||||
if n_comp < 1 or mat.shape[1] < 2:
|
||||
return lf
|
||||
svd = TruncatedSVD(n_components=1, random_state=42).fit(mat)
|
||||
raw_weights = np.abs(svd.components_[0])
|
||||
weights = dict(zip(avail, raw_weights / raw_weights.sum()))
|
||||
|
||||
expr_parts = []
|
||||
for col, w in weights.items():
|
||||
if col == 'modern_tls_rate':
|
||||
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'))
|
||||
|
||||
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()
|
||||
|
||||
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
|
||||
|
||||
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)}
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.info(
|
||||
'[SCORE] adaptive weights from SVD: %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:
|
||||
logging.getLogger(__name__).warning('[SCORE] adaptive scoring skipped: %s', exc)
|
||||
|
||||
return lf
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# _handle_compute_scores
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
async def _handle_compute_scores(dataset_id: str) -> dict:
|
||||
"""Compute adaptive proxy/anomaly/threat scores on entity data."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf = entry['lazyframe']
|
||||
lf = _add_adaptive_scores(lf, [])
|
||||
try:
|
||||
updated_schema = lf.collect_schema()
|
||||
new_schema = {name: str(dtype) for name, dtype in
|
||||
zip(updated_schema.names(), updated_schema.dtypes())}
|
||||
except Exception:
|
||||
new_schema = entry.get('schema', {})
|
||||
store = SessionStore()
|
||||
store.store_dataset(dataset_id, lf, schema=new_schema,
|
||||
metadata=entry.get('metadata', {}))
|
||||
return {'status': 'scored', 'dataset_id': dataset_id}
|
||||
@@ -0,0 +1,78 @@
|
||||
"""evaluate_clustering handler — compute clustering quality metrics."""
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore
|
||||
|
||||
|
||||
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:
|
||||
raise
|
||||
|
||||
return {'scores': scores}
|
||||
@@ -0,0 +1,94 @@
|
||||
"""export_results handler — write datasets/results to disk."""
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore
|
||||
|
||||
|
||||
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:
|
||||
raise
|
||||
|
||||
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}
|
||||
@@ -0,0 +1,309 @@
|
||||
"""extract_features handler — per-cluster feature ranking + ORM persistence."""
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore, _generate_id
|
||||
from ..profile_util import profile
|
||||
|
||||
|
||||
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:
|
||||
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:
|
||||
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,
|
||||
})
|
||||
|
||||
cluster_features.sort(key=lambda x: abs(x['distinguishing_score']), reverse=True)
|
||||
all_features.extend(cluster_features[:top_k])
|
||||
|
||||
# ── Compute UMAP-2D embedding ──────────────────────────────────
|
||||
if len(numeric_cols) >= 2 and len(df) > 2:
|
||||
try:
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
import umap
|
||||
MAX_UMAP_TRAIN = 10000
|
||||
BATCH_SIZE = 1000
|
||||
data_matrix = df.select(numeric_cols).to_numpy()
|
||||
if len(df) > MAX_UMAP_TRAIN:
|
||||
idx_sample = np.random.RandomState(42).choice(
|
||||
len(df), size=MAX_UMAP_TRAIN, replace=False)
|
||||
mat_sample = data_matrix[idx_sample]
|
||||
scaler = StandardScaler().fit(mat_sample)
|
||||
mat_sample_scaled = np.nan_to_num(scaler.transform(mat_sample),
|
||||
nan=0.0, posinf=0.0, neginf=0.0)
|
||||
reducer = umap.UMAP(n_components=2, random_state=42,
|
||||
n_neighbors=15, min_dist=0.1,
|
||||
metric='euclidean')
|
||||
reducer.fit(mat_sample_scaled)
|
||||
coords_list = []
|
||||
for start in range(0, len(df), BATCH_SIZE):
|
||||
end = min(start + BATCH_SIZE, len(df))
|
||||
mat_batch = scaler.transform(data_matrix[start:end])
|
||||
mat_batch = np.nan_to_num(mat_batch, nan=0.0, posinf=0.0, neginf=0.0)
|
||||
coords_list.append(reducer.transform(mat_batch))
|
||||
coords = np.vstack(coords_list)
|
||||
else:
|
||||
data_scaled = StandardScaler().fit_transform(data_matrix)
|
||||
data_scaled = np.nan_to_num(data_scaled, nan=0.0, posinf=0.0, neginf=0.0)
|
||||
reducer = umap.UMAP(n_components=2, random_state=42,
|
||||
n_neighbors=15, min_dist=0.1,
|
||||
metric='euclidean')
|
||||
coords = reducer.fit_transform(data_scaled)
|
||||
non_numeric = [c for c in df.columns if c not in numeric_cols and not c.startswith('_')]
|
||||
entity_col_umap = non_numeric[0] if non_numeric else df.columns[0]
|
||||
labels_list = cluster_entry.get('labels', [])
|
||||
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_umap, ''))
|
||||
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'[UMAP_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
|
||||
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:
|
||||
raise
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
@profile
|
||||
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 not None:
|
||||
from analysis.models import AnalysisRun as AnalysisRun2
|
||||
run = AnalysisRun2.objects.filter(id=run_id).first()
|
||||
if run is None:
|
||||
return False
|
||||
else:
|
||||
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', '')
|
||||
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
|
||||
|
||||
store = SessionStore()
|
||||
cluster_entry = store.get_cluster_result(cluster_result_id)
|
||||
label_counts = {}
|
||||
if cluster_entry:
|
||||
full_labels = cluster_entry.get('labels', [])
|
||||
for lbl in full_labels:
|
||||
label_counts[lbl] = label_counts.get(lbl, 0) + 1
|
||||
|
||||
unique_labels = sorted(set(
|
||||
f['cluster_label'] for f in features
|
||||
))
|
||||
n_clusters = len([l for l in unique_labels if l >= 0])
|
||||
for label in unique_labels:
|
||||
if label == -1:
|
||||
continue
|
||||
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)},
|
||||
)
|
||||
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
|
||||
|
||||
run.cluster_count = n_clusters
|
||||
run.save(update_fields=['cluster_count'])
|
||||
return True
|
||||
@@ -0,0 +1,64 @@
|
||||
"""filter_data handler — apply column-level filters to a dataset."""
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore, _generate_id
|
||||
from ._helpers import _build_filter_expr, _resolve_dataset
|
||||
|
||||
|
||||
async def _handle_filter_data(
|
||||
dataset_id: str,
|
||||
filters: list[dict],
|
||||
logic: str = 'and',
|
||||
) -> dict:
|
||||
"""Apply filters and store filtered dataset."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
full_schema = entry.get('schema', {})
|
||||
|
||||
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.len()).collect(streaming=True).item()
|
||||
except Exception:
|
||||
row_count = None
|
||||
|
||||
new_id = _generate_id('flt')
|
||||
store = SessionStore()
|
||||
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': []},
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
"""load_data handler — load CSV files and store as dataset."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import polars as pl
|
||||
|
||||
from ..data_loader import load_csv_directory
|
||||
from ..session_store import SessionStore, _generate_id
|
||||
|
||||
|
||||
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),
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
"""preprocess_data handler — scale, encode, impute, drop columns."""
|
||||
from typing import Optional
|
||||
|
||||
import polars as pl
|
||||
|
||||
from ..session_store import SessionStore, _generate_id
|
||||
from ._helpers import _resolve_dataset
|
||||
|
||||
|
||||
async def _handle_preprocess_data(
|
||||
dataset_id: str,
|
||||
columns: list[str],
|
||||
config: Optional[dict] = None,
|
||||
) -> dict:
|
||||
"""Preprocess selected columns (scale, encode, impute, drop)."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
cfg = config or {}
|
||||
column_mapping: dict[str, str] = {}
|
||||
|
||||
df = lf.collect(streaming=True)
|
||||
|
||||
# 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'<dropped:{c}>'
|
||||
|
||||
# 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'<onehot:{c}>'
|
||||
|
||||
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 = SessionStore()
|
||||
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),
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
"""profile_data handler — compute per-column statistics and correlations."""
|
||||
from typing import Optional
|
||||
|
||||
import polars as pl
|
||||
|
||||
from ..data_profiler import profile_dataset as _profile_dataset
|
||||
from ._helpers import _resolve_dataset
|
||||
|
||||
|
||||
async def _handle_profile_data(
|
||||
dataset_id: str,
|
||||
sample_size: int = 1000,
|
||||
columns: Optional[list[str]] = None,
|
||||
) -> dict:
|
||||
"""Profile dataset columns and correlations."""
|
||||
entry, err = _resolve_dataset(dataset_id)
|
||||
if err:
|
||||
return err
|
||||
|
||||
lf: pl.LazyFrame = entry['lazyframe']
|
||||
profile = _profile_dataset(lf, sample_size=sample_size, columns=columns)
|
||||
return profile
|
||||
Reference in New Issue
Block a user