autocad-mcp-pro/discovery/transform.py

"""Tool-search transform that speaks AutoCAD.

Subclasses FastMCP's :class:`BM25SearchTransform` rather than replacing it, so
the ``call_tool`` proxy, the auth-filtered catalog accessor
(``_get_visible_tools``) and the re-entrancy bypass all come for free. Two
things are added on top, and only two:

**1. Vocabulary.** Stock BM25 indexes tool name + description + parameter names
+ parameter descriptions. Measured against this server's live catalog, that
leaves ``BPOLY``, ``PEDIT``, ``WBLOCK``, ``OVERKILL``, ``LAYTRANS``,
``MATCHPROP`` and ``QSELECT`` at df=0 — a search for any of them returns
nothing at all. :func:`discovery.aliases.alias_text` supplies the missing
AutoCAD command names and shop-floor synonyms; this module folds them into the
indexed document. (``FILLET`` is *not* in that list: it is already indexed,
because the tool is literally named ``entity_fillet``.)

**2. Two additive boosts**, so an exact alias match outranks an incidental
description word. See :data:`ACAD_MATCH_BOOST` / :data:`SYNONYM_MATCH_BOOST`.

BM25 itself is not reimplemented: the index object the base class already
built is reused, only with different documents.
"""

from __future__ import annotations

import hashlib
import re
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Annotated, Any, Literal

from fastmcp.server.context import Context
from fastmcp.server.transforms.search import BM25SearchTransform
from fastmcp.server.transforms.search.base import SearchResultSerializer
from fastmcp.server.transforms.search.base import (
    _extract_searchable_text as _stock_searchable_text,
)
from fastmcp.tools.base import Tool
from pydantic import Field

from discovery.aliases import alias_text, aliases_for
from discovery.serialize import serialize_search_results

__all__ = [
    "ACAD_MATCH_BOOST",
    "MAX_SEARCH_LIMIT",
    "RISK_CEILINGS",
    "SYNONYM_MATCH_BOOST",
    "CadSearchTransform",
    "SearchOutcome",
    "tool_risk",
]


# ---------------------------------------------------------------------------
# Scoring
# ---------------------------------------------------------------------------
#
# ``_BM25Index.query()`` hands back a *ranking*, not scores, so the two signals
# are fused on a rank scale instead of a raw-score scale: BM25's nth hit is
# worth ``1 / (1 + n)``. That makes the boosts below trivially explainable in
# the only units that matter here — positions.
#
#   ACAD_MATCH_BOOST     the query names an AutoCAD command this tool owns,
#                        *written as a command* — see _names_a_command below.
#                        Worth exactly as much as being BM25's top hit, so a
#                        command match pulls a tool past any number of tools
#                        that merely share a description word.
#   SYNONYM_MATCH_BOOST  the query contains one of the tool's synonym phrases
#                        verbatim ("bill of materials", "round the corner").
#                        Worth BM25's second position: strong, but it loses to
#                        an outright command match.
#
# That is the whole scorer. Both boosts are flat and apply at most once, so a
# tool with a long synonym list cannot out-shout a tool with a short one.
ACAD_MATCH_BOOST = 1.0
SYNONYM_MATCH_BOOST = 0.5

_WORD_RE = re.compile(r"[a-z0-9]+")
_RAW_WORD_RE = re.compile(r"[A-Za-z0-9]+")

# Dropped from both sides before phrase matching, so the synonym "delete a
# layer" still matches "delete the layer called GEOMETRY". Articles only —
# a general stop-word list would start matching phrases that mean other things.
_ARTICLES = frozenset({"a", "an", "the"})


def _words(text: str) -> list[str]:
    """Lowercase word list. Unlike BM25's tokenizer, single characters survive.

    ``_tokenize`` drops ``len <= 1`` tokens, which is right for an inverted
    index and wrong for phrase matching: dropping the "a" out of "round the
    corner" would silently change which phrases can match.
    """
    return _WORD_RE.findall(text.lower())


def _phrase_words(text: str) -> list[str]:
    return [word for word in _words(text) if word not in _ARTICLES]


def _names_a_command(query: str, command: str) -> bool:
    """True when ``query`` uses ``command`` as a command, not as English.

    Several AutoCAD command names are also ordinary nouns — LAYER, TABLE, LINE,
    POINT, TEXT, BOX, ARRAY. Boosting on a bare lowercase "layer" put
    ``layer_create`` / ``layer_modify`` above ``layer_delete`` for "delete the
    layer called GEOMETRY", which is exactly backwards. A drafter reaching for
    the command types it in capitals, or types nothing else:

      "BPOLY", "trace it with BPOLY"  -> command
      "bpoly"                         -> command (nothing else in the query)
      "delete the layer ..."          -> English; BM25 handles it
    """
    words = _words(query)
    if len(words) == 1 and words[0] == command.lower():
        return True
    return command in _RAW_WORD_RE.findall(query)


def _contains_phrase(haystack: Sequence[str], needle: Sequence[str]) -> bool:
    """True when ``needle`` appears as a contiguous run of words in ``haystack``.

    Word-level rather than substring, so the synonym "line" does not match
    "polyline" or "linear".
    """
    n = len(needle)
    if not n or n > len(haystack):
        return False
    first = needle[0]
    for i, word in enumerate(haystack):
        if word == first and list(haystack[i : i + n]) == list(needle):
            return True
    return False


# ---------------------------------------------------------------------------
# Risk
# ---------------------------------------------------------------------------
#
# Derived from the MCP annotations each tool already carries, so there is no
# second risk table to drift out of step with the registrations. The ceiling a
# caller may ask for, most permissive first:
#
#   any    everything (default — a search must not quietly narrow itself)
#   write  hide tools that delete or overwrite (entity_delete_many, layer_delete,
#          transaction_rollback, drawing_finalize, ...)
#   read   only read-only query tools — the right ceiling for "how many ...?"
RISK_CEILINGS = ("any", "write", "read")

MAX_SEARCH_LIMIT = 25


def tool_risk(tool: Tool) -> str:
    """``"read"`` | ``"write"`` | ``"destructive"`` for one tool.

    ``destructiveHint`` wins over ``readOnlyHint``: no tool in this catalog
    sets both, and if one ever did, the safe reading is the destructive one.
    Tools that annotate neither are treated as ``"write"`` — the honest default
    for an unannotated tool on a server whose job is mutating drawings.
    """
    annotations = tool.annotations
    if annotations is not None:
        if annotations.destructiveHint:
            return "destructive"
        if annotations.readOnlyHint:
            return "read"
    return "write"


def _risk_allowed(ceiling: str, level: str) -> bool:
    if ceiling == "any":
        return True
    if ceiling == "write":
        return level != "destructive"
    return level == "read"


@dataclass(frozen=True, slots=True)
class SearchOutcome:
    """One search: the tools shown, plus what the risk ceiling took away.

    ``hidden`` holds only the tools the caller would otherwise have *seen* —
    blocked hits inside the unfiltered top-``limit`` — not every blocked tool
    in the catalog. That keeps the notice a fact about this answer.
    """

    hits: tuple[Tool, ...] = ()
    hidden: tuple[Tool, ...] = ()
    risk: str = "any"

    def notice(self) -> str:
        """One line naming what the risk ceiling removed, or ``""``.

        A filter that silently halves a mixed intent ("list the layers then
        delete GEOMETRY") is worse than no filter: the caller cannot tell the
        difference between "no such tool" and "not shown to you".
        """
        if not self.hidden:
            return ""
        names = ", ".join(tool.name for tool in self.hidden)
        return (
            f"Risk filter: {len(self.hidden)} higher-risk match(es) hidden by "
            f'risk="{self.risk}" ({names}). Re-run with risk="any" to include them.'
        )


class CadSearchTransform(BM25SearchTransform):
    """BM25 tool search with the AutoCAD alias corpus folded into the index.

    Args:
        max_results: Default number of tools returned per search.
        bm25_b: BM25 length-normalisation, left at the 0.75 textbook default —
            the value every measured number in this module was obtained with.
            Lowering it is defensible in principle (these "documents" are
            metadata records of deliberately unequal verbosity, and a tool with
            six synonyms is not a broader tool than one with two), but it was
            not what shipped, so re-measure the golden set before changing it.
    """

    def __init__(
        self,
        *,
        max_results: int = 5,
        bm25_b: float = 0.75,
        always_visible: list[str] | None = None,
        search_tool_name: str = "search_tools",
        call_tool_name: str = "call_tool",
        search_result_serializer: SearchResultSerializer | None = None,
    ) -> None:
        super().__init__(
            max_results=max_results,
            always_visible=always_visible,
            search_tool_name=search_tool_name,
            call_tool_name=call_tool_name,
            search_result_serializer=search_result_serializer or serialize_search_results,
        )
        self._bm25_k1 = self._index.k1
        self._bm25_b = bm25_b

    # ------------------------------------------------------------------
    # Indexed document
    # ------------------------------------------------------------------

    def _extract_searchable_text(self, tool: Tool) -> str:
        """Stock searchable text plus this tool's alias vocabulary.

        The narrowest available hook: upstream's ``_extract_searchable_text``
        is a module-level function, not a method, so it cannot be overridden in
        place — it is called and extended instead. Tools with no alias record
        (anything registered outside this server) keep the stock text
        unchanged.
        """
        extra = alias_text(tool.name)
        stock = _stock_searchable_text(tool)
        return f"{stock} {extra}" if extra else stock

    # ------------------------------------------------------------------
    # Index lifecycle
    # ------------------------------------------------------------------

    def _ensure_index(self, tools: Sequence[Tool]) -> None:
        """Rebuild the index when the catalog's indexed text changes.

        The hash is order-sensitive on purpose: results are mapped back to
        tools positionally, so a reordered catalog with identical text is still
        a different index.

        No lock: this and its only caller (:meth:`rank`) are synchronous and
        await nothing, so two concurrent searches on one event loop cannot
        interleave a rebuild with a query.
        """
        documents = [self._extract_searchable_text(t) for t in tools]
        digest = hashlib.sha256("\x00".join(documents).encode()).hexdigest()
        if digest == self._last_hash:
            return
        # Built via type() rather than importing the private _BM25Index name:
        # the base class already owns an instance, and reusing its class keeps
        # this from depending on a private module path.
        index = type(self._index)(self._bm25_k1, self._bm25_b)
        index.build(documents)
        self._index = index
        self._indexed_tools = tuple(tools)
        self._last_hash = digest

    # ------------------------------------------------------------------
    # Ranking
    # ------------------------------------------------------------------

    def _alias_boost(self, tool: Tool, query: str, phrase_words: Sequence[str]) -> float:
        record = aliases_for(tool.name)
        if record is None:
            return 0.0
        boost = 0.0
        if any(_names_a_command(query, command) for command in record.acad):
            boost += ACAD_MATCH_BOOST
        if any(_contains_phrase(phrase_words, _phrase_words(phrase)) for phrase in record.synonyms):
            boost += SYNONYM_MATCH_BOOST
        return boost

    def rank(
        self,
        tools: Sequence[Tool],
        query: str,
        *,
        limit: int | None = None,
        risk: str = "any",
    ) -> SearchOutcome:
        """Rank ``tools`` against ``query``. Pure — no Context, no I/O."""
        if risk not in RISK_CEILINGS:
            raise ValueError(f"risk must be one of {RISK_CEILINGS}, got {risk!r}")
        cutoff = self._max_results if limit is None else max(1, min(limit, MAX_SEARCH_LIMIT))
        tools = list(tools)
        if not tools:
            return SearchOutcome(risk=risk)
        self._ensure_index(tools)

        ranked = self._index.query(query, len(tools))
        bm25_points = {position: 1.0 / (1.0 + rank) for rank, position in enumerate(ranked)}
        phrase_words = _phrase_words(query)

        scored: list[tuple[float, str, Tool]] = []
        for position, tool in enumerate(tools):
            score = bm25_points.get(position, 0.0) + self._alias_boost(tool, query, phrase_words)
            if score > 0.0:
                scored.append((score, tool.name, tool))
        scored.sort(key=lambda row: (-row[0], row[1]))
        matches = [tool for _, _, tool in scored]

        # Filter before truncating, so a blocked hit frees its slot for the
        # next safe tool instead of just shortening the answer.
        allowed = [tool for tool in matches if _risk_allowed(risk, tool_risk(tool))]
        hidden = tuple(
            tool for tool in matches[:cutoff] if not _risk_allowed(risk, tool_risk(tool))
        )
        return SearchOutcome(hits=tuple(allowed[:cutoff]), hidden=hidden, risk=risk)

    async def _search(self, tools: Sequence[Tool], query: str) -> Sequence[Tool]:
        return list(self.rank(tools, query).hits)

    # ------------------------------------------------------------------
    # The synthetic search tool
    # ------------------------------------------------------------------

    async def render(self, outcome: SearchOutcome) -> Any:
        """Serialize an outcome, keeping the risk notice on its own line.

        With the default (string) serializer the notice is appended as a
        trailing line. A serializer that returns structured data instead gets
        the notice as a sibling key, because there is no line to append to.
        """
        rendered = await self._render_results(outcome.hits)
        notice = outcome.notice()
        if not notice:
            return rendered
        if isinstance(rendered, str):
            return f"{rendered}\n\n{notice}"
        return {"results": rendered, "notice": notice}

    def _make_search_tool(self) -> Tool:
        transform = self
        default_limit = self._max_results

        async def search_tools(
            query: Annotated[
                str,
                Field(
                    description=(
                        "What you want to do, in plain English or as an AutoCAD "
                        "command name (FILLET, BPOLY, QSELECT, DLI)."
                    )
                ),
            ],
            limit: Annotated[
                int,
                Field(description=(f"How many tools to return. Clamped to 1..{MAX_SEARCH_LIMIT}.")),
            ] = default_limit,
            risk: Annotated[
                Literal["any", "write", "read"],
                Field(
                    description=(
                        "Highest risk level to return. 'any' (default) is everything; "
                        "'write' hides tools that delete or overwrite; 'read' returns "
                        "only read-only query tools — use it for questions like 'how "
                        "many ...?'. Anything filtered out is reported in the result."
                    )
                ),
            ] = "any",
            ctx: Context = None,  # type: ignore[assignment]
        ) -> Any:
            """Find the AutoCAD tools that do a job.

            Accepts AutoCAD command names as well as English. Returns one
            compact line per tool: name, risk level, what it does, its required
            arguments and the AutoCAD command it corresponds to. Call a result
            with call_tool.
            """
            visible = await transform._get_visible_tools(ctx)
            outcome = transform.rank(visible, query, limit=limit, risk=risk)
            return await transform.render(outcome)

        return Tool.from_function(fn=search_tools, name=self._search_tool_name)