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}.
+ *
+ * @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.
+ *
+ * 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.
+ *
+ * @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 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.
+ *
+ * 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 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 semanticRerank(
+ String rerankContext,
+ List documents,
+ Map 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 documents,
+ Map 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 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 Mono withRetry(Mono 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 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.
+ *
+ * This class lives in the {@code models} package so it can access package-private setters
+ * on {@link SemanticRerankResult} and {@link SemanticRerankScore}.
+ *
+ * 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 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 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 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 scores;
+ private Map latency;
+ private Map tokenUsage;
+
+ /**
+ * Creates a new instance of SemanticRerankResult.
+ */
+ public SemanticRerankResult() {
+ }
+
+ /**
+ * Gets the list of scored documents.
+ *
+ * @return the list of scored documents.
+ */
+ public List 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 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 getTokenUsage() {
+ return tokenUsage;
+ }
+
+ // Package-private setters used by InferenceResponseParser in the same package
+ void setScores(List scores) {
+ this.scores = scores;
+ }
+
+ void setLatency(Map latency) {
+ this.latency = latency;
+ }
+
+ void setTokenUsage(Map 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-ai/src/samples/java/com/azure/cosmos/ai/SemanticRerankSample.java b/sdk/cosmos/azure-cosmos-ai/src/samples/java/com/azure/cosmos/ai/SemanticRerankSample.java
new file mode 100644
index 000000000000..51bb587a620a
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos-ai/src/samples/java/com/azure/cosmos/ai/SemanticRerankSample.java
@@ -0,0 +1,128 @@
+// 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.cosmos.ai.models.SemanticRerankResult;
+import com.azure.cosmos.ai.models.SemanticRerankScore;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Sample demonstrating semantic rerank functionality using Azure Cosmos DB AI SDK.
+ *
+ * Prerequisites:
+ * 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()}
+ *
+ */
+public class SemanticRerankSample {
+
+ /**
+ * Main entry point.
+ *
+ * @param args command line arguments (not used).
+ */
+ public static void main(String[] args) {
+ 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.
+ *
+ * @param endpoint The inference service endpoint.
+ * @param credential The Azure AD token credential.
+ */
+ static void runSemanticRerankDemo(String endpoint, TokenCredential credential) {
+ System.out.println("Semantic Rerank Sample");
+ System.out.println("======================\n");
+
+ // BEGIN: readme-sample-createClient
+ CosmosAIAsyncClient client = new CosmosAIClientBuilder()
+ .endpoint(endpoint)
+ .credential(credential)
+ .buildAsyncClient();
+ // END: readme-sample-createClient
+
+ List 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
+ client.semanticRerank(rerankContext, documents, null)
+ .subscribe(
+ result -> printResults(result),
+ error -> System.err.println("Error: " + error.getMessage())
+ );
+
+ // Rerank with custom options
+ Map options = new HashMap<>();
+ options.put("return_documents", true);
+ options.put("top_k", 3);
+ options.put("sort", true);
+
+ client.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 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 latency = result.getLatency();
+ if (latency != null) {
+ System.out.println("\nLatency:");
+ latency.forEach((key, value) -> System.out.printf(" %s: %s%n", key, value));
+ }
+
+ Map tokenUsage = result.getTokenUsage();
+ if (tokenUsage != null) {
+ System.out.println("\nToken Usage:");
+ tokenUsage.forEach((key, value) -> System.out.printf(" %s: %s%n", key, value));
+ }
+ }
+}
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 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 requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+ when(mockAzureCoreHttpClient.send(requestCaptor.capture(), any()))
+ .thenReturn(Mono.just(mockResponse));
+
+ InferenceService service = createService();
+
+ List documents = Arrays.asList("doc1", "doc2");
+ Map 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 requestCaptor = ArgumentCaptor.forClass(HttpRequest.class);
+ when(mockAzureCoreHttpClient.send(requestCaptor.capture(), any()))
+ .thenReturn(Mono.just(mockResponse));
+
+ InferenceService service = createService();
+
+ Map 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 getBody() {
+ return reactor.core.publisher.Flux.just(ByteBuffer.wrap(body.getBytes(StandardCharsets.UTF_8)));
+ }
+
+ @Override
+ public Mono getBodyAsByteArray() {
+ return Mono.just(body.getBytes(StandardCharsets.UTF_8));
+ }
+
+ @Override
+ public Mono getBodyAsString() {
+ return Mono.just(body);
+ }
+
+ @Override
+ public Mono getBodyAsString(Charset charset) {
+ return Mono.just(new String(body.getBytes(StandardCharsets.UTF_8), charset));
+ }
+ }
+}
diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md
index 6eb1b76bad4e..106c55521e10 100644
--- a/sdk/cosmos/azure-cosmos/CHANGELOG.md
+++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -3,10 +3,14 @@
### 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 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
#### 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 fe2b36b766e1..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
@@ -108,6 +108,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 +120,7 @@ private static ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccess
private final List requestPolicies;
private final CosmosItemSerializer defaultCustomSerializer;
private final java.util.function.Function containerFactory;
+ private volatile boolean closed;
CosmosAsyncClient(CosmosClientBuilder builder) {
// Async Cosmos client wrapper
@@ -131,7 +133,7 @@ private static ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccess
List 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 +294,16 @@ AzureKeyCredential credential() {
return credential;
}
+ /**
+ * Gets the token credential.
+ *
+ * @return token credential.
+ */
+ TokenCredential tokenCredential() {
+ return this.tokenCredential;
+ }
+
+
/***
* Get the client telemetry config.
*
@@ -560,6 +572,7 @@ public CosmosAsyncDatabase getDatabase(String id) {
*/
@Override
public void close() {
+ this.closed = true;
if (this.clientMetricRegistrySnapshot != null) {
ClientTelemetryMetrics.remove(this.clientMetricRegistrySnapshot);
}
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..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,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.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.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;
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..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,6 +4,7 @@
package com.azure.cosmos;
import com.azure.cosmos.implementation.throughputControl.sdk.config.GlobalThroughputControlGroup;
+import com.azure.cosmos.implementation.RequestTimeoutException;
import com.azure.cosmos.models.CosmosBatch;
import com.azure.cosmos.models.CosmosBatchOperationResult;
import com.azure.cosmos.models.CosmosBatchRequestOptions;
@@ -38,7 +39,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;
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..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
@@ -109,6 +109,8 @@ public enum SinceVersion {
/** v4.74.0 */
V4_74_0,
/** v4.75.0 */
- V4_75_0
+ V4_75_0,
+ /** v4.81.0 */
+ V4_81_0
}
}
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 @@
azure-cosmos
+ azure-cosmos-ai
azure-cosmos-benchmark
azure-cosmos-encryption
azure-cosmos-spark_3