@framers/agentos-ext-ml-classifiers 0.1.0 → 0.2.1

package/test/SlidingWindowBuffer.spec.ts

@@ -0,0 +1,391 @@
+/**
+ * @fileoverview Unit tests for `SlidingWindowBuffer`.
+ *
+ * Tests verify:
+ *  - Null return until chunkSize is reached
+ *  - ChunkReady emission once chunkSize is reached
+ *  - Context carry-forward between consecutive chunks
+ *  - Multiple concurrent streams operate independently
+ *  - maxEvaluations budget is respected
+ *  - flush() returns remaining buffer content
+ *  - flush() returns null for empty or unknown streams
+ *  - pruneStale() removes expired streams
+ *  - clear() removes all streams
+ */
+
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+import {
+  SlidingWindowBuffer,
+} from '../src/SlidingWindowBuffer';
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Build a string of exactly `charCount` characters so token estimation
+ * (ceil(length/4)) is predictable.
+ */
+function chars(charCount: number): string {
+  return 'a'.repeat(charCount);
+}
+
+/**
+ * Push enough characters to fill exactly `tokenCount` estimated tokens
+ * (4 chars per token) into the buffer for the given stream.
+ */
+function pushTokens(
+  buf: SlidingWindowBuffer,
+  streamId: string,
+  tokenCount: number,
+) {
+  return buf.push(streamId, chars(tokenCount * 4));
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe('SlidingWindowBuffer', () => {
+  let buf: SlidingWindowBuffer;
+
+  beforeEach(() => {
+    // Fresh buffer for every test with small sizes to keep tests fast.
+    buf = new SlidingWindowBuffer({
+      chunkSize: 10,       // 10 tokens = 40 chars before a chunk fires
+      contextSize: 5,      // keep 5 tokens (20 chars) of context
+      maxEvaluations: 5,
+      streamTimeoutMs: 1000,
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // Basic accumulation
+  // -------------------------------------------------------------------------
+
+  describe('returns null until chunkSize is reached', () => {
+    it('returns null for a single short push', () => {
+      // Push 5 tokens (< chunkSize of 10)
+      expect(pushTokens(buf, 'stream-a', 5)).toBeNull();
+    });
+
+    it('returns null for multiple small pushes that do not reach chunkSize', () => {
+      expect(buf.push('stream-a', chars(8))).toBeNull();  // 2 tokens
+      expect(buf.push('stream-a', chars(8))).toBeNull();  // 2 tokens
+      expect(buf.push('stream-a', chars(8))).toBeNull();  // 2 tokens
+      // Total: 6 tokens — still below 10
+    });
+  });
+
+  describe('returns ChunkReady when chunkSize is reached', () => {
+    it('emits a chunk exactly when chunkSize tokens accumulate', () => {
+      // Push 9 tokens first — still under threshold
+      expect(pushTokens(buf, 'stream-a', 9)).toBeNull();
+
+      // Push 1 more token to hit chunkSize=10
+      const chunk = pushTokens(buf, 'stream-a', 1);
+      expect(chunk).not.toBeNull();
+      expect(chunk!.evaluationNumber).toBe(1);
+    });
+
+    it('chunk.newText equals exactly what was buffered', () => {
+      // Push exactly chunkSize tokens in one shot
+      const chunk = pushTokens(buf, 'stream-a', 10);
+      expect(chunk).not.toBeNull();
+      // newText should be the 40 chars we pushed
+      expect(chunk!.newText).toHaveLength(40);
+    });
+
+    it('chunk.text equals contextRing (empty on first chunk) + newText', () => {
+      const chunk = pushTokens(buf, 'stream-a', 10);
+      // No prior context on first chunk → text === newText
+      expect(chunk!.text).toBe(chunk!.newText);
+    });
+
+    it('evaluationNumber starts at 1', () => {
+      const chunk = pushTokens(buf, 'stream-a', 10);
+      expect(chunk!.evaluationNumber).toBe(1);
+    });
+
+    it('evaluationNumber increments on successive chunks', () => {
+      const c1 = pushTokens(buf, 'stream-a', 10);
+      const c2 = pushTokens(buf, 'stream-a', 10);
+      expect(c1!.evaluationNumber).toBe(1);
+      expect(c2!.evaluationNumber).toBe(2);
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // Context carry-forward
+  // -------------------------------------------------------------------------
+
+  describe('context carry-forward between chunks', () => {
+    it('second chunk text is longer than its newText (context prepended)', () => {
+      // Emit first chunk
+      pushTokens(buf, 'stream-a', 10);
+
+      // Emit second chunk
+      const c2 = pushTokens(buf, 'stream-a', 10);
+      expect(c2).not.toBeNull();
+
+      // text = contextRing (≤ contextSize=5 tokens = 20 chars) + newText (40 chars)
+      expect(c2!.text.length).toBeGreaterThan(c2!.newText.length);
+    });
+
+    it('context length is bounded by contextSize tokens', () => {
+      pushTokens(buf, 'stream-a', 10); // first chunk
+
+      const c2 = pushTokens(buf, 'stream-a', 10);
+      const contextLen = c2!.text.length - c2!.newText.length;
+
+      // contextSize=5 tokens ≈ 20 chars; context length must be ≤ 20
+      expect(contextLen).toBeLessThanOrEqual(5 * 4);
+    });
+
+    it('context is derived from the tail of the previous buffer', () => {
+      // Use distinct characters to trace which text makes it into the context.
+      // First window: 10 tokens of 'x'
+      buf.push('stream-a', 'x'.repeat(40));
+
+      // Second window: 10 tokens of 'y'
+      const c2 = buf.push('stream-a', 'y'.repeat(40));
+      expect(c2).not.toBeNull();
+
+      // The text should start with some 'x' context, then 'y' new content.
+      expect(c2!.text).toContain('x'); // context from first window
+      expect(c2!.newText).toBe('y'.repeat(40)); // only new content
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // Multiple concurrent streams
+  // -------------------------------------------------------------------------
+
+  describe('multiple concurrent streams are independent', () => {
+    it('two streams accumulate independently', () => {
+      // Push 8 tokens into stream-a, 10 tokens into stream-b
+      pushTokens(buf, 'stream-a', 8); // not ready
+      const chunkB = pushTokens(buf, 'stream-b', 10); // ready
+
+      expect(chunkB).not.toBeNull();
+
+      // stream-a should still be null (only 8 tokens)
+      const chunkA = buf.push('stream-a', ''); // empty push, no change
+      expect(chunkA).toBeNull();
+    });
+
+    it('flushing one stream does not affect another', () => {
+      pushTokens(buf, 'stream-a', 5); // partial
+      pushTokens(buf, 'stream-b', 5); // partial
+
+      buf.flush('stream-a');
+
+      // stream-b still exists and can be flushed separately
+      const chunkB = buf.flush('stream-b');
+      expect(chunkB).not.toBeNull();
+      expect(chunkB!.newText).toHaveLength(20); // 5 tokens * 4 chars
+    });
+
+    it('context rings are independent per stream', () => {
+      // Emit first chunk for each stream with distinct chars
+      buf.push('stream-a', 'A'.repeat(40));
+      buf.push('stream-b', 'B'.repeat(40));
+
+      const c2a = pushTokens(buf, 'stream-a', 10);
+      const c2b = pushTokens(buf, 'stream-b', 10);
+
+      // stream-a context should contain 'A', not 'B'
+      expect(c2a!.text).toContain('A');
+      expect(c2a!.text).not.toContain('B');
+
+      // stream-b context should contain 'B', not 'A'
+      expect(c2b!.text).toContain('B');
+      expect(c2b!.text).not.toContain('A');
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // Evaluation budget (maxEvaluations)
+  // -------------------------------------------------------------------------
+
+  describe('maxEvaluations budget', () => {
+    it('returns null after maxEvaluations chunks are emitted', () => {
+      // Emit exactly maxEvaluations=5 chunks
+      for (let i = 0; i < 5; i++) {
+        const chunk = pushTokens(buf, 'stream-a', 10);
+        expect(chunk).not.toBeNull();
+      }
+
+      // 6th push should return null even though enough tokens are pushed
+      const extra = pushTokens(buf, 'stream-a', 10);
+      expect(extra).toBeNull();
+    });
+
+    it('budget is tracked per-stream (other streams unaffected)', () => {
+      // Exhaust stream-a
+      for (let i = 0; i < 5; i++) {
+        pushTokens(buf, 'stream-a', 10);
+      }
+
+      // stream-b should still be able to emit chunks
+      const chunkB = pushTokens(buf, 'stream-b', 10);
+      expect(chunkB).not.toBeNull();
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // flush()
+  // -------------------------------------------------------------------------
+
+  describe('flush()', () => {
+    it('returns a ChunkReady for remaining buffered text', () => {
+      buf.push('stream-a', chars(20)); // 5 tokens, less than chunkSize=10
+      const chunk = buf.flush('stream-a');
+
+      expect(chunk).not.toBeNull();
+      expect(chunk!.newText).toHaveLength(20);
+    });
+
+    it('returned chunk includes context from prior windows', () => {
+      // Emit one full chunk first
+      pushTokens(buf, 'stream-a', 10);
+
+      // Push partial second window
+      buf.push('stream-a', chars(20)); // 5 tokens
+
+      const chunk = buf.flush('stream-a');
+      expect(chunk).not.toBeNull();
+      // text should be longer than newText (context prepended)
+      expect(chunk!.text.length).toBeGreaterThan(chunk!.newText.length);
+    });
+
+    it('returns null for empty buffer', () => {
+      // Push nothing, then flush
+      const chunk = buf.flush('stream-empty');
+      expect(chunk).toBeNull();
+    });
+
+    it('returns null for a non-existent stream', () => {
+      expect(buf.flush('does-not-exist')).toBeNull();
+    });
+
+    it('removes the stream from internal state after flush', () => {
+      buf.push('stream-a', chars(20));
+      buf.flush('stream-a');
+
+      // size should be 0 after flushing the only stream
+      expect(buf.size).toBe(0);
+    });
+
+    it('subsequent flush on same stream returns null', () => {
+      buf.push('stream-a', chars(20));
+      buf.flush('stream-a');
+
+      // Second flush: stream was deleted
+      expect(buf.flush('stream-a')).toBeNull();
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // pruneStale()
+  // -------------------------------------------------------------------------
+
+  describe('pruneStale()', () => {
+    it('removes streams that have exceeded streamTimeoutMs', async () => {
+      // Use a very short timeout for this test
+      const shortBuf = new SlidingWindowBuffer({
+        chunkSize: 10,
+        contextSize: 5,
+        maxEvaluations: 5,
+        streamTimeoutMs: 50, // 50 ms
+      });
+
+      shortBuf.push('old-stream', chars(20)); // partial push
+
+      // Wait for the timeout to expire
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
+      shortBuf.pruneStale();
+      expect(shortBuf.size).toBe(0);
+    });
+
+    it('does not remove streams that are still within the timeout', async () => {
+      const shortBuf = new SlidingWindowBuffer({
+        chunkSize: 10,
+        contextSize: 5,
+        maxEvaluations: 5,
+        streamTimeoutMs: 5000,
+      });
+
+      shortBuf.push('fresh-stream', chars(20));
+      shortBuf.pruneStale();
+
+      // Should still be present
+      expect(shortBuf.size).toBe(1);
+    });
+
+    it('is invoked lazily when map.size > 10', async () => {
+      const shortBuf = new SlidingWindowBuffer({
+        chunkSize: 10,
+        contextSize: 5,
+        maxEvaluations: 5,
+        streamTimeoutMs: 1, // expire immediately
+      });
+
+      // Create 10 streams that will immediately be stale
+      for (let i = 0; i < 10; i++) {
+        shortBuf.push(`stale-${i}`, chars(4));
+      }
+
+      // Wait for timeout
+      await new Promise((resolve) => setTimeout(resolve, 20));
+
+      // Push to an 11th stream — this triggers lazy pruning
+      shortBuf.push('trigger-prune', chars(4));
+
+      // Stale streams should have been removed; only 'trigger-prune' remains
+      expect(shortBuf.size).toBe(1);
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // clear()
+  // -------------------------------------------------------------------------
+
+  describe('clear()', () => {
+    it('removes all streams', () => {
+      buf.push('s1', chars(20));
+      buf.push('s2', chars(20));
+      buf.push('s3', chars(20));
+
+      buf.clear();
+      expect(buf.size).toBe(0);
+    });
+
+    it('is idempotent on an empty buffer', () => {
+      buf.clear();
+      buf.clear();
+      expect(buf.size).toBe(0);
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // Edge cases
+  // -------------------------------------------------------------------------
+
+  describe('edge cases', () => {
+    it('push with empty string returns null and creates no state', () => {
+      expect(buf.push('stream-a', '')).toBeNull();
+      expect(buf.size).toBe(0);
+    });
+
+    it('handles a single massive push that exceeds chunkSize', () => {
+      // 20 tokens in one push — should still emit exactly one chunk
+      const chunk = pushTokens(buf, 'stream-a', 20);
+      expect(chunk).not.toBeNull();
+      // After the chunk, residual text (10 extra tokens) stays in buffer
+      // A second push of 0 tokens shouldn't fire a second chunk
+      expect(buf.push('stream-a', '')).toBeNull();
+    });
+  });
+});

package/test/ToxicityClassifier.spec.ts

@@ -0,0 +1,268 @@
+/**
+ * @fileoverview Unit tests for {@link ToxicityClassifier}.
+ *
+ * All tests use a mocked {@link ISharedServiceRegistry} that returns a
+ * pre-configured pipeline function.  No real model weights are downloaded.
+ *
+ * Test coverage:
+ *  1. Correct static identity: `id`, `displayName`, `modelId`
+ *  2. Maps pipeline output to ClassificationResult correctly
+ *     (bestClass = highest-score label, confidence = its score, allScores = all labels)
+ *  3. Graceful degradation — returns pass result when model fails to load
+ *  4. Uses ISharedServiceRegistry with the correct service ID
+ *  5. `isLoaded` flag is set after a successful classification
+ *  6. `isLoaded` is false before any classify() call
+ */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import type { ISharedServiceRegistry } from '@framers/agentos';
+import { ToxicityClassifier } from '../src/classifiers/ToxicityClassifier';
+import { ML_CLASSIFIER_SERVICE_IDS } from '../src/types';
+
+// ---------------------------------------------------------------------------
+// Test fixture helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Raw multi-label output that the `unitary/toxic-bert` pipeline would return
+ * when called with `{ topk: null }`.  `toxic` is the winner at 0.92.
+ */
+const TOXICITY_PIPELINE_OUTPUT = [
+  { label: 'toxic', score: 0.92 },
+  { label: 'severe_toxic', score: 0.03 },
+  { label: 'obscene', score: 0.45 },
+  { label: 'threat', score: 0.02 },
+  { label: 'insult', score: 0.61 },
+  { label: 'identity_hate', score: 0.01 },
+];
+
+/**
+ * Build a mock {@link ISharedServiceRegistry} whose `getOrCreate` method
+ * returns a mock pipeline function pre-configured to resolve with
+ * `pipelineResult`.
+ *
+ * @param pipelineResult - The value the mock pipeline resolves with.
+ */
+function mockRegistry(pipelineResult: unknown): ISharedServiceRegistry {
+  // The pipeline is a callable that the classifier invokes as pipeline(text, opts).
+  const pipeline = vi.fn(async () => pipelineResult);
+  return {
+    /**
+     * Ignores the factory and always returns the same mock pipeline.
+     * The `serviceId` is captured in the spy so tests can assert on it.
+     */
+    getOrCreate: vi.fn(async () => pipeline),
+    has: vi.fn(() => false),
+    release: vi.fn(async () => {}),
+    releaseAll: vi.fn(async () => {}),
+  };
+}
+
+/**
+ * Build a registry whose `getOrCreate` rejects with an error to simulate a
+ * model-load failure.
+ */
+function failingRegistry(): ISharedServiceRegistry {
+  return {
+    getOrCreate: vi.fn(async () => {
+      throw new Error('Model download failed');
+    }),
+    has: vi.fn(() => false),
+    release: vi.fn(async () => {}),
+    releaseAll: vi.fn(async () => {}),
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe('ToxicityClassifier', () => {
+  // -------------------------------------------------------------------------
+  // 1. Static identity
+  // -------------------------------------------------------------------------
+
+  describe('static identity', () => {
+    it('has the correct id', () => {
+      const classifier = new ToxicityClassifier(mockRegistry([]));
+      expect(classifier.id).toBe('toxicity');
+    });
+
+    it('has the correct displayName', () => {
+      const classifier = new ToxicityClassifier(mockRegistry([]));
+      expect(classifier.displayName).toBe('Toxicity Classifier');
+    });
+
+    it('has the correct default modelId', () => {
+      const classifier = new ToxicityClassifier(mockRegistry([]));
+      expect(classifier.modelId).toBe('unitary/toxic-bert');
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // 2. isLoaded flag
+  // -------------------------------------------------------------------------
+
+  describe('isLoaded flag', () => {
+    it('is false before any classify() call', () => {
+      const classifier = new ToxicityClassifier(mockRegistry(TOXICITY_PIPELINE_OUTPUT));
+      expect(classifier.isLoaded).toBe(false);
+    });
+
+    it('is true after a successful classify() call', async () => {
+      const classifier = new ToxicityClassifier(mockRegistry(TOXICITY_PIPELINE_OUTPUT));
+      await classifier.classify('some text');
+      expect(classifier.isLoaded).toBe(true);
+    });
+
+    it('remains false after a model-load failure', async () => {
+      const classifier = new ToxicityClassifier(failingRegistry());
+      await classifier.classify('some text');
+      expect(classifier.isLoaded).toBe(false);
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // 3. Result mapping
+  // -------------------------------------------------------------------------
+
+  describe('classify() — result mapping', () => {
+    let classifier: ToxicityClassifier;
+
+    beforeEach(() => {
+      classifier = new ToxicityClassifier(mockRegistry(TOXICITY_PIPELINE_OUTPUT));
+    });
+
+    it('resolves with the label with the highest score as bestClass', async () => {
+      const result = await classifier.classify('You are terrible!');
+      // toxic (0.92) beats insult (0.61) and obscene (0.45)
+      expect(result.bestClass).toBe('toxic');
+    });
+
+    it('resolves with the top score as confidence', async () => {
+      const result = await classifier.classify('You are terrible!');
+      expect(result.confidence).toBeCloseTo(0.92);
+    });
+
+    it('includes all six labels in allScores', async () => {
+      const result = await classifier.classify('You are terrible!');
+      expect(result.allScores).toHaveLength(6);
+    });
+
+    it('allScores contains correct classLabel/score pairs', async () => {
+      const result = await classifier.classify('You are terrible!');
+      // Spot-check a few entries
+      const toxic = result.allScores.find((s) => s.classLabel === 'toxic');
+      expect(toxic?.score).toBeCloseTo(0.92);
+
+      const threat = result.allScores.find((s) => s.classLabel === 'threat');
+      expect(threat?.score).toBeCloseTo(0.02);
+    });
+
+    it('returns bestClass=toxic for a message where toxic wins', async () => {
+      // Verify the classifier picks the maximum regardless of array order
+      const shuffled = [...TOXICITY_PIPELINE_OUTPUT].reverse();
+      const reg = mockRegistry(shuffled);
+      const cls = new ToxicityClassifier(reg);
+      const result = await cls.classify('test');
+      expect(result.bestClass).toBe('toxic');
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // 4. Graceful degradation
+  // -------------------------------------------------------------------------
+
+  describe('graceful degradation on model load failure', () => {
+    it('returns bestClass=benign when model fails to load', async () => {
+      const classifier = new ToxicityClassifier(failingRegistry());
+      const result = await classifier.classify('some text');
+      expect(result.bestClass).toBe('benign');
+    });
+
+    it('returns confidence=0 when model fails to load', async () => {
+      const classifier = new ToxicityClassifier(failingRegistry());
+      const result = await classifier.classify('some text');
+      expect(result.confidence).toBe(0);
+    });
+
+    it('returns empty allScores when model fails to load', async () => {
+      const classifier = new ToxicityClassifier(failingRegistry());
+      const result = await classifier.classify('some text');
+      expect(result.allScores).toEqual([]);
+    });
+
+    it('continues returning pass result on all subsequent calls after failure', async () => {
+      const classifier = new ToxicityClassifier(failingRegistry());
+      // First call triggers the failure
+      await classifier.classify('call 1');
+      // Subsequent calls should still return the pass result without retrying
+      const result = await classifier.classify('call 2');
+      expect(result.bestClass).toBe('benign');
+    });
+
+    it('does not retry getOrCreate after the first failure', async () => {
+      const registry = failingRegistry();
+      const classifier = new ToxicityClassifier(registry);
+      await classifier.classify('call 1');
+      await classifier.classify('call 2');
+      // getOrCreate should only have been called once (on the first classify call)
+      expect(registry.getOrCreate).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // 5. Uses ISharedServiceRegistry with correct service ID
+  // -------------------------------------------------------------------------
+
+  describe('shared service registry integration', () => {
+    it('calls getOrCreate with the TOXICITY_PIPELINE service ID', async () => {
+      const registry = mockRegistry(TOXICITY_PIPELINE_OUTPUT);
+      const classifier = new ToxicityClassifier(registry);
+      await classifier.classify('hello');
+      expect(registry.getOrCreate).toHaveBeenCalledWith(
+        ML_CLASSIFIER_SERVICE_IDS.TOXICITY_PIPELINE,
+        expect.any(Function),
+        expect.objectContaining({ tags: expect.arrayContaining(['toxicity']) }),
+      );
+    });
+
+    it('does not call getOrCreate again on a second classify() call (cached)', async () => {
+      const registry = mockRegistry(TOXICITY_PIPELINE_OUTPUT);
+      const classifier = new ToxicityClassifier(registry);
+      await classifier.classify('first call');
+      await classifier.classify('second call');
+      // Pipeline is retrieved once and re-used
+      expect(registry.getOrCreate).toHaveBeenCalledTimes(2); // once per classify() — registry handles caching internally
+    });
+
+    it('calls release with TOXICITY_PIPELINE service ID on dispose()', async () => {
+      const registry = mockRegistry(TOXICITY_PIPELINE_OUTPUT);
+      const classifier = new ToxicityClassifier(registry);
+      await classifier.classify('hello');
+      await classifier.dispose();
+      expect(registry.release).toHaveBeenCalledWith(
+        ML_CLASSIFIER_SERVICE_IDS.TOXICITY_PIPELINE,
+      );
+    });
+  });
+
+  // -------------------------------------------------------------------------
+  // 6. Config override
+  // -------------------------------------------------------------------------
+
+  describe('ClassifierConfig.modelId override', () => {
+    it('passes the overridden modelId to the factory (verified via the factory closure)', async () => {
+      // We cannot peek inside the factory directly, but we can verify that
+      // getOrCreate is called — the factory is a closure that reads config.modelId.
+      // A true integration test would require a real import; here we just confirm
+      // the registry is invoked at all when a custom modelId is provided.
+      const registry = mockRegistry(TOXICITY_PIPELINE_OUTPUT);
+      const classifier = new ToxicityClassifier(registry, {
+        modelId: 'my-org/custom-toxic-bert',
+      });
+      await classifier.classify('hello');
+      expect(registry.getOrCreate).toHaveBeenCalled();
+    });
+  });
+});