[Cosmos] change retry_write to int, update semantic_reranker api (#43341)

* change `retry_write` to int, update semantic_reranker api, mark provisional

* Update CHANGELOG.md

* changelog, readme, tests

* update offer.py typechecking

* typehint update, set max properly for timeouts,

* Update CHANGELOG with breaking change

Updated CHANGELOG.md to reflect new features, breaking changes, and other modifications.

* Update comments for clarity in set_retry_write method

* Update max write retry count logic

* update tests

Author: simorenoh

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -4,8 +4,10 @@
 
 #### Features Added
 * Added ability to return a tuple of a DatabaseProxy/ContainerProxy with the associated database/container properties when creating or reading databases/containers through `return_properties` parameter. See [PR 41742](https://github.com/Azure/azure-sdk-for-python/pull/41742)
-* Added a new API for Semantic Reranking. See [PR 42991](https://github.com/Azure/azure-sdk-for-python/pull/42991)
+* Added a new **preview feature** API for Semantic Reranking. See [PR 42991](https://github.com/Azure/azure-sdk-for-python/pull/42991)
+
 #### Breaking Changes
+* Changed `retry_write` from `bool` to `int` to match other retryable options. See [PR 43341](https://github.com/Azure/azure-sdk-for-python/pull/43341).
 
 #### Bugs Fixed
 

sdk/cosmos/azure-cosmos/README.md

@@ -686,15 +686,15 @@ The snippet below shows how to enable this feature at the client and request lev
 cosmos_client = CosmosClient(
     url=URL,
     credential=KEY,
-    retry_write=True,  # enables native retryable writes at the client level
+    retry_write=1,  # enables a single native retryable write at the client level
 )
 
 database = cosmos_client.get_database_client(DATABASE_NAME)
 container = database.get_container_client(CONTAINER_NAME)
 
 container.create_item(
     item_body,
-    retry_write=True  # enables native retryable writes at the request level
+    retry_write=1  # enables a single native retryable write at the request level
 )
 ```
 

sdk/cosmos/azure-cosmos/azure/cosmos/_request_object.py

@@ -47,7 +47,7 @@ def __init__(
         self.excluded_locations_circuit_breaker: list[str] = []
         self.healthy_tentative_location: Optional[str] = None
         self.pk_val = pk_val
-        self.retry_write: bool = False
+        self.retry_write: int = 0
 
     def route_to_location_with_preferred_location_flag(  # pylint: disable=name-too-long
         self,
@@ -85,14 +85,16 @@ def set_excluded_location_from_options(self, options: Mapping[str, Any]) -> None
         if self._can_set_excluded_location(options):
             self.excluded_locations = options['excludedLocations']
 
-    def set_retry_write(self, request_options: Mapping[str, Any], client_retry_write: bool) -> None:
+    def set_retry_write(self, request_options: Mapping[str, Any], client_retry_write: int) -> None:
         if self.resource_type == ResourceType.Document:
             if request_options and request_options.get(Constants.Kwargs.RETRY_WRITE):
-                # If request retry write is True, set the option
+                # If request retry write is > 0, set the option
                 self.retry_write = request_options[Constants.Kwargs.RETRY_WRITE]
             elif client_retry_write and self.operation_type != _OperationType.Patch:
-                # If it is not a patch operation and the client config is set, set the retry write to True
+                # If it is not a patch operation and the client config is set, set the retry write to the client value
                 self.retry_write = client_retry_write
+            else:
+                self.retry_write = 0
 
     def set_excluded_locations_from_circuit_breaker(self, excluded_locations: list[str]) -> None: # pylint: disable=name-too-long
         self.excluded_locations_circuit_breaker = excluded_locations

sdk/cosmos/azure-cosmos/azure/cosmos/_retry_utility.py

@@ -297,9 +297,9 @@ def _handle_service_response_retries(request, client, response_retry_policy, exc
         raise exception
 
 def is_write_retryable(request_params, client):
-    return (request_params.retry_write or
-            client.connection_policy.RetryNonIdempotentWrites and
-            not request_params.operation_type == _OperationType.Patch)
+    return (request_params.retry_write > 0 or
+            (client.connection_policy.RetryNonIdempotentWrites > 0 and
+            not request_params.operation_type == _OperationType.Patch))
 
 def _configure_timeout(request: PipelineRequest, absolute: Optional[int], per_request: int) -> None:
     if absolute is not None:

sdk/cosmos/azure-cosmos/azure/cosmos/_service_response_retry_policy.py

@@ -21,13 +21,10 @@ def __init__(self, connection_policy, global_endpoint_manager, pk_range_wrapper,
         self.connection_policy = connection_policy
         self.request = args[0] if args else None
         if self.request:
-            if self.request.retry_write:
+            if self.request.retry_write > 0:
                 # If the request is a write operation, we set the maximum retry count to be the number of
-                # write regional routing contexts available in the global endpoint manager.
-                # This ensures that we retry the write operation across all available regions.
-                # We also ensure that we retry at least once, hence the max is set to 2 by default.
-                self.max_write_retry_count = max(len(self.global_endpoint_manager.
-                                                 location_cache.write_regional_routing_contexts), 2)
+                # write retries provided by the customer.
+                self.max_write_retry_count = self.request.retry_write
             self.location_endpoint = (self.global_endpoint_manager
                                       .resolve_service_endpoint_for_partition(self.request, pk_range_wrapper))
         self.logger = logging.getLogger('azure.cosmos.ServiceResponseRetryPolicy')
@@ -48,9 +45,9 @@ def ShouldRetry(self):
 
         if self.request:
 
-            if not _OperationType.IsReadOnlyOperation(self.request.operation_type) and not self.request.retry_write:
+            if not _OperationType.IsReadOnlyOperation(self.request.operation_type) and not self.request.retry_write > 0:
                 return False
-            if self.request.retry_write and self.failover_retry_count + 1 >= self.max_write_retry_count:
+            if self.request.retry_write > 0 and self.failover_retry_count + 1 >= self.max_write_retry_count:
                 # If we have already retried the write operation to the maximum allowed number of times,
                 # we do not retry further.
                 return False

sdk/cosmos/azure-cosmos/azure/cosmos/_timeout_failover_retry_policy.py

@@ -21,11 +21,9 @@ def __init__(self, connection_policy, global_endpoint_manager, pk_range_wrapper,
         # a request object
         self._max_retry_attempt_count = len(self.global_endpoint_manager.location_cache
                                             .read_regional_routing_contexts) + 1
-       # If the request is a write operation, we only want to retry once if retry write is enabled
+       # If the request is a write operation, we only want to retry as many times as retry_write
         if self.request and _OperationType.IsWriteOperation(self.request.operation_type):
-            self._max_retry_attempt_count = len(
-                self.global_endpoint_manager.location_cache.write_regional_routing_contexts
-            ) + 1
+            self._max_retry_attempt_count = self.request.retry_write
         self.retry_count = 0
         self.connection_policy = connection_policy
 
@@ -70,4 +68,4 @@ def resolve_next_region_service_endpoint(self):
     def is_operation_retryable(self):
         if _OperationType.IsReadOnlyOperation(self.request.operation_type):
             return True
-        return self.request.retry_write
+        return self.request.retry_write > 0

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py

@@ -216,7 +216,7 @@ async def create_item(
         initial_headers: Optional[dict[str, str]] = None,
         priority: Optional[Literal["High", "Low"]] = None,
         no_response: Optional[bool] = None,
-        retry_write: Optional[bool] = None,
+        retry_write: Optional[int] = None,
         throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> CosmosDict:
@@ -246,9 +246,9 @@ async def create_item(
         :keyword bool no_response: Indicates whether service should be instructed to skip
             sending response payloads. When not specified explicitly here, the default value will be determined from
             client-level options.
-        :keyword bool retry_write: Indicates whether the SDK should automatically retry this write operation, even if
+        :keyword int retry_write: Indicates how many times the SDK should automatically retry this write operation, even if
             the operation is not guaranteed to be idempotent. This should only be enabled if the application can
-            tolerate such risks or has logic to safely detect and handle duplicate operations.
+            tolerate such risks or has logic to safely detect and handle duplicate operations. Default is None (no retries).
         :keyword int throughput_bucket: The desired throughput bucket for the client
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: Item with the given ID already exists.
         :returns: A CosmosDict representing the new item. The dict will be empty if `no_response` is specified.
@@ -1103,7 +1103,7 @@ async def upsert_item(
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
         no_response: Optional[bool] = None,
-        retry_write: Optional[bool] = None,
+        retry_write: Optional[int] = None,
         throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> CosmosDict:
@@ -1129,9 +1129,9 @@ async def upsert_item(
         :keyword bool no_response: Indicates whether service should be instructed to skip
             sending response payloads. When not specified explicitly here, the default value will be determined from
             client-level options.
-        :keyword bool retry_write: Indicates whether the SDK should automatically retry this write operation, even if
+        :keyword int retry_write: Indicates how many times the SDK should automatically retry this write operation, even if
             the operation is not guaranteed to be idempotent. This should only be enabled if the application can
-            tolerate such risks or has logic to safely detect and handle duplicate operations.
+            tolerate such risks or has logic to safely detect and handle duplicate operations. Default is None (no retries).
         :keyword int throughput_bucket: The desired throughput bucket for the client
         :keyword Sequence[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations
             in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on.
@@ -1178,29 +1178,28 @@ async def upsert_item(
     @distributed_trace_async
     async def semantic_rerank(
         self,
-        reranking_context: str,
+        *,
+        context: str,
         documents: list[str],
-        semantic_reranking_options: Optional[dict[str, Any]] = None
+        options: Optional[dict[str, Any]] = None
     ) -> CosmosDict:
-        """Rerank a list of documents using semantic reranking.
+        """ **provisional** Rerank a list of documents using semantic reranking.
 
         This method uses a semantic reranker to score and reorder the provided documents
         based on their relevance to the given reranking context.
 
-        :param str reranking_context: The context or query string to use for reranking the documents.
-        :param list[str] documents: A list of documents (as strings) to be reranked.
-        :param dict[str, Any] semantic_reranking_options: Optional dictionary of additional options to customize the semantic reranking process.
+        :keyword str context: The reranking context or query string to use for reranking the documents.
+        :keyword list[str] documents: A list of documents (as strings) to be reranked.
+        :keyword dict[str, Any] options: Optional dictionary of additional request options to customize the semantic reranking process.
 
          Supported options:
-
          * **return_documents** (bool): Whether to return the document text in the response. If False, only scores and indices are returned. Default is True.
          * **top_k** (int): Maximum number of documents to return in the reranked results. If not specified, all documents are returned.
          * **batch_size** (int): Number of documents to process in each batch. Used for optimizing performance with large document sets.
          * **sort** (bool): Whether to sort the results by relevance score in descending order. Default is True.
          * **document_type** (str): Type of documents being reranked. Supported values are "string" and "json".
          * **target_paths** (str): If document_type is "json", the list of JSON paths to extract text from for reranking. Comma-separated string.
 
-        :type semantic_reranking_options: Optional[dict[str, Any]]
         :returns: A CosmosDict containing the reranking results. The structure typically includes results list with reranked documents and their relevance scores. Each result contains index, relevance_score, and optionally document.
         :rtype: ~azure.cosmos.CosmosDict[str, Any]
         :raises ~azure.cosmos.exceptions.CosmosHttpResponseError: If the semantic reranking operation fails.
@@ -1214,9 +1213,9 @@ async def semantic_rerank(
             )
 
         result = await inference_service.rerank(
-            reranking_context=reranking_context,
+            reranking_context=context,
             documents=documents,
-            semantic_reranking_options=semantic_reranking_options
+            semantic_reranking_options=options
         )
 
         return result
@@ -1235,7 +1234,7 @@ async def replace_item(
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
         no_response: Optional[bool] = None,
-        retry_write: Optional[bool] = None,
+        retry_write: Optional[int] = None,
         throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> CosmosDict:
@@ -1262,9 +1261,9 @@ async def replace_item(
         :keyword bool no_response: Indicates whether service should be instructed to skip
             sending response payloads. When not specified explicitly here, the default value will be determined from
             client-level options.
-        :keyword bool retry_write: Indicates whether the SDK should automatically retry this write operation, even if
+        :keyword int retry_write: Indicates how many times the SDK should automatically retry this write operation, even if
             the operation is not guaranteed to be idempotent. This should only be enabled if the application can
-            tolerate such risks or has logic to safely detect and handle duplicate operations.
+            tolerate such risks or has logic to safely detect and handle duplicate operations. Default is None (no retries).
         :keyword int throughput_bucket: The desired throughput bucket for the client
         :keyword Sequence[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations
             in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on.
@@ -1322,7 +1321,7 @@ async def patch_item(
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
         no_response: Optional[bool] = None,
-        retry_write: Optional[bool] = None,
+        retry_write: Optional[int] = None,
         throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> CosmosDict:
@@ -1355,9 +1354,9 @@ async def patch_item(
         :keyword bool no_response: Indicates whether service should be instructed to skip
             sending response payloads. When not specified explicitly here, the default value will be determined from
             client-level options.
-        :keyword bool retry_write: Indicates whether the SDK should automatically retry this write operation, even if
+        :keyword int retry_write: Indicates how many times the SDK should automatically retry this write operation, even if
             the operation is not guaranteed to be idempotent. This should only be enabled if the application can
-            tolerate such risks or has logic to safely detect and handle duplicate operations.
+            tolerate such risks or has logic to safely detect and handle duplicate operations. Default is None (no retries).
         :keyword int throughput_bucket: The desired throughput bucket for the client
         :keyword Sequence[str] excluded_locations: Excluded locations to be skipped from preferred locations. The locations
             in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on.
@@ -1413,7 +1412,7 @@ async def delete_item(
         etag: Optional[str] = None,
         match_condition: Optional[MatchConditions] = None,
         priority: Optional[Literal["High", "Low"]] = None,
-        retry_write: Optional[bool] = None,
+        retry_write: Optional[int] = None,
         throughput_bucket: Optional[int] = None,
         **kwargs: Any
     ) -> None:
@@ -1443,9 +1442,9 @@ async def delete_item(
             in this list are specified as the names of the azure Cosmos locations like, 'West US', 'East US' and so on.
             If all preferred locations were excluded, primary/hub location will be used.
             This excluded_location will override existing excluded_locations in client level.
-        :keyword bool retry_write: Indicates whether the SDK should automatically retry this write operation, even if
+        :keyword int retry_write: Indicates how many times the SDK should automatically retry this write operation, even if
             the operation is not guaranteed to be idempotent. This should only be enabled if the application can
-            tolerate such risks or has logic to safely detect and handle duplicate operations.
+            tolerate such risks or has logic to safely detect and handle duplicate operations. Default is None (no retries).
         :keyword response_hook: A callable invoked with the response metadata.
         :paramtype response_hook: Callable[[Mapping[str, str], None], None]
         :keyword int throughput_bucket: The desired throughput bucket for the client