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(rerankContext: string, documents: string[], options?: SemanticRerankOptions): Promise<SemanticRerankResult>;
     // (undocumented)
     upsert<T, U = T>(input: {
         body: T;
@@ -731,6 +732,7 @@ export class Container {
     readPartitionKeyRanges(feedOptions?: FeedOptions): QueryIterator<PartitionKeyRange>;
     replace(body: ContainerDefinition, options?: RequestOptions): Promise<ContainerResponse>;
     get scripts(): Scripts;
+    semanticRerank(rerankContext: string, documents: string[], options?: SemanticRerankOptions): Promise<SemanticRerankResult>;
     get url(): string;
 }
 
@@ -2133,6 +2135,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 +2386,26 @@ export class Scripts {
     get userDefinedFunctions(): UserDefinedFunctions;
 }
 
+// @public
+export interface SemanticRerankOptions {
+    abortSignal?: AbortSignal;
+    additionalOptions?: Record<string, unknown>;
+    batchSize?: number;
+    documentType?: string;
+    returnDocuments?: boolean;
+    sort?: boolean;
+    targetPaths?: string;
+    topK?: number;
+}
+
+// @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>;
 

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(
+    rerankContext: string,
+    documents: string[],
+    options?: SemanticRerankOptions,
+  ): Promise<SemanticRerankResult> {
+    const service = this.getOrCreateInferenceService();
+    return service.semanticRerank(rerankContext, 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/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,64 @@ 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 reranking 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. Set the `AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT` environment variable
+   *
+   * @param rerankContext - 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 settings for the reranking request.
+   * @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,
+   *   { returnDocuments: true, topK: 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(
+    rerankContext: string,
+    documents: string[],
+    options?: SemanticRerankOptions,
+  ): Promise<SemanticRerankResult> {
+    return this.clientContext.semanticRerank(rerankContext, documents, options);
+  }
+
   /**
    * @internal
    */

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