Bagikan melalui

Mulai cepat: Pencarian vektor dengan Node.js di Azure DocumentDB

Gunakan pencarian vektor di Azure DocumentDB dengan pustaka klien Node.js. Menyimpan dan mengkueri data vektor secara efisien.

Panduan cepat ini menggunakan dataset hotel contoh dalam file JSON dengan vektor dari model text-embedding-3-small. Himpunan data mencakup nama hotel, lokasi, deskripsi, dan penyematan vektor.

Temukan kode sampel di GitHub.

Prasyarat

  • Langganan Azure

    • Jika Anda tidak memiliki langganan Azure, buat akun gratis

Membuat file data dengan vektor

  1. Buat direktori data baru untuk file data hotel:

    mkdir data

  2. Hotels_Vector.json Salin file data mentah dengan vektor ke direktori Andadata.

Membuat proyek Node.js

  1. Buat direktori saudara baru untuk proyek Anda, pada tingkat yang sama dengan direktori data, dan buka di Visual Studio Code:

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

  2. Di terminal, inisialisasi proyek Node.js:

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

  3. Instal paket yang diperlukan:

    npm install mongodb @azure/identity openai @types/node
    • mongodb: MongoDB Driver Node.js
    • @azure/identity: Pustaka Identitas Azure untuk autentikasi tanpa kata sandi
    • openai: Pustaka klien OpenAI untuk membuat vektor
    • @types/node: Definisi tipe untuk Node.js

  4. Buat .env file di akar proyek Anda untuk variabel lingkungan:

    # 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

    Ganti nilai placeholder dalam berkas .env dengan informasi milik Anda.

    • AZURE_OPENAI_EMBEDDING_ENDPOINT: URL titik akhir sumber daya Azure OpenAI Anda
    • MONGO_CLUSTER_NAME: Nama sumber daya Anda

  5. tsconfig.json Tambahkan file untuk mengonfigurasi TypeScript:

    {
    "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"
    ]
}

Membuat skrip npm

package.json Edit file dan tambahkan skrip ini:

Gunakan skrip ini untuk mengkompilasi file TypeScript dan menjalankan implementasi indeks DiskANN.

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

Buat src direktori untuk file TypeScript Anda. Tambahkan dua file: diskann.ts dan utils.ts untuk implementasi indeks DiskANN:

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

Tempelkan kode berikut ke diskann.ts dalam file.

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

Modul utama ini menyediakan fitur-fitur berikut:

  • Menyertakan fungsi utilitas
  • Membuat objek konfigurasi untuk variabel lingkungan
  • Membuat klien untuk Azure OpenAI dan DocumentDB
  • Menyambungkan ke MongoDB, membuat database dan koleksi, menyisipkan data, dan membuat indeks standar
  • Membuat indeks vektor menggunakan IVF, HNSW, atau DiskANN
  • Membuat penyematan untuk contoh teks kueri menggunakan klien OpenAI. Anda bisa mengubah kueri di bagian atas file
  • Menjalankan pencarian vektor menggunakan penyematan dan mencetak hasilnya

Membuat fungsi utilitas

Tempelkan kode berikut ke dalam 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}`);
    });

}

Modul utilitas ini menyediakan fitur-fitur ini:

  • JsonData: Antarmuka untuk struktur data
  • scoreProperty: Lokasi skor dalam hasil kueri berdasarkan metode pencarian vektor
  • getClients: Membuat dan mengembalikan klien untuk Azure OpenAI dan Azure DocumentDB
  • getClientsPasswordless: Membuat dan mengembalikan klien untuk Azure OpenAI dan Azure DocumentDB menggunakan autentikasi tanpa kata sandi. Mengaktifkan RBAC pada kedua sumber daya dan masuk ke Azure CLI
  • readFileReturnJson: Membaca file JSON dan mengembalikan kontennya sebagai array JsonData objek
  • writeFileJson: Menulis array dari objek JsonData ke dalam file JSON
  • insertData: Menyisipkan data dalam batch ke dalam koleksi MongoDB dan membuat indeks standar pada bidang yang ditentukan
  • printSearchResults: Mencetak hasil pencarian vektor, termasuk skor dan nama hotel

Mengautentikasi dengan Azure CLI

Masuk ke Azure CLI sebelum Anda menjalankan aplikasi sehingga aplikasi dapat mengakses sumber daya Azure dengan aman.

az login

Kode ini menggunakan autentikasi pengembang lokal Anda untuk mengakses Azure DocumentDB dan Azure OpenAI menggunakan fungsi getClientsPasswordless dari utils.ts. Saat Anda mengatur AZURE_TOKEN_CREDENTIALS=AzureCliCredential, pengaturan ini memberi tahu fungsi untuk menggunakan kredensial Azure CLI untuk autentikasi secara deterministik. Fungsi ini bergantung pada DefaultAzureCredential dari @azure/identitas untuk menemukan kredensial Azure Anda di lingkungan. Pelajari selengkapnya tentang cara Mengautentikasi aplikasi JavaScript ke layanan Azure menggunakan pustaka Azure Identity.

Membangun dan menjalankan aplikasi

Buat file TypeScript, lalu jalankan aplikasi:

npm run build
npm run start:diskann

Pengelogan dan output aplikasi menunjukkan:

  • Pembuatan koleksi dan status penyisipan data
  • Pembuatan indeks vektor
  • Hasil pencarian dengan nama hotel dan skor kesamaan
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

Menampilkan dan mengelola data di Visual Studio Code

  1. Pilih ekstensi DocumentDB di Visual Studio Code untuk menyambungkan ke akun Azure DocumentDB Anda.

  2. Lihat data dan indeks di database Hotel.

    Cuplikan layar ekstensi DocumentDB memperlihatkan koleksi DocumentDB.

Membersihkan sumber daya

Hapus grup sumber daya, akun DocumentDB, dan sumber daya Azure OpenAI saat Anda tidak memerlukannya untuk menghindari biaya tambahan.

Konten terkait

  • Last updated on