PM-pinou
885da89b6f
- 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)
3369 lines
140 KiB
Python
3369 lines
140 KiB
Python
"""Tool definitions and handler functions for the TLS Analyzer MCP server.
|
||
|
||
Each tool is defined by a :class:`Tool <mcp.types.Tool>` metadata object and
|
||
an async handler function. Handlers interact with the :class:`SessionStore`
|
||
singleton and the Django ORM.
|
||
|
||
All 12 tools:
|
||
1. load_data 7. filter_and_cluster
|
||
2. profile_data 8. build_entity_profiles
|
||
3. filter_data 9. export_results
|
||
4. preprocess_data 10. list_datasets
|
||
5. run_clustering 11. drop_dataset
|
||
6. evaluate_clustering 12. clone_dataset
|
||
"""
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Any, Optional
|
||
|
||
import numpy as np
|
||
import polars as pl
|
||
from mcp.types import Tool
|
||
|
||
from .data_loader import load_csv_directory
|
||
from .data_profiler import profile_dataset as _profile_dataset
|
||
from .session_store import SessionStore, _generate_id
|
||
from .profile_util import profile
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Tool metadata
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
|
||
def _json_size(obj: Any) -> int:
|
||
return len(json.dumps(obj, default=str, ensure_ascii=False))
|
||
|
||
|
||
_MAX_RESPONSE_KB = 64
|
||
_MAX_RESPONSE_BYTES = _MAX_RESPONSE_KB * 1024
|
||
|
||
|
||
def _truncate_response(result: dict, max_kb: int = _MAX_RESPONSE_KB) -> dict:
|
||
"""Truncate large lists/keys in *result* to stay under *max_kb*."""
|
||
if _json_size(result) <= max_kb * 1024:
|
||
result.setdefault('truncated', False)
|
||
return result
|
||
|
||
# Drop verbose keys first
|
||
for costly_key in ('column_stats', 'features', 'correlation_matrix',
|
||
'labels', 'cluster_labels'):
|
||
if costly_key in result and isinstance(result[costly_key], dict):
|
||
result[costly_key] = {'__truncated__': True,
|
||
'key_count': len(result[costly_key])}
|
||
elif costly_key in result and isinstance(result[costly_key], list):
|
||
lst = result[costly_key]
|
||
if len(lst) > 10:
|
||
result[costly_key] = lst[:10] + [{'__truncated__': True,
|
||
'omitted': len(lst) - 10}]
|
||
result['truncated'] = True
|
||
return result
|
||
|
||
|
||
def get_tools_meta() -> list[Tool]:
|
||
"""Return the list of Tool metadata objects for MCP server registration.
|
||
|
||
Call this from :meth:`TLSAnalyzerMCPServer._setup_tools`.
|
||
"""
|
||
return [
|
||
Tool(
|
||
name="load_data",
|
||
description=(
|
||
"加载匹配glob模式的CSV文件。自动检测每个文件的BOM，验证schema一致性，"
|
||
"纵向合并，并将结果作为惰性数据集存储在会话存储中。"
|
||
),
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"csv_glob": {
|
||
"type": "string",
|
||
"description": "Glob pattern for CSV files (e.g. 'data/*.csv')",
|
||
},
|
||
"config_path": {
|
||
"type": "string",
|
||
"description": "Optional path to a YAML config file",
|
||
},
|
||
"encoding": {
|
||
"type": "string",
|
||
"description": "Fallback encoding when no BOM is found",
|
||
"default": "utf-8",
|
||
},
|
||
"delimiter": {
|
||
"type": "string",
|
||
"description": "Column delimiter",
|
||
"default": ",",
|
||
},
|
||
"schema_strict": {
|
||
"type": "boolean",
|
||
"description": "When True, raise on column mismatch; when False, merge leniently with nulls",
|
||
"default": False,
|
||
},
|
||
"recursive": {
|
||
"type": "boolean",
|
||
"description": "When True, use recursive glob for **/*.csv patterns",
|
||
"default": False,
|
||
},
|
||
},
|
||
"required": ["csv_glob"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="profile_data",
|
||
description=(
|
||
"计算数据集的每列统计信息和数值相关矩阵。结果自动截断至约16 KB以"
|
||
"节省LLM上下文窗口空间。"
|
||
),
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"dataset_id": {
|
||
"type": "string",
|
||
"description": "Dataset identifier from a previous load_data call",
|
||
},
|
||
"sample_size": {
|
||
"type": "integer",
|
||
"description": "Maximum rows to profile (sampled if larger)",
|
||
"default": 1000,
|
||
},
|
||
"columns": {
|
||
"type": "array",
|
||
"items": {"type": "string"},
|
||
"description": "Columns to profile (default: all)",
|
||
},
|
||
},
|
||
"required": ["dataset_id"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="filter_data",
|
||
description=(
|
||
"对数据集应用列级过滤器，将结果存储为新的派生数据集。支持eq/neq/gt/gte/"
|
||
"lt/lte/contains/not_contains/in/not_in/is_null/is_not_null运算符。"
|
||
),
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"dataset_id": {
|
||
"type": "string",
|
||
"description": "Source dataset identifier",
|
||
},
|
||
"filters": {
|
||
"type": "array",
|
||
"items": {
|
||
"type": "object",
|
||
"properties": {
|
||
"column": {"type": "string"},
|
||
"op": {
|
||
"type": "string",
|
||
"enum": [
|
||
"eq", "neq", "gt", "gte", "lt", "lte",
|
||
"contains", "not_contains",
|
||
"in", "not_in",
|
||
"is_null", "is_not_null",
|
||
],
|
||
},
|
||
"value": {"description": "Operand value"},
|
||
},
|
||
"required": ["column", "op"],
|
||
},
|
||
"description": "List of filter expressions",
|
||
},
|
||
"logic": {
|
||
"type": "string",
|
||
"enum": ["and", "or"],
|
||
"description": "Combine filters with AND or OR logic",
|
||
"default": "and",
|
||
},
|
||
},
|
||
"required": ["dataset_id", "filters"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="preprocess_data",
|
||
description=(
|
||
"预处理选定列：标准化（StandardScaler）、独热编码、缺失值填充和列删除。"
|
||
),
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"dataset_id": {
|
||
"type": "string",
|
||
"description": "Source dataset identifier",
|
||
},
|
||
"columns": {
|
||
"type": "array",
|
||
"items": {"type": "string"},
|
||
"description": "Columns to preprocess",
|
||
},
|
||
"config": {
|
||
"type": "object",
|
||
"properties": {
|
||
"standardize": {
|
||
"type": "boolean",
|
||
"description": "Apply StandardScaler to numeric columns",
|
||
"default": False,
|
||
},
|
||
"onehot": {
|
||
"type": "boolean",
|
||
"description": "One-hot encode categorical columns",
|
||
"default": False,
|
||
},
|
||
"fillna": {
|
||
"type": "string",
|
||
"enum": ["mean", "median", "zero", "drop"],
|
||
"description": "Missing value strategy",
|
||
},
|
||
"drop": {
|
||
"type": "array",
|
||
"items": {"type": "string"},
|
||
"description": "Columns to drop entirely",
|
||
},
|
||
},
|
||
},
|
||
},
|
||
"required": ["dataset_id", "columns"],
|
||
},
|
||
),
|
||
Tool(
|
||
name="run_clustering",
|
||
description=(
|
||
"对选定列执行聚类。使用HDBSCAN（默认）或KMeans。特征在聚类前进行"
|
||
"标准化缩放。返回聚类标签和质量指标。"
|
||
),
|
||
inputSchema={
|
||
"type": "object",
|
||
"properties": {
|
||
"dataset_id": {
|
||
"type": "string",
|
||
"description": "Dataset identifier",
|
||
},
|
||
"cluster_columns": {
|
||
"type": "array",
|
||
"items": {"type": "string"},
|
||
"description": "Columns to use as clustering features",
|
||
},
|
||
"algorithm": {
|
||
"type": "string",
|
||
"enum": ["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"],
|
||
},
|
||
),
|
||
]
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Handler implementations
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
|
||
async def handle_call(name: str, arguments: dict) -> dict:
|
||
"""Dispatch a tool call to the appropriate handler.
|
||
|
||
Args:
|
||
name: Tool name (must match one of the 11 tools).
|
||
arguments: Tool-specific parameters.
|
||
|
||
Returns:
|
||
A JSON-serialisable dict.
|
||
|
||
Raises:
|
||
ValueError: If *name* is not a recognised tool.
|
||
"""
|
||
_handlers = {
|
||
'load_data': _handle_load_data,
|
||
'profile_data': _handle_profile_data,
|
||
'filter_data': _handle_filter_data,
|
||
'preprocess_data': _handle_preprocess_data,
|
||
'run_clustering': _handle_run_clustering,
|
||
'evaluate_clustering': _handle_evaluate_clustering,
|
||
'extract_features': _handle_extract_features,
|
||
'export_results': _handle_export_results,
|
||
'list_datasets': _handle_list_datasets,
|
||
'drop_dataset': _handle_drop_dataset,
|
||
'clone_dataset': _handle_clone_dataset,
|
||
'build_entity_profiles': _handle_build_entity_profiles,
|
||
'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)
|
||
|
||
|
||
# ── 1. load_data ────────────────────────────────────────────────────────
|
||
|
||
async def _handle_load_data(
|
||
csv_glob: str,
|
||
config_path: Optional[str] = None,
|
||
encoding: str = 'utf-8',
|
||
delimiter: str = ',',
|
||
schema_strict: bool = False,
|
||
recursive: bool = False,
|
||
) -> dict:
|
||
"""Load, validate, merge CSVs, and store in session."""
|
||
try:
|
||
lf, schema, row_count, file_count, memory_mb = load_csv_directory(
|
||
glob_pattern=csv_glob,
|
||
encoding=encoding,
|
||
delimiter=delimiter,
|
||
config_path=config_path,
|
||
sample_rows=10000,
|
||
schema_strict=schema_strict,
|
||
recursive=recursive,
|
||
)
|
||
except FileNotFoundError as e:
|
||
return {'error': f'File not found: {e}', 'truncated': False}
|
||
except ValueError as e:
|
||
return {'error': f'Schema validation failed: {e}', 'truncated': False}
|
||
|
||
store = SessionStore()
|
||
dataset_id = _generate_id('ds')
|
||
store.store_dataset(
|
||
dataset_id=dataset_id,
|
||
lazyframe=lf,
|
||
schema=schema,
|
||
metadata={
|
||
'row_count': row_count,
|
||
'file_count': file_count,
|
||
'memory_mb': round(memory_mb, 2),
|
||
'csv_glob': csv_glob,
|
||
'encoding': encoding,
|
||
},
|
||
)
|
||
|
||
return {
|
||
'dataset_id': dataset_id,
|
||
'schema': schema,
|
||
'row_count': row_count,
|
||
'file_count': file_count,
|
||
'memory_mb': round(memory_mb, 2),
|
||
}
|
||
|
||
|
||
# ── 2. profile_data ─────────────────────────────────────────────────────
|
||
|
||
async def _handle_profile_data(
|
||
dataset_id: str,
|
||
sample_size: int = 1000,
|
||
columns: Optional[list[str]] = None,
|
||
) -> dict:
|
||
"""Profile dataset columns and correlations."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
profile = _profile_dataset(lf, sample_size=sample_size, columns=columns)
|
||
return profile
|
||
|
||
|
||
# ── 3. filter_data ──────────────────────────────────────────────────────
|
||
|
||
def _build_filter_expr(column: str, op: str, value: Any) -> pl.Expr:
|
||
"""Build a single Polars expression from a filter descriptor."""
|
||
col = pl.col(column)
|
||
op_map = {
|
||
'eq': col == value,
|
||
'neq': col != value,
|
||
'gt': col > value,
|
||
'gte': col >= value,
|
||
'lt': col < value,
|
||
'lte': col <= value,
|
||
'contains': col.str.contains(value) if value is not None else col.cast(str).str.contains(''),
|
||
'not_contains': ~(col.str.contains(value) if value is not None else col.cast(str).str.contains('')),
|
||
'in': col.is_in(value) if isinstance(value, (list, tuple)) else col.is_in([value]),
|
||
'not_in': ~(col.is_in(value) if isinstance(value, (list, tuple)) else col.is_in([value])),
|
||
'is_null': col.is_null(),
|
||
'is_not_null': col.is_not_null(),
|
||
}
|
||
expr = op_map.get(op)
|
||
if expr is None:
|
||
raise ValueError(f"Unsupported filter operator: {op}")
|
||
return expr
|
||
|
||
|
||
async def _handle_filter_data(
|
||
dataset_id: str,
|
||
filters: list[dict],
|
||
logic: str = 'and',
|
||
) -> dict:
|
||
"""Apply filters and store filtered dataset."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
full_schema = entry.get('schema', {})
|
||
|
||
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.store_dataset(
|
||
dataset_id=new_id,
|
||
lazyframe=filtered_lf,
|
||
schema=full_schema,
|
||
metadata={
|
||
'row_count': row_count,
|
||
'parent_dataset_id': dataset_id,
|
||
'filter_count': len(filters),
|
||
'logic': logic,
|
||
},
|
||
parent_id=dataset_id,
|
||
)
|
||
|
||
return {
|
||
'filtered_dataset_id': new_id,
|
||
'row_count': row_count,
|
||
'column_deltas': {'kept': list(full_schema.keys()), 'dropped': []},
|
||
}
|
||
|
||
|
||
# ── 4. preprocess_data ──────────────────────────────────────────────────
|
||
|
||
async def _handle_preprocess_data(
|
||
dataset_id: str,
|
||
columns: list[str],
|
||
config: Optional[dict] = None,
|
||
) -> dict:
|
||
"""Preprocess selected columns (scale, encode, impute, drop)."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
cfg = config or {}
|
||
column_mapping: dict[str, str] = {}
|
||
|
||
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.store_dataset(
|
||
dataset_id=new_id,
|
||
lazyframe=new_lf,
|
||
schema=new_schema,
|
||
metadata={
|
||
'row_count': len(df),
|
||
'parent_dataset_id': dataset_id,
|
||
'preprocessing_config': cfg,
|
||
},
|
||
parent_id=dataset_id,
|
||
)
|
||
|
||
return {
|
||
'preprocessed_dataset_id': new_id,
|
||
'column_mapping': column_mapping,
|
||
'row_count': len(df),
|
||
'column_count': len(new_schema),
|
||
}
|
||
|
||
|
||
# ── 5. run_clustering ───────────────────────────────────────────────────
|
||
|
||
async def _handle_run_clustering(
|
||
dataset_id: str,
|
||
cluster_columns: list[str],
|
||
algorithm: str = 'agglomerative',
|
||
params: Optional[dict] = None,
|
||
random_state: int = 42,
|
||
) -> dict:
|
||
"""Run HDBSCAN or KMeans clustering on selected columns (async)."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
p = params or {}
|
||
|
||
try:
|
||
# Collect cluster features - numeric only
|
||
available = lf.collect_schema()
|
||
available_names = available.names()
|
||
available_dtypes = available.dtypes()
|
||
numeric_dtypes = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
||
pl.Float32, pl.Float64)
|
||
feature_cols = [
|
||
c for c in cluster_columns
|
||
if c in available_names
|
||
and available_dtypes[available_names.index(c)] in numeric_dtypes
|
||
]
|
||
if not feature_cols:
|
||
# Fall back: auto-select all numeric columns (exclude internal ones)
|
||
feature_cols = [
|
||
available_names[i] for i, dt in enumerate(available_dtypes)
|
||
if dt in numeric_dtypes and not available_names[i].startswith('_')
|
||
][:10]
|
||
if not feature_cols:
|
||
# Fallback: try coercing Utf8 columns to Float64 (common when data loaded from
|
||
# SQLite TEXT storage without proper dtype restoration).
|
||
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:
|
||
# Attempt per-column cast; filter to columns where coercion succeeded
|
||
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: # at least 10% successfully parsed as floats
|
||
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: compute row count via fast lazy count
|
||
try:
|
||
row_count = lf.select(pl.len()).collect(streaming=True).item()
|
||
except Exception:
|
||
row_count = 0
|
||
if row_count > 100000:
|
||
# Sample to 100K before collect to avoid OOM on 10M+ rows
|
||
lf = lf.collect(streaming=True).sample(n=100000, seed=42).lazy()
|
||
|
||
df_features = lf.select(feature_cols).collect(streaming=True)
|
||
|
||
# Drop all-NaN columns (keep rows, drop useless columns)
|
||
for col in df_features.columns:
|
||
if df_features[col].is_null().all():
|
||
df_features = df_features.drop(col)
|
||
feature_cols = [c for c in feature_cols if c != col]
|
||
|
||
if df_features.width == 0:
|
||
return {'error': 'No valid feature columns after dropping all-NaN columns',
|
||
'truncated': False}
|
||
|
||
data = df_features.to_numpy()
|
||
|
||
# Fill remaining NaN with column mean
|
||
if data.dtype.kind == 'f':
|
||
col_mean = np.nanmean(data, axis=0)
|
||
col_mean = np.nan_to_num(col_mean, nan=0.0)
|
||
data = np.where(np.isnan(data), col_mean, data)
|
||
|
||
# Zero-sample / too-few-samples guard
|
||
if len(data) < 3:
|
||
return {
|
||
'cluster_result_id': _generate_id('clust'),
|
||
'n_clusters': 0,
|
||
'n_noise': 0,
|
||
'quality_metrics': {'note': 'too few samples'},
|
||
'truncated': False,
|
||
}
|
||
|
||
# ── D14: Low memory check (safe fallback if psutil not installed) ──
|
||
low_memory = False
|
||
try:
|
||
import psutil
|
||
mem = psutil.virtual_memory()
|
||
low_memory = mem.available < 2 * 1024**3 # 2 GB
|
||
except ImportError:
|
||
pass
|
||
|
||
# ── D14: Auto-downsample for low memory (>100K rows → 50K) ────────
|
||
if low_memory and len(data) > 100000:
|
||
rng = np.random.default_rng(42)
|
||
idx = rng.choice(len(data), 50000, replace=False)
|
||
data = data[idx]
|
||
|
||
# ── D5: Downsampling → extract matching indices ──────────────
|
||
# Downsample to 50K max, store indices so _handle_extract_features
|
||
# can filter the same rows.
|
||
_ds_idx = None
|
||
if len(data) > 50000:
|
||
rng = np.random.default_rng(42)
|
||
_ds_idx = rng.choice(len(data), 50000, replace=False)
|
||
data = data[_ds_idx]
|
||
|
||
# ── H17: Variance filtering ───────────────────────────────────────
|
||
# Remove features with zero / near-zero variance (< 1e-10)
|
||
if data.shape[1] > 1:
|
||
variances = np.var(data, axis=0)
|
||
keep_var = np.where(variances >= 1e-10)[0]
|
||
if len(keep_var) == 0:
|
||
# All features have zero variance — fall back to unfiltered data
|
||
keep_var = np.arange(data.shape[1])
|
||
if len(keep_var) < data.shape[1]:
|
||
data = data[:, keep_var]
|
||
feature_cols = [feature_cols[i] for i in keep_var]
|
||
|
||
# Log clustering input
|
||
import logging
|
||
_clust_logger = logging.getLogger(__name__)
|
||
_clust_logger.info(f'[CLUSTER_INPUT] rows={len(data)} cols={data.shape[1] if len(data.shape) > 1 else 1} feature_cols={feature_cols}')
|
||
|
||
# Standard scale
|
||
from sklearn.preprocessing import StandardScaler
|
||
_clust_logger.info('[CLUSTER_SCALE] Starting StandardScaler...')
|
||
data_scaled = StandardScaler().fit_transform(data)
|
||
_clust_logger.info(f'[CLUSTER_SCALE] Done. shape={data_scaled.shape}')
|
||
|
||
# ── H17: Correlation filtering ────────────────────────────────────
|
||
# Remove features with Pearson correlation > 0.95 (keep the first one)
|
||
if data_scaled.shape[1] > 1:
|
||
corr = np.corrcoef(data_scaled.T)
|
||
upper = np.triu(np.abs(corr), 1)
|
||
to_drop = [j for j in range(upper.shape[1]) if any(upper[:, j] > 0.95)]
|
||
if to_drop:
|
||
keep = [j for j in range(data_scaled.shape[1]) if j not in to_drop]
|
||
data_scaled = data_scaled[:, keep]
|
||
feature_cols = [feature_cols[j] for j in keep]
|
||
|
||
# ── D14: 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)
|
||
|
||
# ── H17: HDBSCAN hyperparameter auto-tuning based on data size ────
|
||
n = len(data_scaled)
|
||
if algorithm == 'hdbscan':
|
||
min_cluster = p.get('min_cluster_size', max(3, min(50, n // 20)))
|
||
p['min_cluster_size'] = min_cluster
|
||
p['min_samples'] = p.get('min_samples', min_cluster)
|
||
p['cluster_selection_epsilon'] = p.get('cluster_selection_epsilon', 0.0)
|
||
else:
|
||
# KMeans safety cap
|
||
p['min_cluster_size'] = min(
|
||
p.get('min_cluster_size', 5),
|
||
max(n // 2, 2),
|
||
)
|
||
|
||
# Clustering
|
||
from sklearn.cluster import HDBSCAN, MiniBatchKMeans, AgglomerativeClustering
|
||
if algorithm == 'kmeans':
|
||
n_clusters = p.get('n_clusters', 3)
|
||
model = MiniBatchKMeans(
|
||
n_clusters=n_clusters, random_state=random_state,
|
||
batch_size=1024, n_init='auto',
|
||
)
|
||
elif algorithm in ('agglomerative', 'ward'):
|
||
n_clusters = p.get('n_clusters', 5)
|
||
model = AgglomerativeClustering(
|
||
n_clusters=n_clusters,
|
||
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
|
||
|
||
# Map back to original rows
|
||
full_labels = labels_array
|
||
|
||
# Quality metrics
|
||
quality = {
|
||
'n_clusters': n_clusters_found,
|
||
'n_noise': noise_count,
|
||
'noise_ratio': noise_ratio,
|
||
'algorithm': algorithm,
|
||
}
|
||
|
||
if n_clusters_found > 1:
|
||
try:
|
||
from sklearn.metrics import silhouette_score
|
||
score = silhouette_score(data_scaled, labels_array, random_state=random_state)
|
||
quality['silhouette_score'] = round(float(score), 4)
|
||
except Exception:
|
||
quality['silhouette_score'] = None
|
||
try:
|
||
from sklearn.metrics import davies_bouldin_score
|
||
dbs = davies_bouldin_score(data_scaled, labels_array)
|
||
quality['davies_bouldin_score'] = round(float(dbs), 4)
|
||
except Exception:
|
||
quality['davies_bouldin_score'] = None
|
||
try:
|
||
from sklearn.metrics import calinski_harabasz_score
|
||
chs = calinski_harabasz_score(data_scaled, labels_array)
|
||
quality['calinski_harabasz_score'] = round(float(chs), 4)
|
||
except Exception:
|
||
quality['calinski_harabasz_score'] = None
|
||
|
||
# Cluster sizes
|
||
cluster_sizes = {}
|
||
for label in np.unique(labels_array):
|
||
cluster_sizes[int(label)] = int((labels_array == label).sum())
|
||
quality['cluster_sizes'] = cluster_sizes
|
||
|
||
except Exception:
|
||
raise
|
||
|
||
result_id = _generate_id('clust')
|
||
store.store_cluster_result(
|
||
result_id=result_id,
|
||
labels=full_labels.tolist(),
|
||
n_clusters=n_clusters_found,
|
||
params={
|
||
'algorithm': algorithm,
|
||
'params': p,
|
||
'random_state': random_state,
|
||
'feature_columns': feature_cols,
|
||
'_ds_idx': _ds_idx.tolist() if _ds_idx is not None else None,
|
||
},
|
||
parent_dataset_id=dataset_id,
|
||
)
|
||
|
||
return {
|
||
'cluster_result_id': result_id,
|
||
'n_clusters': n_clusters_found,
|
||
'n_noise': noise_count,
|
||
'labels_sample': full_labels[:50].tolist(),
|
||
'quality_metrics': quality,
|
||
}
|
||
|
||
|
||
# ── 6. evaluate_clustering ──────────────────────────────────────────────
|
||
|
||
async def _handle_evaluate_clustering(
|
||
cluster_result_id: str,
|
||
dataset_id: str,
|
||
metrics: Optional[list[str]] = None,
|
||
) -> dict:
|
||
"""Evaluate clustering result against the dataset."""
|
||
store = SessionStore()
|
||
cluster_entry = store.get_cluster_result(cluster_result_id)
|
||
if cluster_entry is None:
|
||
return {'error': f'Cluster result not found: {cluster_result_id}',
|
||
'truncated': False}
|
||
|
||
dataset_entry = store.get_dataset(dataset_id)
|
||
if dataset_entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
labels = cluster_entry['labels']
|
||
default_metrics = ['silhouette', 'davies_bouldin', 'calinski_harabasz',
|
||
'noise_ratio', 'cluster_sizes']
|
||
requested = metrics or default_metrics
|
||
|
||
scores: dict = {'n_clusters': cluster_entry['n_clusters']}
|
||
|
||
try:
|
||
lf: pl.LazyFrame = dataset_entry['lazyframe']
|
||
params = cluster_entry.get('params', {})
|
||
feature_cols = params.get('feature_columns', [])
|
||
if not feature_cols and 'feature_columns' in params:
|
||
pass
|
||
elif feature_cols:
|
||
df_features = lf.select(feature_cols).collect(streaming=True)
|
||
data = df_features.to_numpy()
|
||
|
||
# Align labels with data (handle missing rows)
|
||
valid_mask = ~np.isnan(data).any(axis=1) if data.dtype.kind == 'f' else np.ones(len(data), dtype=bool)
|
||
data_clean = data[valid_mask]
|
||
labels_clean = np.array(labels)[:len(data_clean)]
|
||
|
||
unique_labels = set(labels_clean)
|
||
n_clusters_actual = len(unique_labels - {-1})
|
||
|
||
for m in requested:
|
||
if m == 'silhouette' and n_clusters_actual >= 2:
|
||
from sklearn.metrics import silhouette_score
|
||
scores['silhouette'] = round(float(
|
||
silhouette_score(data_clean, labels_clean, random_state=42)), 4)
|
||
elif m == 'davies_bouldin' and n_clusters_actual >= 2:
|
||
from sklearn.metrics import davies_bouldin_score
|
||
scores['davies_bouldin'] = round(float(
|
||
davies_bouldin_score(data_clean, labels_clean)), 4)
|
||
elif m == 'calinski_harabasz' and n_clusters_actual >= 2:
|
||
from sklearn.metrics import calinski_harabasz_score
|
||
scores['calinski_harabasz'] = round(float(
|
||
calinski_harabasz_score(data_clean, labels_clean)), 4)
|
||
elif m == 'noise_ratio':
|
||
noise = int((labels_clean == -1).sum())
|
||
scores['noise_ratio'] = round(noise / len(labels_clean), 4)
|
||
elif m == 'cluster_sizes':
|
||
sizes = {}
|
||
for label in unique_labels:
|
||
sizes[int(label)] = int((labels_clean == label).sum())
|
||
scores['cluster_sizes'] = sizes
|
||
|
||
scores['n_clusters'] = cluster_entry['n_clusters']
|
||
|
||
except Exception:
|
||
raise
|
||
|
||
return {'scores': scores}
|
||
|
||
|
||
# ── 7. extract_features ─────────────────────────────────────────────────
|
||
|
||
async def _handle_extract_features(
|
||
dataset_id: str,
|
||
cluster_result_id: str,
|
||
top_k: int = 10,
|
||
method: str = 'zscore',
|
||
save_to_db: bool = True,
|
||
run_id: int | None = None,
|
||
) -> dict:
|
||
"""Compute per-cluster distinguishing features and optionally persist."""
|
||
store = SessionStore()
|
||
dataset_entry = store.get_dataset(dataset_id)
|
||
if dataset_entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
cluster_entry = store.get_cluster_result(cluster_result_id)
|
||
if cluster_entry is None:
|
||
return {'error': f'Cluster result not found: {cluster_result_id}',
|
||
'truncated': False}
|
||
|
||
try:
|
||
lf: pl.LazyFrame = dataset_entry['lazyframe']
|
||
labels = cluster_entry['labels']
|
||
|
||
# D5: If clustering downsampled the data, filter matching rows
|
||
_ds_idx = cluster_entry.get('params', {}).get('_ds_idx')
|
||
if _ds_idx is not None:
|
||
# Filter in lazy plan before collect — avoids loading full 10M rows
|
||
df = (lf.with_row_index('__ds_idx')
|
||
.filter(pl.col('__ds_idx').is_in(_ds_idx))
|
||
.drop('__ds_idx')
|
||
.collect(streaming=True))
|
||
else:
|
||
df = lf.collect(streaming=True)
|
||
|
||
# Truncate labels if df was filtered but labels weren't
|
||
if len(labels) != len(df):
|
||
labels = labels[:len(df)]
|
||
|
||
# Only use numeric columns for feature extraction
|
||
numeric_cols = [c for c in df.columns
|
||
if df[c].dtype in (pl.Float32, pl.Float64,
|
||
pl.Int32, pl.Int64,
|
||
pl.UInt32, pl.UInt64)]
|
||
if len(labels) != len(df):
|
||
labels = labels[:len(df)]
|
||
|
||
# Global statistics per column
|
||
global_stats = {}
|
||
for col in numeric_cols:
|
||
s = df[col].drop_nulls()
|
||
if len(s) > 0:
|
||
global_stats[col] = {
|
||
'mean': float(s.mean()),
|
||
'std': float(s.std()) if s.std() is not None else 1.0,
|
||
}
|
||
|
||
unique_labels = sorted(set(labels))
|
||
all_features = []
|
||
|
||
for label in unique_labels:
|
||
mask = np.array(labels) == label
|
||
if mask.sum() < 2:
|
||
continue
|
||
|
||
cluster_df = df.filter(pl.Series('__mask__', mask))
|
||
cluster_features = []
|
||
|
||
for col in numeric_cols:
|
||
c_series = cluster_df[col].drop_nulls()
|
||
gs = global_stats.get(col)
|
||
if gs is None or gs['std'] == 0 or len(c_series) == 0:
|
||
continue
|
||
|
||
c_mean = float(c_series.mean())
|
||
c_std = float(c_series.std()) if c_series.std() is not None else 0.0
|
||
|
||
if method == 'zscore':
|
||
score = (c_mean - gs['mean']) / max(gs['std'], 1e-10)
|
||
else:
|
||
# ANOVA-like: F = (between-group variance) / (within-group variance)
|
||
score = abs(c_mean - gs['mean']) / max(c_std + gs['std'], 1e-10)
|
||
|
||
cluster_features.append({
|
||
'feature_name': col,
|
||
'cluster_label': int(label),
|
||
'mean': round(c_mean, 4),
|
||
'std': round(c_std, 4),
|
||
'global_mean': round(gs['mean'], 4),
|
||
'global_std': round(gs['std'], 4),
|
||
'distinguishing_score': round(float(score), 4),
|
||
'distinguishing_method': method,
|
||
})
|
||
|
||
# Sort by distinguishing score descending, take top_k
|
||
cluster_features.sort(key=lambda x: abs(x['distinguishing_score']), reverse=True)
|
||
all_features.extend(cluster_features[:top_k])
|
||
|
||
# ── Compute 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:
|
||
# Train UMAP on 10K sample, batch-transform full dataset
|
||
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)
|
||
# Batch transform in 1K-row batches
|
||
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)
|
||
# Handle NaN/Inf values (np is imported at module level)
|
||
data_scaled = np.nan_to_num(data_scaled, nan=0.0, posinf=0.0, neginf=0.0)
|
||
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', [])
|
||
# Save: create EntityProfile records with coordinates
|
||
from analysis.models import EntityProfile as EP
|
||
from analysis.models import ClusterResult as CR
|
||
from analysis.models import AnalysisRun as AR
|
||
from asgiref.sync import sync_to_async
|
||
csv_glob = dataset_entry.get('metadata', {}).get('csv_glob', 'manual')
|
||
run_obj = await sync_to_async(
|
||
AR.objects.filter(csv_glob=csv_glob).order_by('-id').first,
|
||
thread_sensitive=True
|
||
)()
|
||
if run_obj:
|
||
run_id_for_ep = run_obj.id
|
||
profiles = []
|
||
cr_cache = {}
|
||
for i, row in enumerate(df.iter_rows(named=True)):
|
||
if i >= len(coords): break
|
||
ev = str(row.get(entity_col_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 # local ref for closure
|
||
async def _do_save(ce_inner):
|
||
if run_id is not None:
|
||
run = await sync_to_async(AnalysisRunModel.objects.get)(id=run_id)
|
||
else:
|
||
csv_glob_fallback = dataset_entry.get('metadata', {}).get('csv_glob', 'manual')
|
||
run = await sync_to_async(
|
||
AnalysisRunModel.objects.filter(csv_glob=csv_glob_fallback).order_by('-id').first,
|
||
thread_sensitive=True
|
||
)()
|
||
if run is None:
|
||
run = await sync_to_async(AnalysisRunModel.objects.create)(
|
||
csv_glob=csv_glob_fallback, status='running', total_flows=0,
|
||
)
|
||
run.status = 'completed'
|
||
run.cluster_count = ce_inner['n_clusters']
|
||
await sync_to_async(run.save)()
|
||
return await sync_to_async(_save_features_to_db)(cluster_result_id, all_features, run_id=run.id)
|
||
db_saved = await _do_save(ce)
|
||
except Exception as e:
|
||
import traceback, sys; print(f'[DB_SAVE_ERROR] {e}\n{traceback.format_exc()}', file=sys.stderr)
|
||
db_saved = False
|
||
|
||
except Exception:
|
||
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 provided, use it directly (skip csv_glob lookup).
|
||
"""
|
||
if run_id is not None:
|
||
from analysis.models import AnalysisRun as AnalysisRun2
|
||
run = AnalysisRun2.objects.filter(id=run_id).first()
|
||
if run is None:
|
||
return False
|
||
else:
|
||
# Try to find the ClusterResult record by run association
|
||
store = SessionStore()
|
||
cluster_entry = store.get_cluster_result(cluster_result_id)
|
||
if cluster_entry is None:
|
||
return False
|
||
|
||
parent_ds_id = cluster_entry.get('parent_dataset_id', '')
|
||
from analysis.models import AnalysisRun
|
||
ds_entry = store.get_dataset(parent_ds_id)
|
||
if ds_entry is None:
|
||
return False
|
||
|
||
csv_glob = ds_entry.get('metadata', {}).get('csv_glob', '')
|
||
# Compute n_clusters from features before using it
|
||
unique_global = sorted(set(f['cluster_label'] for f in features))
|
||
n_clusters_global = len([l for l in unique_global if l >= 0])
|
||
try:
|
||
run = AnalysisRun.objects.filter(csv_glob=csv_glob).order_by('-created_at').first()
|
||
if run is None:
|
||
run = AnalysisRun.objects.create(
|
||
csv_glob=csv_glob,
|
||
status='completed',
|
||
cluster_count=n_clusters_global,
|
||
)
|
||
except Exception:
|
||
return False
|
||
|
||
from analysis.models import ClusterResult as ClusterResultModel
|
||
|
||
# Compute actual cluster sizes from full label list (not from features count!)
|
||
# Use the cluster_entry which has the complete labels array
|
||
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
|
||
))
|
||
# Compute n_clusters from features (excluding noise -1)
|
||
n_clusters = len([l for l in unique_labels if l >= 0])
|
||
for label in unique_labels:
|
||
if label == -1:
|
||
continue # skip noise cluster
|
||
label_features = [f for f in features if f['cluster_label'] == label]
|
||
if not label_features:
|
||
continue
|
||
|
||
cluster_db, _ = ClusterResultModel.objects.get_or_create(
|
||
run=run,
|
||
cluster_label=label,
|
||
defaults={'size': label_counts.get(label, 1)},
|
||
)
|
||
# Update size if already exists
|
||
if label_counts.get(label) is not None:
|
||
ClusterResultModel.objects.filter(
|
||
run=run, cluster_label=label
|
||
).update(size=label_counts[label])
|
||
|
||
from analysis.models import ClusterFeature as ClusterFeatureModel
|
||
for feat in label_features:
|
||
try:
|
||
ClusterFeatureModel.objects.get_or_create(
|
||
cluster=cluster_db,
|
||
feature_name=feat['feature_name'],
|
||
defaults={
|
||
'mean': feat.get('mean'),
|
||
'std': feat.get('std'),
|
||
'median': None,
|
||
'p25': None,
|
||
'p75': None,
|
||
'missing_rate': None,
|
||
'distinguishing_score': feat.get('distinguishing_score'),
|
||
'distinguishing_method': feat.get('distinguishing_method', 'zscore'),
|
||
}
|
||
)
|
||
except Exception:
|
||
continue # skip problematic features
|
||
|
||
# Update run counts
|
||
run.cluster_count = n_clusters
|
||
run.save(update_fields=['cluster_count'])
|
||
return True
|
||
|
||
return True
|
||
|
||
|
||
# ── 8. export_results ───────────────────────────────────────────────────
|
||
|
||
async def _handle_export_results(
|
||
result_id: str,
|
||
output_path: str,
|
||
format: str = 'csv',
|
||
overwrite: bool = False,
|
||
) -> dict:
|
||
"""Export a dataset (or result) to disk."""
|
||
store = SessionStore()
|
||
|
||
out = Path(output_path)
|
||
file_paths: list[str] = []
|
||
|
||
# Check if it's a dataset
|
||
dataset_entry = store.get_dataset(result_id)
|
||
if dataset_entry is not None:
|
||
lf: pl.LazyFrame = dataset_entry['lazyframe']
|
||
try:
|
||
df = lf.collect(streaming=True)
|
||
if out.suffix == '':
|
||
out = out / f"export_{result_id}.{format}"
|
||
|
||
if out.exists() and not overwrite:
|
||
return {'error': f'Output exists and overwrite=False: {out}',
|
||
'truncated': False}
|
||
|
||
out.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
if format == 'csv':
|
||
df.write_csv(out)
|
||
elif format == 'parquet':
|
||
df.write_parquet(out)
|
||
elif format == 'json':
|
||
df.write_json(out)
|
||
else:
|
||
return {'error': f'Unsupported format: {format}',
|
||
'truncated': False}
|
||
|
||
file_paths.append(str(out.resolve()))
|
||
except Exception:
|
||
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}
|
||
|
||
|
||
# ── 9. list_datasets ────────────────────────────────────────────────────
|
||
|
||
async def _handle_list_datasets(
|
||
_filter_type: Optional[str] = None,
|
||
) -> dict:
|
||
"""List all active datasets and results."""
|
||
store = SessionStore()
|
||
datasets = store.list_datasets()
|
||
return {'datasets': datasets, 'count': len(datasets)}
|
||
|
||
|
||
# ── 10. drop_dataset ────────────────────────────────────────────────────
|
||
|
||
async def _handle_drop_dataset(dataset_id: str) -> dict:
|
||
"""Remove a dataset/result and release memory."""
|
||
store = SessionStore()
|
||
|
||
dropped = (store.drop_dataset(dataset_id)
|
||
or store.drop_cluster_result(dataset_id)
|
||
or store.drop_feature_result(dataset_id))
|
||
|
||
if dropped:
|
||
return {'dropped': True, 'dataset_id': dataset_id}
|
||
return {'error': f'No dataset or result found: {dataset_id}',
|
||
'truncated': False}
|
||
|
||
|
||
# ── 11. clone_dataset ───────────────────────────────────────────────────
|
||
|
||
async def _handle_clone_dataset(dataset_id: str) -> dict:
|
||
"""Shallow-clone a dataset."""
|
||
store = SessionStore()
|
||
new_id = store.clone_dataset(dataset_id)
|
||
if new_id is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
return {'original_dataset_id': dataset_id, 'cloned_dataset_id': new_id}
|
||
|
||
|
||
# ── 12. build_entity_profiles ────────────────────────────────────────────
|
||
|
||
_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)
|
||
|
||
|
||
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.
|
||
|
||
Stores the result in the session store and returns the new dataset_id.
|
||
"""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
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:
|
||
# Auto-detect: find columns matching entity keywords
|
||
matched = []
|
||
for kw in _ENTITY_KEYWORDS:
|
||
found = [c for c in col_names if kw.lower() in c.lower()]
|
||
matched.extend(found)
|
||
# Deduplicate while preserving order
|
||
seen = set()
|
||
group_cols = [c for c in matched if not (c in seen or seen.add(c))]
|
||
if not group_cols:
|
||
# Fallback: pick first non-numeric, non-internal column
|
||
for c in col_names:
|
||
if c.startswith('_'):
|
||
continue
|
||
group_cols = [c]
|
||
break
|
||
# Limit to at most 3 entity columns
|
||
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 ──────────────────────────────────
|
||
# Always: count of flows
|
||
agg_exprs: list[pl.Expr] = [
|
||
pl.len().alias('flow_count'),
|
||
pl.col(group_cols[0]).n_unique().alias(f'unique_{group_cols[0]}'),
|
||
]
|
||
|
||
# Numeric aggregations (sum, mean)
|
||
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}'))
|
||
|
||
# Unique counts for string-like entity columns (dst_ip, dst_port, sni, etc.)
|
||
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}'))
|
||
|
||
# Also count distinct values of specific keyword-matching columns
|
||
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)
|
||
# Estimate row count
|
||
n_entities = entity_lf.select(pl.len()).collect(streaming=True).item()
|
||
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.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 anomaly scoring helpers (inlined from entity_aggregator) ─────
|
||
|
||
_ANOMALY_FEATURE_COLS = [
|
||
'non_std_port_rate', 'modern_tls_rate', 'sni_missing_ratio',
|
||
'recoverable_ratio', 'flow_count', 'total_bytes_sent',
|
||
'total_packets', 'unique_dst_ips', 'unique_dst_ports',
|
||
'unique_protocols', 'has_sni', 'avg_sni_length',
|
||
'countries_visited',
|
||
]
|
||
|
||
|
||
@profile
|
||
def _add_adaptive_scores(lf: pl.LazyFrame, feature_names: list[str]) -> pl.LazyFrame:
|
||
"""Add anomaly scores using data-driven weights (SVD + IQR + percentile).
|
||
|
||
Steps:
|
||
1. Find which anomaly feature columns exist in the aggregated data
|
||
2. Collect a sample (max 10K rows) for SVD weight learning
|
||
3. Learn feature weights from SVD first component loadings
|
||
4. Compute weighted proxy_score across ALL rows (vectorized, no per-row ops)
|
||
5. Add percentile rank calibration
|
||
6. Infer risk thresholds from data distribution
|
||
"""
|
||
_agg_cols = lf.collect_schema().names()
|
||
avail = [c for c in _ANOMALY_FEATURE_COLS if c in _agg_cols]
|
||
if not avail:
|
||
return lf
|
||
|
||
try:
|
||
# Step 1-2: Sample up to 10K rows for weight learning
|
||
n_rows = lf.select(pl.len()).collect(streaming=True).item()
|
||
sample_lf = lf if n_rows <= 10000 else lf.sample(n=10000, seed=42)
|
||
df_sample = sample_lf.select(avail).collect(streaming=True)
|
||
mat = df_sample.to_numpy()
|
||
mat = np.nan_to_num(mat, nan=0.0)
|
||
|
||
# Step 3: SVD weight learning
|
||
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])
|
||
# Normalise to sum = 1
|
||
weights = dict(zip(avail, raw_weights / raw_weights.sum()))
|
||
|
||
# Step 4: Weighted sum expression (all vectorised, no loop)
|
||
expr_parts = []
|
||
for col, w in weights.items():
|
||
if col == 'modern_tls_rate':
|
||
# modern_tls_rate is good — invert for anomaly scoring
|
||
expr_parts.append((1.0 - pl.col(col).fill_null(0.5)) * w)
|
||
else:
|
||
expr_parts.append(pl.col(col).fill_null(0.0) * w)
|
||
if expr_parts:
|
||
proxy_expr = expr_parts[0]
|
||
for p in expr_parts[1:]:
|
||
proxy_expr = proxy_expr + p
|
||
else:
|
||
return lf
|
||
|
||
lf = lf.with_columns(proxy_expr.clip(0, 1).alias('_raw_score'))
|
||
|
||
# Step 5: Percentile rank calibration (0-1, relative position)
|
||
sample_scores = df_sample.select(
|
||
pl.sum_horizontal([
|
||
(1.0 - pl.col('modern_tls_rate').fill_null(0.5)) * weights.get('modern_tls_rate', 0)
|
||
if c == 'modern_tls_rate'
|
||
else pl.col(c).fill_null(0.0) * weights.get(c, 0)
|
||
for c in avail
|
||
]).clip(0, 1)
|
||
).to_series().to_numpy()
|
||
|
||
# Compute IQR thresholds from sample
|
||
q25, q50, q75 = np.percentile(sample_scores, [25, 50, 75])
|
||
iqr = q75 - q25
|
||
upper_fence = float(q75 + 1.5 * iqr) if iqr > 0 else 0.9
|
||
extreme_fence = float(q75 + 3.0 * iqr) if iqr > 0 else 0.99
|
||
|
||
# Step 6: Risk level labels (0.0-1.0)
|
||
lf = lf.with_columns([
|
||
pl.col('_raw_score').alias('proxy_score'),
|
||
pl.when(pl.col('_raw_score') >= extreme_fence)
|
||
.then(pl.lit('critical'))
|
||
.when(pl.col('_raw_score') >= upper_fence)
|
||
.then(pl.lit('suspicious'))
|
||
.when(pl.col('_raw_score') >= q50)
|
||
.then(pl.lit('watch'))
|
||
.otherwise(pl.lit('normal'))
|
||
.alias('risk_level'),
|
||
(pl.col('_raw_score') * 0.6 + 0.4 * (pl.col('_raw_score') > q75).cast(pl.Float64))
|
||
.alias('threat_score'),
|
||
]).drop('_raw_score')
|
||
|
||
valid_counts = {c: int((mat[:, i] != 0).sum()) for i, c in enumerate(avail)}
|
||
import logging
|
||
logger = logging.getLogger(__name__)
|
||
logger.info(
|
||
'[SCORE] adaptive weights from 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:
|
||
import logging
|
||
logging.getLogger(__name__).warning('[SCORE] adaptive scoring skipped: %s', exc)
|
||
|
||
return lf
|
||
|
||
|
||
# ── 12. compute_scores ─────────────────────────────────────────────────────
|
||
|
||
async def _handle_compute_scores(dataset_id: str) -> dict:
|
||
"""Compute adaptive proxy/anomaly/threat scores on entity data."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf = entry['lazyframe']
|
||
lf = _add_adaptive_scores(lf, [])
|
||
# Update schema to include new score columns (proxy_score, risk_level, threat_score)
|
||
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.store_dataset(dataset_id, lf, schema=new_schema,
|
||
metadata=entry.get('metadata', {}))
|
||
return {'status': 'scored', 'dataset_id': dataset_id}
|
||
|
||
|
||
# ── 13. filter_and_cluster ─────────────────────────────────────────────────
|
||
|
||
async def _handle_filter_and_cluster(
|
||
dataset_id: str,
|
||
filters: list[dict],
|
||
logic: str = 'and',
|
||
algorithm: str = 'hdbscan',
|
||
params: dict | None = None,
|
||
) -> dict:
|
||
"""Apply column-level filters then run clustering on the filtered result."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
full_schema = entry.get('schema', {})
|
||
|
||
# ── Step 1: Apply filters ─────────────────────────────────────────────
|
||
try:
|
||
exprs = [_build_filter_expr(
|
||
f['column'], f.get('op', 'eq'), f.get('value'),
|
||
) for f in filters]
|
||
|
||
if exprs:
|
||
if logic == 'and':
|
||
combined = exprs[0]
|
||
for e in exprs[1:]:
|
||
combined = combined & e
|
||
else:
|
||
combined = exprs[0]
|
||
for e in exprs[1:]:
|
||
combined = combined | e
|
||
filtered_lf = lf.filter(combined)
|
||
else:
|
||
filtered_lf = lf
|
||
|
||
# Estimate row count (fast streaming count)
|
||
try:
|
||
row_count = filtered_lf.select(pl.len()).collect(streaming=True).item()
|
||
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()
|
||
numeric_dtypes = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
||
pl.Float32, pl.Float64)
|
||
feature_cols = [
|
||
available_names[i] for i, dt in enumerate(available_dtypes)
|
||
if dt in numeric_dtypes and not available_names[i].startswith('_')
|
||
][:10]
|
||
|
||
if not feature_cols:
|
||
return {
|
||
'error': 'No numeric columns found for clustering after filtering',
|
||
'filtered_dataset_id': None,
|
||
'truncated': False,
|
||
}
|
||
except Exception:
|
||
raise
|
||
|
||
# ── Step 3: Store filtered dataset ────────────────────────────────────
|
||
filtered_dataset_id = _generate_id('flt')
|
||
store.store_dataset(
|
||
dataset_id=filtered_dataset_id,
|
||
lazyframe=filtered_lf,
|
||
schema=full_schema,
|
||
metadata={
|
||
'row_count': row_count,
|
||
'parent_dataset_id': dataset_id,
|
||
'filter_count': len(filters),
|
||
'logic': logic,
|
||
},
|
||
parent_id=dataset_id,
|
||
)
|
||
|
||
# ── Step 4: Run clustering (inlined from _handle_run_clustering) ───────
|
||
p = params or {}
|
||
try:
|
||
# Downsample before collect
|
||
try:
|
||
total_count = filtered_lf.select(pl.len()).collect(streaming=True).item()
|
||
except Exception:
|
||
total_count = 0
|
||
|
||
clamp_lf = filtered_lf
|
||
if total_count > 100000:
|
||
clamp_lf = filtered_lf.collect(streaming=True).sample(n=100000, seed=42).lazy()
|
||
|
||
df_features = clamp_lf.select(feature_cols).collect(streaming=True)
|
||
|
||
# Drop all-NaN columns
|
||
for col in df_features.columns:
|
||
if df_features[col].is_null().all():
|
||
df_features = df_features.drop(col)
|
||
feature_cols = [c for c in feature_cols if c != col]
|
||
|
||
if df_features.width == 0:
|
||
return {'error': 'No valid feature columns after dropping all-NaN columns',
|
||
'filtered_dataset_id': filtered_dataset_id, 'truncated': False}
|
||
|
||
data = df_features.to_numpy()
|
||
|
||
# Fill remaining NaN with column mean
|
||
if data.dtype.kind == 'f':
|
||
col_mean = np.nanmean(data, axis=0)
|
||
col_mean = np.nan_to_num(col_mean, nan=0.0)
|
||
data = np.where(np.isnan(data), col_mean, data)
|
||
|
||
# Too-few-samples guard
|
||
if len(data) < 3:
|
||
return {
|
||
'cluster_result_id': _generate_id('clust'),
|
||
'n_clusters': 0,
|
||
'n_noise': 0,
|
||
'quality_metrics': {'note': 'too few samples'},
|
||
'filtered_dataset_id': filtered_dataset_id,
|
||
'truncated': False,
|
||
}
|
||
|
||
# Low memory check
|
||
low_memory = False
|
||
try:
|
||
import psutil
|
||
mem = psutil.virtual_memory()
|
||
low_memory = mem.available < 2 * 1024 ** 3
|
||
except ImportError:
|
||
pass
|
||
|
||
if low_memory and len(data) > 100000:
|
||
rng = np.random.default_rng(42)
|
||
idx = rng.choice(len(data), 50000, replace=False)
|
||
data = data[idx]
|
||
|
||
# Downsample to 50K max
|
||
_ds_idx = None
|
||
if len(data) > 50000:
|
||
rng = np.random.default_rng(42)
|
||
_ds_idx = rng.choice(len(data), 50000, replace=False)
|
||
data = data[_ds_idx]
|
||
|
||
# Variance filtering
|
||
if data.shape[1] > 1:
|
||
variances = np.var(data, axis=0)
|
||
keep_var = np.where(variances >= 1e-10)[0]
|
||
if len(keep_var) == 0:
|
||
keep_var = np.arange(data.shape[1])
|
||
if len(keep_var) < data.shape[1]:
|
||
data = data[:, keep_var]
|
||
feature_cols = [feature_cols[i] for i in keep_var]
|
||
|
||
# Standard scale
|
||
from sklearn.preprocessing import StandardScaler
|
||
data_scaled = StandardScaler().fit_transform(data)
|
||
|
||
# Correlation filtering
|
||
if data_scaled.shape[1] > 1:
|
||
corr = np.corrcoef(data_scaled.T)
|
||
upper = np.triu(np.abs(corr), 1)
|
||
to_drop = [j for j in range(upper.shape[1]) if any(upper[:, j] > 0.95)]
|
||
if to_drop:
|
||
keep = [j for j in range(data_scaled.shape[1]) if j not in to_drop]
|
||
data_scaled = data_scaled[:, keep]
|
||
feature_cols = [feature_cols[j] for j in keep]
|
||
|
||
# SVD dimensionality reduction (if >50 features)
|
||
n_features = data_scaled.shape[1]
|
||
if n_features > 50:
|
||
from sklearn.decomposition import TruncatedSVD
|
||
data_scaled = TruncatedSVD(n_components=50, random_state=42).fit_transform(data_scaled)
|
||
|
||
# Hyperparameter auto-tuning
|
||
n = len(data_scaled)
|
||
if algorithm == 'hdbscan':
|
||
min_cluster = p.get('min_cluster_size', max(3, min(50, n // 20)))
|
||
p['min_cluster_size'] = min_cluster
|
||
p['min_samples'] = p.get('min_samples', min_cluster)
|
||
p['cluster_selection_epsilon'] = p.get('cluster_selection_epsilon', 0.0)
|
||
else:
|
||
p['min_cluster_size'] = min(
|
||
p.get('min_cluster_size', 5),
|
||
max(n // 2, 2),
|
||
)
|
||
|
||
# Fit model
|
||
from sklearn.cluster import HDBSCAN, MiniBatchKMeans
|
||
if algorithm == 'kmeans':
|
||
n_clusters_param = p.get('n_clusters', 3)
|
||
model = MiniBatchKMeans(
|
||
n_clusters=n_clusters_param, random_state=42,
|
||
batch_size=1024, n_init='auto',
|
||
)
|
||
else:
|
||
model = HDBSCAN(
|
||
min_cluster_size=p['min_cluster_size'],
|
||
min_samples=p.get('min_samples', None),
|
||
metric=p.get('metric', 'euclidean'),
|
||
cluster_selection_epsilon=p.get('cluster_selection_epsilon', 0.0),
|
||
)
|
||
|
||
labels_array = model.fit_predict(data_scaled)
|
||
n_clusters_found = int(len(set(labels_array)) - (1 if -1 in labels_array else 0))
|
||
noise_count = int((labels_array == -1).sum()) if algorithm == 'hdbscan' else 0
|
||
noise_ratio = round(noise_count / len(labels_array), 4) if algorithm == 'hdbscan' else 0.0
|
||
|
||
# Quality metrics
|
||
quality = {
|
||
'n_clusters': n_clusters_found,
|
||
'n_noise': noise_count,
|
||
'noise_ratio': noise_ratio,
|
||
'algorithm': algorithm,
|
||
}
|
||
|
||
if n_clusters_found > 1:
|
||
try:
|
||
from sklearn.metrics import silhouette_score
|
||
score = silhouette_score(data_scaled, labels_array, random_state=42)
|
||
quality['silhouette_score'] = round(float(score), 4)
|
||
except Exception:
|
||
quality['silhouette_score'] = None
|
||
try:
|
||
from sklearn.metrics import davies_bouldin_score
|
||
dbs = davies_bouldin_score(data_scaled, labels_array)
|
||
quality['davies_bouldin_score'] = round(float(dbs), 4)
|
||
except Exception:
|
||
quality['davies_bouldin_score'] = None
|
||
try:
|
||
from sklearn.metrics import calinski_harabasz_score
|
||
chs = calinski_harabasz_score(data_scaled, labels_array)
|
||
quality['calinski_harabasz_score'] = round(float(chs), 4)
|
||
except Exception:
|
||
quality['calinski_harabasz_score'] = None
|
||
|
||
# Cluster sizes
|
||
cluster_sizes = {}
|
||
for label in np.unique(labels_array):
|
||
cluster_sizes[int(label)] = int((labels_array == label).sum())
|
||
quality['cluster_sizes'] = cluster_sizes
|
||
|
||
except Exception:
|
||
raise
|
||
|
||
result_id = _generate_id('clust')
|
||
store.store_cluster_result(
|
||
result_id=result_id,
|
||
labels=labels_array.tolist(),
|
||
n_clusters=n_clusters_found,
|
||
params={
|
||
'algorithm': algorithm,
|
||
'params': p,
|
||
'random_state': 42,
|
||
'feature_columns': feature_cols,
|
||
'_ds_idx': _ds_idx.tolist() if _ds_idx is not None else None,
|
||
},
|
||
parent_dataset_id=filtered_dataset_id,
|
||
)
|
||
|
||
return {
|
||
'cluster_result_id': result_id,
|
||
'filtered_dataset_id': filtered_dataset_id,
|
||
'n_clusters': n_clusters_found,
|
||
'n_noise': noise_count,
|
||
'labels_sample': labels_array[:50].tolist(),
|
||
'quality_metrics': quality,
|
||
'feature_columns': feature_cols,
|
||
}
|
||
|
||
|
||
# ── 14. detect_anomalies ──────────────────────────────────────────────────
|
||
|
||
async def _handle_detect_anomalies(dataset_id: str, contamination: float = 0.05) -> dict:
|
||
"""Run Isolation Forest on scored entity data."""
|
||
from sklearn.ensemble import IsolationForest
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
|
||
num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')]
|
||
if not num_cols:
|
||
return {'error': 'No numeric columns for anomaly detection', 'truncated': False}
|
||
|
||
# Collect scored data (already sampled by pipeline head param)
|
||
df = lf.select(num_cols).collect(streaming=True)
|
||
mat = df.to_numpy()
|
||
mat = np.nan_to_num(mat, nan=0.0)
|
||
|
||
# Isolation Forest in chunks
|
||
n = len(mat)
|
||
chunk_size = min(50000, n)
|
||
labels = np.zeros(n, dtype=int)
|
||
for start in range(0, n, chunk_size):
|
||
end = min(start + chunk_size, n)
|
||
iso = IsolationForest(n_estimators=100, max_samples=min(256, chunk_size),
|
||
contamination=contamination, random_state=42, n_jobs=-1)
|
||
labels[start:end] = iso.fit_predict(mat[start:end])
|
||
|
||
anomaly_count = int((labels == -1).sum())
|
||
return {
|
||
'status': 'completed',
|
||
'total_entities': n,
|
||
'anomaly_count': anomaly_count,
|
||
'anomaly_rate': round(anomaly_count / max(n, 1), 4),
|
||
'dataset_id': dataset_id,
|
||
}
|
||
|
||
|
||
# ── 15. visualize_anomalies ───────────────────────────────────────────────
|
||
|
||
async def _handle_visualize_anomalies(dataset_id: str) -> dict:
|
||
"""Generate UMAP-2D embedding for anomaly visualization."""
|
||
from sklearn.preprocessing import StandardScaler
|
||
import umap
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
|
||
num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')]
|
||
if len(num_cols) < 2:
|
||
return {'error': 'Need at least 2 numeric columns for UMAP', 'truncated': False}
|
||
|
||
# Sample for UMAP visualisation (max 10K)
|
||
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], # cap for LLM context
|
||
'dataset_id': dataset_id,
|
||
}
|
||
|
||
|
||
# ── 16. validate_data ──────────────────────────────────────────────────
|
||
|
||
async def _handle_validate_data(dataset_id: str, fix_issues: bool = False) -> dict:
|
||
"""Check data quality: schema consistency, missing values, type conflicts."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
issues = []
|
||
try:
|
||
df_sample = lf.head(1000).collect(streaming=True)
|
||
for col in schema:
|
||
null_count = df_sample[col].is_null().sum()
|
||
if null_count > 500:
|
||
issues.append(f'{col}: {null_count/10:.0f}% null')
|
||
try:
|
||
unique = df_sample[col].n_unique()
|
||
if unique <= 1 and null_count == 0:
|
||
issues.append(f'{col}: constant value (no variation)')
|
||
except Exception:
|
||
pass # n_unique not supported for this column type
|
||
except Exception as exc:
|
||
raise
|
||
total_rows = lf.select(pl.len()).collect(streaming=True).item()
|
||
return {
|
||
'status': 'ok',
|
||
'total_rows': total_rows,
|
||
'columns': len(schema),
|
||
'issues_found': len(issues),
|
||
'issues': issues[:20],
|
||
'recommendation': 'Try explore_distributions for detailed column analysis' if issues else 'No issues found',
|
||
}
|
||
|
||
|
||
# ── 17. explore_distributions ──────────────────────────────────────────
|
||
|
||
async def _handle_explore_distributions(dataset_id: str, columns: list[str] = None,
|
||
max_sample: int = 10000) -> dict:
|
||
"""Per-column distribution statistics."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
|
||
|
||
if not columns:
|
||
columns = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types][:10]
|
||
# Limit columns to prevent excessive computation
|
||
if len(columns) > 20:
|
||
columns = columns[:20]
|
||
|
||
try:
|
||
n = lf.select(pl.len()).collect(streaming=True).item()
|
||
if n > max_sample:
|
||
# Use LazyFrame.sample() on the limited columns to avoid full-materialize
|
||
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)}
|
||
|
||
|
||
# ── 18. find_outliers ─────────────────────────────────────────────────
|
||
|
||
async def _handle_find_outliers(dataset_id: str, columns: list[str] = None,
|
||
threshold: float = 3.0) -> dict:
|
||
"""Find statistical outliers using IQR method."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
|
||
if not columns:
|
||
columns = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types]
|
||
try:
|
||
n = lf.select(pl.len()).collect(streaming=True).item()
|
||
sample_lf = lf if n <= 100000 else lf.collect(streaming=True).sample(n=100000, seed=42).lazy()
|
||
df = sample_lf.select(columns).collect(streaming=True)
|
||
except Exception as exc:
|
||
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 # quantile not supported for this column type
|
||
return {'outlier_counts': outlier_counts, 'total_sampled': len(df), 'threshold_iqr': threshold}
|
||
|
||
|
||
# ── 19. diagnose_clustering ───────────────────────────────────────────
|
||
|
||
async def _handle_diagnose_clustering(dataset_id: str, cluster_result_id: str) -> dict:
|
||
"""Diagnose clustering quality issues."""
|
||
from sklearn.decomposition import 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', {})
|
||
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
|
||
num_cols = [c for c in schema if schema[c].split('(')[0].strip() in numeric_types and not c.startswith('_')]
|
||
|
||
n_clusters = cluster_entry.get('n_clusters', 0)
|
||
n_noise = cluster_entry.get('n_noise', 0)
|
||
quality = cluster_entry.get('quality_metrics', {})
|
||
labels = cluster_entry.get('labels', [])
|
||
params = cluster_entry.get('params', {})
|
||
algo = params.get('algorithm', 'unknown')
|
||
feature_cols = params.get('feature_columns', [])
|
||
|
||
diagnosis = []
|
||
if n_clusters <= 1:
|
||
diagnosis.append('Only 1 cluster — data may lack structure or need feature scaling')
|
||
if n_noise > 0.5 * len(labels) if labels else False:
|
||
diagnosis.append(f'High noise ratio ({n_noise}/{len(labels)}) — try HDBSCAN with larger min_cluster_size')
|
||
sil = quality.get('silhouette_score')
|
||
if sil and sil < 0.2:
|
||
diagnosis.append(f'Low silhouette ({sil:.3f}) — clusters overlap, try KMeans or feature selection')
|
||
if len(feature_cols) < 3:
|
||
diagnosis.append(f'Only {len(feature_cols)} features — add more columns for meaningful clusters')
|
||
|
||
# SVD variance check
|
||
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 # SVD fit failed — likely constant or all-null columns
|
||
|
||
return {
|
||
'diagnosis': diagnosis,
|
||
'n_clusters': n_clusters,
|
||
'n_noise': n_noise,
|
||
'silhouette': sil,
|
||
'algorithm': algo,
|
||
'n_features': len(feature_cols),
|
||
'suggestions': [
|
||
f'Try algorithm={"kmeans" if algo == "hdbscan" else "hdbscan"}',
|
||
'Increase min_cluster_size for fewer, larger clusters',
|
||
'Use explore_distributions to check feature quality',
|
||
],
|
||
}
|
||
|
||
|
||
# ── 20. compare_datasets ─────────────────────────────────────────────
|
||
|
||
async def _handle_compare_datasets(dataset_id_a: str, dataset_id_b: str) -> dict:
|
||
"""Compare two datasets."""
|
||
store = SessionStore()
|
||
entry_a, entry_b = store.get_dataset(dataset_id_a), store.get_dataset(dataset_id_b)
|
||
if entry_a is None or entry_b is None:
|
||
return {'error': 'One or both datasets not found', 'truncated': False}
|
||
|
||
def _get_info(entry):
|
||
schema = entry.get('schema', {})
|
||
lf = entry['lazyframe']
|
||
try:
|
||
n = lf.select(pl.len()).collect(streaming=True).item()
|
||
except Exception:
|
||
n = -1
|
||
return {'schema': set(schema.keys()), 'row_count': n, 'columns': len(schema)}
|
||
|
||
info_a, info_b = _get_info(entry_a), _get_info(entry_b)
|
||
common_cols = info_a['schema'] & info_b['schema']
|
||
only_a = info_a['schema'] - info_b['schema']
|
||
only_b = info_b['schema'] - info_a['schema']
|
||
return {
|
||
'dataset_a': {'id': dataset_id_a, **info_a},
|
||
'dataset_b': {'id': dataset_id_b, **info_b},
|
||
'common_columns': len(common_cols),
|
||
'columns_only_in_a': list(only_a)[:10] if only_a else [],
|
||
'columns_only_in_b': list(only_b)[:10] if only_b else [],
|
||
'row_diff': info_a['row_count'] - info_b['row_count'],
|
||
'identical_schema': len(only_a) == 0 and len(only_b) == 0,
|
||
}
|
||
|
||
|
||
# ── 21. export_debug_sample ───────────────────────────────────────────
|
||
|
||
async def _handle_export_debug_sample(dataset_id: str, max_rows: int = 50) -> dict:
|
||
"""Export raw data sample as JSON for debugging."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
try:
|
||
df = entry['lazyframe'].head(max_rows).collect(streaming=True)
|
||
rows = df.to_dict(as_series=False)
|
||
# Convert to list-of-dicts, serialize safely
|
||
result = []
|
||
for i in range(len(df)):
|
||
row = {}
|
||
for col in df.columns:
|
||
val = rows[col][i]
|
||
if isinstance(val, (float, int, str, bool)):
|
||
row[col] = val
|
||
elif val is None:
|
||
row[col] = None
|
||
else:
|
||
row[col] = str(val)
|
||
result.append(row)
|
||
return {'rows': result, 'count': len(result), 'columns': df.columns}
|
||
except Exception as exc:
|
||
raise
|
||
|
||
|
||
# ── 22. repair_schema ─────────────────────────────────────────────────
|
||
|
||
async def _handle_repair_schema(dataset_id: str) -> dict:
|
||
"""Attempt to repair schema mismatches by normalising column names."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
# Normalise: replace hyphens/underscores/dots, lowercase
|
||
col_map = {}
|
||
for c in schema:
|
||
norm = c.lower().replace('-', '_').replace('.', '_')
|
||
col_map[c] = norm
|
||
if len(set(col_map.values())) < len(col_map):
|
||
# Conflicts exist — can't auto-repair
|
||
return {'status': 'conflict', 'message': 'Normalised names conflict', 'column_map': col_map}
|
||
# Rename columns
|
||
renames = {old: new for old, new in col_map.items() if old != new}
|
||
if renames:
|
||
lf = lf.rename(renames)
|
||
new_schema = {col_map.get(c, c): dtype for c, dtype in schema.items()}
|
||
store.store_dataset(dataset_id, lf, schema=new_schema, metadata=entry.get('metadata', {}))
|
||
return {'status': 'repaired', 'renamed': len(renames), 'column_map': col_map}
|
||
return {'status': 'no_change', 'message': 'Column names already normalised'}
|
||
|
||
|
||
# ── 23. analyze_patterns ──────────────────────────────────────────────
|
||
|
||
async def _handle_analyze_patterns(dataset_id: str, top_n: int = 10) -> dict:
|
||
"""Read-only: traffic pattern analysis."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': 'Dataset not found', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
n = lf.select(pl.len()).collect(streaming=True).item()
|
||
df = lf.head(100000).collect(streaming=True)
|
||
|
||
result = {'total_rows': n, 'sampled': min(n, 100000)}
|
||
|
||
# Top source IPs
|
||
for col_name, label in [(':ips', 'src_ips'), (':ipd', 'dst_ips'),
|
||
('snam', 'snis'), ('cnam', 'cert_names')]:
|
||
if col_name in df.columns:
|
||
try:
|
||
top = df[col_name].value_counts(sort=True).head(top_n)
|
||
result[label] = {
|
||
'top': [{'value': str(r[col_name]), 'count': int(r['count'])}
|
||
for r in top.iter_rows(named=True)],
|
||
'unique': int(df[col_name].n_unique()),
|
||
}
|
||
except Exception:
|
||
pass # value_counts not supported for column, skip
|
||
|
||
# Port distribution
|
||
for port_col in [':prs', ':prd']:
|
||
if port_col in df.columns:
|
||
try:
|
||
top = df[port_col].value_counts(sort=True).head(top_n)
|
||
result[f'{port_col}_dist'] = {
|
||
'top': [{'port': int(r[port_col]), 'count': int(r['count'])}
|
||
for r in top.iter_rows(named=True)],
|
||
'unique': int(df[port_col].n_unique()),
|
||
}
|
||
except Exception:
|
||
pass # port distribution extraction failed — skip
|
||
|
||
# TLS version breakdown
|
||
for tls_col in ['0ver']:
|
||
if tls_col in df.columns:
|
||
try:
|
||
vc = df[tls_col].value_counts(sort=True)
|
||
result['tls_versions'] = [{'version': str(r[tls_col]), 'count': int(r['count'])}
|
||
for r in vc.iter_rows(named=True)]
|
||
except Exception:
|
||
pass # TLS version value_counts not supported — skip
|
||
|
||
return result
|
||
|
||
|
||
# ── 24. analyze_temporal ─────────────────────────────────────────────
|
||
|
||
async def _handle_analyze_temporal(dataset_id: str, timestamp_column: str = None) -> dict:
|
||
"""Read-only: temporal traffic pattern analysis."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': 'Dataset not found', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
|
||
if not timestamp_column:
|
||
for cand in ['timestamp', 'time', '8dbd', '8did']:
|
||
if cand in schema:
|
||
timestamp_column = cand
|
||
break
|
||
if not timestamp_column:
|
||
return {'error': 'No timestamp column found', 'truncated': False}
|
||
|
||
n = lf.select(pl.len()).collect(streaming=True).item()
|
||
return {'total_rows': n, 'timestamp_column': timestamp_column,
|
||
'note': 'Temporal analysis requires parsed timestamps; raw values shown in explore_distributions'}
|
||
|
||
|
||
# ── 24b. analyze_fft ─────────────────────────────────────────────────
|
||
|
||
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.
|
||
|
||
Computes the Fast Fourier Transform of sorted timestamp-derived values,
|
||
returning dominant frequencies, spectral centroid, and a periodicity score.
|
||
"""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
|
||
# ── Resolve timestamp column ──────────────────────────────────────
|
||
col = None
|
||
# Try exact match first
|
||
if timestamp_column in schema:
|
||
col = timestamp_column
|
||
else:
|
||
# Try case-insensitive match
|
||
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}
|
||
|
||
# ── Extract and sort values ───────────────────────────────────────
|
||
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}
|
||
|
||
# Determine dtype and convert to numeric values
|
||
ts_series = df['__ts__']
|
||
dtype = ts_series.dtype
|
||
str_dtypes = (pl.Utf8, pl.String)
|
||
|
||
if dtype in str_dtypes:
|
||
# Try parsing as datetime string, then convert to epoch seconds
|
||
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 # ns → seconds
|
||
except Exception:
|
||
# Fallback: try to cast to float directly
|
||
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:
|
||
# Float, Int types — cast to float
|
||
try:
|
||
ts_numeric = ts_series.cast(pl.Float64)
|
||
except Exception as e:
|
||
return {'error': f'Cannot cast column to float: {e}', 'truncated': False}
|
||
|
||
# Build DataFrame with values, drop nulls, sort
|
||
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)
|
||
|
||
# ── FFT computation ────────────────────────────────────────────────
|
||
# Center by subtracting mean (DC offset removal)
|
||
values_centered = values - np.mean(values)
|
||
fft_complex = np.fft.fft(values_centered)
|
||
fft_magnitudes = np.abs(fft_complex)
|
||
freqs = np.fft.fftfreq(n)
|
||
|
||
# Take positive half (excluding DC at index 0)
|
||
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}
|
||
|
||
# ── Find top-N peaks ──────────────────────────────────────────────
|
||
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])
|
||
# Convert frequency to period (in sample-index units)
|
||
period = float(1.0 / freq) if freq > 1e-10 else float('inf')
|
||
# Map to time-domain period if data is temporal
|
||
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,
|
||
})
|
||
|
||
# ── Spectral centroid (weighted mean frequency) ────────────────────
|
||
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
|
||
|
||
# ── Periodicity score: ratio of top peak to mean magnitude ────────
|
||
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)
|
||
|
||
# ── Return decimated data for LLM context window ──────────────────
|
||
# Decimate to at most 200 frequency-magnitude pairs
|
||
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,
|
||
}
|
||
|
||
|
||
# ── 25. analyze_tls_health ────────────────────────────────────────────
|
||
|
||
async def _handle_analyze_tls_health(dataset_id: str) -> dict:
|
||
"""Read-only: TLS security assessment."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': 'Dataset not found', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
n = lf.select(pl.len()).collect(streaming=True).item()
|
||
df = lf.head(100000).collect(streaming=True)
|
||
|
||
result = {'total_rows': n, 'sampled': min(n, 100000)}
|
||
|
||
if '0ver' in df.columns:
|
||
vc = df['0ver'].value_counts(sort=True)
|
||
versions = {}
|
||
for r in vc.iter_rows(named=True):
|
||
v = str(r['0ver']).replace(' ', '')
|
||
if v in ('0303', '0304'):
|
||
versions['tls_1_2_plus'] = versions.get('tls_1_2_plus', 0) + int(r['count'])
|
||
else:
|
||
versions['tls_legacy'] = versions.get('tls_legacy', 0) + int(r['count'])
|
||
total = sum(versions.values())
|
||
result.update({
|
||
'tls_1_2_plus_pct': round(versions.get('tls_1_2_plus', 0) / max(total, 1) * 100, 1),
|
||
'tls_legacy_pct': round(versions.get('tls_legacy', 0) / max(total, 1) * 100, 1),
|
||
'tls_legacy_versions': [{'version': str(r['0ver']), 'count': int(r['count'])}
|
||
for r in vc.iter_rows(named=True)
|
||
if str(r['0ver']).replace(' ', '') not in ('0303', '0304')],
|
||
})
|
||
|
||
if 'snam' in df.columns:
|
||
empty_sni = df['snam'].is_null().sum() + df['snam'].cast(pl.Utf8).is_in(['', 'null', 'None']).sum()
|
||
result['sni_missing_pct'] = round(float(empty_sni) / max(len(df), 1) * 100, 1)
|
||
|
||
if 'cipher-suite' in df.columns:
|
||
weak_ciphers = ['rc4', 'cbc', 'null', 'export']
|
||
total_ciphers = len(df)
|
||
weak_count = 0
|
||
for w in weak_ciphers:
|
||
try:
|
||
weak_count += df['cipher-suite'].str.to_lowercase().str.contains(w).sum()
|
||
except Exception:
|
||
pass # str.contains not supported for cipher-suite column — skip
|
||
result['weak_cipher_pct'] = round(float(weak_count) / max(total_ciphers, 1) * 100, 1) if total_ciphers else 0
|
||
|
||
if result.get('tls_legacy_pct', 0) > 20:
|
||
result['warning'] = 'High proportion of legacy TLS — possible downgrade attacks'
|
||
if result.get('sni_missing_pct', 0) > 50:
|
||
result['warning'] = 'Most flows lack SNI — possible IP-direct connections (proxy/VPN)'
|
||
return result
|
||
|
||
|
||
# ── 26. analyze_geo_distribution ─────────────────────────────────────
|
||
|
||
async def _handle_analyze_geo_distribution(dataset_id: str) -> dict:
|
||
"""Read-only: geographic traffic analysis."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': 'Dataset not found', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
n = lf.select(pl.len()).collect(streaming=True).item()
|
||
df = lf.head(100000).collect(streaming=True)
|
||
|
||
result = {'total_rows': n, 'sampled': min(n, 100000)}
|
||
|
||
for geo_col, label in [('scnt', 'src_countries'), ('dcnt', 'dst_countries')]:
|
||
if geo_col in df.columns:
|
||
try:
|
||
vc = df[geo_col].value_counts(sort=True).head(15)
|
||
result[label] = [{'country': str(r[geo_col]), 'count': int(r['count'])}
|
||
for r in vc.iter_rows(named=True)]
|
||
result[f'{label}_unique'] = int(df[geo_col].n_unique())
|
||
except Exception:
|
||
pass # country distribution not computable — skip
|
||
|
||
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 # domestic/international comparison failed — skip
|
||
|
||
return result
|
||
|
||
|
||
# ── 27. analyze_entity_detail ────────────────────────────────────────
|
||
|
||
async def _handle_analyze_entity_detail(dataset_id: str, entity_value: str) -> dict:
|
||
"""Read-only: deep-dive into a single entity."""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': 'Dataset not found', 'truncated': False}
|
||
lf = entry['lazyframe']
|
||
schema = entry.get('schema', {})
|
||
|
||
# Find entity column (first non-numeric, non-internal column)
|
||
numeric_types = {'Int8', 'Int16', 'Int32', 'Int64', 'UInt8', 'UInt16', 'UInt32', 'UInt64', 'Float32', 'Float64'}
|
||
entity_col = None
|
||
for c in schema:
|
||
if c.startswith('_') or schema[c].split('(')[0].strip() in numeric_types or '.' in c:
|
||
continue
|
||
entity_col = c
|
||
break
|
||
if not entity_col:
|
||
entity_col = list(schema.keys())[0]
|
||
|
||
try:
|
||
filtered = lf.filter(pl.col(entity_col).cast(pl.Utf8).str.contains(entity_value, literal=True))
|
||
n = filtered.select(pl.len()).collect(streaming=True).item()
|
||
if n == 0:
|
||
return {'error': f"Entity '{entity_value}' not found in column '{entity_col}'", 'truncated': False}
|
||
df = filtered.collect(streaming=True)
|
||
except Exception as exc:
|
||
raise
|
||
|
||
result = {
|
||
'entity_value': entity_value,
|
||
'entity_column': entity_col,
|
||
'flow_count': n,
|
||
'columns': list(df.columns[:15]),
|
||
}
|
||
|
||
# Numeric summary
|
||
for c in df.columns:
|
||
if df[c].dtype in (pl.Float32, pl.Float64, pl.Int32, pl.Int64):
|
||
try:
|
||
result[f'{c}_stats'] = {
|
||
'mean': round(float(df[c].mean()), 2),
|
||
'min': float(df[c].min()) if df[c].min() is not None else None,
|
||
'max': float(df[c].max()) if df[c].max() is not None else None,
|
||
}
|
||
except Exception:
|
||
pass # numeric stats not computable for this column — skip
|
||
|
||
return result
|
||
|
||
|
||
# ── 28. compute_distance_matrix ──────────────────────────────────────────
|
||
|
||
def _indent_function_body(body: str, indent: str = ' ') -> str:
|
||
"""Indent each non-empty line of a code block for wrapping in a def."""
|
||
return '\n'.join(
|
||
indent + line if line.strip() else line
|
||
for line in body.split('\n')
|
||
)
|
||
|
||
|
||
async def _handle_compute_distance_matrix(
|
||
dataset_id: str,
|
||
python_function: str,
|
||
columns: Optional[list[str]] = None,
|
||
max_sample: int = 10000,
|
||
) -> dict:
|
||
"""Evaluate a user-provided Python function on each row of a dataset.
|
||
|
||
Wraps the function body as ``def distance_fn(row): <body>``, executes
|
||
it per-row in a restricted namespace (``math``, ``numpy`` only), and
|
||
returns a distance-score summary plus a sample of scored rows.
|
||
"""
|
||
store = SessionStore()
|
||
entry = store.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return {'error': f'Dataset not found: {dataset_id}', 'truncated': False}
|
||
|
||
lf: pl.LazyFrame = entry['lazyframe']
|
||
|
||
# Collect data (streaming for large datasets)
|
||
df = lf.collect(streaming=True)
|
||
|
||
# Determine columns to pass to the function
|
||
if columns is None:
|
||
numeric_dtypes = (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
|
||
pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64,
|
||
pl.Float32, pl.Float64)
|
||
columns = [c for c in df.columns
|
||
if df[c].dtype in numeric_dtypes]
|
||
|
||
# Filter to columns that actually exist
|
||
available = set(df.columns)
|
||
columns = [c for c in columns if c in available]
|
||
if not columns:
|
||
return {'error': 'No valid columns found in dataset',
|
||
'truncated': False}
|
||
|
||
n_total = len(df)
|
||
|
||
# Sample if larger than max_sample
|
||
if n_total > max_sample:
|
||
df_sampled = df.sample(n=max_sample, seed=42)
|
||
n_evaluated = max_sample
|
||
else:
|
||
df_sampled = df
|
||
n_evaluated = n_total
|
||
|
||
# Build list-of-dicts for the selected columns
|
||
rows = df_sampled.select(columns).to_dicts()
|
||
|
||
# ── Safe execution ────────────────────────────────────────────────────
|
||
safe_builtins = {
|
||
'abs': abs, 'len': len, 'range': range,
|
||
'min': min, 'max': max, 'sum': sum, 'round': round,
|
||
'int': int, 'float': float, 'str': str, 'bool': bool,
|
||
'list': list, 'dict': dict, 'tuple': tuple,
|
||
'isinstance': isinstance, 'type': type,
|
||
'True': True, 'False': False, 'None': None,
|
||
'enumerate': enumerate, 'zip': zip, 'map': map,
|
||
'filter': filter, 'sorted': sorted, 'reversed': reversed,
|
||
}
|
||
|
||
safe_ns = {
|
||
'__builtins__': safe_builtins,
|
||
'math': __import__('math'),
|
||
'np': __import__('numpy'),
|
||
}
|
||
|
||
# Wrap the user's function body in a def statement
|
||
body = python_function.strip()
|
||
if body.startswith('def '):
|
||
# User provided the full function definition
|
||
full_code = body
|
||
else:
|
||
indented = _indent_function_body(body)
|
||
full_code = f'def distance_fn(row):\n{indented}'
|
||
|
||
try:
|
||
exec(full_code, safe_ns)
|
||
# Resolve: the function name might be 'distance_fn' or whatever the user named it
|
||
distance_fn = safe_ns.get('distance_fn')
|
||
if distance_fn is None:
|
||
# Fallback: try to find any callable in safe_ns
|
||
for _name, _obj in safe_ns.items():
|
||
if callable(_obj) and _name not in ('math', 'np', '__builtins__'):
|
||
distance_fn = _obj
|
||
break
|
||
if distance_fn is None:
|
||
return {'error': 'Could not locate the compiled function',
|
||
'truncated': False}
|
||
except SyntaxError as e:
|
||
return {'error': f'Function syntax error: {e}', 'truncated': False}
|
||
except Exception as e:
|
||
return {'error': f'Function compilation failed: {e}',
|
||
'truncated': False}
|
||
|
||
# ── Apply function per row ────────────────────────────────────────────
|
||
scores: list[float | None] = []
|
||
errors = 0
|
||
for row in rows:
|
||
try:
|
||
score = distance_fn(row)
|
||
if isinstance(score, (int, float)):
|
||
scores.append(float(score))
|
||
else:
|
||
scores.append(None)
|
||
errors += 1
|
||
except Exception:
|
||
scores.append(None)
|
||
errors += 1
|
||
|
||
valid_scores = [s for s in scores if s is not None]
|
||
if not valid_scores:
|
||
return {'error': 'No valid scores computed — all rows raised errors',
|
||
'truncated': False}
|
||
|
||
# ── Distance summary ──────────────────────────────────────────────────
|
||
arr = np.array(valid_scores)
|
||
summary = {
|
||
'min': float(np.min(arr)),
|
||
'max': float(np.max(arr)),
|
||
'mean': float(np.mean(arr)),
|
||
'std': float(np.std(arr)),
|
||
'count': len(arr),
|
||
'errors': errors,
|
||
'total_rows_in_sample': len(rows),
|
||
'matrix_shape': [len(rows), len(columns)],
|
||
}
|
||
|
||
# Sample of scored rows (first 30, keep score + first 3 values)
|
||
scored_sample = []
|
||
for i, row in enumerate(rows[:30]):
|
||
if i < len(scores):
|
||
scored_sample.append({
|
||
'index': i,
|
||
'values': {c: row.get(c) for c in columns[:3]},
|
||
'score': scores[i],
|
||
})
|
||
|
||
return {
|
||
'summary': summary,
|
||
'scored_sample': scored_sample,
|
||
'columns_used': columns,
|
||
'total_rows_in_dataset': n_total,
|
||
'rows_evaluated': n_evaluated,
|
||
}
|
||
|
||
|
||
# Auto-profile all public functions in this module
|
||
from analysis.profile_util import auto_profile_module # noqa: E402
|
||
auto_profile_module(__name__)
|