次の方法で共有

クイック スタート: Azure DocumentDB での Node.js を使用したベクター検索

Node.js クライアント ライブラリを使用して、Azure DocumentDB でベクター検索を使用します。 ベクター データを効率的に格納およびクエリします。

このクイック スタートでは、json ファイル内のサンプル ホテル データセットと、 text-embedding-3-small モデルのベクターを使用します。 データセットには、ホテル名、場所、説明、ベクター埋め込みなどが含まれます。

GitHub で サンプル コード を見つけます。

[前提条件]

  • Azure サブスクリプション

    • Azure サブスクリプションをお持ちでない場合は、無料アカウントを作成してください

  • Azure Cloud Shell で Bash 環境を使用します。 詳細については、「Azure Cloud Shell の概要」を参照してください。

  • CLI 参照コマンドをローカルで実行する場合は、Azure CLI を インストール します。 Windows または macOS で実行している場合は、Docker コンテナーで Azure CLI を実行することを検討してください。 詳細については、「Docker コンテナーで Azure CLI を実行する方法」を参照してください。

    • ローカル インストールを使用する場合は、az login コマンドを使用して Azure CLI にサインインします。 認証プロセスを完了するには、ターミナルに表示される手順に従います。 その他のサインイン オプションについては、「 Azure CLI を使用した Azure への認証」を参照してください。

    • 初回使用時にインストールを求められたら、Azure CLI 拡張機能をインストールします。 拡張機能の詳細については、「Azure CLI で拡張機能を使用および管理する」を参照してください。

    • az version を実行し、インストールされているバージョンおよび依存ライブラリを検索します。 最新バージョンにアップグレードするには、az upgrade を実行します。

  • Node.js LTS

  • TypeScript: TypeScript をグローバルにインストールします。

    npm install -g typescript

ベクターを使用してデータ ファイルを作成する

  1. hotels データ ファイルの新しいデータ ディレクトリを作成します。

    mkdir data

  2. ベクターを含む Hotels_Vector.jsonraw データ ファイル を data ディレクトリにコピーします。

Node.js プロジェクトを作成する

  1. プロジェクトの新しい兄弟ディレクトリをデータ ディレクトリと同じレベルで作成し、Visual Studio Code で開きます。

    mkdir vector-search-quickstart
code vector-search-quickstart

  2. ターミナルで、Node.js プロジェクトを初期化します。

    npm init -y
npm pkg set type="module"

  3. 必要なパッケージをインストールします。

    npm install mongodb @azure/identity openai @types/node
    • mongodb: MongoDB Node.js ドライバー
    • @azure/identity: パスワードレス認証用の Azure ID ライブラリ
    • openai: ベクトルを作成するための OpenAI クライアント ライブラリ
    • @types/node: Node.js の型定義

  4. 環境変数のプロジェクト ルートに .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

# MongoDB configuration
MONGO_CLUSTER_NAME=

# Data file
DATA_FILE_WITH_VECTORS=../data/Hotels_Vector.json
FIELD_TO_EMBED=Description
EMBEDDED_FIELD=DescriptionVector
EMBEDDING_DIMENSIONS=1536
LOAD_SIZE_BATCH=100

    .env ファイル内のプレースホルダーの値を独自の情報に置き換えます。

    • AZURE_OPENAI_EMBEDDING_ENDPOINT: Azure OpenAI リソース エンドポイントの URL
    • MONGO_CLUSTER_NAME: 対象のリソース名

  5. TypeScript を構成する tsconfig.json ファイルを追加します。

    {
    "compilerOptions": {
        "target": "ES2020",
        "module": "NodeNext",
        "moduleResolution": "nodenext",
        "declaration": true,
        "outDir": "./dist",
        "strict": true,
        "esModuleInterop": true,
        "skipLibCheck": true,
        "noImplicitAny": false,
        "forceConsistentCasingInFileNames": true,
        "sourceMap": true,
        "resolveJsonModule": true,
    },
    "include": [
        "src/**/*"
    ],
    "exclude": [
        "node_modules",
        "dist"
    ]
}

npm スクリプトを作成する

package.json ファイルを編集し、次のスクリプトを追加します。

これらのスクリプトを使用して TypeScript ファイルをコンパイルし、DiskANN インデックスの実装を実行します。

"scripts": { 
    "build": "tsc",
    "start:diskann": "node --env-file .env dist/diskann.js"
}

TypeScript ファイルの src ディレクトリを作成します。 DiskANN インデックス実装の diskann.ts と utils.ts の 2 つのファイルを追加します。

mkdir src    
touch src/diskann.ts
touch src/utils.ts

次のコードを diskann.ts ファイルに貼り付けます。

import path from 'path';
import { readFileReturnJson, getClientsPasswordless, insertData, printSearchResults } from './utils.js';

// ESM specific features - create __dirname equivalent
import { fileURLToPath } from "node:url";
import { dirname } from "node:path";
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

const config = {
    query: "quintessential lodging near running trails, eateries, retail",
    dbName: "Hotels",
    collectionName: "hotels_diskann",
    indexName: "vectorIndex_diskann",
    dataFile: process.env.DATA_FILE_WITH_VECTORS!,
    batchSize: parseInt(process.env.LOAD_SIZE_BATCH! || '100', 10),
    embeddedField: process.env.EMBEDDED_FIELD!,
    embeddingDimensions: parseInt(process.env.EMBEDDING_DIMENSIONS!, 10),
    deployment: process.env.AZURE_OPENAI_EMBEDDING_MODEL!,
};

async function main() {

    const { aiClient, dbClient } = getClientsPasswordless();

    try {

        if (!aiClient) {
            throw new Error('AI client is not configured. Please check your environment variables.');
        }
        if (!dbClient) {
            throw new Error('Database client is not configured. Please check your environment variables.');
        }

        await dbClient.connect();
        const db = dbClient.db(config.dbName);
        const collection = await db.createCollection(config.collectionName);
        console.log('Created collection:', config.collectionName);
        const data = await readFileReturnJson(path.join(__dirname, "..", config.dataFile));
        const insertSummary = await insertData(config, collection, data);
        console.log('Created vector index:', config.indexName);
        
        // Create the vector index
        const indexOptions = {
            createIndexes: config.collectionName,
            indexes: [
                {
                    name: config.indexName,
                    key: {
                        [config.embeddedField]: 'cosmosSearch'
                    },
                    cosmosSearchOptions: {
                        kind: 'vector-diskann',
                        dimensions: config.embeddingDimensions,
                        similarity: 'COS', // 'COS', 'L2', 'IP'
                        maxDegree: 20, // 20 - 2048,  edges per node
                        lBuild: 10 // 10 - 500, candidate neighbors evaluated
                    }
                }
            ]
        };
        const vectorIndexSummary = await db.command(indexOptions);

        // Create embedding for the query
        const createEmbeddedForQueryResponse = await aiClient.embeddings.create({
            model: config.deployment,
            input: [config.query]
        });

        // Perform the vector similarity search
        const searchResults = await collection.aggregate([
            {
                $search: {
                    cosmosSearch: {
                        vector: createEmbeddedForQueryResponse.data[0].embedding,
                        path: config.embeddedField,
                        k: 5
                    }
                }
            },
            {
                $project: {
                    score: {
                        $meta: "searchScore"
                    },
                    document: "$$ROOT"
                }
            }
        ]).toArray();

        // Print the results
        printSearchResults(insertSummary, vectorIndexSummary, searchResults);

    } catch (error) {
        console.error('App failed:', error);
        process.exitCode = 1;
    } finally {
        console.log('Closing database connection...');
        if (dbClient) await dbClient.close();
        console.log('Database connection closed');
    }
}

// Execute the main function
main().catch(error => {
    console.error('Unhandled error:', error);
    process.exitCode = 1;
});

このメイン モジュールには、次の機能があります。

  • ユーティリティ関数を含む
  • 環境変数の構成オブジェクトを作成します。
  • Azure OpenAI および DocumentDB 用のクライアントを作成します
  • MongoDB に接続し、データベースとコレクションを作成し、データを挿入して、標準インデックスを作成します
  • IVF、HNSW、または DiskANN を使用してベクター インデックスを作成します。
  • OpenAI クライアントを使用して、サンプル クエリ テキストの埋め込みを作成します。 ファイルの先頭にあるクエリを変更できます
  • 埋め込みを使用してベクター検索を実行し、結果を出力します

ユーティリティ関数を作成する

次のコードを utils.tsに貼り付けます。

import { MongoClient, OIDCResponse, OIDCCallbackParams } from 'mongodb';
import { AzureOpenAI } from 'openai/index.js';
import { promises as fs } from "fs";
import { AccessToken, DefaultAzureCredential, TokenCredential, getBearerTokenProvider } from '@azure/identity';

// Define a type for JSON data
export type JsonData = Record<string, any>;

export const AzureIdentityTokenCallback = async (params: OIDCCallbackParams, credential: TokenCredential): Promise<OIDCResponse> => {
    const tokenResponse: AccessToken | null = await credential.getToken(['https://ossrdbms-aad.database.windows.net/.default']);
    return {
        accessToken: tokenResponse?.token || '',
        expiresInSeconds: (tokenResponse?.expiresOnTimestamp || 0) - Math.floor(Date.now() / 1000)
    };
};
export function getClients(): { aiClient: AzureOpenAI; dbClient: MongoClient } {
    const apiKey = process.env.AZURE_OPENAI_EMBEDDING_KEY!;
    const apiVersion = process.env.AZURE_OPENAI_EMBEDDING_API_VERSION!;
    const endpoint = process.env.AZURE_OPENAI_EMBEDDING_ENDPOINT!;
    const deployment = process.env.AZURE_OPENAI_EMBEDDING_MODEL!;
    const mongoConnectionString = process.env.MONGO_CONNECTION_STRING!;

    if (!apiKey || !apiVersion || !endpoint || !deployment || !mongoConnectionString) {
        throw new Error('Missing required environment variables: AZURE_OPENAI_EMBEDDING_KEY, AZURE_OPENAI_EMBEDDING_API_VERSION, AZURE_OPENAI_EMBEDDING_ENDPOINT, AZURE_OPENAI_EMBEDDING_MODEL, MONGO_CONNECTION_STRING');
    }

    const aiClient = new AzureOpenAI({
        apiKey,
        apiVersion,
        endpoint,
        deployment
    });
    const dbClient = new MongoClient(mongoConnectionString, {
        // Performance optimizations
        maxPoolSize: 10,         // Limit concurrent connections
        minPoolSize: 1,          // Maintain at least one connection
        maxIdleTimeMS: 30000,    // Close idle connections after 30 seconds
        connectTimeoutMS: 30000, // Connection timeout
        socketTimeoutMS: 360000, // Socket timeout (for long-running operations)
        writeConcern: {          // Optimize write concern for bulk operations
            w: 1,                // Acknowledge writes after primary has written
            j: false             // Don't wait for journal commit
        }
    });

    return { aiClient, dbClient };
}

export function getClientsPasswordless(): { aiClient: AzureOpenAI | null; dbClient: MongoClient | null } {
    let aiClient: AzureOpenAI | null = null;
    let dbClient: MongoClient | null = null;

    // Validate all required environment variables upfront
    const apiVersion = process.env.AZURE_OPENAI_EMBEDDING_API_VERSION!;
    const endpoint = process.env.AZURE_OPENAI_EMBEDDING_ENDPOINT!;
    const deployment = process.env.AZURE_OPENAI_EMBEDDING_MODEL!;
    const clusterName = process.env.MONGO_CLUSTER_NAME!;

    if (!apiVersion || !endpoint || !deployment || !clusterName) {
        throw new Error('Missing required environment variables: AZURE_OPENAI_EMBEDDING_API_VERSION, AZURE_OPENAI_EMBEDDING_ENDPOINT, AZURE_OPENAI_EMBEDDING_MODEL, MONGO_CLUSTER_NAME');
    }

    console.log(`Using Azure OpenAI Embedding API Version: ${apiVersion}`);
    console.log(`Using Azure OpenAI Embedding Deployment/Model: ${deployment}`);

    const credential = new DefaultAzureCredential();

    // For Azure OpenAI with DefaultAzureCredential
    {
        const scope = "https://cognitiveservices.azure.com/.default";
        const azureADTokenProvider = getBearerTokenProvider(credential, scope);
        aiClient = new AzureOpenAI({
            apiVersion,
            endpoint,
            deployment,
            azureADTokenProvider
        });
    }

    // For DocumentDB with DefaultAzureCredential (uses signed-in user)
    {
        dbClient = new MongoClient(
            `mongodb+srv://${clusterName}.mongocluster.cosmos.azure.com/`, {
            connectTimeoutMS: 120000,
            tls: true,
            retryWrites: false,
            maxIdleTimeMS: 120000,
            authMechanism: 'MONGODB-OIDC',
            authMechanismProperties: {
                OIDC_CALLBACK: (params: OIDCCallbackParams) => AzureIdentityTokenCallback(params, credential),
                ALLOWED_HOSTS: ['*.azure.com']
            }
        }
        );
    }

    return { aiClient, dbClient };
}

export async function readFileReturnJson(filePath: string): Promise<JsonData[]> {

    console.log(`Reading JSON file from ${filePath}`);

    const fileAsString = await fs.readFile(filePath, "utf-8");
    return JSON.parse(fileAsString);
}
export async function writeFileJson(filePath: string, jsonData: JsonData): Promise<void> {
    const jsonString = JSON.stringify(jsonData, null, 2);
    await fs.writeFile(filePath, jsonString, "utf-8");

    console.log(`Wrote JSON file to ${filePath}`);
}
export async function insertData(config, collection, data) {
    console.log(`Processing in batches of ${config.batchSize}...`);
    const totalBatches = Math.ceil(data.length / config.batchSize);

    let inserted = 0;
    let updated = 0;
    let skipped = 0;
    let failed = 0;

    for (let i = 0; i < totalBatches; i++) {
        const start = i * config.batchSize;
        const end = Math.min(start + config.batchSize, data.length);
        const batch = data.slice(start, end);

        try {
            const result = await collection.insertMany(batch, { ordered: false });
            inserted += result.insertedCount || 0;
            console.log(`Batch ${i + 1} complete: ${result.insertedCount} inserted`);
        } catch (error: any) {
            if (error?.writeErrors) {
                // Some documents may have been inserted despite errors
                console.error(`Error in batch ${i + 1}: ${error?.writeErrors.length} failures`);
                failed += error?.writeErrors.length;
                inserted += batch.length - error?.writeErrors.length;
            } else {
                console.error(`Error in batch ${i + 1}:`, error);
                failed += batch.length;
            }
        }

        // Small pause between batches to reduce resource contention
        if (i < totalBatches - 1) {
            await new Promise(resolve => setTimeout(resolve, 100));
        }
    }
    const indexColumns = [
        "HotelId",
        "Category",
        "Description",
        "Description_fr"
    ];
    for (const col of indexColumns) {
        const indexSpec = {};
        indexSpec[col] = 1; // Ascending index
        await collection.createIndex(indexSpec);
    }

    return { total: data.length, inserted, updated, skipped, failed };
}

export function printSearchResults(insertSummary, indexSummary, searchResults) {


    if (!searchResults || searchResults.length === 0) {
        console.log('No search results found.');
        return;
    }

    searchResults.map((result, index) => {

        const { document, score } = result as any;

        console.log(`${index + 1}. HotelName: ${document.HotelName}, Score: ${score.toFixed(4)}`);
        //console.log(`   Description: ${document.Description}`);
    });

}

このユーティリティ モジュールには、次の機能があります。

  • JsonData: データ構造のインターフェイス
  • scoreProperty: ベクトル検索方法に基づくクエリ結果内のスコアの位置
  • getClients: Azure OpenAI と Azure DocumentDB のクライアントを作成して返します。
  • getClientsPasswordless: パスワードレス認証を使用して、Azure OpenAI および Azure DocumentDB のクライアントを作成して返します。 両方のリソースで RBAC を有効にし、Azure CLI にサインインする
  • readFileReturnJson: JSON ファイルを読み取り、その内容を JsonData オブジェクトの配列として返します。
  • writeFileJson: JsonData オブジェクトの配列を JSON ファイルに書き込みます。
  • insertData: MongoDB コレクションにデータをバッチで挿入し、指定されたフィールドに標準インデックスを作成します。
  • printSearchResults: スコアとホテル名を含むベクター検索の結果を出力します。

Azure CLI で認証する

アプリケーションを実行する前に Azure CLI にサインインして、アプリが Azure リソースに安全にアクセスできるようにします。

az login

このコードでは、ローカル開発者認証を使用して、getClientsPasswordlessから utils.ts 関数を使用して Azure DocumentDB と Azure OpenAI にアクセスします。 AZURE_TOKEN_CREDENTIALS=AzureCliCredential設定すると、この設定により、認証に Azure CLI 資格情報を決定的に使用するように関数に指示されます。 この関数は、@azure/ID からの DefaultAzureCredential に依存して、環境内の Azure 資格情報を検索します。 Azure Id ライブラリを使用して Azure サービスに対して JavaScript アプリを認証する方法について説明します。

アプリケーションの構築と実行

TypeScript ファイルをビルドし、アプリケーションを実行します。

npm run build
npm run start:diskann

アプリのログ記録と出力には、次の情報が表示されます。

  • コレクションの作成とデータの挿入の状態
  • ベクター インデックスの作成
  • ホテル名と類似性スコアを含む検索結果
Using Azure OpenAI Embedding API Version: 2023-05-15
Using Azure OpenAI Embedding Deployment/Model: text-embedding-3-small-2
Created collection: hotels_diskann
Reading JSON file from \documentdb-samples\ai\data\Hotels_Vector.json
Processing in batches of 50...
Batch 1 complete: 50 inserted
Created vector index: vectorIndex_diskann
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
Closing database connection...
Database connection closed

Visual Studio Code でデータを表示および管理する

  1. Visual Studio Code で DocumentDB 拡張機能 を選択して、Azure DocumentDB アカウントに接続します。

  2. Hotels データベースのデータとインデックスを表示します。

    DocumentDB コレクションを示す DocumentDB 拡張機能のスクリーンショット。

リソースをクリーンアップする

追加コストを回避するためにリソース グループ、DocumentDB アカウント、Azure OpenAI リソースが不要な場合は、リソース グループ、DocumentDB アカウント、Azure OpenAI リソースを削除します。

関連コンテンツ

  • Last updated on