Fix deferred tool search keyword matching

pydantic_ai_slim/pydantic_ai/toolsets/_tool_search.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass
 from typing import Annotated, Any
 
@@ -18,6 +19,7 @@
 _DISCOVERED_TOOLS_METADATA_KEY = 'discovered_tools'
 
 _MAX_SEARCH_RESULTS = 10
+_SEARCH_TOKEN_RE = re.compile(r'[a-z0-9]+')
 
 
 class _SearchToolArgs(TypedDict):
@@ -42,9 +44,8 @@ class _SearchToolArgs(TypedDict):
 @dataclass(kw_only=True)
 class _SearchIndexEntry:
     name: str
-    name_lower: str
     description: str | None
-    description_lower: str | None
+    search_terms: set[str]
 
 
 @dataclass(kw_only=True)
@@ -90,9 +91,8 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[
         search_index = [
             _SearchIndexEntry(
                 name=name,
-                name_lower=name.lower(),
                 description=tool.tool_def.description,
-                description_lower=tool.tool_def.description.lower() if tool.tool_def.description else None,
+                search_terms=self._search_terms(name, tool.tool_def.description),
             )
             for name, tool in deferred.items()
             if name not in discovered
@@ -148,6 +148,13 @@ async def call_tool(
             return await self._search_tools(tool_args, tool)
         return await self.wrapped.call_tool(name, tool_args, ctx, tool)
 
+    @staticmethod
+    def _search_terms(name: str, description: str | None) -> set[str]:
+        search_terms = set(_SEARCH_TOKEN_RE.findall(name.lower()))
+        if description:
+            search_terms.update(_SEARCH_TOKEN_RE.findall(description.lower()))
+        return search_terms
+
     async def _search_tools(self, tool_args: dict[str, Any], search_tool: _SearchTool[AgentDepsT]) -> ToolReturn:
         """Search for tools matching the keywords.
 
@@ -159,15 +166,16 @@ async def _search_tools(self, tool_args: dict[str, Any], search_tool: _SearchToo
         if not keywords:
             raise ModelRetry('Please provide search keywords.')
 
-        terms = keywords.lower().split()
+        terms = self._search_terms(keywords, None)
 
-        matches: list[dict[str, str | None]] = []
+        scored_matches: list[tuple[int, dict[str, str | None]]] = []
         for entry in search_tool.search_index:
-            searchable = entry.name_lower + (' ' + entry.description_lower if entry.description_lower else '')
-            if any(term in searchable for term in terms):
-                matches.append({'name': entry.name, 'description': entry.description})
-                if len(matches) >= _MAX_SEARCH_RESULTS:
-                    break
+            score = len(terms & entry.search_terms)
+            if score == len(terms):
+                scored_matches.append((score, {'name': entry.name, 'description': entry.description}))
+
+        scored_matches.sort(key=lambda item: item[0], reverse=True)
+        matches = [match for _, match in scored_matches[:_MAX_SEARCH_RESULTS]]
 
         tool_names = [match['name'] for match in matches]
 

tests/test_tool_search.py

@@ -485,6 +485,62 @@ async def test_tool_search_toolset_search_matches_description():
     assert rv[0]['name'] == 'crypto_price'
 
 
+async def test_tool_search_toolset_prefers_specific_term_matches():
+    toolset: FunctionToolset[None] = FunctionToolset()
+
+    @toolset.tool_plain(defer_loading=True)
+    def github_get_me() -> str:  # pragma: no cover
+        """Get the authenticated GitHub profile."""
+        return 'me'
+
+    @toolset.tool_plain(defer_loading=True)
+    def github_create_gist() -> str:  # pragma: no cover
+        """Create a new GitHub gist."""
+        return 'gist'
+
+    searchable = ToolSearchToolset(wrapped=toolset)
+    ctx = _build_run_context(None)
+
+    tools = await searchable.get_tools(ctx)
+    search_tool = tools[_SEARCH_TOOLS_NAME]
+
+    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'keywords': 'github profile'}, ctx, search_tool)
+    assert result == snapshot(
+        ToolReturn(
+            return_value=[{'name': 'github_get_me', 'description': 'Get the authenticated GitHub profile.'}],
+            metadata={'discovered_tools': ['github_get_me']},
+        )
+    )
+
+
+async def test_tool_search_toolset_does_not_match_substrings_inside_words():
+    toolset: FunctionToolset[None] = FunctionToolset()
+
+    @toolset.tool_plain(defer_loading=True)
+    def github_get_me() -> str:  # pragma: no cover
+        """Get my GitHub profile."""
+        return 'me'
+
+    @toolset.tool_plain(defer_loading=True)
+    def github_add_comment_to_pending_review() -> str:  # pragma: no cover
+        """Add a pending review comment on GitHub."""
+        return 'comment'
+
+    searchable = ToolSearchToolset(wrapped=toolset)
+    ctx = _build_run_context(None)
+
+    tools = await searchable.get_tools(ctx)
+    search_tool = tools[_SEARCH_TOOLS_NAME]
+
+    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'keywords': 'get me'}, ctx, search_tool)
+    assert result == snapshot(
+        ToolReturn(
+            return_value=[{'name': 'github_get_me', 'description': 'Get my GitHub profile.'}],
+            metadata={'discovered_tools': ['github_get_me']},
+        )
+    )
+
+
 async def test_tool_search_toolset_search_returns_no_matches():
     """Test that search returns empty list when no matches."""
     toolset = _create_function_toolset()