memory-system/lifecycle/pipeline.py

"""Pipeline Automation - Auto-trigger compression and cleanup based on configurable rules."""
import logging
from typing import Optional

logger = logging.getLogger(__name__)


class PipelineConfig:
    """Pipeline configuration for a team."""

    def __init__(self, team_id: str, compress_every_n: int = 0,
                 compress_target_count: int = 5,
                 cleanup_idle_days: int = 0,
                 cleanup_min_importance: float = 0.2,
                 enabled: bool = True, **kwargs):
        self.team_id = team_id
        self.compress_every_n = compress_every_n  # 0 = disabled
        self.compress_target_count = compress_target_count
        self.cleanup_idle_days = cleanup_idle_days  # 0 = disabled
        self.cleanup_min_importance = cleanup_min_importance
        self.enabled = enabled
        self.warmup_max_memories = kwargs.get("warmup_max_memories", 0)  # 0 = disabled
        self.warmup_compress_every_n = kwargs.get("warmup_compress_every_n", 1)

    def to_dict(self):
        return {
            "team_id": self.team_id,
            "compress_every_n": self.compress_every_n,
            "compress_target_count": self.compress_target_count,
            "cleanup_idle_days": self.cleanup_idle_days,
            "cleanup_min_importance": self.cleanup_min_importance,
            "enabled": self.enabled,
            "warmup_max_memories": self.warmup_max_memories,
            "warmup_compress_every_n": self.warmup_compress_every_n,
        }


class Pipeline:
    """Auto-pipeline: trigger compression and cleanup based on rules."""

    def __init__(self, mysql_store, redis_cache, compressor, cleaner):
        self.mysql = mysql_store
        self.redis = redis_cache
        self.compressor = compressor
        self.cleaner = cleaner

    def get_config(self, team_id: str) -> Optional[PipelineConfig]:
        row = self.mysql.get_pipeline_config(team_id)
        if not row:
            return None
        return PipelineConfig(
            team_id=row["team_id"],
            compress_every_n=row.get("compress_every_n", 0),
            compress_target_count=row.get("compress_target_count", 5),
            cleanup_idle_days=row.get("cleanup_idle_days", 0),
            cleanup_min_importance=row.get("cleanup_min_importance", 0.2),
            enabled=bool(row.get("enabled", 1)),
            warmup_max_memories=row.get("warmup_max_memories", 0),
            warmup_compress_every_n=row.get("warmup_compress_every_n", 1),
        )

    def set_config(self, team_id: str, **kwargs) -> PipelineConfig:
        existing = self.mysql.get_pipeline_config(team_id)
        if existing:
            self.mysql.update_pipeline_config(team_id, **kwargs)
        else:
            self.mysql.create_pipeline_config(team_id, **kwargs)
        return self.get_config(team_id)

    def check_after_wm_write(self, agent_id: str) -> Optional[dict]:
        agent = self.mysql.get_agent(agent_id)
        if not agent:
            return None
        team_id = agent["team_id"]
        config = self.get_config(team_id)
        if not config or not config.enabled:
            return None
        if config.compress_every_n <= 0:
            return None
        wm_count = self.redis.get_working_memory_count(agent_id)
        # Warm-up: use lower threshold for new agents
        effective_threshold = config.compress_every_n
        if config.warmup_max_memories > 0:
            personal_count = len(self.mysql.get_personal_memories_by_agent(agent_id, limit=config.warmup_max_memories + 1))
            if personal_count < config.warmup_max_memories:
                effective_threshold = config.warmup_compress_every_n
        if wm_count < effective_threshold:
            return None
        logger.info(f"Auto-compress triggered for {agent_id}: WM {wm_count} >= {config.compress_every_n}")
        try:
            result = self.compressor.compress(agent_id, max_items=config.compress_target_count)
            result["auto_triggered"] = True
            return result
        except Exception as e:
            logger.error(f"Auto-compress failed for {agent_id}: {e}")
            return None

    def run_cleanup(self, team_id: str = None) -> Optional[dict]:
        configs = []
        if team_id:
            config = self.get_config(team_id)
            if config and config.enabled and config.cleanup_idle_days > 0:
                configs.append(config)
        else:
            all_configs = self.mysql.get_all_pipeline_configs()
            for row in all_configs:
                config = PipelineConfig(
                    team_id=row["team_id"],
                    compress_every_n=row.get("compress_every_n", 0),
                    compress_target_count=row.get("compress_target_count", 5),
                    cleanup_idle_days=row.get("cleanup_idle_days", 0),
                    cleanup_min_importance=row.get("cleanup_min_importance", 0.2),
                    enabled=bool(row.get("enabled", 1)),
                    warmup_max_memories=row.get("warmup_max_memories", 0),
                    warmup_compress_every_n=row.get("warmup_compress_every_n", 1),
                )
                if config.enabled and config.cleanup_idle_days > 0:
                    configs.append(config)
        if not configs:
            return None
        results = []
        for config in configs:
            try:
                result = self.cleaner.cleanup(
                    team_id=config.team_id,
                    max_age_days=config.cleanup_idle_days,
                    min_importance=config.cleanup_min_importance,
                )
                result["team_id"] = config.team_id
                result["auto_triggered"] = True
                results.append(result)
                logger.info(f"Auto-cleanup for {config.team_id}: {result}")
            except Exception as e:
                logger.error(f"Auto-cleanup failed for {config.team_id}: {e}")
                results.append({"team_id": config.team_id, "error": str(e)})
        return {"results": results}