Início Rápido: Pesquisa de vetor com Java no Azure DocumentDB

Saiba como usar a pesquisa de vetor no Azure DocumentDB com o driver Java MongoDB para armazenar e consultar dados de vetor com eficiência.

Este guia de início rápido fornece um tour guiado das principais técnicas de pesquisa de vetor usando um aplicativo de exemplo java no GitHub.

O aplicativo usa um conjunto de dados de hotel de exemplo em um arquivo JSON com vetores pré-calculados do text-embedding-3-small modelo, embora você também possa gerar os vetores por conta própria. Os dados do hotel incluem nomes de hotéis, locais, descrições e inserções de vetor.

Prerequisites

Uma assinatura de Azure
- Se você não tem uma assinatura do Azure, crie uma conta gratuita

Um cluster existente do Azure DocumentDB
- Se você não tiver um cluster, crie um novo cluster
- RBAC (Controle de Acesso Baseado em Função) habilitado
- Firewall configurado para permitir o acesso ao endereço IP do cliente
Recurso do Azure OpenAI
- Domínio personalizado configurado
- RBAC (Controle de Acesso Baseado em Função) habilitado
- text-embedding-3-small modelo implantado
Código do Visual Studio
- Extensão do DocumentDB

Use o ambiente bash em Azure Cloud Shell. Para obter mais informações, confira Introdução ao Azure Cloud Shell.
Se você preferir executar comandos de referência da CLI localmente, instale a CLI do Azure. Se você estiver executando no Windows ou no macOS, considere executar a CLI do Azure em um contêiner do Docker. Para obter mais informações, confira Como executar a CLI do Azure em um contêiner do Docker.
- Se você estiver usando uma instalação local, entre na CLI do Azure usando o comando az login . Para concluir o processo de autenticação, siga as etapas exibidas em seu terminal. Para obter outras opções de entrada, consulte Autenticar no Azure usando a CLI do Azure.
- Quando solicitado, instale a extensão da CLI do Azure no primeiro uso. Para obter mais informações sobre extensões, confira Usar e gerenciar extensões com a CLI do Azure.
- Execute o comando az version para localizar a versão e as bibliotecas dependentes que estão instaladas. Para atualizar para a versão mais recente, execute az upgrade.

Java 21 ou posterior
Maven 3.6 ou posterior

Criar arquivo de dados com vetores

Crie um novo diretório de dados para o arquivo de dados de hotéis:
```
mkdir data
```
Copie o Hotels_Vector.jsonarquivo de dados brutos com vetores para o diretório data .

Criar um projeto Java

Crie um novo diretório irmão para seu projeto, no mesmo nível do diretório de dados e abra-o no Visual Studio Code:
```
mkdir vector-search-quickstart
mkdir vector-search-quickstart/src
code vector-search-quickstart
```

Crie um pom.xml arquivo na raiz do projeto com o seguinte conteúdo:

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

O aplicativo usa as seguintes dependências do Maven especificadas no pom.xml:

mongodb-driver-sync: driver Java oficial do MongoDB para conectividade e operações de banco de dados
azure-identity: biblioteca de identidade Azure para autenticação sem senha com Microsoft Entra ID
azure-ai-openai: Biblioteca de cliente do Azure OpenAI para se comunicar com modelos de IA e criar incorporações de vetores
jackson-databind: biblioteca de serialização e desserialização JSON
slf4j-nop: associação SLF4J sem operação para suprimir a saída de log do driver MongoDB

Crie um .env arquivo na raiz do projeto para variáveis de ambiente:

# 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

Substitua os valores de espaço reservado no arquivo .env por suas próprias informações.

AZURE_OPENAI_EMBEDDING_ENDPOINT: a URL do endpoint do recurso do Azure OpenAI.
MONGO_CLUSTER_NAME: o nome do recurso do Azure DocumentDB.

Carregue as variáveis de ambiente:
```
set -a && source .env && set +a
```

A estrutura do projeto deve ter esta aparência:

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

Adicionar código para pesquisa de vetor

Crie um DiskAnn.java arquivo no src diretório e cole o seguinte código:

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

Crie um IVF.java arquivo no src diretório e cole o seguinte código:

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 IVF index.
 */
public class IVF {
    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_ivf";
    private static final String VECTOR_INDEX_NAME = "vectorIndex_ivf";

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

    public static void main(String[] args) {
        new IVF().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-ivf")
                        .append("dimensions", dimensions)
                        .append("similarity", "COS")
                        .append("numLists", 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;
    }
}

Crie um HNSW.java arquivo no src diretório e cole o seguinte código:

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 HNSW index.
 */
public class HNSW {
    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_hnsw";
    private static final String VECTOR_INDEX_NAME = "vectorIndex_hnsw";

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

    public static void main(String[] args) {
        new HNSW().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-hnsw")
                        .append("dimensions", dimensions)
                        .append("similarity", "COS")
                        .append("m", 16)
                        .append("efConstruction", 64)
                    )
            ));

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

Esse código executa as seguintes tarefas:

Cria uma conexão sem senha com o Azure DocumentDB usando DefaultAzureCredential e o mecanismo OIDC do MongoDB
Cria um cliente do Azure OpenAI para gerar inserções
Descarta e recria a coleção e depois carrega dados de hotel do arquivo JSON em lotes
Cria índices padrão e um índice de vetor com opções específicas de algoritmo
Gera uma inserção para uma consulta de exemplo e executa um pipeline de pesquisa de agregação
Imprime os cinco melhores hotéis correspondentes com pontuações de similaridade

Autenticar no Azure

Entre no Azure antes de executar o aplicativo para que ele possa acessar os recursos do Azure com segurança.

Note

Verifique se a sua identidade de usuário conectado tem as funções necessárias do plano de dados em ambas as contas do recurso Azure DocumentDB e OpenAI do Azure.

az login

O código usa sua autenticação de desenvolvedor local para acessar o Azure DocumentDB e o Azure OpenAI. Quando você define AZURE_TOKEN_CREDENTIALS=AzureCliCredential, essa configuração informa à função para usar as credenciais da CLI do Azure para autenticação deterministicamente. A autenticação depende de DefaultAzureCredential de azure-identity para encontrar suas credenciais do Azure no ambiente. Saiba mais sobre como autenticar aplicativos Java para os serviços do Azure usando a biblioteca Azure Identity.

Compilar o aplicativo

Compile o aplicativo:

mvn clean compile

Execute a pesquisa DiskANN (vizinho mais próximo aproximado baseado em disco):

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

DiskANN é otimizado para grandes conjuntos de dados que não se encaixam na memória, armazenamento baseado em disco eficiente e um bom equilíbrio de velocidade e precisão.

Exemplo de saída:

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

Executar pesquisa de Arquivo Invertido (IVF).

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

Os grupos de IVF organizam vetores por similaridade e fornecem uma pesquisa rápida através dos centróides dos clusters. Oferece compromissos configuráveis entre precisão e velocidade para grandes conjuntos de dados de vetor.

Exemplo de saída:

Created collection: hotels_ivf
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_ivf

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.4785
3. HotelName: Nordick's Valley Motel, Score: 0.4635
4. HotelName: Economy Universe Motel, Score: 0.4461
5. HotelName: Roach Motel, Score: 0.4388

Execute a pesquisa HNSW (Mundo Pequeno Navegável Hierárquico):

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

O HNSW fornece excelente desempenho de pesquisa com altas taxas de recall usando uma estrutura de grafo hierárquica, tornando-o adequado para aplicativos em tempo real.

Exemplo de saída:

Created collection: hotels_hnsw
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_hnsw

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.4785
3. HotelName: Nordick's Valley Motel, Score: 0.4635
4. HotelName: Economy Universe Motel, Score: 0.4461
5. HotelName: Roach Motel, Score: 0.4388

Exibir e gerenciar dados no Visual Studio Code

Instale a extensão do DocumentDB e o Pacote de Extensões para Java no Visual Studio Code.
Conecte-se à sua conta do Azure DocumentDB usando a extensão do DocumentDB.
Exiba os dados e índices no banco de dados Hotéis.

Limpar os recursos

Exclua o grupo de recursos, o cluster do Azure DocumentDB e o recurso do Azure OpenAI quando você não precisar mais deles para evitar custos desnecessários.

Comentários

Esta página foi útil?

Last updated on 2026-04-27

Início Rápido: Pesquisa de vetor com Java no Azure DocumentDB

Prerequisites

Criar arquivo de dados com vetores

Criar um projeto Java

Adicionar código para pesquisa de vetor

Autenticar no Azure

Compilar o aplicativo

Exibir e gerenciar dados no Visual Studio Code

Limpar os recursos

Conteúdo relacionado

Comentários

Recursos adicionais