Allow chat agent to specify the search limit

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

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:
+For each user request:
-1. Call the run_search tool with the original query first
+1. Use the run_search tool with the original query
-2. Then call run_search with 1-2 alternative queries using different keywords
+2. Use run_search again with 1-2 alternative keyword queries
-3. Keep queries SHORT (2-5 words) - use keywords, not full sentences
+3. Keep all queries SHORT (2-5 words)
-4. After all searches complete, respond with "Search complete"
+4. After all tool calls complete, respond "Search complete"
 
-Example workflow for "machine learning":
+You can optionally specify a limit parameter (default 5).
-- run_search("machine learning")
-- run_search("neural networks")
-- run_search("deep learning")
-- "Search complete"
 
-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."""

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
+        # Sort by score descending and apply limit
-        limit = self._config.search.limit
+        effective_limit = limit or self._config.search.limit
-        return sorted(seen.values(), key=lambda r: r.score, reverse=True)[:limit]
+        return sorted(seen.values(), key=lambda r: r.score, reverse=True)[
+            :effective_limit
+        ]