vectra 0.3.0 → 0.4.0

package/package.json

@@ -2,7 +2,7 @@
     "name": "vectra",
     "author": "Steven Ickman",
     "description": "A vector database that uses the local file system for storage.",
-    "version": "0.3.0",
+    "version": "0.4.0",
     "license": "MIT",
     "keywords": [
         "gpt"

package/src/FileFetcher.ts

@@ -0,0 +1,31 @@
+import { TextFetcher } from './types';
+import * as fs from 'fs/promises';
+import * as path from 'path';
+
+export class FileFetcher implements TextFetcher {
+    public async fetch(uri: string, onDocument: (uri: string, text: string, docType?: string | undefined) => Promise<boolean>): Promise<boolean> {
+        // Does path exist and is it a directory?
+        let isDirectory: boolean;
+        try {
+            const stat = await fs.stat(uri);
+            isDirectory = stat.isDirectory();
+        } catch {
+            return true;
+        }
+
+        // If directory, read all files and recurse
+        if (isDirectory) {
+            const files = await fs.readdir(uri);
+            for (const file of files) {
+                const filePath = path.join(uri, file);
+                await this.fetch(filePath, onDocument);
+            }
+            return true;
+        } else {
+            // Read file and call onDocument
+            const text = await fs.readFile(uri, 'utf8');
+            const parts = uri.split('.');
+            return await onDocument(uri, text, parts.length > 0 ? parts[parts.length - 1].toLowerCase() : undefined);
+        }
+    }
+}

package/src/LocalIndex.ts

@@ -20,15 +20,19 @@ export interface CreateIndexConfig {
  */
 export class LocalIndex {
     private readonly _folderPath: string;
+    private readonly _indexName: string;
+
     private _data?: IndexData;
     private _update?: IndexData;
 
     /**
      * 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
      */
-    public constructor(folderPath: string) {
+    public constructor(folderPath: string, indexName?: string) {
         this._folderPath = folderPath;
+        this._indexName = indexName || "index.json";
     }
 
     /**
@@ -38,6 +42,13 @@ export class LocalIndex {
         return this._folderPath;
     }
 
+    /**
+     * Optional name of the index file. 
+     */
+    public get indexName(): string {
+        return this._indexName;
+    }
+
     /**
      * Begins an update to the index.
      * @remarks
@@ -87,7 +98,8 @@ export class LocalIndex {
                 metadata_config: config.metadata_config ?? {},
                 items: []
             };
-            await fs.writeFile(path.join(this._folderPath, 'index.json'), JSON.stringify(this._data));
+
+            await fs.writeFile(path.join(this._folderPath, this._indexName), JSON.stringify(this._data));
         } catch (err: unknown) {
             await this.deleteIndex();
             throw new Error('Error creating index');
@@ -139,7 +151,7 @@ export class LocalIndex {
 
         try {
             // Save index
-            await fs.writeFile(path.join(this._folderPath, 'index.json'), JSON.stringify(this._update));
+            await fs.writeFile(path.join(this._folderPath, this._indexName), JSON.stringify(this._update));
             this._data = this._update;
             this._update = undefined;
         } catch(err: unknown) {
@@ -194,7 +206,7 @@ export class LocalIndex {
      */
     public async isIndexCreated(): Promise<boolean> {
         try {
-            await fs.access(path.join(this._folderPath, 'index.json'));
+            await fs.access(path.join(this._folderPath, this.indexName));
             return true;
         } catch (err: unknown) {
             return false;
@@ -307,7 +319,7 @@ export class LocalIndex {
             throw new Error('Index does not exist');
         }
 
-        const data = await fs.readFile(path.join(this._folderPath, 'index.json'));
+        const data = await fs.readFile(path.join(this._folderPath, this.indexName));
         this._data = JSON.parse(data.toString());
     }
 

package/src/OpenAIEmbeddings.ts

@@ -140,7 +140,8 @@ export class OpenAIEmbeddings implements EmbeddingsModel {
 
         // Process response
         if (response.status < 300) {
-            return { status: 'success', output: response.data.data.sort((a, b) => a.index - b.index).map((item) => item.embedding) };
+            const {data,model,usage} = response.data
+            return { status: 'success', output: data.sort((a, b) => a.index - b.index).map((item) => item.embedding), model, usage };
         } else if (response.status == 429) {
             return { status: 'rate_limited', message: `The embeddings API returned a rate limit error.` }
         } else {
@@ -205,4 +206,4 @@ export class OpenAIEmbeddings implements EmbeddingsModel {
             return response;
         }
     }
-}
+}

package/src/WebFetcher.ts

@@ -3,7 +3,6 @@ import { TextFetcher } from './types';
 import * as cheerio from 'cheerio';
 import TurndownService  from 'turndown';
 
-
 const ALLOWED_CONTENT_TYPES = [
     "text/html",
     "application/json",
@@ -45,7 +44,7 @@ export class WebFetcher implements TextFetcher {
         } as WebFetcherConfig, config);
     }
 
-    public async fetch(uri: string): Promise<{ text: string; docType: string|undefined; }> {
+    public async fetch(uri: string, onDocument: (uri: string, text: string, docType?: string) => Promise<boolean>): Promise<boolean> {
         const httpClient = axios.create({
             validateStatus: () => true,
         });
@@ -78,10 +77,10 @@ export class WebFetcher implements TextFetcher {
         const docType = contentTypeArray[0] != 'text/plain' ? contentTypeArray[0].split('/')[1] : undefined;
         if (docType == 'html' && this._config.htmlToMarkdown) {
             const text = this.htmlToMarkdown(response.data, uri);
-            return {text, docType: 'md'};
+            return await onDocument(uri, text, 'md');
         } else {
             const text = response.data;
-            return {text, docType};
+            return await onDocument(uri, text, docType);
         }
     }
 

package/src/index.ts

@@ -1,3 +1,4 @@
+export * from './FileFetcher';
 export * from './GPT3Tokenizer';
 export * from './ItemSelector';
 export * from './LocalIndex';

package/src/types.ts

@@ -56,7 +56,7 @@ export interface TextChunk {
 }
 
 export interface TextFetcher {
-    fetch(uri: string): Promise<{ text: string; docType: string|undefined; }>;
+    fetch(uri: string, onDocument: (uri: string, text: string, docType?: string) => Promise<boolean>): Promise<boolean>;
 }
 
 export interface IndexStats {

package/src/vectra-cli.ts

@@ -5,6 +5,7 @@ import { LocalDocumentIndex } from "./LocalDocumentIndex";
 import { WebFetcher } from './WebFetcher';
 import { OpenAIEmbeddings } from './OpenAIEmbeddings';
 import { Colorize } from './internals';
+import { FileFetcher } from './FileFetcher';
 
 export async function run() {
     // prettier-ignore
@@ -77,16 +78,20 @@ export async function run() {
             const uris = await getItemList(args.uri as string[], args.list as string, 'web page');
 
             // Fetch web pages
-            const fetcher = new WebFetcher();
-            for (const uri of uris) {
+            const fileFetcher = new FileFetcher();
+            const webFetcher = new WebFetcher();
+            for (const path of uris) {
                 try {
-                    console.log(Colorize.progress(`fetching ${uri}`));
-                    const { text, docType } =  await fetcher.fetch(uri);
-                    console.log(Colorize.replaceLine(Colorize.progress(`indexing ${uri}`)));
-                    await index.upsertDocument(uri, text, docType);
-                    console.log(Colorize.replaceLine(Colorize.success(`added ${uri}`)));
+                    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: ${uri}\n${(err as Error).message}`)));
+                    console.log(Colorize.replaceLine(Colorize.error(`Error adding: ${path}\n${(err as Error).message}`)));
                 }
             }
         })