Semantic Rerank: Adds Semantic Rerank API (#5445)

# Pull Request Template

## Description

This pull request introduces a new semantic reranking feature to the
Azure Cosmos DB .NET SDK, enabling users to rerank documents using an
inference service that leverages Azure Active Directory (AAD)
authentication. The main changes include the addition of the
`InferenceService` class, new API surface for semantic reranking, and
appropriate integration into the SDK's authorization and client context
infrastructure. Notably, this functionality is only available when using
AAD authentication.

**Semantic Reranking Feature Integration:**

* Added the `InferenceService` class, which handles communication with
the Cosmos DB Inference Service for semantic reranking, including HTTP
client configuration, payload construction, and response handling. This
service enforces AAD authentication and manages its own authorization
and disposal.
* Introduced a new public (under `PREVIEW`) or internal API
`SemanticRerankAsync` to the `Container` class, allowing users to rerank
a list of documents based on a context/query string. This is implemented
in `ContainerInlineCore` and routed through the client context.
[[1]](diffhunk://#diff-e3b7704253edcfb63d18e851d9b0a8de3ea60889d704c363c2da1899a3ac39b7R1682-R1702)
[[2]](diffhunk://#diff-d7119df75f749b2d2ebcadc708f02d419663febceeaab4147a898c4be777e33dR700-R708)
* To use this feature, the environment variable
"AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT", must be set with
the inference endpoint from the service.
* Additionally, the environment variable
"AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_SERVICE_MAX_CONNECTION_LIMIT",
can be set to change the inference client's max connection limit.

**Authorization and Token Handling Updates:**

* Extended the `AuthorizationTokenProvider` abstraction and its
implementations to support a new method,
`AddInferenceAuthorizationHeaderAsync`, which is only valid for
AAD-based token providers. Non-AAD providers throw a
`NotImplementedException` for this method.
[[1]](diffhunk://#diff-839842554f21dd8c2180a3c4f5a25abffdc392902a0f10f1130397f064d84db5R55-R60)
[[2]](diffhunk://#diff-93d0d6522a71c823da524b5cdceb024a8aae4497b14b383dcac41a98edf7092aR18-R29)
[[3]](diffhunk://#diff-93d0d6522a71c823da524b5cdceb024a8aae4497b14b383dcac41a98edf7092aR78-R92)
[[4]](diffhunk://#diff-1e0a2b64595baccd2b71f2328bd8f6ca9a81fb8175c5d489e5cdadba1b5c3a05R217-R221)
[[5]](diffhunk://#diff-cca89b75cc8e6136d1f595f20169546258617f88cf0b5e0e345554c073e06e78R95-R99)
[[6]](diffhunk://#diff-3ad037c5217ed55cda63aaa1a0e1897c9b1c6603b6e85b9649cdd8cf457ac6fbR128-R132)

**Client Context and Resource Management:**

* Updated `ClientContextCore` and `CosmosClientContext` to manage the
lifecycle of the `InferenceService`, including creation, caching, and
disposal. Added methods for invoking semantic reranking and for
retrieving or creating the inference service instance.
[[1]](diffhunk://#diff-a41317d6b53931e411d3091f58700758c53e80ca27dccedeafd76226852e58d2R38)
[[2]](diffhunk://#diff-a41317d6b53931e411d3091f58700758c53e80ca27dccedeafd76226852e58d2R472-R497)
[[3]](diffhunk://#diff-a41317d6b53931e411d3091f58700758c53e80ca27dccedeafd76226852e58d2R515)
[[4]](diffhunk://#diff-b0bd965dfed52d866ec53af3bbb3b7afb4f420cdda75ac9071515ac471c0f7baR8)
[[5]](diffhunk://#diff-b0bd965dfed52d866ec53af3bbb3b7afb4f420cdda75ac9071515ac471c0f7baR136-R148)
[[6]](diffhunk://#diff-a41317d6b53931e411d3091f58700758c53e80ca27dccedeafd76226852e58d2R8)

**Dependency Updates:**

* Added a dependency on the `Azure.Identity` package in the test project
to support AAD authentication scenarios.
Please delete options that are not relevant.

**Example**

```csharp
//Sample code to demonstrate Semantic Reranking
// Assume 'container' is an instance of Cosmos.Container
// This example queries items from a fitness store with full-text search and then reranks them semantically.

string search_text = "integrated pull-up bar";

string queryString = $@"
    SELECT TOP 15 c.id, c.Name, c.Brand, c.Description
    FROM c
    WHERE FullTextContains(c.Description, ""{search_text}"")
    ORDER BY RANK FullTextScore(c.Description, ""{search_text}"")
    ";

string reranking_context = "most economical with multiple pulley adjustmnets and ideal for home gyms";

List<string> documents = new List<string>();
FeedIterator<dynamic> resultSetIterator = container.GetItemQueryIterator<dynamic>(
    new QueryDefinition(queryString),
    requestOptions: new QueryRequestOptions()
    {
        MaxItemCount = 15,
    });

while (resultSetIterator.HasMoreResults)
{
    FeedResponse<dynamic> response = await resultSetIterator.ReadNextAsync();
    foreach (JsonElement item in response)
    {
        documents.Add(item.ToString());
    }
}

Dictionary<string, dynamic> options = new Dictionary<string, dynamic>
{
    { "return_documents", true },
    { "top_k", 10 },
    { "batch_size", 32 },
    { "sort", true }
};

SemanticRerankResult results = await container.SemanticRerankAsync(
    reranking_context,
    documents,
    options);

// get the best resulting document from the query
results.RerankScores.First().Document;
// or the index of the document in the original list
results.RerankScores.First().Index;
// or the reranking score 
results.RerankScores.First().Score;

// get the latency information from the reranking operation
Dictonary<string, object. latencyInfo = results.Latency;

// get the token usage information from the reranking operation
Dictonary<string, object> tokenUseageInfo = results.TokenUseage;
```

- [] New feature (non-breaking change which adds functionality)

## Closing issues

To automatically close an issue: closes #IssueNumber

Author: NaluTripician

Microsoft.Azure.Cosmos.Encryption.Custom/src/EncryptionContainer.cs

@@ -1008,6 +1008,21 @@ public override Task<bool> IsFeedRangePartOfAsync(
         }
 #endif
 
+#if PREVIEW && SDKPROJECTREF
+        public override Task<SemanticRerankResult> SemanticRerankAsync(
+            string rerankContext,
+            IEnumerable<string> documents,
+            IDictionary<string, object> options = null,
+            CancellationToken cancellationToken = default)
+        {
+            return this.container.SemanticRerankAsync(
+                rerankContext,
+                documents,
+                options,
+                cancellationToken);
+        }
+#endif
+
         private async Task<ResponseMessage> ReadManyItemsHelperAsync(
            IReadOnlyList<(string id, PartitionKey partitionKey)> items,
            ReadManyRequestOptions readManyRequestOptions = null,

Microsoft.Azure.Cosmos.Encryption/src/EncryptionContainer.cs

@@ -732,6 +732,21 @@ public override FeedIterator<T> GetItemQueryIterator<T>(
         }
 
 #if ENCRYPTIONPREVIEW
+#if SDKPROJECTREF
+        public override Task<SemanticRerankResult> SemanticRerankAsync(
+            string rerankContext,
+            IEnumerable<string> documents,
+            IDictionary<string, object> options = null,
+            CancellationToken cancellationToken = default)
+        {
+            return this.Container.SemanticRerankAsync(
+                rerankContext,
+                documents,
+                options,
+                cancellationToken);
+        }
+
+#endif
         public override async Task<ResponseMessage> DeleteAllItemsByPartitionKeyStreamAsync(
             Cosmos.PartitionKey partitionKey,
             RequestOptions requestOptions = null,

Microsoft.Azure.Cosmos/src/Authorization/AuthorizationTokenProvider.cs

@@ -52,6 +52,12 @@ public abstract ValueTask<string> GetUserAuthorizationTokenAsync(
             AuthorizationTokenType tokenType,
             ITrace trace);
 
+        public abstract ValueTask AddInferenceAuthorizationHeaderAsync(
+            INameValueCollection headersCollection,
+            Uri requestAddress,
+            string verb,
+            AuthorizationTokenType tokenType);
+
         public abstract void TraceUnauthorized(
             DocumentClientException dce,
             string authorizationToken,

Microsoft.Azure.Cosmos/src/Authorization/AuthorizationTokenProviderMasterKey.cs

@@ -214,6 +214,11 @@ private void Dispose(bool disposing)
             this.authKeyHashFunction = null;
         }
 
+        public override ValueTask AddInferenceAuthorizationHeaderAsync(INameValueCollection headersCollection, Uri requestAddress, string verb, AuthorizationTokenType tokenType)
+        {
+            throw new NotImplementedException("AddInferenceAuthorizationHeaderAsync is only valid for AAD");
+        }
+
         // Use C# finalizer syntax for finalization code.
         // This finalizer will run only if the Dispose method does not get called.
         // It gives your base class the opportunity to finalize.

Microsoft.Azure.Cosmos/src/Authorization/AuthorizationTokenProviderResourceToken.cs

@@ -92,6 +92,11 @@ private void Dispose(bool disposing)
             // Do nothing
         }
 
+        public override ValueTask AddInferenceAuthorizationHeaderAsync(INameValueCollection headersCollection, Uri requestAddress, string verb, AuthorizationTokenType tokenType)
+        {
+            throw new NotImplementedException("AddInferenceAuthorizationHeaderAsync is only valid for AAD");
+        }
+
         // Use C# finalizer syntax for finalization code.
         // This finalizer will run only if the Dispose method does not get called.
         // It gives your base class the opportunity to finalize.

Microsoft.Azure.Cosmos/src/Authorization/AuthorizationTokenProviderTokenCredential.cs

@@ -15,14 +15,18 @@ namespace Microsoft.Azure.Cosmos
 
     internal sealed class AuthorizationTokenProviderTokenCredential : AuthorizationTokenProvider
     {
+        private const string InferenceTokenPrefix = "Bearer ";
         internal readonly TokenCredentialCache tokenCredentialCache;
         private bool isDisposed = false;
 
+        internal readonly TokenCredential tokenCredential;
+
         public AuthorizationTokenProviderTokenCredential(
             TokenCredential tokenCredential,
             Uri accountEndpoint,
             TimeSpan? backgroundTokenCredentialRefreshInterval)
         {
+            this.tokenCredential = tokenCredential ?? throw new ArgumentNullException(nameof(tokenCredential));
             this.tokenCredentialCache = new TokenCredentialCache(
                 tokenCredential: tokenCredential,
                 accountEndpoint: accountEndpoint,
@@ -71,6 +75,21 @@ public override async ValueTask AddAuthorizationHeaderAsync(
             }
         }
 
+        public override async ValueTask AddInferenceAuthorizationHeaderAsync(
+            INameValueCollection headersCollection,
+            Uri requestAddress,
+            string verb,
+            AuthorizationTokenType tokenType)
+        {
+            using (Trace trace = Trace.GetRootTrace(nameof(GetUserAuthorizationTokenAsync), TraceComponent.Authorization, TraceLevel.Info))
+            {
+                string token = await this.tokenCredentialCache.GetTokenAsync(trace);
+
+                string inferenceToken = $"{InferenceTokenPrefix}{token}";
+                headersCollection.Add(HttpConstants.HttpHeaders.Authorization, inferenceToken);
+            }
+        }
+
         public override void TraceUnauthorized(
             DocumentClientException dce,
             string authorizationToken,

Microsoft.Azure.Cosmos/src/Authorization/AzureKeyCredentialAuthorizationTokenProvider.cs

@@ -125,5 +125,10 @@ private void CheckAndRefreshTokenProvider()
                 }
             }
         }
+
+        public override ValueTask AddInferenceAuthorizationHeaderAsync(INameValueCollection headersCollection, Uri requestAddress, string verb, AuthorizationTokenType tokenType)
+        {
+            throw new NotImplementedException("AddInferenceAuthorizationHeaderAsync is only valid for AAD");
+        }
     }
 }

Microsoft.Azure.Cosmos/src/Inference/InferenceService.cs

@@ -0,0 +1,209 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+    using System.Linq;
+    using System.Net.Http;
+    using System.Net.Http.Headers;
+    using System.Text;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using global::Azure.Core;
+    using Microsoft.Azure.Documents;
+    using Microsoft.Azure.Documents.Collections;
+
+    /// <summary>
+    /// Provides functionality to interact with the Cosmos DB Inference Service for semantic reranking.
+    /// </summary>
+    internal class InferenceService : IDisposable
+    {
+        // Base path for the inference service endpoint.
+        private const string basePath = "/inference/semanticReranking";
+        // User agent string for inference requests.
+        private const string inferenceUserAgent = "cosmos-inference-dotnet";
+        // Default scope for AAD authentication.
+        private const string inferenceServiceDefaultScope = "https://dbinference.azure.com/.default";
+        private const int inferenceServiceDefaultMaxConnectionLimit = 50;
+
+        private readonly int inferenceServiceMaxConnectionLimit;
+        private readonly string inferenceServiceBaseUrl;
+        private readonly Uri inferenceEndpoint;
+
+        private HttpClient httpClient;
+        private AuthorizationTokenProvider cosmosAuthorization;
+
+        private bool disposedValue;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="InferenceService"/> class.
+        /// </summary>
+        /// <param name="client">The CosmosClient instance.</param>
+        /// <exception cref="InvalidOperationException">Thrown if AAD authentication is not used.</exception>
+        public InferenceService(CosmosClient client)
+        {
+            this.inferenceServiceBaseUrl = ConfigurationManager.GetEnvironmentVariable<string>("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT", null);
+
+            if (string.IsNullOrEmpty(this.inferenceServiceBaseUrl))
+            {
+                throw new ArgumentNullException("Set environment variable AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT to use inference service");
+            }
+
+            this.inferenceServiceMaxConnectionLimit = ConfigurationManager.GetEnvironmentVariable<int?>(
+                "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_SERVICE_MAX_CONNECTION_LIMIT",
+                inferenceServiceDefaultMaxConnectionLimit) ?? inferenceServiceDefaultMaxConnectionLimit;
+
+            // Create and configure HttpClient for inference requests.
+            HttpMessageHandler httpMessageHandler = CosmosHttpClientCore.CreateHttpClientHandler(
+                        gatewayModeMaxConnectionLimit: this.inferenceServiceMaxConnectionLimit,
+                        webProxy: null,
+                        serverCertificateCustomValidationCallback: client.DocumentClient.ConnectionPolicy.ServerCertificateCustomValidationCallback);
+
+            this.httpClient = new HttpClient(httpMessageHandler);
+
+            this.CreateClientHelper(this.httpClient);
+
+            // Construct the inference service endpoint URI.
+            this.inferenceEndpoint = new Uri($"{this.inferenceServiceBaseUrl}/{basePath}");
+
+            // Ensure AAD authentication is used.
+            if (client.DocumentClient.cosmosAuthorization.GetType() != typeof(AuthorizationTokenProviderTokenCredential))
+            {
+                throw new InvalidOperationException("InferenceService only supports AAD authentication.");
+            }
+
+            // Set up token credential for authorization.
+            // This is done to ensure the correct scope, which is different than the scope of the client, is used for the inference service.
+            AuthorizationTokenProviderTokenCredential defaultOperationTokenProvider = client.DocumentClient.cosmosAuthorization as AuthorizationTokenProviderTokenCredential;
+            TokenCredential tokenCredential = defaultOperationTokenProvider.tokenCredential;
+
+            this.cosmosAuthorization = new AuthorizationTokenProviderTokenCredential(
+                tokenCredential: tokenCredential,
+                accountEndpoint: new Uri(inferenceServiceDefaultScope),
+                backgroundTokenCredentialRefreshInterval: client.ClientOptions?.TokenCredentialBackgroundRefreshInterval);
+        }
+
+        /// <summary>
+        /// Sends a semantic rerank request to the inference service.
+        /// </summary>
+        /// <param name="rerankContext">The context/query for reranking.</param>
+        /// <param name="documents">The documents to be reranked.</param>
+        /// <param name="options">Optional additional options for the request.</param>
+        /// <param name="cancellationToken">Cancellation token.</param>
+        /// <returns>A dictionary containing the reranked results.</returns>
+        public async Task<SemanticRerankResult> SemanticRerankAsync(
+            string rerankContext,
+            IEnumerable<string> documents,
+            IDictionary<string, object> options = null,
+            CancellationToken cancellationToken = default)
+        {
+            // Prepare HTTP request for semantic reranking.
+            HttpRequestMessage message = new HttpRequestMessage(HttpMethod.Post, this.inferenceEndpoint);
+            INameValueCollection additionalHeaders = new RequestNameValueCollection();
+            await this.cosmosAuthorization.AddInferenceAuthorizationHeaderAsync(
+                headersCollection: additionalHeaders,
+                this.inferenceEndpoint,
+                HttpConstants.HttpMethods.Post,
+                AuthorizationTokenType.AadToken);
+            additionalHeaders.Add(HttpConstants.HttpHeaders.UserAgent, inferenceUserAgent);
+
+            // Add all headers to the HTTP request.
+            foreach (string key in additionalHeaders.AllKeys())
+            {
+                message.Headers.Add(key, additionalHeaders[key]);
+            }
+
+            // Build the request payload.
+            Dictionary<string, object> body = this.AddSemanticRerankPayload(rerankContext, documents, options);
+
+            message.Content = new StringContent(
+                Newtonsoft.Json.JsonConvert.SerializeObject(body),
+                Encoding.UTF8,
+                RuntimeConstants.MediaTypes.Json);
+
+            // Send the request and ensure success.
+            HttpResponseMessage responseMessage = await this.httpClient.SendAsync(message, cancellationToken);
+            responseMessage.EnsureSuccessStatusCode();
+
+            // Deserialize and return the response content as a dictionary.
+            return await SemanticRerankResult.DeserializeSemanticRerankResultAsync(responseMessage);
+        }
+
+        /// <summary>
+        /// Configures the provided HttpClient with default headers and settings for inference requests.
+        /// </summary>
+        /// <param name="httpClient">The HttpClient to configure.</param>
+        private void CreateClientHelper(HttpClient httpClient)
+        {
+            httpClient.Timeout = TimeSpan.FromSeconds(120);
+            httpClient.DefaultRequestHeaders.CacheControl = new CacheControlHeaderValue { NoCache = true };
+
+            // Set requested API version header for version enforcement.
+            httpClient.DefaultRequestHeaders.Add(HttpConstants.HttpHeaders.Version,
+                HttpConstants.Versions.CurrentVersion);
+
+            httpClient.DefaultRequestHeaders.Add(HttpConstants.HttpHeaders.Accept, RuntimeConstants.MediaTypes.Json);
+        }
+
+        /// <summary>
+        /// Constructs the payload for the semantic rerank request.
+        /// </summary>
+        /// <param name="rerankContext">The context/query for reranking.</param>
+        /// <param name="documents">The documents to be reranked.</param>
+        /// <param name="options">Optional additional options.</param>
+        /// <returns>A dictionary representing the request payload.</returns>
+        private Dictionary<string, object> AddSemanticRerankPayload(string rerankContext, IEnumerable<string> documents, IDictionary<string, object> options)
+        {
+            Dictionary<string, object> payload = new Dictionary<string, object>
+            {
+                { "query", rerankContext },
+                { "documents", documents.ToArray() }
+            };
+
+            if (options == null)
+            {
+                return payload;
+            }
+
+            // Add any additional options to the payload.
+            foreach (string option in options.Keys)
+            {
+                payload.Add(option, options[option]);
+            }
+
+            return payload;
+        }
+
+        /// <summary>
+        /// Disposes managed resources used by the service.
+        /// </summary>
+        /// <param name="disposing">Indicates if called from Dispose.</param>
+        protected void Dispose(bool disposing)
+        {
+            if (!this.disposedValue)
+            {
+                if (disposing)
+                {
+                    this.httpClient.Dispose();
+                    this.cosmosAuthorization.Dispose();
+                    this.httpClient = null;
+                    this.cosmosAuthorization = null;
+                }
+
+                this.disposedValue = true;
+            }
+        }
+
+        /// <summary>
+        /// Disposes the service and its resources.
+        /// </summary>
+        public void Dispose()
+        {
+            this.Dispose(true);
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Inference/RerankScore.cs

@@ -0,0 +1,46 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    /// <summary>
+    /// Represents the score assigned to a document after a reranking operation.
+    /// </summary>
+#if PREVIEW
+    public
+#else
+    internal
+#endif
+
+    class RerankScore
+    {
+        /// <summary>
+        /// Gets the document content or identifier that was reranked.
+        /// </summary>
+        public object Document { get; }
+
+        /// <summary>
+        /// Gets the score assigned to the document after reranking.
+        /// </summary>
+        public double Score { get; }
+
+        /// <summary>
+        /// Gets the original index or position of the document before reranking.
+        /// </summary>
+        public int Index { get; }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="RerankScore"/> class.
+        /// </summary>
+        /// <param name="document">The document content or identifier.</param>
+        /// <param name="score">The reranked score for the document.</param>
+        /// <param name="index">The original index of the document.</param>
+        public RerankScore(object document, double score, int index)
+        {
+            this.Document = document;
+            this.Score = score;
+            this.Index = index;
+        }
+    }
+}