Add effort to Search Mode

test/query_agent/test_query_model.py

@@ -845,6 +845,54 @@ def fake_post_with_capture(url, headers=None, json=None, timeout=None):
     assert captured["json"]["filtering"] is None
 
 
+def test_search_only_mode_with_effort(monkeypatch):
+    captured = {}
+
+    def fake_post_with_capture(url, headers=None, json=None, timeout=None):
+        captured["json"] = json
+        return fake_post_search_only_success()
+
+    monkeypatch.setattr(httpx, "post", fake_post_with_capture)
+    dummy_client = DummyClient()
+    agent = QueryAgent(
+        dummy_client, ["test_collection"], agents_host="http://dummy-agent"
+    )
+    agent._connection = dummy_client
+    agent._headers = dummy_client.additional_headers
+
+    # Test with effort set
+    results = agent.search("test query", limit=2, effort="high")
+    assert isinstance(results, SearchModeResponse)
+    assert captured["json"]["effort"] == "high"
+
+    # Reset captured json, then paginate — effort should persist
+    captured = {}
+    results_2 = results.next(limit=2, offset=1)
+    assert isinstance(results_2, SearchModeResponse)
+    assert captured["json"]["effort"] == "high"
+
+
+def test_search_only_mode_default_effort(monkeypatch):
+    captured = {}
+
+    def fake_post_with_capture(url, headers=None, json=None, timeout=None):
+        captured["json"] = json
+        return fake_post_search_only_success()
+
+    monkeypatch.setattr(httpx, "post", fake_post_with_capture)
+    dummy_client = DummyClient()
+    agent = QueryAgent(
+        dummy_client, ["test_collection"], agents_host="http://dummy-agent"
+    )
+    agent._connection = dummy_client
+    agent._headers = dummy_client.additional_headers
+
+    # Test without effort — should default to None
+    results = agent.search("test query", limit=2)
+    assert isinstance(results, SearchModeResponse)
+    assert captured["json"]["effort"] is None
+
+
 def test_search_only_mode_failure(monkeypatch):
     monkeypatch.setattr(httpx, "post", fake_post_failure)
     dummy_client = DummyClient()
@@ -968,6 +1016,33 @@ async def fake_post_with_capture(self, url, headers=None, json=None, timeout=Non
     assert captured["json"]["filtering"] == "precision"
 
 
+async def test_async_search_only_mode_with_effort(monkeypatch):
+    captured = {}
+
+    async def fake_post_with_capture(self, url, headers=None, json=None, timeout=None):
+        captured["json"] = json
+        return await fake_async_post_search_only_success()
+
+    monkeypatch.setattr(httpx.AsyncClient, "post", fake_post_with_capture)
+    dummy_client = DummyClient()
+    agent = AsyncQueryAgent(
+        dummy_client, ["test_collection"], agents_host="http://dummy-agent"
+    )
+    agent._connection = dummy_client
+    agent._headers = dummy_client.additional_headers
+
+    # Test with effort set
+    results = await agent.search("test query", limit=2, effort="low")
+    assert isinstance(results, AsyncSearchModeResponse)
+    assert captured["json"]["effort"] == "low"
+
+    # Reset captured json, then paginate — effort should persist
+    captured = {}
+    results_2 = await results.next(limit=2, offset=1)
+    assert isinstance(results_2, AsyncSearchModeResponse)
+    assert captured["json"]["effort"] == "low"
+
+
 async def test_async_search_only_mode_failure(monkeypatch):
     monkeypatch.setattr(httpx.AsyncClient, "post", fake_async_post_failure)
     dummy_client = DummyClient()

weaviate_agents/query/classes/request.py

@@ -19,6 +19,7 @@ class SearchModeRequestBase(BaseModel):
     offset: int
     filtering: Optional[Literal["recall", "precision"]] = None
     diversity_weight: Optional[float] = None
+    effort: Optional[Literal["low", "medium", "high"]] = None
 
 
 class SearchModeExecutionRequest(SearchModeRequestBase):

weaviate_agents/query/query_agent.py

@@ -1405,6 +1405,7 @@ def search(
         collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
         filtering: Optional[Literal["recall", "precision"]] = None,
         diversity_weight: Optional[float] = None,
+        effort: Optional[Literal["low", "medium", "high"]] = None,
     ) -> SearchModeResponse:
         """Run the Query Agent search-only mode.
 
@@ -1426,6 +1427,9 @@ def search(
                 results with MMR reranking.
                 Higher values push for more topical variety at the cost of relevance.
                 Defaults to None (no diversity).
+            effort: The amount of effort the agent should put into the search.
+                One of "low", "medium", or "high". Higher effort may improve
+                result quality at the expense of increased latency and cost.
 
         Returns:
             An instance of :class:`~weaviate_agents.query.classes.response.SearchModeResponse` for the first page of results. Use
@@ -1462,6 +1466,7 @@ def search(
             system_prompt=self._system_prompt,
             filtering=filtering,
             diversity_weight=diversity_weight,
+            effort=effort,
         )
         return searcher.run(limit=limit)
 
@@ -2257,6 +2262,7 @@ async def search(
         collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None,
         filtering: Optional[Literal["recall", "precision"]] = None,
         diversity_weight: Optional[float] = None,
+        effort: Optional[Literal["low", "medium", "high"]] = None,
     ) -> AsyncSearchModeResponse:
         """Run the Query Agent search-only mode.
 
@@ -2279,6 +2285,9 @@ async def search(
                 results with MMR reranking.
                 Higher values push for more topical variety at the cost of relevance.
                 Defaults to None (no diversity).
+            effort: The amount of effort the agent should put into the search.
+                One of "low", "medium", or "high". Higher effort may improve
+                result quality at the expense of increased latency and cost.
 
         Returns:
             An instance of :class:`~weaviate_agents.query.classes.response.AsyncSearchModeResponse` for the first page of results. Use
@@ -2315,6 +2324,7 @@ async def search(
             system_prompt=self._system_prompt,
             filtering=filtering,
             diversity_weight=diversity_weight,
+            effort=effort,
         )
         return await searcher.run(limit=limit)
 

weaviate_agents/query/search.py

@@ -37,6 +37,7 @@ def __init__(
         system_prompt: Optional[str],
         filtering: Optional[Literal["recall", "precision"]] = None,
         diversity_weight: Optional[float] = None,
+        effort: Optional[Literal["low", "medium", "high"]] = None,
     ):
         self.headers = headers
         self.connection_headers = connection_headers
@@ -47,6 +48,7 @@ def __init__(
         self.system_prompt = system_prompt
         self.filtering = filtering
         self.diversity_weight = diversity_weight
+        self.effort = effort
         self._cached_searches: Optional[list[QueryResultWithCollectionNormalized]] = (
             None
         )
@@ -67,6 +69,7 @@ def _get_request_body(self, limit: int, offset: int) -> dict[str, Any]:
                 system_prompt=self.system_prompt,
                 filtering=self.filtering,
                 diversity_weight=self.diversity_weight,
+                effort=self.effort,
             ).model_dump(mode="json")
         else:
             return SearchModeExecutionRequest(
@@ -78,6 +81,7 @@ def _get_request_body(self, limit: int, offset: int) -> dict[str, Any]:
                 searches=self._cached_searches,
                 filtering=self.filtering,
                 diversity_weight=self.diversity_weight,
+                effort=self.effort,
             ).model_dump(mode="json")