GEPA 自我进化机制：让 AIOps 系统越用越聪明的秘密，从算法原理到工程实现全解析

"""
GEPA 记忆管理系统
短期记忆（Redis）+ 长期记忆（PostgreSQL + ChromaDB）
"""
import json
import hashlib
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from typing import Optional

import redis
import psycopg2
import chromadb


@dataclass
class MemoryEntry:
    """记忆条目"""
    id: str
    content: str
    category: str  # experience / skill / pitfall
    tags: list[str] = field(default_factory=list)
    success_rate: float = 0.0
    usage_count: int = 0
    created_at: datetime = field(default_factory=datetime.now)
    updated_at: datetime = field(default_factory=datetime.now)


class MemoryManager:
    """双层记忆管理器"""

    def __init__(self, redis_url: str, pg_dsn: str,
                 chroma_path: str):
        # 短期记忆：Redis（TTL 1 小时）
        self.redis = redis.from_url(redis_url)
        self.redis_ttl = timedelta(hours=1)

        # 长期记忆：PostgreSQL（结构化存储）
        self.pg = psycopg2.connect(pg_dsn)

        # 长期记忆：ChromaDB（向量检索）
        self.chroma = chromadb.PersistentClient(path=chroma_path)
        self.collection = self.chroma.get_or_create_collection(
            name="ops_experience",
            metadata={"hnsw:space": "cosine"}
        )

    # ---- 短期记忆操作 ----

    def set_short_term(self, session_id: str,
                       key: str, value: dict):
        """写入短期记忆"""
        redis_key = f"short:{session_id}:{key}"
        self.redis.setex(
            redis_key,
            int(self.redis_ttl.total_seconds()),
            json.dumps(value, ensure_ascii=False)
        )

    def get_short_term(self, session_id: str,
                       key: str) -> Optional[dict]:
        """读取短期记忆"""
        redis_key = f"short:{session_id}:{key}"
        data = self.redis.get(redis_key)
        return json.loads(data) if data else None

    # ---- 长期记忆操作 ----

    def save_long_term(self, entry: MemoryEntry):
        """保存到长期记忆（双写：PG + ChromaDB）"""
        # 写入 PostgreSQL
        with self.pg.cursor() as cur:
            cur.execute("""
                INSERT INTO memory_entries
                    (id, content, category, tags,
                     success_rate, usage_count, created_at, updated_at)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (id) DO UPDATE SET
                    usage_count = memory_entries.usage_count + 1,
                    success_rate = %s,
                    updated_at = %s
            """, (
                entry.id, entry.content, entry.category,
                json.dumps(entry.tags), entry.success_rate,
                entry.usage_count, entry.created_at,
                entry.updated_at,
                entry.success_rate, entry.updated_at
            ))
        self.pg.commit()

        # 写入 ChromaDB（向量索引）
        self.collection.upsert(
            ids=[entry.id],
            documents=[entry.content],
            metadatas=[{
                "category": entry.category,
                "success_rate": entry.success_rate,
                "tags": ",".join(entry.tags),
            }]
        )

    def search_long_term(self, query: str,
                         top_k: int = 5) -> list[MemoryEntry]:
        """向量检索长期记忆"""
        results = self.collection.query(
            query_texts=[query],
            n_results=top_k,
            where={"success_rate": {"$gte": 0.6}}
        )

        entries = []
        for i, doc_id in enumerate(results["ids"][0]):
            meta = results["metadatas"][0][i]
            entries.append(MemoryEntry(
                id=doc_id,
                content=results["documents"][0][i],
                category=meta["category"],
                tags=meta.get("tags", "").split(","),
                success_rate=meta.get("success_rate", 0.0),
            ))
        return entries

class GEPACycle:
    """GEPA 四步闭环引擎"""

    def __init__(self, memory: MemoryManager, llm_client):
        self.memory = memory
        self.llm = llm_client

    def generate(self, session_id: str,
                 alert_context: dict) -> dict:
        """Generate: 生成候选方案"""
        # 1. 检索长期记忆中的相似经验
        similar = self.memory.search_long_term(
            query=alert_context["description"],
            top_k=3
        )

        # 2. 读取短期记忆中的当前会话上下文
        session_ctx = self.memory.get_short_term(
            session_id, "context"
        ) or {}

        # 3. 组装 prompt，让 LLM 基于经验生成方案
        experience_refs = "\n".join([
            f"- [{e.category}] {e.content} "
            f"(成功率: {e.success_rate:.0%})"
            for e in similar
        ])

        prompt = f"""
        当前告警: {alert_context['description']}
        告警级别: {alert_context['severity']}
        服务名称: {alert_context['service']}

        历史相似经验:
        {experience_refs}

        当前会话上下文:
        {json.dumps(session_ctx, ensure_ascii=False)}

        请基于以上经验，生成修复方案。
        输出格式: JSON，包含 steps 和 risk_level。
        """

        response = self.llm.chat(prompt)
        plan = json.loads(response)

        # 4. 写入短期记忆
        self.memory.set_short_term(
            session_id, "generated_plan", plan
        )

        return plan

    def adapt(self, session_id: str,
              execution_result: dict):
        """Adapt: 适应更新，将经验写入长期记忆"""
        plan = self.memory.get_short_term(
            session_id, "generated_plan"
        )
        alert_ctx = self.memory.get_short_term(
            session_id, "alert_context"
        )

        success = execution_result.get("status") == "success"

        # 计算成功率（结合历史数据）
        entry_id = hashlib.md5(
            (alert_ctx["description"] +
             str(plan.get("steps", []))).encode()
        ).hexdigest()[:12]

        # 构建记忆条目
        if success:
            content = (
                f"告警: {alert_ctx['description']}\n"
                f"方案: {json.dumps(plan['steps'], ensure_ascii=False)}\n"
                f"结果: 修复成功，耗时 {execution_result.get('duration_sec')}s"
            )
            category = "experience"
        else:
            content = (
                f"告警: {alert_ctx['description']}\n"
                f"方案: {json.dumps(plan['steps'], ensure_ascii=False)}\n"
                f"结果: 修复失败，原因: {execution_result.get('error')}\n"
                f"避坑建议: {execution_result.get('suggestion', '无')}"
            )
            category = "pitfall"

        entry = MemoryEntry(
            id=entry_id,
            content=content,
            category=category,
            tags=[
                alert_ctx.get("service", ""),
                alert_ctx.get("severity", ""),
                plan.get("risk_level", ""),
            ],
            success_rate=1.0 if success else 0.0,
            usage_count=1,
        )

        self.memory.save_long_term(entry)

class SkillAutoGenerator:
    """技能自动生成器
    当同一场景的经验累积 3 次以上，
    自动提炼为结构化技能
    """

    def __init__(self, memory: MemoryManager, llm_client):
        self.memory = memory
        self.llm = llm_client

    def check_and_generate(self, alert_type: str):
        """检查是否有足够经验生成技能"""
        # 查询该告警类型的成功经验
        similar = self.memory.search_long_term(
            query=alert_type, top_k=10
        )

        success_entries = [
            e for e in similar
            if e.category == "experience"
            and e.success_rate >= 0.8
        ]

        # 累积 3 次以上成功经验，触发技能生成
        if len(success_entries) >= 3:
            return self._generate_skill(
                alert_type, success_entries
            )
        return None

    def _generate_skill(self, alert_type: str,
                        entries: list[MemoryEntry]) -> dict:
        """基于经验生成技能"""
        experience_text = "\n\n".join([
            f"经验{i+1}: {e.content}"
            for i, e in enumerate(entries)
        ])

        prompt = f"""
        告警类型: {alert_type}

        以下为多次成功修复的经验总结:
        {experience_text}

        请提炼为标准化的运维技能，输出 YAML 格式。
        包含: name, description, params, steps, rollback。
        """

        skill_yaml = self.llm.chat(prompt)

        # 保存为技能记忆
        skill_entry = MemoryEntry(
            id=f"skill-{alert_type}",
            content=skill_yaml,
            category="skill",
            tags=[alert_type, "auto-generated"],
            success_rate=1.0,
            usage_count=0,
        )
        self.memory.save_long_term(skill_entry)

        return {"skill_name": alert_type, "yaml": skill_yaml}

def save_long_term(self, entry: MemoryEntry):
    # 准入门槛：相似度 < 0.7 的才入库（避免重复）
    existing = self.collection.query(
        query_texts=[entry.content],
        n_results=1,
    )
    if existing["distances"][0]:
        min_dist = min(existing["distances"][0])
        if min_dist < 0.3:  # 高度相似，跳过
            return

    # 正常保存逻辑
    ...

def cleanup_low_value(self, days: int = 90):
    """定期清理低价值记忆"""
    with self.pg.cursor() as cur:
        cur.execute("""
            DELETE FROM memory_entries
            WHERE success_rate < 0.3
              AND usage_count < 2
              AND updated_at < NOW() - INTERVAL '%s days'
        """, (days,))
        deleted = cur.rowcount
    self.pg.commit()
    return deleted

def get_short_term(self, session_id: str,
                   key: str) -> Optional[dict]:
    redis_key = f"short:{session_id}:{key}"
    data = self.redis.get(redis_key)
    if data:
        # 每次读取时自动续期 1 小时
        self.redis.expire(
            redis_key,
            int(self.redis_ttl.total_seconds())
        )
    return json.loads(data) if data else None

def _generate_skill(self, alert_type, entries):
    skill_yaml = self.llm.chat(prompt)

    # 强制加入安全约束模板
    skill_dict = yaml.safe_load(skill_yaml)
    safety_check = {
        "safety": {
            "require_approval": True,  # 生成的技能默认需审批
            "blocked_actions": ["rm -rf", "DROP TABLE", "DELETE WITHOUT WHERE"],
            "max_risk_level": "medium",
        }
    }
    skill_dict.update(safety_check)

    # 写入待审批队列，而非直接生效
    self._save_to_approval_queue(skill_dict)
    return skill_dict