Add gfa_core/mas/memory module

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-07 09:52:43 +08:00
parent c389e90e64
commit 842455dd78
8 changed files with 1177 additions and 0 deletions
@@ -0,0 +1,151 @@
+from abc import abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+from gfa_core.mas.memory.base import BaseMemory, MemoryEntry, SearchResult
+
+
+@dataclass
+class KnowledgeEntry:
+    entry_id: str
+    content: str
+    embedding: list[float] | None = None
+    source: str | None = None
+    category: str | None = None
+    tags: list[str] = field(default_factory=list)
+    created_at: datetime = field(default_factory=datetime.now)
+    updated_at: datetime = field(default_factory=datetime.now)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class KnowledgeMemory(BaseMemory):
+    def __init__(self, embedding_dim: int = 1536):
+        self._embedding_dim = embedding_dim
+        self._entries: dict[str, MemoryEntry[KnowledgeEntry]] = {}
+        self._embeddings: dict[str, list[float]] = {}
+
+    async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
+        if isinstance(value, KnowledgeEntry):
+            entry = value
+        else:
+            entry = KnowledgeEntry(entry_id=key, **value)
+
+        memory_entry = MemoryEntry(key=key, value=entry)
+        self._entries[key] = memory_entry
+
+        if entry.embedding:
+            self._embeddings[key] = entry.embedding
+
+    async def retrieve(self, key: str) -> Any | None:
+        entry = self._entries.get(key)
+        return entry.value if entry else None
+
+    async def delete(self, key: str) -> bool:
+        if key in self._entries:
+            del self._entries[key]
+            self._embeddings.pop(key, None)
+            return True
+        return False
+
+    async def exists(self, key: str) -> bool:
+        return key in self._entries
+
+    async def clear(self) -> None:
+        self._entries.clear()
+        self._embeddings.clear()
+
+    async def search(
+        self,
+        query_embedding: list[float],
+        top_k: int = 10,
+        threshold: float = 0.0,
+    ) -> list[SearchResult]:
+        if not self._embeddings:
+            return []
+
+        scores = []
+        for key, embedding in self._embeddings.items():
+            score = self._cosine_similarity(query_embedding, embedding)
+            if score >= threshold:
+                scores.append((key, score))
+
+        scores.sort(key=lambda x: x[1], reverse=True)
+        top_results = scores[:top_k]
+
+        results = []
+        for key, score in top_results:
+            entry = self._entries.get(key)
+            if entry:
+                results.append(SearchResult(
+                    key=key,
+                    value=entry.value,
+                    score=score,
+                    metadata=entry.metadata,
+                ))
+
+        return results
+
+    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
+        if len(a) != len(b):
+            return 0.0
+
+        dot_product = sum(x * y for x, y in zip(a, b))
+        norm_a = sum(x * x for x in a) ** 0.5
+        norm_b = sum(x * x for x in b) ** 0.5
+
+        if norm_a == 0 or norm_b == 0:
+            return 0.0
+
+        return dot_product / (norm_a * norm_b)
+
+    async def search_by_tags(self, tags: list[str], top_k: int = 10) -> list[SearchResult]:
+        results = []
+        for key, entry in self._entries.items():
+            if entry.value.tags and any(tag in entry.value.tags for tag in tags):
+                results.append(SearchResult(
+                    key=key,
+                    value=entry.value,
+                    score=1.0,
+                    metadata=entry.metadata,
+                ))
+                if len(results) >= top_k:
+                    break
+
+        return results
+
+    async def search_by_category(self, category: str, top_k: int = 10) -> list[SearchResult]:
+        results = []
+        for key, entry in self._entries.items():
+            if entry.value.category == category:
+                results.append(SearchResult(
+                    key=key,
+                    value=entry.value,
+                    score=1.0,
+                    metadata=entry.metadata,
+                ))
+                if len(results) >= top_k:
+                    break
+
+        return results
+
+    async def add_entry(
+        self,
+        entry_id: str,
+        content: str,
+        embedding: list[float] | None = None,
+        source: str | None = None,
+        category: str | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        entry = KnowledgeEntry(
+            entry_id=entry_id,
+            content=content,
+            embedding=embedding,
+            source=source,
+            category=category,
+            tags=tags or [],
+            metadata=metadata or {},
+        )
+        await self.store(entry_id, entry)