tianxuan/analysis/models.py

from django.db import models, connection


class AnalysisRun(models.Model):
    """Tracks a single analysis pipeline execution."""
    STATUS_CHOICES = [
        ('pending', 'Pending'),
        ('loading', 'Loading Data'),
        ('profiling', 'Profiling'),
        ('aggregating', 'Building Entity Profiles'),
        ('clustering', 'Clustering'),
        ('extracting', 'Extracting Features'),
        ('completed', 'Completed'),
        ('failed', 'Failed'),
    ]

    RUN_TYPE_CHOICES = [
        ('upload', '手动上传'),
        ('manual', '手动分析'),
        ('auto', 'LLM自动'),
    ]

    display_id = models.IntegerField(unique=True, null=True, blank=True,
                                      help_text="User-facing numeric ID (recycles gaps)")
    csv_glob = models.CharField(max_length=1024, help_text="Glob pattern for CSV files")
    config = models.JSONField(default=dict, blank=True, help_text="YAML config snapshot")
    status = models.CharField(max_length=32, choices=STATUS_CHOICES, default='pending')
    run_type = models.CharField(max_length=16, choices=RUN_TYPE_CHOICES, default='upload')
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)
    error_message = models.TextField(blank=True, default='')
    progress_pct = models.IntegerField(default=0)
    progress_msg = models.CharField(max_length=256, default='', blank=True)
    run_log = models.TextField(blank=True, default='')

    # Summary stats populated after completion
    total_flows = models.IntegerField(null=True, blank=True)
    entity_count = models.IntegerField(null=True, blank=True)
    cluster_count = models.IntegerField(null=True, blank=True)
    sqlite_table = models.CharField(max_length=256, blank=True, default='',
                                     help_text="SQLite table name for persisted raw data")

    # LLM auto-analysis fields
    llm_thinking = models.TextField(blank=True, default='',
                                     help_text="Accumulated LLM reasoning / thinking text")
    tool_calls_json = models.JSONField(default=list, blank=True,
                                        help_text="Structured tool call list: [{step, name, input, output}, ...]")

    class Meta:
        ordering = ['-created_at']

    def __str__(self):
        return f"Run #{self.display_id or self.id} ({self.created_at:%Y-%m-%d %H:%M}) - {self.status}"

    def save(self, *args, **kwargs):
        if self._state.adding and self.display_id is None:
            table = self._meta.db_table
            with connection.cursor() as cursor:
                cursor.execute(f"""
                    SELECT COALESCE(MIN(t1.display_id) + 1, 1)
                    FROM {table} t1
                    WHERE t1.display_id IS NOT NULL
                      AND NOT EXISTS (
                        SELECT 1 FROM {table} t2
                        WHERE t2.display_id = t1.display_id + 1
                      )
                """)
                self.display_id = cursor.fetchone()[0]
        super().save(*args, **kwargs)


class ClusterResult(models.Model):
    """Stores clustering results for a single analysis run."""
    run = models.ForeignKey(AnalysisRun, on_delete=models.CASCADE, related_name='clusters')
    cluster_label = models.IntegerField(help_text="Cluster label (-1 for noise)")
    size = models.IntegerField(help_text="Number of entities in this cluster")
    proportion = models.FloatField(null=True, blank=True, help_text="Proportion of total entities")
    noise_ratio = models.FloatField(null=True, blank=True, help_text="Noise ratio of the clustering")
    silhouette_score = models.FloatField(null=True, blank=True, help_text="Silhouette score (sampled if large)")

    class Meta:
        unique_together = ['run', 'cluster_label']

    def __str__(self):
        return f"Cluster #{self.cluster_label} (n={self.size})"


class EntityProfile(models.Model):
    """One row per entity (user/host) per analysis run."""
    run = models.ForeignKey(AnalysisRun, on_delete=models.CASCADE, related_name='entities')
    entity_value = models.CharField(max_length=512, help_text="Entity identifier (IP, domain, etc.)")
    cluster_label = models.IntegerField(null=True, blank=True, help_text="Assigned cluster (-1=noise)")
    cluster = models.ForeignKey(ClusterResult, on_delete=models.SET_NULL, null=True, blank=True,
                                 related_name='entities')
    feature_json = models.JSONField(default=dict, blank=True,
                                     help_text="Aggregated entity features as dict")
    embedding_x = models.FloatField(null=True, blank=True, help_text="PCA-2D X coordinate")
    embedding_y = models.FloatField(null=True, blank=True, help_text="PCA-2D Y coordinate")

    class Meta:
        unique_together = ['run', 'entity_value']

    def __str__(self):
        return f"Entity: {self.entity_value} (Cluster {self.cluster_label})"


class ClusterFeature(models.Model):
    """Per-cluster statistical features and distinguishing scores."""
    cluster = models.ForeignKey(ClusterResult, on_delete=models.CASCADE, related_name='features')
    feature_name = models.CharField(max_length=256)
    mean = models.FloatField(null=True, blank=True)
    std = models.FloatField(null=True, blank=True)
    median = models.FloatField(null=True, blank=True)
    p25 = models.FloatField(null=True, blank=True)
    p75 = models.FloatField(null=True, blank=True)
    missing_rate = models.FloatField(null=True, blank=True)
    distinguishing_score = models.FloatField(null=True, blank=True,
                                              help_text="Z-score or ANOVA F-score")
    distinguishing_method = models.CharField(max_length=16, default='zscore')

    class Meta:
        unique_together = ['cluster', 'feature_name']

    def __str__(self):
        return f"{self.feature_name} (score={self.distinguishing_score:.3f})"