"""MemPalace search and write interface. Wraps the ChromaDB-backed palace so that callers (Evennia commands, harness agents, MCP tools) do not need to know the storage details. ChromaDB is imported lazily; if it is not installed the functions raise ``MemPalaceUnavailable`` with an informative message rather than crashing at import time. """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path from typing import Optional from nexus.mempalace.config import ( MEMPALACE_PATH, FLEET_PALACE_PATH, COLLECTION_NAME, ) class MemPalaceUnavailable(RuntimeError): """Raised when ChromaDB or the palace directory is not accessible.""" @dataclass class MemPalaceResult: """A single memory hit returned by the searcher.""" text: str room: str wing: str score: float = 0.0 source_file: str = "" metadata: dict = field(default_factory=dict) def short(self, max_chars: int = 200) -> str: """Return a truncated preview suitable for MUD output.""" if len(self.text) <= max_chars: return self.text return self.text[:max_chars].rstrip() + "…" def _get_client(palace_path: Path): """Return a ChromaDB persistent client, or raise MemPalaceUnavailable.""" try: import chromadb # type: ignore except ImportError as exc: raise MemPalaceUnavailable( "ChromaDB is not installed. " "Run: pip install chromadb (or: pip install mempalace)" ) from exc if not palace_path.exists(): raise MemPalaceUnavailable( f"Palace directory not found: {palace_path}\n" "Run 'mempalace mine' to initialise the palace." ) return chromadb.PersistentClient(path=str(palace_path)) def search_memories( query: str, *, palace_path: Optional[Path] = None, wing: Optional[str] = None, room: Optional[str] = None, n_results: int = 5, ) -> list[MemPalaceResult]: """Search the palace for memories matching *query*. Args: query: Natural-language search string. palace_path: Override the default palace path. wing: Filter results to a specific wizard's wing. room: Filter results to a specific room (e.g. ``"forge"``). n_results: Maximum number of results to return. Returns: List of :class:`MemPalaceResult`, best-match first. Raises: MemPalaceUnavailable: If ChromaDB is not installed or the palace directory does not exist. """ path = palace_path or MEMPALACE_PATH client = _get_client(path) collection = client.get_or_create_collection(COLLECTION_NAME) where: dict = {} if wing: where["wing"] = wing if room: where["room"] = room kwargs: dict = {"query_texts": [query], "n_results": n_results} if where: kwargs["where"] = where raw = collection.query(**kwargs) results: list[MemPalaceResult] = [] if not raw or not raw.get("documents"): return results docs = raw["documents"][0] metas = raw.get("metadatas", [[]])[0] or [{}] * len(docs) distances = raw.get("distances", [[]])[0] or [0.0] * len(docs) for doc, meta, dist in zip(docs, metas, distances): results.append( MemPalaceResult( text=doc, room=meta.get("room", "general"), wing=meta.get("wing", ""), score=float(1.0 - dist), # cosine similarity from distance source_file=meta.get("source_file", ""), metadata=meta, ) ) return results def search_fleet( query: str, *, room: Optional[str] = None, n_results: int = 10, ) -> list[MemPalaceResult]: """Search the shared fleet palace (closets only, no raw drawers). Args: query: Natural-language search string. room: Optional room filter (e.g. ``"issues"``). n_results: Maximum results. Returns: List of :class:`MemPalaceResult` from all wings. """ return search_memories( query, palace_path=FLEET_PALACE_PATH, room=room, n_results=n_results, ) def add_memory( text: str, *, room: str = "general", wing: Optional[str] = None, palace_path: Optional[Path] = None, source_file: str = "", extra_metadata: Optional[dict] = None, ) -> str: """Add a new memory drawer to the palace. Args: text: The memory text to store. room: Target room (e.g. ``"hall_facts"``). wing: Wing name; defaults to :data:`~nexus.mempalace.config.FLEET_WING`. palace_path: Override the default palace path. source_file: Optional source file attribution. extra_metadata: Additional key/value metadata to store. Returns: The generated document ID. Raises: MemPalaceUnavailable: If ChromaDB is not installed or the palace directory does not exist. """ import uuid from nexus.mempalace.config import FLEET_WING path = palace_path or MEMPALACE_PATH client = _get_client(path) collection = client.get_or_create_collection(COLLECTION_NAME) doc_id = str(uuid.uuid4()) metadata: dict = { "room": room, "wing": wing or FLEET_WING, "source_file": source_file, } if extra_metadata: metadata.update(extra_metadata) collection.add( documents=[text], metadatas=[metadata], ids=[doc_id], ) return doc_id