vectra 0.1.1 → 0.2.0

package/src/LocalDocumentIndex.ts

@@ -0,0 +1,355 @@
+import * as fs from 'fs/promises';
+import * as path from 'path';
+import { v4 } from 'uuid';
+import { GPT3Tokenizer } from "./GPT3Tokenizer";
+import { CreateIndexConfig, LocalIndex } from "./LocalIndex";
+import { TextSplitter, TextSplitterConfig } from "./TextSplitter";
+import { MetadataFilter, EmbeddingsModel, Tokenizer, MetadataTypes, EmbeddingsResponse, QueryResult, DocumentChunkMetadata, DocumentCatalogStats } from "./types";
+import { LocalDocumentResult } from './LocalDocumentResult';
+import { LocalDocument } from './LocalDocument';
+
+const EMBEDDINGS_BATCH_SIZE = 500;
+
+export interface DocumentQueryOptions {
+    maxDocuments?: number;
+    maxChunks?: number;
+    filter?: MetadataFilter;
+}
+
+export interface LocalDocumentIndexConfig {
+    folderPath: string;
+    embeddings?: EmbeddingsModel;
+    tokenizer?: Tokenizer;
+    chunkingConfig?: Partial<TextSplitterConfig>;
+}
+
+export class LocalDocumentIndex extends LocalIndex {
+    private readonly _embeddings?: EmbeddingsModel;
+    private readonly _tokenizer: Tokenizer;
+    private readonly _chunkingConfig?: TextSplitterConfig;
+    private _catalog?: DocumentCatalog;
+    private _newCatalog?: DocumentCatalog;
+
+
+    public constructor(config: LocalDocumentIndexConfig) {
+        super(config.folderPath);
+        this._embeddings = config.embeddings;
+        this._chunkingConfig = Object.assign({
+            keepSeparators: true,
+            chunkSize: 512,
+            chunkOverlap: 0,
+        } as TextSplitterConfig, config.chunkingConfig);
+        this._tokenizer = config.tokenizer ?? this._chunkingConfig.tokenizer ?? new GPT3Tokenizer();
+        this._chunkingConfig.tokenizer = this._tokenizer;
+    }
+
+    /**
+     * Returns true if the document catalog exists.
+     */
+    public async isCatalogCreated(): Promise<boolean> {
+        try {
+            await fs.access(path.join(this.folderPath, 'catalog.json'));
+            return true;
+        } catch (err: unknown) {
+            return false;
+        }
+    }
+
+    public async getDocumentId(uri: string): Promise<string | undefined> {
+        await this.loadIndexData();
+        return this._catalog?.uriToId[uri];
+    }
+
+    public async getDocumentUri(documentId: string): Promise<string | undefined> {
+        await this.loadIndexData();
+        return this._catalog?.idToUri[documentId];
+    }
+
+    public async createIndex(config?: CreateIndexConfig): Promise<void> {
+        await super.createIndex(config);
+        await this.loadIndexData();
+    }
+
+    public async deleteDocument(uri: string): Promise<void> {
+        // Lookup document ID
+        const documentId = await this.getDocumentId(uri);
+        if (documentId == undefined) {
+            return;
+        }
+
+        // Delete document chunks from index and remove from catalog
+        await this.beginUpdate();
+        try {
+            // Get list of chunks for document
+            const chunks = await this.listItemsByMetadata<DocumentChunkMetadata>({ documentId });
+
+            // Delete chunks
+            for (const chunk of chunks) {
+                await this.deleteItem(chunk.id);
+            }
+
+            // Remove entry from catalog
+            delete this._newCatalog!.uriToId[uri];
+            delete this._newCatalog!.idToUri[documentId];
+            this._newCatalog!.count--;
+
+            // Commit changes
+            await this.endUpdate();
+        } catch (err: unknown) {
+            // Cancel update and raise error
+            this.cancelUpdate();
+            throw new Error(`Error deleting document "${uri}": ${(err as any).toString()}`);
+        }
+
+        // Delete text file from disk
+        try {
+            await fs.unlink(path.join(this.folderPath, `${documentId}.txt`));
+        } catch (err: unknown) {
+            throw new Error(`Error removing text file for document "${uri}" from disk: ${(err as any).toString()}`);
+        }
+
+        // Delete metadata file from disk
+        try {
+            await fs.unlink(path.join(this.folderPath, `${documentId}.json`));
+        } catch (err: unknown) {
+            // Ignore error
+        }
+    }
+
+    public async getCatalogStats(): Promise<DocumentCatalogStats> {
+        const stats = await this.getIndexStats()
+        return {
+            version: this._catalog!.version,
+            documents: this._catalog!.count,
+            chunks: stats.items,
+            metadata_config: stats.metadata_config
+        };
+    }
+
+    /**
+     * Adds a document to the catalog.
+     * @remarks
+     * A new update is started if one is not already in progress. If an document with the same uri
+     * already exists, it will be replaced.
+     * @param item Item to insert
+     * @returns Inserted document
+     */
+    public async upsertDocument(uri: string, text: string, metadata?: Record<string, MetadataTypes>): Promise<LocalDocument> {
+        // Ensure embeddings configured
+        if (!this._embeddings) {
+            throw new Error(`Embeddings model not configured.`);
+        }
+
+        // Check for existing document ID
+        let documentId = await this.getDocumentId(uri);
+        if (documentId != undefined) {
+            // Delete existing document
+            await this.deleteDocument(uri);
+        } else {
+            // Generate new document ID
+            documentId = v4();
+        }
+
+        // Populate docType based on extension
+        const config = Object.assign({}, this._chunkingConfig);
+        const pos = uri.lastIndexOf('.');
+        if (pos >= 0) {
+            const ext = uri.substring(pos + 1).toLowerCase();
+            config.docType = ext;
+        }
+
+        // Split text into chunks
+        const splitter = new TextSplitter(config);
+        const chunks = splitter.split(text);
+
+        // Break chunks into batches for embedding generation
+        const chunkBatches: string[][] = [];
+        let currentBatch: string[] = [];
+        for (const chunk of chunks) {
+            currentBatch.push(chunk.text);
+            if (currentBatch.length >= EMBEDDINGS_BATCH_SIZE) {
+                chunkBatches.push(currentBatch);
+                currentBatch = [];
+            }
+        }
+        if (currentBatch.length > 0) {
+            chunkBatches.push(currentBatch);
+        }
+
+        // Generate embeddings for chunks
+        const embeddings: number[][] = [];
+        for (const batch of chunkBatches) {
+            let response: EmbeddingsResponse;
+            try {
+                response = await this._embeddings.createEmbeddings(batch);
+            } catch (err: unknown) {
+                throw new Error(`Error generating embeddings: ${(err as any).toString()}`);
+            }
+
+            // Check for error
+            if (response.status != 'success') {
+                throw new Error(`Error generating embeddings: ${response.message}`);
+            }
+
+            // Add embeddings to output
+            for (const embedding of response.output!) {
+                embeddings.push(embedding);
+            }
+        }
+
+        // Add document chunks to index
+        await this.beginUpdate();
+        try {
+            // Add chunks to index
+            for (let i = 0; i < chunks.length; i++) {
+                const chunk = chunks[i];
+                const embedding = embeddings[i];
+                const chunkMetadata: DocumentChunkMetadata = Object.assign({
+                    documentId,
+                    startPos: chunk.startPos,
+                    endPos: chunk.endPos,
+                }, metadata);
+                await this.insertItem({
+                    id: v4(),
+                    metadata: chunkMetadata,
+                    vector: embedding,
+                });
+            }
+
+            // Save metadata file to disk
+            if (metadata != undefined) {
+                await fs.writeFile(path.join(this.folderPath, `${documentId}.json`), JSON.stringify(metadata));
+            }
+
+            // Save text file to disk
+            await fs.writeFile(path.join(this.folderPath, `${documentId}.txt`), text);
+
+            // Add entry to catalog
+            this._newCatalog!.uriToId[uri] = documentId;
+            this._newCatalog!.idToUri[documentId] = uri;
+            this._newCatalog!.count++;
+
+            // Commit changes
+            await this.endUpdate();
+        } catch (err: unknown) {
+            // Cancel update and raise error
+            this.cancelUpdate();
+            throw new Error(`Error adding document "${uri}": ${(err as any).toString()}`);
+        }
+
+        // Return document
+        return new LocalDocument(this.folderPath, documentId, uri);
+    }
+
+
+    public async queryDocuments(query: string, options?: DocumentQueryOptions): Promise<LocalDocumentResult[]> {
+        // Ensure embeddings configured
+        if (!this._embeddings) {
+            throw new Error(`Embeddings model not configured.`);
+        }
+
+        // Ensure options are defined
+        options = Object.assign({
+            maxDocuments: 10,
+            maxChunks: 50,
+        }, options);
+
+        // Generate embeddings for query
+        let embeddings: EmbeddingsResponse;
+        try {
+            embeddings = await this._embeddings.createEmbeddings(query);
+        } catch (err: unknown) {
+            throw new Error(`Error generating embeddings for query: ${(err as any).toString()}`);
+        }
+
+        // Check for error
+        if (embeddings.status != 'success') {
+            throw new Error(`Error generating embeddings for query: ${embeddings.message}`);
+        }
+
+        // Query index for chunks
+        const results = await this.queryItems<DocumentChunkMetadata>(embeddings.output![0], options.maxChunks!, options.filter);
+
+        // Group chunks by document
+        const documentChunks: { [documentId: string]: QueryResult<DocumentChunkMetadata>[]; } = {};
+        for (const result  of results) {
+            const metadata = result.item.metadata;
+            if (documentChunks[metadata.documentId] == undefined) {
+                documentChunks[metadata.documentId] = [];
+            }
+            documentChunks[metadata.documentId].push(result);
+        }
+
+        // Create a document result for each document
+        const documentResults: LocalDocumentResult[] = [];
+        for (const documentId in documentChunks) {
+            const chunks = documentChunks[documentId];
+            const uri = await this.getDocumentUri(documentId) as string;
+            const documentResult = new LocalDocumentResult(this.folderPath, documentId, uri, chunks, this._tokenizer);
+            documentResults.push(documentResult);
+        }
+
+        // Sort document results by score and return top results
+        return documentResults.sort((a, b) => b.score - a.score).slice(0, options.maxDocuments!);
+    }
+
+    // Overrides
+
+    public async beginUpdate(): Promise<void> {
+        await super.beginUpdate();
+        this._newCatalog = Object.assign({}, this._catalog);
+    }
+
+    public cancelUpdate(): void {
+        super.cancelUpdate();
+        this._newCatalog = undefined;
+    }
+
+    public async endUpdate(): Promise<void> {
+        await super.endUpdate();
+
+        try {
+            // Save catalog
+            await fs.writeFile(path.join(this.folderPath, 'catalog.json'), JSON.stringify(this._newCatalog));
+            this._catalog = this._newCatalog;
+            this._newCatalog = undefined;
+        } catch(err: unknown) {
+            throw new Error(`Error saving document catalog: ${(err as any).toString()}`);
+        }
+    }
+
+    protected async loadIndexData(): Promise<void> {
+        await super.loadIndexData();
+
+        if (this._catalog) {
+            return;
+        }
+
+        const catalogPath = path.join(this.folderPath, 'catalog.json');
+        if (await this.isCatalogCreated()) {
+            // Load catalog
+            const buffer = await fs.readFile(catalogPath);
+            this._catalog = JSON.parse(buffer.toString());
+        } else {
+            try {
+                // Initialize catalog
+                this._catalog = {
+                    version: 1,
+                    count: 0,
+                    uriToId: {},
+                    idToUri: {},
+                };
+                await fs.writeFile(catalogPath, JSON.stringify(this._catalog));
+            } catch(err: unknown) {
+                throw new Error(`Error creating document catalog: ${(err as any).toString()}`);
+            }
+        }
+    }
+}
+
+interface DocumentCatalog {
+    version: number;
+    count: number;
+    uriToId: { [uri: string]: string; };
+    idToUri: { [id: string]: string; };
+}

package/src/LocalDocumentResult.ts

@@ -0,0 +1,206 @@
+import { LocalDocument } from "./LocalDocument";
+import { QueryResult, DocumentChunkMetadata, Tokenizer, DocumentTextSection } from "./types";
+
+export class LocalDocumentResult extends LocalDocument {
+    private readonly _chunks: QueryResult<DocumentChunkMetadata>[];
+    private readonly _tokenizer: Tokenizer;
+    private readonly _score: number;
+
+    public constructor(folderPath: string, id: string, uri: string, chunks: QueryResult<DocumentChunkMetadata>[], tokenizer: Tokenizer) {
+        super(folderPath, id, uri);
+        this._chunks = chunks;
+        this._tokenizer = tokenizer;
+
+        // Compute average score
+        let score = 0;
+        this._chunks.forEach(chunk => score += chunk.score);
+        this._score = score / this._chunks.length;
+    }
+
+    public get chunks(): QueryResult<DocumentChunkMetadata>[] {
+        return this._chunks;
+    }
+
+    public get score(): number {
+        return this._score;
+    }
+
+    public async renderSections(maxTokens: number, maxSections: number): Promise<DocumentTextSection[]> {
+        // Load text from disk
+        const text = await this.loadText();
+
+        // First check to see if the entire document is less than maxTokens
+        const tokens = this._tokenizer.encode(text);
+        if (tokens.length < maxTokens) {
+            return [{
+                text,
+                tokenCount: tokens.length,
+                score: 1.0
+            }];
+        }
+
+        // Otherwise, we need to split the document into sections
+        // - Add each chunk to a temp array and filter out any chunk that's longer then maxTokens.
+        // - Sort the array by startPos to arrange chunks in document order.
+        // - Generate a new array of sections by combining chunks until the maxTokens is reached for each section.
+        // - Generate an aggregate score for each section by averaging the score of each chunk in the section.
+        // - Sort the sections by score and limit to maxSections.
+        // - For each remaining section combine adjacent chunks of text.
+        // - Dynamically add overlapping chunks of text to each section until the maxTokens is reached.
+        const chunks: SectionChunk[] = this._chunks.map(chunk => {
+            const startPos = chunk.item.metadata.startPos;
+            const endPos = chunk.item.metadata.endPos;
+            const chunkText = text.substring(startPos, endPos + 1);
+            return {
+                text: chunkText,
+                startPos,
+                endPos,
+                score: chunk.score,
+                tokenCount: this._tokenizer.encode(chunkText).length
+            };
+        }).filter(chunk => chunk.tokenCount <= maxTokens).sort((a, b) => a.startPos - b.startPos);
+
+        // Check for no chunks
+        if (chunks.length === 0) {
+            // Take the top chunk and return a subset of its text
+            const topChunk = this._chunks[0];
+            const startPos = topChunk.item.metadata.startPos;
+            const endPos = topChunk.item.metadata.endPos;
+            const chunkText = text.substring(startPos, endPos + 1);
+            const tokens = this._tokenizer.encode(chunkText);
+            return [{
+                text: this._tokenizer.decode(tokens.slice(0, maxTokens)),
+                tokenCount: maxTokens,
+                score: topChunk.score
+            }];
+        }
+
+        // Generate sections
+        const sections: Section[] = [{
+            chunks: [],
+            score: 0,
+            tokenCount: 0
+        }];
+        for (let i = 0; i < chunks.length; i++) {
+            const chunk = chunks[i];
+            let section = sections[sections.length - 1];
+            if (section.tokenCount + chunk.tokenCount > maxTokens) {
+                sections.push({
+                    chunks: [],
+                    score: 0,
+                    tokenCount: 0
+                });
+            }
+            sections[sections.length - 1].chunks.push(chunk);
+            sections[sections.length - 1].score += chunk.score;
+            sections[sections.length - 1].tokenCount += chunk.tokenCount;
+        }
+
+        // Normalize section scores
+        sections.forEach(section => section.score /= section.chunks.length);
+
+        // Sort sections by score and limit to maxSections
+        sections.sort((a, b) => b.score - a.score);
+        if (sections.length > maxSections) {
+            sections.splice(maxSections, sections.length - maxSections);
+        }
+
+        // Combine adjacent chunks of text
+        sections.forEach(section => {
+            for (let i = 0; i < section.chunks.length - 1; i++) {
+                const chunk = section.chunks[i];
+                const nextChunk = section.chunks[i + 1];
+                if (chunk.endPos + 1 === nextChunk.startPos) {
+                    chunk.text += nextChunk.text;
+                    chunk.endPos = nextChunk.endPos;
+                    chunk.tokenCount += nextChunk.tokenCount;
+                    section.chunks.splice(i + 1, 1);
+                    i--;
+                }
+            }
+        });
+
+        // Add overlapping chunks of text to each section until the maxTokens is reached
+        const connector: SectionChunk = {
+            text: '\n\n...\n\n',
+            startPos: -1,
+            endPos: -1,
+            score: 0,
+            tokenCount: this._tokenizer.encode('\n\n...\n\n').length
+        };
+        sections.forEach(section => {
+            // Insert connectors between chunks
+            if (section.chunks.length > 1) {
+                for (let i = 0; i < section.chunks.length - 1; i++) {
+                    section.chunks.splice(i + 1, 0, connector);
+                    section.tokenCount += connector.tokenCount;
+                    i++;
+                }
+            }
+
+            // Add chunks to beginning and end of the section until maxTokens is reached
+            let budget = maxTokens - section.tokenCount;
+            if (budget > 40) {
+                const sectionStart = section.chunks[0].startPos;
+                const sectionEnd = section.chunks[section.chunks.length - 1].endPos;
+                if (sectionStart > 0) {
+                    const beforeTex = text.substring(0, section.chunks[0].startPos);
+                    const beforeTokens = this._tokenizer.encode(beforeTex);
+                    const beforeBudget = sectionEnd < text.length - 1 ? Math.min(beforeTokens.length, Math.ceil(budget/2)) : Math.min(beforeTokens.length, budget);
+                    const chunk: SectionChunk = {
+                        text: this._tokenizer.decode(beforeTokens.slice(-beforeBudget)),
+                        startPos: sectionStart - beforeBudget,
+                        endPos: sectionStart - 1,
+                        score: 0,
+                        tokenCount: beforeBudget
+                    };
+                    section.chunks.unshift(chunk);
+                    section.tokenCount += chunk.tokenCount;
+                    budget -= chunk.tokenCount;
+                }
+
+                if (sectionEnd < text.length - 1) {
+                    const afterText = text.substring(sectionEnd + 1);
+                    const afterTokens = this._tokenizer.encode(afterText);
+                    const afterBudget = Math.min(afterTokens.length, budget);
+                    const chunk: SectionChunk = {
+                        text: this._tokenizer.decode(afterTokens.slice(0, afterBudget)),
+                        startPos: sectionEnd + 1,
+                        endPos: sectionEnd + afterBudget,
+                        score: 0,
+                        tokenCount: afterBudget
+                    };
+                    section.chunks.push(chunk);
+                    section.tokenCount += chunk.tokenCount;
+                    budget -= chunk.tokenCount;
+                }
+            }
+        });
+
+        // Return final rendered sections
+        return sections.map(section => {
+            let text = '';
+            section.chunks.forEach(chunk => text += chunk.text);
+            return {
+                text: text,
+                tokenCount: section.tokenCount,
+                score: section.score
+            };
+        });
+    }
+}
+
+interface SectionChunk {
+    text: string;
+    startPos: number;
+    endPos: number;
+    score: number;
+    tokenCount: number;
+}
+
+interface Section {
+    chunks: SectionChunk[];
+    score: number;
+    tokenCount: number;
+}
+

package/src/LocalIndex.ts

@@ -2,6 +2,7 @@ import * as fs from 'fs/promises';
 import * as path from 'path';
 import { v4 } from 'uuid';
 import { ItemSelector } from './ItemSelector';
+import { IndexItem, IndexStats, MetadataFilter, MetadataTypes, QueryResult } from './types';
 
 export interface CreateIndexConfig {
     version: number;
@@ -11,83 +12,6 @@ export interface CreateIndexConfig {
     };
 }
 
-export interface IndexStats {
-    version: number;
-    metadata_config: {
-        indexed?: string[];
-    };
-    items: number;
-}
-
-export interface IndexItem<TMetadata = Record<string,MetadataTypes>> {
-    id: string;
-    metadata: TMetadata;
-    vector: number[];
-    norm: number;
-    metadataFile?: string;
-}
-
-export interface QueryResult<TMetadata = Record<string,MetadataTypes>> {
-    item: IndexItem<TMetadata>;
-    score: number;
-}
-
-export interface MetadataFilter {
-    [key: string]: MetadataTypes|MetadataFilter|(number|string)[]|MetadataFilter[];
-
-    /**
-     * Equal to (number, string, boolean)
-     */
-    '$eq': number|string|boolean;
-
-    /**
-     * Not equal to (number, string, boolean)
-     */
-    '$ne': number|string|boolean;
-
-    /**
-     * Greater than (number)
-     */
-    '$gt': number;
-
-    /**
-     * Greater than or equal to (number)
-     */
-    '$gte': number;
-
-    /**
-     * Less than (number)
-     */
-    '$lt': number;
-
-    /**
-     * Less than or equal to (number)
-     */
-    '$lte': number;
-
-    /**
-     * In array (string or number)
-     */
-    '$in': (number|string)[];
-
-    /**
-     * Not in array (string or number)
-     */
-    '$nin': (number|string)[];
-
-    /**
-     * AND (MetadataFilter[])
-     */
-    '$and': MetadataFilter[];
-
-    /**
-     * OR (MetadataFilter[])
-     */
-    '$or': MetadataFilter[];
-}
-
-export type MetadataTypes = number|string|boolean;
-
 /**
  * Local vector index instance.
  * @remarks
@@ -107,6 +31,13 @@ export class LocalIndex {
         this._folderPath = folderPath;
     }
 
+    /**
+     * Path to the index folder.
+     */
+    public get folderPath(): string {
+        return this._folderPath;
+    }
+
     /**
      * Begins an update to the index.
      * @remarks
@@ -364,7 +295,10 @@ export class LocalIndex {
         }
     }
 
-    private async loadIndexData(): Promise<void> {
+    /**
+     * Ensures that the index has been loaded into memory.
+     */
+    protected async loadIndexData(): Promise<void> {
         if (this._data) {
             return;
         }