Partager via

Démarrage rapide : Recherche vectorielle avec Node.js dans Azure DocumentDB

Utilisez la recherche vectorielle dans Azure DocumentDB avec la bibliothèque cliente Node.js. Stockez et interrogez efficacement les données vectorielles.

Ce guide de démarrage rapide utilise un exemple de jeu de données d’hôtel dans un fichier JSON avec des vecteurs du text-embedding-ada-002 modèle. Le jeu de données comprend des noms d’hôtel, des emplacements, des descriptions et des incorporations vectorielles.

Recherchez l’exemple de code sur GitHub.

Prerequisites

  • Un abonnement Azure

  • Un cluster Azure DocumentDB existant

    • Si vous n’avez pas de cluster, créez un cluster

Créer un projet Node.js

  1. Créez un répertoire pour votre projet et ouvrez-le dans Visual Studio Code :

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

  2. Dans le terminal, initialisez un projet Node.js :

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

  3. Installez les packages nécessaires :

    npm install mongodb @azure/identity openai @types/node
    • mongodb: pilote Node.js MongoDB
    • @azure/identity: Bibliothèque d’identités Azure pour l’authentification sans mot de passe
    • openai: Bibliothèque de client OpenAI pour créer des vecteurs
    • @types/node: définitions de types pour Node.js

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

    # Azure OpenAI Embedding Settings
AZURE_OPENAI_EMBEDDING_MODEL=text-embedding-ada-002
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=HotelsData_toCosmosDB_Vector.json
EMBEDDED_FIELD=text_embedding_ada_002
EMBEDDING_DIMENSIONS=1536
LOAD_SIZE_BATCH=100

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

    • AZURE_OPENAI_EMBEDDING_ENDPOINT: URL de votre point de terminaison de ressource Azure OpenAI
    • MONGO_CLUSTER_NAME: Nom de votre ressource

  5. Ajoutez un tsconfig.json fichier pour configurer 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"
    ]
}

  6. Copiez le HotelsData_toCosmosDB_Vector.jsonfichier de données brutes avec des vecteurs à la racine de votre projet.

Créer des scripts npm

Modifiez le package.json fichier et ajoutez ces scripts :

Utilisez ces scripts pour compiler des fichiers TypeScript et exécuter l’implémentation d’index DiskANN.

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

Créez un src répertoire pour vos fichiers TypeScript. Ajoutez deux fichiers : diskann.ts et utils.ts pour l’implémentation d’index DiskANN :

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

Collez le code suivant dans le diskann.ts fichier.

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

Ce module principal fournit les fonctionnalités suivantes :

  • Inclut les fonctions utilitaires
  • Crée un objet de configuration pour les variables d’environnement
  • Crée des clients pour Azure OpenAI et DocumentDB
  • Se connecte à MongoDB, crée une base de données et une collection, insère des données et crée des index standard
  • Crée un index vectoriel à l’aide d’IVF, HNSW ou DiskANN
  • Crée un encodage pour un texte de requête d'exemple avec le client OpenAI. Vous pouvez modifier la requête en haut du fichier
  • Exécute une recherche vectorielle à l’aide de l’incorporation et imprime les résultats

Créer des fonctions utilitaires

Collez le code suivant dans 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 aiClient = new AzureOpenAI({
        apiKey,
        apiVersion,
        endpoint,
        deployment
    });
    const dbClient = new MongoClient(process.env.MONGO_CONNECTION_STRING!, {
        // 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;

    // For Azure OpenAI with DefaultAzureCredential
    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!;

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

    // For Cosmos DB with DefaultAzureCredential
    const clusterName = process.env.MONGO_CLUSTER_NAME!;

    if (clusterName) {
        const credential = new DefaultAzureCredential();

        dbClient = new MongoClient(
            `mongodb+srv://${clusterName}.global.mongocluster.cosmos.azure.com/`, {
            connectTimeoutMS: 30000,
            tls: true,
            retryWrites: true,
            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}`);
    });

}

Ce module utilitaire fournit les fonctionnalités suivantes :

  • JsonData: interface pour la structure de données
  • scoreProperty: Emplacement du score dans les résultats de requête en fonction de la méthode de recherche vectorielle
  • getClients: crée et retourne des clients pour Azure OpenAI et Azure DocumentDB
  • getClientsPasswordless: crée et retourne des clients pour Azure OpenAI et Azure DocumentDB à l’aide de l’authentification sans mot de passe. Activer RBAC sur les deux ressources et se connecter à Azure CLI
  • readFileReturnJson: lit un fichier JSON et retourne son contenu sous la forme d’un tableau d’objets JsonData
  • writeFileJson: écrit un tableau d’objets JsonData dans un fichier JSON
  • insertData: insère des données dans des lots dans une collection MongoDB et crée des index standard sur des champs spécifiés
  • printSearchResults: imprime les résultats d’une recherche vectorielle, y compris le score et le nom de l’hôtel

S’authentifier auprès d’Azure CLI

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

az login

Générer et exécuter l’application

Générez les fichiers TypeScript, puis exécutez l’application :

npm run build
npm run start:diskann

La journalisation et la sortie de l’application indiquent :

  • Création de regroupements et état d’insertion des données
  • Création d'un index vectoriel
  • Résultats de recherche avec des noms d’hôtel et des scores de similarité
Created collection: hotels_diskann
Reading JSON file from C:\Users\<username>\repos\samples\cosmos-db-vector-samples\data\HotelsData_toCosmosDB_Vector.json
Processing in batches of 100...
Batch 1 complete: 50 inserted
Created vector index: vectorIndex_diskann
1. HotelName: Roach Motel, Score: 0.8399
2. HotelName: Royal Cottage Resort, Score: 0.8385
3. HotelName: Economy Universe Motel, Score: 0.8360
4. HotelName: Foot Happy Suites, Score: 0.8354
5. HotelName: Country Comfort Inn, Score: 0.8346
Closing database connection...
Database connection closed

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

  1. Sélectionnez l’extension DocumentDB dans Visual Studio Code pour vous connecter à votre compte Azure DocumentDB.

  2. 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 compte DocumentDB et la ressource Azure OpenAI lorsque vous n’en avez pas besoin pour éviter des coûts supplémentaires.

Contenu connexe

  • Last updated on