빠른 시작: Azure DocumentDB에서 Java를 사용하여 벡터 검색

Java MongoDB 드라이버와 함께 Azure DocumentDB에서 벡터 검색을 사용하여 벡터 데이터를 효율적으로 저장하고 쿼리하는 방법을 알아봅니다.

이 빠른 시작에서는 GitHub에서 Java 샘플 앱을 사용하여 주요 벡터 검색 기술을 안내합니다.

앱은 모델에서 미리 계산된 벡터가 있는 JSON 파일의 text-embedding-3-small 샘플 호텔 데이터 세트를 사용하지만 직접 벡터를 생성할 수도 있습니다. 호텔 데이터에는 호텔 이름, 위치, 설명 및 벡터 포함이 포함됩니다.

필수 조건

  • Azure 구독

벡터를 사용하여 데이터 파일 만들기

  1. 호텔 데이터 파일에 대한 새 데이터 디렉터리를 만듭니다.

    mkdir data

  2. 벡터를 사용하여 Hotels_Vector.json원시 데이터 파일을 디렉터리에 복사합니다 data .

Java 프로젝트 만들기

  1. 데이터 디렉터리와 동일한 수준에서 프로젝트에 대한 새 형제 디렉터리를 만들고 Visual Studio Code에서 엽니다.

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

  2. 다음 콘텐츠를 사용하여 pom.xml 프로젝트 루트에 파일을 만듭니다.

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

    앱은 다음에 지정된 다음 Maven 종속성을 사용합니다.pom.xml

    • mongodb-driver-sync: 데이터베이스 연결 및 작업에 대한 공식 MongoDB Java 드라이버
    • azure-identity: Microsoft Entra ID를 사용한 암호 없는 인증을 위한 Azure ID 라이브러리
    • azure-ai-openai: Azure OpenAI 클라이언트 라이브러리를 사용하여 AI 모델과 통신하고 벡터 임베딩을 생성합니다.
    • jackson-databind: JSON 직렬화 및 역직렬화 라이브러리
    • slf4j-nop: MongoDB 드라이버의 로깅 출력을 억제하는 연산이 없는 SLF4J 바인딩

  3. 환경 변수에 .env 대한 파일을 프로젝트 루트에 만듭니다.

    # 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

    .env 파일의 자리 표시자 값을 사용자 정보로 바꿉니다.

    • AZURE_OPENAI_EMBEDDING_ENDPOINT: Azure OpenAI 리소스 엔드포인트 URL입니다.
    • MONGO_CLUSTER_NAME: Azure DocumentDB 리소스 이름입니다.

  4. 환경 변수를 로드합니다.

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

  5. 프로젝트 구조는 다음과 같습니다.

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

DiskAnn.java 디렉터리에 src 파일을 만들고 다음 코드를 붙여넣습니다.

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

이 코드는 다음 작업을 수행합니다.

  • MongoDB OIDC 메커니즘을 사용하여 DefaultAzureCredential Azure DocumentDB에 암호 없는 연결을 만듭니다.
  • 임베딩을 생성하기 위한 Azure OpenAI 클라이언트를 만듭니다.
  • 컬렉션을 삭제하고 다시 생성한 다음 JSON 파일의 호텔 데이터를 일괄 처리로 로드합니다.
  • 알고리즘별 옵션을 사용하여 표준 인덱스 및 벡터 인덱스를 만듭니다.
  • 샘플 쿼리에 대한 포함을 생성하고 집계 검색 파이프라인을 실행합니다.
  • 유사성 점수가 있는 상위 5개 호텔 인쇄

Azure에 인증하기

Azure 리소스에 안전하게 액세스할 수 있도록 애플리케이션을 실행하기 전에 Azure에 로그인합니다.

비고

로그인한 ID에 Azure DocumentDB 계정과 Azure OpenAI 리소스 모두에서 필요한 데이터 평면 역할이 있는지 확인합니다.

az login

이 코드는 로컬 개발자 인증을 사용하여 Azure DocumentDB 및 Azure OpenAI에 액세스합니다. 설정할 AZURE_TOKEN_CREDENTIALS=AzureCliCredential때 이 설정은 함수에 인증을 위해 Azure CLI 자격 증명을 결정적으로 사용하도록 지시합니다. 인증은 Azure ID의 DefaultAzureCredential을 사용하여 환경에서 Azure 자격 증명을 찾습니다. Azure ID 라이브러리를 사용하여 Azure 서비스에 Java 앱을 인증하는 방법에 대해 자세히 알아봅니다.

애플리케이션 빌드

애플리케이션을 컴파일합니다.

mvn clean compile

DiskANN(디스크 기반 근사 최근접 이웃) 검색을 수행합니다.

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

DiskANN은 메모리에 맞지 않는 큰 데이터 세트, 효율적인 디스크 기반 스토리지 및 속도와 정확도의 균형에 최적화되어 있습니다.

예시 출력:

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

Visual Studio Code에서 데이터 보기 및 관리

  1. Visual Studio Code에서 Java용 DocumentDB 확장 및 확장 팩 을 설치합니다.

  2. DocumentDB 확장을 사용하여 Azure DocumentDB 계정에 연결합니다.

  3. Hotels 데이터베이스에서 데이터 및 인덱스를 봅니다.

    DocumentDB 컬렉션을 보여 주는 DocumentDB 확장의 스크린샷

자원을 정리하세요

불필요한 비용을 방지하기 위해 더 이상 필요하지 않은 경우 리소스 그룹, Azure DocumentDB 클러스터 및 Azure OpenAI 리소스를 삭제합니다.

관련 콘텐츠

  • Last updated on