diff --git a/nitrostack/transforms/search/__init__.py b/nitrostack/transforms/search/__init__.py index 0779c25..1cf6512 100644 --- a/nitrostack/transforms/search/__init__.py +++ b/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", ] \ No newline at end of file diff --git a/nitrostack/transforms/search/bm25.py b/nitrostack/transforms/search/bm25.py new file mode 100644 index 0000000..1a747d9 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/nitrostack/transforms/search/bm25_engine.py b/nitrostack/transforms/search/bm25_engine.py new file mode 100644 index 0000000..8c61a4a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/nitrostack/transforms/search/tokenizer.py b/nitrostack/transforms/search/tokenizer.py new file mode 100644 index 0000000..c97d1c4 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/test_transforms_bm25.py b/tests/test_transforms_bm25.py new file mode 100644 index 0000000..3cb2d21 --- /dev/null +++ b/tests/test_transforms_bm25.py @@ -0,0 +1,209 @@ +import asyncio +import math + +import mcp.types as types +import pytest + +from nitrostack.transforms import TransformContext +from nitrostack.transforms.search import ( + Bm25Engine, Bm25SearchTransform, SearchToolsInput, tokenize, +) + + +def run(coro): + return asyncio.run(coro) + + +def make_tool(name, description="", props=None): + return types.Tool( + name=name, + description=description, + inputSchema={"type": "object", "properties": props or {}}, + ) + + +CATALOG = [ + make_tool("get_user", "Look up a user profile by id", {"user_id": {"type": "string"}}), + make_tool("get_order", "Fetch an order and its line items", {"order_id": {"type": "string"}}), + make_tool("send_email", "Send an email message to a recipient", {"to": {"type": "string"}}), + make_tool("search_genes", "Look up a gene by symbol and return annotations", {"symbol": {"type": "string"}}), + make_tool("get_rna_by_id", "Look up an RNA record by identifier", {"urs_id": {"type": "string"}}), + make_tool("list_invoices", "List invoices for a customer", {}), + make_tool("create_invoice", "Create an invoice for a customer order", {}), + make_tool("delete_user", "Delete a user account", {}), +] + + +class FakeRegistry: + def __init__(self, tools): + self.tools = tools + + def list_tools(self): + return list(self.tools) + + +def names(outcome): + return [tool.name for tool, _score in outcome.ranked] + + +# ---- tokenizer ------------------------------------------------------------- + +@pytest.mark.parametrize("text,expected", [ + ("getRnaById", ["get", "rna", "id"]), # "by" is a stop word + ("get_user_profile", ["get", "user", "profile"]), + ("HTTPServer", ["http", "server"]), + ("sha256Hash", ["sha256", "hash"]), # digits stay attached + ("find---genes!!", ["find", "genes"]), # repeated separators + ("Ünïcode Straße", ["ünïcode", "straße"]), + ("", []), + ("the of and", []), # stop words only +]) +def test_tokenizer(text, expected): + assert tokenize(text) == expected + + +def test_tokenizer_can_keep_stop_words(): + assert tokenize("the gene", drop_stop_words=False) == ["the", "gene"] + + +# ---- engine ---------------------------------------------------------------- + +def test_bm25_score_matches_hand_computation(): + engine = Bm25Engine(k1=1.5, b=0.75) + engine.build([("A", [(1.0, ["apple", "banana"])]), ("B", [(1.0, ["banana"])])]) + # N=2, n=1: idf = ln(1 + 1.5/1.5) = ln 2. avgdl = 1.5. A: tf=1, dl=2. + # denom = 1 + 1.5 * (0.25 + 0.75 * 2/1.5) = 2.875 ; num = 2.5 + expected = math.log(2) * 2.5 / 2.875 + result = engine.search(["apple"]) + assert [doc for doc, _ in result] == ["A"] + assert result[0][1] == pytest.approx(expected) + + +def test_idf_is_never_negative_for_common_terms(): + engine = Bm25Engine() + engine.build([(f"d{i}", [(1.0, ["common"])]) for i in range(5)]) + assert engine.idf("common") > 0 + + +def test_equal_scores_sort_by_id(): + engine = Bm25Engine() + engine.build([("b", [(1.0, ["x"])]), ("a", [(1.0, ["x"])])]) + assert [doc for doc, _ in engine.search(["x"])] == ["a", "b"] + + +def test_repeated_query_terms_count_once(): + engine = Bm25Engine() + engine.build([("a", [(1.0, ["x", "y"])])]) + assert engine.search(["x", "x", "x"]) == engine.search(["x"]) + + +@pytest.mark.parametrize("kwargs", [{"k1": -1}, {"b": 1.5}, {"b": -0.1}]) +def test_invalid_parameters_rejected(kwargs): + with pytest.raises(ValueError): + Bm25Engine(**kwargs) + + +# ---- ranking (golden) ------------------------------------------------------ +# These expectations are reasoned, not generated. Run once, and if one fails, +# check the reasoning before changing the expected value. + +def test_golden_single_match_queries(): + t = Bm25SearchTransform() + assert names(run(t.rank("email", CATALOG))) == ["send_email"] + assert names(run(t.rank("gene", CATALOG))) == ["search_genes"] # via description + assert names(run(t.rank("invoice", CATALOG))) == ["create_invoice"] + assert names(run(t.rank("invoices", CATALOG))) == ["list_invoices"] # no stemming + + +def test_golden_name_match_outranks_description_match(): + result = names(run(Bm25SearchTransform().rank("order", CATALOG))) + assert result[0] == "get_order" + assert "create_invoice" in result + + +def test_golden_user_query_returns_both_user_tools(): + result = names(run(Bm25SearchTransform().rank("user", CATALOG))) + assert set(result[:2]) == {"get_user", "delete_user"} + + +def test_ranking_is_deterministic_across_instances_and_catalog_order(): + first = names(run(Bm25SearchTransform().rank("look up", CATALOG))) + second = names(run(Bm25SearchTransform().rank("look up", list(reversed(CATALOG))))) + assert first == second + + +def test_no_match_returns_empty_list_not_error(): + outcome = run(Bm25SearchTransform().rank("zzzzzz", CATALOG)) + assert outcome.ranked == [] + + +def test_stop_word_only_query_explains_itself(): + outcome = run(Bm25SearchTransform().rank("of the", CATALOG)) + assert outcome.ranked == [] + assert "no searchable words" in outcome.note + + +# ---- index lifecycle ------------------------------------------------------- + +def test_index_is_not_rebuilt_on_every_search(): + t = Bm25SearchTransform() + assert t.build_count == 0 # lazy + for q in ("email", "gene", "order", "user"): + run(t.rank(q, CATALOG)) + assert t.build_count == 1 + + +def test_catalog_order_does_not_trigger_rebuild(): + t = Bm25SearchTransform() + run(t.rank("email", CATALOG)) + run(t.rank("email", list(reversed(CATALOG)))) + assert t.build_count == 1 + + +def test_adding_removing_or_editing_a_tool_rebuilds_once(): + t = Bm25SearchTransform() + run(t.rank("email", CATALOG)) + + added = CATALOG + [make_tool("send_sms", "Send a text message")] + assert "send_sms" in names(run(t.rank("sms", added))) + assert t.build_count == 2 + + run(t.rank("sms", added)) + assert t.build_count == 2 + + edited = [make_tool(x.name, "Totally new text") if x.name == "get_user" else x for x in added] + run(t.rank("totally", edited)) + assert t.build_count == 3 + + run(t.rank("email", [x for x in edited if x.name != "send_email"])) + assert t.build_count == 4 + + +# ---- through the search_tools handler -------------------------------------- + +async def _search(transform, **kw): + ctx = TransformContext() + + async def none(_name): + return None + + tool = await transform.resolve_tool(transform.search_tool_name, none, ctx) + return await tool.handler(SearchToolsInput(**kw), ctx) + + +def test_search_tools_handler_returns_ranked_results_and_clamps_limit(): + t = Bm25SearchTransform() + t.on_register(FakeRegistry(CATALOG)) + out = run(_search(t, query="email")) + assert [r["name"] for r in out["results"]] == ["send_email"] + big = run(_search(t, query="look", limit=500)) + assert len(big["results"]) <= 20 + assert "reduced" in big["note"] + + +def test_search_tools_handler_no_match_is_not_an_error(): + t = Bm25SearchTransform() + t.on_register(FakeRegistry(CATALOG)) + out = run(_search(t, query="zzzzzz")) + assert out["results"] == [] + assert "No tools matched" in out["note"] \ No newline at end of file