Compartir a través de

Inicio rápido: Búsqueda de vectores con Node.js en Azure DocumentDB

Use la búsqueda vectorial en Azure DocumentDB con la biblioteca cliente de Node.js. Almacene y consulte los datos vectoriales de forma eficaz.

En este inicio rápido se usa un conjunto de datos de hotel de ejemplo en un archivo JSON con vectores del text-embedding-ada-002 modelo. El conjunto de datos incluye nombres de hotel, ubicaciones, descripciones e incrustaciones vectoriales.

Busque el código de ejemplo en GitHub.

Prerrequisitos

  • Una suscripción de Azure

  • Un clúster de Azure DocumentDB existente

Creación de un proyecto de Node.js

  1. Cree un directorio para el proyecto y ábralo en Visual Studio Code:

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

  2. En el terminal, inicialice un proyecto de Node.js:

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

  3. Instale los paquetes necesarios:

    npm install mongodb @azure/identity openai @types/node
    • mongodb: controlador de Node.js de MongoDB
    • @azure/identity: biblioteca de identidades de Azure para la autenticación sin contraseña
    • openai: biblioteca cliente de OpenAI para crear vectores
    • @types/node: definiciones de tipo para Node.js

  4. Cree un .env archivo en la raíz del proyecto para las variables de entorno:

    # 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

    Reemplace los valores de marcador de posición del .env archivo por su propia información:

    • AZURE_OPENAI_EMBEDDING_ENDPOINT: dirección URL del punto de conexión de recursos de Azure OpenAI
    • MONGO_CLUSTER_NAME: nombre del recurso

  5. Agregue un tsconfig.json archivo para configurar 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. Copie el HotelsData_toCosmosDB_Vector.jsonarchivo de datos sin procesar con vectores en la raíz del proyecto.

Creación de scripts de npm

Edite el package.json archivo y agregue estos scripts:

Use estos scripts para compilar archivos TypeScript y ejecutar la implementación del índice DiskANN.

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

Cree un src directorio para los archivos de TypeScript. Agregue dos archivos: diskann.ts y utils.ts para la implementación del índice DiskANN:

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

Pegue el código siguiente en el diskann.ts archivo.

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

Este módulo principal proporciona estas características:

  • Incluye funciones de utilidad
  • Crea un objeto de configuración para variables de entorno
  • Crea clientes para Azure OpenAI y DocumentDB
  • Se conecta a MongoDB, crea una base de datos y una colección, inserta datos y crea índices estándar.
  • Crea un índice vectorial mediante IVF, HNSW o DiskANN
  • Crea una inserción para un texto de consulta de ejemplo mediante el cliente openAI. Puede cambiar la consulta en la parte superior del archivo.
  • Ejecuta una búsqueda de vectores mediante la inserción e imprime los resultados.

Creación de funciones de utilidad

Pegue el código siguiente en 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}`);
    });

}

Este módulo de utilidad proporciona estas características:

  • JsonData: interfaz para la estructura de datos
  • scoreProperty: ubicación de la puntuación en los resultados de la consulta en función del método de búsqueda vectorial.
  • getClients: crea y devuelve clientes para Azure OpenAI y Azure DocumentDB.
  • getClientsPasswordless: crea y devuelve clientes para Azure OpenAI y Azure DocumentDB mediante la autenticación sin contraseña. Habilitación de RBAC en los recursos e inicio de sesión en la CLI de Azure
  • readFileReturnJson: lee un archivo JSON y devuelve su contenido como una matriz de JsonData objetos.
  • writeFileJson: escribe una matriz de JsonData objetos en un archivo JSON.
  • insertData: inserta datos en lotes en una colección de MongoDB y crea índices estándar en campos especificados.
  • printSearchResults: imprime los resultados de una búsqueda vectorial, incluida la puntuación y el nombre del hotel.

Autenticación con la CLI de Azure

Inicie sesión en la CLI de Azure antes de ejecutar la aplicación para que pueda acceder a los recursos de Azure de forma segura.

az login

Compilación y ejecución de la aplicación

Compile los archivos TypeScript y, a continuación, ejecute la aplicación:

npm run build
npm run start:diskann

El registro y la salida de la aplicación muestran:

  • Estado de creación e inserción de datos de colección
  • Creación de índices vectoriales
  • Resultados de búsqueda con nombres de hotel y puntuaciones de similitud
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

Visualización y administración de datos en Visual Studio Code

  1. Seleccione la extensión DocumentDB en Visual Studio Code para conectarse a la cuenta de Azure DocumentDB.

  2. Vea los datos e índices en la base de datos Hotels.

    Captura de pantalla de la extensión de DocumentDB que muestra la colección DocumentDB.

Limpieza de recursos

Elimine el grupo de recursos, la cuenta de DocumentDB y el recurso de Azure OpenAI cuando no los necesite para evitar costos adicionales.

Contenido relacionado

  • Last updated on