Implement inference service (semantic reranking) into new azure-cosmos-ai module

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 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
 

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();
     }
 

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,46 @@ 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,

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_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.
      *

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java

@@ -421,6 +421,12 @@ 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;
 
@@ -531,6 +537,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()) {

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);
+        }
+    }
 }