fix: split underscore terms in BM25 description and param tokens

livekit-agents/livekit/agents/beta/toolsets/tool_search.py

@@ -196,6 +196,33 @@ def _get_tool_params(tool: FunctionTool | RawFunctionTool) -> dict[str, str]:
     }
 
 
+# a plain word, optionally followed by sentence punctuation; anything else may
+# carry regex syntax (^get_weather$) or term-identifying punctuation (c++, c#)
+# and must be kept verbatim when tokenizing keyword queries
+_PROSE_TOKEN = re.compile(r"([^\W_]+)[.,!?;:']*")
+# terms whose punctuation is part of their identity and must not collapse to
+# the bare word (c++ vs c#)
+_SPECIAL_TERM = re.compile(r"[^\W_]+[+#]+")
+
+
+def _keywords(query: str) -> list[str]:
+    """Keyword query tokens: plain words drop trailing sentence punctuation,
+    other tokens keep their matching rules (regex syntax, term punctuation)."""
+    tokens: list[str] = []
+    for token in query.lower().split():
+        m = _PROSE_TOKEN.fullmatch(token)
+        tokens.append(m.group(1) if m else token)
+    return list(set(tokens))
+
+
+def _bm25_tokens(text: str) -> list[str]:
+    """Lowercase word tokens (punctuation dropped, underscores split) plus
+    punctuation-bearing terms (c++, c#) kept verbatim so distinct terms stay
+    distinct."""
+    lowered = text.lower()
+    return re.findall(r"[^\W_]+", lowered) + _SPECIAL_TERM.findall(lowered)
+
+
 class KeywordSearchStrategy:
     """Keyword search using regex matching.
 
@@ -211,7 +238,7 @@ def build_index(self, items: list[SearchItem]) -> None:
             }
 
     def search(self, query: str, items: list[SearchItem], max_results: int) -> list[SearchItem]:
-        keywords = list(set(query.lower().split()))
+        keywords = _keywords(query)
         if not keywords:
             return []
 
@@ -304,7 +331,7 @@ def build_index(self, items: list[SearchItem]) -> None:
             self._idf[term] = math.log((n - freq + 0.5) / (freq + 0.5) + 1.0)
 
     def search(self, query: str, items: list[SearchItem], max_results: int) -> list[SearchItem]:
-        query_terms = query.lower().replace("_", " ").split()
+        query_terms = _bm25_tokens(query)
         if not query_terms:
             return []
 
@@ -323,12 +350,12 @@ def cleanup(self) -> None:
 
     def _tokenize(self, item: SearchItem) -> list[str]:
         """Tokenize with field weighting: name 3x, description 2x, parameters 1x."""
-        name_tokens = item.name.lower().replace("_", " ").split()
-        desc_tokens = item.description.lower().split()
+        name_tokens = _bm25_tokens(item.name)
+        desc_tokens = _bm25_tokens(item.description)
         param_tokens = []
         for k, v in item.parameters.items():
-            param_tokens.extend(k.lower().replace("_", " ").split())
-            param_tokens.extend(v.lower().split())
+            param_tokens.extend(_bm25_tokens(k))
+            param_tokens.extend(_bm25_tokens(v))
 
         # weight by repeating tokens
         return name_tokens * 3 + desc_tokens * 2 + param_tokens

tests/test_tool_search.py

@@ -135,6 +135,176 @@ def test_match_by_parameter(self, strategy: SearchStrategy):
         assert len(results) >= 1
         assert results[0].name == "get_stock_price"
 
+    @pytest.mark.parametrize("strategy", STRATEGIES)
+    @pytest.mark.parametrize("query", ["checking_balance", "checking balance"])
+    def test_match_by_underscore_term_in_description(self, strategy: SearchStrategy, query: str):
+        """Underscore-joined terms in the description must match, like the
+        space-separated query normalization applied by search()."""
+        items = [
+            SearchItem(
+                name="get_records",
+                description="Fetch the user's checking_balance from the ledger",
+                parameters={"record_type": "Which records to fetch"},
+                source=calculator,
+            ),
+            SearchItem(
+                name="get_weather",
+                description="Get current weather for a city",
+                parameters={"city": "City name"},
+                source=weather_tool,
+            ),
+        ]
+        strategy.build_index(items)
+        results = strategy.search(query, items, 5)
+        assert len(results) == 1
+        assert results[0].name == "get_records"
+
+    @pytest.mark.parametrize("strategy", STRATEGIES)
+    @pytest.mark.parametrize("query", ["pending_refunds", "pending refunds"])
+    def test_match_by_underscore_term_in_param_description(
+        self, strategy: SearchStrategy, query: str
+    ):
+        """Underscore-joined terms in parameter descriptions must match, like
+        the space-separated query normalization applied by search()."""
+        items = [
+            SearchItem(
+                name="get_records",
+                description="Fetch ledger records for the user",
+                parameters={"filter": "One of: pending_refunds only, settled only, all"},
+                source=calculator,
+            ),
+            SearchItem(
+                name="get_weather",
+                description="Get current weather for a city",
+                parameters={"city": "City name"},
+                source=weather_tool,
+            ),
+        ]
+        strategy.build_index(items)
+        results = strategy.search(query, items, 5)
+        assert len(results) == 1
+        assert results[0].name == "get_records"
+
+    @pytest.mark.parametrize("strategy", STRATEGIES)
+    @pytest.mark.parametrize("query", ["balance", "ledger"])
+    def test_match_by_punctuation_in_description(self, strategy: SearchStrategy, query: str):
+        """Punctuation attached to description words (commas, trailing period)
+        must not prevent matches, like the query normalization applied by
+        search()."""
+        items = [
+            SearchItem(
+                name="get_records",
+                description="Fetch the ledger, then the balance.",
+                parameters={"record_type": "Which records to fetch"},
+                source=calculator,
+            ),
+            SearchItem(
+                name="get_weather",
+                description="Get current weather for a city",
+                parameters={"city": "City name"},
+                source=weather_tool,
+            ),
+        ]
+        strategy.build_index(items)
+        results = strategy.search(query, items, 5)
+        assert len(results) == 1
+        assert results[0].name == "get_records"
+
+    @pytest.mark.parametrize("strategy", STRATEGIES)
+    def test_match_by_mid_sentence_punctuation_in_description(self, strategy: SearchStrategy):
+        """A comma attached to a mid-sentence description word must not prevent
+        matches, like the query normalization applied by search()."""
+        items = [
+            SearchItem(
+                name="get_records",
+                description="Track ledger, journal and receipts",
+                parameters={"record_type": "Which records to track"},
+                source=calculator,
+            ),
+            SearchItem(
+                name="get_weather",
+                description="Get current weather for a city",
+                parameters={"city": "City name"},
+                source=weather_tool,
+            ),
+        ]
+        strategy.build_index(items)
+        results = strategy.search("ledger", items, 5)
+        assert len(results) == 1
+        assert results[0].name == "get_records"
+
+    @pytest.mark.parametrize("strategy", STRATEGIES)
+    @pytest.mark.parametrize("query", ["real", "time"])
+    def test_match_by_hyphenated_word_in_description(self, strategy: SearchStrategy, query: str):
+        """Hyphenated description words must match their parts, like the query
+        normalization applied by search()."""
+        items = [
+            SearchItem(
+                name="get_records",
+                description="Get real-time balance updates",
+                parameters={"metric": "Which metric to report"},
+                source=calculator,
+            ),
+            SearchItem(
+                name="get_weather",
+                description="Get current weather for a city",
+                parameters={"city": "City name"},
+                source=weather_tool,
+            ),
+        ]
+        strategy.build_index(items)
+        results = strategy.search(query, items, 5)
+        assert len(results) == 1
+        assert results[0].name == "get_records"
+
+    @pytest.mark.parametrize("strategy", STRATEGIES)
+    @pytest.mark.parametrize("query", ["the balance.", "balance?"])
+    def test_match_by_punctuation_in_query(self, strategy: SearchStrategy, query: str):
+        """Punctuation attached to query words must not prevent matches against
+        clean description words, like the index normalization applied by
+        build_index()."""
+        items = [
+            SearchItem(
+                name="get_records",
+                description="Fetch account balance",
+                parameters={"account": "Account identifier"},
+                source=calculator,
+            ),
+            SearchItem(
+                name="get_stock_price",
+                description="Get stock price for a symbol",
+                parameters={"symbol": "Stock ticker symbol"},
+                source=stock_tool,
+            ),
+        ]
+        strategy.build_index(items)
+        results = strategy.search(query, items, 5)
+        assert len(results) == 1
+        assert results[0].name == "get_records"
+
+    @pytest.mark.parametrize("strategy", STRATEGIES)
+    def test_match_by_period_terminated_param_description(self, strategy: SearchStrategy):
+        """A trailing period on a parameter description must not prevent
+        matches, like the query normalization applied by search()."""
+        items = [
+            SearchItem(
+                name="get_records",
+                description="Fetch ledger records",
+                parameters={"filter": "Only include settled refunds."},
+                source=calculator,
+            ),
+            SearchItem(
+                name="get_weather",
+                description="Get current weather for a city",
+                parameters={"city": "City name"},
+                source=weather_tool,
+            ),
+        ]
+        strategy.build_index(items)
+        results = strategy.search("refunds", items, 5)
+        assert len(results) == 1
+        assert results[0].name == "get_records"
+
     @pytest.mark.parametrize("strategy", STRATEGIES)
     def test_max_results_limit(self, strategy: SearchStrategy):
         items = _make_items()
@@ -180,6 +350,35 @@ def test_index_data_preserved_through_search(self, strategy: SearchStrategy):
         assert results[0].index_data is not None
 
 
+class TestKeywordSpecific:
+    """Tests specific to keyword (regex) behavior."""
+
+    def test_regex_query_keeps_matching_rules(self):
+        """Regex syntax in a keyword query must keep its matching rules: the
+        anchored query matches only the exact name, not every tool containing
+        the same words. The forecast tool is indexed first so an anchor-stripping
+        regression that ties both tools returns the wrong one."""
+        items = [
+            SearchItem(
+                name="get_weather_forecast",
+                description="Get the weather forecast for a city",
+                parameters={"city": "City name"},
+                source=forecast_tool,
+            ),
+            SearchItem(
+                name="get_weather",
+                description="Get current weather for a city",
+                parameters={"city": "City name"},
+                source=weather_tool,
+            ),
+        ]
+        strategy = KeywordSearchStrategy()
+        strategy.build_index(items)
+        results = strategy.search("^get_weather$", items, 1)
+        assert len(results) == 1
+        assert results[0].name == "get_weather"
+
+
 class TestBM25Specific:
     """Tests specific to BM25 behavior."""
 
@@ -193,6 +392,34 @@ def test_specific_term_ranks_higher_than_common(self):
         assert len(results) >= 1
         assert results[0].name == "search_papers"
 
+    @pytest.mark.parametrize(
+        ("query", "expected"),
+        [("c++", "compile_cpp"), ("c#", "compile_csharp")],
+    )
+    def test_programming_language_terms_stay_distinct(self, query: str, expected: str):
+        """Punctuation-bearing terms like c++ and c# must not collapse to the
+        same token, or the wrong-language tool can be returned. The C# tool is
+        indexed first so a token collision returns it for both queries."""
+        items = [
+            SearchItem(
+                name="compile_csharp",
+                description="Compile C# code",
+                parameters={"source": "C# source file"},
+                source=calculator,
+            ),
+            SearchItem(
+                name="compile_cpp",
+                description="Compile C++ code",
+                parameters={"source": "C++ source file"},
+                source=stock_tool,
+            ),
+        ]
+        strategy = BM25SearchStrategy()
+        strategy.build_index(items)
+        results = strategy.search(query, items, 1)
+        assert len(results) == 1
+        assert results[0].name == expected
+
     def test_cleanup_clears_state(self):
         strategy = BM25SearchStrategy()
         items = _make_items()