long_articles_agentic_supply/supply/domain/agent/memory.py

"""Memory 抽象 —— Agent 的记忆系统

两种实现：
  - ConversationMemory: 列表内存储存最近 N 轮对话
  - VectorMemory (future): 基于 Milvus 的长期记忆，语义检索
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


# ==================== 消息模型 ====================


@dataclass
class Message:
    """单条对话消息"""

    role: str  # "system" | "user" | "assistant" | "tool"
    content: str
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {"role": self.role, "content": self.content}


# ==================== Memory 抽象 ====================


class Memory(ABC):
    """Memory 抽象基类

    每种 Memory 实现管理 Agent 的对话历史。
    短期记忆在单次会话内有效，长期记忆可跨会话检索。
    """

    @abstractmethod
    async def add(self, message: Message) -> None:
        """添加一条消息到记忆"""
        ...

    @abstractmethod
    async def get(self, n: int = 10) -> List[Message]:
        """获取最近 n 条消息"""
        ...

    @abstractmethod
    async def clear(self) -> None:
        """清空记忆"""
        ...

    @abstractmethod
    def to_messages(self, n: Optional[int] = None) -> List[dict]:
        """导出为 LLM 可用的 messages 格式"""
        ...

    @property
    @abstractmethod
    def size(self) -> int:
        """当前记忆中的消息数"""
        ...


# ==================== 对话记忆实现 ====================


class ConversationMemory(Memory):
    """基于列表的短期对话记忆

    适用于单次会话内的上下文窗口管理。
    超出 max_turns 时自动滚动（FIFO 裁剪）。
    """

    def __init__(self, max_turns: int = 20):
        self._messages: List[Message] = []
        self.max_turns = max_turns

    async def add(self, message: Message) -> None:
        self._messages.append(message)
        # FIFO 裁剪：超出限制时去掉最早的非 system 消息
        if len(self._messages) > self.max_turns:
            for i, msg in enumerate(self._messages):
                if msg.role != "system":
                    self._messages.pop(i)
                    break

    async def get(self, n: int = 10) -> List[Message]:
        return self._messages[-n:]

    async def clear(self) -> None:
        self._messages.clear()

    def to_messages(self, n: Optional[int] = None) -> List[dict]:
        msgs = self._messages if n is None else self._messages[-n:]
        return [m.to_dict() for m in msgs]

    @property
    def size(self) -> int:
        return len(self._messages)


# ==================== 向量记忆（Future） ====================


class VectorMemory(Memory):
    """基于向量数据库的长期记忆

    通过 Milvus 存储语义向量，支持跨会话检索。
    目前为桩实现，后续通过 MilvusBackend 对接。
    """

    def __init__(self, collection: str = "agent_memory"):
        self.collection = collection
        self._messages: List[Message] = []  # 当前会话缓存
        self._backend: Any = None  # MilvusBackend 实例，由 DI 注入

    async def add(self, message: Message) -> None:
        self._messages.append(message)
        # TODO: embed + insert to Milvus

    async def get(self, n: int = 10) -> List[Message]:
        return self._messages[-n:]

    async def search(self, query: str, top_k: int = 5) -> List[Message]:
        """语义检索相关记忆"""
        # TODO: embed query → Milvus search → return messages
        return []

    async def clear(self) -> None:
        self._messages.clear()

    def to_messages(self, n: Optional[int] = None) -> List[dict]:
        msgs = self._messages if n is None else self._messages[-n:]
        return [m.to_dict() for m in msgs]

    @property
    def size(self) -> int:
        return len(self._messages)