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