vectra 0.1.1 → 0.2.0

package/src/WebFetcher.ts

@@ -0,0 +1,128 @@
+import axios, { AxiosRequestConfig } from "axios";
+import * as cheerio from "cheerio";
+import { TextFetcher } from './types';
+
+
+const ALLOWED_CONTENT_TYPES = [
+    "text/html",
+    "application/json",
+    "application/xml",
+    "application/javascript",
+    "text/plain",
+];
+
+
+const DEFAULT_HEADERS = {
+    Accept: "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
+    "Accept-Encoding": "gzip, deflate",
+    "Accept-Language": "en-US,en;q=0.5",
+    "Alt-Used": "LEAVE-THIS-KEY-SET-BY-TOOL",
+    Connection: "keep-alive",
+    Host: "LEAVE-THIS-KEY-SET-BY-TOOL",
+    Referer: "https://www.google.com/",
+    "Sec-Fetch-Dest": "document",
+    "Sec-Fetch-Mode": "navigate",
+    "Sec-Fetch-Site": "cross-site",
+    "Upgrade-Insecure-Requests": "1",
+    "User-Agent": "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:109.0) Gecko/20100101 Firefox/111.0",
+};
+
+export interface WebFetcherConfig {
+    headers?: Record<string,string>;
+    requestConfig?: AxiosRequestConfig;
+    htmlToText: boolean;
+    summarizeHtml: boolean;
+}
+
+export class WebFetcher implements TextFetcher {
+    private readonly _config: WebFetcherConfig;
+
+    public constructor(config?: Partial<WebFetcherConfig>) {
+        this._config = Object.assign({
+            htmlToText: true,
+            summarizeHtml: false,
+        } as WebFetcherConfig, config);
+    }
+
+    public async fetch(uri: string): Promise<string> {
+        const {data, contentType} = await this.fetchPage(uri);
+        if (contentType === "text/html" && this._config.htmlToText) {
+            return this.extractText(data, uri, this._config.summarizeHtml);
+        } else {
+            return data;
+        }
+    }
+
+    private extractText(html: string, baseUrl: string, summarize: boolean): string {
+        // Parse all elements including <noscript> tags
+        const $ = cheerio.load(html, { scriptingEnabled: true });
+
+        // If we want a summary, just get use the <body/>
+        let text = '';
+        $(`${summarize ? 'body ' : '*'}:not(style):not(script):not(svg)`).each((i, elem: any) => {
+            // Remove any children to avoid duplicate text
+            let content = $(elem).clone().children().remove().end().text().trim();
+            const $el = $(elem);
+
+            // Print links in markdown format
+            let href = $el.attr("href");
+            if ($el.prop("tagName")?.toLowerCase() === "a" && href) {
+                if (!href.startsWith("http")) {
+                    // Try converting to a relevant link
+                    try {
+                        href = new URL(href, baseUrl).toString();
+                    } catch {
+                        // Leave as is
+                    }
+                }
+
+                // If the link has content, use that as the text
+                const altText = $el.find("img[alt]").attr("alt")?.trim();
+                if (altText) {
+                    content += ` ${altText}`;
+                }
+
+                text += ` [${content}](${href})`;
+            }
+            // otherwise just print the content
+            else if (content !== "") {
+                text += ` ${content}`;
+            }
+        });
+
+        // Remove newlines
+        return text.trim().replace(/\n+/g, ' ');
+    }
+
+    private async fetchPage(baseUrl: string): Promise<{data: string; contentType: string;}> {
+        const httpClient = axios.create({
+            validateStatus: () => true,
+        });
+
+        // Clone headers to avoid mutating the original
+        const headers = Object.assign({}, DEFAULT_HEADERS, this._config.headers)
+
+        // get hostname from url
+        const host = new URL(baseUrl).hostname;
+        headers['Host'] = host;
+        headers['Alt-Used'] = host;
+
+        // Fetch page and check for errors
+        const response = await httpClient.get(baseUrl, {
+            headers,
+            ...this._config.requestConfig,
+        });
+        if (response.status >= 400) {
+            throw new Error(`Site returned an HTTP status of ${response.status}`);
+        }
+
+        // Check for valid content type
+        const contentType = response.headers['content-type'];
+        const contentTypeArray = contentType.split(';');
+        if (!contentTypeArray[0] || !ALLOWED_CONTENT_TYPES.includes(contentTypeArray[0])) {
+            throw new Error(`Site returned an invalid content type of ${contentType}`);
+        }
+
+        return {data: response.data, contentType: contentTypeArray[0]};
+    }
+}

package/src/index.ts

@@ -1,2 +1,10 @@
+export * from './GPT3Tokenizer';
 export * from './ItemSelector';
 export * from './LocalIndex';
+export * from './LocalDocument';
+export * from './LocalDocumentIndex';
+export * from './LocalDocumentResult';
+export * from './OpenAIEmbeddings';
+export * from './TextSplitter';
+export * from './types';
+export * from './WebFetcher';

package/src/internals/Colorize.ts

@@ -0,0 +1,64 @@
+const colorizer = require('json-colorizer');
+
+/**
+ * @private
+ */
+export class Colorize {
+    public static replaceLine(text: string): string {
+        return '\x1b[A\x1b[2K' + text;
+    }
+
+    public static error(error: Error|string): string {
+        if (typeof error === 'string') {
+            return `\x1b[31;1m${error}\x1b[0m`;
+        } else {
+            return `\x1b[31;1m${error.message}\x1b[0m`;
+        }
+    }
+
+    public static output(output: object | string, quote: string = '', units: string = ''): string {
+        if (typeof output === 'string') {
+            return `\x1b[32m${quote}${output}${quote}\x1b[0m`;
+        } else if (typeof output === 'object' && output !== null) {
+            return colorizer(output, {
+                pretty: true,
+                colors: {
+                    BRACE: 'white',
+                    BRACKET: 'white',
+                    COLON: 'white',
+                    COMMA: 'white',
+                    STRING_KEY: 'white',
+                    STRING_LITERAL: 'green',
+                    NUMBER_LITERAL: 'blue',
+                    BOOLEAN_LITERAL: 'blue',
+                    NULL_LITERAL: 'blue'
+                }
+            });
+        } else if (typeof output == 'number') {
+            return `\x1b[34m${output}${units}\x1b[0m`;
+        } else {
+            return `\x1b[34m${output}\x1b[0m`;
+        }
+    }
+
+    public static progress(message: string): string {
+        return message;
+    }
+
+    public static success(message: string): string {
+        return `\x1b[32;1m${message}\x1b[0m`;
+    }
+
+    public static title(title: string): string {
+        return `\x1b[35;1m${title}\x1b[0m`;
+    }
+
+    public static value(field: string, value: any, units: string = ''): string {
+        return `${field}: ${Colorize.output(value, '"', units)}`;
+    }
+
+    public static warning(warning: string): string {
+        return `\x1b[33m${warning}\x1b[0m`;
+
+    }
+}

package/src/internals/index.ts

@@ -0,0 +1,2 @@
+export * from "./Colorize";
+export * from "./types";

package/src/internals/types.ts

@@ -0,0 +1,46 @@
+/**
+ * @private
+ */
+export interface CreateEmbeddingRequest {
+    input: CreateEmbeddingRequestInput;
+    user?: string;
+}
+
+/**
+ * @private
+ */
+export interface OpenAICreateEmbeddingRequest extends CreateEmbeddingRequest {
+    model: string;
+}
+
+/**
+ * @private
+ */
+export interface CreateEmbeddingResponse {
+    object: string;
+    model: string;
+    data: Array<CreateEmbeddingResponseDataInner>;
+    usage: CreateEmbeddingResponseUsage;
+}
+
+/**
+ * @private
+ */
+export interface CreateEmbeddingResponseDataInner {
+    index: number;
+    object: string;
+    embedding: Array<number>;
+}
+
+/**
+ * @private
+ */
+export interface CreateEmbeddingResponseUsage {
+    prompt_tokens: number;
+    total_tokens: number;
+}
+
+/**
+ * @private
+ */
+export type CreateEmbeddingRequestInput = Array<any> | Array<number> | Array<string> | string;

package/src/types.ts

@@ -0,0 +1,160 @@
+
+
+/**
+ * An AI model that can be used to create embeddings.
+ */
+export interface EmbeddingsModel {
+    /**
+     * Creates embeddings for the given inputs.
+     * @param inputs Text inputs to create embeddings for.
+     * @returns A `EmbeddingsResponse` with a status and the generated embeddings or a message when an error occurs.
+     */
+    createEmbeddings(inputs: string|string[]): Promise<EmbeddingsResponse>;
+}
+
+/**
+ * Status of the embeddings response.
+ * @remarks
+ * `success` - The embeddings were successfully created.
+ * `error` - An error occurred while creating the embeddings.
+ * `rate_limited` - The request was rate limited.
+ */
+export type EmbeddingsResponseStatus = 'success' | 'error' | 'rate_limited';
+
+/**
+ * Response returned by a `EmbeddingsClient`.
+ */
+export interface EmbeddingsResponse {
+    /**
+     * Status of the embeddings response.
+     */
+    status: EmbeddingsResponseStatus;
+
+    /**
+     * Optional. Embeddings for the given inputs.
+     */
+    output?: number[][];
+
+    /**
+     * Optional. Message when status is not equal to `success`.
+     */
+    message?: string;
+}
+
+export interface TextChunk {
+    text: string;
+    tokens: number[];
+    startPos: number;
+    endPos: number;
+    startOverlap: number[];
+    endOverlap: number[];
+}
+
+export interface TextFetcher {
+    fetch(uri: string): Promise<string>;
+}
+
+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 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[];
+
+    [key: string]: unknown;
+}
+
+export type MetadataTypes = number|string|boolean;
+
+export interface QueryResult<TMetadata = Record<string,MetadataTypes>> {
+    item: IndexItem<TMetadata>;
+    score: number;
+}
+
+export interface Tokenizer {
+    decode(tokens: number[]): string;
+    encode(text: string): number[];
+}
+
+export interface DocumentChunkMetadata {
+    documentId: string;
+    startPos: number;
+    endPos: number;
+    [key: string]: MetadataTypes;
+}
+
+export interface DocumentCatalogStats {
+    version: number;
+    documents: number;
+    chunks: number;
+    metadata_config: {
+        indexed?: string[];
+    };
+}
+
+export interface DocumentTextSection {
+    text: string;
+    tokenCount: number;
+    score: number;
+}

package/src/vectra-cli.ts

@@ -0,0 +1,238 @@
+import * as fs from 'fs/promises';
+import yargs from "yargs/yargs";
+import { hideBin } from "yargs/helpers";
+import { LocalDocumentIndex } from "./LocalDocumentIndex";
+import { WebFetcher } from './WebFetcher';
+import { OpenAIEmbeddings } from './OpenAIEmbeddings';
+import { Colorize } from './internals';
+import { Console } from 'console';
+
+export async function run() {
+    // prettier-ignore
+    const args = await yargs(hideBin(process.argv))
+        .command('create <index>', `create a new local index`, {}, async (args) => {
+            const folderPath = args.index as string;
+            const index = new LocalDocumentIndex({ folderPath });
+            console.log(Colorize.output(`creating index at ${folderPath}`));
+            await index.createIndex({ version: 1, deleteIfExists: true });
+        })
+        .command('delete <index>', `delete an existing local index`, {}, async (args) => {
+            const folderPath = args.index as string;
+            console.log(Colorize.output(`deleting index at ${folderPath}`));
+            const index = new LocalDocumentIndex({ folderPath });
+            await index.deleteIndex();
+        })
+        .command('add-web <index>', `adds one or more web pages to an index`, (yargs) => {
+            return yargs
+                .option('keys', {
+                    alias: 'k',
+                    describe: 'path of a JSON file containing the model keys to use for generating embeddings',
+                    type: 'string'
+                })
+                .option('uri', {
+                    alias: 'u',
+                    array: true,
+                    describe: 'http/https link to a web page to add',
+                    type: 'string'
+                })
+                .option('list', {
+                    alias: 'l',
+                    describe: 'path to a file containing a list of web pages to add',
+                    type: 'string'
+                })
+                .option('chunk-size', {
+                    alias: 'cs',
+                    describe: 'size of the generated chunks in tokens (defaults to 512)',
+                    type: 'number',
+                    default: 512
+                })
+                .check((argv) => {
+                    if (Array.isArray(argv.uri) && argv.uri.length > 0) {
+                        return true;
+                    } else if (typeof argv.list == 'string' && argv.list.trim().length > 0) {
+                        return true;
+                    } else {
+                        throw new Error(`you must specify either one or more "--uri <link>" for the pages to add or a "--list <file path>" for a file containing the list of pages to add.`);
+                    }
+                })
+                .demandOption(['keys']);
+        }, async (args) => {
+            console.log(Colorize.title('Adding Web Pages to Index'));
+
+            // Create embeddings
+            const keys = JSON.parse(await fs.readFile(args.keys as string, 'utf-8'));
+            const embeddings = new OpenAIEmbeddings(Object.assign({ model: 'text-embedding-ada-002' }, keys));
+
+            // Initialize index
+            const folderPath = args.index as string;
+            const index = new LocalDocumentIndex({
+                folderPath,
+                embeddings,
+                chunkingConfig: {
+                    chunkSize: args.chunkSize
+                }
+            });
+
+            // Get list of url's
+            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) {
+                try {
+                    console.log(Colorize.progress(`fetching ${uri}`));
+                    const content =  await fetcher.fetch(uri);
+                    console.log(Colorize.replaceLine(Colorize.progress(`indexing ${uri}`)));
+                    await index.upsertDocument(uri, content);
+                    console.log(Colorize.replaceLine(Colorize.success(`added ${uri}`)));
+                } catch (err: unknown) {
+                    console.log(Colorize.replaceLine(Colorize.error(`Error adding: ${uri}\n${(err as Error).message}`)));
+                }
+            }
+        })
+        .command('remove <index>', `removes one or more documents from an index`, (yargs) => {
+            return yargs
+                .option('uri', {
+                    alias: 'u',
+                    array: true,
+                    describe: 'uri of a document to remove',
+                    type: 'string'
+                })
+                .option('list', {
+                    alias: 'l',
+                    describe: 'path to a file containing a list of documents to remove',
+                    type: 'string'
+                })
+                .check((argv) => {
+                    if (Array.isArray(argv.uri) && argv.uri.length > 0) {
+                        return true;
+                    } else if (typeof argv.list == 'string' && argv.list.trim().length > 0) {
+                        return true;
+                    } else {
+                        throw new Error(`you must specify either one or more "--uri <link>" for the pages to add or a "--list <file path>" for a file containing the list of pages to add.`);
+                    }
+                });
+        }, async (args) => {
+            // Initialize index
+            const folderPath = args.index as string;
+            const index = new LocalDocumentIndex({ folderPath });
+
+            // Get list of uri's
+            const uris = await getItemList(args.uri as string[], args.list as string, 'document');
+
+            // Remove documents
+            for (const uri of uris) {
+                console.log(`removing ${uri}`);
+                await index.deleteDocument(uri);
+            }
+        })
+        .command('stats <index>', `prints the stats for a local index`, {}, async (args) => {
+            const folderPath = args.index as string;
+            const index = new LocalDocumentIndex({ folderPath });
+            const stats = await index.getCatalogStats();
+            console.log(Colorize.title('Index Stats'));
+            console.log(Colorize.output(stats));
+        })
+        .command('query <index> <query>', `queries a local index`, (yargs) => {
+            return yargs
+                .option('keys', {
+                    alias: 'k',
+                    describe: 'path of a JSON file containing the model keys to use for generating embeddings'
+                })
+                .option('document-count', {
+                    alias: 'dc',
+                    describe: 'max number of documents to return (defaults to 10)',
+                    type: 'count',
+                    default: 10
+                })
+                .option('chunk-count', {
+                    alias: 'cc',
+                    describe: 'max number of chunks to return (defaults to 50)',
+                    type: 'count',
+                    default: 50
+                })
+                .option('section-count', {
+                    alias: 'sc',
+                    describe: 'max number of document sections to render (defaults to 1)',
+                    type: 'count',
+                    default: 1
+                })
+                .option('tokens', {
+                    alias: 't',
+                    describe: 'max number of tokens to render for each document section (defaults to 2000)',
+                    type: 'count',
+                    default: 2000
+                })
+                .option('format', {
+                    alias: 'f',
+                    describe: `format of the rendered results. Defaults to 'sections'`,
+                    choices: ['sections', 'stats', 'chunks'],
+                    default: 'sections'
+                })
+                .demandOption(['keys']);
+        }, async (args) => {
+            console.log(Colorize.title('Querying Index'));
+
+            // Create embeddings
+            const keys = JSON.parse(await fs.readFile(args.keys as string, 'utf-8'));
+            const embeddings = new OpenAIEmbeddings(Object.assign({ model: 'text-embedding-ada-002' }, keys));
+
+            // Initialize index
+            const folderPath = args.index as string;
+            const index = new LocalDocumentIndex({
+                folderPath,
+                embeddings
+            });
+
+            // Query index
+            const query = args.query as string;
+            const results = await index.queryDocuments(query, {
+                maxDocuments: args.documentCount,
+                maxChunks: args.chunkCount,
+            });
+
+            // Render results
+            for (const result of results) {
+                console.log(Colorize.output(result.uri));
+                console.log(Colorize.value('score', result.score));
+                console.log(Colorize.value('chunks', result.chunks.length));
+                if (args.format == 'sections') {
+                    const sections = await result.renderSections(args.tokens, args.sectionCount);
+                    for (let i = 0; i < sections.length; i++) {
+                        const section = sections[i];
+                        console.log(Colorize.title(args.sectionCount > 1 ? 'Section' : `Section ${1}`));
+                        console.log(Colorize.value('score', section.score));
+                        console.log(Colorize.value('tokens', section.tokenCount));
+                        console.log(Colorize.output(section.text));
+                    }
+                } else if (args.format == 'chunks') {
+                    const text = await result.loadText();
+                    for (let i = 0; i < result.chunks.length; i++) {
+                        const chunk = result.chunks[i];
+                        const startPos = chunk.item.metadata.startPos;
+                        const endPos = chunk.item.metadata.endPos;
+                        console.log(Colorize.title(`Chunk ${i + 1}`));
+                        console.log(Colorize.value('score', chunk.score));
+                        console.log(Colorize.value('startPos', startPos));
+                        console.log(Colorize.value('endPos', endPos));
+                        console.log(Colorize.output(text.substring(startPos, endPos + 1)));
+                    }
+                }
+            }
+        })
+        .help()
+        .demandCommand()
+        .parseAsync();
+}
+
+
+async function getItemList(items: string[], listFile: string, uriType: string): Promise<string[]> {
+    if (Array.isArray(items) && items.length > 0) {
+        return items;
+    } else if (typeof listFile == 'string' && listFile.trim().length > 0) {
+        const list = await fs.readFile(listFile, 'utf-8');
+        return list.split('\n').map((item) => item.trim()).filter((item) => item.length > 0);
+    } else {
+        throw new Error(`you must specify either one or more "--uri <${uriType}>" for the items or a "--list <file path>" for a file containing the items.`)
+    }
+}