feat: Improve find_related API, fix chunk boundary bug, update MCP instructions

README.md

@@ -40,8 +40,8 @@ result.chunk.start_line  # 127
 result.chunk.end_line    # 150
 result.chunk.content     # "def save_pretrained(self, path: PathLike, ..."
 
-# Find code similar to a specific location in the codebase
-related = index.find_related("model2vec/model.py", line=127, top_k=3)
+# Find code similar to a specific result
+related = index.find_related(results[0], top_k=3)
 ```
 
 ## Main Features

src/semble/index/index.py

@@ -154,23 +154,14 @@ def from_git(
 
             return index
 
-    def find_related(self, file_path: str, line: int, top_k: int = 5) -> list[SearchResult]:
-        """Return chunks semantically similar to the chunk at the given file location.
-
-        :param file_path: Path to the file, in the same format stored by the index.
-            For both `from_path` and `from_git` this is a repo-relative path
-            (e.g. ``src/foo.py``).  Use `chunk.file_path` from a prior search result
-            to guarantee the correct format.
-        :param line: Line number (1-indexed) used to identify the source chunk.
+    def find_related(self, source: Chunk | SearchResult, *, top_k: int = 5) -> list[SearchResult]:
+        """Return chunks semantically similar to the given chunk or search result.
+
+        :param source: A SearchResult or Chunk to use as the seed.
         :param top_k: Number of similar chunks to return.
         :return: Ranked list of SearchResult objects, most similar first.
         """
-        target = next(
-            (c for c in self.chunks if c.file_path == file_path and c.start_line <= line <= c.end_line),
-            None,
-        )
-        if target is None:
-            return []
+        target = source.chunk if isinstance(source, SearchResult) else source
         selector = self._get_selector_vector(filter_languages=[target.language]) if target.language else None
         results = search_semantic(target.content, self.model, self._semantic_index, self.chunks, top_k + 1, selector)
         return [r for r in results if r.chunk != target][:top_k]

src/semble/mcp.py

@@ -10,26 +10,29 @@
 
 from semble.index import SembleIndex
 from semble.index.dense import load_model
-from semble.types import Encoder, SearchResult
+from semble.types import Chunk, Encoder, SearchResult
 
 _REPO_DESCRIPTION = (
     "Git URL (e.g. https://github.com/org/repo) or local path to index and search. "
     "Required when no default index was configured at startup. "
     "The index is cached after the first call, so repeat queries are fast."
 )
 
+_GIT_URL_SCHEMES = ("https://", "http://", "ssh://", "git://", "git+ssh://", "file://")
+# scp-like syntax: [user@]host:path, where host has no '/' before the ':'.
+_SCP_GIT_URL_RE = re.compile(r"^[\w.-]+@[\w.-]+:(?!/)")
+
 
 def create_server(cache: _IndexCache, default_source: str | None = None) -> FastMCP:
     """Build and return a configured FastMCP server backed by the given cache."""
     server = FastMCP(
         "semble",
         instructions=(
-            "Use this server to search any codebase by source code. "
-            "When the user asks how a library or project works, call `search` with the "
-            "GitHub URL of the relevant repository as `repo` and a natural-language query. "
-            "Resolve the GitHub URL from your training knowledge (e.g. a PyPI package name "
-            "maps to its source repo). Always prefer `search` over Grep, Glob, or Read for "
-            "any question about how code works."
+            "Instant code search for any local or GitHub repository. "
+            "Call `search` to find relevant code; call `find_related` on a result to discover similar code elsewhere. "
+            "For questions about a library (e.g. a PyPI/npm package), resolve the GitHub URL from your training "
+            "knowledge and pass it as `repo`. "
+            "Prefer these tools over Grep, Glob, or Read for any question about how code works."
         ),
     )
 
@@ -41,13 +44,12 @@ async def search(
             Literal["hybrid", "semantic", "bm25"],
             Field(description="Search mode. 'hybrid' is best for most queries."),
         ] = "hybrid",
-        top_k: Annotated[int, Field(description="Number of results to return.", ge=1, le=20)] = 5,
+        top_k: Annotated[int, Field(description="Number of results to return.", ge=1)] = 5,
     ) -> str:
         """Search a codebase with a natural-language or code query.
 
-        Pass a git URL or local path as `repo` to clone and index it on demand.
-        The index is cached so subsequent searches on the same repo are instant.
-        Returns the most relevant code chunks with file paths and line numbers.
+        Pass a git URL or local path as `repo` to index it on demand; indexes are cached for the session.
+        Use this to find where something is implemented, understand a library, or locate related code.
         """
         source = repo or default_source
         if not source:
@@ -72,12 +74,12 @@ async def find_related(
         ],
         line: Annotated[int, Field(description="Line number (1-indexed).")],
         repo: Annotated[str | None, Field(description=_REPO_DESCRIPTION)] = None,
-        top_k: Annotated[int, Field(description="Number of similar chunks to return.", ge=1, le=10)] = 5,
+        top_k: Annotated[int, Field(description="Number of similar chunks to return.", ge=1)] = 5,
     ) -> str:
         """Find code chunks semantically similar to a specific location in a file.
 
-        Useful for discovering related logic elsewhere in the codebase.
-        Pass the same `repo` used in the original `search` call.
+        Use after `search` to explore related implementations or callers.
+        Pass file_path and line from a prior search result.
         """
         source = repo or default_source
         if not source:
@@ -89,12 +91,15 @@ async def find_related(
             index = await cache.get(source)
         except Exception as exc:
             return f"Failed to index {source!r}: {exc}"
-        results = index.find_related(file_path, line, top_k=top_k)
-        if not results:
+        chunk = _resolve_chunk(index.chunks, file_path, line)
+        if chunk is None:
             return (
-                f"No related chunks found for {file_path}:{line}. "
+                f"No chunk found at {file_path}:{line}. "
                 "Make sure the file is indexed and the line number is within a known chunk."
             )
+        results = index.find_related(chunk, top_k=top_k)
+        if not results:
+            return f"No related chunks found for {file_path}:{line}."
         return _format_results(f"Chunks related to {file_path}:{line}", results)
 
     return server
@@ -148,9 +153,26 @@ async def get(self, source: str, ref: str | None = None) -> SembleIndex:
             raise
 
 
-_GIT_URL_SCHEMES = ("https://", "http://", "ssh://", "git://", "git+ssh://", "file://")
-# scp-like syntax: [user@]host:path, where host has no '/' before the ':'.
-_SCP_GIT_URL_RE = re.compile(r"^[\w.-]+@[\w.-]+:(?!/)")
+def _resolve_chunk(chunks: list[Chunk], file_path: str, line: int) -> Chunk | None:
+    """Return the chunk that contains *line* in *file_path*, or None.
+
+    MCP tool arguments are JSON primitives (strings and ints), so the agent
+    passes file_path + line rather than a Chunk object. This function
+    reconstructs the Chunk at the MCP boundary before calling into the library.
+
+    :param chunks: All indexed chunks to search.
+    :param file_path: File path as stored in the index.
+    :param line: 1-indexed line number to resolve.
+    :return: The best-matching Chunk, or None if not found.
+    """
+    fallback = None
+    for chunk in chunks:
+        if chunk.file_path == file_path and chunk.start_line <= line <= chunk.end_line:
+            if line < chunk.end_line:
+                return chunk
+            if fallback is None:  # line == end_line: boundary; keep as fallback for end-of-file chunks
+                fallback = chunk
+    return fallback
 
 
 def _is_git_url(path: str) -> bool:

tests/test_index.py

@@ -84,14 +84,18 @@ def test_search_empty_query_returns_empty(indexed_index: SembleIndex, mode: str,
 
 
 def test_find_related(indexed_index: SembleIndex) -> None:
-    """find_related: returns similar chunks for a known location; returns [] for an unknown file."""
+    """find_related returns related chunks for a Chunk or SearchResult seed."""
     chunk = indexed_index.chunks[0]
-    results = indexed_index.find_related(chunk.file_path, chunk.start_line, top_k=3)
-    assert isinstance(results, list)
-    assert all(r.chunk != chunk for r in results)
-    assert len(results) <= 3
-
-    assert indexed_index.find_related("/does/not/exist.py", 1) == []
+    via_chunk = indexed_index.find_related(chunk, top_k=3)
+    assert isinstance(via_chunk, list)
+    assert len(via_chunk) <= 3
+    assert all(r.chunk != chunk for r in via_chunk)
+
+    # SearchResult form returns the same results as Chunk form.
+    result = indexed_index.search("authenticate", top_k=1)[0]
+    assert [r.chunk for r in indexed_index.find_related(result, top_k=3)] == [
+        r.chunk for r in indexed_index.find_related(result.chunk, top_k=3)
+    ]
 
 
 _GIT_ENV = {

tests/test_mcp.py

@@ -5,8 +5,8 @@
 
 import pytest
 
-from semble.mcp import _format_results, _IndexCache, _is_git_url, create_server, main, serve
-from semble.types import Encoder, SearchMode, SearchResult
+from semble.mcp import _format_results, _IndexCache, _is_git_url, _resolve_chunk, create_server, main, serve
+from semble.types import Chunk, Encoder, SearchMode, SearchResult
 from tests.conftest import make_chunk
 
 
@@ -22,11 +22,14 @@ async def _call_tool(
     *,
     index_method: str,
     index_return: list[SearchResult],
+    index_chunks: list[Chunk] | None = None,
     default_source: str | None = "/some/path",
 ) -> str:
     """Patch SembleIndex.from_path with a fake index and invoke the tool, returning the text."""
     fake_index = MagicMock()
     getattr(fake_index, index_method).return_value = index_return
+    if index_chunks is not None:
+        fake_index.chunks = index_chunks
     with patch("semble.mcp.SembleIndex.from_path", return_value=fake_index):
         server = create_server(cache, default_source=default_source)
         result = await server.call_tool(tool, args)
@@ -39,6 +42,24 @@ def cache() -> _IndexCache:
     return _IndexCache(model=MagicMock(spec=Encoder))
 
 
+def test_resolve_chunk() -> None:
+    """_resolve_chunk returns the correct chunk and handles boundary and miss cases."""
+    interior = make_chunk("line1\nline2\nline3", "src/a.py")  # start=1, end=3
+    boundary = make_chunk("last line", "src/a.py")  # start=1, end=1 (single-line)
+
+    # Line strictly inside a multi-line chunk hits the early-return path.
+    assert _resolve_chunk([interior], "src/a.py", 2) is interior
+
+    # Line equal to end_line of a single-line chunk hits the fallback path.
+    assert _resolve_chunk([boundary], "src/a.py", 1) is boundary
+
+    # Unknown file returns None.
+    assert _resolve_chunk([interior], "src/other.py", 1) is None
+
+    # Line out of range returns None.
+    assert _resolve_chunk([interior], "src/a.py", 99) is None
+
+
 @pytest.mark.parametrize(
     ("path", "expected"),
     [
@@ -156,13 +177,14 @@ async def test_tool_index_failure(cache: _IndexCache, tool: str, args: dict[str,
 
 @pytest.mark.anyio
 @pytest.mark.parametrize(
-    ("tool", "args", "method", "results", "expected_substrings"),
+    ("tool", "args", "method", "results", "chunks", "expected_substrings"),
     [
         pytest.param(
             "search",
             {"query": "bar"},
             "search",
             [SearchResult(chunk=make_chunk("def bar(): pass", "src/bar.py"), score=0.9, source=SearchMode.HYBRID)],
+            None,
             ["bar", "0.900"],
             id="search_with_results",
         ),
@@ -171,6 +193,7 @@ async def test_tool_index_failure(cache: _IndexCache, tool: str, args: dict[str,
             {"query": "nothing"},
             "search",
             [],
+            None,
             ["No results found"],
             id="search_no_results",
         ),
@@ -179,17 +202,28 @@ async def test_tool_index_failure(cache: _IndexCache, tool: str, args: dict[str,
             {"file_path": "src/foo.py", "line": 1},
             "find_related",
             [SearchResult(chunk=make_chunk("class Foo: pass", "src/foo.py"), score=0.8, source=SearchMode.SEMANTIC)],
+            [make_chunk("class Foo: pass", "src/foo.py")],
             ["src/foo.py:1", "0.800"],
             id="find_related_with_results",
         ),
         pytest.param(
             "find_related",
-            {"file_path": "src/foo.py", "line": 99},
+            {"file_path": "src/foo.py", "line": 1},
             "find_related",
             [],
+            [make_chunk("class Foo: pass", "src/foo.py")],
             ["No related chunks found"],
             id="find_related_no_results",
         ),
+        pytest.param(
+            "find_related",
+            {"file_path": "src/unknown.py", "line": 1},
+            "find_related",
+            [],
+            [],
+            ["No chunk found"],
+            id="find_related_unknown_file",
+        ),
     ],
 )
 async def test_tool_output(
@@ -198,10 +232,11 @@ async def test_tool_output(
     args: dict[str, Any],
     method: str,
     results: list[SearchResult],
+    chunks: list[Chunk] | None,
     expected_substrings: list[str],
 ) -> None:
     """Search and find_related format results (or an empty-state message) through the server."""
-    text = await _call_tool(cache, tool, args, index_method=method, index_return=results)
+    text = await _call_tool(cache, tool, args, index_method=method, index_return=results, index_chunks=chunks)
     for substring in expected_substrings:
         assert substring in text