feat(cosmos): add semantic rerank API

.vscode/cspell.json

@@ -320,6 +320,8 @@
         "Parition",
         "colls",
         "pkranges",
+        "rerank",
+        "Rerank",
         "sproc",
         "sprocs",
         "udfs",

sdk/cosmosdb/cosmos/review/cosmos-node.api.md

@@ -317,6 +317,7 @@ export class ClientContext {
         diagnosticNode: DiagnosticNodeInternal;
         partitionKeyRangeId?: string;
     }): Promise<Response_2<T & Resource>>;
+    semanticRerank(context: string, documents: string[], options?: SemanticRerankOptions): Promise<SemanticRerankResult>;
     // (undocumented)
     upsert<T, U = T>(input: {
         body: T;
@@ -700,6 +701,11 @@ export const Constants: {
     DefaultEncryptionCacheTimeToLiveInSeconds: number;
     EncryptionCacheRefreshIntervalInMs: number;
     RequestTimeoutForReadsInMs: number;
+    InferenceBasePath: string;
+    InferenceUserAgent: string;
+    InferenceDefaultScope: string;
+    InferenceDefaultTimeoutMs: number;
+    InferenceEndpointEnvVar: string;
 };
 
 // @public
@@ -731,6 +737,7 @@ export class Container {
     readPartitionKeyRanges(feedOptions?: FeedOptions): QueryIterator<PartitionKeyRange>;
     replace(body: ContainerDefinition, options?: RequestOptions): Promise<ContainerResponse>;
     get scripts(): Scripts;
+    semanticRerank(context: string, documents: string[], options?: SemanticRerankOptions): Promise<SemanticRerankResult>;
     get url(): string;
 }
 
@@ -824,6 +831,7 @@ export interface CosmosClientOptions {
     diagnosticLevel?: CosmosDbDiagnosticLevel;
     endpoint?: string;
     httpClient?: HttpClient;
+    inferenceEndpoint?: string;
     key?: string;
     permissionFeed?: PermissionDefinition[];
     resourceTokens?: {
@@ -2133,6 +2141,13 @@ export interface RequestOptions extends SharedOptions {
     urlConnection?: string;
 }
 
+// @public
+export interface RerankScore {
+    document: string | null;
+    index: number;
+    score: number;
+}
+
 // @public (undocumented)
 export interface Resource {
     _etag: string;
@@ -2377,6 +2392,17 @@ export class Scripts {
     get userDefinedFunctions(): UserDefinedFunctions;
 }
 
+// @public
+export type SemanticRerankOptions = Record<string, unknown>;
+
+// @public
+export interface SemanticRerankResult {
+    headers: Record<string, string>;
+    latency: Record<string, unknown> | undefined;
+    rerankScores: RerankScore[];
+    tokenUsage: Record<string, unknown> | undefined;
+}
+
 // @public
 export function setAuthorizationTokenHeaderUsingMasterKey(verb: HTTPMethod, resourceId: string, resourceType: ResourceType, headers: CosmosHeaders, masterKey: string): Promise<void>;
 
@@ -2462,6 +2488,8 @@ export interface StatusCodesType {
     // (undocumented)
     MethodNotAllowed: 405;
     // (undocumented)
+    MultipleChoices: 300;
+    // (undocumented)
     MultiStatus: 207;
     // (undocumented)
     NoContent: 204;

sdk/cosmosdb/cosmos/src/ClientContext.ts

@@ -51,6 +51,9 @@ import {
   AAD_AUTH_PREFIX,
   AAD_RESOURCE_NOT_FOUND_ERROR,
 } from "./common/constants.js";
+import { InferenceService } from "./inference/InferenceService.js";
+import type { SemanticRerankOptions } from "./inference/SemanticRerankOptions.js";
+import type { SemanticRerankResult } from "./inference/SemanticRerankResult.js";
 
 const logger: AzureLogger = createClientLogger("ClientContext");
 
@@ -70,6 +73,7 @@ export class ClientContext {
   public partitionKeyRangeCache: PartitionKeyRangeCache;
   /** boolean flag to support operations with client-side encryption */
   public enableEncryption: boolean = false;
+  private inferenceService: InferenceService | null = null;
 
   public constructor(
     private cosmosClientOptions: CosmosClientOptions,
@@ -1108,4 +1112,45 @@ export class ClientContext {
       this.globalEndpointManager.lastKnownPPCBEnabled
     );
   }
+
+  /**
+   * Rerank a list of documents using semantic reranking via the Cosmos DB Inference Service.
+   * This method uses a semantic reranker to score and reorder the provided documents
+   * based on their relevance to the given reranking context.
+   *
+   * The semantic reranking requests use a separate HTTP pipeline and do not use
+   * the default SDK retry policies.
+   *
+   * @param rerankContext - The context (e.g. query string) to use for reranking.
+   * @param documents - The documents to be reranked.
+   * @param options - Optional settings for the reranking request.
+   * @returns The reranking results including scores, latency, and token usage.
+   */
+  public async semanticRerank(
+    context: string,
+    documents: string[],
+    options?: SemanticRerankOptions,
+  ): Promise<SemanticRerankResult> {
+    const service = this.getOrCreateInferenceService();
+    return service.semanticRerank(context, documents, options);
+  }
+
+  /**
+   * Gets or lazily creates the InferenceService instance.
+   * @internal
+   */
+  private getOrCreateInferenceService(): InferenceService {
+    if (!this.inferenceService) {
+      this.inferenceService = new InferenceService(this.cosmosClientOptions);
+    }
+    return this.inferenceService;
+  }
+
+  /**
+   * Disposes the InferenceService if it was created.
+   * @internal
+   */
+  public disposeInferenceService(): void {
+    this.inferenceService = null;
+  }
 }

sdk/cosmosdb/cosmos/src/CosmosClient.ts

@@ -355,6 +355,7 @@ export class CosmosClient {
     if (this.globalPartitionEndpointManager) {
       this.globalPartitionEndpointManager.dispose();
     }
+    this.clientContext.disposeInferenceService();
   }
 
   private async backgroundRefreshEndpointList(

sdk/cosmosdb/cosmos/src/CosmosClientOptions.ts

@@ -81,4 +81,10 @@ export interface CosmosClientOptions {
 
   /** An optional parameter that represents the connection string. Your database connection string can be found in the Azure Portal. */
   connectionString?: string;
+
+  /**
+   * The endpoint URL for the Cosmos DB Inference Service, used for features such as semantic reranking.
+   * If not provided, the SDK falls back to the `AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT` environment variable.
+   */
+  inferenceEndpoint?: string;
 }

sdk/cosmosdb/cosmos/src/client/Container/Container.ts

@@ -44,6 +44,8 @@ import { MetadataLookUpType } from "../../CosmosDiagnostics.js";
 import type { EncryptionSettingForProperty } from "../../encryption/index.js";
 import { EncryptionProcessor } from "../../encryption/index.js";
 import type { EncryptionManager } from "../../encryption/EncryptionManager.js";
+import type { SemanticRerankOptions } from "../../inference/SemanticRerankOptions.js";
+import type { SemanticRerankResult } from "../../inference/SemanticRerankResult.js";
 
 /**
  * Operations for reading, replacing, or deleting a specific, existing container by id.
@@ -691,6 +693,74 @@ export class Container {
     }
   }
 
+  /**
+   * Rerank a list of documents using semantic reranking via the Cosmos DB Inference Service.
+   * This method uses a semantic reranker to score and reorder the provided documents
+   * based on their relevance to the given context.
+   *
+   * The semantic reranking requests use a separate HTTP pipeline from the main Cosmos DB client
+   * and do not use the default SDK retry policies.
+   *
+   * To use this feature, you must:
+   * 1. Configure AAD authentication via `aadCredentials` in `CosmosClientOptions`
+   * 2. Provide the inference endpoint via `inferenceEndpoint` in `CosmosClientOptions`,
+   *    or set the `AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT` environment variable
+   *
+   * @param context - The context (e.g. query string) to use for reranking the documents.
+   * @param documents - A list of documents (as JSON strings) to be reranked.
+   * @param options - Optional dictionary of settings for the reranking request.
+   *   Known service options:
+   *   - `return_documents` (boolean) — include reranked documents in the response.
+   *   - `top_k` (number) — max number of top-ranked documents to return.
+   *   - `batch_size` (number) — batch size for processing documents.
+   *   - `sort` (boolean) — sort results by relevance score in descending order.
+   *   - `document_type` (`"string"` | `"json"`) — type of documents being reranked.
+   *   - `target_paths` (string) — comma-separated JSON paths (when document_type is `"json"`).
+   *   - `abortSignal` (AbortSignal) — signal to cancel the request.
+   *   Any additional keys are forwarded as-is to the inference service.
+   * @returns The reranking results including scored documents, latency, and token usage.
+   *
+   * @example Semantic reranking of query results
+   * ```ts snippet:ContainerSemanticRerank
+   * import { DefaultAzureCredential } from "@azure/identity";
+   * import { CosmosClient } from "@azure/cosmos";
+   *
+   * const endpoint = "https://your-account.documents.azure.com";
+   * const aadCredentials = new DefaultAzureCredential();
+   * const client = new CosmosClient({
+   *   endpoint,
+   *   aadCredentials,
+   * });
+   *
+   * const { database } = await client.databases.createIfNotExists({ id: "Test Database" });
+   * const { container } = await database.containers.createIfNotExists({ id: "Test Container" });
+   *
+   * const queryResults = ["doc1 JSON", "doc2 JSON", "doc3 JSON"];
+   * const result = await container.semanticRerank(
+   *   "most economical with multiple adjustments",
+   *   queryResults,
+   *   { return_documents: true, top_k: 10, sort: true },
+   * );
+   * // Access the top-ranked document
+   * if (result.rerankScores.length > 0) {
+   *   const topResult = result.rerankScores[0];
+   *   const topScore = topResult.score;
+   *   const topDocument = topResult.document;
+   *   if (topDocument !== null) {
+   *     console.log("Top-ranked document:", topDocument);
+   *   }
+   *   console.log("Top score:", topScore);
+   * }
+   * ```
+   */
+  public async semanticRerank(
+    context: string,
+    documents: string[],
+    options?: SemanticRerankOptions,
+  ): Promise<SemanticRerankResult> {
+    return this.clientContext.semanticRerank(context, documents, options);
+  }
+
   /**
    * @internal
    */

sdk/cosmosdb/cosmos/src/common/constants.ts

@@ -304,6 +304,13 @@ export const Constants = {
   EncryptionCacheRefreshIntervalInMs: 60000, // 1 minute
 
   RequestTimeoutForReadsInMs: 2000, // 2 seconds
+
+  // Inference Service
+  InferenceBasePath: "/inference/semanticReranking",
+  InferenceUserAgent: "cosmos-inference-js",
+  InferenceDefaultScope: "https://dbinference.azure.com/.default",
+  InferenceDefaultTimeoutMs: 120_000, // 120 seconds
+  InferenceEndpointEnvVar: "AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT",
 };
 
 export const AAD_DEFAULT_SCOPE = "https://cosmos.azure.com/.default";

sdk/cosmosdb/cosmos/src/common/statusCodes.ts

@@ -11,6 +11,9 @@ export interface StatusCodesType {
   Accepted: 202;
   NoContent: 204;
   MultiStatus: 207;
+
+  // Redirection
+  MultipleChoices: 300;
   NotModified: 304;
 
   // Client error
@@ -50,6 +53,9 @@ export const StatusCodes: StatusCodesType = {
   Accepted: 202,
   NoContent: 204,
   MultiStatus: 207,
+
+  // Redirection
+  MultipleChoices: 300,
   NotModified: 304,
 
   // Client error

sdk/cosmosdb/cosmos/src/index.ts

@@ -164,3 +164,9 @@ export {
   type CosmosEncryptedNumber,
   CosmosEncryptedNumberType,
 } from "./encryption/index.js";
+
+export type {
+  RerankScore,
+  SemanticRerankResult,
+  SemanticRerankOptions,
+} from "./inference/index.js";