chore: import upstream snapshot with attribution

runtime/node/agent/memory/memory_base.py

@@ -0,0 +1,304 @@
+"""Base memory abstractions with multimodal snapshots."""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+import time
+
+from entity.configs import MemoryAttachmentConfig, MemoryStoreConfig
+from entity.configs.node.memory import FileMemoryConfig, SimpleMemoryConfig
+from entity.enums import AgentExecFlowStage
+from entity.messages import Message, MessageBlock
+from runtime.node.agent.memory.embedding import EmbeddingBase, EmbeddingFactory
+
+
+@dataclass
+class MemoryContentSnapshot:
+    """Lightweight serialization of a multimodal payload."""
+
+    text: str
+    blocks: List[Dict[str, Any]] = field(default_factory=list)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"text": self.text, "blocks": self.blocks}
+
+    @classmethod
+    def from_dict(cls, payload: Dict[str, Any] | None) -> "MemoryContentSnapshot | None":
+        if not payload:
+            return None
+        text = payload.get("text", "")
+        blocks = payload.get("blocks") or []
+        return cls(text=text, blocks=list(blocks))
+
+    @classmethod
+    def from_message(cls, message: Message | str | None) -> "MemoryContentSnapshot | None":
+        if message is None:
+            return None
+        if isinstance(message, Message):
+            return cls(
+                text=message.text_content(),
+                blocks=[
+                    {
+                        "role": message.role.value,
+                        "block": block.to_dict(include_data=True),
+                    }
+                    for block in message.blocks()
+                ],
+            )
+        if isinstance(message, str):
+            return cls(text=message, blocks=[])
+        return cls(text=str(message), blocks=[])
+
+    @classmethod
+    def from_messages(cls, messages: List[Message]) -> "MemoryContentSnapshot | None":
+        if not messages:
+            return None
+        parts: List[str] = []
+        blocks: List[Dict[str, Any]] = []
+        for message in messages:
+            parts.append(f"({message.role.value}) {message.text_content()}")
+            for block in message.blocks():
+                blocks.append(
+                    {
+                        "role": message.role.value,
+                        "block": block.to_dict(include_data=True),
+                    }
+                )
+        return cls(text="\n\n".join(parts), blocks=blocks)
+
+    def to_message_blocks(self) -> List[MessageBlock]:
+        blocks: List[MessageBlock] = []
+        for payload in self.blocks:
+            block_data = payload.get("block") if isinstance(payload, dict) else None
+            if not isinstance(block_data, dict):
+                continue
+            try:
+                blocks.append(MessageBlock.from_dict(block_data))
+            except Exception:
+                continue
+        return blocks
+
+    def attachment_overview(self) -> List[Dict[str, Any]]:
+        attachments: List[Dict[str, Any]] = []
+        for payload in self.blocks:
+            block_data = payload.get("block") if isinstance(payload, dict) else None
+            if not isinstance(block_data, dict):
+                continue
+            attachment = block_data.get("attachment")
+            if attachment:
+                attachments.append(
+                    {
+                        "role": payload.get("role"),
+                        "attachment_id": attachment.get("attachment_id"),
+                        "mime_type": attachment.get("mime_type"),
+                        "name": attachment.get("name"),
+                        "size": attachment.get("size"),
+                    }
+                )
+        return attachments
+
+    @classmethod
+    def from_blocks(
+        cls,
+        *,
+        text: str,
+        blocks: List[MessageBlock],
+        role: str = "input",
+    ) -> "MemoryContentSnapshot":
+        serialized = [
+            {
+                "role": role,
+                "block": block.to_dict(include_data=True),
+            }
+            for block in blocks
+        ]
+        return cls(text=text, blocks=serialized)
+
+
+@dataclass
+class MemoryItem:
+    id: str
+    content_summary: str
+    metadata: Dict[str, Any]
+    embedding: Optional[List[float]] = None
+    timestamp: float | None = None
+    input_snapshot: MemoryContentSnapshot | None = None
+    output_snapshot: MemoryContentSnapshot | None = None
+
+    def __post_init__(self) -> None:
+        if self.timestamp is None:
+            self.timestamp = time.time()
+
+    def to_dict(self) -> Dict[str, Any]:
+        payload: Dict[str, Any] = {
+            "id": self.id,
+            "content_summary": self.content_summary,
+            "metadata": self.metadata,
+            "embedding": self.embedding,
+            "timestamp": self.timestamp,
+        }
+        if self.input_snapshot:
+            payload["input_snapshot"] = self.input_snapshot.to_dict()
+        if self.output_snapshot:
+            payload["output_snapshot"] = self.output_snapshot.to_dict()
+        return payload
+
+    @classmethod
+    def from_dict(cls, payload: Dict[str, Any]) -> "MemoryItem":
+        return cls(
+            id=payload["id"],
+            content_summary=payload.get("content_summary", ""),
+            metadata=payload.get("metadata") or {},
+            embedding=payload.get("embedding"),
+            timestamp=payload.get("timestamp"),
+            input_snapshot=MemoryContentSnapshot.from_dict(payload.get("input_snapshot")),
+            output_snapshot=MemoryContentSnapshot.from_dict(payload.get("output_snapshot")),
+        )
+
+    def attachments(self) -> List[Dict[str, Any]]:
+        attachments: List[Dict[str, Any]] = []
+        if self.input_snapshot:
+            attachments.extend(self.input_snapshot.attachment_overview())
+        if self.output_snapshot:
+            attachments.extend(self.output_snapshot.attachment_overview())
+        return attachments
+
+
+@dataclass
+class MemoryWritePayload:
+    agent_role: str
+    inputs_text: str
+    input_snapshot: MemoryContentSnapshot | None
+    output_snapshot: MemoryContentSnapshot | None
+
+
+@dataclass
+class MemoryRetrievalResult:
+    formatted_text: str
+    items: List[MemoryItem]
+
+    def has_multimodal(self) -> bool:
+        return any(
+            (item.input_snapshot and item.input_snapshot.blocks)
+            or (item.output_snapshot and item.output_snapshot.blocks)
+            for item in self.items
+        )
+
+    def attachment_overview(self) -> List[Dict[str, Any]]:
+        attachments: List[Dict[str, Any]] = []
+        for item in self.items:
+            attachments.extend(item.attachments())
+        return attachments
+
+
+class MemoryBase:
+    def __init__(self, store: MemoryStoreConfig):
+        self.store = store
+        self.name = store.name
+        self.contents: List[MemoryItem] = []
+
+        embedding_cfg = None
+        simple_cfg = store.as_config(SimpleMemoryConfig)
+        file_cfg = store.as_config(FileMemoryConfig)
+        if simple_cfg and simple_cfg.embedding:
+            embedding_cfg = simple_cfg.embedding
+        elif file_cfg and file_cfg.embedding:
+            embedding_cfg = file_cfg.embedding
+
+        self.embedding: EmbeddingBase | None = (
+            EmbeddingFactory.create_embedding(embedding_cfg) if embedding_cfg else None
+        )
+
+    def count_memories(self) -> int:
+        return len(self.contents)
+
+    def load(self) -> None:  # pragma: no cover - implemented by subclasses
+        raise NotImplementedError
+
+    def save(self) -> None:  # pragma: no cover - implemented by subclasses
+        raise NotImplementedError
+
+    def retrieve(
+        self,
+        agent_role: str,
+        query: MemoryContentSnapshot,
+        top_k: int,
+        similarity_threshold: float,
+    ) -> List[MemoryItem]:
+        raise NotImplementedError
+
+    def update(self, payload: MemoryWritePayload) -> None:
+        raise NotImplementedError
+
+
+class MemoryManager:
+    def __init__(self, attachments: List[MemoryAttachmentConfig], stores: Dict[str, MemoryBase]):
+        self.attachments = attachments
+        self.memories: Dict[str, MemoryBase] = {}
+        for attachment in attachments:
+            memory = stores.get(attachment.name)
+            if not memory:
+                raise ValueError(f"memory store {attachment.name} not found")
+            self.memories[attachment.name] = memory
+
+    def retrieve(
+        self,
+        agent_role: str,
+        query: MemoryContentSnapshot,
+        current_stage: AgentExecFlowStage,
+    ) -> MemoryRetrievalResult | None:
+        results: List[tuple[str, MemoryItem, float]] = []
+        for attachment in self.attachments:
+            if attachment.retrieve_stage and current_stage not in attachment.retrieve_stage:
+                continue
+            if not attachment.read:
+                continue
+            memory = self.memories.get(attachment.name)
+            if not memory:
+                continue
+            items = memory.retrieve(agent_role, query, attachment.top_k, attachment.similarity_threshold)
+            for item in items:
+                combined_score = self._score_memory(item, query.text)
+                results.append((attachment.name, item, combined_score))
+
+        if not results:
+            return None
+
+        results.sort(key=lambda entry: entry[2], reverse=True)
+        formatted = ["===== Related Memories ====="]
+        grouped: Dict[str, List[MemoryItem]] = {}
+        for name, item, _ in results:
+            grouped.setdefault(name, []).append(item)
+        for name, items in grouped.items():
+            formatted.append(f"\n--- {name} ---")
+            for idx, item in enumerate(items, 1):
+                formatted.append(f"{idx}. {item.content_summary}")
+        formatted.append("\n===== End of Memory =====")
+
+        ordered_items = [item for _, item, _ in results]
+        return MemoryRetrievalResult(formatted_text="\n".join(formatted), items=ordered_items)
+
+    def update(self, payload: MemoryWritePayload) -> None:
+        for attachment in self.attachments:
+            if not attachment.write:
+                continue
+            memory = self.memories.get(attachment.name)
+            if not memory:
+                continue
+            memory.update(payload)
+            memory.save()
+
+    def _score_memory(self, memory_item: MemoryItem, query: str) -> float:
+        current_time = time.time()
+        age_hours = (current_time - (memory_item.timestamp or current_time)) / 3600
+        time_decay = max(0.1, 1.0 - age_hours / (24 * 30))
+        length = len(memory_item.content_summary)
+        if length < 20:
+            length_factor = 0.5
+        elif length > 200:
+            length_factor = 0.8
+        else:
+            length_factor = 1.0
+        query_words = set(query.lower().split())
+        content_words = set(memory_item.content_summary.lower().split())
+        relevance = len(query_words & content_words) / len(query_words) if query_words else 0.0
+        return 0.7 * time_decay * length_factor + 0.3 * relevance