diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index 3bb0c76..02eb7bc 100644 --- a/test/query_agent/test_query_model.py +++ b/test/query_agent/test_query_model.py @@ -797,6 +797,54 @@ def fake_post_with_capture(url, headers=None, json=None, timeout=None): assert captured["json"]["diversity_weight"] is None +def test_search_only_mode_with_filtering(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 filtering set + results = agent.search("test query", limit=2, filtering="recall") + assert isinstance(results, SearchModeResponse) + assert captured["json"]["filtering"] == "recall" + + # Reset captured json, then paginate — filtering should persist + captured = {} + results_2 = results.next(limit=2, offset=1) + assert isinstance(results_2, SearchModeResponse) + assert captured["json"]["filtering"] == "recall" + + +def test_search_only_mode_default_filtering(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 filtering — should default to None + results = agent.search("test query", limit=2) + assert isinstance(results, SearchModeResponse) + assert captured["json"]["filtering"] is None + + def test_search_only_mode_failure(monkeypatch): monkeypatch.setattr(httpx, "post", fake_post_failure) dummy_client = DummyClient() @@ -893,6 +941,33 @@ async def fake_post_with_capture(self, url, headers=None, json=None, timeout=Non assert captured["json"]["diversity_weight"] == 0.7 +async def test_async_search_only_mode_with_filtering(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 filtering set + results = await agent.search("test query", limit=2, filtering="precision") + assert isinstance(results, AsyncSearchModeResponse) + assert captured["json"]["filtering"] == "precision" + + # Reset captured json, then paginate — filtering should persist + captured = {} + results_2 = await results.next(limit=2, offset=1) + assert isinstance(results_2, AsyncSearchModeResponse) + assert captured["json"]["filtering"] == "precision" + + 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 aee9ee7..1247aa7 100644 --- a/weaviate_agents/query/classes/request.py +++ b/weaviate_agents/query/classes/request.py @@ -17,6 +17,7 @@ class SearchModeRequestBase(BaseModel): collections: list[Union[str, QueryAgentCollectionConfig]] limit: int offset: int + filtering: Optional[Literal["recall", "precision"]] = None diversity_weight: Optional[float] = None diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index d3e9a64..7f0756c 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -498,6 +498,7 @@ def search( query: Union[str, list[ChatMessage]], limit: int = 20, collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, + filtering: Optional[Literal["recall", "precision"]] = None, diversity_weight: Optional[float] = None, ) -> Union[SearchModeResponse, Coroutine[Any, Any, AsyncSearchModeResponse]]: pass @@ -1013,6 +1014,7 @@ def search( query: Union[str, list[ChatMessage]], limit: int = 20, collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, + filtering: Optional[Literal["recall", "precision"]] = None, diversity_weight: Optional[float] = None, ) -> SearchModeResponse: """Run the Query Agent search-only mode. @@ -1027,6 +1029,10 @@ def search( limit: The maximum number of results to return for the first page. collections: The collections to query. Either a list of strings, or a list of :class:`~weaviate_agents.query.classes.QueryAgentCollectionConfig` objects. Overrides any collections provided in the constructor when set. + filtering: The filtering strategy to use for this search. + Use "recall" to optimize for finding all relevant results, + or "precision" to optimize for the accuracy of returned results. + Defaults to "recall". diversity_weight: Optional float between 0.0 and 1.0 to diversify results with MMR reranking. Higher values push for more topical variety at the cost of relevance. @@ -1065,6 +1071,7 @@ def search( query=query, collections=collections, system_prompt=self._system_prompt, + filtering=filtering, diversity_weight=diversity_weight, ) return searcher.run(limit=limit) @@ -1652,6 +1659,7 @@ async def search( query: Union[str, list[ChatMessage]], limit: int = 20, collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, + filtering: Optional[Literal["recall", "precision"]] = None, diversity_weight: Optional[float] = None, ) -> AsyncSearchModeResponse: """Run the Query Agent search-only mode. @@ -1667,6 +1675,10 @@ async def search( limit: The maximum number of results to return for the first page. collections: The collections to query. Overrides any collections provided in the constructor when set. + filtering: The filtering strategy to use for this search. + Use "recall" to optimize for finding all relevant results, + or "precision" to optimize for the accuracy of returned results. + Defaults to "recall". diversity_weight: Optional float between 0.0 and 1.0 to diversify results with MMR reranking. Higher values push for more topical variety at the cost of relevance. @@ -1705,6 +1717,7 @@ async def search( query=query, collections=collections, system_prompt=self._system_prompt, + filtering=filtering, diversity_weight=diversity_weight, ) return await searcher.run(limit=limit) diff --git a/weaviate_agents/query/search.py b/weaviate_agents/query/search.py index 77c2365..6a2feb4 100644 --- a/weaviate_agents/query/search.py +++ b/weaviate_agents/query/search.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Optional, Union +from typing import Any, Literal, Optional, Union import httpx @@ -35,6 +35,7 @@ def __init__( query: Union[str, list[ChatMessage]], collections: list[Union[str, QueryAgentCollectionConfig]], system_prompt: Optional[str], + filtering: Optional[Literal["recall", "precision"]] = None, diversity_weight: Optional[float] = None, ): self.headers = headers @@ -44,6 +45,7 @@ def __init__( self.query = query self.collections = collections self.system_prompt = system_prompt + self.filtering = filtering self.diversity_weight = diversity_weight self._cached_searches: Optional[list[QueryResultWithCollectionNormalized]] = ( None @@ -63,6 +65,7 @@ def _get_request_body(self, limit: int, offset: int) -> dict[str, Any]: limit=limit, offset=offset, system_prompt=self.system_prompt, + filtering=self.filtering, diversity_weight=self.diversity_weight, ).model_dump(mode="json") else: @@ -73,6 +76,7 @@ def _get_request_body(self, limit: int, offset: int) -> dict[str, Any]: limit=limit, offset=offset, searches=self._cached_searches, + filtering=self.filtering, diversity_weight=self.diversity_weight, ).model_dump(mode="json") @@ -93,7 +97,7 @@ def _handle_response(self, response: httpx.Response) -> SearchModeResponse: raise Exception(response.text) parsed_response = SearchModeResponse(**response.json()) - if parsed_response.searches: + if parsed_response.searches is not None: self._cached_searches = parsed_response.searches parsed_response._searcher = self return parsed_response @@ -138,7 +142,7 @@ def _handle_response(self, response: httpx.Response) -> AsyncSearchModeResponse: raise Exception(response.text) parsed_response = AsyncSearchModeResponse(**response.json()) - if parsed_response.searches: + if parsed_response.searches is not None: self._cached_searches = parsed_response.searches parsed_response._searcher = self return parsed_response