Fix the answer query bug (#35259)

* Fix the answer query bug

* Add tests

* Update CHANGELOG.md

* Update CHANGELOG.md

* update

* Update sdk/search/azure-search-documents/azure/search/documents/_utils.py

Co-authored-by: Paul Van Eck <[email protected]>

---------

Co-authored-by: Paul Van Eck <[email protected]>

Author: xiangyan99

sdk/search/azure-search-documents/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fixed the bug that SearchClient failed when both answer count and answer threshold applied.
+
 ### Other Changes
 
 ## 11.6.0b3 (2024-04-09)

sdk/search/azure-search-documents/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/search/azure-search-documents",
-  "Tag": "python/search/azure-search-documents_a3db22f661"
+  "Tag": "python/search/azure-search-documents_dd0ce4c420"
 }

sdk/search/azure-search-documents/azure/search/documents/_search_client.py

@@ -34,7 +34,7 @@
 from ._paging import SearchItemPaged, SearchPageIterator
 from ._queries import AutocompleteQuery, SearchQuery, SuggestQuery
 from ._headers_mixin import HeadersMixin
-from ._utils import get_authentication_policy
+from ._utils import get_authentication_policy, get_answer_query
 from ._version import SDK_MONIKER
 
 
@@ -326,8 +326,7 @@ def search(
         filter_arg = filter
         search_fields_str = ",".join(search_fields) if search_fields else None
 
-        answers = query_answer if not query_answer_count else "{}|count-{}".format(query_answer, query_answer_count)
-        answers = answers if not query_answer_threshold else "{}|threshold-{}".format(answers, query_answer_threshold)
+        answers = get_answer_query(query_answer, query_answer_count, query_answer_threshold)
 
         captions = (
             query_caption

sdk/search/azure-search-documents/azure/search/documents/_utils.py

@@ -3,15 +3,31 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from typing import Any, Optional
+from typing import Any, Optional, Union
 from azure.core.pipeline.policies import (
     BearerTokenCredentialPolicy,
     AsyncBearerTokenCredentialPolicy,
 )
+from ._generated.models import QueryAnswerType
 
 DEFAULT_AUDIENCE = "https://search.azure.com"
 
 
+def get_answer_query(
+    query_answer: Optional[Union[str, QueryAnswerType]] = None,
+    query_answer_count: Optional[int] = None,
+    query_answer_threshold: Optional[float] = None,
+) -> Optional[Union[str, QueryAnswerType]]:
+    answers = query_answer
+    separator = "|"
+    if query_answer_count:
+        answers = f"{answers}{separator}count-{query_answer_count}"
+        separator = ","
+    if query_answer_threshold:
+        answers = f"{answers}{separator}threshold-{query_answer_threshold}"
+    return answers
+
+
 def is_retryable_status_code(status_code: Optional[int]) -> bool:
     if not status_code:
         return False

sdk/search/azure-search-documents/azure/search/documents/aio/_search_client_async.py

@@ -10,7 +10,7 @@
 from azure.core.credentials_async import AsyncTokenCredential
 from azure.core.tracing.decorator_async import distributed_trace_async
 from ._paging import AsyncSearchItemPaged, AsyncSearchPageIterator
-from .._utils import get_authentication_policy
+from .._utils import get_authentication_policy, get_answer_query
 from .._generated.aio import SearchIndexClient
 from .._generated.models import (
     AutocompleteMode,
@@ -329,8 +329,7 @@ async def search(
         include_total_result_count = include_total_count
         filter_arg = filter
         search_fields_str = ",".join(search_fields) if search_fields else None
-        answers = query_answer if not query_answer_count else "{}|count-{}".format(query_answer, query_answer_count)
-        answers = answers if not query_answer_threshold else "{}|threshold-{}".format(answers, query_answer_threshold)
+        answers = get_answer_query(query_answer, query_answer_count, query_answer_threshold)
         captions = (
             query_caption
             if not query_caption_highlight_enabled

sdk/search/azure-search-documents/tests/test_search_client.py

@@ -20,7 +20,7 @@
     RequestEntityTooLargeError,
     ApiVersion,
 )
-from azure.search.documents._utils import odata
+from azure.search.documents._utils import odata, get_answer_query
 
 CREDENTIAL = AzureKeyCredential(key="test_api_key")
 
@@ -56,6 +56,23 @@ def test_prevent_double_quoting(self):
         assert odata("foo eq '{foo}'", foo="a string") == "foo eq 'a string'"
 
 
+class TestAnswerQuery:
+    def test_no_args(self):
+        assert get_answer_query() is None
+
+    def test_query_answer(self):
+        assert get_answer_query("query") == "query"
+
+    def test_query_answer_count(self):
+        assert get_answer_query("query", 5) == "query|count-5"
+
+    def test_query_answer_threshold(self):
+        assert get_answer_query("query", query_answer_threshold=0.5) == "query|threshold-0.5"
+
+    def test_query_answer_count_threshold(self):
+        assert get_answer_query("query", 5, 0.5) == "query|count-5,threshold-0.5"
+
+
 class TestSearchClient:
     def test_init(self):
         client = SearchClient("endpoint", "index name", CREDENTIAL)