[transforms] BM25 ranked search for tool discovery: Issue #64

nitrostack/transforms/search/__init__.py

@@ -4,9 +4,13 @@
     CallToolInput, RankOutcome, SearchToolsInput,
 )
 from nitrostack.transforms.search.regex import RegexSearchTransform
+from nitrostack.transforms.search.bm25 import Bm25SearchTransform
+from nitrostack.transforms.search.bm25_engine import Bm25Engine
+from nitrostack.transforms.search.tokenizer import tokenize
 
 __all__ = [
     "BaseSearchTransform", "RegexSearchTransform", "RankOutcome",
     "SearchToolsInput", "CallToolInput", "is_browse_query",
     "MAX_SEARCH_RESULTS", "DEFAULT_SEARCH_LIMIT", "BROWSE_MAX_RESULTS",
+    "Bm25SearchTransform", "Bm25Engine", "tokenize",
 ]

nitrostack/transforms/search/bm25.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import hashlib
+from typing import Any, Optional, Sequence
+
+from nitrostack.transforms.search.base import BaseSearchTransform
+from nitrostack.transforms.search.bm25_engine import Bm25Engine
+from nitrostack.transforms.search.models import RankOutcome
+from nitrostack.transforms.search.serializer import FIELD_WEIGHTS, tool_fields
+from nitrostack.transforms.search.tokenizer import tokenize
+
+
+def catalog_fingerprint(tools: Sequence[Any]) -> str:
+    """Order-independent hash of everything the index is built from."""
+    digest = hashlib.sha256()
+    for fields in sorted(tool_fields(t) for t in tools):
+        for part in fields:
+            digest.update(part.encode("utf-8"))
+            digest.update(b"\x00")
+        digest.update(b"\x01")
+    return digest.hexdigest()
+
+
+class Bm25SearchTransform(BaseSearchTransform):
+    """Ranks tools by BM25 over name, title, description and parameter text.
+
+    Field weights (name > title > description > parameters) come from FIELD_WEIGHTS.
+    The index is built lazily and rebuilt only when the catalog fingerprint changes.
+    """
+
+    name = "bm25_search"
+
+    def __init__(self, *, k1: float = 1.5, b: float = 0.75, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._k1 = k1
+        self._b = b
+        Bm25Engine(k1=k1, b=b)  # validate parameters now, not on the first search
+        self._engine: Optional[Bm25Engine] = None
+        self._fingerprint: Optional[str] = None
+        self.build_count = 0
+
+    def _ensure_index(self, tools: Sequence[Any]) -> Bm25Engine:
+        fingerprint = catalog_fingerprint(tools)
+        if self._engine is None or fingerprint != self._fingerprint:
+            # Build a complete new engine, then swap. There is no await in between,
+            # so no search can observe a half-built index.
+            engine = Bm25Engine(k1=self._k1, b=self._b)
+            engine.build(
+                (
+                    tool.name,
+                    [
+                        (weight, tokenize(text))
+                        for weight, text in zip(FIELD_WEIGHTS, tool_fields(tool))
+                    ],
+                )
+                for tool in tools
+            )
+            self._engine = engine
+            self._fingerprint = fingerprint
+            self.build_count += 1
+        return self._engine
+
+    async def rank(self, query: str, tools: Sequence[Any]) -> RankOutcome:
+        tokens = tokenize(query)
+        if not tokens:
+            return RankOutcome(
+                ranked=[],
+                note="The query has no searchable words. Try a specific term, or search with an empty query to browse.",
+            )
+        engine = self._ensure_index(tools)
+        by_name = {tool.name: tool for tool in tools}
+        ranked = [(by_name[name], score) for name, score in engine.search(tokens) if name in by_name]
+        return RankOutcome(ranked=ranked)

nitrostack/transforms/search/bm25_engine.py

@@ -0,0 +1,71 @@
+"""Pure-Python Okapi BM25 with weighted fields. No dependencies."""
+from __future__ import annotations
+
+import math
+from collections import Counter
+from typing import Dict, Iterable, List, Sequence, Tuple
+
+# One document: (doc_id, [(field_weight, tokens), ...])
+IndexDocument = Tuple[str, Sequence[Tuple[float, Sequence[str]]]]
+
+
+class Bm25Engine:
+    def __init__(self, k1: float = 1.5, b: float = 0.75) -> None:
+        if k1 < 0:
+            raise ValueError("k1 must be >= 0.")
+        if not 0.0 <= b <= 1.0:
+            raise ValueError("b must be between 0 and 1.")
+        self.k1 = k1
+        self.b = b
+        self._ids: List[str] = []
+        self._lengths: List[float] = []
+        self._postings: Dict[str, Dict[int, float]] = {}
+        self._avg_len = 0.0
+
+    @property
+    def size(self) -> int:
+        return len(self._ids)
+
+    def build(self, documents: Iterable[IndexDocument]) -> None:
+        ids: List[str] = []
+        lengths: List[float] = []
+        postings: Dict[str, Dict[int, float]] = {}
+        for doc_index, (doc_id, fields) in enumerate(documents):
+            tf: Counter = Counter()
+            length = 0.0
+            for weight, tokens in fields:
+                for token in tokens:
+                    tf[token] += weight
+                    length += weight
+            ids.append(doc_id)
+            lengths.append(length)
+            for term, freq in tf.items():
+                postings.setdefault(term, {})[doc_index] = freq
+        self._ids = ids
+        self._lengths = lengths
+        self._postings = postings
+        self._avg_len = (sum(lengths) / len(lengths)) if lengths else 0.0
+
+    def idf(self, term: str) -> float:
+        n = len(self._postings.get(term, ()))
+        total = len(self._ids)
+        # The "1 +" keeps IDF positive even for terms that appear in most documents.
+        return math.log(1.0 + (total - n + 0.5) / (n + 0.5))
+
+    def search(self, query_tokens: Sequence[str]) -> List[Tuple[str, float]]:
+        """Return (doc_id, score) for documents with score > 0, best first, ties by id."""
+        scores: Dict[int, float] = {}
+        for term in set(query_tokens):  # a repeated query word counts once
+            docs = self._postings.get(term)
+            if not docs:
+                continue
+            idf = self.idf(term)
+            for doc_index, freq in docs.items():
+                norm = 1.0
+                if self._avg_len > 0:
+                    norm = 1.0 - self.b + self.b * (self._lengths[doc_index] / self._avg_len)
+                denom = freq + self.k1 * norm
+                scores[doc_index] = scores.get(doc_index, 0.0) + idf * (freq * (self.k1 + 1.0)) / denom
+        ranked = [(self._ids[i], s) for i, s in scores.items() if s > 0]
+        ranked.sort(key=lambda pair: (-pair[1], pair[0]))
+        return ranked

nitrostack/transforms/search/tokenizer.py

@@ -0,0 +1,30 @@
+"""Tokenizer for tool search: lowercase, split on punctuation, snake_case and camelCase."""
+from __future__ import annotations
+
+import re
+import unicodedata
+from typing import List
+
+# Words that carry no search meaning. Removed from documents and queries alike.
+STOP_WORDS = frozenset({
+    "a", "an", "the", "of", "to", "for", "and", "or", "in", "on", "by",
+    "with", "is", "it", "this", "that",
+})
+
+# getRnaById -> get Rna By Id ; HTTPServer -> HTTP Server ; v2Api -> v2 Api
+_CAMEL_LOWER_UPPER = re.compile(r"(?<=[a-z0-9])(?=[A-Z])")
+_CAMEL_ACRONYM = re.compile(r"(?<=[A-Z])(?=[A-Z][a-z])")
+# Letters and digits in any script; underscore is a separator, not a word character.
+_WORD = re.compile(r"[^\W_]+", re.UNICODE)
+
+
+def tokenize(text: str, *, drop_stop_words: bool = True) -> List[str]:
+    if not text:
+        return []
+    text = unicodedata.normalize("NFC", text)
+    text = _CAMEL_LOWER_UPPER.sub(" ", text)
+    text = _CAMEL_ACRONYM.sub(" ", text)
+    tokens = [t.lower() for t in _WORD.findall(text)]
+    if drop_stop_words:
+        tokens = [t for t in tokens if t not in STOP_WORDS]
+    return tokens