the-nexus/nexus/mempalace/searcher.py

"""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