vectra 0.5.4 → 0.6.0

package/src/LocalDocumentIndex.ts

@@ -8,19 +8,58 @@ import { MetadataFilter, EmbeddingsModel, Tokenizer, MetadataTypes, EmbeddingsRe
 import { LocalDocumentResult } from './LocalDocumentResult';
 import { LocalDocument } from './LocalDocument';
 
+/**
+ * Options for querying documents in the index.
+ */
 export interface DocumentQueryOptions {
+    /**
+     * Optional. Maximum number of documents to return.
+     * @remarks
+     * Default is 10.
+     */
     maxDocuments?: number;
+    
+    /**
+     * Maximum number of chunks to return per document.
+     * @remarks
+     * Default is 50.
+     */
     maxChunks?: number;
+
+    /**
+     * Optional. Filter to apply to the document metadata.
+     */
     filter?: MetadataFilter;
 }
 
+/**
+ * Configuration settings for a local document index.
+ */
 export interface LocalDocumentIndexConfig {
+    /**
+     * Folder path where the index is stored.
+     */
     folderPath: string;
+
+    /**
+     * Optional. Embeddings model to use for generating document embeddings.
+     */
     embeddings?: EmbeddingsModel;
+
+    /**
+     * Optional. Tokenizer to use for splitting text into tokens.
+     */
     tokenizer?: Tokenizer;
+
+    /**
+     * Optional. Configuration settings for splitting text into chunks.
+     */
     chunkingConfig?: Partial<TextSplitterConfig>;
 }
 
+/**
+ * Represents a local index of documents stored on disk.
+ */
 export class LocalDocumentIndex extends LocalIndex {
     private readonly _embeddings?: EmbeddingsModel;
     private readonly _tokenizer: Tokenizer;
@@ -28,7 +67,10 @@ export class LocalDocumentIndex extends LocalIndex {
     private _catalog?: DocumentCatalog;
     private _newCatalog?: DocumentCatalog;
 
-
+    /**
+     * Creates a new `LocalDocumentIndex` instance.
+     * @param config Configuration settings for the document index.
+     */
     public constructor(config: LocalDocumentIndexConfig) {
         super(config.folderPath);
         this._embeddings = config.embeddings;
@@ -41,6 +83,20 @@ export class LocalDocumentIndex extends LocalIndex {
         this._chunkingConfig.tokenizer = this._tokenizer;
     }
 
+    /**
+     * Returns the embeddings model used by the index (if configured.)
+     */
+    public get embeddings(): EmbeddingsModel | undefined {
+        return this._embeddings;
+    }
+
+    /**
+     * Returns the tokenizer used by the index.
+     */
+    public get tokenizer(): Tokenizer {
+        return this._tokenizer;
+    }
+
     /**
      * Returns true if the document catalog exists.
      */
@@ -53,21 +109,44 @@ export class LocalDocumentIndex extends LocalIndex {
         }
     }
 
+    /**
+     * Returns the document ID for the given URI.
+     * @param uri URI of the document to lookup.
+     * @returns Document ID or undefined if not found.
+     */
     public async getDocumentId(uri: string): Promise<string | undefined> {
         await this.loadIndexData();
         return this._catalog?.uriToId[uri];
     }
 
+    /**
+     * Returns the document URI for the given ID.
+     * @param documentId ID of the document to lookup.
+     * @returns Document URI or undefined if not found.
+     */
     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();
+    /**
+     * Loads the document catalog from disk and returns its stats.
+     * @returns Catalog stats.
+     */
+    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
+        };
     }
 
+    /**
+     * Deletes a document from the index.
+     * @param uri URI of the document to delete.
+     */
     public async deleteDocument(uri: string): Promise<void> {
         // Lookup document ID
         const documentId = await this.getDocumentId(uri);
@@ -114,16 +193,6 @@ export class LocalDocumentIndex extends LocalIndex {
         }
     }
 
-    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
@@ -245,10 +314,44 @@ export class LocalDocumentIndex extends LocalIndex {
         }
 
         // Return document
-        return new LocalDocument(this.folderPath, documentId, uri);
+        return new LocalDocument(this, documentId, uri);
     }
+    
+    /**
+     * Returns all documents in the index.
+     * @remarks
+     * Each document will contain all of the documents indexed chunks.
+     * @returns Array of documents.
+     */
+    public async listDocuments(): Promise<LocalDocumentResult[]> {
+        // Sort chunks by document ID
+        const docs: { [documentId: string]: QueryResult<DocumentChunkMetadata>[]; } = {};
+        const chunks = await this.listItems<DocumentChunkMetadata>();
+        chunks.forEach(chunk => {
+            const metadata = chunk.metadata;
+            if (docs[metadata.documentId] == undefined) {
+                docs[metadata.documentId] = [];
+            }
+            docs[metadata.documentId].push({ item: chunk, score: 1.0 });
+        });
+
+        // Create document results
+        const results: LocalDocumentResult[] = [];
+        for (const documentId in docs) {
+            const uri = await this.getDocumentUri(documentId) as string;
+            const documentResult = new LocalDocumentResult(this, documentId, uri, docs[documentId], this._tokenizer);
+            results.push(documentResult);
+        }
 
+        return results;
+    }
 
+    /**
+     * Queries the index for documents similar to the given query.
+     * @param query Text to query for.
+     * @param options Optional. Query options.
+     * @returns Array of document results.
+     */
     public async queryDocuments(query: string, options?: DocumentQueryOptions): Promise<LocalDocumentResult[]> {
         // Ensure embeddings configured
         if (!this._embeddings) {
@@ -292,7 +395,7 @@ export class LocalDocumentIndex extends LocalIndex {
         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);
+            const documentResult = new LocalDocumentResult(this, documentId, uri, chunks, this._tokenizer);
             documentResults.push(documentResult);
         }
 
@@ -312,6 +415,11 @@ export class LocalDocumentIndex extends LocalIndex {
         this._newCatalog = undefined;
     }
 
+    public async createIndex(config?: CreateIndexConfig): Promise<void> {
+        await super.createIndex(config);
+        await this.loadIndexData();
+    }
+
     public async endUpdate(): Promise<void> {
         await super.endUpdate();
 

package/src/LocalDocumentResult.ts

@@ -1,13 +1,21 @@
 import { LocalDocument } from "./LocalDocument";
+import { LocalDocumentIndex } from "./LocalDocumentIndex";
 import { QueryResult, DocumentChunkMetadata, Tokenizer, DocumentTextSection } from "./types";
 
+/**
+ * Represents a search result for a document stored on disk.
+ */
 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);
+    /**
+     * @private
+     * Internal constructor for `LocalDocumentResult` instances.
+     */
+    public constructor(index: LocalDocumentIndex, id: string, uri: string, chunks: QueryResult<DocumentChunkMetadata>[], tokenizer: Tokenizer) {
+        super(index, id, uri);
         this._chunks = chunks;
         this._tokenizer = tokenizer;
 
@@ -17,30 +25,112 @@ export class LocalDocumentResult extends LocalDocument {
         this._score = score / this._chunks.length;
     }
 
+    /**
+     * Returns the chunks of the document that matched the query.
+     */
     public get chunks(): QueryResult<DocumentChunkMetadata>[] {
         return this._chunks;
     }
 
+    /**
+     * Returns the average score of the document result.
+     */
     public get score(): number {
         return this._score;
     }
 
-    public async renderSections(maxTokens: number, maxSections: number, overlappingChunks = true): Promise<DocumentTextSection[]> {
+    /**
+     * Renders all of the results chunks as spans of text (sections.)
+     * @remarks
+     * The returned sections will be sorted by document order and limited to maxTokens in length.
+     * @param maxTokens Maximum number of tokens per section.
+     * @returns Array of rendered text sections.
+     */
+    public async renderAllSections(maxTokens: number): Promise<DocumentTextSection[]> {
         // Load text from disk
         const text = await this.loadText();
 
-        // First check to see if the entire document is less than maxTokens
-        if (text.length <= (maxTokens * 8)) {
-            const tokens = this._tokenizer.encode(text);
-            if (tokens.length < maxTokens) {
-                return [{
-                    text,
-                    tokenCount: tokens.length,
-                    score: 1.0
-                }];
+        // Add chunks to a temp array and split any chunks that are longer than maxTokens.
+        const chunks: SectionChunk[]  = [];
+        for (let i = 0; i < this._chunks.length; i++) {
+            const chunk = this._chunks[i];
+            const startPos = chunk.item.metadata.startPos;
+            const endPos = chunk.item.metadata.endPos;
+            const chunkText = text.substring(startPos, endPos + 1);
+            const tokens = this._tokenizer.encode(chunkText);
+            let offset = 0;
+            while (offset < tokens.length) {
+                const chunkLength = Math.min(maxTokens, tokens.length - offset);
+                chunks.push({
+                    text: this._tokenizer.decode(tokens.slice(offset, offset + chunkLength)),
+                    startPos: startPos + offset,
+                    endPos: startPos + offset + chunkLength - 1,
+                    score: chunk.score,
+                    tokenCount: chunkLength
+                });
+                offset += chunkLength;
             }
         }
 
+        // Sort chunks by startPos
+        const sorted = chunks.sort((a, b) => a.startPos - b.startPos);
+
+        // Generate sections
+        const sections: Section[] = [];
+        for (let i = 0; i < sorted.length; i++) {
+            const chunk = sorted[i];
+            let section = sections[sections.length - 1];
+            if (!section || section.tokenCount + chunk.tokenCount > maxTokens) {
+                section = {
+                    chunks: [],
+                    score: 0,
+                    tokenCount: 0
+                };
+                sections.push(section);
+            }
+            section.chunks.push(chunk);
+            section.score += chunk.score;
+            section.tokenCount += chunk.tokenCount;
+        }
+        
+        // Normalize section scores
+        sections.forEach(section => section.score /= section.chunks.length);
+
+        // 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
+            };
+        });
+    }
+
+    /**
+     * Renders the top spans of text (sections) of the document based on the query result.
+     * @remarks
+     * The returned sections will be sorted by relevance and limited to the top `maxSections`.
+     * @param maxTokens Maximum number of tokens per section.
+     * @param maxSections Maximum number of sections to return.
+     * @param overlappingChunks Optional. If true, overlapping chunks of text will be added to each section until the maxTokens is reached.
+     * @returns Array of rendered text sections.
+     */
+    public async renderSections(maxTokens: number, maxSections: number, overlappingChunks = true): Promise<DocumentTextSection[]> {
+        // Load text from disk
+        const text = await this.loadText();
+
+        // First check to see if the entire document is shorter than maxTokens
+        const length = await this.getLength();
+        if (length <= maxTokens) {
+            return [{
+                text,
+                tokenCount: 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.
@@ -78,24 +168,21 @@ export class LocalDocumentResult extends LocalDocument {
         }
 
         // Generate sections
-        const sections: Section[] = [{
-            chunks: [],
-            score: 0,
-            tokenCount: 0
-        }];
+        const sections: Section[] = [];
         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({
+            if (!section || section.tokenCount + chunk.tokenCount > maxTokens) {
+                section = {
                     chunks: [],
                     score: 0,
                     tokenCount: 0
-                });
+                };
+                sections.push(section);
             }
-            sections[sections.length - 1].chunks.push(chunk);
-            sections[sections.length - 1].score += chunk.score;
-            sections[sections.length - 1].tokenCount += chunk.tokenCount;
+            section.chunks.push(chunk);
+            section.score += chunk.score;
+            section.tokenCount += chunk.tokenCount;
         }
 
         // Normalize section scores

package/src/LocalIndex.ts

@@ -27,8 +27,8 @@ export class LocalIndex {
 
     /**
      * Creates a new instance of LocalIndex.
-     * @param folderPath - Path to the index folder
-     * @param indexName - Optional name of the index file. Defaults to index.json
+     * @param folderPath Path to the index folder.
+     * @param indexName Optional name of the index file. Defaults to index.json.
      */
     public constructor(folderPath: string, indexName?: string) {
         this._folderPath = folderPath;
@@ -76,7 +76,7 @@ export class LocalIndex {
      * Creates a new index.
      * @remarks
      * This method creates a new folder on disk containing an index.json file.
-     * @param config - Index configuration
+     * @param config Index configuration.
      */
     public async createIndex(config: CreateIndexConfig = {version: 1}): Promise<void> {
         // Delete if exists
@@ -121,7 +121,7 @@ export class LocalIndex {
 
     /**
      * Deletes an item from the index.
-     * @param id - Item id
+     * @param id ID of item to delete.
      */
     public async deleteItem(id: string): Promise<void> {
         if (this._update) {
@@ -161,7 +161,7 @@ export class LocalIndex {
 
     /**
      * Loads an index from disk and returns its stats.
-     * @returns Index stats
+     * @returns Index stats.
      */
     public async getIndexStats(): Promise<IndexStats> {
         await this.loadIndexData();
@@ -174,8 +174,8 @@ export class LocalIndex {
 
     /**
      * Returns an item from the index given its ID.
-     * @param id Item id
-     * @returns Item or undefined if not found
+     * @param id ID of the item to retrieve.
+     * @returns Item or undefined if not found.
      */
     public async getItem<TMetadata = Record<string,MetadataTypes>>(id: string): Promise<IndexItem<TMetadata> | undefined> {
         await this.loadIndexData();
@@ -187,8 +187,8 @@ export class LocalIndex {
      * @remarks
      * A new update is started if one is not already in progress. If an item with the same ID
      * already exists, an error will be thrown.
-     * @param item Item to insert
-     * @returns Inserted item
+     * @param item Item to insert.
+     * @returns Inserted item.
      */
     public async insertItem<TMetadata = Record<string,MetadataTypes>>(item: Partial<IndexItem<TMetadata>>): Promise<IndexItem<TMetadata>> {
         if (this._update) {
@@ -218,7 +218,7 @@ export class LocalIndex {
      * @remarks
      * This method loads the index into memory and returns all its items. A copy of the items
      * array is returned so no modifications should be made to the array.
-     * @returns All items in the index
+     * @returns Array of all items in the index.
      */
     public async listItems<TMetadata = Record<string,MetadataTypes>>(): Promise<IndexItem<TMetadata>[]> {
         await this.loadIndexData();
@@ -229,8 +229,8 @@ export class LocalIndex {
      * Returns all items in the index matching the filter.
      * @remarks
      * This method loads the index into memory and returns all its items matching the filter.
-     * @param filter Filter to apply
-     * @returns Items matching the filter
+     * @param filter Filter to apply.
+     * @returns Array of items matching the filter.
      */
     public async listItemsByMetadata<TMetadata = Record<string,MetadataTypes>>(filter: MetadataFilter): Promise<IndexItem<TMetadata>[]> {
         await this.loadIndexData();
@@ -242,10 +242,10 @@ export class LocalIndex {
      * @remarks
      * This method loads the index into memory and returns the top k items that are most similar.
      * An optional filter can be applied to the metadata of the items.
-     * @param vector Vector to query against
-     * @param topK Number of items to return
-     * @param filter Optional filter to apply
-     * @returns Similar items to the vector that matches the filter
+     * @param vector Vector to query against.
+     * @param topK Number of items to return.
+     * @param filter Optional. Filter to apply.
+     * @returns Similar items to the vector that matche the supplied filter.
      */
     public async queryItems<TMetadata = Record<string,MetadataTypes>>(vector: number[], topK: number, filter?: MetadataFilter): Promise<QueryResult<TMetadata>[]> {
         await this.loadIndexData();
@@ -293,8 +293,8 @@ export class LocalIndex {
      * @remarks
      * A new update is started if one is not already in progress. If an item with the same ID
      * already exists, it will be replaced.
-     * @param item Item to insert or replace
-     * @returns Upserted item
+     * @param item Item to insert or replace.
+     * @returns Upserted item.
      */
     public async upsertItem<TMetadata = Record<string,MetadataTypes>>(item: Partial<IndexItem<TMetadata>>): Promise<IndexItem<TMetadata>> {
         if (this._update) {

package/src/vectra-cli.ts

@@ -82,27 +82,22 @@ export async function run() {
             // Get list of url's
             const uris = await getItemList(args.uri as string[], args.list as string, 'web page');
 
-            // Fetch web pages
-            await index.beginUpdate();
-            try {
-                const fileFetcher = new FileFetcher();
-                const webFetcher = args.cookie ? new WebFetcher({ headers: { "cookie": args.cookie }}) : new WebFetcher();
-                for (const path of uris) {
-                    try {
-                        console.log(Colorize.progress(`fetching ${path}`));
-                        const fetcher = path.startsWith('http') ? webFetcher : fileFetcher;
-                        await fetcher.fetch(path, async (uri, text, docType) => {
-                            console.log(Colorize.replaceLine(Colorize.progress(`indexing ${uri}`)));
-                            await index.upsertDocument(uri, text, docType);
-                            console.log(Colorize.replaceLine(Colorize.success(`added ${uri}`)));
-                            return true;
-                        });
-                    } catch (err: unknown) {
-                        console.log(Colorize.replaceLine(Colorize.error(`Error adding: ${path}\n${(err as Error).message}`)));
-                    }
+            // Fetch documents
+            const fileFetcher = new FileFetcher();
+            const webFetcher = args.cookie ? new WebFetcher({ headers: { "cookie": args.cookie }}) : new WebFetcher();
+            for (const path of uris) {
+                try {
+                    console.log(Colorize.progress(`fetching ${path}`));
+                    const fetcher = path.startsWith('http') ? webFetcher : fileFetcher;
+                    await fetcher.fetch(path, async (uri, text, docType) => {
+                        console.log(Colorize.replaceLine(Colorize.progress(`indexing ${uri}`)));
+                        await index.upsertDocument(uri, text, docType);
+                        console.log(Colorize.replaceLine(Colorize.success(`added ${uri}`)));
+                        return true;
+                    });
+                } catch (err: unknown) {
+                    console.log(Colorize.replaceLine(Colorize.error(`Error adding: ${path}\n${(err as Error).message}`)));
                 }
-            } finally {
-                await index.endUpdate();
             }
         })
         .command('remove <index>', `removes one or more documents from an index`, (yargs) => {