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

eng/versioning/version_client.txt

@@ -111,6 +111,7 @@ com.azure:azure-cosmos-benchmark;4.0.1-beta.1;4.0.1-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark_3;0.0.1-beta.1;0.0.1-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark_3-5;0.0.1-beta.1;0.0.1-beta.1
 com.azure:azure-cosmos-encryption;2.29.0;2.30.0-beta.1
+com.azure:azure-cosmos-ai;1.0.0-beta.1;1.0.0-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark-account-data-resolver-sample;1.0.0-beta.1;1.0.0-beta.1
 com.azure:azure-cosmos-test;1.0.0-beta.18;1.0.0-beta.19
 com.azure.cosmos.spark:azure-cosmos-spark_3-3_2-12;4.48.0;4.49.0-beta.1

sdk/cosmos/azure-cosmos-ai/pom.xml

@@ -0,0 +1,183 @@
+<!--
+Copyright (c) Microsoft Corporation. All rights reserved.
+Licensed under the MIT License.
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>com.azure</groupId>
+    <artifactId>azure-client-sdk-parent</artifactId>
+    <version>1.7.0</version> <!-- {x-version-update;com.azure:azure-client-sdk-parent;current} -->
+    <relativePath>../../parents/azure-client-sdk-parent</relativePath>
+  </parent>
+
+  <groupId>com.azure</groupId>
+  <artifactId>azure-cosmos-ai</artifactId>
+  <version>1.0.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos-ai;current} -->
+  <name>Microsoft Azure Cosmos DB AI Services SDK</name>
+  <description>This Package contains Microsoft Azure Cosmos DB AI Services SDK for semantic reranking and inference operations</description>
+  <packaging>jar</packaging>
+  <url>https://github.com/Azure/azure-sdk-for-java</url>
+
+  <distributionManagement>
+    <site>
+      <id>azure-java-build-docs</id>
+      <url>${site.url}/site/${project.artifactId}</url>
+    </site>
+  </distributionManagement>
+
+  <scm>
+    <url>scm:git:https://github.com/Azure/azure-sdk-for-java</url>
+    <connection>scm:git:git@github.com:Azure/azure-sdk-for-java.git</connection>
+    <tag>HEAD</tag>
+  </scm>
+
+  <properties>
+    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+    <jacoco.min.linecoverage>0.09</jacoco.min.linecoverage>
+    <jacoco.min.branchcoverage>0.02</jacoco.min.branchcoverage>
+
+    <!-- CosmosSkip - This can only be enabled once we release GA, as it needs a stable version to check for breaking changes. -->
+    <revapi.skip>true</revapi.skip>
+    <checkstyle.excludes>**/implementation/**/*.java</checkstyle.excludes>
+    <checkstyle.includeTestSourceDirectory>false</checkstyle.includeTestSourceDirectory>
+    <doclintMissingInclusion>-</doclintMissingInclusion>
+    <spotless.skip>true</spotless.skip>
+  </properties>
+
+  <dependencies>
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-core</artifactId>
+      <version>1.58.0</version> <!-- {x-version-update;com.azure:azure-core;dependency} -->
+    </dependency>
+
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-core-http-netty</artifactId>
+      <version>1.16.4</version> <!-- {x-version-update;com.azure:azure-core-http-netty;dependency} -->
+    </dependency>
+
+    <!-- Added this provided dependency to include necessary annotations used by "reactor-core".
+         Without this dependency, javadoc throws a warning as it cannot find enum When.MAYBE
+         which is used in @Nullable annotation in reactor core classes.
+
+         Similar provided dependency exits for "azure-core" as well.
+    -->
+    <dependency>
+      <groupId>com.google.code.findbugs</groupId>
+      <artifactId>jsr305</artifactId>
+      <version>3.0.2</version> <!-- {x-version-update;com.google.code.findbugs:jsr305;external_dependency} -->
+      <scope>provided</scope>
+    </dependency>
+
+    <!-- Test dependencies -->
+    <dependency>
+      <groupId>io.projectreactor</groupId>
+      <artifactId>reactor-test</artifactId>
+      <version>3.7.17</version> <!-- {x-version-update;io.projectreactor:reactor-test;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.testng</groupId>
+      <artifactId>testng</artifactId>
+      <version>7.5.1</version> <!-- {x-version-update;org.testng:testng;external_dependency} -->
+      <scope>test</scope>
+      <exclusions>
+        <exclusion>
+          <groupId>org.apache.ant</groupId>
+          <artifactId>ant</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>org.yaml</groupId>
+          <artifactId>snakeyaml</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+
+    <dependency>
+      <groupId>org.assertj</groupId>
+      <artifactId>assertj-core</artifactId>
+      <version>3.27.7</version> <!-- {x-version-update;org.assertj:assertj-core;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <version>4.11.0</version> <!-- {x-version-update;org.mockito:mockito-core;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy</artifactId>
+      <version>1.17.7</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy-agent</artifactId>
+      <version>1.17.7</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy-agent;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.logging.log4j</groupId>
+      <artifactId>log4j-slf4j-impl</artifactId>
+      <version>2.17.2</version> <!-- {x-version-update;org.apache.logging.log4j:log4j-slf4j-impl;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.logging.log4j</groupId>
+      <artifactId>log4j-api</artifactId>
+      <version>2.17.2</version> <!-- {x-version-update;org.apache.logging.log4j:log4j-api;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+
+    <dependency>
+      <groupId>org.apache.logging.log4j</groupId>
+      <artifactId>log4j-core</artifactId>
+      <version>2.17.2</version> <!-- {x-version-update;org.apache.logging.log4j:log4j-core;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+
+  <build>
+    <plugins>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-surefire-plugin</artifactId>
+        <version>3.5.3</version> <!-- {x-version-update;org.apache.maven.plugins:maven-surefire-plugin;external_dependency} -->
+        <configuration>
+          <groups>unit</groups>
+          <includes>
+            <include>%regex[.*]</include>
+          </includes>
+          <properties>
+            <property>
+              <name>surefire.testng.verbose</name>
+              <value>2</value>
+            </property>
+          </properties>
+        </configuration>
+      </plugin>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-enforcer-plugin</artifactId>
+        <version>3.6.1</version> <!-- {x-version-update;org.apache.maven.plugins:maven-enforcer-plugin;external_dependency} -->
+        <configuration>
+          <rules>
+            <bannedDependencies>
+              <includes>
+                <include>com.azure:*</include>
+              </includes>
+            </bannedDependencies>
+          </rules>
+        </configuration>
+      </plugin>
+    </plugins>
+  </build>
+</project>

sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIAsyncClient.java

@@ -0,0 +1,78 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai;
+
+import com.azure.cosmos.ai.implementation.InferenceService;
+import com.azure.cosmos.ai.models.SemanticRerankResult;
+import reactor.core.publisher.Mono;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Asynchronous client for Azure Cosmos DB AI services.
+ *
+ * <p>Provides semantic reranking of documents using the Azure Cosmos DB inference service.
+ * Use {@link CosmosAIClientBuilder} to create an instance.</p>
+ *
+ * <p><strong>Example:</strong></p>
+ * <pre>
+ * CosmosAIAsyncClient client = new CosmosAIClientBuilder()
+ *     .endpoint("https://my-inference.dbinference.azure.com")
+ *     .credential(new DefaultAzureCredentialBuilder().build())
+ *     .buildAsyncClient();
+ *
+ * client.semanticRerank("What is the capital of France?", documents, options)
+ *     .subscribe(result -&gt; {
+ *         result.getScores().forEach(s -&gt; System.out.println(s.getIndex() + ": " + s.getScore()));
+ *     });
+ * </pre>
+ */
+public final class CosmosAIAsyncClient implements AutoCloseable {
+
+    private final InferenceService inferenceService;
+
+    /**
+     * Creates a new CosmosAIAsyncClient.
+     *
+     * @param inferenceService The inference service instance.
+     */
+    CosmosAIAsyncClient(InferenceService inferenceService) {
+        this.inferenceService = inferenceService;
+    }
+
+    /**
+     * Performs semantic reranking of documents using the inference service.
+     *
+     * @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}.
+     */
+    public Mono<SemanticRerankResult> semanticRerank(
+        String rerankContext,
+        List<String> documents,
+        Map<String, Object> options) {
+
+        return inferenceService.semanticRerank(rerankContext, documents, options);
+    }
+
+    /**
+     * Closes this client and releases resources.
+     */
+    @Override
+    public void close() {
+        inferenceService.close();
+    }
+}

sdk/cosmos/azure-cosmos-ai/src/main/java/com/azure/cosmos/ai/CosmosAIClient.java

@@ -0,0 +1,90 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.ai;
+
+import com.azure.cosmos.ai.implementation.InferenceService;
+import com.azure.cosmos.ai.models.SemanticRerankResult;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Synchronous client for Azure Cosmos DB AI services.
+ *
+ * <p>Provides semantic reranking of documents using the Azure Cosmos DB inference service.
+ * Use {@link CosmosAIClientBuilder} to create an instance.</p>
+ *
+ * <p><strong>Example:</strong></p>
+ * <pre>
+ * CosmosAIClient client = new CosmosAIClientBuilder()
+ *     .endpoint("https://my-inference.dbinference.azure.com")
+ *     .credential(new DefaultAzureCredentialBuilder().build())
+ *     .buildClient();
+ *
+ * SemanticRerankResult result = client.semanticRerank("What is the capital of France?", documents, options);
+ * result.getScores().forEach(s -&gt; System.out.println(s.getIndex() + ": " + s.getScore()));
+ * </pre>
+ */
+public final class CosmosAIClient implements AutoCloseable {
+
+    private final CosmosAIAsyncClient asyncClient;
+
+    /**
+     * Creates a new CosmosAIClient wrapping the given async client.
+     *
+     * @param asyncClient The async client to wrap.
+     */
+    CosmosAIClient(CosmosAIAsyncClient asyncClient) {
+        this.asyncClient = asyncClient;
+    }
+
+    /**
+     * 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)}.</p>
+     *
+     * @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.
+     * @throws RuntimeException if the operation fails or times out.
+     */
+    public SemanticRerankResult semanticRerank(
+        String rerankContext,
+        List<String> documents,
+        Map<String, Object> options) {
+
+        Duration blockTimeout = resolveBlockTimeout(options);
+        return asyncClient.semanticRerank(rerankContext, documents, options)
+            .block(blockTimeout);
+    }
+
+    /**
+     * Closes this client and releases resources.
+     */
+    @Override
+    public void close() {
+        asyncClient.close();
+    }
+
+    private static Duration resolveBlockTimeout(Map<String, Object> options) {
+        if (options != null) {
+            Object timeoutVal = options.get(InferenceService.OPTION_TIMEOUT_SECONDS);
+            if (timeoutVal instanceof Number) {
+                double seconds = ((Number) timeoutVal).doubleValue();
+                if (seconds > 0) {
+                    return Duration.ofMillis((long) (seconds * 1000));
+                }
+            }
+        }
+        return InferenceService.DEFAULT_REQUEST_TIMEOUT;
+    }
+}