fix(azure/search): reject bool and non-positive max_results for bing_grounding count

Extract a shared _valid_max_results predicate that rejects bools (an int
subclass) and non-positive values, and reuse it from both the connection-mode
request count and the response-side cap so both paths honor the same contract.
This commit is contained in:
mateo-berri 2026-08-24 12:53:44 -07:00
parent a0c319878b
commit aef7424637
2 changed files with 30 additions and 5 deletions

View file

@ -178,6 +178,16 @@ def _citation_results(envelope: _ResponsesEnvelope) -> tuple[SearchResult, ...]:
return tuple(first_by_url[url] for url in dict.fromkeys(result.url for result in cited))
def _valid_max_results(max_results: object) -> int | None:
"""A positive-int `max_results`, else None. Rejects bools, an `int` subclass, and
non-positive values so neither the request-side `count` nor the response-side cap
forwards a value the other would silently ignore.
"""
if isinstance(max_results, bool) or not isinstance(max_results, int):
return None
return max_results if max_results > 0 else None
def _requested_max_results(response_kwargs: Mapping[str, object]) -> int | None:
"""The unified `max_results` cap the caller asked for, if any.
@ -188,10 +198,7 @@ def _requested_max_results(response_kwargs: Mapping[str, object]) -> int | None:
optional_params: Final = response_kwargs.get("optional_params")
if not isinstance(optional_params, Mapping):
return None
max_results: Final = optional_params.get("max_results")
return (
max_results if isinstance(max_results, int) and not isinstance(max_results, bool) and max_results > 0 else None
)
return _valid_max_results(optional_params.get("max_results"))
def _capped(results: tuple[SearchResult, ...], max_results: int | None) -> tuple[SearchResult, ...]:
@ -247,7 +254,7 @@ def _search_tool(optional_params: Mapping[str, object]) -> _BingGroundingTool |
if connection_id:
configuration: Final = _SearchConfiguration(
project_connection_id=connection_id,
count=max_results if isinstance(max_results, int) else None,
count=_valid_max_results(max_results),
)
return _BingGroundingTool(bing_grounding=_BingGroundingParams(search_configurations=(configuration,)))
location: Final = _UserLocation(country=country.upper()) if isinstance(country, str) else None

View file

@ -195,6 +195,24 @@ def test_transform_search_request_connection_mode_omits_count_without_max_result
assert body["tools"][0]["bing_grounding"]["search_configurations"] == [{"project_connection_id": "conn-id"}]
@pytest.mark.parametrize("max_results", [True, False, 0, -1])
def test_transform_search_request_connection_mode_omits_count_for_invalid_max_results(
monkeypatch: pytest.MonkeyPatch, max_results: object
):
monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1")
monkeypatch.setenv("BING_GROUNDING_CONNECTION_ID", "conn-id")
body = _config().transform_search_request("q", {"max_results": max_results})
assert body["tools"][0]["bing_grounding"]["search_configurations"] == [{"project_connection_id": "conn-id"}]
def test_transform_search_response_ignores_invalid_max_results_cap():
annotations = [_citation(f"https://example.com/{i}", f"T{i}", 0, 5) for i in range(3)]
resp = _config().transform_search_response(
_resp(_message_response("claim", annotations)), logging_obj=Mock(), optional_params={"max_results": True}
)
assert [r.url for r in resp.results] == [f"https://example.com/{i}" for i in range(3)]
def test_transform_search_request_joins_list_query(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1")
assert _config().transform_search_request(["foo", "bar"], {})["input"] == "foo bar"