tianxuan/analysis/management/commands/run_pipeline.py

"""Django management command: ``python manage.py run_pipeline``

Runs the complete TLS analysis pipeline from CLI, no LLM required.
Usage:  uv run python manage.py run_pipeline <csv_glob> [--entity-col NAME] [--algo hdbscan|kmeans]
"""
import os, sys, json, asyncio
from django.core.management.base import BaseCommand, CommandError
from django.utils import timezone

class Command(BaseCommand):
    help = 'Run the full TLS analysis pipeline from CLI (no LLM needed)'

    def add_arguments(self, parser):
        parser.add_argument('csv_glob', type=str, help='Glob pattern for CSV files')
        parser.add_argument('--entity-col', type=str, default=None,
                            help='Entity column (auto-detect if omitted)')
        parser.add_argument('--algo', type=str, default='hdbscan', choices=['hdbscan', 'kmeans'],
                            help='Clustering algorithm (default: hdbscan)')
        parser.add_argument('--output', type=str, default=None,
                            help='Output directory for exports (default: no export)')

    def handle(self, *args, **options):
        import django; django.setup()
        from analysis.models import AnalysisRun
        from analysis.session_store import SessionStore
        from analysis.data_loader import load_csv_directory
        from analysis.entity_detector import detect_entity_column
        from analysis.entity_aggregator import aggregate_by_entity

        csv_glob = options['csv_glob']
        entity_col = options['entity_col']
        algo = options['algo']
        output = options['output']

        store = SessionStore()
        store.drop_all()

        # ── Step 1: Load CSV ─────────────────────────────────────────────
        self.stdout.write(f'[1/5] 加载 CSV: {csv_glob}')
        lf, schema, row_count, file_count, memory_mb = load_csv_directory(csv_glob)
        ds_id = 'manual_ds'
        store.store_dataset(ds_id, lf, schema=schema, metadata={
            'row_count': row_count, 'file_count': file_count, 'csv_glob': csv_glob,
        })
        self.stdout.write(self.style.SUCCESS(f'  → {row_count} rows, {file_count} files, {memory_mb:.1f} MB'))

        # ── Step 2: Detect entity column ────────────────────────────────
        self.stdout.write('[2/5] 检测实体列')
        if entity_col:
            used_col = entity_col
            self.stdout.write(f'  → 使用指定列: {used_col}')
        else:
            result = detect_entity_column(ds_id)
            candidates = result.get('candidates', [])
            used_col = result.get('recommended')
            if not used_col:
                self.stderr.write(self.style.ERROR('  ✗ 未检测到实体列，请用 --entity-col 手动指定'))
                return
            self.stdout.write(f'  → 自动检测: {used_col} (共 {len(candidates)} 候选)')

        # ── Step 3: Aggregate by entity ──────────────────────────────────
        self.stdout.write('[3/5] 实体聚合')
        agg_lf, feature_columns = aggregate_by_entity(lf, used_col, schema)
        df_agg = agg_lf.collect(streaming=True)
        entity_ds_id = 'manual_entity'
        store.store_dataset(entity_ds_id, agg_lf, schema=dict(zip(feature_columns, ['']*len(feature_columns))),
                            metadata={'row_count': len(df_agg), 'entity_column': used_col, 'csv_glob': csv_glob})
        self.stdout.write(self.style.SUCCESS(f'  → {len(df_agg)} 个实体, {len(feature_columns)} 维特征'))

        # ── Step 4: Clustering ──────────────────────────────────────────
        self.stdout.write(f'[4/5] 聚类 ({algo})')
        from analysis.tool_registry import _handle_run_clustering
        # 只选用户面特征列（排除内部 _ 开头列 + 时间列）
        user_features = [
            c for c in feature_columns
            if not c.startswith('_') and c not in ('first_seen', 'last_seen', 'src_ip')
        ][:10]
        self.stdout.write(f'  → 聚类特征: {user_features}')
        result = asyncio.run(_handle_run_clustering(
            dataset_id=entity_ds_id,
            cluster_columns=user_features,
            algorithm=algo, params={}, random_state=42,
        ))
        if 'error' in result:
            self.stderr.write(self.style.ERROR(f'  ✗ 聚类失败: {result["error"]}'))
            return
        cluster_id = result.get('cluster_result_id', '')
        n_clusters = result.get('n_clusters', 0)
        quality = result.get('quality_metrics', {})
        self.stdout.write(self.style.SUCCESS(
            f'  → {n_clusters} 个簇, 噪声比 {quality.get("noise_ratio", "N/A")}'
        ))
        if quality.get('silhouette_score'):
            self.stdout.write(f'  → Silhouette: {quality["silhouette_score"]:.4f}')

        # ── Step 5: Feature extraction ───────────────────────────────────
        self.stdout.write('[5/5] 特征提取')
        from analysis.tool_registry import _handle_extract_features
        feat_result = asyncio.run(_handle_extract_features(
            dataset_id=entity_ds_id, cluster_result_id=cluster_id,
            top_k=10, method='zscore', save_to_db=True,
        ))
        if 'error' in feat_result:
            self.stderr.write(self.style.ERROR(f'  ✗ 特征提取失败: {feat_result["error"]}'))
        else:
            self.stdout.write(self.style.SUCCESS(f'  → {feat_result.get("n_features", 0)} 个特征已保存到数据库'))

        # ── Export ───────────────────────────────────────────────────────
        if output:
            from analysis.tool_registry import _handle_export_results
            os.makedirs(output, exist_ok=True)
            asyncio.run(_handle_export_results(result_id=entity_ds_id, output_path=output, format='csv', overwrite=True))
            self.stdout.write(self.style.SUCCESS(f'  → 导出到 {output}'))

        # ── Summary ──────────────────────────────────────────────────────
        run_record = AnalysisRun.objects.create(
            csv_glob=csv_glob, status='completed',
            total_flows=row_count, entity_count=len(df_agg),
            cluster_count=n_clusters, entity_column=used_col,
        )

        # ── Step 6: PCA-2D embedding (after run_record exists) ─────────────
        try:
            from sklearn.preprocessing import StandardScaler
            from sklearn.decomposition import PCA
            import polars as pl, numpy as np
            num_cols = [c for c in df_agg.columns if df_agg[c].dtype in (
                pl.Int8,pl.Int16,pl.Int32,pl.Int64,pl.UInt8,pl.UInt16,pl.UInt32,pl.UInt64,pl.Float32,pl.Float64)]
            if len(num_cols) >= 2 and len(df_agg) > 2:
                mat = np.nan_to_num(StandardScaler().fit_transform(df_agg.select(num_cols).to_numpy()), nan=0.0)
                coords = PCA(n_components=2, random_state=42).fit_transform(mat)
                non_num = [c for c in df_agg.columns if c not in num_cols and not c.startswith('_')]
                ecol = non_num[0] if non_num else df_agg.columns[0]
                from analysis.models import EntityProfile as EP, ClusterResult as CR
                for i, row in enumerate(df_agg.iter_rows(named=True)):
                    if i >= len(coords): break
                    ev = str(row.get(ecol, ''))
                    if ev:
                        EP.objects.get_or_create(entity_value=ev, defaults={
                            'run_id': run_record.id,
                            'cluster_label': 0,
                            'embedding_x': float(coords[i,0]),
                            'embedding_y': float(coords[i,1])})
        except Exception as e:
            self.stderr.write(f'  ⚠ PCA降维跳过: {e}')
        self.stdout.write(self.style.SUCCESS(
            f'\n*** 分析完成! Run ID: #{run_record.id}'
            f' 查看 Django: http://127.0.0.1:8000/runs/{run_record.id}/'
        ))