diff --git a/test/query_agent/test_query_model.py b/test/query_agent/test_query_model.py index 02eb7bc..64b56b5 100644 --- a/test/query_agent/test_query_model.py +++ b/test/query_agent/test_query_model.py @@ -845,6 +845,55 @@ def fake_post_with_capture(url, headers=None, json=None, timeout=None): assert captured["json"]["filtering"] is None +def test_search_only_mode_with_ranking_instructions(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 ranking_instructions set — sent verbatim in the payload + instructions = "Prioritize recent documents over older ones." + results = agent.search("test query", limit=2, ranking_instructions=instructions) + assert isinstance(results, SearchModeResponse) + assert captured["json"]["ranking_instructions"] == instructions + + # Reset captured json, then paginate — ranking_instructions should persist + captured = {} + results_2 = results.next(limit=2, offset=1) + assert isinstance(results_2, SearchModeResponse) + assert captured["json"]["ranking_instructions"] == instructions + + +def test_search_only_mode_without_ranking_instructions(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 ranking_instructions — should default to None + results = agent.search("test query", limit=2) + assert isinstance(results, SearchModeResponse) + assert captured["json"]["ranking_instructions"] is None + + def test_search_only_mode_failure(monkeypatch): monkeypatch.setattr(httpx, "post", fake_post_failure) dummy_client = DummyClient() @@ -968,6 +1017,57 @@ 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_ranking_instructions(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 ranking_instructions set — sent verbatim in the payload + instructions = "Prioritize recent documents over older ones." + results = await agent.search( + "test query", limit=2, ranking_instructions=instructions + ) + assert isinstance(results, AsyncSearchModeResponse) + assert captured["json"]["ranking_instructions"] == instructions + + # Reset captured json, then paginate — ranking_instructions should persist + captured = {} + results_2 = await results.next(limit=2, offset=1) + assert isinstance(results_2, AsyncSearchModeResponse) + assert captured["json"]["ranking_instructions"] == instructions + + +async def test_async_search_only_mode_without_ranking_instructions(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 without ranking_instructions — should default to None + results = await agent.search("test query", limit=2) + assert isinstance(results, AsyncSearchModeResponse) + assert captured["json"]["ranking_instructions"] is None + + 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..13fb0d8 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 + ranking_instructions: Optional[str] = None class SearchModeExecutionRequest(SearchModeRequestBase): diff --git a/weaviate_agents/query/query_agent.py b/weaviate_agents/query/query_agent.py index 7f0756c..bffbbea 100644 --- a/weaviate_agents/query/query_agent.py +++ b/weaviate_agents/query/query_agent.py @@ -500,6 +500,7 @@ def search( collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, filtering: Optional[Literal["recall", "precision"]] = None, diversity_weight: Optional[float] = None, + ranking_instructions: Optional[str] = None, ) -> Union[SearchModeResponse, Coroutine[Any, Any, AsyncSearchModeResponse]]: pass @@ -1016,6 +1017,7 @@ def search( collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, filtering: Optional[Literal["recall", "precision"]] = None, diversity_weight: Optional[float] = None, + ranking_instructions: Optional[str] = None, ) -> SearchModeResponse: """Run the Query Agent search-only mode. @@ -1037,6 +1039,11 @@ def search( results with MMR reranking. Higher values push for more topical variety at the cost of relevance. Defaults to None (no diversity). + ranking_instructions: Optional natural language instructions to guide the + instruction-following reranker on how to prioritize relevance + (e.g. "Prioritize recent documents over older ones"). + Only affects the ordering of results, not which results are retrieved. + Limited to ~500 tokens, enforced server-side. Returns: An instance of :class:`~weaviate_agents.query.classes.response.SearchModeResponse` for the first page of results. Use @@ -1073,6 +1080,7 @@ def search( system_prompt=self._system_prompt, filtering=filtering, diversity_weight=diversity_weight, + ranking_instructions=ranking_instructions, ) return searcher.run(limit=limit) @@ -1661,6 +1669,7 @@ async def search( collections: Union[list[Union[str, QueryAgentCollectionConfig]], None] = None, filtering: Optional[Literal["recall", "precision"]] = None, diversity_weight: Optional[float] = None, + ranking_instructions: Optional[str] = None, ) -> AsyncSearchModeResponse: """Run the Query Agent search-only mode. @@ -1683,6 +1692,11 @@ async def search( results with MMR reranking. Higher values push for more topical variety at the cost of relevance. Defaults to None (no diversity). + ranking_instructions: Optional natural language instructions to guide the + instruction-following reranker on how to prioritize relevance + (e.g. "Prioritize recent documents over older ones"). + Only affects the ordering of results, not which results are retrieved. + Limited to ~500 tokens, enforced server-side. Returns: An instance of :class:`~weaviate_agents.query.classes.response.AsyncSearchModeResponse` for the first page of results. Use @@ -1719,6 +1733,7 @@ async def search( system_prompt=self._system_prompt, filtering=filtering, diversity_weight=diversity_weight, + ranking_instructions=ranking_instructions, ) return await searcher.run(limit=limit) diff --git a/weaviate_agents/query/search.py b/weaviate_agents/query/search.py index 6a2feb4..3c8c4ea 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, + ranking_instructions: Optional[str] = None, ): self.headers = headers self.connection_headers = connection_headers @@ -45,8 +46,9 @@ def __init__( self.query = query self.collections = collections self.system_prompt = system_prompt - self.filtering = filtering + self.filtering: Optional[Literal["recall", "precision"]] = filtering self.diversity_weight = diversity_weight + self.ranking_instructions = ranking_instructions 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, + ranking_instructions=self.ranking_instructions, ).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, + ranking_instructions=self.ranking_instructions, ).model_dump(mode="json")