Démarrage rapide : Recherche vectorielle avec Java dans Azure DocumentDB

Apprenez à utiliser la recherche vectorielle dans Azure DocumentDB avec le pilote Java MongoDB pour stocker et interroger efficacement les données vectorielles.

Ce guide de démarrage rapide fournit une visite guidée des techniques de recherche vectorielle clés à l’aide d’un exemple d’application Java sur GitHub.

L’application utilise un exemple de jeu de données d’hôtel dans un fichier JSON avec des vecteurs précalculés à partir du text-embedding-3-small modèle, bien que vous puissiez également générer les vecteurs vous-même. Les données de l’hôtel incluent des noms d’hôtel, des emplacements, des descriptions et des incorporations vectorielles.

Prerequisites

  • Un abonnement Azure

Créer un fichier de données avec des vecteurs

  1. Créez un nouveau répertoire de données pour le fichier de données des hôtels :

    mkdir data

  2. Copiez le Hotels_Vector.jsonfichier de données brutes avec des vecteurs dans votre data répertoire.

Créer un projet Java

  1. Créez un répertoire frère pour votre projet, au même niveau que le répertoire de données, puis ouvrez-le dans Visual Studio Code :

    mkdir vector-search-quickstart
mkdir vector-search-quickstart/src
code vector-search-quickstart

  2. Créez un pom.xml fichier à la racine du projet avec le contenu suivant :

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

    <groupId>com.azure.documentdb.samples</groupId>
    <artifactId>vector-search-quickstart</artifactId>
    <version>1.0-SNAPSHOT</version>
        
    <properties>
        <maven.compiler.release>21</maven.compiler.release>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    </properties>
        
    <dependencies>
        <dependency>
            <groupId>org.mongodb</groupId>
            <artifactId>mongodb-driver-sync</artifactId>
            <version>5.6.2</version>
        </dependency>
        <dependency>
            <groupId>com.azure</groupId>
            <artifactId>azure-identity</artifactId>
            <version>1.18.1</version>
        </dependency>
        <dependency>
            <groupId>com.azure</groupId>
            <artifactId>azure-ai-openai</artifactId>
            <version>1.0.0-beta.16</version>
        </dependency>
        <dependency>
            <groupId>tools.jackson.core</groupId>
            <artifactId>jackson-databind</artifactId>
            <version>3.0.3</version>
        </dependency>
        <dependency>
            <groupId>org.slf4j</groupId>
            <artifactId>slf4j-nop</artifactId>
            <version>2.0.17</version>
            <scope>runtime</scope>
        </dependency>
    </dependencies>
</project>

    L’application utilise les dépendances Maven suivantes spécifiées dans le pom.xml:

    • mongodb-driver-sync: Pilote Java MongoDB officiel pour la connectivité et les opérations de base de données
    • azure-identity: Bibliothèque d’identités Azure pour l’authentification sans mot de passe avec l’ID Microsoft Entra
    • azure-ai-openai: Bibliothèque de client Azure OpenAI pour communiquer avec des modèles IA et créer des incorporations vectorielles
    • jackson-databind: bibliothèque de sérialisation et de désérialisation JSON
    • slf4j-nop: Liaison SLF4J sans opération pour supprimer la sortie de journalisation du pilote MongoDB

  3. Créez un .env fichier à la racine de votre projet pour les variables d’environnement :

    # Identity for local developer authentication with Azure CLI
AZURE_TOKEN_CREDENTIALS=AzureCliCredential

# Azure OpenAI Embedding Settings
AZURE_OPENAI_EMBEDDING_MODEL=text-embedding-3-small
AZURE_OPENAI_EMBEDDING_API_VERSION=2023-05-15
AZURE_OPENAI_EMBEDDING_ENDPOINT=
EMBEDDING_SIZE_BATCH=16

# Azure DocumentDB configuration
MONGO_CLUSTER_NAME=

# Data file
DATA_FILE_WITH_VECTORS=../data/Hotels_Vector.json
EMBEDDED_FIELD=DescriptionVector
EMBEDDING_DIMENSIONS=1536
LOAD_SIZE_BATCH=50

    Remplacez les valeurs de substitution dans le fichier .env par vos propres informations.

    • AZURE_OPENAI_EMBEDDING_ENDPOINT: URL du point de terminaison de la ressource Azure OpenAI.
    • MONGO_CLUSTER_NAME: Nom de votre ressource Azure DocumentDB.

  4. Chargez les variables d’environnement :

    set -a && source .env && set +a

  5. La structure du projet doit ressembler à ceci :

    data
└── Hotels_Vector.json
vector-search-quickstart
├── .env
├── pom.xml
└── src

Créez un DiskAnn.java fichier dans le src répertoire et collez le code suivant :

package com.azure.documentdb.samples;

import com.azure.ai.openai.OpenAIClient;
import com.azure.ai.openai.OpenAIClientBuilder;
import com.azure.ai.openai.models.EmbeddingsOptions;
import com.azure.identity.DefaultAzureCredentialBuilder;
import com.mongodb.ConnectionString;
import com.mongodb.MongoClientSettings;
import com.mongodb.MongoCredential;
import com.mongodb.client.AggregateIterable;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoClients;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.Indexes;
import org.bson.Document;
import tools.jackson.core.type.TypeReference;
import tools.jackson.databind.json.JsonMapper;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * Vector search sample using DiskANN index.
 */
public class DiskAnn {
    private static final String SAMPLE_QUERY = "quintessential lodging near running trails, eateries, retail";
    private static final String DATABASE_NAME = "Hotels";
    private static final String COLLECTION_NAME = "hotels_diskann";
    private static final String VECTOR_INDEX_NAME = "vectorIndex_diskann";

    private final JsonMapper jsonMapper = JsonMapper.builder().build();

    public static void main(String[] args) {
        new DiskAnn().run();
        System.exit(0);
    }

    public void run() {
        try (var mongoClient = createMongoClient()) {
            var openAIClient = createOpenAIClient();

            var database = mongoClient.getDatabase(DATABASE_NAME);
            var collection = database.getCollection(COLLECTION_NAME, Document.class);

            // Drop and recreate collection
            collection.drop();
            database.createCollection(COLLECTION_NAME);
            System.out.println("Created collection: " + COLLECTION_NAME);

            // Load and insert data
            var hotelData = loadHotelData();
            insertDataInBatches(collection, hotelData);

            // Create standard indexes
            createStandardIndexes(collection);

            // Create vector index
            createVectorIndex(database);

            // Perform vector search
            var queryEmbedding = createEmbedding(openAIClient, SAMPLE_QUERY);
            performVectorSearch(collection, queryEmbedding);

        } catch (Exception e) {
            System.err.println("Error: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private MongoClient createMongoClient() {
        var clusterName = System.getenv("MONGO_CLUSTER_NAME");
        var managedIdentityPrincipalId = System.getenv("AZURE_MANAGED_IDENTITY_PRINCIPAL_ID");
        var azureCredential = new DefaultAzureCredentialBuilder().build();

        MongoCredential.OidcCallback callback = (MongoCredential.OidcCallbackContext context) -> {
            var token = azureCredential.getToken(
                new com.azure.core.credential.TokenRequestContext()
                    .addScopes("https://ossrdbms-aad.database.windows.net/.default")
            ).block();

            if (token == null) {
                throw new RuntimeException("Failed to obtain Azure AD token");
            }

            return new MongoCredential.OidcCallbackResult(token.getToken());
        };

        var credential = MongoCredential.createOidcCredential(null)
            .withMechanismProperty("OIDC_CALLBACK", callback);

        var connectionString = new ConnectionString(
            String.format("mongodb+srv://%s@%s.mongocluster.cosmos.azure.com/?authMechanism=MONGODB-OIDC&tls=true&retrywrites=false&maxIdleTimeMS=120000",
                managedIdentityPrincipalId, clusterName)
        );

        var settings = MongoClientSettings.builder()
            .applyConnectionString(connectionString)
            .credential(credential)
            .build();

        return MongoClients.create(settings);
    }

    private OpenAIClient createOpenAIClient() {
        var endpoint = System.getenv("AZURE_OPENAI_EMBEDDING_ENDPOINT");
        var credential = new DefaultAzureCredentialBuilder().build();

        return new OpenAIClientBuilder()
            .endpoint(endpoint)
            .credential(credential)
            .buildClient();
    }

    private List<Map<String, Object>> loadHotelData() throws IOException {
        var dataFile = System.getenv("DATA_FILE_WITH_VECTORS");
        var filePath = Path.of(dataFile);

        System.out.println("Reading JSON file from " + filePath.toAbsolutePath());
        var jsonContent = Files.readString(filePath);

        return jsonMapper.readValue(jsonContent, new TypeReference<List<Map<String, Object>>>() {});
    }

    private void insertDataInBatches(MongoCollection<Document> collection, List<Map<String, Object>> hotelData) {
        var batchSizeStr = System.getenv("LOAD_SIZE_BATCH");
        var batchSize = batchSizeStr != null ? Integer.parseInt(batchSizeStr) : 100;
        var batches = partitionList(hotelData, batchSize);

        System.out.println("Processing in batches of " + batchSize + "...");

        for (int i = 0; i < batches.size(); i++) {
            var batch = batches.get(i);
            var documents = batch.stream()
                .map(Document::new)
                .toList();

            collection.insertMany(documents);
            System.out.println("Batch " + (i + 1) + " complete: " + documents.size() + " inserted");
        }
    }

    private void createStandardIndexes(MongoCollection<Document> collection) {
        collection.createIndex(Indexes.ascending("HotelId"));
        collection.createIndex(Indexes.ascending("Category"));
        collection.createIndex(Indexes.ascending("Description"));
        collection.createIndex(Indexes.ascending("Description_fr"));
    }

    private void createVectorIndex(MongoDatabase database) {
        var embeddedField = System.getenv("EMBEDDED_FIELD");
        var dimensionsStr = System.getenv("EMBEDDING_DIMENSIONS");
        var dimensions = dimensionsStr != null ? Integer.parseInt(dimensionsStr) : 1536;

        var indexDefinition = new Document()
            .append("createIndexes", COLLECTION_NAME)
            .append("indexes", List.of(
                new Document()
                    .append("name", VECTOR_INDEX_NAME)
                    .append("key", new Document(embeddedField, "cosmosSearch"))
                    .append("cosmosSearchOptions", new Document()
                        .append("kind", "vector-diskann")
                        .append("dimensions", dimensions)
                        .append("similarity", "COS")
                        .append("maxDegree", 20)
                        .append("lBuild", 10)
                    )
            ));

        database.runCommand(indexDefinition);
        System.out.println("Created vector index: " + VECTOR_INDEX_NAME);
    }

    private List<Double> createEmbedding(OpenAIClient openAIClient, String text) {
        var model = System.getenv("AZURE_OPENAI_EMBEDDING_MODEL");
        var options = new EmbeddingsOptions(List.of(text));

        var response = openAIClient.getEmbeddings(model, options);
        return response.getData().get(0).getEmbedding().stream()
                .map(Float::doubleValue)
                .toList();
    }

    private void performVectorSearch(MongoCollection<Document> collection, List<Double> queryEmbedding) {
        var embeddedField = System.getenv("EMBEDDED_FIELD");

        var searchStage = new Document("$search", new Document()
            .append("cosmosSearch", new Document()
                .append("vector", queryEmbedding)
                .append("path", embeddedField)
                .append("k", 5)
            )
        );

        var projectStage = new Document("$project", new Document()
            .append("score", new Document("$meta", "searchScore"))
            .append("document", "$$ROOT")
        );

        var pipeline = List.of(searchStage, projectStage);

        System.out.println("\nVector search results for: \"" + SAMPLE_QUERY + "\"");

        AggregateIterable<Document> results = collection.aggregate(pipeline);
        var rank = 1;

        for (var result : results) {
            var document = result.get("document", Document.class);
            var hotelName = document.getString("HotelName");
            var score = result.getDouble("score");
            System.out.printf("%d. HotelName: %s, Score: %.4f%n", rank++, hotelName, score);
        }
    }

    private static <T> List<List<T>> partitionList(List<T> list, int batchSize) {
        var partitions = new ArrayList<List<T>>();
        for (int i = 0; i < list.size(); i += batchSize) {
            partitions.add(list.subList(i, Math.min(i + batchSize, list.size())));
        }
        return partitions;
    }
}

Ce code effectue les tâches suivantes :

  • Crée une connexion sans mot de passe à Azure DocumentDB à l’aide DefaultAzureCredential du mécanisme OIDC MongoDB
  • Crée un client Azure OpenAI pour générer des incorporations
  • Supprime et recrée la collection, puis charge les données de l'hôtel à partir du fichier JSON par lots
  • Crée des index standard et un index vectoriel avec des options spécifiques à l’algorithme
  • Génère une incorporation pour un exemple de requête et exécute un pipeline de recherche d’agrégation
  • Imprime les cinq hôtels les plus correspondants avec leurs scores de similarité

S’authentifier auprès de Azure

Connectez-vous à Azure avant d’exécuter l’application afin qu’elle puisse accéder en toute sécurité aux ressources Azure.

Note

Vérifiez que vous êtes connecté avec les rôles de plan de données requis sur le compte Azure DocumentDB et la ressource Azure OpenAI.

az login

Le code utilise l’authentification de votre développeur local pour accéder à Azure DocumentDB et Azure OpenAI. Lorsque vous définissez AZURE_TOKEN_CREDENTIALS=AzureCliCredential, ce paramètre indique à la fonction d’utiliser les informations d’identification Azure CLI pour l’authentification de manière déterministe. L’authentification s’appuie sur DefaultAzureCredential à partir d’azure-identity pour rechercher vos informations d’identification Azure dans l’environnement. Apprenez-en davantage sur l’authentification des applications Java auprès des services Azure à l’aide de la bibliothèque d’identités Azure.

Créer l’application

Compilez l’application :

mvn clean compile

Exécutez la recherche DiskANN (Voisin proche le plus proche basé sur disque) :

mvn exec:java -Dexec.mainClass="com.azure.documentdb.samples.DiskAnn"

DiskANN est optimisé pour les jeux de données volumineux qui ne tiennent pas en mémoire, un stockage sur disque efficace, et un bon équilibre entre vitesse et précision.

Exemple de sortie :

Created collection: hotels_diskann
Reading JSON file from /workspaces/documentdb-samples/ai/vector-search-java/../data/Hotels_Vector.json
Processing in batches of 50...
Batch 1 complete: 50 inserted
Created vector index: vectorIndex_diskann

Vector search results for: "quintessential lodging near running trails, eateries, retail"
1. HotelName: Royal Cottage Resort, Score: 0.4991
2. HotelName: Country Comfort Inn, Score: 0.4786
3. HotelName: Nordick's Valley Motel, Score: 0.4635
4. HotelName: Economy Universe Motel, Score: 0.4462
5. HotelName: Roach Motel, Score: 0.4389

Afficher et gérer des données dans Visual Studio Code

  1. Installez l’extension DocumentDB et le pack d’extension pour Java dans Visual Studio Code.

  2. Connectez-vous à votre compte Azure DocumentDB à l’aide de l’extension DocumentDB.

  3. Affichez les données et les index dans la base de données Hotels.

    Capture d’écran de l’extension DocumentDB montrant la collection DocumentDB.

Nettoyer les ressources

Supprimez le groupe de ressources, le cluster Azure DocumentDB et la ressource Azure OpenAI lorsque vous n’en avez plus besoin pour éviter les coûts inutiles.

Contenu connexe

  • Last updated on