""" In-process vector store — drop-in replacement for WeaviateClient. Requires no external service. Uses sentence-transformers for embeddings and numpy for cosine similarity. Lore memories are loaded from the markdown file at startup and are always included in search results regardless of player_id filters. Player-specific memories are persisted to a JSON file. """ import asyncio import json import logging import re import uuid from concurrent.futures import ThreadPoolExecutor from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional import numpy as np from sentence_transformers import SentenceTransformer from memory.schemas import MemoryType logger = logging.getLogger(__name__) # Map lore-file type strings to MemoryType enum values. # The enum values are already lowercase strings, so direct construction works # for well-formed entries; this explicit map covers any aliases. _LORE_TYPE_MAP: Dict[str, MemoryType] = { "lore": MemoryType.LORE, "emotional": MemoryType.EMOTIONAL, "streak": MemoryType.STREAK, "personal_note": MemoryType.PERSONAL_NOTE, "pattern": MemoryType.PATTERN, "player_behavior": MemoryType.PLAYER_BEHAVIOR, "player_observation": MemoryType.PLAYER_OBSERVATION, "game_context": MemoryType.GAME_CONTEXT, } class _MemoryEntry: """Internal record holding a single memory and its embedding.""" __slots__ = ( "id", "content", "memory_type", "related_player_id", "metadata", "timestamp", "created_by", "embedding", # unit-normalised numpy vector (float32) ) def __init__( self, *, id: str, content: str, memory_type: MemoryType, related_player_id: Optional[str], metadata: Optional[Dict[str, Any]], timestamp: datetime, created_by: str, embedding: np.ndarray, ) -> None: self.id = id self.content = content self.memory_type = memory_type self.related_player_id = related_player_id self.metadata = metadata self.timestamp = timestamp self.created_by = created_by self.embedding = embedding def to_result(self, distance: float) -> Dict[str, Any]: """Return caller-facing dict with the standard keys.""" return { "id": self.id, "content": self.content, "memory_type": self.memory_type.value, "distance": distance, "metadata": self.metadata or {}, "related_player_id": self.related_player_id, "timestamp": self.timestamp.isoformat(), } def to_json_dict(self) -> Dict[str, Any]: """Serialise to a JSON-safe dict (for dynamic_memories.json).""" return { "id": self.id, "content": self.content, "memory_type": self.memory_type.value, "related_player_id": self.related_player_id, "metadata": self.metadata, "timestamp": self.timestamp.isoformat(), "created_by": self.created_by, "embedding": self.embedding.tolist(), } @classmethod def from_json_dict(cls, data: Dict[str, Any]) -> "_MemoryEntry": """Deserialise from a JSON dict (loaded from dynamic_memories.json).""" embedding = np.array(data["embedding"], dtype=np.float32) embedding = _normalise(embedding) return cls( id=data["id"], content=data["content"], memory_type=MemoryType(data["memory_type"]), related_player_id=data.get("related_player_id"), metadata=data.get("metadata"), timestamp=datetime.fromisoformat(data["timestamp"]), created_by=data.get("created_by", "main_agent"), embedding=embedding, ) def _normalise(vec: np.ndarray) -> np.ndarray: """Return a unit-normalised copy of *vec* (float32).""" vec = vec.astype(np.float32) norm = np.linalg.norm(vec) if norm > 0.0: vec = vec / norm return vec class SimpleVectorStore: """ In-process semantic memory store — drop-in replacement for WeaviateClient. Public interface: await store(content, memory_type, related_player_id, metadata) -> str await retrieve(query, related_player_id, limit, memory_types) -> List[Dict] Lore memories (related_player_id=None) are always returned in searches, regardless of any player_id filter. Player memories are only returned when their player_id matches the filter (or no filter is applied). The ``distance`` field equals ``1.0 - cosine_similarity``; lower is better. """ EMBEDDING_MODEL = "all-MiniLM-L6-v2" def __init__( self, lore_file: str = "./data/chess_master_lore.md", dynamic_file: str = "./data/dynamic_memories.json", embedding_model: str = EMBEDDING_MODEL, max_workers: int = 2, ) -> None: """ Initialise the store. Args: lore_file: Path to the Chess Master lore markdown file. dynamic_file: Path where runtime memories are persisted as JSON. embedding_model: HuggingFace sentence-transformers model name. max_workers: Thread-pool size used for embedding operations. """ self.lore_path = Path(lore_file) self.dynamic_path = Path(dynamic_file) self._executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="svs-embed") logger.info("Loading embedding model: %s", embedding_model) self._model = SentenceTransformer(embedding_model) # Two separate lists make player-id filtering cheap. self._lore_memories: List[_MemoryEntry] = [] self._dynamic_memories: List[_MemoryEntry] = [] self._load_lore() self._load_dynamic_memories() logger.info( "SimpleVectorStore ready — %d lore + %d dynamic memories", len(self._lore_memories), len(self._dynamic_memories), ) # ------------------------------------------------------------------ # Public async API # ------------------------------------------------------------------ async def store( self, content: str, memory_type: MemoryType, related_player_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, timestamp: Optional[datetime] = None, created_by: str = "main_agent", ) -> str: """ Embed and persist a new memory. Args: content: The memory text. memory_type: Category tag. related_player_id: Player this memory belongs to (or None for global). metadata: Arbitrary key/value context. timestamp: Creation time; defaults to now. created_by: Originating agent label. Returns: UUID string of the stored memory. """ if not content or not content.strip(): raise ValueError("Memory content must not be empty.") loop = asyncio.get_running_loop() embedding = await loop.run_in_executor(self._executor, self._embed, content) entry = _MemoryEntry( id=str(uuid.uuid4()), content=content, memory_type=memory_type, related_player_id=related_player_id, metadata=metadata, timestamp=timestamp or datetime.now(), created_by=created_by, embedding=embedding, ) self._dynamic_memories.append(entry) self._persist_dynamic_memories() logger.info("Stored memory %s (type=%s, player=%s)", entry.id, memory_type.value, related_player_id) return entry.id async def retrieve( self, query: str, related_player_id: Optional[str] = None, limit: int = 5, memory_types: Optional[List[MemoryType]] = None, ) -> List[Dict[str, Any]]: """ Semantic search over stored memories. Lore memories (related_player_id=None) are always candidates. Player memories are included only when their player_id matches ``related_player_id``, or when ``related_player_id`` is None. Args: query: Natural-language search string. related_player_id: Restrict player memories to this player. limit: Maximum number of results to return. memory_types: Optional allow-list of MemoryType values. Returns: List of memory dicts sorted by distance ascending (best first). Each dict has: id, content, memory_type, distance, metadata, related_player_id, timestamp. """ if not query or not query.strip(): return [] loop = asyncio.get_running_loop() query_vec = await loop.run_in_executor(self._executor, self._embed, query) candidates = self._gather_candidates(related_player_id, memory_types) if not candidates: return [] # Stack all embeddings for a single batched dot product. matrix = np.stack([e.embedding for e in candidates]) # (N, D) similarities = matrix @ query_vec # cosine similarity (vectors are unit-norm) distances = 1.0 - similarities # Sort ascending by distance and take top-N. order = np.argsort(distances)[:limit] results = [candidates[i].to_result(float(distances[i])) for i in order] logger.debug("Retrieved %d/%d memories for query '%s…'", len(results), len(candidates), query[:40]) return results # ------------------------------------------------------------------ # Initialisation helpers # ------------------------------------------------------------------ def _load_lore(self) -> None: """Parse the lore markdown file and embed all entries.""" if not self.lore_path.exists(): logger.warning("Lore file not found: %s", self.lore_path) return logger.info("Parsing lore file: %s", self.lore_path) raw = self.lore_path.read_text(encoding="utf-8") blocks = re.split(r"^## Memory \d+:", raw, flags=re.MULTILINE) parsed_blocks = [self._parse_lore_block(b) for b in blocks[1:]] valid = [p for p in parsed_blocks if p is not None] # Batch-embed all lore content at once — much faster than one-by-one if valid: texts = [p["content"] for p in valid] vectors = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False) for p, vec in zip(valid, vectors): entry = _MemoryEntry( id=str(uuid.uuid4()), content=p["content"], memory_type=p["memory_type"], related_player_id=None, metadata=p.get("metadata"), timestamp=p["timestamp"], created_by="lore_seed", embedding=_normalise(np.array(vec, dtype=np.float32)), ) self._lore_memories.append(entry) logger.info("Loaded %d lore memories", len(self._lore_memories)) def _load_dynamic_memories(self) -> None: """Load previously persisted dynamic memories from JSON.""" if not self.dynamic_path.exists(): logger.info("No dynamic memories file found at %s — starting fresh.", self.dynamic_path) return try: data = json.loads(self.dynamic_path.read_text(encoding="utf-8")) for record in data: try: self._dynamic_memories.append(_MemoryEntry.from_json_dict(record)) except Exception as exc: logger.warning("Skipping malformed dynamic memory record: %s", exc) logger.info("Loaded %d dynamic memories from %s", len(self._dynamic_memories), self.dynamic_path) except Exception as exc: logger.error("Failed to load dynamic memories: %s", exc) # ------------------------------------------------------------------ # Persistence # ------------------------------------------------------------------ def _persist_dynamic_memories(self) -> None: """Write all dynamic memories to JSON, atomically via a temp file.""" self.dynamic_path.parent.mkdir(parents=True, exist_ok=True) tmp_path = self.dynamic_path.with_suffix(".tmp") payload = json.dumps( [m.to_json_dict() for m in self._dynamic_memories], ensure_ascii=False, indent=2, ) tmp_path.write_text(payload, encoding="utf-8") tmp_path.replace(self.dynamic_path) # ------------------------------------------------------------------ # Retrieval helpers # ------------------------------------------------------------------ def _gather_candidates( self, related_player_id: Optional[str], memory_types: Optional[List[MemoryType]], ) -> List[_MemoryEntry]: """ Collect the candidate memories for a search. Lore memories are always included. Dynamic memories are included when: - related_player_id is None (no filter), OR - the memory's player_id matches the filter. """ type_set = set(memory_types) if memory_types else None def _type_ok(entry: _MemoryEntry) -> bool: return type_set is None or entry.memory_type in type_set # All lore that passes the type filter. candidates: List[_MemoryEntry] = [e for e in self._lore_memories if _type_ok(e)] # Dynamic memories: apply both player and type filters. for entry in self._dynamic_memories: if not _type_ok(entry): continue if related_player_id is None or entry.related_player_id == related_player_id: candidates.append(entry) return candidates # ------------------------------------------------------------------ # Embedding # ------------------------------------------------------------------ def _embed(self, text: str) -> np.ndarray: """Encode *text* and return a unit-normalised float32 vector.""" vec = self._model.encode(text, convert_to_numpy=True) return _normalise(np.array(vec, dtype=np.float32)) # ------------------------------------------------------------------ # Lore parsing # ------------------------------------------------------------------ def _parse_lore_block(self, block: str) -> Optional[Dict[str, Any]]: """ Parse a single ``## Memory N: ...`` block from the lore file. Returns a dict with keys ``content``, ``memory_type``, ``timestamp``, and optionally ``metadata``, or None if the block is invalid. """ result: Dict[str, Any] = {} for line in block.splitlines(): stripped = line.strip() if not stripped.startswith("- "): continue stripped = stripped[2:] # Drop "- " if stripped.startswith("Type: "): type_str = stripped[len("Type: "):].strip().lower() memory_type = _LORE_TYPE_MAP.get(type_str) if memory_type is None: try: memory_type = MemoryType(type_str) except ValueError: logger.warning("Unknown memory type in lore: %r", type_str) return None result["memory_type"] = memory_type elif stripped.startswith("Timestamp: "): ts_str = stripped[len("Timestamp: "):].strip() try: result["timestamp"] = datetime.fromisoformat(ts_str) except ValueError: logger.warning("Invalid timestamp in lore: %r", ts_str) return None elif stripped.startswith("Content: "): content = stripped[len("Content: "):].strip() # Strip surrounding quotes if present. if len(content) >= 2 and content[0] == '"' and content[-1] == '"': content = content[1:-1] result["content"] = content elif stripped.startswith("Metadata: "): metadata_str = stripped[len("Metadata: "):].strip() result["metadata"] = self._parse_metadata_str(metadata_str) if "content" not in result or "memory_type" not in result: return None if "timestamp" not in result: result["timestamp"] = datetime.now() return result @staticmethod def _parse_metadata_str(metadata_str: str) -> Dict[str, str]: """ Parse a ``key=value, key2=value2`` metadata string. Values may be quoted with double-quotes; quotes are stripped. """ metadata: Dict[str, str] = {} for pair in metadata_str.split(","): pair = pair.strip() if "=" not in pair: continue key, _, value = pair.partition("=") key = key.strip() value = value.strip() if len(value) >= 2 and value[0] == '"' and value[-1] == '"': value = value[1:-1] if key: metadata[key] = value return metadata # ------------------------------------------------------------------ # Diagnostics # ------------------------------------------------------------------ def get_stats(self) -> Dict[str, Any]: """Return a summary of the current store state.""" return { "lore_memory_count": len(self._lore_memories), "dynamic_memory_count": len(self._dynamic_memories), "total_memory_count": len(self._lore_memories) + len(self._dynamic_memories), "embedding_model": self.EMBEDDING_MODEL, "dynamic_file": str(self.dynamic_path), } __all__ = ["SimpleVectorStore"]