From 1809f1d1e511f2c61e0fc577104b1e240b49441f Mon Sep 17 00:00:00 2001 From: Connor Shorten Date: Tue, 28 Jul 2026 09:14:48 -0400 Subject: [PATCH] first look --- test/query_agent/test_query_model.py | 75 ++++++++++++++++++++++++ weaviate_agents/query/classes/request.py | 1 + weaviate_agents/query/query_agent.py | 10 ++++ weaviate_agents/query/search.py | 4 ++ 4 files changed, 90 insertions(+) diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index 587b40e..7ec2e7c 100644 --- a/test/query_agent/test_query_model.py +++ b/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() diff --git a/weaviate_agents/query/classes/request.py b/weaviate_agents/query/classes/request.py index 1247aa7..55a8fdc 100644 --- a/weaviate_agents/query/classes/request.py +++ b/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): diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index b9fdc53..b38975c 100644 --- a/weaviate_agents/query/query_agent.py +++ b/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) diff --git a/weaviate_agents/query/search.py b/weaviate_agents/query/search.py index 6a2feb4..28d9d48 100644 --- a/weaviate_agents/query/search.py +++ b/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")