94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
"""纯内存向量存储后端，供测试用。
|
|
|
|
不持久化、不依赖外部服务，每次进程重启清空。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
|
|
class MemoryVectorStore:
|
|
"""纯内存向量存储（测试用）。"""
|
|
|
|
def __init__(self) -> None:
|
|
self._data: dict[str, dict[str, dict[str, Any]]] = {}
|
|
"""{character_id: {memory_id: {"text": ..., "embedding": [...], "metadata": {...}}}}"""
|
|
|
|
# ── 增 ────────────────────────────────────────────────────────────────
|
|
|
|
def add_memory(
|
|
self,
|
|
character_id: str,
|
|
memory_id: str,
|
|
text: str,
|
|
embedding: list[float],
|
|
*,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> None:
|
|
char_store = self._data.setdefault(character_id, {})
|
|
if memory_id not in char_store:
|
|
char_store[memory_id] = {}
|
|
char_store[memory_id] = {
|
|
"text": text,
|
|
"embedding": embedding,
|
|
"metadata": metadata or {},
|
|
}
|
|
|
|
def update_memory(
|
|
self,
|
|
character_id: str,
|
|
memory_id: str,
|
|
text: str,
|
|
embedding: list[float],
|
|
*,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> None:
|
|
self.add_memory(character_id, memory_id, text, embedding, metadata=metadata)
|
|
|
|
def delete_memory(self, character_id: str, memory_id: str) -> None:
|
|
char_store = self._data.get(character_id)
|
|
if char_store:
|
|
char_store.pop(memory_id, None)
|
|
|
|
# ── 查 ────────────────────────────────────────────────────────────────
|
|
|
|
def search(
|
|
self,
|
|
character_id: str,
|
|
query: list[float],
|
|
*,
|
|
limit: int = 5,
|
|
) -> list[dict[str, Any]]:
|
|
"""余弦相似度暴力搜索（O(n)，仅测试用）。"""
|
|
char_store = self._data.get(character_id)
|
|
if not char_store:
|
|
return []
|
|
|
|
scored: list[tuple[float, str, dict[str, Any]]] = []
|
|
for mid, entry in char_store.items():
|
|
emb = entry.get("embedding", [])
|
|
if not emb:
|
|
continue
|
|
score = _cosine_similarity(query, emb)
|
|
scored.append((score, mid, entry.get("metadata", {})))
|
|
|
|
scored.sort(key=lambda x: x[0], reverse=True)
|
|
return [
|
|
{"id": mid, "score": score, "metadata": meta}
|
|
for score, mid, meta in scored[:limit]
|
|
]
|
|
|
|
def count(self, character_id: str) -> int:
|
|
char_store = self._data.get(character_id)
|
|
return len(char_store) if char_store else 0
|
|
|
|
|
|
def _cosine_similarity(a: list[float], b: list[float]) -> float:
|
|
dot = sum(x * y for x, y in zip(a, b, strict=False))
|
|
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 / (norm_a * norm_b)
|