From 34d3b38730a1b9fed0882243713d4bbacdd6663e Mon Sep 17 00:00:00 2001 From: StanPonomarev Date: Mon, 13 Jul 2026 18:49:24 +0200 Subject: [PATCH] optimize plain search --- README.md | 4 ++- src/corpus_mcp/database.py | 72 +++++++++++++------------------------- tests/test_search.py | 32 +++++++++++++---- 3 files changed, 52 insertions(+), 56 deletions(-) diff --git a/README.md b/README.md index cb04381..4d9e6f2 100644 --- a/README.md +++ b/README.md @@ -71,7 +71,9 @@ returns ranked, highlighted excerpts grouped by document. Plain searches match all whitespace-separated terms without interpreting operators. Set `syntax` to `fts5` for phrases, `OR`, `NOT`, and prefix expressions. Each result includes the containing chunk's character offsets, which can be passed to -`read_document_excerpt` to retrieve more surrounding context. +`read_document_excerpt` to retrieve more surrounding context. Search ranks a +bounded pool of chunk matches so broad queries remain responsive; as a result, +a document with many highly ranked passages may occupy multiple candidates. `read_document_excerpt(path, offset=0, max_chars=4000)` reads a bounded range and returns previous and next offsets. Responses are capped at 20,000 diff --git a/src/corpus_mcp/database.py b/src/corpus_mcp/database.py index 909e747..e47d73d 100644 --- a/src/corpus_mcp/database.py +++ b/src/corpus_mcp/database.py @@ -15,6 +15,7 @@ CHUNK_OVERLAP = 400 MAX_SEARCH_DOCUMENTS = 20 MAX_EXCERPTS_PER_DOCUMENT = 5 MAX_EXCERPT_CHARS = 20_000 +SEARCH_CANDIDATE_MULTIPLIER = 10 class CorpusError(Exception): @@ -352,58 +353,33 @@ def search_database( connection = _connect_read_only(database_path) try: try: - documents = connection.execute( + rows = connection.execute( """ - WITH chunk_matches AS ( - SELECT d.id AS document_id, - d.path, - chunks_fts.rank AS score - FROM chunks_fts - JOIN chunks AS c ON c.id = chunks_fts.rowid - JOIN documents AS d ON d.id = c.document_id - WHERE chunks_fts MATCH ? - ) - SELECT document_id, path, min(score) AS document_score - FROM chunk_matches - GROUP BY document_id, path - ORDER BY document_score ASC, path ASC + SELECT d.path, + c.ordinal, + c.start_char, + c.end_char, + snippet( + chunks_fts, + 0, + '', + '', + ' ... ', + 48 + ) AS excerpt, + chunks_fts.rank AS score + FROM chunks_fts + JOIN chunks AS c ON c.id = chunks_fts.rowid + JOIN documents AS d ON d.id = c.document_id + WHERE chunks_fts MATCH ? + ORDER BY score ASC, d.path ASC, c.ordinal ASC LIMIT ? """, - (match_query, limit), + ( + match_query, + limit * excerpts_per_document * SEARCH_CANDIDATE_MULTIPLIER, + ), ).fetchall() - - rows: list[sqlite3.Row] = [] - for document in documents: - rows.extend( - connection.execute( - """ - SELECT d.path, - c.ordinal, - c.start_char, - c.end_char, - snippet( - chunks_fts, - 0, - '', - '', - ' ... ', - 48 - ) AS excerpt, - chunks_fts.rank AS score - FROM chunks_fts - JOIN chunks AS c ON c.id = chunks_fts.rowid - JOIN documents AS d ON d.id = c.document_id - WHERE chunks_fts MATCH ? AND d.id = ? - ORDER BY score ASC, c.ordinal ASC - LIMIT ? - """, - ( - match_query, - document["document_id"], - excerpts_per_document * 3, - ), - ).fetchall() - ) except sqlite3.OperationalError as exc: if syntax == "fts5": raise CorpusError(f"Invalid FTS5 query: {exc}") from exc diff --git a/tests/test_search.py b/tests/test_search.py index 83d7e18..bd51244 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -1,9 +1,12 @@ from pathlib import Path +from typing import Any import pytest +import corpus_mcp.database as database_module from corpus_mcp.database import ( CorpusError, + SEARCH_CANDIDATE_MULTIPLIER, build_database, read_document_excerpt, search_database, @@ -50,8 +53,8 @@ def test_search_returns_separate_excerpts_from_a_large_document(tmp_path: Path) assert excerpts[1]["chunk_end_char"] <= len(content) -def test_large_document_does_not_hide_other_matching_documents( - tmp_path: Path, +def test_common_term_search_uses_a_bounded_candidate_pool( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: source = tmp_path / "source" source.mkdir() @@ -59,13 +62,28 @@ def test_large_document_does_not_hide_other_matching_documents( (source / "small.txt").write_text("common term", encoding="utf-8") database = tmp_path / "corpus.sqlite" build_database(source, database) + connection = database_module._connect_read_only(database) + executed: list[tuple[str, tuple[Any, ...]]] = [] - result = search_database(database, "common", limit=2) + class RecordingConnection: + def execute(self, statement: str, parameters: tuple[Any, ...] = ()) -> Any: + executed.append((statement, parameters)) + return connection.execute(statement, parameters) - assert {match["path"] for match in result["matches"]} == { - "large.txt", - "small.txt", - } + def close(self) -> None: + connection.close() + + monkeypatch.setattr( + database_module, + "_connect_read_only", + lambda _database: RecordingConnection(), + ) + + search_database(database, "common", limit=2, excerpts_per_document=3) + + assert len(executed) == 1 + assert "GROUP BY" not in executed[0][0] + assert executed[0][1][-1] == 2 * 3 * SEARCH_CANDIDATE_MULTIPLIER def test_read_document_excerpt_is_bounded_and_navigable(