Condividi tramite

Avvio rapido: Ricerca vettoriale con Node.js in Azure DocumentDB

Usare la ricerca vettoriale in Azure DocumentDB con la libreria client Node.js. Archiviare ed eseguire query sui dati vettoriali in modo efficiente.

Questa guida di avvio rapido utilizza un set di dati di esempio di hotel in un file JSON con vettori del modello text-embedding-3-small. Il set di dati include nomi di hotel, località, descrizioni e incorporamenti vettoriali.

Trovare il codice di esempio in GitHub.

Prerequisiti

  • Una sottoscrizione di Azure

Creare un file di dati con vettori

  1. Creare una nuova cartella dati per il file di dati degli hotel:

    mkdir data

  2. Copiare il Hotels_Vector.jsonfile di dati non elaborato con vettori nella data directory.

Creare un progetto di Node.js

  1. Creare una nuova directory di pari livello per il progetto, allo stesso livello della directory dati e aprirla in Visual Studio Code:

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

  2. Nel terminale inizializzare un progetto Node.js:

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

  3. Installare i pacchetti necessari:

    npm install mongodb @azure/identity openai @types/node
    • mongodb: driver Node.js di MongoDB
    • @azure/identity: libreria di identità di Azure per l'autenticazione senza password
    • openai: libreria client OpenAI per creare vettori
    • @types/node: definizioni di tipi per Node.js

  4. Creare un .env file nella radice del progetto per le variabili di ambiente:

    # 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

    Sostituire i valori segnaposto nel .env file con le proprie informazioni:

    • AZURE_OPENAI_EMBEDDING_ENDPOINT: URL dell'endpoint della risorsa OpenAI di Azure
    • MONGO_CLUSTER_NAME: nome della risorsa

  5. Aggiungere un tsconfig.json file per configurare 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"
    ]
}

Creare script npm

Modificare il package.json file e aggiungere gli script seguenti:

Usare questi script per compilare i file TypeScript ed eseguire l'implementazione dell'indice DiskANN.

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

Crea una cartella src per i file TypeScript. Aggiungere due file: diskann.ts e utils.ts per l'implementazione dell'indice DiskANN:

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

Incollare il codice seguente nel diskann.ts 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;
});

Questo modulo principale offre queste funzionalità:

  • Include funzioni di utilità
  • Crea un oggetto di configurazione per le variabili di ambiente
  • Crea client per Azure OpenAI e DocumentDB
  • Si connette a MongoDB, crea un database e una raccolta, inserisce dati e crea indici standard
  • Crea un indice vettoriale usando IVF, HNSW o DiskANN
  • Crea un incorporamento per un testo di query di esempio usando il client OpenAI. È possibile modificare la query all'inizio del file
  • Esegue una ricerca vettoriale usando l'incorporamento e stampa i risultati

Creare funzioni di utilità

Incollare il codice seguente 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}`);
    });

}

Questo modulo di utilità offre queste funzionalità:

  • JsonData: interfaccia per la struttura dei dati
  • scoreProperty: posizione del punteggio nei risultati della query in base al metodo di ricerca vettoriale
  • getClients: crea e restituisce i client per Azure OpenAI e Azure DocumentDB
  • getClientsPasswordless: crea e restituisce i client per Azure OpenAI e Azure DocumentDB usando l'autenticazione senza password. Abilitare RBAC su entrambe le risorse ed effettuare l'accesso al CLI di Azure
  • readFileReturnJson: legge un file JSON e ne restituisce il contenuto come matrice di JsonData oggetti
  • writeFileJson: scrive una matrice di JsonData oggetti in un file JSON
  • insertData: inserisce i dati in batch in una raccolta MongoDB e crea indici standard nei campi specificati
  • printSearchResults: stampa i risultati di una ricerca vettoriale, inclusi il punteggio e il nome dell'hotel

Eseguire l'autenticazione con l'interfaccia della riga di comando di Azure

Accedere all'interfaccia della riga di comando di Azure prima di eseguire l'applicazione in modo che l'app possa accedere in modo sicuro alle risorse di Azure.

az login

Il codice usa l'autenticazione dello sviluppatore locale per accedere ad Azure DocumentDB e Azure OpenAI con la getClientsPasswordless funzione da utils.ts. Quando si imposta AZURE_TOKEN_CREDENTIALS=AzureCliCredential, questa impostazione indica alla funzione di usare le credenziali dell'interfaccia della riga di comando di Azure per l'autenticazione in modo deterministico. La funzione si basa su DefaultAzureCredential da @azure/identità per trovare le credenziali di Azure nell'ambiente. Altre informazioni su come autenticare le app JavaScript nei servizi di Azure usando la libreria di identità di Azure.

Compilare ed eseguire l'applicazione

Compilare i file TypeScript, quindi eseguire l'applicazione:

npm run build
npm run start:diskann

La registrazione e l'output dell'app mostrano:

  • Stato di creazione e inserimento dati della raccolta
  • Creazione dell'indice vettoriale
  • Risultati della ricerca con nomi di hotel e punteggi di somiglianza
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

Visualizzare e gestire i dati in Visual Studio Code

  1. Selezionare l'estensione DocumentDB in Visual Studio Code per connettersi all'account Azure DocumentDB.

  2. Visualizzare i dati e gli indici nel database Hotels.

    Screenshot dell'estensione DocumentDB che mostra l'insieme DocumentDB.

Pulire le risorse

Eliminare il gruppo di risorse, l'account DocumentDB e la risorsa OpenAI di Azure quando non sono necessari per evitare costi aggiuntivi.

Contenuti correlati

  • Last updated on