From 653280626b838d481ba18b261b1ac57055ccad34 Mon Sep 17 00:00:00 2001 From: Nemanja Date: Sat, 25 Jul 2026 22:26:16 +0200 Subject: [PATCH] feat: expose chunk_tier as a search filter chunk_tier is computed for every indexed symbol and registered as a Qdrant keyword payload index, but no code path applied it as a filter, so clients couldn't scope a query to just classes or just methods. Adds chunk_tier to QdrantStore.search() / find_by_name() as a FieldCondition, and threads it through the search_code and find_symbol MCP tools. Closes #95 Co-Authored-By: Claude Sonnet 5 --- README.md | 4 +++- docs/retrieval-rrf.md | 8 +++---- server/store/qdrant.py | 19 ++++++++++----- server/tools/search.py | 11 ++++++++- tests/test_store.py | 48 +++++++++++++++++++++++++++++++++++++- tests/tools/test_search.py | 46 +++++++++++++++++++++++++++++++++++- 6 files changed, 122 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 80fa787..7210f70 100644 --- a/README.md +++ b/README.md @@ -360,7 +360,9 @@ in `.env`, then start without the `jina` profile (`docker-compose up` / `make do `search_code` queries both via a Qdrant `query_points` call with `FusionQuery(fusion=RRF)`. Indexed payload fields (`language`, `service`, `symbol_type`, `chunk_tier`, `parent_name`, `file_path`) are -usable as filters. The full payload also includes `signature`, `docstring`, `annotations`, `package`, +usable as filters. `search_code` and `find_symbol` expose `chunk_tier` (`"method"` or `"class"`) +directly, so a query can be scoped to just classes or just methods. The full payload also includes +`signature`, `docstring`, `annotations`, `package`, `start_line`, `end_line`, `file_hash`, `indexed_at`, and language-specific extras (`http_method`, `http_route`, `spring_stereotype`, `lombok_annotations`, `is_async`, `uses_memo`, …). diff --git a/docs/retrieval-rrf.md b/docs/retrieval-rrf.md index d37fc31..84c3d9b 100644 --- a/docs/retrieval-rrf.md +++ b/docs/retrieval-rrf.md @@ -100,13 +100,13 @@ semcode exposes four search tools to AI clients via the MCP protocol (`server/to ### `search_code` ``` -search_code(query: str, service: str | None, limit: int = 10) -> str +search_code(query: str, service: str | None, chunk_tier: str | None, limit: int = 10) -> str ``` The primary semantic search tool. At query time: 1. Embeds the query string with both the dense provider (`embed_query`) and the sparse provider (`embed_query`) -2. Calls `store.search()` with both vectors → RRF fusion +2. Calls `store.search()` with both vectors → RRF fusion, optionally scoped to `chunk_tier` (`"method"` or `"class"`) 3. Returns a formatted Markdown string with up to `limit` results Each result includes: symbol name and type, RRF score, file location (path + line range), service, language, annotations, HTTP route (if present), and the symbol's signature or source (first 500 characters from the payload). @@ -114,10 +114,10 @@ Each result includes: symbol name and type, RRF score, file location (path + lin ### `find_symbol` ``` -find_symbol(name: str, symbol_type: str | None, service: str | None, exact: bool = False) -> str +find_symbol(name: str, symbol_type: str | None, service: str | None, chunk_tier: str | None, exact: bool = False) -> str ``` -Name-based lookup via `store.find_by_name()`. Does not use vectors or RRF. Returns up to 20 (exact) or 50 (substring) matches. Each result includes: name, type, location, package, parent class, and source (first 800 characters). +Name-based lookup via `store.find_by_name()`. Does not use vectors or RRF. Supports filtering by `chunk_tier` (`"method"` or `"class"`) in addition to `symbol_type` and `service`. Returns up to 20 (exact) or 50 (substring) matches. Each result includes: name, type, location, package, parent class, and source (first 800 characters). ### `find_usages` diff --git a/server/store/qdrant.py b/server/store/qdrant.py index ff057e5..1655377 100644 --- a/server/store/qdrant.py +++ b/server/store/qdrant.py @@ -216,14 +216,16 @@ async def search( sparse_vector: SparseVector, limit: int = 10, service: str | None = None, + chunk_tier: str | None = None, ) -> list[ScoredPoint]: - query_filter = ( - Filter( - must=[FieldCondition(key="service", match=MatchValue(value=service))] + must = [] + if service: + must.append(FieldCondition(key="service", match=MatchValue(value=service))) + if chunk_tier: + must.append( + FieldCondition(key="chunk_tier", match=MatchValue(value=chunk_tier)) ) - if service - else None - ) + query_filter = Filter(must=must) if must else None result = await self._client.query_points( collection_name=self._collection, @@ -252,6 +254,7 @@ async def find_by_name( name: str, symbol_type: str | None = None, service: str | None = None, + chunk_tier: str | None = None, exact: bool = False, ) -> list[ScoredPoint]: must = [] @@ -263,6 +266,10 @@ async def find_by_name( ) if service: must.append(FieldCondition(key="service", match=MatchValue(value=service))) + if chunk_tier: + must.append( + FieldCondition(key="chunk_tier", match=MatchValue(value=chunk_tier)) + ) base_filter = Filter(must=must) if must else None diff --git a/server/tools/search.py b/server/tools/search.py index ed92c33..abeb954 100644 --- a/server/tools/search.py +++ b/server/tools/search.py @@ -19,6 +19,7 @@ def register_search_tools(mcp: FastMCP) -> None: async def search_code( query: str, service: str | None = None, + chunk_tier: str | None = None, limit: int = 10, ) -> str: """Semantically search code across indexed services using natural language. @@ -26,6 +27,7 @@ async def search_code( Args: query: Natural language description of what you're looking for. service: Filter by service name + chunk_tier: Filter by chunk tier: "method" or "class" limit: Maximum number of results (default 10) """ embedder = get_embedding_provider() @@ -39,6 +41,7 @@ async def search_code( sparse_vector=sparse_vector, limit=limit, service=service, + chunk_tier=chunk_tier, ) if not results: @@ -74,6 +77,7 @@ async def find_symbol( name: str, symbol_type: str | None = None, service: str | None = None, + chunk_tier: str | None = None, exact: bool = False, ) -> str: """Find a class, method, interface, or function by name. @@ -82,11 +86,16 @@ async def find_symbol( name: Symbol name to search for symbol_type: Optional type filter: class, method, interface, enum, record, function, etc. service: Optional service filter + chunk_tier: Optional chunk tier filter: "method" or "class" exact: If true, only exact name matches. If false (default), partial/fuzzy matching. """ store = get_store() results = await store.find_by_name( - name=name, symbol_type=symbol_type, service=service, exact=exact + name=name, + symbol_type=symbol_type, + service=service, + chunk_tier=chunk_tier, + exact=exact, ) if not results: diff --git a/tests/test_store.py b/tests/test_store.py index 82feed1..8d1f442 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -4,7 +4,7 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock -from qdrant_client.models import Fusion, FusionQuery, SparseVector +from qdrant_client.models import FieldCondition, Fusion, FusionQuery, SparseVector from server.store.qdrant import QdrantStore @@ -99,3 +99,49 @@ async def test_search_uses_prefetch_and_rrf() -> None: assert isinstance(kwargs["query"], FusionQuery) assert kwargs["query"].fusion == Fusion.RRF + + +async def test_search_filters_by_chunk_tier() -> None: + store = QdrantStore.__new__(QdrantStore) + store._collection = "test" + + fake_result = MagicMock() + fake_result.points = [] + store._client = MagicMock() + store._client.query_points = AsyncMock(return_value=fake_result) + + dense = [0.1] * 768 + sparse = SparseVector(indices=[1, 2], values=[0.5, 0.3]) + + await store.search( + dense_vector=dense, sparse_vector=sparse, limit=5, chunk_tier="method" + ) + + kwargs = store._client.query_points.call_args.kwargs + for prefetch in kwargs["prefetch"]: + conditions = prefetch.filter.must + assert any( + isinstance(c, FieldCondition) + and c.key == "chunk_tier" + and c.match.value == "method" + for c in conditions + ) + + +async def test_find_by_name_filters_by_chunk_tier() -> None: + store = QdrantStore.__new__(QdrantStore) + store._collection = "test" + + record = _make_record("MyService") + store._client = MagicMock() + store._client.scroll = AsyncMock(return_value=([record], None)) + + await store.find_by_name("MyService", exact=True, chunk_tier="class") + + scroll_filter = store._client.scroll.call_args.kwargs["scroll_filter"] + assert any( + isinstance(c, FieldCondition) + and c.key == "chunk_tier" + and c.match.value == "class" + for c in scroll_filter.must + ) diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index 9632380..e43fae0 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -67,6 +67,46 @@ async def test_search_code_formats_hits() -> None: assert "0.870" in result +async def test_search_code_passes_chunk_tier_to_store() -> None: + search_code = _tool("search_code") + store = AsyncMock() + store.search.return_value = [] + + with ( + patch("server.tools.search.get_embedding_provider") as mock_embedder, + patch("server.tools.search.get_sparse_provider") as mock_sparse, + patch("server.tools.search.get_store", return_value=store), + ): + mock_embedder.return_value.embed_query = AsyncMock(return_value=[0.1]) + mock_sparse.return_value.embed_query = AsyncMock(return_value={}) + await search_code("find the order service", chunk_tier="method") + + store.search.assert_awaited_once_with( + dense_vector=[0.1], + sparse_vector={}, + limit=10, + service=None, + chunk_tier="method", + ) + + +async def test_find_symbol_passes_chunk_tier_to_store() -> None: + find_symbol = _tool("find_symbol") + store = AsyncMock() + store.find_by_name.return_value = [] + + with patch("server.tools.search.get_store", return_value=store): + await find_symbol("OrderService", chunk_tier="class") + + store.find_by_name.assert_awaited_once_with( + name="OrderService", + symbol_type=None, + service=None, + chunk_tier="class", + exact=False, + ) + + async def test_find_symbol_reports_no_match() -> None: find_symbol = _tool("find_symbol") store = AsyncMock() @@ -103,7 +143,11 @@ async def test_find_symbol_formats_match_with_parent() -> None: assert "`placeOrder`" in result assert "**Parent**: `OrderService`" in result store.find_by_name.assert_awaited_once_with( - name="placeOrder", symbol_type=None, service=None, exact=True + name="placeOrder", + symbol_type=None, + service=None, + chunk_tier=None, + exact=True, )