From 4da50873bd5ebf141e94e6968b64bca6fcc05521 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 9 Dec 2025 11:46:43 +0200 Subject: [PATCH] Make search limit configurable --- docs/configuration/index.md | 1 + docs/configuration/storage.md | 6 ++++-- haiku_rag_slim/haiku/rag/app.py | 4 +++- haiku_rag_slim/haiku/rag/cli.py | 6 +++--- haiku_rag_slim/haiku/rag/client.py | 7 +++++-- haiku_rag_slim/haiku/rag/config/models.py | 1 + haiku_rag_slim/haiku/rag/graph/common/nodes.py | 4 ++-- haiku_rag_slim/haiku/rag/mcp.py | 6 ++++-- haiku_rag_slim/haiku/rag/qa/agent.py | 2 +- tests/test_app.py | 4 ++-- tests/test_cli.py | 2 +- tests/test_mcp.py | 2 +- tests/test_search.py | 3 ++- 13 files changed, 30 insertions(+), 18 deletions(-) diff --git a/docs/configuration/index.md b/docs/configuration/index.md index 01c6081e..8b58db58 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -99,6 +99,7 @@ research: max_concurrency: 1 search: + limit: 5 # Default number of results to return vector_index_metric: cosine # cosine, l2, or dot vector_refine_factor: 30 diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index ca3a4733..f8182bdf 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -57,16 +57,18 @@ haiku.rag intelligently handles database creation based on operation type: This prevents the common mistake where a search query accidentally creates an empty database. To initialize your database, simply add your first document using `haiku-rag add` or `haiku-rag add-src`. -## Vector Indexing +## Search Settings -Configure vector indexing behavior for efficient similarity search: +Configure search behavior: ```yaml search: + limit: 5 # Default number of results to return vector_index_metric: cosine # cosine, l2, or dot vector_refine_factor: 30 # Re-ranking factor for accuracy ``` +- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and API. Default: 5 - **vector_index_metric**: Distance metric for vector similarity: - `cosine`: Cosine similarity (default, best for most embeddings) - `l2`: Euclidean distance diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 330a3410..37d87eac 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -265,7 +265,9 @@ class HaikuRAGApp: f"[yellow]Document with id {doc_id} not found.[/yellow]" ) - async def search(self, query: str, limit: int = 5, filter: str | None = None): + async def search( + self, query: str, limit: int | None = None, filter: str | None = None + ): async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client: results = await self.client.search(query, limit=limit, filter=filter) if not results: diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index e02e2ff9..1ebb45da 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -239,11 +239,11 @@ def search( query: str = typer.Argument( help="The search query to use", ), - limit: int = typer.Option( - 5, + limit: int | None = typer.Option( + None, "--limit", "-l", - help="Maximum number of results to return", + help="Maximum number of results to return (default: config search.default_limit)", ), filter: str | None = typer.Option( None, diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 9df0be72..3466bdbc 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -834,7 +834,7 @@ class HaikuRAG: async def search( self, query: str, - limit: int = 5, + limit: int | None = None, search_type: str = "hybrid", filter: str | None = None, ) -> list[SearchResult]: @@ -842,13 +842,16 @@ class HaikuRAG: Args: query: The search query string. - limit: Maximum number of results to return. + limit: Maximum number of results to return. Defaults to config.search.default_limit. search_type: Type of search - "vector", "fts", or "hybrid" (default). filter: Optional SQL WHERE clause to filter documents before searching chunks. Returns: List of SearchResult objects ordered by relevance. """ + if limit is None: + limit = self._config.search.limit + reranker = get_reranker(config=self._config) if reranker is None: diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 30c10943..8809ea81 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -123,6 +123,7 @@ class ProcessingConfig(BaseModel): class SearchConfig(BaseModel): + limit: int = 5 vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine" vector_refine_factor: int = 30 diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index 885036d0..62d2e9e6 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/haiku_rag_slim/haiku/rag/graph/common/nodes.py @@ -101,7 +101,7 @@ def create_plan_node[AgentDepsT: GraphAgentDeps]( @plan_agent.tool async def gather_context( - ctx2: RunContext[AgentDepsT], query: str, limit: int = 6 + ctx2: RunContext[AgentDepsT], query: str, limit: int | None = None ) -> str: results = await ctx2.deps.client.search(query, limit=limit) results = await ctx2.deps.client.expand_context(results) @@ -227,7 +227,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( @agent.tool async def search_and_answer( - ctx2: RunContext[AgentDepsT], query: str, limit: int = 5 + ctx2: RunContext[AgentDepsT], query: str, limit: int | None = None ) -> str: """Search the knowledge base for relevant documents. diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 676fa82d..48dbaaa3 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -79,11 +79,13 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP: return None @mcp.tool() - async def search_documents(query: str, limit: int = 5) -> list[SearchResult]: + async def search_documents( + query: str, limit: int | None = None + ) -> list[SearchResult]: """Search the RAG system for documents using hybrid search (vector similarity + full-text search).""" try: async with HaikuRAG(db_path, config=config) as rag: - return await rag.search(query, limit) + return await rag.search(query, limit=limit) except Exception: return [] diff --git a/haiku_rag_slim/haiku/rag/qa/agent.py b/haiku_rag_slim/haiku/rag/qa/agent.py index c180e57f..5c37ffeb 100644 --- a/haiku_rag_slim/haiku/rag/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/qa/agent.py @@ -39,7 +39,7 @@ class QuestionAnswerAgent: async def search_documents( ctx: RunContext[Dependencies], query: str, - limit: int = 5, + limit: int | None = None, ) -> str: """Search the knowledge base for relevant documents. diff --git a/tests/test_app.py b/tests/test_app.py index ec5d581c..b7743605 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -159,7 +159,7 @@ async def test_search(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.search("query") - mock_client.search.assert_called_once_with("query", limit=5, filter=None) + mock_client.search.assert_called_once_with("query", limit=None, filter=None) assert mock_rich_print_search.call_count == len(mock_results) @@ -176,7 +176,7 @@ async def test_search_no_results(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.search("query") - mock_client.search.assert_called_once_with("query", limit=5, filter=None) + mock_client.search.assert_called_once_with("query", limit=None, filter=None) mock_print.assert_called_once_with("[yellow]No results found.[/yellow]") diff --git a/tests/test_cli.py b/tests/test_cli.py index 4ef0521f..0fbdc05e 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -171,7 +171,7 @@ def test_search(): assert result.exit_code == 0 mock_app_instance.search.assert_called_once_with( - query="query", limit=5, filter=None + query="query", limit=None, filter=None ) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8ddf0d7a..418ea339 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -148,7 +148,7 @@ async def test_mcp_search_documents(): assert result[0].content == "Result 1" assert result[0].score == 0.9 assert result[1].document_id == "doc2" - mock_rag.search.assert_called_once_with("test query", 5) + mock_rag.search.assert_called_once_with("test query", limit=5) @pytest.mark.asyncio diff --git a/tests/test_search.py b/tests/test_search.py index b8b8d012..cf62845d 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -225,8 +225,9 @@ async def test_search_graceful_degradation(temp_db_path): custom_chunks = [ Chunk(content="Custom chunk without docling metadata", metadata={}), ] + docling_doc = await client.convert("Document with custom chunks") await client.import_document( - content="Document with custom chunks", + docling_document=docling_doc, chunks=custom_chunks, uri="https://example.com/custom.html", )