Cap QA agent search iterations to reduce response time

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Changed
+
+- **QA search cap**: Replace dead `max_iterations`/`max_concurrency` config with `max_searches` (default: 3). The QA agent now enforces a per-run search limit, reducing average response time from ~30s to ~15s while maintaining accuracy. The limit resets per agent run so toolsets can be safely reused.
+
 ## [0.33.1] - 2026-03-06
 
 ### Changed

docs/configuration/index.md

@@ -86,8 +86,7 @@ qa:
     name: gpt-oss
     enable_thinking: true
     temperature: 0.3
-  max_iterations: 2
-  max_concurrency: 1
+  max_searches: 3
 
 research:
   model:

docs/configuration/qa-research.md

@@ -33,13 +33,11 @@ qa:
     name: gpt-oss
     enable_thinking: true
     temperature: 0.3          # Default: 0.3
-  max_iterations: 2     # Maximum search iterations
-  max_concurrency: 1    # Concurrent search operations
+  max_searches: 3       # Maximum search tool calls per question
 ```
 
 - **model**: LLM configuration (see [Providers](providers.md#model-settings))
-- **max_iterations**: Maximum search iterations (default: 2)
-- **max_concurrency**: Number of concurrent search operations (default: 1)
+- **max_searches**: Maximum number of search tool calls the QA agent can make per question (default: 3)
 
 ## Research Configuration
 

evaluations/evaluations/benchmark.py

@@ -61,6 +61,7 @@ def build_experiment_metadata(
         "qa_temperature": config.qa.model.temperature,
         "qa_max_tokens": config.qa.model.max_tokens,
         "qa_enable_thinking": config.qa.model.enable_thinking,
+        "qa_max_searches": config.qa.max_searches,
     }
     if judge_config is not None:
         metadata.update(

haiku_rag_slim/haiku/rag/agents/qa/agent.py

@@ -47,21 +47,27 @@ class QuestionAnswerAgent:
             Tuple of (answer text, list of resolved citations)
         """
         accumulated_results: list[SearchResult] = []
+        max_searches = self._config.qa.max_searches
         search_toolset = create_search_toolset(
             self._config,
             base_filter=filter,
-            tool_name="search_documents",
+            tool_name="search",
             on_results=accumulated_results.extend,
+            max_searches=max_searches,
         )
 
         # Agent created per-call: toolset varies with filter, and Agent
         # construction is pure Python (no IO).
         model = get_model(self._model_config, self._config)
+        try:
+            system_prompt = self._system_prompt.format(max_searches=max_searches)
+        except KeyError:
+            system_prompt = self._system_prompt
         agent: Agent[_QARunDeps, RawSearchAnswer] = Agent(  # ty: ignore[invalid-assignment]
             model=model,
             deps_type=_QARunDeps,
             output_type=structured_output_type(RawSearchAnswer, model),
-            instructions=self._system_prompt,
+            instructions=system_prompt,
             toolsets=[search_toolset],
             retries=3,
         )

haiku_rag_slim/haiku/rag/agents/qa/prompts.py

@@ -1,9 +1,9 @@
 QA_SYSTEM_PROMPT = """You are a knowledgeable assistant that answers questions using a document knowledge base.
 
 Process:
-1. Call search_documents with relevant keywords from the question
+1. Call search with relevant keywords from the question
 2. Review the results ordered by relevance
-3. If needed, perform follow-up searches with different keywords (max 3 total)
+3. If needed, perform follow-up searches with different keywords (max {max_searches} total)
 4. Provide a concise answer based strictly on the retrieved content
 
 The search tool returns results like:
@@ -35,4 +35,5 @@ Guidelines:
 - If information is insufficient, say: "I cannot find enough information in the knowledge base to answer this question."
 - Be concise and direct - avoid elaboration unless asked
 - Results are ordered by relevance, with rank 1 being most relevant
+- If the search tool tells you the search limit is reached, stop searching immediately and answer with what you have
 """

haiku_rag_slim/haiku/rag/config/models.py

@@ -79,8 +79,7 @@ class QAConfig(BaseModel):
             temperature=0.3,
         )
     )
-    max_iterations: int = 2
-    max_concurrency: int = 1
+    max_searches: int = 3
 
 
 class ResearchConfig(BaseModel):

haiku_rag_slim/haiku/rag/tools/search.py

@@ -13,6 +13,7 @@ def create_search_toolset(
     base_filter: str | None = None,
     tool_name: str = "search",
     on_results: Callable[[list[SearchResult]], None] | None = None,
+    max_searches: int | None = None,
 ) -> FunctionToolset[RAGDeps]:
     """Create a toolset with search capabilities.
 
@@ -25,10 +26,16 @@ def create_search_toolset(
         tool_name: Name for the search tool. Defaults to "search".
         on_results: Optional callback invoked with search results after each search.
             Useful for accumulating results externally (e.g., for citation resolution).
+        max_searches: Maximum number of searches allowed. When exceeded, returns
+            a message directing the agent to answer with existing results.
 
     Returns:
         FunctionToolset with a search tool.
     """
+    # Per-run search counter. Resets when run_id changes so the toolset
+    # can safely be reused across multiple agent.run() calls.
+    search_count = 0
+    current_run_id: str | None = None
 
     async def search(
         ctx: RunContext[RAGDeps],
@@ -44,6 +51,17 @@ def create_search_toolset(
         Returns:
             Formatted search results with content and metadata.
         """
+        nonlocal search_count, current_run_id
+        if ctx.run_id != current_run_id:
+            current_run_id = ctx.run_id
+            search_count = 0
+        search_count += 1
+        if max_searches is not None and search_count > max_searches:
+            return (
+                "Search limit reached. "
+                "Answer the question using the results you already have."
+            )
+
         client = ctx.deps.client
 
         effective_filter = base_filter
@@ -71,5 +89,5 @@ def create_search_toolset(
         return "\n\n".join(formatted)
 
     toolset: FunctionToolset[RAGDeps] = FunctionToolset()
-    toolset.add_function(search, name=tool_name)
+    toolset.add_function(search, name=tool_name, retries=3)
     return toolset

tests/tools/test_search.py

@@ -12,9 +12,9 @@ def vcr_cassette_dir():
     return str(Path(__file__).parent.parent / "cassettes" / "test_search_tools")
 
 
-def make_ctx(client):
+def make_ctx(client, run_id="test-run"):
     """Create a lightweight RunContext-like object for direct tool function calls."""
-    return SimpleNamespace(deps=SimpleNamespace(client=client))
+    return SimpleNamespace(deps=SimpleNamespace(client=client), run_id=run_id)
 
 
 @pytest.mark.vcr()
@@ -126,6 +126,69 @@ class TestSearchToolExecution:
         assert "Python" in result or "programming" in result
 
 
+@pytest.mark.vcr()
+class TestSearchMaxSearches:
+    """Tests for max_searches cap on search toolset."""
+
+    @pytest.mark.asyncio
+    async def test_searches_within_limit_return_results(
+        self, search_client, search_config
+    ):
+        """Searches within max_searches return normal results."""
+        toolset = create_search_toolset(search_config, max_searches=2)
+        search_tool = toolset.tools["search"]
+        ctx = make_ctx(search_client)
+
+        result1 = await search_tool.function(ctx, "Python")
+        assert "Search limit reached" not in result1
+
+        result2 = await search_tool.function(ctx, "JavaScript")
+        assert "Search limit reached" not in result2
+
+    @pytest.mark.asyncio
+    async def test_searches_beyond_limit_return_cap_message(
+        self, search_client, search_config
+    ):
+        """Searches beyond max_searches return limit message."""
+        toolset = create_search_toolset(search_config, max_searches=1)
+        search_tool = toolset.tools["search"]
+        ctx = make_ctx(search_client)
+
+        result1 = await search_tool.function(ctx, "Python")
+        assert "Search limit reached" not in result1
+
+        result2 = await search_tool.function(ctx, "JavaScript")
+        assert "Search limit reached" in result2
+
+    @pytest.mark.asyncio
+    async def test_counter_resets_across_runs(self, search_client, search_config):
+        """Search counter resets when run_id changes (new agent run)."""
+        toolset = create_search_toolset(search_config, max_searches=1)
+        search_tool = toolset.tools["search"]
+
+        ctx_run1 = make_ctx(search_client, run_id="run-1")
+        result = await search_tool.function(ctx_run1, "Python")
+        assert "Search limit reached" not in result
+
+        result2 = await search_tool.function(ctx_run1, "JavaScript")
+        assert "Search limit reached" in result2
+
+        ctx_run2 = make_ctx(search_client, run_id="run-2")
+        result3 = await search_tool.function(ctx_run2, "Python")
+        assert "Search limit reached" not in result3
+
+    @pytest.mark.asyncio
+    async def test_no_limit_by_default(self, search_client, search_config):
+        """Without max_searches, searches are unlimited."""
+        toolset = create_search_toolset(search_config)
+        search_tool = toolset.tools["search"]
+        ctx = make_ctx(search_client)
+
+        for _ in range(5):
+            result = await search_tool.function(ctx, "Python")
+            assert "Search limit reached" not in result
+
+
 @pytest.fixture
 async def search_client(temp_db_path):
     """Create a HaikuRAG client with test data for search tests."""