[ML] Enable retrying on 500 error response from Cohere text embedding API (#105797)

* Retrying on 500

* Update docs/changelog/105797.yaml

---------

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

docs/changelog/105797.yaml

@@ -0,0 +1,5 @@
+pr: 105797
+summary: Enable retrying on 500 error response from Cohere text embedding API
+area: Machine Learning
+type: enhancement
+issues: []

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandler.java

@@ -59,7 +59,9 @@ public class CohereResponseHandler extends BaseResponseHandler {
         }
 
         // handle error codes
-        if (statusCode >= 500) {
+        if (statusCode == 500) {
+            throw new RetryException(true, buildError(SERVER_ERROR, request, result));
+        } else if (statusCode > 500) {
             throw new RetryException(false, buildError(SERVER_ERROR, request, result));
         } else if (statusCode == 429) {
             throw new RetryException(true, buildError(RATE_LIMIT, request, result));

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandlerTests.java

@@ -44,6 +44,16 @@ public class CohereResponseHandlerTests extends ESTestCase {
         MatcherAssert.assertThat(((ElasticsearchStatusException) exception.getCause()).status(), is(RestStatus.BAD_REQUEST));
     }
 
+    public void testCheckForFailureStatusCode_ThrowsFor500_WithShouldRetryTrue() {
+        var exception = expectThrows(RetryException.class, () -> callCheckForFailureStatusCode(500, "id"));
+        assertTrue(exception.shouldRetry());
+        MatcherAssert.assertThat(
+            exception.getCause().getMessage(),
+            containsString("Received a server error status code for request from inference entity id [id] status [500]")
+        );
+        MatcherAssert.assertThat(((ElasticsearchStatusException) exception.getCause()).status(), is(RestStatus.BAD_REQUEST));
+    }
+
     public void testCheckForFailureStatusCode_ThrowsFor429() {
         var exception = expectThrows(RetryException.class, () -> callCheckForFailureStatusCode(429, "id"));
         assertTrue(exception.shouldRetry());