Freigeben über

Schnellstart: Vektorsuche mit Node.js in Azure DocumentDB

Verwenden Sie die Vektorsuche in Azure DocumentDB mit der Node.js-Clientbibliothek. Effizientes Speichern und Abfragen von Vektordaten.

Diese Schnellstartanleitung verwendet ein Beispiel-Hotel-Dataset in einer JSON-Datei mit Vektoren aus dem text-embedding-3-small Modell. Das Dataset umfasst Hotelnamen, Standorte, Beschreibungen und Vektoreinbettungen.

Suchen Sie den Beispielcode auf GitHub.

Voraussetzungen

  • Ein Azure-Abonnement

    • Wenn Sie nicht über ein Azure-Abonnement verfügen, erstellen Sie ein kostenloses Konto

Erstellen einer Datendatei mit Vektoren

  1. Erstellen Sie ein neues Datenverzeichnis für die Hotelsdatendatei:

    mkdir data

  2. Kopieren Sie die Hotels_Vector.jsonRohdatendatei mit Vektoren in Ihr data Verzeichnis.

Erstellen eines Node.js Projekts

  1. Erstellen Sie ein neues gleichgeordnetes Verzeichnis für Ihr Projekt auf derselben Ebene wie das Datenverzeichnis, und öffnen Sie es in Visual Studio Code:

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

  2. Initialisieren Sie im Terminal ein Node.js Projekt:

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

  3. Installieren Sie die erforderlichen Pakete:

    npm install mongodb @azure/identity openai @types/node
    • mongodb: MongoDB Node.js Treiber
    • @azure/identity: Azure Identity-Bibliothek für kennwortlose Authentifizierung
    • openai: OpenAI-Clientbibliothek zum Erstellen von Vektoren
    • @types/node: Typdefinitionen für Node.js

  4. Erstellen Sie eine .env Datei im Projektstamm für Umgebungsvariablen:

    # 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

    Ersetzen Sie die Platzhalterwerte in der .env Datei durch Ihre eigenen Informationen:

    • AZURE_OPENAI_EMBEDDING_ENDPOINT: Ihre Azure OpenAI-Ressourcenendpunkt-URL
    • MONGO_CLUSTER_NAME: Ihr Ressourcenname

  5. Fügen Sie eine tsconfig.json Datei zum Konfigurieren von TypeScript hinzu:

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

Erstellen von npm-Skripts

Bearbeiten Sie die package.json Datei, und fügen Sie diese Skripts hinzu:

Verwenden Sie diese Skripts, um TypeScript-Dateien zu kompilieren und die DiskANN-Indeximplementierung auszuführen.

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

Erstellen Sie ein src Verzeichnis für Ihre TypeScript-Dateien. Fügen Sie zwei Dateien hinzu: diskann.ts und utils.ts für die DiskANN-Indeximplementierung:

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

Fügen Sie den folgenden Code in die diskann.ts Datei ein.

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

Dieses Hauptmodul bietet die folgenden Features:

  • Enthält Hilfsfunktionen
  • Erstellt ein Konfigurationsobjekt für Umgebungsvariablen.
  • Erstellt Clients für Azure OpenAI und DocumentDB
  • Stellt eine Verbindung mit MongoDB bereit, erstellt eine Datenbank und Sammlung, fügt Daten ein und erstellt Standardindizes.
  • Erstellt einen Vektorindex mithilfe von IVF, HNSW oder DiskANN
  • Erstellt eine Einbettung für einen Beispielabfragetext mithilfe des OpenAI-Clients. Sie können die Abfrage am Anfang der Datei ändern.
  • Führt eine Vektorsuche mithilfe der Einbettung aus und druckt die Ergebnisse.

Erstellen von Hilfsfunktionen

Fügen Sie den folgenden Code in 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}`);
    });

}

Dieses Hilfsmodul bietet die folgenden Features:

  • JsonData: Schnittstelle für die Datenstruktur
  • scoreProperty: Position der Ergebnisbewertung in Abfrageergebnissen basierend auf der Vektorsuchmethode
  • getClients: Erstellt Clients für Azure OpenAI und Azure DocumentDB und gibt sie zurück
  • getClientsPasswordless: Erstellt und gibt Clients für Azure OpenAI und Azure DocumentDB mit kennwortloser Authentifizierung zurück. RBAC auf beiden Ressourcen aktivieren und bei Azure CLI anmelden
  • readFileReturnJson: Liest eine JSON-Datei und gibt den Inhalt als Array von JsonData Objekten zurück.
  • writeFileJson: Schreibt ein Array von JsonData Objekten in eine JSON-Datei.
  • insertData: Fügt Daten in Batches in eine MongoDB-Auflistung ein und erstellt Standardindizes für angegebene Felder.
  • printSearchResults: Druckt die Ergebnisse einer Vektorsuche, einschließlich der Bewertung und des Hotelnamens.

Authentifizieren mit Azure CLI

Melden Sie sich bei Azure CLI an, bevor Sie die Anwendung ausführen, damit die App sicher auf Azure-Ressourcen zugreifen kann.

az login

Der Code verwendet Ihre lokale Entwicklerauthentifizierung für den Zugriff auf Azure DocumentDB und Azure OpenAI mit der getClientsPasswordless Funktion von utils.ts. Wenn Sie diese Einstellung festlegen AZURE_TOKEN_CREDENTIALS=AzureCliCredential, weist diese Einstellung die Funktion an, Azure CLI-Anmeldeinformationen für die Authentifizierung deterministisch zu verwenden. Die Funktion basiert auf DefaultAzureCredential aus @azure/Identität , um Ihre Azure-Anmeldeinformationen in der Umgebung zu finden. Erfahren Sie mehr darüber, wie Sie JavaScript-Apps mit Azure-Diensten mithilfe der Azure Identity-Bibliothek authentifizieren.

Erstellen und Ausführen der Anwendung

Erstellen Sie die TypeScript-Dateien, und führen Sie dann die Anwendung aus:

npm run build
npm run start:diskann

Die App-Protokollierung und -Ausgabe zeigen Folgendes an:

  • Sammlungserstellungs- und Dateneinfügungsstatus
  • Vektorindexerstellung
  • Suchergebnisse mit Hotelnamen und Ähnlichkeitsbewertungen
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

Anzeigen und Verwalten von Daten in Visual Studio Code

  1. Wählen Sie die DocumentDB-Erweiterung in Visual Studio Code aus, um eine Verbindung mit Ihrem Azure DocumentDB-Konto herzustellen.

  2. Zeigen Sie die Daten und Indizes der Datenbank "Hotels" an.

    Screenshot der DocumentDB-Erweiterung mit der DocumentDB-Auflistung.

Bereinigen von Ressourcen

Löschen Sie die Ressourcengruppe, das DocumentDB-Konto und die Azure OpenAI-Ressource, wenn Sie sie nicht benötigen, um zusätzliche Kosten zu vermeiden.

Verwandte Inhalte

  • Last updated on