Agent/agent/memory/stores.py

"""
Memory Implementation - 内存存储实现

用于测试和简单场景，数据不持久化

MemoryTraceStore 已移动到 agent.execution.store
"""

from typing import Dict, List, Optional, Any
from datetime import datetime

from agent.memory.models import Experience, Skill


class MemoryMemoryStore:
    """内存 Memory 存储（Experience + Skill）"""

    def __init__(self):
        self._experiences: Dict[str, Experience] = {}
        self._skills: Dict[str, Skill] = {}

    # ===== Experience =====

    async def add_experience(self, exp: Experience) -> str:
        self._experiences[exp.exp_id] = exp
        return exp.exp_id

    async def get_experience(self, exp_id: str) -> Optional[Experience]:
        return self._experiences.get(exp_id)

    async def search_experiences(
        self,
        scope: str,
        context: str,
        limit: int = 10
    ) -> List[Experience]:
        # 简单实现：按 scope 过滤，按 confidence 排序
        experiences = [
            e for e in self._experiences.values()
            if e.scope == scope
        ]
        experiences.sort(key=lambda e: e.confidence, reverse=True)
        return experiences[:limit]

    async def update_experience_stats(
        self,
        exp_id: str,
        success: bool
    ) -> None:
        exp = self._experiences.get(exp_id)
        if exp:
            exp.usage_count += 1
            if success:
                # 更新成功率
                total_success = exp.success_rate * (exp.usage_count - 1) + (1 if success else 0)
                exp.success_rate = total_success / exp.usage_count
            exp.updated_at = datetime.now()

    # ===== Skill =====

    async def add_skill(self, skill: Skill) -> str:
        self._skills[skill.skill_id] = skill
        return skill.skill_id

    async def get_skill(self, skill_id: str) -> Optional[Skill]:
        return self._skills.get(skill_id)

    async def get_skill_tree(self, scope: str) -> List[Skill]:
        return [s for s in self._skills.values() if s.scope == scope]

    async def search_skills(
        self,
        scope: str,
        context: str,
        limit: int = 5
    ) -> List[Skill]:
        # 简单实现：按 scope 过滤
        skills = [s for s in self._skills.values() if s.scope == scope]
        return skills[:limit]


class MemoryStateStore:
    """内存状态存储"""

    def __init__(self):
        self._state: Dict[str, Dict[str, Any]] = {}

    async def get(self, key: str) -> Optional[Dict[str, Any]]:
        return self._state.get(key)

    async def set(
        self,
        key: str,
        value: Dict[str, Any],
        ttl: Optional[int] = None
    ) -> None:
        # 内存实现忽略 ttl
        self._state[key] = value

    async def update(self, key: str, **updates) -> None:
        if key in self._state:
            self._state[key].update(updates)

    async def delete(self, key: str) -> None:
        self._state.pop(key, None)