vectra 0.1.1 → 0.2.0

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

package/src/TextSplitter.ts

@@ -0,0 +1,480 @@
+import { GPT3Tokenizer } from "./GPT3Tokenizer";
+import { TextChunk, Tokenizer } from "./types";
+
+export interface TextSplitterConfig {
+    separators: string[];
+    keepSeparators: boolean;
+    chunkSize: number;
+    chunkOverlap: number;
+    tokenizer: Tokenizer;
+    docType?: string;
+}
+
+export class TextSplitter {
+    private readonly _config: TextSplitterConfig;
+
+    public constructor(config?: Partial<TextSplitterConfig>) {
+        this._config = Object.assign({
+            separators: ["\n\n", "\n", " ", ""],
+            keepSeparators: false,
+            chunkSize: 400,
+            chunkOverlap: 40,
+        } as TextSplitterConfig, config);
+
+        // Create a default tokenizer if none is provided
+        if (!this._config.tokenizer) {
+            this._config.tokenizer = new GPT3Tokenizer();
+        }
+
+        // Use default separators if none are provided
+        if (!this._config.separators || this._config.separators.length === 0) {
+            this._config.separators = this.getSeparators(this._config.docType);
+        }
+
+        // Validate the config settings
+        if (this._config.chunkSize < 1) {
+            throw new Error("chunkSize must be >= 1");
+        } else if (this._config.chunkOverlap < 0) {
+            throw new Error("chunkOverlap must be >= 0");
+        } else if (this._config.chunkOverlap > this._config.chunkSize) {
+            throw new Error("chunkOverlap must be <= chunkSize");
+        }
+    }
+
+    public split(text: string): TextChunk[] {
+        // Get basic chunks
+        const chunks = this.recursiveSplit(text, this._config.separators, 0);
+
+        const that = this;
+        function getOverlapTokens(tokens?: number[]): number[] {
+            if (tokens != undefined) {
+                const len = tokens.length > that._config.chunkOverlap ? that._config.chunkOverlap : tokens.length;
+                return tokens.slice(tokens.length);
+            } else {
+                return [];
+            }
+        }
+
+        // Add overlap tokens and text to the start and end of each chunk
+        if (this._config.chunkOverlap > 0) {
+            for (let i = 1; i < chunks.length; i++) {
+                const previousChunk = chunks[i - 1];
+                const chunk = chunks[i];
+                const nextChunk = i < chunks.length - 1 ? chunks[i + 1] : undefined;
+                chunk.startOverlap = getOverlapTokens(previousChunk.tokens.reverse()).reverse();
+                chunk.endOverlap = getOverlapTokens(nextChunk?.tokens);
+            }
+        }
+
+        return chunks;
+    }
+
+    private recursiveSplit(text: string, separators: string[], startPos: number): TextChunk[] {
+        const chunks: TextChunk[] = [];
+        if (text.length > 0 && separators.length > 0) {
+            const separator = separators[0];
+            const nextSeparators = separators.length > 1 ? separators.slice(1) : [];
+            const parts = text.split(separator);
+            for (let i = 0; i < parts.length; i++) {
+                const lastChunk = (i === parts.length - 1);
+
+                // Get chunk text and endPos
+                let chunk = parts[i];
+                const endPos = (startPos + (chunk.length - 1)) + (lastChunk ? 0 : separator.length);
+                if (this._config.keepSeparators && !lastChunk) {
+                    chunk += separator;
+                }
+
+                // Encode chunk text
+                const tokens = this._config.tokenizer.encode(chunk);
+                if (tokens.length > this._config.chunkSize) {
+                    // Break the text into smaller chunks
+                    const subChunks = this.recursiveSplit(chunk, nextSeparators, startPos);
+                    chunks.push(...subChunks);
+                } else {
+                    // Append chunk to output
+                    chunks.push({
+                        text: chunk,
+                        tokens: tokens,
+                        startPos: startPos,
+                        endPos: endPos,
+                        startOverlap: [],
+                        endOverlap: [],
+                    });
+                }
+
+                // Update startPos
+                startPos = endPos + 1;
+            }
+        }
+
+        return chunks;
+    }
+
+    private getSeparators(docType?: string): string[] {
+        switch (docType ?? '') {
+            case "cpp":
+                return [
+                    // Split along class definitions
+                    "\nclass ",
+                    // Split along function definitions
+                    "\nvoid ",
+                    "\nint ",
+                    "\nfloat ",
+                    "\ndouble ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nfor ",
+                    "\nwhile ",
+                    "\nswitch ",
+                    "\ncase ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "go":
+                return [
+                    // Split along function definitions
+                    "\nfunc ",
+                    "\nvar ",
+                    "\nconst ",
+                    "\ntype ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nfor ",
+                    "\nswitch ",
+                    "\ncase ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "java":
+            case "c#":
+            case "csharp":
+            case "cs":
+            case "ts":
+            case "tsx":
+            case "typescript":
+                return [
+                    // Split along class definitions
+                    "\nclass ",
+                    // Split along method definitions
+                    "\npublic ",
+                    "\nprotected ",
+                    "\nprivate ",
+                    "\nstatic ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nfor ",
+                    "\nwhile ",
+                    "\nswitch ",
+                    "\ncase ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "js":
+            case "jsx":
+            case "javascript":
+                return [
+                    // Split along class definitions
+                    "\nclass ",
+                    // Split along function definitions
+                    "\nfunction ",
+                    "\nconst ",
+                    "\nlet ",
+                    "\nvar ",
+                    "\nclass ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nfor ",
+                    "\nwhile ",
+                    "\nswitch ",
+                    "\ncase ",
+                    "\ndefault ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "php":
+                return [
+                    // Split along function definitions
+                    "\nfunction ",
+                    // Split along class definitions
+                    "\nclass ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nforeach ",
+                    "\nwhile ",
+                    "\ndo ",
+                    "\nswitch ",
+                    "\ncase ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "proto":
+                return [
+                    // Split along message definitions
+                    "\nmessage ",
+                    // Split along service definitions
+                    "\nservice ",
+                    // Split along enum definitions
+                    "\nenum ",
+                    // Split along option definitions
+                    "\noption ",
+                    // Split along import statements
+                    "\nimport ",
+                    // Split along syntax declarations
+                    "\nsyntax ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "python":
+            case "py":
+                return [
+                    // First, try to split along class definitions
+                    "\nclass ",
+                    "\ndef ",
+                    "\n\tdef ",
+                    // Now split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "rst":
+                return [
+                    // Split along section titles
+                    "\n===\n",
+                    "\n---\n",
+                    "\n***\n",
+                    // Split along directive markers
+                    "\n.. ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "ruby":
+                return [
+                    // Split along method definitions
+                    "\ndef ",
+                    "\nclass ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nunless ",
+                    "\nwhile ",
+                    "\nfor ",
+                    "\ndo ",
+                    "\nbegin ",
+                    "\nrescue ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "rust":
+                return [
+                    // Split along function definitions
+                    "\nfn ",
+                    "\nconst ",
+                    "\nlet ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nwhile ",
+                    "\nfor ",
+                    "\nloop ",
+                    "\nmatch ",
+                    "\nconst ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "scala":
+                return [
+                    // Split along class definitions
+                    "\nclass ",
+                    "\nobject ",
+                    // Split along method definitions
+                    "\ndef ",
+                    "\nval ",
+                    "\nvar ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nfor ",
+                    "\nwhile ",
+                    "\nmatch ",
+                    "\ncase ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "swift":
+                return [
+                    // Split along function definitions
+                    "\nfunc ",
+                    // Split along class definitions
+                    "\nclass ",
+                    "\nstruct ",
+                    "\nenum ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nfor ",
+                    "\nwhile ",
+                    "\ndo ",
+                    "\nswitch ",
+                    "\ncase ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "markdown":
+                return [
+                    // First, try to split along Markdown headings (starting with level 2)
+                    "\n## ",
+                    "\n### ",
+                    "\n#### ",
+                    "\n##### ",
+                    "\n###### ",
+                    // Note the alternative syntax for headings (below) is not handled here
+                    // Heading level 2
+                    // ---------------
+                    // End of code block
+                    "```\n\n",
+                    // Horizontal lines
+                    "\n\n***\n\n",
+                    "\n\n---\n\n",
+                    "\n\n___\n\n",
+                    // Note that this splitter doesn't handle horizontal lines defined
+                    // by *three or more* of ***, ---, or ___, but this is not handled
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "latex":
+                return [
+                    // First, try to split along Latex sections
+                    "\n\\chapter{",
+                    "\n\\section{",
+                    "\n\\subsection{",
+                    "\n\\subsubsection{",
+
+                    // Now split by environments
+                    "\n\\begin{enumerate}",
+                    "\n\\begin{itemize}",
+                    "\n\\begin{description}",
+                    "\n\\begin{list}",
+                    "\n\\begin{quote}",
+                    "\n\\begin{quotation}",
+                    "\n\\begin{verse}",
+                    "\n\\begin{verbatim}",
+
+                    // Now split by math environments
+                    "\n\\begin{align}",
+                    "$$",
+                    "$",
+
+                    // Now split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            case "html":
+                return [
+                    // First, try to split along HTML tags
+                    "<body>",
+                    "<div>",
+                    "<p>",
+                    "<br>",
+                    "<li>",
+                    "<h1>",
+                    "<h2>",
+                    "<h3>",
+                    "<h4>",
+                    "<h5>",
+                    "<h6>",
+                    "<span>",
+                    "<table>",
+                    "<tr>",
+                    "<td>",
+                    "<th>",
+                    "<ul>",
+                    "<ol>",
+                    "<header>",
+                    "<footer>",
+                    "<nav>",
+                    // Head
+                    "<head>",
+                    "<style>",
+                    "<script>",
+                    "<meta>",
+                    "<title>",
+                    // Normal type of lines
+                    " ",
+                    "",
+                ];
+            case "sol":
+                return [
+                    // Split along compiler informations definitions
+                    "\npragma ",
+                    "\nusing ",
+                    // Split along contract definitions
+                    "\ncontract ",
+                    "\ninterface ",
+                    "\nlibrary ",
+                    // Split along method definitions
+                    "\nconstructor ",
+                    "\ntype ",
+                    "\nfunction ",
+                    "\nevent ",
+                    "\nmodifier ",
+                    "\nerror ",
+                    "\nstruct ",
+                    "\nenum ",
+                    // Split along control flow statements
+                    "\nif ",
+                    "\nfor ",
+                    "\nwhile ",
+                    "\ndo while ",
+                    "\nassembly ",
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+            default:
+                return [
+                    // Split by the normal type of lines
+                    "\n\n",
+                    "\n",
+                    " ",
+                    "",
+                ];
+        }
+    }
+}