vectra 0.12.3 → 0.14.0

package/src/LocalEmbeddings.spec.ts

@@ -0,0 +1,138 @@
+import * as assert from 'assert';
+import * as sinon from 'sinon';
+import { LocalEmbeddings } from './LocalEmbeddings';
+
+describe('LocalEmbeddings', () => {
+    let requireStub: sinon.SinonStub;
+    let fakePipeline: sinon.SinonStub;
+    let fakePipelineFn: sinon.SinonStub;
+
+    beforeEach(() => {
+        // Create a fake pipeline function that returns tensor-like results
+        fakePipelineFn = sinon.stub();
+        fakePipeline = sinon.stub().resolves(fakePipelineFn);
+
+        // Stub require to intercept @huggingface/transformers
+        requireStub = sinon.stub(module.constructor.prototype, 'require');
+        requireStub.callThrough(); // Allow all other requires to pass through
+        requireStub.withArgs('@huggingface/transformers').returns({
+            pipeline: fakePipeline,
+        });
+    });
+
+    afterEach(() => {
+        sinon.restore();
+    });
+
+    describe('constructor', () => {
+        it('should use default model and maxTokens', () => {
+            const embeddings = new LocalEmbeddings();
+            assert.strictEqual(embeddings.model, 'Xenova/all-MiniLM-L6-v2');
+            assert.strictEqual(embeddings.maxTokens, 256);
+        });
+
+        it('should accept custom model name', () => {
+            const embeddings = new LocalEmbeddings({ model: 'custom/model' });
+            assert.strictEqual(embeddings.model, 'custom/model');
+        });
+
+        it('should accept custom maxTokens', () => {
+            const embeddings = new LocalEmbeddings({ maxTokens: 512 });
+            assert.strictEqual(embeddings.maxTokens, 512);
+        });
+    });
+
+    describe('createEmbeddings', () => {
+        it('should create embeddings for a single string input', async () => {
+            const fakeVector = new Float32Array([0.1, 0.2, 0.3]);
+            fakePipelineFn.resolves({ data: fakeVector });
+
+            const embeddings = new LocalEmbeddings();
+            const result = await embeddings.createEmbeddings('hello world');
+
+            assert.strictEqual(result.status, 'success');
+            assert.ok(result.output);
+            assert.strictEqual(result.output!.length, 1);
+            assert.deepStrictEqual(result.output![0], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896]);
+        });
+
+        it('should create embeddings for an array of inputs', async () => {
+            const vec1 = new Float32Array([0.1, 0.2, 0.3]);
+            const vec2 = new Float32Array([0.4, 0.5, 0.6]);
+            fakePipelineFn.onFirstCall().resolves({ data: vec1 });
+            fakePipelineFn.onSecondCall().resolves({ data: vec2 });
+
+            const embeddings = new LocalEmbeddings();
+            const result = await embeddings.createEmbeddings(['hello', 'world']);
+
+            assert.strictEqual(result.status, 'success');
+            assert.ok(result.output);
+            assert.strictEqual(result.output!.length, 2);
+        });
+
+        it('should call pipeline with mean pooling and normalize', async () => {
+            const fakeVector = new Float32Array([0.5, 0.5]);
+            fakePipelineFn.resolves({ data: fakeVector });
+
+            const embeddings = new LocalEmbeddings();
+            await embeddings.createEmbeddings('test');
+
+            assert.ok(fakePipelineFn.calledOnce);
+            const [input, options] = fakePipelineFn.firstCall.args;
+            assert.strictEqual(input, 'test');
+            assert.deepStrictEqual(options, { pooling: 'mean', normalize: true });
+        });
+
+        it('should initialize the pipeline with correct model name', async () => {
+            const fakeVector = new Float32Array([0.5]);
+            fakePipelineFn.resolves({ data: fakeVector });
+
+            const embeddings = new LocalEmbeddings({ model: 'my/model' });
+            await embeddings.createEmbeddings('test');
+
+            assert.ok(fakePipeline.calledOnce);
+            const [task, model] = fakePipeline.firstCall.args;
+            assert.strictEqual(task, 'feature-extraction');
+            assert.strictEqual(model, 'my/model');
+        });
+
+        it('should reuse the pipeline across calls', async () => {
+            const fakeVector = new Float32Array([0.5]);
+            fakePipelineFn.resolves({ data: fakeVector });
+
+            const embeddings = new LocalEmbeddings();
+            await embeddings.createEmbeddings('first');
+            await embeddings.createEmbeddings('second');
+
+            // Pipeline should only be created once
+            assert.ok(fakePipeline.calledOnce);
+            // But the pipeline function should be called twice
+            assert.strictEqual(fakePipelineFn.callCount, 2);
+        });
+
+        it('should return error status when pipeline fails', async () => {
+            fakePipelineFn.rejects(new Error('Model not found'));
+
+            const embeddings = new LocalEmbeddings();
+            const result = await embeddings.createEmbeddings('test');
+
+            assert.strictEqual(result.status, 'error');
+            assert.ok(result.message);
+            assert.ok(result.message!.includes('Model not found'));
+        });
+    });
+
+    describe('optional dependency', () => {
+        it('should return error when @huggingface/transformers is not installed', async () => {
+            requireStub.withArgs('@huggingface/transformers').throws(new Error('Cannot find module'));
+
+            const embeddings = new LocalEmbeddings();
+            const result = await embeddings.createEmbeddings('test');
+
+            assert.strictEqual(result.status, 'error');
+            assert.ok(result.message);
+            assert.ok(result.message!.includes('@huggingface/transformers'));
+            assert.ok(result.message!.includes('npm install'));
+        });
+    });
+});

package/src/LocalEmbeddings.ts

@@ -0,0 +1,120 @@
+import { EmbeddingsModel, EmbeddingsResponse } from "./types";
+
+/**
+ * Options for configuring a `LocalEmbeddings` instance.
+ */
+export interface LocalEmbeddingsOptions {
+    /**
+     * Optional. Model name to use for embeddings.
+     * @remarks
+     * Defaults to `Xenova/all-MiniLM-L6-v2`. Any model compatible with the
+     * `feature-extraction` pipeline from `@huggingface/transformers` can be used.
+     */
+    model?: string;
+
+    /**
+     * Optional. Maximum number of tokens the model supports.
+     * @remarks
+     * Defaults to 256 for the default model. Adjust when using a model with a
+     * different context window.
+     */
+    maxTokens?: number;
+}
+
+/**
+ * An `EmbeddingsModel` that runs locally using `@huggingface/transformers`.
+ * @remarks
+ * Requires the `@huggingface/transformers` package to be installed.
+ * The pipeline is lazily initialized on the first call to `createEmbeddings()`.
+ * Models are downloaded and cached locally on first use.
+ */
+export class LocalEmbeddings implements EmbeddingsModel {
+    public readonly maxTokens: number;
+
+    private readonly _modelName: string;
+    private _pipeline: any = null;
+    private _pipelinePromise: Promise<any> | null = null;
+
+    /**
+     * Creates a new `LocalEmbeddings` instance.
+     * @param options Optional configuration.
+     */
+    public constructor(options?: LocalEmbeddingsOptions) {
+        this._modelName = options?.model ?? 'Xenova/all-MiniLM-L6-v2';
+        this.maxTokens = options?.maxTokens ?? 256;
+    }
+
+    /**
+     * The model name used for embeddings.
+     */
+    public get model(): string {
+        return this._modelName;
+    }
+
+    /**
+     * 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.
+     */
+    public async createEmbeddings(inputs: string | string[]): Promise<EmbeddingsResponse> {
+        try {
+            const pipe = await this.getPipeline();
+            const inputArray = Array.isArray(inputs) ? inputs : [inputs];
+            const output: number[][] = [];
+
+            for (const input of inputArray) {
+                const result = await pipe(input, {
+                    pooling: 'mean',
+                    normalize: true,
+                });
+                output.push(Array.from(result.data as Float32Array));
+            }
+
+            return { status: 'success', output };
+        } catch (err: unknown) {
+            return {
+                status: 'error',
+                message: err instanceof Error ? err.message : String(err),
+            };
+        }
+    }
+
+    /**
+     * @private
+     * Lazily initializes and returns the transformer.js pipeline.
+     * Uses a singleton promise to prevent duplicate initialization.
+     */
+    private async getPipeline(): Promise<any> {
+        if (this._pipeline) {
+            return this._pipeline;
+        }
+
+        if (!this._pipelinePromise) {
+            this._pipelinePromise = this.initPipeline();
+        }
+
+        return this._pipelinePromise;
+    }
+
+    /**
+     * @private
+     */
+    private async initPipeline(): Promise<any> {
+        let transformers: any;
+        try {
+            transformers = require('@huggingface/transformers');
+        } catch {
+            throw new Error(
+                'The @huggingface/transformers package is required for local embeddings. ' +
+                'Install it with: npm install @huggingface/transformers'
+            );
+        }
+
+        this._pipeline = await transformers.pipeline(
+            'feature-extraction',
+            this._modelName,
+        );
+
+        return this._pipeline;
+    }
+}