From 097616053b25df55cc88283b3857ae4a9c902a35 Mon Sep 17 00:00:00 2001
From: GitHub Copilot CLI <223556219+Copilot@users.noreply.github.com>
Date: Tue, 12 May 2026 16:59:31 -0700
Subject: [PATCH 1/6] Adding InferenceService for semantic reranking (rebased;
 review feedback addressed)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Squashed rebase of PR #48258 onto upstream/main. Addresses outstanding review comments from kushagraThapar and FabianMeiswinkel:

- Configs.java: INFERENCE_ENDPOINT_PROPERTY / INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE are now private (single source of truth, with comment); removed unused getDefaultInferenceServiceConfig().

- InferenceService.java: consolidated duplicate MAX_CONNECTION_LIMIT constant (MAX_CONNECTION_LIMIT_VARIABLE = MAX_CONNECTION_LIMIT_PROPERTY); rewrote withRetry to honour the server-supplied Retry-After header on 429 (delay-seconds form per RFC 7231, capped at RETRY_MAX_BACKOFF), falling back to exponential backoff with jitter — matches Python SDK behaviour.

- SemanticRerankResult / SemanticRerankScore: removed @JsonProperty annotations (POJOs are populated via ModelBridgeInternal setters, not via Jackson deserialization).

- SemanticRerankSample: removed import of internal InferenceService; redefined option-key constants as raw string literals locally in the sample.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../inference/InferenceServiceTest.java       | 628 ++++++++++++++++++
 sdk/cosmos/azure-cosmos/CHANGELOG.md          |   1 +
 .../com/azure/cosmos/CosmosAsyncClient.java   |  44 +-
 .../azure/cosmos/CosmosAsyncContainer.java    |  57 +-
 .../com/azure/cosmos/CosmosContainer.java     |  57 ++
 .../azure/cosmos/implementation/Configs.java  |  18 +
 .../inference/InferenceService.java           | 513 ++++++++++++++
 .../cosmos/models/ModelBridgeInternal.java    |  33 +
 .../cosmos/models/SemanticRerankResult.java   |  78 +++
 .../cosmos/models/SemanticRerankScore.java    |  75 +++
 .../main/java/com/azure/cosmos/util/Beta.java |   4 +-
 .../semanticrerank/SemanticRerankSample.java  | 130 ++++
 12 files changed, 1635 insertions(+), 3 deletions(-)
 create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
 create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
 create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
 create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java
 create mode 100644 sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/semanticrerank/SemanticRerankSample.java

diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
new file mode 100644
index 000000000000..24bbcb1d6da2
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
@@ -0,0 +1,628 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.implementation.inference;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.cosmos.implementation.BadRequestException;
+import com.azure.cosmos.implementation.HttpConstants;
+import com.azure.cosmos.implementation.http.HttpClient;
+import com.azure.cosmos.implementation.http.HttpClientConfig;
+import com.azure.cosmos.implementation.http.HttpRequest;
+import com.azure.cosmos.implementation.http.HttpResponse;
+import org.mockito.ArgumentCaptor;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+import reactor.core.Exceptions;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.time.OffsetDateTime;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.*;
+
+public class InferenceServiceTest {
+
+    private static final String MOCK_TOKEN = "mock-access-token";
+    private static final String TEST_ENDPOINT = "https://test-inference.westus.dbinference.azure.com";
+    private static final String INFERENCE_ENDPOINT_PROPERTY = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT";
+
+    // Total calls expected for a fully-retried retryable failure: 1 initial + RETRY_MAX_ATTEMPTS retries
+    private static final int TOTAL_CALLS_AFTER_EXHAUSTION = 1 + InferenceService.RETRY_MAX_ATTEMPTS;
+
+    private TokenCredential mockTokenCredential;
+    private HttpClient mockHttpClient;
+    private MockedStatic<HttpClient> httpClientStaticMock;
+
+    @BeforeMethod
+    public void setUp() {
+        // Set the inference endpoint system property required by InferenceService constructor
+        System.setProperty(INFERENCE_ENDPOINT_PROPERTY, TEST_ENDPOINT);
+
+        mockTokenCredential = mock(TokenCredential.class);
+        mockHttpClient = mock(HttpClient.class);
+
+        // Mock token credential to return a valid token
+        AccessToken accessToken = new AccessToken(MOCK_TOKEN, OffsetDateTime.now().plusHours(1));
+        when(mockTokenCredential.getToken(any(TokenRequestContext.class)))
+            .thenReturn(Mono.just(accessToken));
+
+        // Mock the static HttpClient.createFixed method
+        httpClientStaticMock = Mockito.mockStatic(HttpClient.class);
+        httpClientStaticMock.when(() -> HttpClient.createFixed(any(HttpClientConfig.class)))
+            .thenReturn(mockHttpClient);
+    }
+
+    @AfterMethod
+    public void tearDown() {
+        System.clearProperty(INFERENCE_ENDPOINT_PROPERTY);
+        if (httpClientStaticMock != null) {
+            httpClientStaticMock.close();
+        }
+    }
+
+    // -------------------------------------------------------------------------
+    // Constructor / validation tests (unchanged by retry logic)
+    // -------------------------------------------------------------------------
+
+    @Test(groups = {"unit"})
+    public void constructorShouldThrowWhenTokenCredentialIsNull() {
+        assertThatThrownBy(() -> new InferenceService(null))
+            .isInstanceOf(NullPointerException.class)
+            .hasMessageContaining("Semantic reranking requires AAD authentication")
+            .hasMessageContaining("key-based auth");
+    }
+
+    @Test(groups = {"unit"})
+    public void constructorShouldThrowWhenEndpointNotConfigured() {
+        System.clearProperty(INFERENCE_ENDPOINT_PROPERTY);
+
+        assertThatThrownBy(() -> new InferenceService(mockTokenCredential))
+            .isInstanceOf(IllegalArgumentException.class)
+            .hasMessageContaining("must be set to use semantic reranking");
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldThrowWhenRerankContextIsNull() {
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        assertThatThrownBy(() -> service.semanticRerank(null, Arrays.asList("doc1", "doc2"), null))
+            .isInstanceOf(NullPointerException.class)
+            .hasMessageContaining("Rerank context cannot be null");
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldThrowWhenDocumentsIsNull() {
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        assertThatThrownBy(() -> service.semanticRerank("query", null, null))
+            .isInstanceOf(NullPointerException.class)
+            .hasMessageContaining("Documents list cannot be null");
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldThrowWhenRerankContextIsEmpty() {
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("   ", Arrays.asList("doc1", "doc2"), null))
+            .expectErrorMatches(error -> error instanceof IllegalArgumentException
+                && error.getMessage().contains("Rerank context cannot be empty"))
+            .verify();
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldThrowWhenDocumentsListIsEmpty() {
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("query", Collections.emptyList(), null))
+            .expectErrorMatches(error -> error instanceof IllegalArgumentException
+                && error.getMessage().contains("Documents list cannot be empty"))
+            .verify();
+    }
+
+    // -------------------------------------------------------------------------
+    // Happy-path tests (unchanged by retry logic)
+    // -------------------------------------------------------------------------
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldSucceedWithValidResponse() {
+        // Create a mock response
+        String responseBody = createSuccessResponseBody();
+        HttpResponse mockResponse = createMockResponse(200, responseBody);
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        List<String> documents = Arrays.asList(
+            "This is document 1",
+            "This is document 2",
+            "This is document 3"
+        );
+
+        StepVerifier.create(service.semanticRerank("search query", documents, null))
+            .assertNext(result -> {
+                assertThat(result).isNotNull();
+                assertThat(result.getScores()).isNotNull();
+                assertThat(result.getScores()).hasSize(3);
+                assertThat(result.getScores().get(0).getIndex()).isEqualTo(0);
+                assertThat(result.getScores().get(0).getScore()).isEqualTo(0.95);
+                assertThat(result.getScores().get(0).getDocument()).isEqualTo("This is document 1");
+                assertThat(result.getLatency()).isNotNull();
+                assertThat(result.getLatency().get("inference_time")).isEqualTo(0.5);
+                assertThat(result.getTokenUsage()).isNotNull();
+                assertThat(result.getTokenUsage().get("total_tokens")).isEqualTo(100);
+            })
+            .verifyComplete();
+
+        // Success on first attempt — exactly 1 call
+        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldIncludeOptionsInRequest() {
+        String responseBody = createSuccessResponseBody();
+        HttpResponse mockResponse = createMockResponse(200, responseBody);
+
+        ArgumentCaptor<HttpRequest> requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+        when(mockHttpClient.send(requestCaptor.capture(), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        List<String> documents = Arrays.asList("doc1", "doc2");
+        Map<String, Object> options = new HashMap<>();
+        options.put("return_documents", true);
+        options.put("top_k", 5);
+        options.put("batch_size", 16);
+        options.put("sort", true);
+
+        StepVerifier.create(service.semanticRerank("query", documents, options))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        // Verify request body contains options
+        HttpRequest capturedRequest = requestCaptor.getValue();
+        assertThat(capturedRequest).isNotNull();
+
+        // Collect the Flux<byte[]> body into a single string
+        String requestBody = capturedRequest.body()
+            .map(bytes -> new String(bytes, StandardCharsets.UTF_8))
+            .collectList()
+            .map(list -> String.join("", list))
+            .block();
+
+        assertThat(requestBody).contains("\"return_documents\":true");
+        assertThat(requestBody).contains("\"top_k\":5");
+        assertThat(requestBody).contains("\"batch_size\":16");
+        assertThat(requestBody).contains("\"sort\":true");
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldNotForwardTimeoutSecondsToPayload() {
+        String responseBody = createSuccessResponseBody();
+        HttpResponse mockResponse = createMockResponse(200, responseBody);
+
+        ArgumentCaptor<HttpRequest> requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+        when(mockHttpClient.send(requestCaptor.capture(), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        Map<String, Object> options = new HashMap<>();
+        options.put("top_k", 3);
+        options.put(InferenceService.OPTION_TIMEOUT_SECONDS, 30); // SDK-local, must NOT reach the endpoint
+
+        StepVerifier.create(service.semanticRerank("query", Collections.singletonList("doc1"), options))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        // Verify that the request body does not contain the timeout seconds
+        HttpRequest capturedRequest = requestCaptor.getValue();
+        String requestBody = capturedRequest.body()
+            .map(bytes -> new String(bytes, StandardCharsets.UTF_8))
+            .collectList()
+            .map(list -> String.join("", list))
+            .block();
+
+        assertThat(requestBody).doesNotContain(InferenceService.OPTION_TIMEOUT_SECONDS);
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldIncludeAuthorizationHeader() {
+        String responseBody = createSuccessResponseBody();
+        HttpResponse mockResponse = createMockResponse(200, responseBody);
+
+        ArgumentCaptor<HttpRequest> requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+        when(mockHttpClient.send(requestCaptor.capture(), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        HttpRequest capturedRequest = requestCaptor.getValue();
+        assertThat(capturedRequest.headers().value("Authorization")).isEqualTo("Bearer " + MOCK_TOKEN);
+    }
+
+    // -------------------------------------------------------------------------
+    // Non-retryable error tests — expect exactly 1 HTTP call, immediate failure
+    // -------------------------------------------------------------------------
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldHandle403Forbidden() {
+        // 403 is non-retryable — should fail immediately with 1 call
+        String errorBody = "{\"error\": \"Forbidden - insufficient permissions\"}";
+        HttpResponse mockResponse = createMockResponse(403, errorBody);
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .expectErrorMatches(error -> error.getMessage().contains("Semantic rerank request failed")
+                && error.getMessage().contains("Forbidden"))
+            .verify();
+
+        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldHandleTokenAcquisitionFailure() {
+        // Token errors are not CosmosException — should not be retried
+        when(mockTokenCredential.getToken(any(TokenRequestContext.class)))
+            .thenReturn(Mono.error(new RuntimeException("Token acquisition failed")));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .expectErrorMatches(error -> error instanceof RuntimeException
+                && error.getMessage().contains("Token acquisition failed"))
+            .verify();
+
+        // No HTTP call should be made if token acquisition fails
+        verify(mockHttpClient, never()).send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldNotRetryNetworkErrors() {
+        // Raw RuntimeException (not CosmosException) — not retryable
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.error(new RuntimeException("Network connection failed")));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .expectErrorMatches(error -> error instanceof RuntimeException
+                && error.getMessage().contains("Network connection failed"))
+            .verify();
+
+        // Non-CosmosException errors bypass the retry filter — exactly 1 call
+        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldHandleMalformedResponse() {
+        // Parse failure uses BadRequestException (400) + CUSTOM_SERIALIZER_EXCEPTION sub-status,
+        // matching CosmosItemSerializer's pattern — not retryable, fails immediately with 1 call.
+        String malformedBody = "{ invalid json }";
+        HttpResponse mockResponse = createMockResponse(200, malformedBody);
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .expectErrorMatches(error ->
+                error instanceof BadRequestException
+                    && error.getMessage().contains("Failed to parse semantic rerank response")
+                    && ((BadRequestException) error).getStatusCode() == HttpConstants.StatusCodes.BADREQUEST
+                    && ((BadRequestException) error).getSubStatusCode()
+                        == HttpConstants.SubStatusCodes.CUSTOM_SERIALIZER_EXCEPTION)
+            .verify();
+
+        // Not retryable — exactly 1 call
+        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldHandleEmptyScoresResponse() {
+        String responseBody = "{\"Scores\": [], \"latency\": {}, \"token_usage\": {}}";
+        HttpResponse mockResponse = createMockResponse(200, responseBody);
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .assertNext(result -> {
+                assertThat(result).isNotNull();
+                assertThat(result.getScores()).isEmpty();
+            })
+            .verifyComplete();
+    }
+
+    // -------------------------------------------------------------------------
+    // Retry behaviour tests — use virtual time to skip backoff delays
+    // -------------------------------------------------------------------------
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldRetryOn429AndSucceedOnSecondAttempt() {
+        // First call returns 429, second returns 200
+        HttpResponse throttledResponse = createMockResponse(429, "{\"error\": \"Too Many Requests\"}");
+        HttpResponse successResponse = createMockResponse(200, createSuccessResponseBody());
+
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(throttledResponse))
+            .thenReturn(Mono.just(successResponse));
+
+        // Inject VirtualTimeScheduler so Retry.backoff delays run under virtual time.
+        // The Mono is returned inside the supplier so it is assembled after the
+        // virtual scheduler is installed.
+        final InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        // 1 initial + 1 retry = 2 total calls
+        verify(mockHttpClient, times(2)).send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldRetryOn500AndSucceedOnThirdAttempt() {
+        // First two calls return 500, third returns 200
+        HttpResponse errorResponse = createMockResponse(500, "{\"error\": \"Internal Server Error\"}");
+        HttpResponse successResponse = createMockResponse(200, createSuccessResponseBody());
+
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(errorResponse))
+            .thenReturn(Mono.just(errorResponse))
+            .thenReturn(Mono.just(successResponse));
+
+        final InferenceService service = new InferenceService(mockTokenCredential);
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        // 1 initial + 2 retries = 3 total calls
+        verify(mockHttpClient, times(3)).send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldExhaustRetriesOn429AndFail() {
+        HttpResponse throttledResponse = createMockResponse(429, "{\"error\": \"Too Many Requests\"}");
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(throttledResponse));
+
+        final InferenceService service = new InferenceService(mockTokenCredential);
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
+            .verify();
+
+        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldExhaustRetriesOn500AndFail() {
+        HttpResponse errorResponse = createMockResponse(500, "{\"error\": \"Internal Server Error\"}");
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(errorResponse));
+
+        final InferenceService service = new InferenceService(mockTokenCredential);
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
+            .verify();
+
+        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldExhaustRetriesOn502AndFail() {
+        HttpResponse errorResponse = createMockResponse(502, "{\"error\": \"Bad Gateway\"}");
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(errorResponse));
+
+        final InferenceService service = new InferenceService(mockTokenCredential);
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
+            .verify();
+
+        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldExhaustRetriesOn503AndFail() {
+        HttpResponse errorResponse = createMockResponse(503, "{\"error\": \"Service Unavailable\"}");
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(errorResponse));
+
+        final InferenceService service = new InferenceService(mockTokenCredential);
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
+            .verify();
+
+        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    // -------------------------------------------------------------------------
+    // DataProvider — non-retryable codes fail immediately (1 call),
+    //                retryable codes exhaust retries (TOTAL_CALLS_AFTER_EXHAUSTION calls)
+    // -------------------------------------------------------------------------
+
+    @DataProvider(name = "nonRetryableStatusCodes")
+    public Object[][] nonRetryableStatusCodeProvider() {
+        return new Object[][] {
+            {400, "Bad Request"},
+            {401, "Unauthorized"},
+            {403, "Forbidden"},
+            {404, "Not Found"},
+        };
+    }
+
+    @Test(groups = {"unit"}, dataProvider = "nonRetryableStatusCodes")
+    public void semanticRerankShouldNotRetryNonRetryableStatusCodes(int statusCode, String errorMessage) {
+        String errorBody = String.format("{\"error\": \"%s\"}", errorMessage);
+        HttpResponse mockResponse = createMockResponse(statusCode, errorBody);
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = new InferenceService(mockTokenCredential);
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .expectErrorMatches(error -> containsMessage(error, "Semantic rerank request failed"))
+            .verify();
+
+        // Non-retryable — exactly 1 call regardless of status code
+        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    @DataProvider(name = "retryableStatusCodes")
+    public Object[][] retryableStatusCodeProvider() {
+        return new Object[][] {
+            {429, "Too Many Requests"},
+            {500, "Internal Server Error"},
+            {502, "Bad Gateway"},
+            {503, "Service Unavailable"},
+        };
+    }
+
+    @Test(groups = {"unit"}, dataProvider = "retryableStatusCodes")
+    public void semanticRerankShouldRetryRetryableStatusCodesAndExhaust(int statusCode, String errorMessage) {
+        String errorBody = String.format("{\"error\": \"%s\"}", errorMessage);
+        HttpResponse mockResponse = createMockResponse(statusCode, errorBody);
+        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
+            .thenReturn(Mono.just(mockResponse));
+
+        final InferenceService service = new InferenceService(mockTokenCredential);
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
+            .verify();
+
+        // 1 initial attempt + RETRY_MAX_ATTEMPTS retries
+        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .send(any(HttpRequest.class), any(Duration.class));
+    }
+
+    // -------------------------------------------------------------------------
+    // Helper methods
+    // -------------------------------------------------------------------------
+
+    /**
+     * Checks that an error represents exhausted retries on a retryable inference failure.
+     *
+     * When {@code Retry.backoff} exhausts all attempts it wraps the last error in a
+     * {@code RetryExhaustedException}. {@link com.azure.cosmos.CosmosException#getMessage()}
+     * returns a JSON diagnostic blob, not the plain constructor message, so we use
+     * {@link com.azure.cosmos.CosmosException#getShortMessage()} for plain-string matching.
+     * As a fallback we also walk the cause chain.
+     */
+    private boolean isRetryExhaustedWithMessage(Throwable error, String messageSubstring) {
+        if (!Exceptions.isRetryExhausted(error)) {
+            // Not wrapped — check the error itself
+            return containsMessage(error, messageSubstring);
+        }
+        // Walk the cause chain from the RetryExhaustedException
+        Throwable cause = error.getCause();
+        while (cause != null) {
+            if (containsMessage(cause, messageSubstring)) {
+                return true;
+            }
+            cause = cause.getCause();
+        }
+        return false;
+    }
+
+    private boolean containsMessage(Throwable t, String substring) {
+        if (t instanceof com.azure.cosmos.CosmosException) {
+            // getShortMessage() returns the plain constructor message without JSON wrapping
+            String short_ = ((com.azure.cosmos.CosmosException) t).getShortMessage();
+            if (short_ != null && short_.contains(substring)) {
+                return true;
+            }
+        }
+        return t.getMessage() != null && t.getMessage().contains(substring);
+    }
+
+    private String createSuccessResponseBody() {
+        return "{"
+            + "\"Scores\": ["
+            + "  {\"index\": 0, \"score\": 0.95, \"document\": \"This is document 1\"},"
+            + "  {\"index\": 1, \"score\": 0.85, \"document\": \"This is document 2\"},"
+            + "  {\"index\": 2, \"score\": 0.75, \"document\": \"This is document 3\"}"
+            + "],"
+            + "\"latency\": {"
+            + "  \"data_preprocess_time\": 0.1,"
+            + "  \"inference_time\": 0.5,"
+            + "  \"postprocess_time\": 0.05"
+            + "},"
+            + "\"token_usage\": {"
+            + "  \"total_tokens\": 100"
+            + "}"
+            + "}";
+    }
+
+    private HttpResponse createMockResponse(int statusCode, String body) {
+        HttpResponse mockResponse = mock(HttpResponse.class);
+        when(mockResponse.statusCode()).thenReturn(statusCode);
+        when(mockResponse.bodyAsString()).thenReturn(Mono.just(body));
+        when(mockResponse.headers()).thenReturn(new com.azure.cosmos.implementation.http.HttpHeaders());
+        return mockResponse;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md
index 6eb1b76bad4e..bcf8ab5c3547 100644
--- a/sdk/cosmos/azure-cosmos/CHANGELOG.md
+++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -3,6 +3,7 @@
 ### 4.81.0-beta.1 (Unreleased)
 
 #### Features Added
+* Added beta `semanticRerank` API on `CosmosAsyncContainer` and `CosmosContainer` for semantic reranking of documents using the Azure Cosmos DB inference service. Requires AAD authentication (`TokenCredential`) and the `AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT` system property or environment variable to be set. Includes retry logic for transient failures (429, 500, 502, 503) with exponential backoff. - [PR 48258](https://github.com/Azure/azure-sdk-for-java/pull/48258)
 
 #### Breaking Changes
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
index fe2b36b766e1..cb52abd347ee 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
@@ -22,6 +22,7 @@
 import com.azure.cosmos.implementation.Strings;
 import com.azure.cosmos.implementation.WriteRetryPolicy;
 import com.azure.cosmos.implementation.clienttelemetry.ClientMetricsDiagnosticsHandler;
+import com.azure.cosmos.implementation.inference.InferenceService;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetryMetrics;
 import com.azure.cosmos.implementation.clienttelemetry.CosmosMeterOptions;
@@ -61,6 +62,7 @@
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -108,6 +110,7 @@ private static ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccess
     private final ConsistencyLevel desiredConsistencyLevel;
     private final ReadConsistencyStrategy readConsistencyStrategy;
     private final AzureKeyCredential credential;
+    private final TokenCredential tokenCredential;
     private final CosmosClientTelemetryConfig clientTelemetryConfig;
     private final DiagnosticsProvider diagnosticsProvider;
     private final Tag clientCorrelationTag;
@@ -119,6 +122,7 @@ private static ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccess
     private final List<CosmosOperationPolicy> requestPolicies;
     private final CosmosItemSerializer defaultCustomSerializer;
     private final java.util.function.Function<CosmosAsyncContainer, CosmosAsyncContainer> containerFactory;
+    private final AtomicReference<InferenceService> inferenceService = new AtomicReference<>();
 
     CosmosAsyncClient(CosmosClientBuilder builder) {
         // Async Cosmos client wrapper
@@ -131,7 +135,7 @@ private static ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccess
         List<CosmosPermissionProperties> permissions = builder.getPermissions();
         CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver = builder.getAuthorizationTokenResolver();
         this.credential = builder.getCredential();
-        TokenCredential tokenCredential = builder.getTokenCredential();
+        tokenCredential = builder.getTokenCredential();
         boolean sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled();
         boolean enableTransportClientSharing = builder.isConnectionSharingAcrossClientsEnabled();
         this.proactiveContainerInitConfig = builder.getProactiveContainerInitConfig();
@@ -292,6 +296,40 @@ AzureKeyCredential credential() {
         return credential;
     }
 
+    /**
+     * Gets the token credential.
+     *
+     * @return token credential.
+     */
+    TokenCredential tokenCredential() {
+        return this.tokenCredential;
+    }
+
+    /**
+     * Returns the shared InferenceService instance for this client, creating it lazily on first use.
+     * The instance is tied to this client's lifecycle and will be shut down in {@link #close()}.
+     *
+     * @throws IllegalStateException if the client was built with key-based auth instead of AAD.
+     */
+    InferenceService getOrCreateInferenceService() {
+        if (this.tokenCredential == null) {
+            throw new IllegalStateException(
+                "Semantic reranking requires AAD authentication. "
+                    + "The CosmosClient was built with key-based auth (master key or AzureKeyCredential), "
+                    + "which is not supported for this operation. "
+                    + "Rebuild the client using .credential(TokenCredential) with a DefaultAzureCredential "
+                    + "or another TokenCredential implementation.");
+        }
+        if (this.inferenceService.get() == null) {
+            InferenceService newSvc = new InferenceService(this.tokenCredential);
+            if (!this.inferenceService.compareAndSet(null, newSvc)) {
+                // Another thread already set the instance; close the losing one to prevent resource leak
+                newSvc.close();
+            }
+        }
+        return this.inferenceService.get();
+    }
+
     /***
      * Get the client telemetry config.
      *
@@ -563,6 +601,10 @@ public void close() {
         if (this.clientMetricRegistrySnapshot != null) {
             ClientTelemetryMetrics.remove(this.clientMetricRegistrySnapshot);
         }
+        InferenceService svc = this.inferenceService.get();
+        if (svc != null) {
+            svc.close();
+        }
         asyncDocumentClient.close();
     }
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
index d64c671f3249..585b7da76c41 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
@@ -35,6 +35,7 @@
 import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider;
 import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
 import com.azure.cosmos.implementation.feedranges.FeedRangeInternal;
+import com.azure.cosmos.implementation.inference.InferenceService;
 import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
 import com.azure.cosmos.implementation.routing.Range;
 import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup;
@@ -57,6 +58,7 @@
 import com.azure.cosmos.models.CosmosItemOperation;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
 import com.azure.cosmos.models.CosmosItemResponse;
+import com.azure.cosmos.models.CosmosOperationDetails;
 import com.azure.cosmos.models.CosmosPatchItemRequestOptions;
 import com.azure.cosmos.models.CosmosPatchOperations;
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
@@ -67,7 +69,7 @@
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.models.PartitionKeyDefinition;
-import com.azure.cosmos.models.CosmosOperationDetails;
+import com.azure.cosmos.models.SemanticRerankResult;
 import com.azure.cosmos.models.ShowQueryMode;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.models.ThroughputProperties;
@@ -84,9 +86,11 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -1944,6 +1948,57 @@ public <T> CosmosPagedFlux<T> readAllItems(
         });
     }
 
+    /**
+     * Performs semantic reranking of documents using the Azure Cosmos DB inference service.
+     * <p>
+     * Requires the {@link CosmosAsyncClient} to have been built with AAD authentication
+     * ({@code .credential(TokenCredential)}). Key-based auth is not supported.
+     * <p>
+     * The request timeout defaults to 120 seconds. To override it, pass
+     * {@code "timeout_seconds"} (as a {@link Number}) in the {@code options} map.
+     * You can also apply {@code .timeout(Duration)} to the returned {@code Mono} directly.
+     *
+     * @param rerankContext The query or context string used to score documents.
+     * @param documents The list of document strings to rerank (must be non-null and non-empty).
+     * @param options Optional reranking parameters as a {@code Map<String, Object>}.
+     *                Supported keys:
+     *                <ul>
+     *                  <li>{@code "return_documents"} (Boolean) — include document text in the response</li>
+     *                  <li>{@code "top_k"} (Integer) — maximum number of results to return</li>
+     *                  <li>{@code "batch_size"} (Integer) — documents per inference batch</li>
+     *                  <li>{@code "sort"} (Boolean) — sort results by relevance score</li>
+     *                  <li>{@code "document_type"} (String) — {@code "string"} or {@code "json"}</li>
+     *                  <li>{@code "target_paths"} (String) — JSON paths for extraction when
+     *                      {@code document_type} is {@code "json"}</li>
+     *                  <li>{@code "timeout_seconds"} (Number) — per-request timeout override</li>
+     *                </ul>
+     * @return A {@link Mono} emitting the {@link SemanticRerankResult}.
+     */
+    @Beta(value = Beta.SinceVersion.V4_78_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+    public Mono<SemanticRerankResult> semanticRerank(
+        String rerankContext,
+        List<String> documents,
+        Map<String, Object> options) {
+
+        if (rerankContext == null) {
+            return Mono.error(new IllegalArgumentException("Rerank context cannot be null"));
+        }
+        if (documents == null) {
+            return Mono.error(new IllegalArgumentException("Documents list cannot be null"));
+        }
+
+        if (rerankContext.trim().isEmpty()) {
+            return Mono.error(new IllegalArgumentException("Rerank context cannot be empty"));
+        }
+
+        if (documents.isEmpty()) {
+            return Mono.error(new IllegalArgumentException("Documents list cannot be empty"));
+        }
+
+        return this.database.getClient().getOrCreateInferenceService().semanticRerank(rerankContext, documents, options);
+    }
+
+
     /**
      * Replaces an existing item in a container with a new item.
      * It performs a complete replacement of the item,
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
index d0146a75157b..61f59f10030b 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
@@ -4,6 +4,7 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup;
+import com.azure.cosmos.implementation.inference.InferenceService;
 import com.azure.cosmos.models.CosmosBatch;
 import com.azure.cosmos.models.CosmosBatchOperationResult;
 import com.azure.cosmos.models.CosmosBatchRequestOptions;
@@ -26,6 +27,7 @@
 import com.azure.cosmos.models.FeedRange;
 import com.azure.cosmos.models.FeedResponse;
 import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.models.SemanticRerankResult;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.models.ThroughputProperties;
 import com.azure.cosmos.models.ThroughputResponse;
@@ -38,7 +40,9 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import java.time.Duration;
 import java.util.List;
+import java.util.Map;
 
 import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
 
@@ -312,6 +316,34 @@ private CosmosItemResponse<Object> blockDeleteItemResponse(Mono<CosmosItemRespon
         }
     }
 
+    private SemanticRerankResult blockSemanticRerankResponse(
+        Mono<SemanticRerankResult> semanticRerankResultMono,
+        Map<String, Object> options) {
+        // Derive the blocking timeout from the same option key used by InferenceService,
+        // so the sync API respects timeout_seconds just like the async path.
+        // Default is InferenceService.DEFAULT_REQUEST_TIMEOUT (120 s) if not supplied.
+        Duration blockTimeout = InferenceService.DEFAULT_REQUEST_TIMEOUT;
+        if (options != null) {
+            Object timeoutVal = options.get(InferenceService.OPTION_TIMEOUT_SECONDS);
+            if (timeoutVal instanceof Number) {
+                double seconds = ((Number) timeoutVal).doubleValue();
+                if (seconds > 0) {
+                    blockTimeout = Duration.ofMillis((long) (seconds * 1000));
+                }
+            }
+        }
+        try {
+            return semanticRerankResultMono.block(blockTimeout);
+        } catch (Exception ex) {
+            final Throwable throwable = Exceptions.unwrap(ex);
+            if (throwable instanceof CosmosException) {
+                throw (CosmosException) throwable;
+            } else {
+                throw ex;
+            }
+        }
+    }
+
     private CosmosBatchResponse blockBatchResponse(Mono<CosmosBatchResponse> batchResponseMono) {
         try {
             return batchResponseMono.block();
@@ -1037,6 +1069,31 @@ public <TContext> Iterable<CosmosBulkOperationResponse<TContext>> executeBulkOpe
         return this.blockBulkResponse(asyncContainer.executeBulkOperations(Flux.fromIterable(operations), bulkOptions));
     }
 
+    /**
+     * Performs semantic reranking of documents using the inference service.
+     *
+     * <p><strong>Timeout:</strong> This method blocks with a default timeout of 120 seconds.
+     * To override, pass {@code "timeout_seconds"} (as a {@link Number}) in the {@code options} map,
+     * for example: {@code options.put("timeout_seconds", 30)}.
+     *
+     * @param rerankContext The query or context string used to score documents.
+     * @param documents The list of document strings to rerank.
+     * @param options Optional reranking parameters as a map. Supported keys include:
+     *                {@code return_documents} (Boolean) - whether to return document text in the response,
+     *                {@code top_k} (Integer) - maximum number of documents to return,
+     *                {@code batch_size} (Integer) - number of documents per batch,
+     *                {@code sort} (Boolean) - whether to sort results by relevance score,
+     *                {@code timeout_seconds} (Number) - per-request timeout in seconds (default: 120).
+     * @return The semantic rerank result.
+     */
+    @Beta(value = Beta.SinceVersion.V4_78_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+    public SemanticRerankResult semanticRerank(
+        String rerankContext,
+        List<String> documents,
+        Map<String, Object> options) {
+        return blockSemanticRerankResponse(this.asyncContainer.semanticRerank(rerankContext, documents, options), options);
+    }
+
     /**
      * Gets the Cosmos scripts using the current container as context.
      *
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
index 18eef0544e18..477eb483e87a 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
@@ -421,6 +421,11 @@ public class Configs {
     private static final String CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT = "COSMOS.CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT";
     private static final int DEFAULT_CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT = 120;
 
+    // Inference service related configs
+    // Single source of truth: same name is used for both the system property and the environment variable lookup.
+    private static final String INFERENCE_ENDPOINT_PROPERTY = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT";
+    private static final String INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE = INFERENCE_ENDPOINT_PROPERTY;
+
     private static final Object lockObject = new Object();
     private static Boolean cachedIsHostnameValidationDisabled = null;
 
@@ -531,6 +536,19 @@ public URI getThinclientEndpoint() {
         return URI.create(DEFAULT_THINCLIENT_ENDPOINT);
     }
 
+    public URI getInferenceServiceEndpoint() {
+        String valueFromSystemProperty = System.getProperty(INFERENCE_ENDPOINT_PROPERTY);
+        if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) {
+            return URI.create(valueFromSystemProperty);
+        }
+
+        String valueFromEnvVariable = System.getenv(INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE);
+        if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) {
+            return URI.create(valueFromEnvVariable);
+        }
+        return null;
+    }
+
     public static boolean isThinClientEnabled() {
         String valueFromSystemProperty = System.getProperty(THINCLIENT_ENABLED);
         if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) {
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
new file mode 100644
index 000000000000..7474d62017b8
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
@@ -0,0 +1,513 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.implementation.inference;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.SimpleTokenCache;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.cosmos.BridgeInternal;
+import com.azure.cosmos.CosmosException;
+import com.azure.cosmos.implementation.BadRequestException;
+import com.azure.cosmos.implementation.Configs;
+import com.azure.cosmos.implementation.HttpConstants;
+import com.azure.cosmos.implementation.LifeCycleUtils;
+import com.azure.cosmos.implementation.Utils;
+import com.azure.cosmos.implementation.http.HttpClient;
+import com.azure.cosmos.implementation.http.HttpClientConfig;
+import com.azure.cosmos.implementation.http.HttpHeaders;
+import com.azure.cosmos.implementation.http.HttpRequest;
+import com.azure.cosmos.implementation.http.HttpResponse;
+import com.azure.cosmos.models.ModelBridgeInternal;
+import com.azure.cosmos.models.SemanticRerankResult;
+import com.azure.cosmos.models.SemanticRerankScore;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import io.netty.handler.codec.http.HttpMethod;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Scheduler;
+import reactor.core.scheduler.Schedulers;
+import reactor.util.retry.Retry;
+
+import java.io.IOException;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
+
+/**
+ * Internal service for semantic reranking operations.
+ * <p>
+ * While this class is public, it is not part of our published public APIs.
+ * This is meant to be internally used only by our SDK.
+ */
+public class InferenceService implements AutoCloseable {
+    private static final Logger logger = LoggerFactory.getLogger(InferenceService.class);
+    private static final String INFERENCE_SCOPE = "https://dbinference.azure.com/.default";
+    private static final String BASE_PATH = "/inference/semanticReranking";
+    private static final String INFERENCE_USER_AGENT =
+        "cosmos-inference-java/" + HttpConstants.Versions.getSdkVersion();
+    private static final ObjectMapper OBJECT_MAPPER = Utils.getSimpleObjectMapper();
+
+    // System property / environment variable name (single string used for both lookup mechanisms)
+    private static final String MAX_CONNECTION_LIMIT_PROPERTY = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_SERVICE_MAX_CONNECTION_LIMIT";
+    private static final String MAX_CONNECTION_LIMIT_VARIABLE = MAX_CONNECTION_LIMIT_PROPERTY;
+
+    // Option keys that callers can pass in the options map
+    public static final String OPTION_TIMEOUT_SECONDS = "timeout_seconds";
+    public static final String OPTION_RETURN_DOCUMENTS = "return_documents";
+    public static final String OPTION_TOP_K = "top_k";
+    public static final String OPTION_BATCH_SIZE = "batch_size";
+    public static final String OPTION_SORT = "sort";
+    public static final String OPTION_DOCUMENT_TYPE = "document_type";
+    public static final String OPTION_TARGET_PATHS = "target_paths";
+
+    /**
+     * Default network-level (connection + read) timeout for inference service HTTP calls.
+     * Set to match {@link #DEFAULT_REQUEST_TIMEOUT} because inference calls can take tens of
+     * seconds for large document sets. If this were shorter than the operation timeout,
+     * the socket would time out before the 120-second operation deadline could fire.
+     */
+    public static final Duration DEFAULT_NETWORK_REQUEST_TIMEOUT = Duration.ofSeconds(120);
+    public static final Duration DEFAULT_IDLE_CONNECTION_TIMEOUT = Duration.ofSeconds(60);
+    public static final Duration DEFAULT_CONNECTION_ACQUIRE_TIMEOUT = Duration.ofSeconds(5);
+    /**
+     * Default per-request timeout for semantic rerank calls (120 seconds).
+     * Callers using the async API can override downstream with {@code .timeout(Duration)}.
+     * Callers using the sync API can override by passing {@value #OPTION_TIMEOUT_SECONDS} in the options map.
+     */
+    public static final Duration DEFAULT_REQUEST_TIMEOUT = Duration.ofSeconds(120);
+    /**
+     * Default maximum number of pooled connections to the inference endpoint.
+     * Matches .NET's {@code inferenceServiceDefaultMaxConnectionLimit = 50}.
+     * Override via system property or environment variable
+     * {@value #MAX_CONNECTION_LIMIT_PROPERTY}.
+     */
+    public static final int DEFAULT_MAX_CONNECTION_POOL_SIZE = 50;
+
+    // Retry policy — matches Python: TOTAL_RETRIES=3, RETRY_BACKOFF_FACTOR=0.8s, RETRY_BACKOFF_MAX=120s
+    static final int RETRY_MAX_ATTEMPTS = 3;
+    static final Duration RETRY_INITIAL_BACKOFF = Duration.ofMillis(800);
+    static final Duration RETRY_MAX_BACKOFF = Duration.ofSeconds(120);
+    // Status codes that are safe to retry (transient failures and rate limiting)
+    private static final Set<Integer> RETRYABLE_STATUS_CODES = Collections.unmodifiableSet(
+        new HashSet<>(Arrays.asList(
+            HttpConstants.StatusCodes.TOO_MANY_REQUESTS,      // 429 — rate limited
+            HttpConstants.StatusCodes.INTERNAL_SERVER_ERROR,  // 500 — transient server error
+            502,                                               // 502 — bad gateway (no constant in HttpConstants)
+            HttpConstants.StatusCodes.SERVICE_UNAVAILABLE     // 503 — transient unavailability
+        ))
+    );
+
+    private final URI inferenceEndpoint;
+    private final HttpClient httpClient;
+    private SimpleTokenCache tokenCache = null;
+    // Package-private so tests can inject a VirtualTimeScheduler to make backoff delays deterministic
+    Scheduler retryScheduler = Schedulers.parallel();
+    // Package-private so tests can disable jitter for deterministic virtual-time delays
+    double retryJitter = 0.5;
+
+    /**
+     * Creates a new InferenceService instance.
+     *
+     * @param tokenCredential The Azure AD token credential.
+     * @throws IllegalArgumentException if inference endpoint is not configured or token credential is null.
+     */
+    public InferenceService(TokenCredential tokenCredential) {
+        checkNotNull(tokenCredential,
+            "Semantic reranking requires AAD authentication. "
+                + "Rebuild the CosmosClient using .credential(TokenCredential) — "
+                + "key-based auth (master key or AzureKeyCredential) is not supported for this operation.");
+
+        Configs configs = new Configs();
+        URI inferenceBaseUrl = configs.getInferenceServiceEndpoint();
+
+        if (inferenceBaseUrl == null || inferenceBaseUrl.toString().trim().isEmpty()) {
+            throw new IllegalArgumentException("Inference endpoint property must be set to use semantic reranking");
+        }
+
+        this.inferenceEndpoint = URI.create(inferenceBaseUrl + BASE_PATH);
+        HttpClientConfig httpClientConfig = new HttpClientConfig(configs)
+            .withNetworkRequestTimeout(DEFAULT_NETWORK_REQUEST_TIMEOUT)
+            .withConnectionAcquireTimeout(DEFAULT_CONNECTION_ACQUIRE_TIMEOUT)
+            .withMaxIdleConnectionTimeout(DEFAULT_IDLE_CONNECTION_TIMEOUT)
+            .withPoolSize(resolveMaxConnectionPoolSize());
+        this.httpClient = HttpClient.createFixed(httpClientConfig);
+
+        // Create token cache for inference service scope
+        this.tokenCache = new SimpleTokenCache(() -> {
+            TokenRequestContext context = new TokenRequestContext().addScopes(INFERENCE_SCOPE);
+            return tokenCredential.getToken(context)
+                .doOnNext(token -> {
+                    if (logger.isDebugEnabled()) {
+                        logger.debug("Acquired AAD token for inference service scope: {}", INFERENCE_SCOPE);
+                    }
+                });
+        });
+
+        if (logger.isInfoEnabled()) {
+            logger.info("InferenceService initialized with endpoint: {}", inferenceEndpoint);
+        }
+    }
+
+    /**
+     * Closes the InferenceService and releases the underlying HTTP client connection pool.
+     */
+    @Override
+    public void close() {
+        logger.info("Shutting down InferenceService httpClient...");
+        LifeCycleUtils.closeQuietly(this.httpClient);
+    }
+
+    /**
+     * Resolves the maximum connection pool size from system property, environment variable,
+     * or the default ({@value #DEFAULT_MAX_CONNECTION_POOL_SIZE}).
+     * Matches .NET's AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_SERVICE_MAX_CONNECTION_LIMIT override.
+     */
+    private static int resolveMaxConnectionPoolSize() {
+        String fromProperty = System.getProperty(MAX_CONNECTION_LIMIT_PROPERTY);
+        if (fromProperty != null && !fromProperty.isEmpty()) {
+            try {
+                return Integer.parseInt(fromProperty);
+            } catch (NumberFormatException e) {
+                logger.warn("Invalid value for system property {}: '{}'. Using default {}.",
+                    MAX_CONNECTION_LIMIT_PROPERTY, fromProperty, DEFAULT_MAX_CONNECTION_POOL_SIZE);
+            }
+        }
+
+        String fromEnv = System.getenv(MAX_CONNECTION_LIMIT_VARIABLE);
+        if (fromEnv != null && !fromEnv.isEmpty()) {
+            try {
+                return Integer.parseInt(fromEnv);
+            } catch (NumberFormatException e) {
+                logger.warn("Invalid value for environment variable {}: '{}'. Using default {}.",
+                    MAX_CONNECTION_LIMIT_VARIABLE, fromEnv, DEFAULT_MAX_CONNECTION_POOL_SIZE);
+            }
+        }
+
+        return DEFAULT_MAX_CONNECTION_POOL_SIZE;
+    }
+
+    /**
+     * Performs semantic reranking of documents.
+     *
+     * <p>The request timeout defaults to 120 seconds ({@link #DEFAULT_REQUEST_TIMEOUT}).
+     * To override it, pass {@code "timeout_seconds"} (as a {@link Number}) in the {@code options} map.
+     * Callers using the async {@link reactor.core.publisher.Mono} result can also apply
+     * {@code .timeout(Duration)} downstream without needing to set the option.
+     *
+     * @param rerankContext The query or context string used to score documents.
+     * @param documents The list of document strings to rerank.
+     * @param options Optional reranking parameters. SDK-local keys ({@link #OPTION_TIMEOUT_SECONDS})
+     *                are consumed locally and not forwarded to the inference endpoint.
+     * @return A Mono emitting the semantic rerank result.
+     */
+    public Mono<SemanticRerankResult> semanticRerank(
+        String rerankContext,
+        List<String> documents,
+        Map<String, Object> options) {
+
+        // Resolve the per-request timeout: options > default
+        final Duration requestTimeout = resolveRequestTimeout(options);
+
+        return this.tokenCache.getToken()
+            .flatMap(accessToken -> {
+                try {
+                    // Build request payload
+                    ObjectNode payload = OBJECT_MAPPER.createObjectNode();
+                    payload.put("query", rerankContext);
+
+                    ArrayNode documentsArray = payload.putArray("documents");
+                    for (String doc : documents) {
+                        documentsArray.add(doc);
+                    }
+
+                    if (options != null) {
+                        options.forEach((key, value) -> {
+                            // timeout_seconds is a SDK-local option — do not forward to the endpoint
+                            if (value != null && !OPTION_TIMEOUT_SECONDS.equals(key)) {
+                                payload.set(key, OBJECT_MAPPER.valueToTree(value));
+                            }
+                        });
+                    }
+
+                    String requestBody = OBJECT_MAPPER.writeValueAsString(payload);
+
+                    HttpRequest httpRequest = getHttpRequest(accessToken);
+                    httpRequest.withBody(requestBody.getBytes(StandardCharsets.UTF_8));
+
+                    if (logger.isDebugEnabled()) {
+                        logger.debug("Sending semantic rerank request to: {} (timeout: {})", inferenceEndpoint, requestTimeout);
+                    }
+
+                    return httpClient.send(httpRequest, requestTimeout)
+                        .flatMap(response -> parseResponse(response));
+
+                } catch (IOException e) {
+                    logger.error("Failed to serialize request payload", e);
+                    return Mono.error(BridgeInternal.createCosmosException(500, "Failed to serialize semantic rerank request"));
+                }
+            })
+            .as(this::withRetry)
+            .doOnError(error -> logger.error("Semantic rerank operation failed", error));
+    }
+
+    /**
+     * Resolves the effective request timeout.
+     * If the caller supplied {@value #OPTION_TIMEOUT_SECONDS} in the options map, that value is used;
+     * otherwise {@link #DEFAULT_REQUEST_TIMEOUT} applies.
+     */
+    private static Duration resolveRequestTimeout(Map<String, Object> options) {
+        if (options != null) {
+            Object value = options.get(OPTION_TIMEOUT_SECONDS);
+            if (value instanceof Number) {
+                double seconds = ((Number) value).doubleValue();
+                if (seconds > 0) {
+                    return Duration.ofMillis((long) (seconds * 1000));
+                } else {
+                    logger.warn("Invalid '{}' value: {}. Must be > 0. Using default timeout {}.",
+                        OPTION_TIMEOUT_SECONDS, seconds, DEFAULT_REQUEST_TIMEOUT);
+                }
+            }
+        }
+        return DEFAULT_REQUEST_TIMEOUT;
+    }
+
+    /**
+     * Wraps a {@link Mono} with retry logic for transient inference endpoint failures.
+     *
+     * <p>Retries up to {@value #RETRY_MAX_ATTEMPTS} times on retryable {@link com.azure.cosmos.CosmosException}
+     * status codes (429, 500, 502, 503). For HTTP 429 responses, the server-supplied
+     * {@code Retry-After} header is honoured (delay-seconds form per RFC 7231 § 7.1.3) and capped at
+     * {@link #RETRY_MAX_BACKOFF}. If the header is missing or in HTTP-date form, exponential
+     * backoff with jitter is used instead, starting at {@link #RETRY_INITIAL_BACKOFF} and capped at
+     * {@link #RETRY_MAX_BACKOFF}. All other exceptions (non-retryable status codes, serialisation
+     * errors, etc.) propagate immediately.
+     *
+     * <p>Matches the Python SDK's retry policy:
+     * {@code TOTAL_RETRIES=3, RETRY_BACKOFF_FACTOR=0.8, RETRY_BACKOFF_MAX=120s, honour Retry-After on 429}.
+     */
+    private <T> Mono<T> withRetry(Mono<T> source) {
+        return source.retryWhen(
+            Retry.from(retrySignals -> retrySignals.concatMap(signal -> {
+                Throwable error = signal.failure();
+                if (!(error instanceof CosmosException)) {
+                    return Mono.error(error);
+                }
+                CosmosException cex = (CosmosException) error;
+                int statusCode = cex.getStatusCode();
+                if (!RETRYABLE_STATUS_CODES.contains(statusCode)) {
+                    return Mono.error(error);
+                }
+                long attempt = signal.totalRetries();
+                if (attempt >= RETRY_MAX_ATTEMPTS) {
+                    return Mono.error(error);
+                }
+
+                Duration delay = computeRetryDelay(cex, statusCode, attempt);
+                logger.warn(
+                    "Semantic rerank transient failure (status={}, attempt={}/{}), retrying after {} ms...",
+                    statusCode, attempt + 1, RETRY_MAX_ATTEMPTS, delay.toMillis());
+                return Mono.delay(delay, retryScheduler);
+            }))
+        );
+    }
+
+    /**
+     * Computes the delay before the next retry attempt.
+     * Prefers the {@code Retry-After} response header on 429 (delay-seconds form, capped at
+     * {@link #RETRY_MAX_BACKOFF}). Otherwise falls back to exponential backoff with jitter.
+     */
+    private Duration computeRetryDelay(CosmosException cex, int statusCode, long attempt) {
+        if (statusCode == HttpConstants.StatusCodes.TOO_MANY_REQUESTS) {
+            Duration serverDelay = parseRetryAfterSeconds(cex.getResponseHeaders());
+            if (serverDelay != null) {
+                return serverDelay.compareTo(RETRY_MAX_BACKOFF) > 0 ? RETRY_MAX_BACKOFF : serverDelay;
+            }
+        }
+        return exponentialBackoffWithJitter(attempt);
+    }
+
+    /**
+     * Parses the RFC 7231 {@code Retry-After} header as delay-seconds.
+     * Returns {@code null} if the header is absent, blank, negative, or in HTTP-date form
+     * (the latter is not supported here — exponential backoff is used in that case).
+     * Lookup is case-insensitive per RFC 7230 § 3.2.
+     */
+    private static Duration parseRetryAfterSeconds(Map<String, String> headers) {
+        if (headers == null || headers.isEmpty()) {
+            return null;
+        }
+        String value = null;
+        for (Map.Entry<String, String> entry : headers.entrySet()) {
+            if (entry.getKey() != null
+                && HttpConstants.HttpHeaders.RETRY_AFTER.equalsIgnoreCase(entry.getKey())) {
+                value = entry.getValue();
+                break;
+            }
+        }
+        if (value == null || value.trim().isEmpty()) {
+            return null;
+        }
+        try {
+            long seconds = Long.parseLong(value.trim());
+            if (seconds < 0) {
+                return null;
+            }
+            return Duration.ofSeconds(seconds);
+        } catch (NumberFormatException ignored) {
+            // HTTP-date form (e.g. "Fri, 31 Dec 1999 23:59:59 GMT") is not parsed.
+            return null;
+        }
+    }
+
+    /**
+     * Exponential backoff with jitter, matching the Python SDK semantics.
+     * {@code RETRY_INITIAL_BACKOFF * 2^attempt}, capped at {@link #RETRY_MAX_BACKOFF},
+     * with +/- {@code retryJitter * base} random jitter applied.
+     */
+    private Duration exponentialBackoffWithJitter(long attempt) {
+        long baseMillis = RETRY_INITIAL_BACKOFF.toMillis() << Math.min(attempt, 30);
+        long cappedMillis = Math.min(baseMillis, RETRY_MAX_BACKOFF.toMillis());
+        long jitterRange = (long) (cappedMillis * retryJitter);
+        long jitter = jitterRange > 0
+            ? ThreadLocalRandom.current().nextLong(-jitterRange, jitterRange + 1)
+            : 0L;
+        return Duration.ofMillis(Math.max(0L, cappedMillis + jitter));
+    }
+
+    private HttpRequest getHttpRequest(AccessToken accessToken) {
+        HttpHeaders headers = new HttpHeaders();
+        headers.set(HttpConstants.HttpHeaders.CONTENT_TYPE, "application/json");
+        headers.set(HttpConstants.HttpHeaders.AUTHORIZATION, "Bearer " + accessToken.getToken());
+        headers.set(HttpConstants.HttpHeaders.USER_AGENT, INFERENCE_USER_AGENT);
+
+        return new HttpRequest(
+            HttpMethod.POST,
+            inferenceEndpoint,
+            inferenceEndpoint.getPort(),
+            headers);
+    }
+
+    /**
+     * Parses the HTTP response into a SemanticRerankResult.
+     *
+     * @param response The HTTP response.
+     * @return A Mono emitting the parsed result.
+     */
+    private Mono<SemanticRerankResult> parseResponse(HttpResponse response) {
+        int statusCode = response.statusCode();
+
+        return response.bodyAsString()
+            .flatMap(bodyString -> {
+                if (statusCode >= 400) {
+                    logger.error("Semantic rerank request failed with status {}: {}", statusCode, bodyString);
+                    Map<String, String> headersMap = convertHeaders(response.headers());
+                    return Mono.error(BridgeInternal.createCosmosException(
+                        String.format("Semantic rerank request failed: %s", bodyString),
+                        null,
+                        headersMap,
+                        statusCode,
+                        null));
+                }
+
+                try {
+                    JsonNode rootNode = OBJECT_MAPPER.readTree(bodyString);
+                    SemanticRerankResult result = new SemanticRerankResult();
+
+                    // Parse scores
+                    if (rootNode.has("Scores")) {
+                        JsonNode scoresNode = rootNode.get("Scores");
+                        List<SemanticRerankScore> scores = new ArrayList<>();
+
+                        if (scoresNode.isArray()) {
+                            for (JsonNode scoreNode : scoresNode) {
+                                SemanticRerankScore score = new SemanticRerankScore();
+                                JsonNode indexNode = scoreNode.get("index");
+                                JsonNode scoreValNode = scoreNode.get("score");
+                                if (indexNode != null) {
+                                    ModelBridgeInternal.setSemanticRerankScoreIndex(score, indexNode.asInt());
+                                }
+                                if (scoreValNode != null) {
+                                    ModelBridgeInternal.setSemanticRerankScoreScore(score, scoreValNode.asDouble());
+                                }
+
+                                if (scoreNode.has("document")) {
+                                    ModelBridgeInternal.setSemanticRerankScoreDocument(score, scoreNode.get("document").asText());
+                                }
+
+                                scores.add(score);
+                            }
+                        }
+                        ModelBridgeInternal.setSemanticRerankResultScores(result, scores);
+                    }
+
+                    // Parse latency
+                    if (rootNode.has("latency")) {
+                        Map<String, Object> latency = new HashMap<>();
+                        rootNode.get("latency").fields().forEachRemaining(
+                            entry -> latency.put(entry.getKey(), entry.getValue().asDouble()));
+                        ModelBridgeInternal.setSemanticRerankResultLatency(result, latency);
+                    }
+
+                    // Parse token usage
+                    if (rootNode.has("token_usage")) {
+                        Map<String, Object> tokenUsage = new HashMap<>();
+                        rootNode.get("token_usage").fields().forEachRemaining(
+                            entry -> tokenUsage.put(entry.getKey(), entry.getValue().asInt()));
+                        ModelBridgeInternal.setSemanticRerankResultTokenUsage(result, tokenUsage);
+                    }
+
+                    if (logger.isDebugEnabled()) {
+                        logger.debug("Successfully parsed semantic rerank response with {} scores",
+                            result.getScores() != null ? result.getScores().size() : 0);
+                    }
+
+                    return Mono.just(result);
+
+                } catch (Exception e) {
+                    logger.error("Failed to parse semantic rerank response", e);
+                    // Use BadRequestException (400) + CUSTOM_SERIALIZER_EXCEPTION sub-status —
+                    // the same pattern CosmosItemSerializer uses for client-side deserialization
+                    // failures. 400 is not in RETRYABLE_STATUS_CODES, so no retry will be attempted.
+                    BadRequestException parseException = new BadRequestException(
+                        "Failed to parse semantic rerank response: " + e.getMessage(), e);
+                    BridgeInternal.setSubStatusCode(parseException,
+                        HttpConstants.SubStatusCodes.CUSTOM_SERIALIZER_EXCEPTION);
+                    return Mono.error(parseException);
+                }
+            });
+    }
+
+    /**
+     * Converts HttpHeaders to a Map.
+     *
+     * @param headers The HTTP headers.
+     * @return A map of header names to values.
+     */
+    private Map<String, String> convertHeaders(HttpHeaders headers) {
+        Map<String, String> headersMap = new HashMap<>();
+        if (headers != null) {
+            headers.toMap().forEach((key, value) -> {
+                if (value != null) {
+                    headersMap.put(key, value);
+                }
+            });
+        }
+        return headersMap;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
index 7dde81e06e98..e26b48115787 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
@@ -744,6 +744,39 @@ public static List<PatchOperation> getPatchOperationsFromCosmosPatch(CosmosPatch
         return cosmosPatchOperations.getPatchOperations();
     }
 
+    // Bridge methods for SemanticRerankResult and SemanticRerankScore — setters are package-private
+    // to prevent mutation by callers; InferenceService uses these bridge methods to populate instances.
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static void setSemanticRerankResultScores(SemanticRerankResult result, java.util.List<SemanticRerankScore> scores) {
+        result.setScores(scores);
+    }
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static void setSemanticRerankResultLatency(SemanticRerankResult result, Map<String, Object> latency) {
+        result.setLatency(latency);
+    }
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static void setSemanticRerankResultTokenUsage(SemanticRerankResult result, Map<String, Object> tokenUsage) {
+        result.setTokenUsage(tokenUsage);
+    }
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static void setSemanticRerankScoreIndex(SemanticRerankScore score, int index) {
+        score.setIndex(index);
+    }
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static void setSemanticRerankScoreScore(SemanticRerankScore score, double scoreValue) {
+        score.setScore(scoreValue);
+    }
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static void setSemanticRerankScoreDocument(SemanticRerankScore score, String document) {
+        score.setDocument(document);
+    }
+
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static void  initializeAllAccessors() {
         CosmosBatch.initialize();
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
new file mode 100644
index 000000000000..7f40dc5fb646
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
@@ -0,0 +1,78 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.models;
+
+import com.azure.cosmos.util.Beta;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Represents the result of a semantic rerank operation.
+ */
+@Beta(value = Beta.SinceVersion.V4_78_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+public final class SemanticRerankResult {
+    private List<SemanticRerankScore> scores;
+    private Map<String, Object> latency;
+    private Map<String, Object> tokenUsage;
+
+    /**
+     * Creates a new instance of SemanticRerankResult.
+     */
+    public SemanticRerankResult() {
+    }
+
+    /**
+     * Gets the list of scored documents.
+     *
+     * @return the list of scored documents.
+     */
+    public List<SemanticRerankScore> getScores() {
+        return scores;
+    }
+
+    /**
+     * Sets the list of scored documents.
+     *
+     * @param scores the list of scored documents.
+     */
+    void setScores(List<SemanticRerankScore> scores) {
+        this.scores = scores;
+    }
+
+    /**
+     * Gets the latency information for the operation as a map of metric names to values.
+     *
+     * @return the latency information map.
+     */
+    public Map<String, Object> getLatency() {
+        return latency;
+    }
+
+    /**
+     * Sets the latency information for the operation.
+     *
+     * @param latency the latency information map.
+     */
+    void setLatency(Map<String, Object> latency) {
+        this.latency = latency;
+    }
+
+    /**
+     * Gets the token usage information for the operation as a map of metric names to values.
+     *
+     * @return the token usage information map.
+     */
+    public Map<String, Object> getTokenUsage() {
+        return tokenUsage;
+    }
+
+    /**
+     * Sets the token usage information for the operation.
+     *
+     * @param tokenUsage the token usage information map.
+     */
+    void setTokenUsage(Map<String, Object> tokenUsage) {
+        this.tokenUsage = tokenUsage;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java
new file mode 100644
index 000000000000..625b403b0ce3
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java
@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.models;
+
+import com.azure.cosmos.util.Beta;
+
+/**
+ * Represents a single scored document in the semantic rerank result.
+ */
+@Beta(value = Beta.SinceVersion.V4_78_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+public final class SemanticRerankScore {
+    private int index;
+    private String document;
+    private double score;
+
+    /**
+     * Creates a new instance of SemanticRerankScore.
+     */
+    public SemanticRerankScore() {
+    }
+
+    /**
+     * Gets the index of the document in the original input list.
+     *
+     * @return the document index.
+     */
+    public int getIndex() {
+        return index;
+    }
+
+    /**
+     * Sets the index of the document in the original input list.
+     *
+     * @param index the document index.
+     */
+    void setIndex(int index) {
+        this.index = index;
+    }
+
+    /**
+     * Gets the document text (if returnDocuments was true in the request).
+     *
+     * @return the document text, or null if not included.
+     */
+    public String getDocument() {
+        return document;
+    }
+
+    /**
+     * Sets the document text.
+     *
+     * @param document the document text.
+     */
+    void setDocument(String document) {
+        this.document = document;
+    }
+
+    /**
+     * Gets the semantic relevance score for this document.
+     *
+     * @return the relevance score.
+     */
+    public double getScore() {
+        return score;
+    }
+
+    /**
+     * Sets the semantic relevance score for this document.
+     *
+     * @param score the relevance score.
+     */
+    void setScore(double score) {
+        this.score = score;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java
index 53072b09326d..ca01b8cc9f65 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java
@@ -109,6 +109,8 @@ public enum SinceVersion {
         /** v4.74.0 */
         V4_74_0,
         /** v4.75.0 */
-        V4_75_0
+        V4_75_0,
+        /** v4.78.0 */
+        V4_78_0
     }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/semanticrerank/SemanticRerankSample.java b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/semanticrerank/SemanticRerankSample.java
new file mode 100644
index 000000000000..ea47468f1cfe
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/semanticrerank/SemanticRerankSample.java
@@ -0,0 +1,130 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.semanticrerank;
+
+import com.azure.core.credential.TokenCredential;
+import com.azure.cosmos.CosmosAsyncClient;
+import com.azure.cosmos.CosmosAsyncContainer;
+import com.azure.cosmos.CosmosClientBuilder;
+import com.azure.cosmos.models.SemanticRerankResult;
+import com.azure.cosmos.models.SemanticRerankScore;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Sample demonstrating semantic rerank functionality in Azure Cosmos DB.
+ * <p>
+ * Prerequisites:
+ * 1. Set environment variable {@code COSMOS_ENDPOINT} with your Cosmos DB account endpoint.
+ * 2. Set environment variable {@code AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT} with
+ *    the inference service endpoint.
+ * 3. Use Azure AD authentication — provide a {@link TokenCredential} implementation.
+ *    If you have {@code azure-identity} on your classpath, use:
+ *    {@code new com.azure.identity.DefaultAzureCredentialBuilder().build()}
+ * </p>
+ */
+public class SemanticRerankSample {
+
+    public static void main(String[] args) {
+        String endpoint = System.getenv("COSMOS_ENDPOINT");
+        if (endpoint == null || endpoint.isEmpty()) {
+            System.err.println("Please set COSMOS_ENDPOINT environment variable");
+            return;
+        }
+
+        String inferenceEndpoint = System.getenv("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT");
+        if (inferenceEndpoint == null || inferenceEndpoint.isEmpty()) {
+            System.err.println("Please set AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT environment variable");
+            return;
+        }
+
+        // Provide an Azure AD TokenCredential. If azure-identity is on the classpath:
+        //   TokenCredential credential = new com.azure.identity.DefaultAzureCredentialBuilder().build();
+        // Key-based authentication is not supported for semantic reranking.
+        throw new UnsupportedOperationException(
+            "Supply a TokenCredential (e.g. DefaultAzureCredentialBuilder) and remove this line to run the sample.");
+    }
+
+    /**
+     * Runs the semantic rerank demo using the provided client and a container in the given database.
+     *
+     * @param client       Cosmos async client built with AAD credentials.
+     * @param databaseName Name of the database.
+     * @param containerName Name of the container.
+     */
+    static void runSemanticRerankDemo(CosmosAsyncClient client, String databaseName, String containerName) {
+        System.out.println("Semantic Rerank Sample");
+        System.out.println("======================\n");
+
+        CosmosAsyncContainer container = client.getDatabase(databaseName).getContainer(containerName);
+
+        List<String> documents = Arrays.asList(
+            "Berlin is the capital of Germany.",
+            "Paris is the capital of France.",
+            "Madrid is the capital of Spain.",
+            "Rome is the capital of Italy.",
+            "London is the capital of England."
+        );
+
+        String rerankContext = "What is the capital of France?";
+
+        System.out.println("Query: " + rerankContext);
+        System.out.println("\nDocuments to rerank:");
+        for (int i = 0; i < documents.size(); i++) {
+            System.out.println("  " + i + ": " + documents.get(i));
+        }
+        System.out.println();
+
+        // Basic rerank with default options
+        container.semanticRerank(rerankContext, documents, null)
+            .subscribe(
+                result -> printResults(result),
+                error -> System.err.println("Error: " + error.getMessage())
+            );
+
+        // Rerank with custom options. Keys are passed through to the inference endpoint as-is.
+        // See the JavaDoc on CosmosAsyncContainer.semanticRerank for the list of supported keys.
+        Map<String, Object> options = new HashMap<>();
+        options.put("return_documents", true);
+        options.put("top_k", 3);
+        options.put("sort", true);
+
+        container.semanticRerank(rerankContext, documents, options)
+            .subscribe(
+                result -> printResults(result),
+                error -> System.err.println("Error: " + error.getMessage())
+            );
+    }
+
+    private static void printResults(SemanticRerankResult result) {
+        System.out.println("\nSemantic Rerank Results:");
+        System.out.println("========================\n");
+
+        List<SemanticRerankScore> scores = result.getScores();
+        if (scores != null && !scores.isEmpty()) {
+            System.out.println("Scores (ranked by relevance):");
+            for (int i = 0; i < scores.size(); i++) {
+                SemanticRerankScore score = scores.get(i);
+                System.out.printf("  [%d] index=%d, score=%.7f%n", i, score.getIndex(), score.getScore());
+                if (score.getDocument() != null) {
+                    System.out.println("      document: " + score.getDocument());
+                }
+            }
+        }
+
+        Map<String, Object> latency = result.getLatency();
+        if (latency != null) {
+            System.out.println("\nLatency:");
+            latency.forEach((key, value) -> System.out.printf("  %s: %s%n", key, value));
+        }
+
+        Map<String, Object> tokenUsage = result.getTokenUsage();
+        if (tokenUsage != null) {
+            System.out.println("\nToken Usage:");
+            tokenUsage.forEach((key, value) -> System.out.printf("  %s: %s%n", key, value));
+        }
+    }
+}

From 1f4628c6dea55c2a764b70c8a039ec675624411f Mon Sep 17 00:00:00 2001
From: Bhaskar Mallapragada <mbhaskar@gmail.com>
Date: Mon, 18 May 2026 11:19:53 -0700
Subject: [PATCH 2/6] Address review feedback: accessor pattern, Preconditions,
 Beta version bump

- Bump @Beta annotations from V4_78_0 to V4_81_0 to match the current
  unreleased version (4.81.0-beta.1).
- Use Preconditions.checkNotNull / checkArgument in
  CosmosAsyncContainer.semanticRerank for consistency with the rest of the
  class (addresses kushagraThapar feedback).
- Fix duplicate constant strings in Configs and InferenceService: system
  property uses dot.notation and environment variable uses UPPER_SNAKE_CASE,
  matching the existing convention in Configs.
- Replace the new ModelBridgeInternal bridge methods with the standard
  ImplementationBridgeHelpers accessor pattern for SemanticRerankResult and
  SemanticRerankScore. ModelBridgeInternal.java now has zero diff vs main
  (addresses FabianMeiswinkel API-review feedback). Response POJO setters
  are now private.
- Update CHANGELOG to reflect the correct system property name versus
  environment variable name and to note that Retry-After is honoured on 429.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 sdk/cosmos/azure-cosmos/CHANGELOG.md          |  2 +-
 .../azure/cosmos/CosmosAsyncContainer.java    | 21 ++----
 .../com/azure/cosmos/CosmosContainer.java     |  2 +-
 .../azure/cosmos/implementation/Configs.java  |  9 ++-
 .../ImplementationBridgeHelpers.java          | 74 +++++++++++++++++++
 .../inference/InferenceService.java           | 24 +++---
 .../cosmos/models/ModelBridgeInternal.java    | 33 ---------
 .../cosmos/models/SemanticRerankResult.java   | 50 ++++++++-----
 .../cosmos/models/SemanticRerankScore.java    | 50 ++++++++-----
 .../main/java/com/azure/cosmos/util/Beta.java |  4 +-
 10 files changed, 164 insertions(+), 105 deletions(-)

diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md
index bcf8ab5c3547..66f1ca90c7a9 100644
--- a/sdk/cosmos/azure-cosmos/CHANGELOG.md
+++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -3,7 +3,7 @@
 ### 4.81.0-beta.1 (Unreleased)
 
 #### Features Added
-* Added beta `semanticRerank` API on `CosmosAsyncContainer` and `CosmosContainer` for semantic reranking of documents using the Azure Cosmos DB inference service. Requires AAD authentication (`TokenCredential`) and the `AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT` system property or environment variable to be set. Includes retry logic for transient failures (429, 500, 502, 503) with exponential backoff. - [PR 48258](https://github.com/Azure/azure-sdk-for-java/pull/48258)
+* Added beta `semanticRerank` API on `CosmosAsyncContainer` and `CosmosContainer` for semantic reranking of documents using the Azure Cosmos DB inference service. Requires AAD authentication (`TokenCredential`) and the inference endpoint to be configured via the `azure.cosmos.semanticReranker.inferenceEndpoint` system property or the `AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT` environment variable. Includes retry logic for transient failures (429, 500, 502, 503) with exponential backoff (the server-supplied `Retry-After` header is honoured on 429 responses when present). - [PR 48258](https://github.com/Azure/azure-sdk-for-java/pull/48258)
 
 #### Breaking Changes
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
index 585b7da76c41..7f03ddd552b3 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
@@ -1974,31 +1974,20 @@ public <T> CosmosPagedFlux<T> readAllItems(
      *                </ul>
      * @return A {@link Mono} emitting the {@link SemanticRerankResult}.
      */
-    @Beta(value = Beta.SinceVersion.V4_78_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+    @Beta(value = Beta.SinceVersion.V4_81_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
     public Mono<SemanticRerankResult> semanticRerank(
         String rerankContext,
         List<String> documents,
         Map<String, Object> options) {
 
-        if (rerankContext == null) {
-            return Mono.error(new IllegalArgumentException("Rerank context cannot be null"));
-        }
-        if (documents == null) {
-            return Mono.error(new IllegalArgumentException("Documents list cannot be null"));
-        }
-
-        if (rerankContext.trim().isEmpty()) {
-            return Mono.error(new IllegalArgumentException("Rerank context cannot be empty"));
-        }
-
-        if (documents.isEmpty()) {
-            return Mono.error(new IllegalArgumentException("Documents list cannot be empty"));
-        }
+        checkNotNull(rerankContext, "Argument 'rerankContext' must not be null.");
+        checkNotNull(documents, "Argument 'documents' must not be null.");
+        checkArgument(!rerankContext.trim().isEmpty(), "Argument 'rerankContext' must not be empty.");
+        checkArgument(!documents.isEmpty(), "Argument 'documents' must not be empty.");
 
         return this.database.getClient().getOrCreateInferenceService().semanticRerank(rerankContext, documents, options);
     }
 
-
     /**
      * Replaces an existing item in a container with a new item.
      * It performs a complete replacement of the item,
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
index 61f59f10030b..6d4adc068f37 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
@@ -1086,7 +1086,7 @@ public <TContext> Iterable<CosmosBulkOperationResponse<TContext>> executeBulkOpe
      *                {@code timeout_seconds} (Number) - per-request timeout in seconds (default: 120).
      * @return The semantic rerank result.
      */
-    @Beta(value = Beta.SinceVersion.V4_78_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+    @Beta(value = Beta.SinceVersion.V4_81_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
     public SemanticRerankResult semanticRerank(
         String rerankContext,
         List<String> documents,
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
index 477eb483e87a..32b7af3a3bc0 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
@@ -421,10 +421,11 @@ public class Configs {
     private static final String CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT = "COSMOS.CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT";
     private static final int DEFAULT_CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT = 120;
 
-    // Inference service related configs
-    // Single source of truth: same name is used for both the system property and the environment variable lookup.
-    private static final String INFERENCE_ENDPOINT_PROPERTY = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT";
-    private static final String INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE = INFERENCE_ENDPOINT_PROPERTY;
+    // Inference service related configs.
+    // Following the convention used elsewhere in this file: the system property uses dot.notation
+    // and the environment variable uses UPPER_SNAKE_CASE.
+    private static final String INFERENCE_ENDPOINT_PROPERTY = "azure.cosmos.semanticReranker.inferenceEndpoint";
+    private static final String INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT";
 
     private static final Object lockObject = new Object();
     private static Boolean cachedIsHostnameValidationDisabled = null;
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java
index c48555496c29..38d19b112ac3 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java
@@ -79,6 +79,8 @@
 import com.azure.cosmos.models.PriorityLevel;
 import com.azure.cosmos.models.ShowQueryMode;
 import com.azure.cosmos.models.SqlParameter;
+import com.azure.cosmos.models.SemanticRerankResult;
+import com.azure.cosmos.models.SemanticRerankScore;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.util.CosmosPagedFlux;
 import com.azure.cosmos.util.UtilBridgeInternal;
@@ -2044,4 +2046,76 @@ public interface SqlParameterAccessor {
             SqlParameter createCopy(SqlParameter original);
         }
     }
+
+    public static final class SemanticRerankResultHelper {
+        private static final AtomicReference<SemanticRerankResultAccessor> accessor = new AtomicReference<>();
+        private static final AtomicBoolean semanticRerankResultClassLoaded = new AtomicBoolean(false);
+
+        private SemanticRerankResultHelper() {}
+
+        public static void setSemanticRerankResultAccessor(final SemanticRerankResultAccessor newAccessor) {
+            if (!accessor.compareAndSet(null, newAccessor)) {
+                logger.debug("SemanticRerankResultAccessor already initialized!");
+            } else {
+                logger.debug("Setting SemanticRerankResultAccessor...");
+                semanticRerankResultClassLoaded.set(true);
+            }
+        }
+
+        public static SemanticRerankResultAccessor getSemanticRerankResultAccessor() {
+            if (!semanticRerankResultClassLoaded.get()) {
+                logger.debug("Initializing SemanticRerankResultAccessor...");
+                initializeAllAccessors();
+            }
+
+            SemanticRerankResultAccessor snapshot = accessor.get();
+            if (snapshot == null) {
+                logger.error("SemanticRerankResultAccessor is not initialized yet!");
+            }
+
+            return snapshot;
+        }
+
+        public interface SemanticRerankResultAccessor {
+            void setScores(SemanticRerankResult result, java.util.List<SemanticRerankScore> scores);
+            void setLatency(SemanticRerankResult result, java.util.Map<String, Object> latency);
+            void setTokenUsage(SemanticRerankResult result, java.util.Map<String, Object> tokenUsage);
+        }
+    }
+
+    public static final class SemanticRerankScoreHelper {
+        private static final AtomicReference<SemanticRerankScoreAccessor> accessor = new AtomicReference<>();
+        private static final AtomicBoolean semanticRerankScoreClassLoaded = new AtomicBoolean(false);
+
+        private SemanticRerankScoreHelper() {}
+
+        public static void setSemanticRerankScoreAccessor(final SemanticRerankScoreAccessor newAccessor) {
+            if (!accessor.compareAndSet(null, newAccessor)) {
+                logger.debug("SemanticRerankScoreAccessor already initialized!");
+            } else {
+                logger.debug("Setting SemanticRerankScoreAccessor...");
+                semanticRerankScoreClassLoaded.set(true);
+            }
+        }
+
+        public static SemanticRerankScoreAccessor getSemanticRerankScoreAccessor() {
+            if (!semanticRerankScoreClassLoaded.get()) {
+                logger.debug("Initializing SemanticRerankScoreAccessor...");
+                initializeAllAccessors();
+            }
+
+            SemanticRerankScoreAccessor snapshot = accessor.get();
+            if (snapshot == null) {
+                logger.error("SemanticRerankScoreAccessor is not initialized yet!");
+            }
+
+            return snapshot;
+        }
+
+        public interface SemanticRerankScoreAccessor {
+            void setIndex(SemanticRerankScore score, int index);
+            void setScore(SemanticRerankScore score, double scoreValue);
+            void setDocument(SemanticRerankScore score, String document);
+        }
+    }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
index 7474d62017b8..ab387daa36b4 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
@@ -11,6 +11,7 @@
 import com.azure.cosmos.implementation.BadRequestException;
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.HttpConstants;
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.LifeCycleUtils;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.http.HttpClient;
@@ -18,7 +19,6 @@
 import com.azure.cosmos.implementation.http.HttpHeaders;
 import com.azure.cosmos.implementation.http.HttpRequest;
 import com.azure.cosmos.implementation.http.HttpResponse;
-import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.SemanticRerankResult;
 import com.azure.cosmos.models.SemanticRerankScore;
 import com.fasterxml.jackson.databind.JsonNode;
@@ -63,9 +63,9 @@ public class InferenceService implements AutoCloseable {
         "cosmos-inference-java/" + HttpConstants.Versions.getSdkVersion();
     private static final ObjectMapper OBJECT_MAPPER = Utils.getSimpleObjectMapper();
 
-    // System property / environment variable name (single string used for both lookup mechanisms)
-    private static final String MAX_CONNECTION_LIMIT_PROPERTY = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_SERVICE_MAX_CONNECTION_LIMIT";
-    private static final String MAX_CONNECTION_LIMIT_VARIABLE = MAX_CONNECTION_LIMIT_PROPERTY;
+    // Following the convention used in Configs: system property is dot.notation; environment variable is UPPER_SNAKE_CASE.
+    private static final String MAX_CONNECTION_LIMIT_PROPERTY = "azure.cosmos.semanticReranker.inferenceService.maxConnectionLimit";
+    private static final String MAX_CONNECTION_LIMIT_VARIABLE = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_SERVICE_MAX_CONNECTION_LIMIT";
 
     // Option keys that callers can pass in the options map
     public static final String OPTION_TIMEOUT_SECONDS = "timeout_seconds";
@@ -428,6 +428,10 @@ private Mono<SemanticRerankResult> parseResponse(HttpResponse response) {
                 try {
                     JsonNode rootNode = OBJECT_MAPPER.readTree(bodyString);
                     SemanticRerankResult result = new SemanticRerankResult();
+                    ImplementationBridgeHelpers.SemanticRerankResultHelper.SemanticRerankResultAccessor resultAccessor =
+                        ImplementationBridgeHelpers.SemanticRerankResultHelper.getSemanticRerankResultAccessor();
+                    ImplementationBridgeHelpers.SemanticRerankScoreHelper.SemanticRerankScoreAccessor scoreAccessor =
+                        ImplementationBridgeHelpers.SemanticRerankScoreHelper.getSemanticRerankScoreAccessor();
 
                     // Parse scores
                     if (rootNode.has("Scores")) {
@@ -440,20 +444,20 @@ private Mono<SemanticRerankResult> parseResponse(HttpResponse response) {
                                 JsonNode indexNode = scoreNode.get("index");
                                 JsonNode scoreValNode = scoreNode.get("score");
                                 if (indexNode != null) {
-                                    ModelBridgeInternal.setSemanticRerankScoreIndex(score, indexNode.asInt());
+                                    scoreAccessor.setIndex(score, indexNode.asInt());
                                 }
                                 if (scoreValNode != null) {
-                                    ModelBridgeInternal.setSemanticRerankScoreScore(score, scoreValNode.asDouble());
+                                    scoreAccessor.setScore(score, scoreValNode.asDouble());
                                 }
 
                                 if (scoreNode.has("document")) {
-                                    ModelBridgeInternal.setSemanticRerankScoreDocument(score, scoreNode.get("document").asText());
+                                    scoreAccessor.setDocument(score, scoreNode.get("document").asText());
                                 }
 
                                 scores.add(score);
                             }
                         }
-                        ModelBridgeInternal.setSemanticRerankResultScores(result, scores);
+                        resultAccessor.setScores(result, scores);
                     }
 
                     // Parse latency
@@ -461,7 +465,7 @@ private Mono<SemanticRerankResult> parseResponse(HttpResponse response) {
                         Map<String, Object> latency = new HashMap<>();
                         rootNode.get("latency").fields().forEachRemaining(
                             entry -> latency.put(entry.getKey(), entry.getValue().asDouble()));
-                        ModelBridgeInternal.setSemanticRerankResultLatency(result, latency);
+                        resultAccessor.setLatency(result, latency);
                     }
 
                     // Parse token usage
@@ -469,7 +473,7 @@ private Mono<SemanticRerankResult> parseResponse(HttpResponse response) {
                         Map<String, Object> tokenUsage = new HashMap<>();
                         rootNode.get("token_usage").fields().forEachRemaining(
                             entry -> tokenUsage.put(entry.getKey(), entry.getValue().asInt()));
-                        ModelBridgeInternal.setSemanticRerankResultTokenUsage(result, tokenUsage);
+                        resultAccessor.setTokenUsage(result, tokenUsage);
                     }
 
                     if (logger.isDebugEnabled()) {
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
index e26b48115787..7dde81e06e98 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
@@ -744,39 +744,6 @@ public static List<PatchOperation> getPatchOperationsFromCosmosPatch(CosmosPatch
         return cosmosPatchOperations.getPatchOperations();
     }
 
-    // Bridge methods for SemanticRerankResult and SemanticRerankScore — setters are package-private
-    // to prevent mutation by callers; InferenceService uses these bridge methods to populate instances.
-
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static void setSemanticRerankResultScores(SemanticRerankResult result, java.util.List<SemanticRerankScore> scores) {
-        result.setScores(scores);
-    }
-
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static void setSemanticRerankResultLatency(SemanticRerankResult result, Map<String, Object> latency) {
-        result.setLatency(latency);
-    }
-
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static void setSemanticRerankResultTokenUsage(SemanticRerankResult result, Map<String, Object> tokenUsage) {
-        result.setTokenUsage(tokenUsage);
-    }
-
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static void setSemanticRerankScoreIndex(SemanticRerankScore score, int index) {
-        score.setIndex(index);
-    }
-
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static void setSemanticRerankScoreScore(SemanticRerankScore score, double scoreValue) {
-        score.setScore(scoreValue);
-    }
-
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static void setSemanticRerankScoreDocument(SemanticRerankScore score, String document) {
-        score.setDocument(document);
-    }
-
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static void  initializeAllAccessors() {
         CosmosBatch.initialize();
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
index 7f40dc5fb646..bfbe0330b909 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.models;
 
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.util.Beta;
 
 import java.util.List;
@@ -10,7 +11,7 @@
 /**
  * Represents the result of a semantic rerank operation.
  */
-@Beta(value = Beta.SinceVersion.V4_78_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+@Beta(value = Beta.SinceVersion.V4_81_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
 public final class SemanticRerankResult {
     private List<SemanticRerankScore> scores;
     private Map<String, Object> latency;
@@ -31,12 +32,7 @@ public List<SemanticRerankScore> getScores() {
         return scores;
     }
 
-    /**
-     * Sets the list of scored documents.
-     *
-     * @param scores the list of scored documents.
-     */
-    void setScores(List<SemanticRerankScore> scores) {
+    private void setScores(List<SemanticRerankScore> scores) {
         this.scores = scores;
     }
 
@@ -49,12 +45,7 @@ public Map<String, Object> getLatency() {
         return latency;
     }
 
-    /**
-     * Sets the latency information for the operation.
-     *
-     * @param latency the latency information map.
-     */
-    void setLatency(Map<String, Object> latency) {
+    private void setLatency(Map<String, Object> latency) {
         this.latency = latency;
     }
 
@@ -67,12 +58,33 @@ public Map<String, Object> getTokenUsage() {
         return tokenUsage;
     }
 
-    /**
-     * Sets the token usage information for the operation.
-     *
-     * @param tokenUsage the token usage information map.
-     */
-    void setTokenUsage(Map<String, Object> tokenUsage) {
+    private void setTokenUsage(Map<String, Object> tokenUsage) {
         this.tokenUsage = tokenUsage;
     }
+
+    ///////////////////////////////////////////////////////////////////////////////////////////
+    // the following helper/accessor only helps to access this class outside of this package.//
+    ///////////////////////////////////////////////////////////////////////////////////////////
+    static void initialize() {
+        ImplementationBridgeHelpers.SemanticRerankResultHelper.setSemanticRerankResultAccessor(
+            new ImplementationBridgeHelpers.SemanticRerankResultHelper.SemanticRerankResultAccessor() {
+                @Override
+                public void setScores(SemanticRerankResult result, List<SemanticRerankScore> scores) {
+                    result.setScores(scores);
+                }
+
+                @Override
+                public void setLatency(SemanticRerankResult result, Map<String, Object> latency) {
+                    result.setLatency(latency);
+                }
+
+                @Override
+                public void setTokenUsage(SemanticRerankResult result, Map<String, Object> tokenUsage) {
+                    result.setTokenUsage(tokenUsage);
+                }
+            }
+        );
+    }
+
+    static { initialize(); }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java
index 625b403b0ce3..96ae57ce8389 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java
@@ -2,12 +2,13 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.models;
 
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.util.Beta;
 
 /**
  * Represents a single scored document in the semantic rerank result.
  */
-@Beta(value = Beta.SinceVersion.V4_78_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
+@Beta(value = Beta.SinceVersion.V4_81_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
 public final class SemanticRerankScore {
     private int index;
     private String document;
@@ -28,12 +29,7 @@ public int getIndex() {
         return index;
     }
 
-    /**
-     * Sets the index of the document in the original input list.
-     *
-     * @param index the document index.
-     */
-    void setIndex(int index) {
+    private void setIndex(int index) {
         this.index = index;
     }
 
@@ -46,12 +42,7 @@ public String getDocument() {
         return document;
     }
 
-    /**
-     * Sets the document text.
-     *
-     * @param document the document text.
-     */
-    void setDocument(String document) {
+    private void setDocument(String document) {
         this.document = document;
     }
 
@@ -64,12 +55,33 @@ public double getScore() {
         return score;
     }
 
-    /**
-     * Sets the semantic relevance score for this document.
-     *
-     * @param score the relevance score.
-     */
-    void setScore(double score) {
+    private void setScore(double score) {
         this.score = score;
     }
+
+    ///////////////////////////////////////////////////////////////////////////////////////////
+    // the following helper/accessor only helps to access this class outside of this package.//
+    ///////////////////////////////////////////////////////////////////////////////////////////
+    static void initialize() {
+        ImplementationBridgeHelpers.SemanticRerankScoreHelper.setSemanticRerankScoreAccessor(
+            new ImplementationBridgeHelpers.SemanticRerankScoreHelper.SemanticRerankScoreAccessor() {
+                @Override
+                public void setIndex(SemanticRerankScore score, int index) {
+                    score.setIndex(index);
+                }
+
+                @Override
+                public void setScore(SemanticRerankScore score, double scoreValue) {
+                    score.setScore(scoreValue);
+                }
+
+                @Override
+                public void setDocument(SemanticRerankScore score, String document) {
+                    score.setDocument(document);
+                }
+            }
+        );
+    }
+
+    static { initialize(); }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java
index ca01b8cc9f65..ae90b5314418 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java
@@ -110,7 +110,7 @@ public enum SinceVersion {
         V4_74_0,
         /** v4.75.0 */
         V4_75_0,
-        /** v4.78.0 */
-        V4_78_0
+        /** v4.81.0 */
+        V4_81_0
     }
 }

From 228b26965a79890bf9596c173a6d47d04b8195a3 Mon Sep 17 00:00:00 2001
From: Bhaskar Mallapragada <mbhaskar@gmail.com>
Date: Mon, 18 May 2026 15:49:58 -0700
Subject: [PATCH 3/6] Fix InferenceServiceTest failures

- Update test constant INFERENCE_ENDPOINT_PROPERTY to match the new system-property name ('azure.cosmos.semanticReranker.inferenceEndpoint') used by Configs.getInferenceServiceEndpoint(). The previous value was the env-variable name and no longer matched after the property/env constants were split.

- Add input validation in InferenceService.semanticRerank(): synchronously throw NPE for null rerankContext/documents and emit Mono.error(IllegalArgumentException) for empty values. This restores the contract the unit tests assert and provides defense-in-depth alongside the container-level checks.

- Register SemanticRerankResult and SemanticRerankScore accessors via ModelBridgeInternal.initializeAllAccessors() so the ImplementationBridgeHelpers accessor pattern works when classes are first reached via the helper (fixes NPE 'scoreAccessor is null' in semanticRerankShouldIncludeAuthorizationHeader).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../implementation/inference/InferenceServiceTest.java |  2 +-
 .../implementation/inference/InferenceService.java     | 10 ++++++++++
 .../com/azure/cosmos/models/ModelBridgeInternal.java   |  2 ++
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
index 24bbcb1d6da2..b151acc13a61 100644
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
@@ -40,7 +40,7 @@ public class InferenceServiceTest {
 
     private static final String MOCK_TOKEN = "mock-access-token";
     private static final String TEST_ENDPOINT = "https://test-inference.westus.dbinference.azure.com";
-    private static final String INFERENCE_ENDPOINT_PROPERTY = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT";
+    private static final String INFERENCE_ENDPOINT_PROPERTY = "azure.cosmos.semanticReranker.inferenceEndpoint";
 
     // Total calls expected for a fully-retried retryable failure: 1 initial + RETRY_MAX_ATTEMPTS retries
     private static final int TOTAL_CALLS_AFTER_EXHAUSTION = 1 + InferenceService.RETRY_MAX_ATTEMPTS;
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
index ab387daa36b4..0b117613eeb3 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
@@ -221,6 +221,16 @@ public Mono<SemanticRerankResult> semanticRerank(
         List<String> documents,
         Map<String, Object> options) {
 
+        checkNotNull(rerankContext, "Rerank context cannot be null");
+        checkNotNull(documents, "Documents list cannot be null");
+
+        if (rerankContext.trim().isEmpty()) {
+            return Mono.error(new IllegalArgumentException("Rerank context cannot be empty"));
+        }
+        if (documents.isEmpty()) {
+            return Mono.error(new IllegalArgumentException("Documents list cannot be empty"));
+        }
+
         // Resolve the per-request timeout: options > default
         final Duration requestTimeout = resolveRequestTimeout(options);
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
index 7dde81e06e98..f0253ad99486 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
@@ -769,5 +769,7 @@ public static void  initializeAllAccessors() {
         PriorityLevel.initialize();
         SqlQuerySpec.initialize();
         SqlParameter.initialize();
+        SemanticRerankResult.initialize();
+        SemanticRerankScore.initialize();
     }
 }

From 56e30bcacfa4ab6c2075e1424ad3ecca54722479 Mon Sep 17 00:00:00 2001
From: Bhaskar Mallapragada <mbhaskar@gmail.com>
Date: Mon, 18 May 2026 17:00:49 -0700
Subject: [PATCH 4/6] Fix InferenceServiceTest @BeforeMethod not firing under
 TestNG group filter

The unit-testng.xml suite filters tests to the 'unit' group. TestNG excludes
@BeforeMethod / @AfterMethod configuration methods that are not in an
included group (and don't have alwaysRun=true), so the bare @BeforeMethod
and @AfterMethod annotations were skipped in CI.

This caused mockTokenCredential to remain null, producing NPEs at
'new InferenceService(mockTokenCredential)' with the message 'Semantic
reranking requires AAD authentication' for all 28 tests that depend on the
mocked credential.

Annotate @BeforeMethod/@AfterMethod with groups = {"unit"} to match the
rest of the codebase (RxPartitionKeyRangeCacheTest, SqlQuerySpecLoggerTest,
etc.) and ensure the setup/teardown fire under the 'unit' group filter.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../cosmos/implementation/inference/InferenceServiceTest.java | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
index b151acc13a61..68dfee640d18 100644
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
@@ -49,7 +49,7 @@ public class InferenceServiceTest {
     private HttpClient mockHttpClient;
     private MockedStatic<HttpClient> httpClientStaticMock;
 
-    @BeforeMethod
+    @BeforeMethod(groups = {"unit"})
     public void setUp() {
         // Set the inference endpoint system property required by InferenceService constructor
         System.setProperty(INFERENCE_ENDPOINT_PROPERTY, TEST_ENDPOINT);
@@ -68,7 +68,7 @@ public void setUp() {
             .thenReturn(mockHttpClient);
     }
 
-    @AfterMethod
+    @AfterMethod(groups = {"unit"})
     public void tearDown() {
         System.clearProperty(INFERENCE_ENDPOINT_PROPERTY);
         if (httpClientStaticMock != null) {

From 4fe0e753e141e80dcbfa27738daa7518b8d4e479 Mon Sep 17 00:00:00 2001
From: Bhaskar Mallapragada <mbhaskar@gmail.com>
Date: Tue, 19 May 2026 13:48:49 -0700
Subject: [PATCH 5/6] Cosmos: address PR #48258 review feedback (#1 #2 #3)

Three small but correctness-relevant fixes raised by @kushagraThapar:

1. InferenceService: payload serialization (Jackson) failures are now
   surfaced as BadRequestException(400) with sub-status
   CUSTOM_SERIALIZER_EXCEPTION instead of a generic CosmosException(500).
   400 is not in RETRYABLE_STATUS_CODES, so the retry policy no longer
   wastes attempts on a deterministic payload failure. Matches the
   pattern already used for response-parse failures.

2. CosmosContainer.semanticRerank (sync): apply timeout_seconds as a
   whole-operation deadline via Mono.timeout(...) and map
   java.util.concurrent.TimeoutException to RequestTimeoutException (408)
   so callers see a typed CosmosException instead of Reactor's raw
   RuntimeException("Timeout on blocking read"). Javadoc updated to
   document the sync-vs-async timeout semantics.

3. CosmosAsyncClient: fix a race between close() and lazy initialization
   of InferenceService that could permanently leak a Netty event-loop
   group and pooled connections. close() now sets a volatile closed
   flag first, then atomically transfers the reference via getAndSet(null).
   getOrCreateInferenceService re-checks closed after a successful CAS
   and tears down any late-arriving instance, raising IllegalStateException.

CHANGELOG entries added for all three under 4.81.0-beta.1.

Verified: mvn install on azure-cosmos, checkstyle+spotbugs clean,
InferenceServiceTest 29/29 passing.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 sdk/cosmos/azure-cosmos/CHANGELOG.md          |  3 +++
 .../com/azure/cosmos/CosmosAsyncClient.java   | 20 +++++++++++++++++-
 .../com/azure/cosmos/CosmosContainer.java     | 21 ++++++++++++++++++-
 .../inference/InferenceService.java           |  9 +++++++-
 4 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md
index 66f1ca90c7a9..106c55521e10 100644
--- a/sdk/cosmos/azure-cosmos/CHANGELOG.md
+++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -8,6 +8,9 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Semantic rerank: serialization failures are now surfaced as a typed `BadRequestException` (HTTP 400, sub-status `CUSTOM_SERIALIZER_EXCEPTION`) instead of a generic `500`, so the SDK retry policy does not retry deterministic payload failures. - [PR 48258](https://github.com/Azure/azure-sdk-for-java/pull/48258)
+* Semantic rerank: the synchronous `CosmosContainer.semanticRerank` API now applies `timeout_seconds` as a whole-operation deadline and maps Reactor's blocking timeout to `RequestTimeoutException` (HTTP 408), matching the typed `CosmosException` contract callers expect. - [PR 48258](https://github.com/Azure/azure-sdk-for-java/pull/48258)
+* Semantic rerank: fixed a race between `CosmosAsyncClient.close()` and lazy initialization of the inference HTTP client that could leak a Netty event-loop group and pooled connections if a `semanticRerank` call arrived while another thread was closing the client. `close()` now sets a `closed` flag, atomically transfers the reference, and any late-arriving instance is torn down. - [PR 48258](https://github.com/Azure/azure-sdk-for-java/pull/48258)
 
 #### Other Changes
 * Replaced per-client `Schedulers.newSingle()` schedulers in `GlobalEndpointManager` and `GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker` with shared `BoundedElastic` schedulers in `CosmosSchedulers` to prevent thread count from scaling linearly with client/tenant count. - See [PR 49062](https://github.com/Azure/azure-sdk-for-java/pull/49062)
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
index cb52abd347ee..d52d4f3c44e9 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
@@ -123,6 +123,7 @@ private static ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccess
     private final CosmosItemSerializer defaultCustomSerializer;
     private final java.util.function.Function<CosmosAsyncContainer, CosmosAsyncContainer> containerFactory;
     private final AtomicReference<InferenceService> inferenceService = new AtomicReference<>();
+    private volatile boolean closed;
 
     CosmosAsyncClient(CosmosClientBuilder builder) {
         // Async Cosmos client wrapper
@@ -312,6 +313,9 @@ TokenCredential tokenCredential() {
      * @throws IllegalStateException if the client was built with key-based auth instead of AAD.
      */
     InferenceService getOrCreateInferenceService() {
+        if (this.closed) {
+            throw new IllegalStateException("CosmosAsyncClient has been closed.");
+        }
         if (this.tokenCredential == null) {
             throw new IllegalStateException(
                 "Semantic reranking requires AAD authentication. "
@@ -325,6 +329,15 @@ InferenceService getOrCreateInferenceService() {
             if (!this.inferenceService.compareAndSet(null, newSvc)) {
                 // Another thread already set the instance; close the losing one to prevent resource leak
                 newSvc.close();
+            } else if (this.closed) {
+                // close() ran concurrently and already saw a null reference; ensure the late-arriving
+                // instance we just published is torn down so we don't leak the Netty event-loop group
+                // and pooled connections held by the underlying HttpClient.
+                InferenceService leaked = this.inferenceService.getAndSet(null);
+                if (leaked != null) {
+                    leaked.close();
+                }
+                throw new IllegalStateException("CosmosAsyncClient has been closed.");
             }
         }
         return this.inferenceService.get();
@@ -598,10 +611,15 @@ public CosmosAsyncDatabase getDatabase(String id) {
      */
     @Override
     public void close() {
+        // Set the closed flag BEFORE reading the inference reference. getOrCreateInferenceService()
+        // re-checks this flag after publishing a new instance and tears it down if it raced with us,
+        // so the combination of (closed = true) + getAndSet(null) here guarantees that any
+        // InferenceService that is ever published is also closed exactly once.
+        this.closed = true;
         if (this.clientMetricRegistrySnapshot != null) {
             ClientTelemetryMetrics.remove(this.clientMetricRegistrySnapshot);
         }
-        InferenceService svc = this.inferenceService.get();
+        InferenceService svc = this.inferenceService.getAndSet(null);
         if (svc != null) {
             svc.close();
         }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
index 6d4adc068f37..368f4d43ca4f 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
@@ -5,6 +5,7 @@
 
 import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup;
 import com.azure.cosmos.implementation.inference.InferenceService;
+import com.azure.cosmos.implementation.RequestTimeoutException;
 import com.azure.cosmos.models.CosmosBatch;
 import com.azure.cosmos.models.CosmosBatchOperationResult;
 import com.azure.cosmos.models.CosmosBatchRequestOptions;
@@ -333,7 +334,16 @@ private SemanticRerankResult blockSemanticRerankResponse(
             }
         }
         try {
-            return semanticRerankResultMono.block(blockTimeout);
+            // Apply the same timeout semantics on the sync wrapper that InferenceService applies per
+            // HTTP attempt: convert a timeout to a typed RequestTimeoutException (408) instead of
+            // letting Reactor's raw RuntimeException("Timeout on blocking read") leak to the caller.
+            final Duration timeout = blockTimeout;
+            return semanticRerankResultMono
+                .timeout(timeout)
+                .onErrorMap(java.util.concurrent.TimeoutException.class,
+                    t -> new RequestTimeoutException(
+                        "Semantic rerank request timed out after " + timeout, (java.net.URI) null))
+                .block(timeout);
         } catch (Exception ex) {
             final Throwable throwable = Exceptions.unwrap(ex);
             if (throwable instanceof CosmosException) {
@@ -1076,6 +1086,15 @@ public <TContext> Iterable<CosmosBulkOperationResponse<TContext>> executeBulkOpe
      * To override, pass {@code "timeout_seconds"} (as a {@link Number}) in the {@code options} map,
      * for example: {@code options.put("timeout_seconds", 30)}.
      *
+     * <p><strong>Sync vs async timeout semantics:</strong> on the async API
+     * ({@code CosmosAsyncContainer.semanticRerank}) {@code timeout_seconds} is applied per HTTP attempt,
+     * so the total wall-clock can stretch across retries and backoff. On this sync API the same value
+     * is also applied as the {@link Mono#timeout(java.time.Duration) whole-operation deadline}
+     * around the inner {@link Mono}, so the call returns once the overall budget elapses regardless of
+     * how many retries the async path would otherwise perform. When the budget elapses the sync API
+     * surfaces a {@link CosmosException} with status {@code 408 Request Timeout} rather than a raw
+     * Reactor {@link RuntimeException}.
+     *
      * @param rerankContext The query or context string used to score documents.
      * @param documents The list of document strings to rerank.
      * @param options Optional reranking parameters as a map. Supported keys include:
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
index 0b117613eeb3..36271e3172df 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
@@ -269,7 +269,14 @@ public Mono<SemanticRerankResult> semanticRerank(
 
                 } catch (IOException e) {
                     logger.error("Failed to serialize request payload", e);
-                    return Mono.error(BridgeInternal.createCosmosException(500, "Failed to serialize semantic rerank request"));
+                    // Use BadRequestException (400) + CUSTOM_SERIALIZER_EXCEPTION sub-status — same
+                    // pattern as the response-parse path below. 400 is not in RETRYABLE_STATUS_CODES,
+                    // so withRetry will not waste attempts on a deterministic serialization failure.
+                    BadRequestException badRequest = new BadRequestException(
+                        "Failed to serialize semantic rerank request: " + e.getMessage(), e);
+                    BridgeInternal.setSubStatusCode(badRequest,
+                        HttpConstants.SubStatusCodes.CUSTOM_SERIALIZER_EXCEPTION);
+                    return Mono.error(badRequest);
                 }
             })
             .as(this::withRetry)

From f40594447f64e4e7495a3e5ea9c1d341282cc694 Mon Sep 17 00:00:00 2001
From: Bhaskar Mallapragada <mbhaskar@gmail.com>
Date: Thu, 4 Jun 2026 15:23:22 -0700
Subject: [PATCH 6/6] Extract semantic reranking into new azure-cosmos-ai
 module

Create new standalone Maven module sdk/cosmos/azure-cosmos-ai with:
- CosmosAIClientBuilder, CosmosAIAsyncClient, CosmosAIClient public API
- InferenceService rewritten to use azure-core HttpPipeline
- SemanticRerankResult/SemanticRerankScore models with package-private setters
- InferenceResponseParser for JSON parsing (in models package)
- Unit tests (22 passing) and sample code

Remove all semantic rerank code from azure-cosmos:
- CosmosAsyncContainer.semanticRerank() and CosmosContainer.semanticRerank()
- CosmosAsyncClient.getOrCreateInferenceService() and close() cleanup
- ImplementationBridgeHelpers SemanticRerankResult/ScoreHelper classes
- ModelBridgeInternal.initialize() calls for SemanticRerank types
- Configs.java inference endpoint properties
- Old InferenceService, model classes, test, and sample files
---
 eng/versioning/version_client.txt             |   1 +
 sdk/cosmos/azure-cosmos-ai/pom.xml            | 183 +++++
 .../azure/cosmos/ai/CosmosAIAsyncClient.java  |  78 +++
 .../com/azure/cosmos/ai/CosmosAIClient.java   |  90 +++
 .../cosmos/ai/CosmosAIClientBuilder.java      | 207 ++++++
 .../ai/implementation/InferenceService.java   | 302 +++++++++
 .../ai/models/InferenceResponseParser.java    |  84 +++
 .../ai/models/SemanticRerankResult.java       |  61 ++
 .../cosmos/ai/models/SemanticRerankScore.java |  58 ++
 .../cosmos/ai}/SemanticRerankSample.java      |  52 +-
 .../implementation/InferenceServiceTest.java  | 474 +++++++++++++
 .../inference/InferenceServiceTest.java       | 628 ------------------
 .../com/azure/cosmos/CosmosAsyncClient.java   |  47 --
 .../azure/cosmos/CosmosAsyncContainer.java    |  44 +-
 .../com/azure/cosmos/CosmosContainer.java     |  73 --
 .../azure/cosmos/implementation/Configs.java  |  19 -
 .../ImplementationBridgeHelpers.java          |  74 ---
 .../inference/InferenceService.java           | 534 ---------------
 .../cosmos/models/ModelBridgeInternal.java    |   2 -
 .../cosmos/models/SemanticRerankResult.java   |  90 ---
 .../cosmos/models/SemanticRerankScore.java    |  87 ---
 sdk/cosmos/pom.xml                            |   1 +
 22 files changed, 1566 insertions(+), 1623 deletions(-)
 create mode 100644 sdk/cosmos/azure-cosmos-ai/pom.xml
 create mode 100644 sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIAsyncClient.java
 create mode 100644 sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClient.java
 create mode 100644 sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClientBuilder.java
 create mode 100644 sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/implementation/InferenceService.java
 create mode 100644 sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/InferenceResponseParser.java
 create mode 100644 sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/SemanticRerankResult.java
 create mode 100644 sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/SemanticRerankScore.java
 rename sdk/cosmos/{azure-cosmos/src/samples/java/com/azure/cosmos/semanticrerank => azure-cosmos-ai/src/samples/java/com/azure/cosmos/ai}/SemanticRerankSample.java (70%)
 create mode 100644 sdk/cosmos/azure-cosmos-ai/src/test/java/com/azure/cosmos/ai/implementation/InferenceServiceTest.java
 delete mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
 delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
 delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
 delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java

diff --git a/eng/versioning/version_client.txt b/eng/versioning/version_client.txt
index 1c7aaef057b5..2976187cd70e 100644
--- a/eng/versioning/version_client.txt
+++ b/eng/versioning/version_client.txt
@@ -111,6 +111,7 @@ com.azure:azure-cosmos-benchmark;4.0.1-beta.1;4.0.1-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark_3;0.0.1-beta.1;0.0.1-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark_3-5;0.0.1-beta.1;0.0.1-beta.1
 com.azure:azure-cosmos-encryption;2.29.0;2.30.0-beta.1
+com.azure:azure-cosmos-ai;1.0.0-beta.1;1.0.0-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark-account-data-resolver-sample;1.0.0-beta.1;1.0.0-beta.1
 com.azure:azure-cosmos-test;1.0.0-beta.18;1.0.0-beta.19
 com.azure.cosmos.spark:azure-cosmos-spark_3-3_2-12;4.48.0;4.49.0-beta.1
diff --git a/sdk/cosmos/azure-cosmos-ai/pom.xml b/sdk/cosmos/azure-cosmos-ai/pom.xml
new file mode 100644
index 000000000000..c78e1fb490eb
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/pom.xml
@@ -0,0 +1,183 @@
+<!--
+Copyright (c) Microsoft Corporation. All rights reserved.
+Licensed under the MIT License.
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>com.azure</groupId>
+    <artifactId>azure-client-sdk-parent</artifactId>
+    <version>1.7.0</version> <!-- {x-version-update;com.azure:azure-client-sdk-parent;current} -->
+    <relativePath>../../parents/azure-client-sdk-parent</relativePath>
+  </parent>
+
+  <groupId>com.azure</groupId>
+  <artifactId>azure-cosmos-ai</artifactId>
+  <version>1.0.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos-ai;current} -->
+  <name>Microsoft Azure Cosmos DB AI Services SDK</name>
+  <description>This Package contains Microsoft Azure Cosmos DB AI Services SDK for semantic reranking and inference operations</description>
+  <packaging>jar</packaging>
+  <url>https://github.com/Azure/azure-sdk-for-java</url>
+
+  <distributionManagement>
+    <site>
+      <id>azure-java-build-docs</id>
+      <url>${site.url}/site/${project.artifactId}</url>
+    </site>
+  </distributionManagement>
+
+  <scm>
+    <url>scm:git:https://github.com/Azure/azure-sdk-for-java</url>
+    <connection>scm:git:git@github.com:Azure/azure-sdk-for-java.git</connection>
+    <tag>HEAD</tag>
+  </scm>
+
+  <properties>
+    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+    <jacoco.min.linecoverage>0.09</jacoco.min.linecoverage>
+    <jacoco.min.branchcoverage>0.02</jacoco.min.branchcoverage>
+
+    <!-- CosmosSkip - This can only be enabled once we release GA, as it needs a stable version to check for breaking changes. -->
+    <revapi.skip>true</revapi.skip>
+    <checkstyle.excludes>**/implementation/**/*.java</checkstyle.excludes>
+    <checkstyle.includeTestSourceDirectory>false</checkstyle.includeTestSourceDirectory>
+    <doclintMissingInclusion>-</doclintMissingInclusion>
+    <spotless.skip>true</spotless.skip>
+  </properties>
+
+  <dependencies>
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-core</artifactId>
+      <version>1.58.0</version> <!-- {x-version-update;com.azure:azure-core;dependency} -->
+    </dependency>
+
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-core-http-netty</artifactId>
+      <version>1.16.4</version> <!-- {x-version-update;com.azure:azure-core-http-netty;dependency} -->
+    </dependency>
+
+    <!-- Added this provided dependency to include necessary annotations used by "reactor-core".
+         Without this dependency, javadoc throws a warning as it cannot find enum When.MAYBE
+         which is used in @Nullable annotation in reactor core classes.
+
+         Similar provided dependency exits for "azure-core" as well.
+    -->
+    <dependency>
+      <groupId>com.google.code.findbugs</groupId>
+      <artifactId>jsr305</artifactId>
+      <version>3.0.2</version> <!-- {x-version-update;com.google.code.findbugs:jsr305;external_dependency} -->
+      <scope>provided</scope>
+    </dependency>
+
+    <!-- Test dependencies -->
+    <dependency>
+      <groupId>io.projectreactor</groupId>
+      <artifactId>reactor-test</artifactId>
+      <version>3.7.17</version> <!-- {x-version-update;io.projectreactor:reactor-test;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.testng</groupId>
+      <artifactId>testng</artifactId>
+      <version>7.5.1</version> <!-- {x-version-update;org.testng:testng;external_dependency} -->
+      <scope>test</scope>
+      <exclusions>
+        <exclusion>
+          <groupId>org.apache.ant</groupId>
+          <artifactId>ant</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>org.yaml</groupId>
+          <artifactId>snakeyaml</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+
+    <dependency>
+      <groupId>org.assertj</groupId>
+      <artifactId>assertj-core</artifactId>
+      <version>3.27.7</version> <!-- {x-version-update;org.assertj:assertj-core;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <version>4.11.0</version> <!-- {x-version-update;org.mockito:mockito-core;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy</artifactId>
+      <version>1.17.7</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy-agent</artifactId>
+      <version>1.17.7</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy-agent;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.logging.log4j</groupId>
+      <artifactId>log4j-slf4j-impl</artifactId>
+      <version>2.17.2</version> <!-- {x-version-update;org.apache.logging.log4j:log4j-slf4j-impl;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.logging.log4j</groupId>
+      <artifactId>log4j-api</artifactId>
+      <version>2.17.2</version> <!-- {x-version-update;org.apache.logging.log4j:log4j-api;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.logging.log4j</groupId>
+      <artifactId>log4j-core</artifactId>
+      <version>2.17.2</version> <!-- {x-version-update;org.apache.logging.log4j:log4j-core;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+
+  <build>
+    <plugins>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-surefire-plugin</artifactId>
+        <version>3.5.3</version> <!-- {x-version-update;org.apache.maven.plugins:maven-surefire-plugin;external_dependency} -->
+        <configuration>
+          <groups>unit</groups>
+          <includes>
+            <include>%regex[.*]</include>
+          </includes>
+          <properties>
+            <property>
+              <name>surefire.testng.verbose</name>
+              <value>2</value>
+            </property>
+          </properties>
+        </configuration>
+      </plugin>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-enforcer-plugin</artifactId>
+        <version>3.6.1</version> <!-- {x-version-update;org.apache.maven.plugins:maven-enforcer-plugin;external_dependency} -->
+        <configuration>
+          <rules>
+            <bannedDependencies>
+              <includes>
+                <include>com.azure:*</include>
+              </includes>
+            </bannedDependencies>
+          </rules>
+        </configuration>
+      </plugin>
+    </plugins>
+  </build>
+</project>
diff --git a/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIAsyncClient.java b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIAsyncClient.java
new file mode 100644
index 000000000000..2ad03ba13898
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIAsyncClient.java
@@ -0,0 +1,78 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai;
+
+import com.azure.cosmos.ai.implementation.InferenceService;
+import com.azure.cosmos.ai.models.SemanticRerankResult;
+import reactor.core.publisher.Mono;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Asynchronous client for Azure Cosmos DB AI services.
+ *
+ * <p>Provides semantic reranking of documents using the Azure Cosmos DB inference service.
+ * Use {@link CosmosAIClientBuilder} to create an instance.</p>
+ *
+ * <p><strong>Example:</strong></p>
+ * <pre>
+ * CosmosAIAsyncClient client = new CosmosAIClientBuilder()
+ *     .endpoint("https://my-inference.dbinference.azure.com")
+ *     .credential(new DefaultAzureCredentialBuilder().build())
+ *     .buildAsyncClient();
+ *
+ * client.semanticRerank("What is the capital of France?", documents, options)
+ *     .subscribe(result -&gt; {
+ *         result.getScores().forEach(s -&gt; System.out.println(s.getIndex() + ": " + s.getScore()));
+ *     });
+ * </pre>
+ */
+public final class CosmosAIAsyncClient implements AutoCloseable {
+
+    private final InferenceService inferenceService;
+
+    /**
+     * Creates a new CosmosAIAsyncClient.
+     *
+     * @param inferenceService The inference service instance.
+     */
+    CosmosAIAsyncClient(InferenceService inferenceService) {
+        this.inferenceService = inferenceService;
+    }
+
+    /**
+     * Performs semantic reranking of documents using the inference service.
+     *
+     * @param rerankContext The query or context string used to score documents.
+     * @param documents The list of document strings to rerank (must be non-null and non-empty).
+     * @param options Optional reranking parameters as a {@code Map<String, Object>}.
+     *                Supported keys:
+     *                <ul>
+     *                  <li>{@code "return_documents"} (Boolean) — include document text in the response</li>
+     *                  <li>{@code "top_k"} (Integer) — maximum number of results to return</li>
+     *                  <li>{@code "batch_size"} (Integer) — documents per inference batch</li>
+     *                  <li>{@code "sort"} (Boolean) — sort results by relevance score</li>
+     *                  <li>{@code "document_type"} (String) — {@code "string"} or {@code "json"}</li>
+     *                  <li>{@code "target_paths"} (String) — JSON paths for extraction when
+     *                      {@code document_type} is {@code "json"}</li>
+     *                  <li>{@code "timeout_seconds"} (Number) — per-request timeout override</li>
+     *                </ul>
+     * @return A {@link Mono} emitting the {@link SemanticRerankResult}.
+     */
+    public Mono<SemanticRerankResult> semanticRerank(
+        String rerankContext,
+        List<String> documents,
+        Map<String, Object> options) {
+
+        return inferenceService.semanticRerank(rerankContext, documents, options);
+    }
+
+    /**
+     * Closes this client and releases resources.
+     */
+    @Override
+    public void close() {
+        inferenceService.close();
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClient.java b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClient.java
new file mode 100644
index 000000000000..a7abcae93a15
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClient.java
@@ -0,0 +1,90 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai;
+
+import com.azure.cosmos.ai.implementation.InferenceService;
+import com.azure.cosmos.ai.models.SemanticRerankResult;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Synchronous client for Azure Cosmos DB AI services.
+ *
+ * <p>Provides semantic reranking of documents using the Azure Cosmos DB inference service.
+ * Use {@link CosmosAIClientBuilder} to create an instance.</p>
+ *
+ * <p><strong>Example:</strong></p>
+ * <pre>
+ * CosmosAIClient client = new CosmosAIClientBuilder()
+ *     .endpoint("https://my-inference.dbinference.azure.com")
+ *     .credential(new DefaultAzureCredentialBuilder().build())
+ *     .buildClient();
+ *
+ * SemanticRerankResult result = client.semanticRerank("What is the capital of France?", documents, options);
+ * result.getScores().forEach(s -&gt; System.out.println(s.getIndex() + ": " + s.getScore()));
+ * </pre>
+ */
+public final class CosmosAIClient implements AutoCloseable {
+
+    private final CosmosAIAsyncClient asyncClient;
+
+    /**
+     * Creates a new CosmosAIClient wrapping the given async client.
+     *
+     * @param asyncClient The async client to wrap.
+     */
+    CosmosAIClient(CosmosAIAsyncClient asyncClient) {
+        this.asyncClient = asyncClient;
+    }
+
+    /**
+     * Performs semantic reranking of documents using the inference service.
+     *
+     * <p><strong>Timeout:</strong> This method blocks with a default timeout of 120 seconds.
+     * To override, pass {@code "timeout_seconds"} (as a {@link Number}) in the {@code options} map,
+     * for example: {@code options.put("timeout_seconds", 30)}.</p>
+     *
+     * @param rerankContext The query or context string used to score documents.
+     * @param documents The list of document strings to rerank.
+     * @param options Optional reranking parameters as a map. Supported keys include:
+     *                {@code return_documents} (Boolean) - whether to return document text in the response,
+     *                {@code top_k} (Integer) - maximum number of documents to return,
+     *                {@code batch_size} (Integer) - number of documents per batch,
+     *                {@code sort} (Boolean) - whether to sort results by relevance score,
+     *                {@code timeout_seconds} (Number) - per-request timeout in seconds (default: 120).
+     * @return The semantic rerank result.
+     * @throws RuntimeException if the operation fails or times out.
+     */
+    public SemanticRerankResult semanticRerank(
+        String rerankContext,
+        List<String> documents,
+        Map<String, Object> options) {
+
+        Duration blockTimeout = resolveBlockTimeout(options);
+        return asyncClient.semanticRerank(rerankContext, documents, options)
+            .block(blockTimeout);
+    }
+
+    /**
+     * Closes this client and releases resources.
+     */
+    @Override
+    public void close() {
+        asyncClient.close();
+    }
+
+    private static Duration resolveBlockTimeout(Map<String, Object> options) {
+        if (options != null) {
+            Object timeoutVal = options.get(InferenceService.OPTION_TIMEOUT_SECONDS);
+            if (timeoutVal instanceof Number) {
+                double seconds = ((Number) timeoutVal).doubleValue();
+                if (seconds > 0) {
+                    return Duration.ofMillis((long) (seconds * 1000));
+                }
+            }
+        }
+        return InferenceService.DEFAULT_REQUEST_TIMEOUT;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClientBuilder.java b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClientBuilder.java
new file mode 100644
index 000000000000..6d527cb123d9
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClientBuilder.java
@@ -0,0 +1,207 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai;
+
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.http.HttpClient;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.policy.BearerTokenAuthenticationPolicy;
+import com.azure.core.http.policy.HttpLogDetailLevel;
+import com.azure.core.http.policy.HttpLogOptions;
+import com.azure.core.http.policy.HttpLoggingPolicy;
+import com.azure.core.http.policy.HttpPipelinePolicy;
+import com.azure.core.http.policy.RequestIdPolicy;
+import com.azure.core.http.policy.UserAgentPolicy;
+import com.azure.cosmos.ai.implementation.InferenceService;
+
+import java.net.URI;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Builder for creating {@link CosmosAIAsyncClient} and {@link CosmosAIClient} instances.
+ *
+ * <p>Minimum configuration requires {@link #endpoint(String)} and {@link #credential(TokenCredential)}.
+ * The endpoint can also be supplied via the system property {@code azure.cosmos.semanticReranker.inferenceEndpoint}
+ * or the environment variable {@code AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT}.</p>
+ *
+ * <p><strong>Example:</strong></p>
+ * <pre>
+ * CosmosAIAsyncClient asyncClient = new CosmosAIClientBuilder()
+ *     .endpoint("https://my-inference.dbinference.azure.com")
+ *     .credential(new DefaultAzureCredentialBuilder().build())
+ *     .buildAsyncClient();
+ * </pre>
+ */
+public final class CosmosAIClientBuilder {
+
+    private static final String INFERENCE_SCOPE = "https://dbinference.azure.com/.default";
+    private static final String SDK_NAME = "azure-cosmos-ai";
+    private static final String SDK_VERSION = "1.0.0-beta.1";
+
+    private static final String INFERENCE_ENDPOINT_PROPERTY = "azure.cosmos.semanticReranker.inferenceEndpoint";
+    private static final String INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT";
+
+    private String endpoint;
+    private TokenCredential credential;
+    private HttpPipeline httpPipeline;
+    private HttpClient httpClient;
+    private HttpLogOptions httpLogOptions;
+    private final List<HttpPipelinePolicy> additionalPolicies = new ArrayList<>();
+
+    /**
+     * Creates a new instance of CosmosAIClientBuilder.
+     */
+    public CosmosAIClientBuilder() {
+    }
+
+    /**
+     * Sets the inference service endpoint.
+     *
+     * <p>If not set, the builder falls back to the system property
+     * {@code azure.cosmos.semanticReranker.inferenceEndpoint} or the environment variable
+     * {@code AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT}.</p>
+     *
+     * @param endpoint The inference service endpoint URL.
+     * @return this builder.
+     */
+    public CosmosAIClientBuilder endpoint(String endpoint) {
+        this.endpoint = endpoint;
+        return this;
+    }
+
+    /**
+     * Sets the Azure AD token credential for authentication.
+     *
+     * @param credential The token credential.
+     * @return this builder.
+     */
+    public CosmosAIClientBuilder credential(TokenCredential credential) {
+        this.credential = credential;
+        return this;
+    }
+
+    /**
+     * Sets a custom HTTP pipeline.
+     *
+     * <p>When a pipeline is provided, the builder ignores {@link #credential(TokenCredential)},
+     * {@link #httpClient(HttpClient)}, and any additional policies — the caller is fully in
+     * control of the pipeline configuration.</p>
+     *
+     * @param httpPipeline The HTTP pipeline.
+     * @return this builder.
+     */
+    public CosmosAIClientBuilder pipeline(HttpPipeline httpPipeline) {
+        this.httpPipeline = httpPipeline;
+        return this;
+    }
+
+    /**
+     * Sets the HTTP client implementation to use.
+     *
+     * @param httpClient The HTTP client.
+     * @return this builder.
+     */
+    public CosmosAIClientBuilder httpClient(HttpClient httpClient) {
+        this.httpClient = httpClient;
+        return this;
+    }
+
+    /**
+     * Sets the HTTP log options for request and response logging.
+     *
+     * @param httpLogOptions The HTTP log options.
+     * @return this builder.
+     */
+    public CosmosAIClientBuilder httpLogOptions(HttpLogOptions httpLogOptions) {
+        this.httpLogOptions = httpLogOptions;
+        return this;
+    }
+
+    /**
+     * Adds an additional pipeline policy.
+     *
+     * @param policy The policy to add.
+     * @return this builder.
+     */
+    public CosmosAIClientBuilder addPolicy(HttpPipelinePolicy policy) {
+        Objects.requireNonNull(policy, "'policy' must not be null.");
+        this.additionalPolicies.add(policy);
+        return this;
+    }
+
+    /**
+     * Builds a new {@link CosmosAIAsyncClient} instance.
+     *
+     * @return A new asynchronous client.
+     * @throws IllegalStateException if the endpoint cannot be resolved or credential is missing.
+     */
+    public CosmosAIAsyncClient buildAsyncClient() {
+        URI resolvedEndpoint = resolveEndpoint();
+        HttpPipeline pipeline = buildPipeline();
+        InferenceService inferenceService = new InferenceService(resolvedEndpoint, pipeline);
+        return new CosmosAIAsyncClient(inferenceService);
+    }
+
+    /**
+     * Builds a new {@link CosmosAIClient} instance.
+     *
+     * @return A new synchronous client.
+     * @throws IllegalStateException if the endpoint cannot be resolved or credential is missing.
+     */
+    public CosmosAIClient buildClient() {
+        return new CosmosAIClient(buildAsyncClient());
+    }
+
+    private URI resolveEndpoint() {
+        if (endpoint != null && !endpoint.trim().isEmpty()) {
+            return URI.create(endpoint);
+        }
+
+        String fromProperty = System.getProperty(INFERENCE_ENDPOINT_PROPERTY);
+        if (fromProperty != null && !fromProperty.trim().isEmpty()) {
+            return URI.create(fromProperty);
+        }
+
+        String fromEnv = System.getenv(INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE);
+        if (fromEnv != null && !fromEnv.trim().isEmpty()) {
+            return URI.create(fromEnv);
+        }
+
+        throw new IllegalStateException(
+            "Inference endpoint must be set via .endpoint(), system property '"
+                + INFERENCE_ENDPOINT_PROPERTY + "', or environment variable '"
+                + INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE + "'.");
+    }
+
+    private HttpPipeline buildPipeline() {
+        if (this.httpPipeline != null) {
+            return this.httpPipeline;
+        }
+
+        if (this.credential == null) {
+            throw new IllegalStateException(
+                "Semantic reranking requires AAD authentication. "
+                    + "Provide a TokenCredential via .credential() or a pre-built pipeline via .pipeline().");
+        }
+
+        List<HttpPipelinePolicy> policies = new ArrayList<>();
+        policies.add(new UserAgentPolicy(null, SDK_NAME, SDK_VERSION, null));
+        policies.add(new RequestIdPolicy());
+        policies.add(new BearerTokenAuthenticationPolicy(this.credential, INFERENCE_SCOPE));
+        policies.addAll(this.additionalPolicies);
+        policies.add(new HttpLoggingPolicy(
+            this.httpLogOptions != null ? this.httpLogOptions : new HttpLogOptions().setLogLevel(HttpLogDetailLevel.NONE)));
+
+        HttpPipelineBuilder pipelineBuilder = new HttpPipelineBuilder()
+            .policies(policies.toArray(new HttpPipelinePolicy[0]));
+
+        if (this.httpClient != null) {
+            pipelineBuilder.httpClient(this.httpClient);
+        }
+
+        return pipelineBuilder.build();
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/implementation/InferenceService.java b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/implementation/InferenceService.java
new file mode 100644
index 000000000000..b61d1c437008
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/implementation/InferenceService.java
@@ -0,0 +1,302 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai.implementation;
+
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.exception.HttpResponseException;
+import com.azure.cosmos.ai.models.InferenceResponseParser;
+import com.azure.cosmos.ai.models.SemanticRerankResult;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Scheduler;
+import reactor.core.scheduler.Schedulers;
+import reactor.util.retry.Retry;
+
+import java.io.IOException;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.ThreadLocalRandom;
+
+/**
+ * Internal service for semantic reranking operations using azure-core HTTP pipeline.
+ * <p>
+ * While this class is public, it is not part of our published public APIs.
+ * This is meant to be internally used only by our SDK.
+ */
+public class InferenceService implements AutoCloseable {
+    private static final Logger LOGGER = LoggerFactory.getLogger(InferenceService.class);
+    private static final String BASE_PATH = "/inference/semanticReranking";
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+    // Option keys that callers can pass in the options map
+    /** Timeout option key. */
+    public static final String OPTION_TIMEOUT_SECONDS = "timeout_seconds";
+    /** Return documents option key. */
+    public static final String OPTION_RETURN_DOCUMENTS = "return_documents";
+    /** Top K option key. */
+    public static final String OPTION_TOP_K = "top_k";
+    /** Batch size option key. */
+    public static final String OPTION_BATCH_SIZE = "batch_size";
+    /** Sort option key. */
+    public static final String OPTION_SORT = "sort";
+    /** Document type option key. */
+    public static final String OPTION_DOCUMENT_TYPE = "document_type";
+    /** Target paths option key. */
+    public static final String OPTION_TARGET_PATHS = "target_paths";
+
+    /**
+     * Default per-request timeout for semantic rerank calls (120 seconds).
+     */
+    public static final Duration DEFAULT_REQUEST_TIMEOUT = Duration.ofSeconds(120);
+
+    // Retry policy — matches Python: TOTAL_RETRIES=3, RETRY_BACKOFF_FACTOR=0.8s, RETRY_BACKOFF_MAX=120s
+    static final int RETRY_MAX_ATTEMPTS = 3;
+    static final Duration RETRY_INITIAL_BACKOFF = Duration.ofMillis(800);
+    static final Duration RETRY_MAX_BACKOFF = Duration.ofSeconds(120);
+
+    private static final Set<Integer> RETRYABLE_STATUS_CODES = Collections.unmodifiableSet(
+        new HashSet<>(Arrays.asList(429, 500, 502, 503))
+    );
+
+    private static final HttpHeaderName RETRY_AFTER_HEADER = HttpHeaderName.fromString("Retry-After");
+
+    private final URI inferenceEndpoint;
+    private final HttpPipeline httpPipeline;
+
+    // Package-private so tests can inject a VirtualTimeScheduler to make backoff delays deterministic
+    Scheduler retryScheduler = Schedulers.parallel();
+    // Package-private so tests can disable jitter for deterministic virtual-time delays
+    double retryJitter = 0.5;
+
+    /**
+     * Creates a new InferenceService instance.
+     *
+     * @param endpoint The inference service base endpoint.
+     * @param httpPipeline The HTTP pipeline (with auth, retry, logging policies).
+     * @throws NullPointerException if endpoint or httpPipeline is null.
+     */
+    public InferenceService(URI endpoint, HttpPipeline httpPipeline) {
+        Objects.requireNonNull(endpoint, "'endpoint' must not be null.");
+        Objects.requireNonNull(httpPipeline, "'httpPipeline' must not be null.");
+
+        this.inferenceEndpoint = URI.create(endpoint.toString() + BASE_PATH);
+        this.httpPipeline = httpPipeline;
+
+        if (LOGGER.isInfoEnabled()) {
+            LOGGER.info("InferenceService initialized with endpoint: {}", this.inferenceEndpoint);
+        }
+    }
+
+    /**
+     * Closes the InferenceService.
+     */
+    @Override
+    public void close() {
+        LOGGER.info("Shutting down InferenceService...");
+    }
+
+    /**
+     * Performs semantic reranking of documents.
+     *
+     * @param rerankContext The query or context string used to score documents.
+     * @param documents The list of document strings to rerank.
+     * @param options Optional reranking parameters.
+     * @return A Mono emitting the semantic rerank result.
+     */
+    public Mono<SemanticRerankResult> semanticRerank(
+        String rerankContext,
+        List<String> documents,
+        Map<String, Object> options) {
+
+        Objects.requireNonNull(rerankContext, "Rerank context cannot be null");
+        Objects.requireNonNull(documents, "Documents list cannot be null");
+
+        if (rerankContext.trim().isEmpty()) {
+            return Mono.error(new IllegalArgumentException("Rerank context cannot be empty"));
+        }
+        if (documents.isEmpty()) {
+            return Mono.error(new IllegalArgumentException("Documents list cannot be empty"));
+        }
+
+        final Duration requestTimeout = resolveRequestTimeout(options);
+
+        try {
+            String requestBody = buildRequestPayload(rerankContext, documents, options);
+
+            HttpRequest httpRequest = new HttpRequest(HttpMethod.POST, inferenceEndpoint.toURL())
+                .setBody(requestBody.getBytes(StandardCharsets.UTF_8))
+                .setHeader(HttpHeaderName.CONTENT_TYPE, "application/json");
+
+            if (LOGGER.isDebugEnabled()) {
+                LOGGER.debug("Sending semantic rerank request to: {} (timeout: {})", inferenceEndpoint, requestTimeout);
+            }
+
+            return httpPipeline.send(httpRequest)
+                .timeout(requestTimeout)
+                .flatMap(this::parseResponse)
+                .as(this::withRetry)
+                .doOnError(error -> LOGGER.error("Semantic rerank operation failed", error));
+
+        } catch (Exception e) {
+            LOGGER.error("Failed to create request", e);
+            return Mono.error(new IllegalArgumentException(
+                "Failed to create semantic rerank request: " + e.getMessage(), e));
+        }
+    }
+
+    private String buildRequestPayload(String rerankContext, List<String> documents,
+                                       Map<String, Object> options) throws IOException {
+        ObjectNode payload = OBJECT_MAPPER.createObjectNode();
+        payload.put("query", rerankContext);
+
+        ArrayNode documentsArray = payload.putArray("documents");
+        for (String doc : documents) {
+            documentsArray.add(doc);
+        }
+
+        if (options != null) {
+            options.forEach((key, value) -> {
+                // timeout_seconds is a SDK-local option — do not forward to the endpoint
+                if (value != null && !OPTION_TIMEOUT_SECONDS.equals(key)) {
+                    payload.set(key, OBJECT_MAPPER.valueToTree(value));
+                }
+            });
+        }
+
+        return OBJECT_MAPPER.writeValueAsString(payload);
+    }
+
+    /**
+     * Resolves the effective request timeout.
+     */
+    private static Duration resolveRequestTimeout(Map<String, Object> options) {
+        if (options != null) {
+            Object value = options.get(OPTION_TIMEOUT_SECONDS);
+            if (value instanceof Number) {
+                double seconds = ((Number) value).doubleValue();
+                if (seconds > 0) {
+                    return Duration.ofMillis((long) (seconds * 1000));
+                } else {
+                    LOGGER.warn("Invalid '{}' value: {}. Must be > 0. Using default timeout {}.",
+                        OPTION_TIMEOUT_SECONDS, seconds, DEFAULT_REQUEST_TIMEOUT);
+                }
+            }
+        }
+        return DEFAULT_REQUEST_TIMEOUT;
+    }
+
+    /**
+     * Wraps a Mono with retry logic for transient inference endpoint failures.
+     */
+    private <T> Mono<T> withRetry(Mono<T> source) {
+        return source.retryWhen(
+            Retry.from(retrySignals -> retrySignals.concatMap(signal -> {
+                Throwable error = signal.failure();
+                if (!(error instanceof HttpResponseException)) {
+                    return Mono.error(error);
+                }
+                HttpResponseException httpEx = (HttpResponseException) error;
+                com.azure.core.http.HttpResponse response = httpEx.getResponse();
+                int statusCode = response != null ? response.getStatusCode() : 0;
+                if (!RETRYABLE_STATUS_CODES.contains(statusCode)) {
+                    return Mono.error(error);
+                }
+                long attempt = signal.totalRetries();
+                if (attempt >= RETRY_MAX_ATTEMPTS) {
+                    return Mono.error(error);
+                }
+
+                Duration delay = computeRetryDelay(response, statusCode, attempt);
+                LOGGER.warn(
+                    "Semantic rerank transient failure (status={}, attempt={}/{}), retrying after {} ms...",
+                    statusCode, attempt + 1, RETRY_MAX_ATTEMPTS, delay.toMillis());
+                return Mono.delay(delay, retryScheduler);
+            }))
+        );
+    }
+
+    private Duration computeRetryDelay(HttpResponse response, int statusCode, long attempt) {
+        if (statusCode == 429 && response != null) {
+            Duration serverDelay = parseRetryAfterSeconds(response.getHeaders());
+            if (serverDelay != null) {
+                return serverDelay.compareTo(RETRY_MAX_BACKOFF) > 0 ? RETRY_MAX_BACKOFF : serverDelay;
+            }
+        }
+        return exponentialBackoffWithJitter(attempt);
+    }
+
+    private static Duration parseRetryAfterSeconds(HttpHeaders headers) {
+        if (headers == null) {
+            return null;
+        }
+        String value = headers.getValue(RETRY_AFTER_HEADER);
+        if (value == null || value.trim().isEmpty()) {
+            return null;
+        }
+        try {
+            long seconds = Long.parseLong(value.trim());
+            if (seconds < 0) {
+                return null;
+            }
+            return Duration.ofSeconds(seconds);
+        } catch (NumberFormatException ignored) {
+            return null;
+        }
+    }
+
+    private Duration exponentialBackoffWithJitter(long attempt) {
+        long baseMillis = RETRY_INITIAL_BACKOFF.toMillis() << Math.min(attempt, 30);
+        long cappedMillis = Math.min(baseMillis, RETRY_MAX_BACKOFF.toMillis());
+        long jitterRange = (long) (cappedMillis * retryJitter);
+        long jitter = jitterRange > 0
+            ? ThreadLocalRandom.current().nextLong(-jitterRange, jitterRange + 1)
+            : 0L;
+        return Duration.ofMillis(Math.max(0L, cappedMillis + jitter));
+    }
+
+    private Mono<SemanticRerankResult> parseResponse(HttpResponse response) {
+        int statusCode = response.getStatusCode();
+
+        return response.getBodyAsString()
+            .flatMap(bodyString -> {
+                if (statusCode >= 400) {
+                    LOGGER.error("Semantic rerank request failed with status {}: {}", statusCode, bodyString);
+                    return Mono.error(new HttpResponseException(
+                        String.format("Semantic rerank request failed with status %d: %s", statusCode, bodyString),
+                        response));
+                }
+
+                try {
+                    SemanticRerankResult result = InferenceResponseParser.parseRerankResponse(bodyString);
+
+                    if (LOGGER.isDebugEnabled()) {
+                        LOGGER.debug("Successfully parsed semantic rerank response with {} scores",
+                            result.getScores() != null ? result.getScores().size() : 0);
+                    }
+
+                    return Mono.just(result);
+                } catch (Exception e) {
+                    LOGGER.error("Failed to parse semantic rerank response", e);
+                    return Mono.error(new IllegalStateException(
+                        "Failed to parse semantic rerank response: " + e.getMessage(), e));
+                }
+            });
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/InferenceResponseParser.java b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/InferenceResponseParser.java
new file mode 100644
index 000000000000..88ed99097ac1
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/InferenceResponseParser.java
@@ -0,0 +1,84 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai.models;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Internal parser for semantic rerank responses.
+ * <p>
+ * This class lives in the {@code models} package so it can access package-private setters
+ * on {@link SemanticRerankResult} and {@link SemanticRerankScore}.
+ * <p>
+ * While this class is public, it is not part of our published public APIs.
+ * This is meant to be internally used only by our SDK.
+ */
+public final class InferenceResponseParser {
+
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+    private InferenceResponseParser() {
+    }
+
+    /**
+     * Parses the JSON response body into a {@link SemanticRerankResult}.
+     *
+     * @param responseBody the JSON response body string.
+     * @return the parsed {@link SemanticRerankResult}.
+     * @throws IOException if parsing fails.
+     */
+    public static SemanticRerankResult parseRerankResponse(String responseBody) throws IOException {
+        JsonNode rootNode = OBJECT_MAPPER.readTree(responseBody);
+        SemanticRerankResult result = new SemanticRerankResult();
+
+        // Parse scores
+        if (rootNode.has("Scores")) {
+            JsonNode scoresNode = rootNode.get("Scores");
+            List<SemanticRerankScore> scores = new ArrayList<>();
+
+            if (scoresNode.isArray()) {
+                for (JsonNode scoreNode : scoresNode) {
+                    SemanticRerankScore score = new SemanticRerankScore();
+                    JsonNode indexNode = scoreNode.get("index");
+                    JsonNode scoreValNode = scoreNode.get("score");
+                    if (indexNode != null) {
+                        score.setIndex(indexNode.asInt());
+                    }
+                    if (scoreValNode != null) {
+                        score.setScore(scoreValNode.asDouble());
+                    }
+                    if (scoreNode.has("document")) {
+                        score.setDocument(scoreNode.get("document").asText());
+                    }
+                    scores.add(score);
+                }
+            }
+            result.setScores(scores);
+        }
+
+        // Parse latency
+        if (rootNode.has("latency")) {
+            Map<String, Object> latency = new HashMap<>();
+            rootNode.get("latency").fields().forEachRemaining(
+                entry -> latency.put(entry.getKey(), entry.getValue().asDouble()));
+            result.setLatency(latency);
+        }
+
+        // Parse token usage
+        if (rootNode.has("token_usage")) {
+            Map<String, Object> tokenUsage = new HashMap<>();
+            rootNode.get("token_usage").fields().forEachRemaining(
+                entry -> tokenUsage.put(entry.getKey(), entry.getValue().asInt()));
+            result.setTokenUsage(tokenUsage);
+        }
+
+        return result;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/SemanticRerankResult.java b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/SemanticRerankResult.java
new file mode 100644
index 000000000000..1205856c765e
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/SemanticRerankResult.java
@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai.models;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Represents the result of a semantic rerank operation.
+ */
+public final class SemanticRerankResult {
+    private List<SemanticRerankScore> scores;
+    private Map<String, Object> latency;
+    private Map<String, Object> tokenUsage;
+
+    /**
+     * Creates a new instance of SemanticRerankResult.
+     */
+    public SemanticRerankResult() {
+    }
+
+    /**
+     * Gets the list of scored documents.
+     *
+     * @return the list of scored documents.
+     */
+    public List<SemanticRerankScore> getScores() {
+        return scores;
+    }
+
+    /**
+     * Gets the latency information for the operation as a map of metric names to values.
+     *
+     * @return the latency information map.
+     */
+    public Map<String, Object> getLatency() {
+        return latency;
+    }
+
+    /**
+     * Gets the token usage information for the operation as a map of metric names to values.
+     *
+     * @return the token usage information map.
+     */
+    public Map<String, Object> getTokenUsage() {
+        return tokenUsage;
+    }
+
+    // Package-private setters used by InferenceResponseParser in the same package
+    void setScores(List<SemanticRerankScore> scores) {
+        this.scores = scores;
+    }
+
+    void setLatency(Map<String, Object> latency) {
+        this.latency = latency;
+    }
+
+    void setTokenUsage(Map<String, Object> tokenUsage) {
+        this.tokenUsage = tokenUsage;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/SemanticRerankScore.java b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/SemanticRerankScore.java
new file mode 100644
index 000000000000..69ea72399dad
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/models/SemanticRerankScore.java
@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai.models;
+
+/**
+ * Represents a single scored document in the semantic rerank result.
+ */
+public final class SemanticRerankScore {
+    private int index;
+    private String document;
+    private double score;
+
+    /**
+     * Creates a new instance of SemanticRerankScore.
+     */
+    public SemanticRerankScore() {
+    }
+
+    /**
+     * Gets the index of the document in the original input list.
+     *
+     * @return the document index.
+     */
+    public int getIndex() {
+        return index;
+    }
+
+    /**
+     * Gets the document text (if returnDocuments was true in the request).
+     *
+     * @return the document text, or null if not included.
+     */
+    public String getDocument() {
+        return document;
+    }
+
+    /**
+     * Gets the semantic relevance score for this document.
+     *
+     * @return the relevance score.
+     */
+    public double getScore() {
+        return score;
+    }
+
+    // Package-private setters used by InferenceResponseParser in the same package
+    void setIndex(int index) {
+        this.index = index;
+    }
+
+    void setDocument(String document) {
+        this.document = document;
+    }
+
+    void setScore(double score) {
+        this.score = score;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/semanticrerank/SemanticRerankSample.java b/sdk/cosmos/azure-cosmos-ai/src/samples/java/com/azure/cosmos/ai/SemanticRerankSample.java
similarity index 70%
rename from sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/semanticrerank/SemanticRerankSample.java
rename to sdk/cosmos/azure-cosmos-ai/src/samples/java/com/azure/cosmos/ai/SemanticRerankSample.java
index ea47468f1cfe..51bb587a620a 100644
--- a/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/semanticrerank/SemanticRerankSample.java
+++ b/sdk/cosmos/azure-cosmos-ai/src/samples/java/com/azure/cosmos/ai/SemanticRerankSample.java
@@ -1,13 +1,10 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
-package com.azure.cosmos.semanticrerank;
+package com.azure.cosmos.ai;
 
 import com.azure.core.credential.TokenCredential;
-import com.azure.cosmos.CosmosAsyncClient;
-import com.azure.cosmos.CosmosAsyncContainer;
-import com.azure.cosmos.CosmosClientBuilder;
-import com.azure.cosmos.models.SemanticRerankResult;
-import com.azure.cosmos.models.SemanticRerankScore;
+import com.azure.cosmos.ai.models.SemanticRerankResult;
+import com.azure.cosmos.ai.models.SemanticRerankScore;
 
 import java.util.Arrays;
 import java.util.HashMap;
@@ -15,26 +12,24 @@
 import java.util.Map;
 
 /**
- * Sample demonstrating semantic rerank functionality in Azure Cosmos DB.
+ * Sample demonstrating semantic rerank functionality using Azure Cosmos DB AI SDK.
  * <p>
  * Prerequisites:
- * 1. Set environment variable {@code COSMOS_ENDPOINT} with your Cosmos DB account endpoint.
- * 2. Set environment variable {@code AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT} with
- *    the inference service endpoint.
- * 3. Use Azure AD authentication — provide a {@link TokenCredential} implementation.
+ * 1. Set environment variable {@code AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT} with
+ *    the inference service endpoint, or pass it to the builder via {@code .endpoint(...)}.
+ * 2. Use Azure AD authentication — provide a {@link TokenCredential} implementation.
  *    If you have {@code azure-identity} on your classpath, use:
  *    {@code new com.azure.identity.DefaultAzureCredentialBuilder().build()}
  * </p>
  */
 public class SemanticRerankSample {
 
+    /**
+     * Main entry point.
+     *
+     * @param args command line arguments (not used).
+     */
     public static void main(String[] args) {
-        String endpoint = System.getenv("COSMOS_ENDPOINT");
-        if (endpoint == null || endpoint.isEmpty()) {
-            System.err.println("Please set COSMOS_ENDPOINT environment variable");
-            return;
-        }
-
         String inferenceEndpoint = System.getenv("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT");
         if (inferenceEndpoint == null || inferenceEndpoint.isEmpty()) {
             System.err.println("Please set AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT environment variable");
@@ -49,17 +44,21 @@ public static void main(String[] args) {
     }
 
     /**
-     * Runs the semantic rerank demo using the provided client and a container in the given database.
+     * Runs the semantic rerank demo.
      *
-     * @param client       Cosmos async client built with AAD credentials.
-     * @param databaseName Name of the database.
-     * @param containerName Name of the container.
+     * @param endpoint The inference service endpoint.
+     * @param credential The Azure AD token credential.
      */
-    static void runSemanticRerankDemo(CosmosAsyncClient client, String databaseName, String containerName) {
+    static void runSemanticRerankDemo(String endpoint, TokenCredential credential) {
         System.out.println("Semantic Rerank Sample");
         System.out.println("======================\n");
 
-        CosmosAsyncContainer container = client.getDatabase(databaseName).getContainer(containerName);
+        // BEGIN: readme-sample-createClient
+        CosmosAIAsyncClient client = new CosmosAIClientBuilder()
+            .endpoint(endpoint)
+            .credential(credential)
+            .buildAsyncClient();
+        // END: readme-sample-createClient
 
         List<String> documents = Arrays.asList(
             "Berlin is the capital of Germany.",
@@ -79,20 +78,19 @@ static void runSemanticRerankDemo(CosmosAsyncClient client, String databaseName,
         System.out.println();
 
         // Basic rerank with default options
-        container.semanticRerank(rerankContext, documents, null)
+        client.semanticRerank(rerankContext, documents, null)
             .subscribe(
                 result -> printResults(result),
                 error -> System.err.println("Error: " + error.getMessage())
             );
 
-        // Rerank with custom options. Keys are passed through to the inference endpoint as-is.
-        // See the JavaDoc on CosmosAsyncContainer.semanticRerank for the list of supported keys.
+        // Rerank with custom options
         Map<String, Object> options = new HashMap<>();
         options.put("return_documents", true);
         options.put("top_k", 3);
         options.put("sort", true);
 
-        container.semanticRerank(rerankContext, documents, options)
+        client.semanticRerank(rerankContext, documents, options)
             .subscribe(
                 result -> printResults(result),
                 error -> System.err.println("Error: " + error.getMessage())
diff --git a/sdk/cosmos/azure-cosmos-ai/src/test/java/com/azure/cosmos/ai/implementation/InferenceServiceTest.java b/sdk/cosmos/azure-cosmos-ai/src/test/java/com/azure/cosmos/ai/implementation/InferenceServiceTest.java
new file mode 100644
index 000000000000..f979e5d7ed49
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/test/java/com/azure/cosmos/ai/implementation/InferenceServiceTest.java
@@ -0,0 +1,474 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai.implementation;
+
+import com.azure.core.exception.HttpResponseException;
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import org.mockito.ArgumentCaptor;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+import reactor.core.Exceptions;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import java.net.URI;
+import java.net.URL;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class InferenceServiceTest {
+
+    private static final String TEST_ENDPOINT = "https://test-inference.westus.dbinference.azure.com";
+    // Total calls expected for a fully-retried retryable failure: 1 initial + RETRY_MAX_ATTEMPTS retries
+    private static final int TOTAL_CALLS_AFTER_EXHAUSTION = 1 + InferenceService.RETRY_MAX_ATTEMPTS;
+
+    private com.azure.core.http.HttpClient mockAzureCoreHttpClient;
+    private HttpPipeline httpPipeline;
+
+    @BeforeMethod(groups = {"unit"})
+    public void setUp() {
+        mockAzureCoreHttpClient = mock(com.azure.core.http.HttpClient.class);
+        httpPipeline = new HttpPipelineBuilder()
+            .httpClient(mockAzureCoreHttpClient)
+            .build();
+    }
+
+    // -------------------------------------------------------------------------
+    // Constructor / validation tests
+    // -------------------------------------------------------------------------
+
+    @Test(groups = {"unit"})
+    public void constructorShouldThrowWhenEndpointIsNull() {
+        assertThatThrownBy(() -> new InferenceService(null, httpPipeline))
+            .isInstanceOf(NullPointerException.class)
+            .hasMessageContaining("endpoint");
+    }
+
+    @Test(groups = {"unit"})
+    public void constructorShouldThrowWhenPipelineIsNull() {
+        assertThatThrownBy(() -> new InferenceService(URI.create(TEST_ENDPOINT), null))
+            .isInstanceOf(NullPointerException.class)
+            .hasMessageContaining("httpPipeline");
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldThrowWhenRerankContextIsNull() {
+        InferenceService service = createService();
+
+        assertThatThrownBy(() -> service.semanticRerank(null, Arrays.asList("doc1", "doc2"), null))
+            .isInstanceOf(NullPointerException.class)
+            .hasMessageContaining("Rerank context cannot be null");
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldThrowWhenDocumentsIsNull() {
+        InferenceService service = createService();
+
+        assertThatThrownBy(() -> service.semanticRerank("query", null, null))
+            .isInstanceOf(NullPointerException.class)
+            .hasMessageContaining("Documents list cannot be null");
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldThrowWhenRerankContextIsEmpty() {
+        InferenceService service = createService();
+
+        StepVerifier.create(service.semanticRerank("   ", Arrays.asList("doc1", "doc2"), null))
+            .expectErrorMatches(error -> error instanceof IllegalArgumentException
+                && error.getMessage().contains("Rerank context cannot be empty"))
+            .verify();
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldThrowWhenDocumentsListIsEmpty() {
+        InferenceService service = createService();
+
+        StepVerifier.create(service.semanticRerank("query", Collections.emptyList(), null))
+            .expectErrorMatches(error -> error instanceof IllegalArgumentException
+                && error.getMessage().contains("Documents list cannot be empty"))
+            .verify();
+    }
+
+    // -------------------------------------------------------------------------
+    // Happy-path tests
+    // -------------------------------------------------------------------------
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldSucceedWithValidResponse() {
+        String responseBody = createSuccessResponseBody();
+        MockHttpResponse mockResponse = new MockHttpResponse(null, 200, responseBody);
+        when(mockAzureCoreHttpClient.send(any(HttpRequest.class), any()))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = createService();
+
+        List<String> documents = Arrays.asList(
+            "This is document 1",
+            "This is document 2",
+            "This is document 3"
+        );
+
+        StepVerifier.create(service.semanticRerank("search query", documents, null))
+            .assertNext(result -> {
+                assertThat(result).isNotNull();
+                assertThat(result.getScores()).isNotNull();
+                assertThat(result.getScores()).hasSize(3);
+                assertThat(result.getScores().get(0).getIndex()).isEqualTo(0);
+                assertThat(result.getScores().get(0).getScore()).isEqualTo(0.95);
+                assertThat(result.getScores().get(0).getDocument()).isEqualTo("This is document 1");
+                assertThat(result.getLatency()).isNotNull();
+                assertThat(result.getLatency().get("inference_time")).isEqualTo(0.5);
+                assertThat(result.getTokenUsage()).isNotNull();
+                assertThat(result.getTokenUsage().get("total_tokens")).isEqualTo(100);
+            })
+            .verifyComplete();
+
+        verify(mockAzureCoreHttpClient, times(1)).send(any(HttpRequest.class), any());
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldIncludeOptionsInRequest() {
+        String responseBody = createSuccessResponseBody();
+        MockHttpResponse mockResponse = new MockHttpResponse(null, 200, responseBody);
+
+        ArgumentCaptor<HttpRequest> requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+        when(mockAzureCoreHttpClient.send(requestCaptor.capture(), any()))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = createService();
+
+        List<String> documents = Arrays.asList("doc1", "doc2");
+        Map<String, Object> options = new HashMap<>();
+        options.put("return_documents", true);
+        options.put("top_k", 5);
+        options.put("batch_size", 16);
+        options.put("sort", true);
+
+        StepVerifier.create(service.semanticRerank("query", documents, options))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        HttpRequest capturedRequest = requestCaptor.getValue();
+        assertThat(capturedRequest).isNotNull();
+
+        String requestBody = capturedRequest.getBodyAsBinaryData().toString();
+        assertThat(requestBody).contains("\"return_documents\":true");
+        assertThat(requestBody).contains("\"top_k\":5");
+        assertThat(requestBody).contains("\"batch_size\":16");
+        assertThat(requestBody).contains("\"sort\":true");
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldNotForwardTimeoutSecondsToPayload() {
+        String responseBody = createSuccessResponseBody();
+        MockHttpResponse mockResponse = new MockHttpResponse(null, 200, responseBody);
+
+        ArgumentCaptor<HttpRequest> requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+        when(mockAzureCoreHttpClient.send(requestCaptor.capture(), any()))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = createService();
+
+        Map<String, Object> options = new HashMap<>();
+        options.put("top_k", 3);
+        options.put(InferenceService.OPTION_TIMEOUT_SECONDS, 30);
+
+        StepVerifier.create(service.semanticRerank("query", Collections.singletonList("doc1"), options))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        HttpRequest capturedRequest = requestCaptor.getValue();
+        String requestBody = capturedRequest.getBodyAsBinaryData().toString();
+        assertThat(requestBody).doesNotContain(InferenceService.OPTION_TIMEOUT_SECONDS);
+    }
+
+    // -------------------------------------------------------------------------
+    // Non-retryable error tests — expect exactly 1 HTTP call, immediate failure
+    // -------------------------------------------------------------------------
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldHandle403Forbidden() {
+        String errorBody = "{\"error\": \"Forbidden - insufficient permissions\"}";
+        MockHttpResponse mockResponse = new MockHttpResponse(null, 403, errorBody);
+        when(mockAzureCoreHttpClient.send(any(HttpRequest.class), any()))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = createService();
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .expectErrorMatches(error -> error instanceof HttpResponseException
+                && error.getMessage().contains("Semantic rerank request failed")
+                && error.getMessage().contains("Forbidden"))
+            .verify();
+
+        verify(mockAzureCoreHttpClient, times(1)).send(any(HttpRequest.class), any());
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldHandleMalformedResponse() {
+        String malformedBody = "{ invalid json }";
+        MockHttpResponse mockResponse = new MockHttpResponse(null, 200, malformedBody);
+        when(mockAzureCoreHttpClient.send(any(HttpRequest.class), any()))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = createService();
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .expectErrorMatches(error ->
+                error instanceof IllegalStateException
+                    && error.getMessage().contains("Failed to parse semantic rerank response"))
+            .verify();
+
+        verify(mockAzureCoreHttpClient, times(1)).send(any(HttpRequest.class), any());
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldHandleEmptyScoresResponse() {
+        String responseBody = "{\"Scores\": [], \"latency\": {}, \"token_usage\": {}}";
+        MockHttpResponse mockResponse = new MockHttpResponse(null, 200, responseBody);
+        when(mockAzureCoreHttpClient.send(any(HttpRequest.class), any()))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = createService();
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .assertNext(result -> {
+                assertThat(result).isNotNull();
+                assertThat(result.getScores()).isEmpty();
+            })
+            .verifyComplete();
+    }
+
+    // -------------------------------------------------------------------------
+    // Retry behaviour tests — use virtual time to skip backoff delays
+    // -------------------------------------------------------------------------
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldRetryOn429AndSucceedOnSecondAttempt() {
+        MockHttpResponse throttledResponse = new MockHttpResponse(null, 429, "{\"error\": \"Too Many Requests\"}");
+        MockHttpResponse successResponse = new MockHttpResponse(null, 200, createSuccessResponseBody());
+
+        when(mockAzureCoreHttpClient.send(any(HttpRequest.class), any()))
+            .thenReturn(Mono.just(throttledResponse))
+            .thenReturn(Mono.just(successResponse));
+
+        final InferenceService service = createService();
+
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        verify(mockAzureCoreHttpClient, times(2)).send(any(HttpRequest.class), any());
+    }
+
+    @Test(groups = {"unit"})
+    public void semanticRerankShouldRetryOn500AndSucceedOnThirdAttempt() {
+        MockHttpResponse errorResponse = new MockHttpResponse(null, 500, "{\"error\": \"Internal Server Error\"}");
+        MockHttpResponse successResponse = new MockHttpResponse(null, 200, createSuccessResponseBody());
+
+        when(mockAzureCoreHttpClient.send(any(HttpRequest.class), any()))
+            .thenReturn(Mono.just(errorResponse))
+            .thenReturn(Mono.just(errorResponse))
+            .thenReturn(Mono.just(successResponse));
+
+        final InferenceService service = createService();
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .assertNext(result -> assertThat(result).isNotNull())
+            .verifyComplete();
+
+        verify(mockAzureCoreHttpClient, times(3)).send(any(HttpRequest.class), any());
+    }
+
+    @DataProvider(name = "nonRetryableStatusCodes")
+    public Object[][] nonRetryableStatusCodeProvider() {
+        return new Object[][] {
+            {400, "Bad Request"},
+            {401, "Unauthorized"},
+            {403, "Forbidden"},
+            {404, "Not Found"},
+        };
+    }
+
+    @Test(groups = {"unit"}, dataProvider = "nonRetryableStatusCodes")
+    public void semanticRerankShouldNotRetryNonRetryableStatusCodes(int statusCode, String errorMessage) {
+        String errorBody = String.format("{\"error\": \"%s\"}", errorMessage);
+        MockHttpResponse mockResponse = new MockHttpResponse(null, statusCode, errorBody);
+        when(mockAzureCoreHttpClient.send(any(HttpRequest.class), any()))
+            .thenReturn(Mono.just(mockResponse));
+
+        InferenceService service = createService();
+
+        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
+            .expectErrorMatches(error -> error.getMessage().contains("Semantic rerank request failed"))
+            .verify();
+
+        verify(mockAzureCoreHttpClient, times(1)).send(any(HttpRequest.class), any());
+    }
+
+    @DataProvider(name = "retryableStatusCodes")
+    public Object[][] retryableStatusCodeProvider() {
+        return new Object[][] {
+            {429, "Too Many Requests"},
+            {500, "Internal Server Error"},
+            {502, "Bad Gateway"},
+            {503, "Service Unavailable"},
+        };
+    }
+
+    @Test(groups = {"unit"}, dataProvider = "retryableStatusCodes")
+    public void semanticRerankShouldRetryRetryableStatusCodesAndExhaust(int statusCode, String errorMessage) {
+        String errorBody = String.format("{\"error\": \"%s\"}", errorMessage);
+        MockHttpResponse mockResponse = new MockHttpResponse(null, statusCode, errorBody);
+        when(mockAzureCoreHttpClient.send(any(HttpRequest.class), any()))
+            .thenReturn(Mono.just(mockResponse));
+
+        final InferenceService service = createService();
+        StepVerifier.withVirtualTime(() -> {
+                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
+                service.retryJitter = 0.0;
+                return service.semanticRerank("query", Arrays.asList("doc1"), null);
+            })
+            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
+            .verify();
+
+        verify(mockAzureCoreHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
+            .send(any(HttpRequest.class), any());
+    }
+
+    // -------------------------------------------------------------------------
+    // Helper methods
+    // -------------------------------------------------------------------------
+
+    private InferenceService createService() {
+        return new InferenceService(URI.create(TEST_ENDPOINT), httpPipeline);
+    }
+
+    private boolean isRetryExhaustedWithMessage(Throwable error, String messageSubstring) {
+        if (!Exceptions.isRetryExhausted(error)) {
+            return error.getMessage() != null && error.getMessage().contains(messageSubstring);
+        }
+        Throwable cause = error.getCause();
+        while (cause != null) {
+            if (cause.getMessage() != null && cause.getMessage().contains(messageSubstring)) {
+                return true;
+            }
+            cause = cause.getCause();
+        }
+        return false;
+    }
+
+    private String createSuccessResponseBody() {
+        return "{"
+            + "\"Scores\": ["
+            + "  {\"index\": 0, \"score\": 0.95, \"document\": \"This is document 1\"},"
+            + "  {\"index\": 1, \"score\": 0.85, \"document\": \"This is document 2\"},"
+            + "  {\"index\": 2, \"score\": 0.75, \"document\": \"This is document 3\"}"
+            + "],"
+            + "\"latency\": {"
+            + "  \"data_preprocess_time\": 0.1,"
+            + "  \"inference_time\": 0.5,"
+            + "  \"postprocess_time\": 0.05"
+            + "},"
+            + "\"token_usage\": {"
+            + "  \"total_tokens\": 100"
+            + "}"
+            + "}";
+    }
+
+    /**
+     * Simple mock implementation of azure-core HttpResponse for testing.
+     */
+    private static class MockHttpResponse extends HttpResponse {
+        private final int statusCode;
+        private final String body;
+        private final HttpHeaders headers;
+
+        MockHttpResponse(HttpRequest request, int statusCode, String body) {
+            super(request);
+            this.statusCode = statusCode;
+            this.body = body;
+            this.headers = new HttpHeaders();
+        }
+
+        MockHttpResponse(HttpRequest request, int statusCode, String body, HttpHeaders headers) {
+            super(request);
+            this.statusCode = statusCode;
+            this.body = body;
+            this.headers = headers;
+        }
+
+        @Override
+        public int getStatusCode() {
+            return statusCode;
+        }
+
+        @Override
+        @Deprecated
+        public String getHeaderValue(String name) {
+            return headers.getValue(HttpHeaderName.fromString(name));
+        }
+
+        @Override
+        public String getHeaderValue(HttpHeaderName headerName) {
+            return headers.getValue(headerName);
+        }
+
+        @Override
+        public HttpHeaders getHeaders() {
+            return headers;
+        }
+
+        @Override
+        public reactor.core.publisher.Flux<ByteBuffer> getBody() {
+            return reactor.core.publisher.Flux.just(ByteBuffer.wrap(body.getBytes(StandardCharsets.UTF_8)));
+        }
+
+        @Override
+        public Mono<byte[]> getBodyAsByteArray() {
+            return Mono.just(body.getBytes(StandardCharsets.UTF_8));
+        }
+
+        @Override
+        public Mono<String> getBodyAsString() {
+            return Mono.just(body);
+        }
+
+        @Override
+        public Mono<String> getBodyAsString(Charset charset) {
+            return Mono.just(new String(body.getBytes(StandardCharsets.UTF_8), charset));
+        }
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
deleted file mode 100644
index 68dfee640d18..000000000000
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/inference/InferenceServiceTest.java
+++ /dev/null
@@ -1,628 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-package com.azure.cosmos.implementation.inference;
-
-import com.azure.core.credential.AccessToken;
-import com.azure.core.credential.TokenCredential;
-import com.azure.core.credential.TokenRequestContext;
-import com.azure.cosmos.implementation.BadRequestException;
-import com.azure.cosmos.implementation.HttpConstants;
-import com.azure.cosmos.implementation.http.HttpClient;
-import com.azure.cosmos.implementation.http.HttpClientConfig;
-import com.azure.cosmos.implementation.http.HttpRequest;
-import com.azure.cosmos.implementation.http.HttpResponse;
-import org.mockito.ArgumentCaptor;
-import org.mockito.MockedStatic;
-import org.mockito.Mockito;
-import org.testng.annotations.AfterMethod;
-import org.testng.annotations.BeforeMethod;
-import org.testng.annotations.DataProvider;
-import org.testng.annotations.Test;
-import reactor.core.Exceptions;
-import reactor.core.publisher.Mono;
-import reactor.test.StepVerifier;
-
-import java.nio.charset.StandardCharsets;
-import java.time.Duration;
-import java.time.OffsetDateTime;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.assertThatThrownBy;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.*;
-
-public class InferenceServiceTest {
-
-    private static final String MOCK_TOKEN = "mock-access-token";
-    private static final String TEST_ENDPOINT = "https://test-inference.westus.dbinference.azure.com";
-    private static final String INFERENCE_ENDPOINT_PROPERTY = "azure.cosmos.semanticReranker.inferenceEndpoint";
-
-    // Total calls expected for a fully-retried retryable failure: 1 initial + RETRY_MAX_ATTEMPTS retries
-    private static final int TOTAL_CALLS_AFTER_EXHAUSTION = 1 + InferenceService.RETRY_MAX_ATTEMPTS;
-
-    private TokenCredential mockTokenCredential;
-    private HttpClient mockHttpClient;
-    private MockedStatic<HttpClient> httpClientStaticMock;
-
-    @BeforeMethod(groups = {"unit"})
-    public void setUp() {
-        // Set the inference endpoint system property required by InferenceService constructor
-        System.setProperty(INFERENCE_ENDPOINT_PROPERTY, TEST_ENDPOINT);
-
-        mockTokenCredential = mock(TokenCredential.class);
-        mockHttpClient = mock(HttpClient.class);
-
-        // Mock token credential to return a valid token
-        AccessToken accessToken = new AccessToken(MOCK_TOKEN, OffsetDateTime.now().plusHours(1));
-        when(mockTokenCredential.getToken(any(TokenRequestContext.class)))
-            .thenReturn(Mono.just(accessToken));
-
-        // Mock the static HttpClient.createFixed method
-        httpClientStaticMock = Mockito.mockStatic(HttpClient.class);
-        httpClientStaticMock.when(() -> HttpClient.createFixed(any(HttpClientConfig.class)))
-            .thenReturn(mockHttpClient);
-    }
-
-    @AfterMethod(groups = {"unit"})
-    public void tearDown() {
-        System.clearProperty(INFERENCE_ENDPOINT_PROPERTY);
-        if (httpClientStaticMock != null) {
-            httpClientStaticMock.close();
-        }
-    }
-
-    // -------------------------------------------------------------------------
-    // Constructor / validation tests (unchanged by retry logic)
-    // -------------------------------------------------------------------------
-
-    @Test(groups = {"unit"})
-    public void constructorShouldThrowWhenTokenCredentialIsNull() {
-        assertThatThrownBy(() -> new InferenceService(null))
-            .isInstanceOf(NullPointerException.class)
-            .hasMessageContaining("Semantic reranking requires AAD authentication")
-            .hasMessageContaining("key-based auth");
-    }
-
-    @Test(groups = {"unit"})
-    public void constructorShouldThrowWhenEndpointNotConfigured() {
-        System.clearProperty(INFERENCE_ENDPOINT_PROPERTY);
-
-        assertThatThrownBy(() -> new InferenceService(mockTokenCredential))
-            .isInstanceOf(IllegalArgumentException.class)
-            .hasMessageContaining("must be set to use semantic reranking");
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldThrowWhenRerankContextIsNull() {
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        assertThatThrownBy(() -> service.semanticRerank(null, Arrays.asList("doc1", "doc2"), null))
-            .isInstanceOf(NullPointerException.class)
-            .hasMessageContaining("Rerank context cannot be null");
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldThrowWhenDocumentsIsNull() {
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        assertThatThrownBy(() -> service.semanticRerank("query", null, null))
-            .isInstanceOf(NullPointerException.class)
-            .hasMessageContaining("Documents list cannot be null");
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldThrowWhenRerankContextIsEmpty() {
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("   ", Arrays.asList("doc1", "doc2"), null))
-            .expectErrorMatches(error -> error instanceof IllegalArgumentException
-                && error.getMessage().contains("Rerank context cannot be empty"))
-            .verify();
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldThrowWhenDocumentsListIsEmpty() {
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("query", Collections.emptyList(), null))
-            .expectErrorMatches(error -> error instanceof IllegalArgumentException
-                && error.getMessage().contains("Documents list cannot be empty"))
-            .verify();
-    }
-
-    // -------------------------------------------------------------------------
-    // Happy-path tests (unchanged by retry logic)
-    // -------------------------------------------------------------------------
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldSucceedWithValidResponse() {
-        // Create a mock response
-        String responseBody = createSuccessResponseBody();
-        HttpResponse mockResponse = createMockResponse(200, responseBody);
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        List<String> documents = Arrays.asList(
-            "This is document 1",
-            "This is document 2",
-            "This is document 3"
-        );
-
-        StepVerifier.create(service.semanticRerank("search query", documents, null))
-            .assertNext(result -> {
-                assertThat(result).isNotNull();
-                assertThat(result.getScores()).isNotNull();
-                assertThat(result.getScores()).hasSize(3);
-                assertThat(result.getScores().get(0).getIndex()).isEqualTo(0);
-                assertThat(result.getScores().get(0).getScore()).isEqualTo(0.95);
-                assertThat(result.getScores().get(0).getDocument()).isEqualTo("This is document 1");
-                assertThat(result.getLatency()).isNotNull();
-                assertThat(result.getLatency().get("inference_time")).isEqualTo(0.5);
-                assertThat(result.getTokenUsage()).isNotNull();
-                assertThat(result.getTokenUsage().get("total_tokens")).isEqualTo(100);
-            })
-            .verifyComplete();
-
-        // Success on first attempt — exactly 1 call
-        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldIncludeOptionsInRequest() {
-        String responseBody = createSuccessResponseBody();
-        HttpResponse mockResponse = createMockResponse(200, responseBody);
-
-        ArgumentCaptor<HttpRequest> requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
-        when(mockHttpClient.send(requestCaptor.capture(), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        List<String> documents = Arrays.asList("doc1", "doc2");
-        Map<String, Object> options = new HashMap<>();
-        options.put("return_documents", true);
-        options.put("top_k", 5);
-        options.put("batch_size", 16);
-        options.put("sort", true);
-
-        StepVerifier.create(service.semanticRerank("query", documents, options))
-            .assertNext(result -> assertThat(result).isNotNull())
-            .verifyComplete();
-
-        // Verify request body contains options
-        HttpRequest capturedRequest = requestCaptor.getValue();
-        assertThat(capturedRequest).isNotNull();
-
-        // Collect the Flux<byte[]> body into a single string
-        String requestBody = capturedRequest.body()
-            .map(bytes -> new String(bytes, StandardCharsets.UTF_8))
-            .collectList()
-            .map(list -> String.join("", list))
-            .block();
-
-        assertThat(requestBody).contains("\"return_documents\":true");
-        assertThat(requestBody).contains("\"top_k\":5");
-        assertThat(requestBody).contains("\"batch_size\":16");
-        assertThat(requestBody).contains("\"sort\":true");
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldNotForwardTimeoutSecondsToPayload() {
-        String responseBody = createSuccessResponseBody();
-        HttpResponse mockResponse = createMockResponse(200, responseBody);
-
-        ArgumentCaptor<HttpRequest> requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
-        when(mockHttpClient.send(requestCaptor.capture(), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        Map<String, Object> options = new HashMap<>();
-        options.put("top_k", 3);
-        options.put(InferenceService.OPTION_TIMEOUT_SECONDS, 30); // SDK-local, must NOT reach the endpoint
-
-        StepVerifier.create(service.semanticRerank("query", Collections.singletonList("doc1"), options))
-            .assertNext(result -> assertThat(result).isNotNull())
-            .verifyComplete();
-
-        // Verify that the request body does not contain the timeout seconds
-        HttpRequest capturedRequest = requestCaptor.getValue();
-        String requestBody = capturedRequest.body()
-            .map(bytes -> new String(bytes, StandardCharsets.UTF_8))
-            .collectList()
-            .map(list -> String.join("", list))
-            .block();
-
-        assertThat(requestBody).doesNotContain(InferenceService.OPTION_TIMEOUT_SECONDS);
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldIncludeAuthorizationHeader() {
-        String responseBody = createSuccessResponseBody();
-        HttpResponse mockResponse = createMockResponse(200, responseBody);
-
-        ArgumentCaptor<HttpRequest> requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
-        when(mockHttpClient.send(requestCaptor.capture(), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
-            .assertNext(result -> assertThat(result).isNotNull())
-            .verifyComplete();
-
-        HttpRequest capturedRequest = requestCaptor.getValue();
-        assertThat(capturedRequest.headers().value("Authorization")).isEqualTo("Bearer " + MOCK_TOKEN);
-    }
-
-    // -------------------------------------------------------------------------
-    // Non-retryable error tests — expect exactly 1 HTTP call, immediate failure
-    // -------------------------------------------------------------------------
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldHandle403Forbidden() {
-        // 403 is non-retryable — should fail immediately with 1 call
-        String errorBody = "{\"error\": \"Forbidden - insufficient permissions\"}";
-        HttpResponse mockResponse = createMockResponse(403, errorBody);
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
-            .expectErrorMatches(error -> error.getMessage().contains("Semantic rerank request failed")
-                && error.getMessage().contains("Forbidden"))
-            .verify();
-
-        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldHandleTokenAcquisitionFailure() {
-        // Token errors are not CosmosException — should not be retried
-        when(mockTokenCredential.getToken(any(TokenRequestContext.class)))
-            .thenReturn(Mono.error(new RuntimeException("Token acquisition failed")));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
-            .expectErrorMatches(error -> error instanceof RuntimeException
-                && error.getMessage().contains("Token acquisition failed"))
-            .verify();
-
-        // No HTTP call should be made if token acquisition fails
-        verify(mockHttpClient, never()).send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldNotRetryNetworkErrors() {
-        // Raw RuntimeException (not CosmosException) — not retryable
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.error(new RuntimeException("Network connection failed")));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
-            .expectErrorMatches(error -> error instanceof RuntimeException
-                && error.getMessage().contains("Network connection failed"))
-            .verify();
-
-        // Non-CosmosException errors bypass the retry filter — exactly 1 call
-        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldHandleMalformedResponse() {
-        // Parse failure uses BadRequestException (400) + CUSTOM_SERIALIZER_EXCEPTION sub-status,
-        // matching CosmosItemSerializer's pattern — not retryable, fails immediately with 1 call.
-        String malformedBody = "{ invalid json }";
-        HttpResponse mockResponse = createMockResponse(200, malformedBody);
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
-            .expectErrorMatches(error ->
-                error instanceof BadRequestException
-                    && error.getMessage().contains("Failed to parse semantic rerank response")
-                    && ((BadRequestException) error).getStatusCode() == HttpConstants.StatusCodes.BADREQUEST
-                    && ((BadRequestException) error).getSubStatusCode()
-                        == HttpConstants.SubStatusCodes.CUSTOM_SERIALIZER_EXCEPTION)
-            .verify();
-
-        // Not retryable — exactly 1 call
-        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldHandleEmptyScoresResponse() {
-        String responseBody = "{\"Scores\": [], \"latency\": {}, \"token_usage\": {}}";
-        HttpResponse mockResponse = createMockResponse(200, responseBody);
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
-            .assertNext(result -> {
-                assertThat(result).isNotNull();
-                assertThat(result.getScores()).isEmpty();
-            })
-            .verifyComplete();
-    }
-
-    // -------------------------------------------------------------------------
-    // Retry behaviour tests — use virtual time to skip backoff delays
-    // -------------------------------------------------------------------------
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldRetryOn429AndSucceedOnSecondAttempt() {
-        // First call returns 429, second returns 200
-        HttpResponse throttledResponse = createMockResponse(429, "{\"error\": \"Too Many Requests\"}");
-        HttpResponse successResponse = createMockResponse(200, createSuccessResponseBody());
-
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(throttledResponse))
-            .thenReturn(Mono.just(successResponse));
-
-        // Inject VirtualTimeScheduler so Retry.backoff delays run under virtual time.
-        // The Mono is returned inside the supplier so it is assembled after the
-        // virtual scheduler is installed.
-        final InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.withVirtualTime(() -> {
-                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
-                service.retryJitter = 0.0;
-                return service.semanticRerank("query", Arrays.asList("doc1"), null);
-            })
-            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .assertNext(result -> assertThat(result).isNotNull())
-            .verifyComplete();
-
-        // 1 initial + 1 retry = 2 total calls
-        verify(mockHttpClient, times(2)).send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldRetryOn500AndSucceedOnThirdAttempt() {
-        // First two calls return 500, third returns 200
-        HttpResponse errorResponse = createMockResponse(500, "{\"error\": \"Internal Server Error\"}");
-        HttpResponse successResponse = createMockResponse(200, createSuccessResponseBody());
-
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(errorResponse))
-            .thenReturn(Mono.just(errorResponse))
-            .thenReturn(Mono.just(successResponse));
-
-        final InferenceService service = new InferenceService(mockTokenCredential);
-        StepVerifier.withVirtualTime(() -> {
-                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
-                service.retryJitter = 0.0;
-                return service.semanticRerank("query", Arrays.asList("doc1"), null);
-            })
-            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .assertNext(result -> assertThat(result).isNotNull())
-            .verifyComplete();
-
-        // 1 initial + 2 retries = 3 total calls
-        verify(mockHttpClient, times(3)).send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldExhaustRetriesOn429AndFail() {
-        HttpResponse throttledResponse = createMockResponse(429, "{\"error\": \"Too Many Requests\"}");
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(throttledResponse));
-
-        final InferenceService service = new InferenceService(mockTokenCredential);
-        StepVerifier.withVirtualTime(() -> {
-                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
-                service.retryJitter = 0.0;
-                return service.semanticRerank("query", Arrays.asList("doc1"), null);
-            })
-            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
-            .verify();
-
-        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldExhaustRetriesOn500AndFail() {
-        HttpResponse errorResponse = createMockResponse(500, "{\"error\": \"Internal Server Error\"}");
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(errorResponse));
-
-        final InferenceService service = new InferenceService(mockTokenCredential);
-        StepVerifier.withVirtualTime(() -> {
-                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
-                service.retryJitter = 0.0;
-                return service.semanticRerank("query", Arrays.asList("doc1"), null);
-            })
-            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
-            .verify();
-
-        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldExhaustRetriesOn502AndFail() {
-        HttpResponse errorResponse = createMockResponse(502, "{\"error\": \"Bad Gateway\"}");
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(errorResponse));
-
-        final InferenceService service = new InferenceService(mockTokenCredential);
-        StepVerifier.withVirtualTime(() -> {
-                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
-                service.retryJitter = 0.0;
-                return service.semanticRerank("query", Arrays.asList("doc1"), null);
-            })
-            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
-            .verify();
-
-        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @Test(groups = {"unit"})
-    public void semanticRerankShouldExhaustRetriesOn503AndFail() {
-        HttpResponse errorResponse = createMockResponse(503, "{\"error\": \"Service Unavailable\"}");
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(errorResponse));
-
-        final InferenceService service = new InferenceService(mockTokenCredential);
-        StepVerifier.withVirtualTime(() -> {
-                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
-                service.retryJitter = 0.0;
-                return service.semanticRerank("query", Arrays.asList("doc1"), null);
-            })
-            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
-            .verify();
-
-        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    // -------------------------------------------------------------------------
-    // DataProvider — non-retryable codes fail immediately (1 call),
-    //                retryable codes exhaust retries (TOTAL_CALLS_AFTER_EXHAUSTION calls)
-    // -------------------------------------------------------------------------
-
-    @DataProvider(name = "nonRetryableStatusCodes")
-    public Object[][] nonRetryableStatusCodeProvider() {
-        return new Object[][] {
-            {400, "Bad Request"},
-            {401, "Unauthorized"},
-            {403, "Forbidden"},
-            {404, "Not Found"},
-        };
-    }
-
-    @Test(groups = {"unit"}, dataProvider = "nonRetryableStatusCodes")
-    public void semanticRerankShouldNotRetryNonRetryableStatusCodes(int statusCode, String errorMessage) {
-        String errorBody = String.format("{\"error\": \"%s\"}", errorMessage);
-        HttpResponse mockResponse = createMockResponse(statusCode, errorBody);
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        InferenceService service = new InferenceService(mockTokenCredential);
-
-        StepVerifier.create(service.semanticRerank("query", Arrays.asList("doc1"), null))
-            .expectErrorMatches(error -> containsMessage(error, "Semantic rerank request failed"))
-            .verify();
-
-        // Non-retryable — exactly 1 call regardless of status code
-        verify(mockHttpClient, times(1)).send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    @DataProvider(name = "retryableStatusCodes")
-    public Object[][] retryableStatusCodeProvider() {
-        return new Object[][] {
-            {429, "Too Many Requests"},
-            {500, "Internal Server Error"},
-            {502, "Bad Gateway"},
-            {503, "Service Unavailable"},
-        };
-    }
-
-    @Test(groups = {"unit"}, dataProvider = "retryableStatusCodes")
-    public void semanticRerankShouldRetryRetryableStatusCodesAndExhaust(int statusCode, String errorMessage) {
-        String errorBody = String.format("{\"error\": \"%s\"}", errorMessage);
-        HttpResponse mockResponse = createMockResponse(statusCode, errorBody);
-        when(mockHttpClient.send(any(HttpRequest.class), any(Duration.class)))
-            .thenReturn(Mono.just(mockResponse));
-
-        final InferenceService service = new InferenceService(mockTokenCredential);
-        StepVerifier.withVirtualTime(() -> {
-                service.retryScheduler = reactor.test.scheduler.VirtualTimeScheduler.get();
-                service.retryJitter = 0.0;
-                return service.semanticRerank("query", Arrays.asList("doc1"), null);
-            })
-            .thenAwait(InferenceService.RETRY_MAX_BACKOFF.multipliedBy(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .expectErrorMatches(error -> isRetryExhaustedWithMessage(error, "Semantic rerank request failed"))
-            .verify();
-
-        // 1 initial attempt + RETRY_MAX_ATTEMPTS retries
-        verify(mockHttpClient, times(TOTAL_CALLS_AFTER_EXHAUSTION))
-            .send(any(HttpRequest.class), any(Duration.class));
-    }
-
-    // -------------------------------------------------------------------------
-    // Helper methods
-    // -------------------------------------------------------------------------
-
-    /**
-     * Checks that an error represents exhausted retries on a retryable inference failure.
-     *
-     * When {@code Retry.backoff} exhausts all attempts it wraps the last error in a
-     * {@code RetryExhaustedException}. {@link com.azure.cosmos.CosmosException#getMessage()}
-     * returns a JSON diagnostic blob, not the plain constructor message, so we use
-     * {@link com.azure.cosmos.CosmosException#getShortMessage()} for plain-string matching.
-     * As a fallback we also walk the cause chain.
-     */
-    private boolean isRetryExhaustedWithMessage(Throwable error, String messageSubstring) {
-        if (!Exceptions.isRetryExhausted(error)) {
-            // Not wrapped — check the error itself
-            return containsMessage(error, messageSubstring);
-        }
-        // Walk the cause chain from the RetryExhaustedException
-        Throwable cause = error.getCause();
-        while (cause != null) {
-            if (containsMessage(cause, messageSubstring)) {
-                return true;
-            }
-            cause = cause.getCause();
-        }
-        return false;
-    }
-
-    private boolean containsMessage(Throwable t, String substring) {
-        if (t instanceof com.azure.cosmos.CosmosException) {
-            // getShortMessage() returns the plain constructor message without JSON wrapping
-            String short_ = ((com.azure.cosmos.CosmosException) t).getShortMessage();
-            if (short_ != null && short_.contains(substring)) {
-                return true;
-            }
-        }
-        return t.getMessage() != null && t.getMessage().contains(substring);
-    }
-
-    private String createSuccessResponseBody() {
-        return "{"
-            + "\"Scores\": ["
-            + "  {\"index\": 0, \"score\": 0.95, \"document\": \"This is document 1\"},"
-            + "  {\"index\": 1, \"score\": 0.85, \"document\": \"This is document 2\"},"
-            + "  {\"index\": 2, \"score\": 0.75, \"document\": \"This is document 3\"}"
-            + "],"
-            + "\"latency\": {"
-            + "  \"data_preprocess_time\": 0.1,"
-            + "  \"inference_time\": 0.5,"
-            + "  \"postprocess_time\": 0.05"
-            + "},"
-            + "\"token_usage\": {"
-            + "  \"total_tokens\": 100"
-            + "}"
-            + "}";
-    }
-
-    private HttpResponse createMockResponse(int statusCode, String body) {
-        HttpResponse mockResponse = mock(HttpResponse.class);
-        when(mockResponse.statusCode()).thenReturn(statusCode);
-        when(mockResponse.bodyAsString()).thenReturn(Mono.just(body));
-        when(mockResponse.headers()).thenReturn(new com.azure.cosmos.implementation.http.HttpHeaders());
-        return mockResponse;
-    }
-}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
index d52d4f3c44e9..30f14089e9b5 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java
@@ -22,7 +22,6 @@
 import com.azure.cosmos.implementation.Strings;
 import com.azure.cosmos.implementation.WriteRetryPolicy;
 import com.azure.cosmos.implementation.clienttelemetry.ClientMetricsDiagnosticsHandler;
-import com.azure.cosmos.implementation.inference.InferenceService;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetryMetrics;
 import com.azure.cosmos.implementation.clienttelemetry.CosmosMeterOptions;
@@ -62,7 +61,6 @@
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Objects;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -122,7 +120,6 @@ private static ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccess
     private final List<CosmosOperationPolicy> requestPolicies;
     private final CosmosItemSerializer defaultCustomSerializer;
     private final java.util.function.Function<CosmosAsyncContainer, CosmosAsyncContainer> containerFactory;
-    private final AtomicReference<InferenceService> inferenceService = new AtomicReference<>();
     private volatile boolean closed;
 
     CosmosAsyncClient(CosmosClientBuilder builder) {
@@ -306,42 +303,6 @@ TokenCredential tokenCredential() {
         return this.tokenCredential;
     }
 
-    /**
-     * Returns the shared InferenceService instance for this client, creating it lazily on first use.
-     * The instance is tied to this client's lifecycle and will be shut down in {@link #close()}.
-     *
-     * @throws IllegalStateException if the client was built with key-based auth instead of AAD.
-     */
-    InferenceService getOrCreateInferenceService() {
-        if (this.closed) {
-            throw new IllegalStateException("CosmosAsyncClient has been closed.");
-        }
-        if (this.tokenCredential == null) {
-            throw new IllegalStateException(
-                "Semantic reranking requires AAD authentication. "
-                    + "The CosmosClient was built with key-based auth (master key or AzureKeyCredential), "
-                    + "which is not supported for this operation. "
-                    + "Rebuild the client using .credential(TokenCredential) with a DefaultAzureCredential "
-                    + "or another TokenCredential implementation.");
-        }
-        if (this.inferenceService.get() == null) {
-            InferenceService newSvc = new InferenceService(this.tokenCredential);
-            if (!this.inferenceService.compareAndSet(null, newSvc)) {
-                // Another thread already set the instance; close the losing one to prevent resource leak
-                newSvc.close();
-            } else if (this.closed) {
-                // close() ran concurrently and already saw a null reference; ensure the late-arriving
-                // instance we just published is torn down so we don't leak the Netty event-loop group
-                // and pooled connections held by the underlying HttpClient.
-                InferenceService leaked = this.inferenceService.getAndSet(null);
-                if (leaked != null) {
-                    leaked.close();
-                }
-                throw new IllegalStateException("CosmosAsyncClient has been closed.");
-            }
-        }
-        return this.inferenceService.get();
-    }
 
     /***
      * Get the client telemetry config.
@@ -611,18 +572,10 @@ public CosmosAsyncDatabase getDatabase(String id) {
      */
     @Override
     public void close() {
-        // Set the closed flag BEFORE reading the inference reference. getOrCreateInferenceService()
-        // re-checks this flag after publishing a new instance and tears it down if it raced with us,
-        // so the combination of (closed = true) + getAndSet(null) here guarantees that any
-        // InferenceService that is ever published is also closed exactly once.
         this.closed = true;
         if (this.clientMetricRegistrySnapshot != null) {
             ClientTelemetryMetrics.remove(this.clientMetricRegistrySnapshot);
         }
-        InferenceService svc = this.inferenceService.getAndSet(null);
-        if (svc != null) {
-            svc.close();
-        }
         asyncDocumentClient.close();
     }
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
index 7f03ddd552b3..b30316a7ba8f 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java
@@ -35,7 +35,7 @@
 import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider;
 import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
 import com.azure.cosmos.implementation.feedranges.FeedRangeInternal;
-import com.azure.cosmos.implementation.inference.InferenceService;
+
 import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
 import com.azure.cosmos.implementation.routing.Range;
 import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup;
@@ -69,7 +69,7 @@
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.models.PartitionKeyDefinition;
-import com.azure.cosmos.models.SemanticRerankResult;
+
 import com.azure.cosmos.models.ShowQueryMode;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.models.ThroughputProperties;
@@ -1948,46 +1948,6 @@ public <T> CosmosPagedFlux<T> readAllItems(
         });
     }
 
-    /**
-     * Performs semantic reranking of documents using the Azure Cosmos DB inference service.
-     * <p>
-     * Requires the {@link CosmosAsyncClient} to have been built with AAD authentication
-     * ({@code .credential(TokenCredential)}). Key-based auth is not supported.
-     * <p>
-     * The request timeout defaults to 120 seconds. To override it, pass
-     * {@code "timeout_seconds"} (as a {@link Number}) in the {@code options} map.
-     * You can also apply {@code .timeout(Duration)} to the returned {@code Mono} directly.
-     *
-     * @param rerankContext The query or context string used to score documents.
-     * @param documents The list of document strings to rerank (must be non-null and non-empty).
-     * @param options Optional reranking parameters as a {@code Map<String, Object>}.
-     *                Supported keys:
-     *                <ul>
-     *                  <li>{@code "return_documents"} (Boolean) — include document text in the response</li>
-     *                  <li>{@code "top_k"} (Integer) — maximum number of results to return</li>
-     *                  <li>{@code "batch_size"} (Integer) — documents per inference batch</li>
-     *                  <li>{@code "sort"} (Boolean) — sort results by relevance score</li>
-     *                  <li>{@code "document_type"} (String) — {@code "string"} or {@code "json"}</li>
-     *                  <li>{@code "target_paths"} (String) — JSON paths for extraction when
-     *                      {@code document_type} is {@code "json"}</li>
-     *                  <li>{@code "timeout_seconds"} (Number) — per-request timeout override</li>
-     *                </ul>
-     * @return A {@link Mono} emitting the {@link SemanticRerankResult}.
-     */
-    @Beta(value = Beta.SinceVersion.V4_81_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
-    public Mono<SemanticRerankResult> semanticRerank(
-        String rerankContext,
-        List<String> documents,
-        Map<String, Object> options) {
-
-        checkNotNull(rerankContext, "Argument 'rerankContext' must not be null.");
-        checkNotNull(documents, "Argument 'documents' must not be null.");
-        checkArgument(!rerankContext.trim().isEmpty(), "Argument 'rerankContext' must not be empty.");
-        checkArgument(!documents.isEmpty(), "Argument 'documents' must not be empty.");
-
-        return this.database.getClient().getOrCreateInferenceService().semanticRerank(rerankContext, documents, options);
-    }
-
     /**
      * Replaces an existing item in a container with a new item.
      * It performs a complete replacement of the item,
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
index 368f4d43ca4f..30645812cfab 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
@@ -4,7 +4,6 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup;
-import com.azure.cosmos.implementation.inference.InferenceService;
 import com.azure.cosmos.implementation.RequestTimeoutException;
 import com.azure.cosmos.models.CosmosBatch;
 import com.azure.cosmos.models.CosmosBatchOperationResult;
@@ -28,7 +27,6 @@
 import com.azure.cosmos.models.FeedRange;
 import com.azure.cosmos.models.FeedResponse;
 import com.azure.cosmos.models.PartitionKey;
-import com.azure.cosmos.models.SemanticRerankResult;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.models.ThroughputProperties;
 import com.azure.cosmos.models.ThroughputResponse;
@@ -317,43 +315,6 @@ private CosmosItemResponse<Object> blockDeleteItemResponse(Mono<CosmosItemRespon
         }
     }
 
-    private SemanticRerankResult blockSemanticRerankResponse(
-        Mono<SemanticRerankResult> semanticRerankResultMono,
-        Map<String, Object> options) {
-        // Derive the blocking timeout from the same option key used by InferenceService,
-        // so the sync API respects timeout_seconds just like the async path.
-        // Default is InferenceService.DEFAULT_REQUEST_TIMEOUT (120 s) if not supplied.
-        Duration blockTimeout = InferenceService.DEFAULT_REQUEST_TIMEOUT;
-        if (options != null) {
-            Object timeoutVal = options.get(InferenceService.OPTION_TIMEOUT_SECONDS);
-            if (timeoutVal instanceof Number) {
-                double seconds = ((Number) timeoutVal).doubleValue();
-                if (seconds > 0) {
-                    blockTimeout = Duration.ofMillis((long) (seconds * 1000));
-                }
-            }
-        }
-        try {
-            // Apply the same timeout semantics on the sync wrapper that InferenceService applies per
-            // HTTP attempt: convert a timeout to a typed RequestTimeoutException (408) instead of
-            // letting Reactor's raw RuntimeException("Timeout on blocking read") leak to the caller.
-            final Duration timeout = blockTimeout;
-            return semanticRerankResultMono
-                .timeout(timeout)
-                .onErrorMap(java.util.concurrent.TimeoutException.class,
-                    t -> new RequestTimeoutException(
-                        "Semantic rerank request timed out after " + timeout, (java.net.URI) null))
-                .block(timeout);
-        } catch (Exception ex) {
-            final Throwable throwable = Exceptions.unwrap(ex);
-            if (throwable instanceof CosmosException) {
-                throw (CosmosException) throwable;
-            } else {
-                throw ex;
-            }
-        }
-    }
-
     private CosmosBatchResponse blockBatchResponse(Mono<CosmosBatchResponse> batchResponseMono) {
         try {
             return batchResponseMono.block();
@@ -1079,40 +1040,6 @@ public <TContext> Iterable<CosmosBulkOperationResponse<TContext>> executeBulkOpe
         return this.blockBulkResponse(asyncContainer.executeBulkOperations(Flux.fromIterable(operations), bulkOptions));
     }
 
-    /**
-     * Performs semantic reranking of documents using the inference service.
-     *
-     * <p><strong>Timeout:</strong> This method blocks with a default timeout of 120 seconds.
-     * To override, pass {@code "timeout_seconds"} (as a {@link Number}) in the {@code options} map,
-     * for example: {@code options.put("timeout_seconds", 30)}.
-     *
-     * <p><strong>Sync vs async timeout semantics:</strong> on the async API
-     * ({@code CosmosAsyncContainer.semanticRerank}) {@code timeout_seconds} is applied per HTTP attempt,
-     * so the total wall-clock can stretch across retries and backoff. On this sync API the same value
-     * is also applied as the {@link Mono#timeout(java.time.Duration) whole-operation deadline}
-     * around the inner {@link Mono}, so the call returns once the overall budget elapses regardless of
-     * how many retries the async path would otherwise perform. When the budget elapses the sync API
-     * surfaces a {@link CosmosException} with status {@code 408 Request Timeout} rather than a raw
-     * Reactor {@link RuntimeException}.
-     *
-     * @param rerankContext The query or context string used to score documents.
-     * @param documents The list of document strings to rerank.
-     * @param options Optional reranking parameters as a map. Supported keys include:
-     *                {@code return_documents} (Boolean) - whether to return document text in the response,
-     *                {@code top_k} (Integer) - maximum number of documents to return,
-     *                {@code batch_size} (Integer) - number of documents per batch,
-     *                {@code sort} (Boolean) - whether to sort results by relevance score,
-     *                {@code timeout_seconds} (Number) - per-request timeout in seconds (default: 120).
-     * @return The semantic rerank result.
-     */
-    @Beta(value = Beta.SinceVersion.V4_81_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
-    public SemanticRerankResult semanticRerank(
-        String rerankContext,
-        List<String> documents,
-        Map<String, Object> options) {
-        return blockSemanticRerankResponse(this.asyncContainer.semanticRerank(rerankContext, documents, options), options);
-    }
-
     /**
      * Gets the Cosmos scripts using the current container as context.
      *
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
index 32b7af3a3bc0..18eef0544e18 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java
@@ -421,12 +421,6 @@ public class Configs {
     private static final String CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT = "COSMOS.CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT";
     private static final int DEFAULT_CLIENT_ENDPOINT_FAILOVER_MAX_RETRY_COUNT = 120;
 
-    // Inference service related configs.
-    // Following the convention used elsewhere in this file: the system property uses dot.notation
-    // and the environment variable uses UPPER_SNAKE_CASE.
-    private static final String INFERENCE_ENDPOINT_PROPERTY = "azure.cosmos.semanticReranker.inferenceEndpoint";
-    private static final String INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT";
-
     private static final Object lockObject = new Object();
     private static Boolean cachedIsHostnameValidationDisabled = null;
 
@@ -537,19 +531,6 @@ public URI getThinclientEndpoint() {
         return URI.create(DEFAULT_THINCLIENT_ENDPOINT);
     }
 
-    public URI getInferenceServiceEndpoint() {
-        String valueFromSystemProperty = System.getProperty(INFERENCE_ENDPOINT_PROPERTY);
-        if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) {
-            return URI.create(valueFromSystemProperty);
-        }
-
-        String valueFromEnvVariable = System.getenv(INFERENCE_ENDPOINT_ENVIRONMENT_VARIABLE);
-        if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) {
-            return URI.create(valueFromEnvVariable);
-        }
-        return null;
-    }
-
     public static boolean isThinClientEnabled() {
         String valueFromSystemProperty = System.getProperty(THINCLIENT_ENABLED);
         if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) {
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java
index 38d19b112ac3..c48555496c29 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java
@@ -79,8 +79,6 @@
 import com.azure.cosmos.models.PriorityLevel;
 import com.azure.cosmos.models.ShowQueryMode;
 import com.azure.cosmos.models.SqlParameter;
-import com.azure.cosmos.models.SemanticRerankResult;
-import com.azure.cosmos.models.SemanticRerankScore;
 import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.util.CosmosPagedFlux;
 import com.azure.cosmos.util.UtilBridgeInternal;
@@ -2046,76 +2044,4 @@ public interface SqlParameterAccessor {
             SqlParameter createCopy(SqlParameter original);
         }
     }
-
-    public static final class SemanticRerankResultHelper {
-        private static final AtomicReference<SemanticRerankResultAccessor> accessor = new AtomicReference<>();
-        private static final AtomicBoolean semanticRerankResultClassLoaded = new AtomicBoolean(false);
-
-        private SemanticRerankResultHelper() {}
-
-        public static void setSemanticRerankResultAccessor(final SemanticRerankResultAccessor newAccessor) {
-            if (!accessor.compareAndSet(null, newAccessor)) {
-                logger.debug("SemanticRerankResultAccessor already initialized!");
-            } else {
-                logger.debug("Setting SemanticRerankResultAccessor...");
-                semanticRerankResultClassLoaded.set(true);
-            }
-        }
-
-        public static SemanticRerankResultAccessor getSemanticRerankResultAccessor() {
-            if (!semanticRerankResultClassLoaded.get()) {
-                logger.debug("Initializing SemanticRerankResultAccessor...");
-                initializeAllAccessors();
-            }
-
-            SemanticRerankResultAccessor snapshot = accessor.get();
-            if (snapshot == null) {
-                logger.error("SemanticRerankResultAccessor is not initialized yet!");
-            }
-
-            return snapshot;
-        }
-
-        public interface SemanticRerankResultAccessor {
-            void setScores(SemanticRerankResult result, java.util.List<SemanticRerankScore> scores);
-            void setLatency(SemanticRerankResult result, java.util.Map<String, Object> latency);
-            void setTokenUsage(SemanticRerankResult result, java.util.Map<String, Object> tokenUsage);
-        }
-    }
-
-    public static final class SemanticRerankScoreHelper {
-        private static final AtomicReference<SemanticRerankScoreAccessor> accessor = new AtomicReference<>();
-        private static final AtomicBoolean semanticRerankScoreClassLoaded = new AtomicBoolean(false);
-
-        private SemanticRerankScoreHelper() {}
-
-        public static void setSemanticRerankScoreAccessor(final SemanticRerankScoreAccessor newAccessor) {
-            if (!accessor.compareAndSet(null, newAccessor)) {
-                logger.debug("SemanticRerankScoreAccessor already initialized!");
-            } else {
-                logger.debug("Setting SemanticRerankScoreAccessor...");
-                semanticRerankScoreClassLoaded.set(true);
-            }
-        }
-
-        public static SemanticRerankScoreAccessor getSemanticRerankScoreAccessor() {
-            if (!semanticRerankScoreClassLoaded.get()) {
-                logger.debug("Initializing SemanticRerankScoreAccessor...");
-                initializeAllAccessors();
-            }
-
-            SemanticRerankScoreAccessor snapshot = accessor.get();
-            if (snapshot == null) {
-                logger.error("SemanticRerankScoreAccessor is not initialized yet!");
-            }
-
-            return snapshot;
-        }
-
-        public interface SemanticRerankScoreAccessor {
-            void setIndex(SemanticRerankScore score, int index);
-            void setScore(SemanticRerankScore score, double scoreValue);
-            void setDocument(SemanticRerankScore score, String document);
-        }
-    }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
deleted file mode 100644
index 36271e3172df..000000000000
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/inference/InferenceService.java
+++ /dev/null
@@ -1,534 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-package com.azure.cosmos.implementation.inference;
-
-import com.azure.core.credential.AccessToken;
-import com.azure.core.credential.SimpleTokenCache;
-import com.azure.core.credential.TokenCredential;
-import com.azure.core.credential.TokenRequestContext;
-import com.azure.cosmos.BridgeInternal;
-import com.azure.cosmos.CosmosException;
-import com.azure.cosmos.implementation.BadRequestException;
-import com.azure.cosmos.implementation.Configs;
-import com.azure.cosmos.implementation.HttpConstants;
-import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
-import com.azure.cosmos.implementation.LifeCycleUtils;
-import com.azure.cosmos.implementation.Utils;
-import com.azure.cosmos.implementation.http.HttpClient;
-import com.azure.cosmos.implementation.http.HttpClientConfig;
-import com.azure.cosmos.implementation.http.HttpHeaders;
-import com.azure.cosmos.implementation.http.HttpRequest;
-import com.azure.cosmos.implementation.http.HttpResponse;
-import com.azure.cosmos.models.SemanticRerankResult;
-import com.azure.cosmos.models.SemanticRerankScore;
-import com.fasterxml.jackson.databind.JsonNode;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.node.ArrayNode;
-import com.fasterxml.jackson.databind.node.ObjectNode;
-import io.netty.handler.codec.http.HttpMethod;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import reactor.core.publisher.Mono;
-import reactor.core.scheduler.Scheduler;
-import reactor.core.scheduler.Schedulers;
-import reactor.util.retry.Retry;
-
-import java.io.IOException;
-import java.net.URI;
-import java.nio.charset.StandardCharsets;
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ThreadLocalRandom;
-
-import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
-
-/**
- * Internal service for semantic reranking operations.
- * <p>
- * While this class is public, it is not part of our published public APIs.
- * This is meant to be internally used only by our SDK.
- */
-public class InferenceService implements AutoCloseable {
-    private static final Logger logger = LoggerFactory.getLogger(InferenceService.class);
-    private static final String INFERENCE_SCOPE = "https://dbinference.azure.com/.default";
-    private static final String BASE_PATH = "/inference/semanticReranking";
-    private static final String INFERENCE_USER_AGENT =
-        "cosmos-inference-java/" + HttpConstants.Versions.getSdkVersion();
-    private static final ObjectMapper OBJECT_MAPPER = Utils.getSimpleObjectMapper();
-
-    // Following the convention used in Configs: system property is dot.notation; environment variable is UPPER_SNAKE_CASE.
-    private static final String MAX_CONNECTION_LIMIT_PROPERTY = "azure.cosmos.semanticReranker.inferenceService.maxConnectionLimit";
-    private static final String MAX_CONNECTION_LIMIT_VARIABLE = "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_SERVICE_MAX_CONNECTION_LIMIT";
-
-    // Option keys that callers can pass in the options map
-    public static final String OPTION_TIMEOUT_SECONDS = "timeout_seconds";
-    public static final String OPTION_RETURN_DOCUMENTS = "return_documents";
-    public static final String OPTION_TOP_K = "top_k";
-    public static final String OPTION_BATCH_SIZE = "batch_size";
-    public static final String OPTION_SORT = "sort";
-    public static final String OPTION_DOCUMENT_TYPE = "document_type";
-    public static final String OPTION_TARGET_PATHS = "target_paths";
-
-    /**
-     * Default network-level (connection + read) timeout for inference service HTTP calls.
-     * Set to match {@link #DEFAULT_REQUEST_TIMEOUT} because inference calls can take tens of
-     * seconds for large document sets. If this were shorter than the operation timeout,
-     * the socket would time out before the 120-second operation deadline could fire.
-     */
-    public static final Duration DEFAULT_NETWORK_REQUEST_TIMEOUT = Duration.ofSeconds(120);
-    public static final Duration DEFAULT_IDLE_CONNECTION_TIMEOUT = Duration.ofSeconds(60);
-    public static final Duration DEFAULT_CONNECTION_ACQUIRE_TIMEOUT = Duration.ofSeconds(5);
-    /**
-     * Default per-request timeout for semantic rerank calls (120 seconds).
-     * Callers using the async API can override downstream with {@code .timeout(Duration)}.
-     * Callers using the sync API can override by passing {@value #OPTION_TIMEOUT_SECONDS} in the options map.
-     */
-    public static final Duration DEFAULT_REQUEST_TIMEOUT = Duration.ofSeconds(120);
-    /**
-     * Default maximum number of pooled connections to the inference endpoint.
-     * Matches .NET's {@code inferenceServiceDefaultMaxConnectionLimit = 50}.
-     * Override via system property or environment variable
-     * {@value #MAX_CONNECTION_LIMIT_PROPERTY}.
-     */
-    public static final int DEFAULT_MAX_CONNECTION_POOL_SIZE = 50;
-
-    // Retry policy — matches Python: TOTAL_RETRIES=3, RETRY_BACKOFF_FACTOR=0.8s, RETRY_BACKOFF_MAX=120s
-    static final int RETRY_MAX_ATTEMPTS = 3;
-    static final Duration RETRY_INITIAL_BACKOFF = Duration.ofMillis(800);
-    static final Duration RETRY_MAX_BACKOFF = Duration.ofSeconds(120);
-    // Status codes that are safe to retry (transient failures and rate limiting)
-    private static final Set<Integer> RETRYABLE_STATUS_CODES = Collections.unmodifiableSet(
-        new HashSet<>(Arrays.asList(
-            HttpConstants.StatusCodes.TOO_MANY_REQUESTS,      // 429 — rate limited
-            HttpConstants.StatusCodes.INTERNAL_SERVER_ERROR,  // 500 — transient server error
-            502,                                               // 502 — bad gateway (no constant in HttpConstants)
-            HttpConstants.StatusCodes.SERVICE_UNAVAILABLE     // 503 — transient unavailability
-        ))
-    );
-
-    private final URI inferenceEndpoint;
-    private final HttpClient httpClient;
-    private SimpleTokenCache tokenCache = null;
-    // Package-private so tests can inject a VirtualTimeScheduler to make backoff delays deterministic
-    Scheduler retryScheduler = Schedulers.parallel();
-    // Package-private so tests can disable jitter for deterministic virtual-time delays
-    double retryJitter = 0.5;
-
-    /**
-     * Creates a new InferenceService instance.
-     *
-     * @param tokenCredential The Azure AD token credential.
-     * @throws IllegalArgumentException if inference endpoint is not configured or token credential is null.
-     */
-    public InferenceService(TokenCredential tokenCredential) {
-        checkNotNull(tokenCredential,
-            "Semantic reranking requires AAD authentication. "
-                + "Rebuild the CosmosClient using .credential(TokenCredential) — "
-                + "key-based auth (master key or AzureKeyCredential) is not supported for this operation.");
-
-        Configs configs = new Configs();
-        URI inferenceBaseUrl = configs.getInferenceServiceEndpoint();
-
-        if (inferenceBaseUrl == null || inferenceBaseUrl.toString().trim().isEmpty()) {
-            throw new IllegalArgumentException("Inference endpoint property must be set to use semantic reranking");
-        }
-
-        this.inferenceEndpoint = URI.create(inferenceBaseUrl + BASE_PATH);
-        HttpClientConfig httpClientConfig = new HttpClientConfig(configs)
-            .withNetworkRequestTimeout(DEFAULT_NETWORK_REQUEST_TIMEOUT)
-            .withConnectionAcquireTimeout(DEFAULT_CONNECTION_ACQUIRE_TIMEOUT)
-            .withMaxIdleConnectionTimeout(DEFAULT_IDLE_CONNECTION_TIMEOUT)
-            .withPoolSize(resolveMaxConnectionPoolSize());
-        this.httpClient = HttpClient.createFixed(httpClientConfig);
-
-        // Create token cache for inference service scope
-        this.tokenCache = new SimpleTokenCache(() -> {
-            TokenRequestContext context = new TokenRequestContext().addScopes(INFERENCE_SCOPE);
-            return tokenCredential.getToken(context)
-                .doOnNext(token -> {
-                    if (logger.isDebugEnabled()) {
-                        logger.debug("Acquired AAD token for inference service scope: {}", INFERENCE_SCOPE);
-                    }
-                });
-        });
-
-        if (logger.isInfoEnabled()) {
-            logger.info("InferenceService initialized with endpoint: {}", inferenceEndpoint);
-        }
-    }
-
-    /**
-     * Closes the InferenceService and releases the underlying HTTP client connection pool.
-     */
-    @Override
-    public void close() {
-        logger.info("Shutting down InferenceService httpClient...");
-        LifeCycleUtils.closeQuietly(this.httpClient);
-    }
-
-    /**
-     * Resolves the maximum connection pool size from system property, environment variable,
-     * or the default ({@value #DEFAULT_MAX_CONNECTION_POOL_SIZE}).
-     * Matches .NET's AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_SERVICE_MAX_CONNECTION_LIMIT override.
-     */
-    private static int resolveMaxConnectionPoolSize() {
-        String fromProperty = System.getProperty(MAX_CONNECTION_LIMIT_PROPERTY);
-        if (fromProperty != null && !fromProperty.isEmpty()) {
-            try {
-                return Integer.parseInt(fromProperty);
-            } catch (NumberFormatException e) {
-                logger.warn("Invalid value for system property {}: '{}'. Using default {}.",
-                    MAX_CONNECTION_LIMIT_PROPERTY, fromProperty, DEFAULT_MAX_CONNECTION_POOL_SIZE);
-            }
-        }
-
-        String fromEnv = System.getenv(MAX_CONNECTION_LIMIT_VARIABLE);
-        if (fromEnv != null && !fromEnv.isEmpty()) {
-            try {
-                return Integer.parseInt(fromEnv);
-            } catch (NumberFormatException e) {
-                logger.warn("Invalid value for environment variable {}: '{}'. Using default {}.",
-                    MAX_CONNECTION_LIMIT_VARIABLE, fromEnv, DEFAULT_MAX_CONNECTION_POOL_SIZE);
-            }
-        }
-
-        return DEFAULT_MAX_CONNECTION_POOL_SIZE;
-    }
-
-    /**
-     * Performs semantic reranking of documents.
-     *
-     * <p>The request timeout defaults to 120 seconds ({@link #DEFAULT_REQUEST_TIMEOUT}).
-     * To override it, pass {@code "timeout_seconds"} (as a {@link Number}) in the {@code options} map.
-     * Callers using the async {@link reactor.core.publisher.Mono} result can also apply
-     * {@code .timeout(Duration)} downstream without needing to set the option.
-     *
-     * @param rerankContext The query or context string used to score documents.
-     * @param documents The list of document strings to rerank.
-     * @param options Optional reranking parameters. SDK-local keys ({@link #OPTION_TIMEOUT_SECONDS})
-     *                are consumed locally and not forwarded to the inference endpoint.
-     * @return A Mono emitting the semantic rerank result.
-     */
-    public Mono<SemanticRerankResult> semanticRerank(
-        String rerankContext,
-        List<String> documents,
-        Map<String, Object> options) {
-
-        checkNotNull(rerankContext, "Rerank context cannot be null");
-        checkNotNull(documents, "Documents list cannot be null");
-
-        if (rerankContext.trim().isEmpty()) {
-            return Mono.error(new IllegalArgumentException("Rerank context cannot be empty"));
-        }
-        if (documents.isEmpty()) {
-            return Mono.error(new IllegalArgumentException("Documents list cannot be empty"));
-        }
-
-        // Resolve the per-request timeout: options > default
-        final Duration requestTimeout = resolveRequestTimeout(options);
-
-        return this.tokenCache.getToken()
-            .flatMap(accessToken -> {
-                try {
-                    // Build request payload
-                    ObjectNode payload = OBJECT_MAPPER.createObjectNode();
-                    payload.put("query", rerankContext);
-
-                    ArrayNode documentsArray = payload.putArray("documents");
-                    for (String doc : documents) {
-                        documentsArray.add(doc);
-                    }
-
-                    if (options != null) {
-                        options.forEach((key, value) -> {
-                            // timeout_seconds is a SDK-local option — do not forward to the endpoint
-                            if (value != null && !OPTION_TIMEOUT_SECONDS.equals(key)) {
-                                payload.set(key, OBJECT_MAPPER.valueToTree(value));
-                            }
-                        });
-                    }
-
-                    String requestBody = OBJECT_MAPPER.writeValueAsString(payload);
-
-                    HttpRequest httpRequest = getHttpRequest(accessToken);
-                    httpRequest.withBody(requestBody.getBytes(StandardCharsets.UTF_8));
-
-                    if (logger.isDebugEnabled()) {
-                        logger.debug("Sending semantic rerank request to: {} (timeout: {})", inferenceEndpoint, requestTimeout);
-                    }
-
-                    return httpClient.send(httpRequest, requestTimeout)
-                        .flatMap(response -> parseResponse(response));
-
-                } catch (IOException e) {
-                    logger.error("Failed to serialize request payload", e);
-                    // Use BadRequestException (400) + CUSTOM_SERIALIZER_EXCEPTION sub-status — same
-                    // pattern as the response-parse path below. 400 is not in RETRYABLE_STATUS_CODES,
-                    // so withRetry will not waste attempts on a deterministic serialization failure.
-                    BadRequestException badRequest = new BadRequestException(
-                        "Failed to serialize semantic rerank request: " + e.getMessage(), e);
-                    BridgeInternal.setSubStatusCode(badRequest,
-                        HttpConstants.SubStatusCodes.CUSTOM_SERIALIZER_EXCEPTION);
-                    return Mono.error(badRequest);
-                }
-            })
-            .as(this::withRetry)
-            .doOnError(error -> logger.error("Semantic rerank operation failed", error));
-    }
-
-    /**
-     * Resolves the effective request timeout.
-     * If the caller supplied {@value #OPTION_TIMEOUT_SECONDS} in the options map, that value is used;
-     * otherwise {@link #DEFAULT_REQUEST_TIMEOUT} applies.
-     */
-    private static Duration resolveRequestTimeout(Map<String, Object> options) {
-        if (options != null) {
-            Object value = options.get(OPTION_TIMEOUT_SECONDS);
-            if (value instanceof Number) {
-                double seconds = ((Number) value).doubleValue();
-                if (seconds > 0) {
-                    return Duration.ofMillis((long) (seconds * 1000));
-                } else {
-                    logger.warn("Invalid '{}' value: {}. Must be > 0. Using default timeout {}.",
-                        OPTION_TIMEOUT_SECONDS, seconds, DEFAULT_REQUEST_TIMEOUT);
-                }
-            }
-        }
-        return DEFAULT_REQUEST_TIMEOUT;
-    }
-
-    /**
-     * Wraps a {@link Mono} with retry logic for transient inference endpoint failures.
-     *
-     * <p>Retries up to {@value #RETRY_MAX_ATTEMPTS} times on retryable {@link com.azure.cosmos.CosmosException}
-     * status codes (429, 500, 502, 503). For HTTP 429 responses, the server-supplied
-     * {@code Retry-After} header is honoured (delay-seconds form per RFC 7231 § 7.1.3) and capped at
-     * {@link #RETRY_MAX_BACKOFF}. If the header is missing or in HTTP-date form, exponential
-     * backoff with jitter is used instead, starting at {@link #RETRY_INITIAL_BACKOFF} and capped at
-     * {@link #RETRY_MAX_BACKOFF}. All other exceptions (non-retryable status codes, serialisation
-     * errors, etc.) propagate immediately.
-     *
-     * <p>Matches the Python SDK's retry policy:
-     * {@code TOTAL_RETRIES=3, RETRY_BACKOFF_FACTOR=0.8, RETRY_BACKOFF_MAX=120s, honour Retry-After on 429}.
-     */
-    private <T> Mono<T> withRetry(Mono<T> source) {
-        return source.retryWhen(
-            Retry.from(retrySignals -> retrySignals.concatMap(signal -> {
-                Throwable error = signal.failure();
-                if (!(error instanceof CosmosException)) {
-                    return Mono.error(error);
-                }
-                CosmosException cex = (CosmosException) error;
-                int statusCode = cex.getStatusCode();
-                if (!RETRYABLE_STATUS_CODES.contains(statusCode)) {
-                    return Mono.error(error);
-                }
-                long attempt = signal.totalRetries();
-                if (attempt >= RETRY_MAX_ATTEMPTS) {
-                    return Mono.error(error);
-                }
-
-                Duration delay = computeRetryDelay(cex, statusCode, attempt);
-                logger.warn(
-                    "Semantic rerank transient failure (status={}, attempt={}/{}), retrying after {} ms...",
-                    statusCode, attempt + 1, RETRY_MAX_ATTEMPTS, delay.toMillis());
-                return Mono.delay(delay, retryScheduler);
-            }))
-        );
-    }
-
-    /**
-     * Computes the delay before the next retry attempt.
-     * Prefers the {@code Retry-After} response header on 429 (delay-seconds form, capped at
-     * {@link #RETRY_MAX_BACKOFF}). Otherwise falls back to exponential backoff with jitter.
-     */
-    private Duration computeRetryDelay(CosmosException cex, int statusCode, long attempt) {
-        if (statusCode == HttpConstants.StatusCodes.TOO_MANY_REQUESTS) {
-            Duration serverDelay = parseRetryAfterSeconds(cex.getResponseHeaders());
-            if (serverDelay != null) {
-                return serverDelay.compareTo(RETRY_MAX_BACKOFF) > 0 ? RETRY_MAX_BACKOFF : serverDelay;
-            }
-        }
-        return exponentialBackoffWithJitter(attempt);
-    }
-
-    /**
-     * Parses the RFC 7231 {@code Retry-After} header as delay-seconds.
-     * Returns {@code null} if the header is absent, blank, negative, or in HTTP-date form
-     * (the latter is not supported here — exponential backoff is used in that case).
-     * Lookup is case-insensitive per RFC 7230 § 3.2.
-     */
-    private static Duration parseRetryAfterSeconds(Map<String, String> headers) {
-        if (headers == null || headers.isEmpty()) {
-            return null;
-        }
-        String value = null;
-        for (Map.Entry<String, String> entry : headers.entrySet()) {
-            if (entry.getKey() != null
-                && HttpConstants.HttpHeaders.RETRY_AFTER.equalsIgnoreCase(entry.getKey())) {
-                value = entry.getValue();
-                break;
-            }
-        }
-        if (value == null || value.trim().isEmpty()) {
-            return null;
-        }
-        try {
-            long seconds = Long.parseLong(value.trim());
-            if (seconds < 0) {
-                return null;
-            }
-            return Duration.ofSeconds(seconds);
-        } catch (NumberFormatException ignored) {
-            // HTTP-date form (e.g. "Fri, 31 Dec 1999 23:59:59 GMT") is not parsed.
-            return null;
-        }
-    }
-
-    /**
-     * Exponential backoff with jitter, matching the Python SDK semantics.
-     * {@code RETRY_INITIAL_BACKOFF * 2^attempt}, capped at {@link #RETRY_MAX_BACKOFF},
-     * with +/- {@code retryJitter * base} random jitter applied.
-     */
-    private Duration exponentialBackoffWithJitter(long attempt) {
-        long baseMillis = RETRY_INITIAL_BACKOFF.toMillis() << Math.min(attempt, 30);
-        long cappedMillis = Math.min(baseMillis, RETRY_MAX_BACKOFF.toMillis());
-        long jitterRange = (long) (cappedMillis * retryJitter);
-        long jitter = jitterRange > 0
-            ? ThreadLocalRandom.current().nextLong(-jitterRange, jitterRange + 1)
-            : 0L;
-        return Duration.ofMillis(Math.max(0L, cappedMillis + jitter));
-    }
-
-    private HttpRequest getHttpRequest(AccessToken accessToken) {
-        HttpHeaders headers = new HttpHeaders();
-        headers.set(HttpConstants.HttpHeaders.CONTENT_TYPE, "application/json");
-        headers.set(HttpConstants.HttpHeaders.AUTHORIZATION, "Bearer " + accessToken.getToken());
-        headers.set(HttpConstants.HttpHeaders.USER_AGENT, INFERENCE_USER_AGENT);
-
-        return new HttpRequest(
-            HttpMethod.POST,
-            inferenceEndpoint,
-            inferenceEndpoint.getPort(),
-            headers);
-    }
-
-    /**
-     * Parses the HTTP response into a SemanticRerankResult.
-     *
-     * @param response The HTTP response.
-     * @return A Mono emitting the parsed result.
-     */
-    private Mono<SemanticRerankResult> parseResponse(HttpResponse response) {
-        int statusCode = response.statusCode();
-
-        return response.bodyAsString()
-            .flatMap(bodyString -> {
-                if (statusCode >= 400) {
-                    logger.error("Semantic rerank request failed with status {}: {}", statusCode, bodyString);
-                    Map<String, String> headersMap = convertHeaders(response.headers());
-                    return Mono.error(BridgeInternal.createCosmosException(
-                        String.format("Semantic rerank request failed: %s", bodyString),
-                        null,
-                        headersMap,
-                        statusCode,
-                        null));
-                }
-
-                try {
-                    JsonNode rootNode = OBJECT_MAPPER.readTree(bodyString);
-                    SemanticRerankResult result = new SemanticRerankResult();
-                    ImplementationBridgeHelpers.SemanticRerankResultHelper.SemanticRerankResultAccessor resultAccessor =
-                        ImplementationBridgeHelpers.SemanticRerankResultHelper.getSemanticRerankResultAccessor();
-                    ImplementationBridgeHelpers.SemanticRerankScoreHelper.SemanticRerankScoreAccessor scoreAccessor =
-                        ImplementationBridgeHelpers.SemanticRerankScoreHelper.getSemanticRerankScoreAccessor();
-
-                    // Parse scores
-                    if (rootNode.has("Scores")) {
-                        JsonNode scoresNode = rootNode.get("Scores");
-                        List<SemanticRerankScore> scores = new ArrayList<>();
-
-                        if (scoresNode.isArray()) {
-                            for (JsonNode scoreNode : scoresNode) {
-                                SemanticRerankScore score = new SemanticRerankScore();
-                                JsonNode indexNode = scoreNode.get("index");
-                                JsonNode scoreValNode = scoreNode.get("score");
-                                if (indexNode != null) {
-                                    scoreAccessor.setIndex(score, indexNode.asInt());
-                                }
-                                if (scoreValNode != null) {
-                                    scoreAccessor.setScore(score, scoreValNode.asDouble());
-                                }
-
-                                if (scoreNode.has("document")) {
-                                    scoreAccessor.setDocument(score, scoreNode.get("document").asText());
-                                }
-
-                                scores.add(score);
-                            }
-                        }
-                        resultAccessor.setScores(result, scores);
-                    }
-
-                    // Parse latency
-                    if (rootNode.has("latency")) {
-                        Map<String, Object> latency = new HashMap<>();
-                        rootNode.get("latency").fields().forEachRemaining(
-                            entry -> latency.put(entry.getKey(), entry.getValue().asDouble()));
-                        resultAccessor.setLatency(result, latency);
-                    }
-
-                    // Parse token usage
-                    if (rootNode.has("token_usage")) {
-                        Map<String, Object> tokenUsage = new HashMap<>();
-                        rootNode.get("token_usage").fields().forEachRemaining(
-                            entry -> tokenUsage.put(entry.getKey(), entry.getValue().asInt()));
-                        resultAccessor.setTokenUsage(result, tokenUsage);
-                    }
-
-                    if (logger.isDebugEnabled()) {
-                        logger.debug("Successfully parsed semantic rerank response with {} scores",
-                            result.getScores() != null ? result.getScores().size() : 0);
-                    }
-
-                    return Mono.just(result);
-
-                } catch (Exception e) {
-                    logger.error("Failed to parse semantic rerank response", e);
-                    // Use BadRequestException (400) + CUSTOM_SERIALIZER_EXCEPTION sub-status —
-                    // the same pattern CosmosItemSerializer uses for client-side deserialization
-                    // failures. 400 is not in RETRYABLE_STATUS_CODES, so no retry will be attempted.
-                    BadRequestException parseException = new BadRequestException(
-                        "Failed to parse semantic rerank response: " + e.getMessage(), e);
-                    BridgeInternal.setSubStatusCode(parseException,
-                        HttpConstants.SubStatusCodes.CUSTOM_SERIALIZER_EXCEPTION);
-                    return Mono.error(parseException);
-                }
-            });
-    }
-
-    /**
-     * Converts HttpHeaders to a Map.
-     *
-     * @param headers The HTTP headers.
-     * @return A map of header names to values.
-     */
-    private Map<String, String> convertHeaders(HttpHeaders headers) {
-        Map<String, String> headersMap = new HashMap<>();
-        if (headers != null) {
-            headers.toMap().forEach((key, value) -> {
-                if (value != null) {
-                    headersMap.put(key, value);
-                }
-            });
-        }
-        return headersMap;
-    }
-}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
index f0253ad99486..7dde81e06e98 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
@@ -769,7 +769,5 @@ public static void  initializeAllAccessors() {
         PriorityLevel.initialize();
         SqlQuerySpec.initialize();
         SqlParameter.initialize();
-        SemanticRerankResult.initialize();
-        SemanticRerankScore.initialize();
     }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
deleted file mode 100644
index bfbe0330b909..000000000000
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankResult.java
+++ /dev/null
@@ -1,90 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-package com.azure.cosmos.models;
-
-import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
-import com.azure.cosmos.util.Beta;
-
-import java.util.List;
-import java.util.Map;
-
-/**
- * Represents the result of a semantic rerank operation.
- */
-@Beta(value = Beta.SinceVersion.V4_81_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
-public final class SemanticRerankResult {
-    private List<SemanticRerankScore> scores;
-    private Map<String, Object> latency;
-    private Map<String, Object> tokenUsage;
-
-    /**
-     * Creates a new instance of SemanticRerankResult.
-     */
-    public SemanticRerankResult() {
-    }
-
-    /**
-     * Gets the list of scored documents.
-     *
-     * @return the list of scored documents.
-     */
-    public List<SemanticRerankScore> getScores() {
-        return scores;
-    }
-
-    private void setScores(List<SemanticRerankScore> scores) {
-        this.scores = scores;
-    }
-
-    /**
-     * Gets the latency information for the operation as a map of metric names to values.
-     *
-     * @return the latency information map.
-     */
-    public Map<String, Object> getLatency() {
-        return latency;
-    }
-
-    private void setLatency(Map<String, Object> latency) {
-        this.latency = latency;
-    }
-
-    /**
-     * Gets the token usage information for the operation as a map of metric names to values.
-     *
-     * @return the token usage information map.
-     */
-    public Map<String, Object> getTokenUsage() {
-        return tokenUsage;
-    }
-
-    private void setTokenUsage(Map<String, Object> tokenUsage) {
-        this.tokenUsage = tokenUsage;
-    }
-
-    ///////////////////////////////////////////////////////////////////////////////////////////
-    // the following helper/accessor only helps to access this class outside of this package.//
-    ///////////////////////////////////////////////////////////////////////////////////////////
-    static void initialize() {
-        ImplementationBridgeHelpers.SemanticRerankResultHelper.setSemanticRerankResultAccessor(
-            new ImplementationBridgeHelpers.SemanticRerankResultHelper.SemanticRerankResultAccessor() {
-                @Override
-                public void setScores(SemanticRerankResult result, List<SemanticRerankScore> scores) {
-                    result.setScores(scores);
-                }
-
-                @Override
-                public void setLatency(SemanticRerankResult result, Map<String, Object> latency) {
-                    result.setLatency(latency);
-                }
-
-                @Override
-                public void setTokenUsage(SemanticRerankResult result, Map<String, Object> tokenUsage) {
-                    result.setTokenUsage(tokenUsage);
-                }
-            }
-        );
-    }
-
-    static { initialize(); }
-}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java
deleted file mode 100644
index 96ae57ce8389..000000000000
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/SemanticRerankScore.java
+++ /dev/null
@@ -1,87 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-package com.azure.cosmos.models;
-
-import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
-import com.azure.cosmos.util.Beta;
-
-/**
- * Represents a single scored document in the semantic rerank result.
- */
-@Beta(value = Beta.SinceVersion.V4_81_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING)
-public final class SemanticRerankScore {
-    private int index;
-    private String document;
-    private double score;
-
-    /**
-     * Creates a new instance of SemanticRerankScore.
-     */
-    public SemanticRerankScore() {
-    }
-
-    /**
-     * Gets the index of the document in the original input list.
-     *
-     * @return the document index.
-     */
-    public int getIndex() {
-        return index;
-    }
-
-    private void setIndex(int index) {
-        this.index = index;
-    }
-
-    /**
-     * Gets the document text (if returnDocuments was true in the request).
-     *
-     * @return the document text, or null if not included.
-     */
-    public String getDocument() {
-        return document;
-    }
-
-    private void setDocument(String document) {
-        this.document = document;
-    }
-
-    /**
-     * Gets the semantic relevance score for this document.
-     *
-     * @return the relevance score.
-     */
-    public double getScore() {
-        return score;
-    }
-
-    private void setScore(double score) {
-        this.score = score;
-    }
-
-    ///////////////////////////////////////////////////////////////////////////////////////////
-    // the following helper/accessor only helps to access this class outside of this package.//
-    ///////////////////////////////////////////////////////////////////////////////////////////
-    static void initialize() {
-        ImplementationBridgeHelpers.SemanticRerankScoreHelper.setSemanticRerankScoreAccessor(
-            new ImplementationBridgeHelpers.SemanticRerankScoreHelper.SemanticRerankScoreAccessor() {
-                @Override
-                public void setIndex(SemanticRerankScore score, int index) {
-                    score.setIndex(index);
-                }
-
-                @Override
-                public void setScore(SemanticRerankScore score, double scoreValue) {
-                    score.setScore(scoreValue);
-                }
-
-                @Override
-                public void setDocument(SemanticRerankScore score, String document) {
-                    score.setDocument(document);
-                }
-            }
-        );
-    }
-
-    static { initialize(); }
-}
diff --git a/sdk/cosmos/pom.xml b/sdk/cosmos/pom.xml
index c53f1e7d00b2..0c59f51ce4b2 100644
--- a/sdk/cosmos/pom.xml
+++ b/sdk/cosmos/pom.xml
@@ -11,6 +11,7 @@
 
   <modules>
     <module>azure-cosmos</module>
+    <module>azure-cosmos-ai</module>
     <module>azure-cosmos-benchmark</module>
     <module>azure-cosmos-encryption</module>
     <module>azure-cosmos-spark_3</module>