From f4ddb0f5c64bb40b164fa0891eb2677f5e98667b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 12 Jan 2026 15:52:40 +0200 Subject: [PATCH] Allow chat agent to specify the search limit --- haiku_rag_slim/haiku/rag/agents/chat/agent.py | 6 ++++-- .../haiku/rag/agents/chat/prompts.py | 20 ++++++++----------- .../haiku/rag/agents/chat/search.py | 16 ++++++++++----- 3 files changed, 23 insertions(+), 19 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index a857f731..2ae7cb2c 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -35,6 +35,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: ctx: RunContext[ChatDeps], query: str, document_name: str | None = None, + limit: int | None = None, ) -> ToolReturn: """Search the knowledge base for relevant documents. @@ -43,14 +44,15 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: Args: query: The search query (what to search for) - document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual") + document_name: Optional document name/title to search within + limit: Number of results to return (default: 5) """ # Build filter from document_name doc_filter = build_document_filter(document_name) if document_name else None # Use search agent for query expansion and deduplication search_agent = SearchAgent(ctx.deps.client, ctx.deps.config) - results = await search_agent.search(query, filter=doc_filter) + results = await search_agent.search(query, filter=doc_filter, limit=limit) # Store for potential citation resolution ctx.deps.search_results = results diff --git a/haiku_rag_slim/haiku/rag/agents/chat/prompts.py b/haiku_rag_slim/haiku/rag/agents/chat/prompts.py index 9fd91aa7..a1213f48 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/prompts.py @@ -27,18 +27,14 @@ IMPORTANT - When user mentions a document in search/ask: Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user.""" -SEARCH_SYSTEM_PROMPT = """You are a search query optimizer for a document knowledge base. +SEARCH_SYSTEM_PROMPT = """You are a search query optimizer. You MUST use the run_search tool to execute searches. -Given a user's search request: -1. Call the run_search tool with the original query first -2. Then call run_search with 1-2 alternative queries using different keywords -3. Keep queries SHORT (2-5 words) - use keywords, not full sentences -4. After all searches complete, respond with "Search complete" +For each user request: +1. Use the run_search tool with the original query +2. Use run_search again with 1-2 alternative keyword queries +3. Keep all queries SHORT (2-5 words) +4. After all tool calls complete, respond "Search complete" -Example workflow for "machine learning": -- run_search("machine learning") -- run_search("neural networks") -- run_search("deep learning") -- "Search complete" +You can optionally specify a limit parameter (default 5). -Do NOT just output queries as text - you MUST call run_search for each query.""" +IMPORTANT: You must make actual tool calls. Do not output "run_search(...)" as text.""" diff --git a/haiku_rag_slim/haiku/rag/agents/chat/search.py b/haiku_rag_slim/haiku/rag/agents/chat/search.py index 689101e2..5c71f31b 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/search.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/search.py @@ -27,15 +27,17 @@ class SearchAgent: async def run_search( ctx: RunContext[SearchDeps], query: str, + limit: int | None = None, ) -> str: """Run a single search query against the knowledge base. Args: query: The search query + limit: Number of results to fetch (default: 5) """ - limit = ctx.deps.config.search.limit + effective_limit = limit or 5 results = await ctx.deps.client.search( - query, limit=limit, filter=ctx.deps.filter + query, limit=effective_limit, filter=ctx.deps.filter ) results = await ctx.deps.client.expand_context(results) ctx.deps.search_results.extend(results) @@ -49,6 +51,7 @@ class SearchAgent: query: str, context: str | None = None, filter: str | None = None, + limit: int | None = None, ) -> list[SearchResult]: """Execute search with query expansion and deduplication. @@ -56,6 +59,7 @@ class SearchAgent: query: The user's search request context: Optional conversation context filter: Optional SQL WHERE clause to filter documents + limit: Maximum number of results to return (default: config limit) Returns: Deduplicated list of SearchResult sorted by score @@ -74,6 +78,8 @@ class SearchAgent: if chunk_id not in seen or result.score > seen[chunk_id].score: seen[chunk_id] = result - # Sort by score descending and limit to config - limit = self._config.search.limit - return sorted(seen.values(), key=lambda r: r.score, reverse=True)[:limit] + # Sort by score descending and apply limit + effective_limit = limit or self._config.search.limit + return sorted(seen.values(), key=lambda r: r.score, reverse=True)[ + :effective_limit + ]