vectra 0.1.2 → 0.2.1

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;
         }

package/src/OpenAIEmbeddings.ts

@@ -0,0 +1,205 @@
+import axios, { AxiosInstance, AxiosResponse, AxiosRequestConfig } from 'axios';
+import { EmbeddingsModel, EmbeddingsResponse } from "./types";
+import { CreateEmbeddingRequest, CreateEmbeddingResponse, OpenAICreateEmbeddingRequest } from "./internals";
+
+export interface BaseOpenAIEmbeddingsOptions {
+    /**
+     * Optional. Retry policy to use when calling the OpenAI API.
+     * @remarks
+     * The default retry policy is `[2000, 5000]` which means that the first retry will be after
+     * 2 seconds and the second retry will be after 5 seconds.
+     */
+    retryPolicy?: number[];
+
+    /**
+     * Optional. Request options to use when calling the OpenAI API.
+     */
+    requestConfig?: AxiosRequestConfig;
+}
+
+/**
+ * Options for configuring an `OpenAIEmbeddings` to generate embeddings using an OpenAI hosted model.
+ */
+export interface OpenAIEmbeddingsOptions extends BaseOpenAIEmbeddingsOptions {
+    /**
+     * API key to use when calling the OpenAI API.
+     * @remarks
+     * A new API key can be created at https://platform.openai.com/account/api-keys.
+     */
+    apiKey: string;
+
+    /**
+     * Model to use for completion.
+     * @remarks
+     * For Azure OpenAI this is the name of the deployment to use.
+     */
+    model: string;
+
+    /**
+     * Optional. Organization to use when calling the OpenAI API.
+     */
+    organization?: string;
+
+    /**
+     * Optional. Endpoint to use when calling the OpenAI API.
+     * @remarks
+     * For Azure OpenAI this is the deployment endpoint.
+     */
+    endpoint?: string;
+}
+
+/**
+ * Options for configuring an `OpenAIEmbeddings` to generate embeddings using an Azure OpenAI hosted model.
+ */
+export interface AzureOpenAIEmbeddingsOptions extends BaseOpenAIEmbeddingsOptions {
+    /**
+     * API key to use when making requests to Azure OpenAI.
+     */
+    azureApiKey: string;
+
+    /**
+     * Deployment endpoint to use.
+     */
+    azureEndpoint: string;
+
+    /**
+     * Name of the Azure OpenAI deployment (model) to use.
+     */
+    azureDeployment: string;
+
+    /**
+     * Optional. Version of the API being called. Defaults to `2023-05-15`.
+     */
+    azureApiVersion?: string;
+}
+
+/**
+ * A `PromptCompletionModel` for calling OpenAI and Azure OpenAI hosted models.
+ * @remarks
+ */
+export class OpenAIEmbeddings implements EmbeddingsModel {
+    private readonly _httpClient: AxiosInstance;
+    private readonly _useAzure: boolean;
+
+    private readonly UserAgent = 'AlphaWave';
+
+    /**
+     * Options the client was configured with.
+     */
+    public readonly options: OpenAIEmbeddingsOptions|AzureOpenAIEmbeddingsOptions;
+
+    /**
+     * Creates a new `OpenAIClient` instance.
+     * @param options Options for configuring an `OpenAIClient`.
+     */
+    public constructor(options: OpenAIEmbeddingsOptions|AzureOpenAIEmbeddingsOptions) {
+        // Check for azure config
+        if ((options as AzureOpenAIEmbeddingsOptions).azureApiKey) {
+            this._useAzure = true;
+            this.options = Object.assign({
+                retryPolicy: [2000, 5000],
+                azureApiVersion: '2023-05-15',
+            }, options) as AzureOpenAIEmbeddingsOptions;
+
+            // Cleanup and validate endpoint
+            let endpoint = this.options.azureEndpoint.trim();
+            if (endpoint.endsWith('/')) {
+                endpoint = endpoint.substring(0, endpoint.length - 1);
+            }
+
+            if (!endpoint.toLowerCase().startsWith('https://')) {
+                throw new Error(`Client created with an invalid endpoint of '${endpoint}'. The endpoint must be a valid HTTPS url.`);
+            }
+
+            this.options.azureEndpoint = endpoint;
+        } else {
+            this._useAzure = false;
+            this.options = Object.assign({
+                retryPolicy: [2000, 5000]
+            }, options) as OpenAIEmbeddingsOptions;
+        }
+
+        // Create client
+        this._httpClient = axios.create({
+            validateStatus: (status) => status < 400 || status == 429
+        });
+    }
+
+    /**
+     * Creates embeddings for the given inputs using the OpenAI API.
+     * @param model Name of the model to use (or deployment for Azure).
+     * @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.
+     */
+    public async createEmbeddings(inputs: string | string[]): Promise<EmbeddingsResponse> {
+        const response = await this.createEmbeddingRequest({
+            input: inputs,
+        });
+
+        // Process response
+        if (response.status < 300) {
+            return { status: 'success', output: response.data.data.sort((a, b) => a.index - b.index).map((item) => item.embedding) };
+        } else if (response.status == 429) {
+            return { status: 'rate_limited', message: `The embeddings API returned a rate limit error.` }
+        } else {
+            return { status: 'error', message: `The embeddings API returned an error status of ${response.status}: ${response.statusText}` };
+        }
+    }
+
+    /**
+     * @private
+     */
+    protected createEmbeddingRequest(request: CreateEmbeddingRequest): Promise<AxiosResponse<CreateEmbeddingResponse>> {
+        if (this._useAzure) {
+            const options = this.options as AzureOpenAIEmbeddingsOptions;
+            const url = `${options.azureEndpoint}/openai/deployments/${options.azureDeployment}/embeddings?api-version=${options.azureApiVersion!}`;
+            return this.post(url, request);
+        } else {
+            const options = this.options as OpenAIEmbeddingsOptions;
+            const url = `${options.endpoint ?? 'https://api.openai.com'}/v1/embeddings`;
+            (request as OpenAICreateEmbeddingRequest).model = options.model;
+            return this.post(url, request);
+        }
+    }
+
+    /**
+     * @private
+     */
+    protected async post<TData>(url: string, body: object, retryCount = 0): Promise<AxiosResponse<TData>> {
+        // Initialize request config
+        const requestConfig: AxiosRequestConfig = Object.assign({}, this.options.requestConfig);
+
+        // Initialize request headers
+        if (!requestConfig.headers) {
+            requestConfig.headers = {};
+        }
+        if (!requestConfig.headers['Content-Type']) {
+            requestConfig.headers['Content-Type'] = 'application/json';
+        }
+        if (!requestConfig.headers['User-Agent']) {
+            requestConfig.headers['User-Agent'] = this.UserAgent;
+        }
+        if (this._useAzure) {
+            const options = this.options as AzureOpenAIEmbeddingsOptions;
+            requestConfig.headers['api-key'] = options.azureApiKey;
+        } else {
+            const options = this.options as OpenAIEmbeddingsOptions;
+            requestConfig.headers['Authorization'] = `Bearer ${options.apiKey}`;
+            if (options.organization) {
+                requestConfig.headers['OpenAI-Organization'] = options.organization;
+            }
+        }
+
+        // Send request
+        const response = await this._httpClient.post(url, body, requestConfig);
+
+        // Check for rate limit error
+        if (response.status == 429 && Array.isArray(this.options.retryPolicy) && retryCount < this.options.retryPolicy.length) {
+            const delay = this.options.retryPolicy[retryCount];
+            await new Promise((resolve) => setTimeout(resolve, delay));
+            return this.post(url, body, retryCount + 1);
+        } else {
+            return response;
+        }
+    }
+}