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

[mbhaskar]

Summary

Implements the semantic reranking/inference functionality ina new standalone module azure-cosmos-ai (com.azure:azure-cosmos-ai:1.0.0-beta.1) with package com.azure.cosmos.ai.

Motivation

The semantic reranking feature uses an external inference endpoint (not the Cosmos DB service itself) and relies on AAD authentication. Keeping it as a separate module:

• Avoids adding inference-specific dependencies to the core Cosmos SDK
• Follows Azure SDK conventions for standalone client libraries
• Uses azure-core HTTP pipeline (standard policies, retry, logging) instead of Cosmos-internal HTTP primitives
• Allows independent versioning and release cadence

Changes

New module: sdk/cosmos/azure-cosmos-ai

File	Purpose
CosmosAIClientBuilder	Builder with endpoint(), credential(), pipeline(), httpClient(), addPolicy()
CosmosAIAsyncClient	Async client exposing semanticRerank() returning Mono<SemanticRerankResult>
CosmosAIClient	Sync wrapper with configurable timeout
InferenceService	Internal implementation using HttpPipeline with retry logic (429/5xx, exponential backoff)
SemanticRerankResult / SemanticRerankScore	Public model classes
InferenceResponseParser	Internal JSON response parser

This is a change from earlier implemented design in azure-cosmos. So removed the below which were part of earlier commits

• CosmosAsyncContainer.semanticRerank() and CosmosContainer.semanticRerank()
• CosmosAsyncClient.getOrCreateInferenceService() and inference cleanup in close()
• ImplementationBridgeHelpers.SemanticRerankResultHelper / SemanticRerankScoreHelper
• ModelBridgeInternal initialize calls for semantic rerank types
• Configs.java inference endpoint properties
• Old InferenceService, model classes, test, and sample files

How to use the new module

Maven dependency

<dependency>
  <groupId>com.azure</groupId>
  <artifactId>azure-cosmos-ai</artifactId>
  <version>1.0.0-beta.1</version>
</dependency>

Quick start (async)

import com.azure.cosmos.ai.CosmosAIClientBuilder;
import com.azure.cosmos.ai.CosmosAIAsyncClient;
import com.azure.cosmos.ai.models.SemanticRerankResult;
import com.azure.identity.DefaultAzureCredentialBuilder;

// Build the client
CosmosAIAsyncClient client = new CosmosAIClientBuilder()
    .endpoint("https://<your-inference-endpoint>.dbinference.azure.com")
    .credential(new DefaultAzureCredentialBuilder().build())
    .buildAsyncClient();

// Rerank documents against a query
List<String> documents = Arrays.asList(
    "Azure Cosmos DB is a globally distributed database",
    "Azure Blob Storage stores unstructured data",
    "Azure Cosmos DB supports multiple consistency levels"
);

SemanticRerankResult result = client.semanticRerank(
    "How does Cosmos DB handle consistency?",
    documents,
    null  // optional options map
).block();

// Results are sorted by relevance score (highest first)
result.getScores().forEach(score ->
    System.out.printf("Index: %d, Score: %.4f, Doc: %s%n",
        score.getIndex(), score.getScore(), score.getDocument())
);

Quick start (sync)

import com.azure.cosmos.ai.CosmosAIClientBuilder;
import com.azure.cosmos.ai.CosmosAIClient;

CosmosAIClient client = new CosmosAIClientBuilder()
    .endpoint("https://<your-inference-endpoint>.dbinference.azure.com")
    .credential(new DefaultAzureCredentialBuilder().build())
    .buildClient();

SemanticRerankResult result = client.semanticRerank(
    "How does Cosmos DB handle consistency?",
    documents,
    null
);

Endpoint resolution

The endpoint can be provided via (in priority order):

1. builder.endpoint("...") — explicit
2. System property azure.cosmos.semanticReranker.inferenceEndpoint
3. Environment variable AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT

Testing

• 22 unit tests pass covering: constructor validation, happy-path parsing, retry logic (429/5xx), non-retryable errors (400/401/403), timeout handling, and request payload structure.

Verification

# Compile both modules
mvn compile -pl sdk/cosmos/azure-cosmos -DskipTests
mvn compile -pl sdk/cosmos/azure-cosmos-ai -DskipTests

# Run unit tests
mvn test -pl sdk/cosmos/azure-cosmos-ai -Dgroups=unit

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[aayush3011]

Overall it looks good, left a few comments.

Approved by mistake.

[aayush3011]

Deep Review — Semantic Reranking InferenceService

Summary: This PR adds semantic reranking via a new InferenceService class, with semanticRerank() on both CosmosAsyncContainer and CosmosContainer, retry logic, and unit tests. The core functionality is solid and the 22 prior review comments appear addressed. Below are new findings from a deeper review pass.

Severity	Count	Areas
🔴 Blocking	3	Resource leak, NPE, misleading Javadoc
🟡 Recommendation	5	API surface, validation, User-Agent, sample, test leak
🟢 Suggestion	1	Option constants
💬 Observation	1	Minor consolidation

[aayush3011]

LGTM. Thanks Bhaskar.

[kushagraThapar]

Thanks for the iteration here — caught up on the diff and verified that 41 of the 46 prior inline comments are fully resolved at HEAD 228b26965a7, including all the headline correctness items from earlier rounds (lazy-init race, missing close(), retry policy, AAD enforcement, leaked test executor, options key constants, etc.). Nice work landing those.

I have three correctness items I'd like to flag inline before this lands — all in the retry / lifecycle area. Happy to discuss any of them.

[ananth7592]

The four entries added to 4.81.0-beta.1 describe the prior in-tree design that was deleted later in this same PR (commit f405944):

Markdown

• Added beta semanticRerank API on CosmosAsyncContainer and CosmosContainer …
• Semantic rerank: serialization failures are now surfaced as a typed BadRequestException (HTTP 400, sub-status CUSTOM_SERIALIZER_EXCEPTION) …
• 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) …
• Semantic rerank: fixed a race between CosmosAsyncClient.close() and lazy initialization of the inference HTTP client …
  None of these are true at the head SHA:

There is no semanticRerank method on CosmosAsyncContainer / CosmosContainer (grep confirms).
The new azure-cosmos-ai InferenceService throws IllegalArgumentException on serialization failure, not BadRequestException(400, CUSTOM_SERIALIZER_EXCEPTION).
CosmosAIClient.semanticRerank calls Mono.block(timeout), which throws IllegalStateException, not RequestTimeoutException(408).
CosmosAsyncClient.close() no longer manages an inference HTTP client at all (the lazy-init path was deleted with the old design).

[mbhaskar]

AI quirks! Will refresh the changelong once I get the changes approved

[ananth7592]

Is this main() functional i.e. usable by customer for async semantics?

[ananth7592]

Do we need to add explicit Jackson dependency in this package or we rely on azure-core's transitive version?
The reason I ask is InferenceService.java, InferenceResponseParser.java

[ananth7592]

Tests on the Public API builder surface?

• CosmosAIClientBuilder endpoint resolution priority (explicit > sysprop > env > error)
• Missing-credential error path
• Custom pipeline() overriding credential() / httpClient() / additionalPolicies
• InferenceResponseParser edge cases (missing Scores, malformed score entries with missing index/score, empty body)
• CosmosAIClient (sync) block(timeout) behavior

[NaluTripician]

Review — azure-cosmos-ai semantic reranking module

Reviewed locally with a deep PR reviewer + an adversarial "skeptic" pass, plus a cross-SDK comparison against the .NET (azure-cosmos-dotnet-v3) and Python (azure-sdk-for-python) implementations. Highest-risk findings were independently verified against this PR's head and the peer-SDK source.

Overall: I don't think this is ready to merge yet. The inference/retry/parse core is reasonable and reasonably tested, but there's build-breaking residue from the abandoned in-core design, a changelog that documents things that don't ship, missing module scaffolding, and a public API surface that's below the Azure SDK bar for something that freezes at GA.

🔴 Blocking

B1 — Dead code + unused imports left in core azure-cosmos (will fail CI). The description says the container-level design was removed, but the diff still mutates core files with orphaned changes:

• CosmosAsyncClient.java: volatile boolean closed is written in close() but never read; tokenCredential is promoted to a field with a tokenCredential() accessor that nothing calls.
• CosmosAsyncContainer.java: unused imports (java.util.Map, AtomicReference) + whitespace churn.
• CosmosContainer.java: unused import java.util.Map.

UnusedImports is a build-failing checkstyle rule here, so the core module almost certainly fails CI. Suggest fully reverting all three core files — the new module shouldn't need any edits to azure-cosmos.

B2 — azure-cosmos/CHANGELOG.md documents features/bugfixes that don't exist in this PR. It advertises container-level semanticRerank APIs, a BadRequestException/CUSTOM_SERIALIZER_EXCEPTION mapping, a sync timeout→RequestTimeoutException (408) mapping, and a close() race fix — none of which ship in the new-module design. These are stale entries from the abandoned approach and will mislead users/release auditing. Remove them and document the new module in its own CHANGELOG.md.

B3 — New module missing required scaffolding. No CHANGELOG.md, README.md, module-info.java, or package-info.java for azure-cosmos-ai / its packages → fails repo CI gates.

🟡 High

• H1 — Stringly-typed Map<String,Object> options forwarded verbatim to the wire (InferenceService.buildRequestPayload): no compile-time safety, typos silently become wire fields, and a caller can overwrite reserved query/documents keys. This is the most consequential item since the surface freezes at GA. (Note: .NET and Python also use an untyped options bag, so it's a cross-SDK pattern — but a new azure-core-style module should prefer a typed SemanticRerankOptions with an allow-listed mapping.)
• H2 — Missing Azure SDK annotations/traits: no @ServiceClientBuilder / @ServiceClient / @ServiceMethod, no HttpTrait/TokenCredentialTrait/EndpointTrait; builder lacks retryPolicy/retryOptions/clientOptions. Normally an API-review blocker for a new client.
• H3 — No RetryPolicy in the pipeline; the hand-rolled Reactor retry only fires on HttpResponseException, so transient transport/IO errors (connection reset, Netty TimeoutException, DNS) are never retried — exactly the failures retry should absorb. Prefer azure-core's RetryPolicy, or extend the custom retry to cover transport errors.
• H4 — Sync block(timeout) uses the same value as the per-attempt .timeout(), so it can expire mid-retry and surfaces a raw IllegalStateException (not a typed 408). The changelog even claims a 408 mapping that isn't implemented. Make timeout_seconds a single whole-operation deadline and map timeout to a typed exception.
• H5 — close() is a misleading no-op while the clients are AutoCloseable and the Javadoc says "releases resources." A self-created Netty HttpClient is never disposed (leak in long-lived multi-client scenarios). Either track ownership and dispose what the SDK created, or drop AutoCloseable and fix the docs.
• H6 — Public "internal" parser: InferenceResponseParser is public in the public models package, so it's part of the shipped surface (APIView/revapi will flag it). Move it to implementation.
• H7 — getScores() returns null when Scores is absent — the PR's own sample (result.getScores().forEach(...)) would NPE. Default to an empty immutable list.

🔵 Cross-SDK comparison (.NET / Python)

Shared contract is identical across all three SDKs (verified): request {query, documents, ...options}; response top-level Scores (PascalCase) with index/score/document + lowercase latency/token_usage; AAD scope https://dbinference.azure.com/.default; path /inference/semanticReranking; env var AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT.

• ✅ The capital "Scores" key is correct, not a bug — confirmed against .NET's SemanticRerankResult.DeserializeSemanticRerankResultAsync (also reads "Scores") and Python test mocks. (Python's docstring mentioning relevance_score is just inaccurate docs; the wire field is score.)
• ✅ Score→document mapping is safe — the parser honors each item's explicit index, it does not positionally zip.

Behavioral drift to close before GA:

1. JSON-object document parsing loses data (Med/High): Java uses JsonNode.asText(), which returns "" for object nodes; .NET uses GetRawText(). This breaks the document_type="json" path. Use .toString() for non-textual nodes and add a JSON-document test.
2. Timeout not mapped to a typed error (Med): Java leaks a raw reactor TimeoutException; both .NET and Python normalize to 408.
3. Retry set diverges while claiming parity (Med): Java retries {429,500,502,503}; the comment says "matches Python" but Python is {429,500}; .NET does no retries. Pick a canonical set and fix the comment.
4. Default timeout mismatch (Med): Java 120s vs .NET 5s vs Python (from connection policy).
5. Design divergence (intentional): .NET/Python are container-scoped methods reusing the client credential; Java is a clean standalone module with explicit TokenCredential + azure-core pipeline — more idiomatic, but the source of the drift above.
6. Java omits Headers on the result (both peers expose them) and the connection-limit knob (.NET). Don't copy .NET's public typo TokenUseage — Java's getTokenUsage() is correct.

🟢 Lower-priority

• Use azure-core ClientLogger instead of SLF4J; avoid logging full error bodies / the endpoint at info level.
• Endpoint validation is weak: URI.create throws IllegalArgumentException (Javadoc says IllegalStateException); no https/absolute check; endpoint + BASE_PATH breaks on a trailing slash.
• pom.xml: jacoco thresholds (~0.09 line / 0.02 branch) effectively disable coverage; spotless.skip=true; the bannedDependencies rule is a no-op as written; azure-core-http-netty is at compile scope (consider runtime).

Test gaps

22 unit tests give strong retry coverage but zero public-surface coverage: no CosmosAIClientBuilder tests (endpoint priority, missing endpoint/credential, pipeline() override), no sync CosmosAIClient/resolveBlockTimeout tests, no Retry-After header test, no test that the per-request timeout fires, and no document_type="json" round-trip (which would catch finding #1). No integration/emulator test (.NET and Python both have them).

---

_{🤖 Generated with a multi-agent local review (deep reviewer + skeptic + cross-SDK).}

[kushagraThapar]

Suggestion: have azure-cosmos-ai wrap an existing CosmosAsyncClient rather than standing up a second client + credential

As implemented, CosmosAIClientBuilder builds a fully standalone client that needs its own endpoint and its own TokenCredential (CosmosAIClientBuilder.java#L141-L205), with no link to a CosmosAsyncClient. So a customer who already has a Cosmos client configured (credential, endpoint, HTTP pipeline, retry/throttling, telemetry) has to construct a second, unrelated client and re-supply a second credential just to call semanticRerank.

That runs against the consistent convention across the Java + Cosmos ecosystem, where every integration that layers functionality on top of Cosmos takes the user's existing CosmosAsyncClient rather than building its own:

Library	How it takes the client
azure-spring-data-cosmos (this repo)	CosmosFactory(CosmosAsyncClient cosmosAsyncClient, String databaseName)
Spring AI — CosmosDBVectorStore	CosmosDBVectorStore.builder(CosmosAsyncClient cosmosClient, EmbeddingModel embeddingModel)
Spring AI — Cosmos chat-memory repository	CosmosDBChatMemoryRepositoryConfig.Builder.withCosmosClient(CosmosAsyncClient cosmosClient)
LangChain4j — AzureCosmosDbNoSqlEmbeddingStore	.builder().cosmosClient(CosmosAsyncClient cosmosClient)

Suggested approach: keep azure-cosmos-ai as a separate module, but have the builder accept the user's existing client so it can reuse the already-configured credential, pipeline, and endpoint resolution:

CosmosAIAsyncClient ai = new CosmosAIClientBuilder()
    .cosmosAsyncClient(existingClient)   // reuse credential + pipeline
    .buildAsyncClient();

ai.semanticRerank(query, documents);

This keeps the dependency-isolation and independent-versioning benefits that motivated the separate module, while restoring the "wrap the existing CosmosClient" convention the rest of the Java ecosystem follows — so customers configure one credential, in one place. It would require azure-cosmos to expose the credential/pipeline to the companion module via a bridge, which is an established pattern in this codebase.

Two smaller, related notes:

• The endpoint() + credential() + system-property / env-var fallback resolution is largely re-deriving configuration the Cosmos client already holds; wrapping the client removes most of that surface (and the associated tests).