Merge pull request #204 from keboola/AI-1172-search-tools-ref

AI-1172 Move Find Component Id tool to Search tool module
Do we want to cover this tool in the search tool module?

Author: mariankrotil

integtests/tools/components/test_tools.py

@@ -5,7 +5,7 @@
 from mcp.server.fastmcp import Context
 
 from integtests.conftest import ConfigDef, ProjectDef
-from keboola_mcp_server.client import KeboolaClient, SuggestedComponent
+from keboola_mcp_server.client import KeboolaClient
 from keboola_mcp_server.config import MetadataField
 from keboola_mcp_server.links import Link
 from keboola_mcp_server.tools.components.model import (
@@ -21,7 +21,6 @@
     add_config_row,
     create_config,
     create_sql_transformation,
-    find_component_id,
     get_component,
     get_config,
     get_config_examples,
@@ -972,22 +971,6 @@ async def test_get_config_examples(mcp_context: Context, configs: list[ConfigDef
     assert 'parameters' in result
 
 
-@pytest.mark.asyncio
-async def test_find_component_id(mcp_context: Context):
-    """Tests that `find_component_id` returns relevant component IDs for a query."""
-    query = 'generic extractor'
-    generic_extractor_id = 'ex-generic-v2'
-
-    result = await find_component_id(query=query, ctx=mcp_context)
-
-    assert isinstance(result, list)
-    assert len(result) > 0
-    assert generic_extractor_id in [component.component_id for component in result]
-
-    for component in result:
-        assert isinstance(component, SuggestedComponent)
-
-
 @pytest.mark.asyncio
 async def test_get_config_examples_with_invalid_component(mcp_context: Context):
     """Tests that `get_config_examples` handles non-existent components properly."""

integtests/tools/test_search.py

@@ -4,8 +4,8 @@
 from fastmcp import Context
 
 from integtests.conftest import BucketDef, ConfigDef, TableDef
-from keboola_mcp_server.client import KeboolaClient
-from keboola_mcp_server.tools.search import GlobalSearchOutput, find_ids_by_name
+from keboola_mcp_server.client import KeboolaClient, SuggestedComponent
+from keboola_mcp_server.tools.search import GlobalSearchOutput, find_component_id, find_ids_by_name
 
 LOG = logging.getLogger(__name__)
 
@@ -70,3 +70,24 @@ async def test_global_search_end_to_end(
     found_config_ids = {item.id for item in config_group.items}
     # At least some test configurations should be found
     assert found_config_ids.intersection(expected_config_ids), 'Should find at least one test configuration'
+
+    config_groups = [group for group in result.groups.values() if group.type == 'configuration']
+    assert len(config_groups) == 1
+    config_group = config_groups[0]
+    found_config_ids = {item.id for item in config_group.items}
+    # At least some test configurations should be found
+    assert found_config_ids.intersection(expected_config_ids), 'Should find at least one test configuration'
+
+
+@pytest.mark.asyncio
+async def test_find_component_id(mcp_context: Context):
+    """Tests that `find_component_id` returns relevant component IDs for a query."""
+    query = 'generic extractor'
+    generic_extractor_id = 'ex-generic-v2'
+
+    result = await find_component_id(query=query, ctx=mcp_context)
+
+    assert isinstance(result, list)
+    assert len(result) > 0
+    assert all(isinstance(component, SuggestedComponent) for component in result)
+    assert generic_extractor_id in [component.component_id for component in result]

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "keboola-mcp-server"
-version = "1.10.1"
+version = "1.10.2"
 description = "MCP server for interacting with Keboola Connection"
 readme = "README.md"
 requires-python = ">=3.10"

src/keboola_mcp_server/tools/components/tools.py

@@ -8,7 +8,7 @@
 from httpx import HTTPStatusError
 from pydantic import Field
 
-from keboola_mcp_server.client import JsonDict, KeboolaClient, SuggestedComponent
+from keboola_mcp_server.client import JsonDict, KeboolaClient
 from keboola_mcp_server.config import MetadataField
 from keboola_mcp_server.errors import tool_errors
 from keboola_mcp_server.links import ProjectLinksManager
@@ -56,7 +56,6 @@ def add_component_tools(mcp: KeboolaMcpServer) -> None:
     mcp.add_tool(FunctionTool.from_function(update_config))
     mcp.add_tool(FunctionTool.from_function(update_config_row))
     mcp.add_tool(FunctionTool.from_function(get_config_examples))
-    mcp.add_tool(FunctionTool.from_function(find_component_id))
     mcp.add_tool(FunctionTool.from_function(create_sql_transformation))
     mcp.add_tool(FunctionTool.from_function(update_sql_transformation))
     mcp.add_tool(FunctionTool.from_function(list_transformations))
@@ -1022,24 +1021,3 @@ async def get_config_examples(
             markdown += f'{i}. Row Configuration:\n```json\n{json.dumps(example, indent=2)}\n```\n\n'
 
     return markdown
-
-
-@tool_errors()
-@with_session_state()
-async def find_component_id(
-    ctx: Context,
-    query: Annotated[str, Field(description='Natural language query to find the requested component.')],
-) -> list[SuggestedComponent]:
-    """
-    Returns list of component IDs that match the given query.
-
-    USAGE:
-    - Use when you want to find the component for a specific purpose.
-
-    EXAMPLES:
-    - user_input: `I am looking for a salesforce extractor component`
-        - returns a list of component IDs that match the query, ordered by relevance/best match.
-    """
-    client = KeboolaClient.from_state(ctx.session.state)
-    suggestion_response = await client.ai_service_client.suggest_component(query)
-    return suggestion_response.components

src/keboola_mcp_server/tools/search.py

@@ -7,7 +7,7 @@
 from fastmcp.tools import FunctionTool
 from pydantic import BaseModel, Field
 
-from keboola_mcp_server.client import GlobalSearchResponse, ItemType, KeboolaClient
+from keboola_mcp_server.client import GlobalSearchResponse, ItemType, KeboolaClient, SuggestedComponent
 from keboola_mcp_server.errors import tool_errors
 from keboola_mcp_server.mcp import with_session_state
 
@@ -20,6 +20,7 @@
 def add_search_tools(mcp: FastMCP) -> None:
     """Add tools to the MCP server."""
     search_tools = [
+        find_component_id,
         find_ids_by_name,
     ]
     for tool in search_tools:
@@ -146,3 +147,24 @@ async def find_ids_by_name(
         query=joined_prefixes, types=item_types, limit=limit, offset=offset
     )
     return GlobalSearchOutput.from_api_responses(response)
+
+
+@tool_errors()
+@with_session_state()
+async def find_component_id(
+    ctx: Context,
+    query: Annotated[str, Field(description='Natural language query to find the requested component.')],
+) -> list[SuggestedComponent]:
+    """
+    Returns list of component IDs that match the given query.
+
+    USAGE:
+    - Use when you want to find the component for a specific purpose.
+
+    EXAMPLES:
+    - user_input: `I am looking for a salesforce extractor component`
+        - returns a list of component IDs that match the query, ordered by relevance/best match.
+    """
+    client = KeboolaClient.from_state(ctx.session.state)
+    suggestion_response = await client.ai_service_client.suggest_component(query)
+    return suggestion_response.components