ai-utils.js/vector-index/memory/MemoryVectorIndex.js

Build AI applications, chatbots, and agents with JavaScript and TypeScript.

import z from "zod";
import { cosineSimilarity } from "../../util/cosineSimilarity.js";
/**
 * A very simple vector index that stores all entries in memory. Useful when you only have
 * a small number of entries and don't want to set up a real database, e.g. for conversational memory
 * that does not need to be persisted.
 */
export class MemoryVectorIndex {
    constructor() {
        Object.defineProperty(this, "entries", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: new Map()
        });
    }
    static async deserialize({ serializedData, schema, }) {
        let json = JSON.parse(serializedData);
        if (schema != null) {
            json = z
                .array(z.object({
                id: z.string(),
                vector: z.array(z.number()),
                data: schema,
            }))
                .parse(json);
        }
        const vectorIndex = new MemoryVectorIndex();
        vectorIndex.upsertMany(json);
        return vectorIndex;
    }
    async upsertMany(data) {
        for (const entry of data) {
            this.entries.set(entry.id, entry);
        }
    }
    async queryByVector({ queryVector, similarityThreshold, maxResults, }) {
        const results = [...this.entries.values()]
            .map((entry) => ({
            id: entry.id,
            similarity: cosineSimilarity(entry.vector, queryVector),
            data: entry.data,
        }))
            .filter((entry) => similarityThreshold == undefined ||
            entry.similarity == undefined ||
            entry.similarity > similarityThreshold);
        results.sort((a, b) => b.similarity - a.similarity);
        return results.slice(0, maxResults);
    }
    serialize() {
        return JSON.stringify([...this.entries.values()]);
    }
    asIndex() {
        return this;
    }
}