@framers/agentos-ext-diarization 0.2.0

package/dist/types.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"types.js","sourceRoot":"","sources":["../src/types.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;GAUG"}

package/manifest.json

@@ -0,0 +1,8 @@
+{
+  "name": "@framers/agentos-ext-diarization",
+  "version": "0.1.0",
+  "description": "Speaker diarization with provider delegation and local spectral-centroid clustering",
+  "kind": "diarization-provider",
+  "extensionId": "diarization",
+  "entryPoint": "./dist/index.js"
+}

package/package.json

@@ -0,0 +1,48 @@
+{
+  "name": "@framers/agentos-ext-diarization",
+  "version": "0.2.0",
+  "description": "Speaker diarization extension pack for AgentOS voice pipeline — provider delegation and local clustering",
+  "type": "module",
+  "main": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": {
+    ".": {
+      "import": "./dist/index.js",
+      "types": "./dist/index.d.ts"
+    }
+  },
+  "files": [
+    "dist",
+    "src",
+    "SKILL.md",
+    "manifest.json"
+  ],
+  "peerDependencies": {
+    "@framers/agentos": "^0.1.0",
+    "onnxruntime-node": "^1.18.0"
+  },
+  "peerDependenciesMeta": {
+    "onnxruntime-node": {
+      "optional": true
+    }
+  },
+  "devDependencies": {
+    "typescript": "^5.5.0",
+    "vitest": "^1.6.0",
+    "@framers/agentos": "0.1.94"
+  },
+  "license": "MIT",
+  "author": "Frame.dev",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/framersai/agentos-extensions.git",
+    "directory": "registry/curated/voice/diarization"
+  },
+  "publishConfig": {
+    "access": "public"
+  },
+  "scripts": {
+    "build": "tsc -p tsconfig.json",
+    "test": "vitest run"
+  }
+}

package/src/ClusteringStrategy.ts

@@ -0,0 +1,139 @@
+/**
+ * @file ClusteringStrategy.ts
+ * @description Agglomerative clustering to merge drifted speaker centroids.
+ *
+ * Over the course of a long session, a single speaker's vocal characteristics
+ * may drift enough that two separate centroids are created for them.
+ * {@link ClusteringStrategy.mergeClusters} detects this by computing pairwise
+ * cosine similarity between all centroids and iteratively merging the closest
+ * pair until either no pair exceeds the merge threshold or the centroid count
+ * equals `expectedSpeakers`.
+ *
+ * @module diarization/ClusteringStrategy
+ */
+
+import { cosineSimilarity } from './SpeakerEmbeddingCache.js';
+
+// ---------------------------------------------------------------------------
+// Main class
+// ---------------------------------------------------------------------------
+
+/**
+ * Agglomerative speaker-centroid merging strategy.
+ *
+ * This is an optional post-processing step applied by {@link DiarizationSession}
+ * when the number of tracked centroids exceeds the expected speaker count.
+ *
+ * @example
+ * ```ts
+ * const strategy = new ClusteringStrategy(0.85);
+ * const mapping = strategy.mergeClusters(cache.centroids, 2);
+ * // mapping: Map<oldId, canonicalId>
+ * ```
+ */
+export class ClusteringStrategy {
+  // -------------------------------------------------------------------------
+  // Constructor
+  // -------------------------------------------------------------------------
+
+  /**
+   * @param mergeThreshold - Minimum cosine similarity between two centroids
+   *   for them to be considered the same speaker and merged.
+   *   @defaultValue 0.85
+   */
+  constructor(private readonly mergeThreshold: number = 0.85) {}
+
+  // -------------------------------------------------------------------------
+  // Public API
+  // -------------------------------------------------------------------------
+
+  /**
+   * Identify centroid pairs that should be merged and return a renaming map.
+   *
+   * The algorithm:
+   * 1. Compute all pairwise cosine similarities.
+   * 2. If `expectedSpeakers` is set and the current count exceeds it, merge
+   *    the closest pair regardless of the threshold.
+   * 3. Otherwise merge pairs that exceed `mergeThreshold`.
+   * 4. Repeat until no further merges are possible or the count matches
+   *    `expectedSpeakers`.
+   *
+   * The returned `Map<string, string>` maps every old centroid ID that was
+   * subsumed into a canonical ID.  IDs that were not merged are not present in
+   * the map.  Callers should rename all occurrences of a key to its value.
+   *
+   * @param centroids - Current centroid snapshot (id → embedding).
+   * @param expectedSpeakers - Optional upper bound on speaker count.
+   * @returns Rename map: `oldId → canonicalId`.
+   */
+  mergeClusters(
+    centroids: Map<string, Float32Array>,
+    expectedSpeakers?: number,
+  ): Map<string, string> {
+    // Build a mutable working copy so we can iteratively merge.
+    const working = new Map<string, Float32Array>(centroids);
+    // Accumulated rename mapping.
+    const renameMap = new Map<string, string>();
+
+    while (true) {
+      const ids = Array.from(working.keys());
+      const count = ids.length;
+
+      // Nothing to merge.
+      if (count < 2) break;
+
+      // Find the closest pair.
+      let bestSim = -Infinity;
+      let bestI = 0;
+      let bestJ = 1;
+
+      for (let i = 0; i < count; i++) {
+        for (let j = i + 1; j < count; j++) {
+          const sim = cosineSimilarity(working.get(ids[i]!)!, working.get(ids[j]!)!);
+          if (sim > bestSim) {
+            bestSim = sim;
+            bestI = i;
+            bestJ = j;
+          }
+        }
+      }
+
+      // Decide whether to merge.
+      const shouldMergeDueToThreshold = bestSim >= this.mergeThreshold;
+      const shouldMergeDueToCount =
+        expectedSpeakers !== undefined && count > expectedSpeakers;
+
+      if (!shouldMergeDueToThreshold && !shouldMergeDueToCount) break;
+
+      // Merge ids[bestJ] into ids[bestI] (keep the lexicographically earlier
+      // ID as the canonical one for determinism).
+      const keepId = ids[bestI]!;
+      const dropId = ids[bestJ]!;
+
+      // Average the two centroids (equal weight — simple heuristic).
+      const keepEmb = working.get(keepId)!;
+      const dropEmb = working.get(dropId)!;
+      const merged = new Float32Array(keepEmb.length);
+      for (let k = 0; k < keepEmb.length; k++) {
+        merged[k] = (keepEmb[k]! + dropEmb[k]!) / 2;
+      }
+
+      working.set(keepId, merged);
+      working.delete(dropId);
+
+      // Record the rename. Chase any existing mappings so the final map is
+      // transitively resolved.
+      renameMap.set(dropId, keepId);
+
+      // Resolve transitive renames: if dropId was itself a canonical target of
+      // an earlier merge, update those entries to point to keepId.
+      for (const [old, target] of renameMap) {
+        if (target === dropId) {
+          renameMap.set(old, keepId);
+        }
+      }
+    }
+
+    return renameMap;
+  }
+}

package/src/DiarizationEngine.ts

@@ -0,0 +1,94 @@
+/**
+ * @file DiarizationEngine.ts
+ * @description Factory that creates {@link DiarizationSession} instances,
+ * selecting the appropriate backend based on the supplied configuration.
+ *
+ * When a provider backend is requested (`config.backend === 'provider'`), the
+ * engine creates a {@link ProviderDiarizationBackend} and wraps it in a
+ * session.  Otherwise it constructs the full local pipeline:
+ * {@link SlidingWindowExtractor} → {@link LocalDiarizationBackend} →
+ * {@link SpeakerEmbeddingCache}.
+ *
+ * @module diarization/DiarizationEngine
+ */
+
+import type { IDiarizationEngine, IDiarizationSession, DiarizationConfig } from './types.js';
+import { DiarizationSession } from './DiarizationSession.js';
+import { SpeakerEmbeddingCache } from './SpeakerEmbeddingCache.js';
+import { SlidingWindowExtractor } from './SlidingWindowExtractor.js';
+import { LocalDiarizationBackend } from './LocalDiarizationBackend.js';
+import { ProviderDiarizationBackend } from './ProviderDiarizationBackend.js';
+
+// ---------------------------------------------------------------------------
+// Default values
+// ---------------------------------------------------------------------------
+
+const DEFAULT_SIMILARITY_THRESHOLD = 0.7;
+const DEFAULT_CHUNK_SIZE_MS = 1500;
+const DEFAULT_OVERLAP_MS = 500;
+const DEFAULT_SAMPLE_RATE = 16_000;
+
+// ---------------------------------------------------------------------------
+// Main class
+// ---------------------------------------------------------------------------
+
+/**
+ * Factory for diarization sessions.
+ *
+ * @example
+ * ```ts
+ * const engine = new DiarizationEngine();
+ *
+ * // Local mode (default)
+ * const session = engine.startSession();
+ * session.on('speaker_identified', ({ speakerId }) => console.log(speakerId));
+ *
+ * // Provider mode
+ * const providerSession = engine.startSession({ backend: 'provider' });
+ * sttSession.on('transcript', (e) => providerSession.labelTranscript(e));
+ * ```
+ */
+export class DiarizationEngine implements IDiarizationEngine {
+  // -------------------------------------------------------------------------
+  // IDiarizationEngine implementation
+  // -------------------------------------------------------------------------
+
+  /**
+   * Create and return a new {@link IDiarizationSession}.
+   *
+   * The backend is selected based on `config.backend`:
+   * - `'provider'` — use {@link ProviderDiarizationBackend}
+   * - `'local'` (default) — use {@link LocalDiarizationBackend} with a fresh
+   *   {@link SpeakerEmbeddingCache} and {@link SlidingWindowExtractor}
+   *
+   * @param config - Optional session configuration.
+   * @returns A ready-to-use {@link DiarizationSession}.
+   */
+  startSession(config: DiarizationConfig = {}): IDiarizationSession {
+    const backend = config.backend ?? 'local';
+
+    if (backend === 'provider') {
+      return new DiarizationSession({
+        kind: 'provider',
+        backend: new ProviderDiarizationBackend(),
+      });
+    }
+
+    // Local mode — wire up the full pipeline.
+    const threshold = config.similarityThreshold ?? DEFAULT_SIMILARITY_THRESHOLD;
+    const chunkSizeMs = config.chunkSizeMs ?? DEFAULT_CHUNK_SIZE_MS;
+    const overlapMs = config.overlapMs ?? DEFAULT_OVERLAP_MS;
+    const sampleRate = config.sampleRate ?? DEFAULT_SAMPLE_RATE;
+
+    const cache = new SpeakerEmbeddingCache(threshold);
+    const extractor = new SlidingWindowExtractor(chunkSizeMs, overlapMs, sampleRate);
+    const localBackend = new LocalDiarizationBackend(cache, extractor);
+
+    return new DiarizationSession({
+      kind: 'local',
+      backend: localBackend,
+      cache,
+      extractor,
+    });
+  }
+}

package/src/DiarizationSession.ts

@@ -0,0 +1,178 @@
+/**
+ * @file DiarizationSession.ts
+ * @description Live diarization session that wraps either a
+ * {@link ProviderDiarizationBackend} or a {@link LocalDiarizationBackend}.
+ *
+ * The session implements {@link IDiarizationSession} and bridges the raw audio
+ * / transcript inputs to the selected backend, forwarding `speaker_identified`
+ * and `segment_ready` events to callers.
+ *
+ * @module diarization/DiarizationSession
+ */
+
+import { EventEmitter } from 'node:events';
+import type { IDiarizationSession, AudioFrame, TranscriptEvent, DiarizedSegment, SpeakerIdentified } from './types.js';
+import type { ProviderDiarizationBackend } from './ProviderDiarizationBackend.js';
+import type { LocalDiarizationBackend } from './LocalDiarizationBackend.js';
+import type { SpeakerEmbeddingCache } from './SpeakerEmbeddingCache.js';
+import type { SlidingWindowExtractor } from './SlidingWindowExtractor.js';
+
+// ---------------------------------------------------------------------------
+// Discriminated union for the backend
+// ---------------------------------------------------------------------------
+
+/** Provider-delegated backend descriptor. */
+interface ProviderBackendDescriptor {
+  kind: 'provider';
+  backend: ProviderDiarizationBackend;
+}
+
+/** Local-clustering backend descriptor. */
+interface LocalBackendDescriptor {
+  kind: 'local';
+  backend: LocalDiarizationBackend;
+  cache: SpeakerEmbeddingCache;
+  extractor: SlidingWindowExtractor;
+}
+
+type BackendDescriptor = ProviderBackendDescriptor | LocalBackendDescriptor;
+
+// ---------------------------------------------------------------------------
+// Main class
+// ---------------------------------------------------------------------------
+
+/**
+ * A live diarization session.
+ *
+ * Create instances via {@link DiarizationEngine.startSession} rather than
+ * constructing directly.
+ *
+ * ### Events
+ * - `speaker_identified` — payload: {@link SpeakerIdentified}
+ * - `segment_ready`      — payload: {@link DiarizedSegment}
+ * - `error`              — payload: `Error`
+ * - `close`              — no payload
+ */
+export class DiarizationSession extends EventEmitter implements IDiarizationSession {
+  // -------------------------------------------------------------------------
+  // Private state
+  // -------------------------------------------------------------------------
+
+  private closed = false;
+
+  /** Most recently identified speaker (used for transcript labelling in local mode). */
+  private lastSpeakerId: string | null = null;
+
+  // -------------------------------------------------------------------------
+  // Constructor
+  // -------------------------------------------------------------------------
+
+  /**
+   * @param descriptor - Resolved backend descriptor.
+   */
+  constructor(private readonly descriptor: BackendDescriptor) {
+    super();
+
+    if (descriptor.kind === 'local') {
+      // Forward speaker_identified events from the local backend.
+      descriptor.backend.on('speaker_identified', (payload: SpeakerIdentified) => {
+        this.lastSpeakerId = payload.speakerId;
+        this.emit('speaker_identified', payload);
+      });
+
+      descriptor.backend.start();
+    }
+  }
+
+  // -------------------------------------------------------------------------
+  // IDiarizationSession implementation
+  // -------------------------------------------------------------------------
+
+  /**
+   * Push a raw audio frame into the pipeline.
+   *
+   * In local mode the frame is forwarded to the sliding-window extractor.
+   * In provider mode this is a no-op.
+   *
+   * @param frame - Audio frame to process.
+   */
+  pushAudio(frame: AudioFrame): void {
+    if (this.closed) return;
+
+    if (this.descriptor.kind === 'local') {
+      this.descriptor.extractor.pushAudio(frame);
+    }
+    // Provider mode: audio processing happens on the provider side.
+  }
+
+  /**
+   * Attach speaker labels to a transcript event and emit a segment.
+   *
+   * In provider mode, word-level speaker labels are extracted directly.  In
+   * local mode, the most recently identified speaker is used.
+   *
+   * @param event - Transcript event from the STT provider.
+   * @returns The labelled {@link DiarizedSegment}, or `null` if no speaker
+   *   information is available yet.
+   */
+  labelTranscript(event: TranscriptEvent): DiarizedSegment | null {
+    if (this.closed) return null;
+
+    let segment: DiarizedSegment | null = null;
+
+    if (this.descriptor.kind === 'provider') {
+      segment = this.descriptor.backend.pushTranscriptEvent(event);
+    } else {
+      // Local mode — use the last identified speaker.
+      if (this.lastSpeakerId !== null) {
+        segment = {
+          speakerId: this.lastSpeakerId,
+          text: event.text,
+          startMs: event.timestamp,
+          endMs: event.timestamp + event.text.length * 50, // rough estimate
+          confidence: 0.7,
+        };
+      }
+    }
+
+    if (segment) {
+      this.emit('segment_ready', segment);
+    }
+
+    return segment;
+  }
+
+  /**
+   * Pre-register a known speaker voiceprint.
+   *
+   * Only has effect in local mode.  In provider mode, speaker IDs come
+   * directly from the STT provider and this call is a no-op.
+   *
+   * @param id - Human-readable speaker name.
+   * @param voiceprint - Reference embedding for this speaker.
+   */
+  enrollSpeaker(id: string, voiceprint: Float32Array): void {
+    if (this.descriptor.kind === 'local') {
+      this.descriptor.cache.enrollSpeaker(id, voiceprint);
+    }
+    // Provider mode: enrolment is handled by the provider; ignore here.
+  }
+
+  /**
+   * Terminate the session and release all resources.
+   *
+   * After `close()` no further events are emitted.
+   */
+  close(): void {
+    if (this.closed) return;
+    this.closed = true;
+
+    if (this.descriptor.kind === 'local') {
+      this.descriptor.backend.stop();
+      this.descriptor.extractor.flush();
+    }
+
+    this.emit('close');
+    this.removeAllListeners();
+  }
+}

package/src/LocalDiarizationBackend.ts

@@ -0,0 +1,201 @@
+/**
+ * @file LocalDiarizationBackend.ts
+ * @description Offline speaker diarization backend using a sliding-window
+ * spectral-centroid voiceprint and cosine similarity matching.
+ *
+ * Audio is accumulated in a {@link SlidingWindowExtractor}.  When a full chunk
+ * is ready, a 16-dimensional feature vector is extracted and matched against
+ * cached speaker centroids.  If no centroid exceeds the similarity threshold a
+ * new `Speaker_N` identity is created.
+ *
+ * The built-in feature extractor is intentionally lightweight.  An ONNX
+ * x-vector model can replace {@link LocalDiarizationBackend.extractSimpleEmbedding}
+ * later without changing the surrounding API.
+ *
+ * @module diarization/LocalDiarizationBackend
+ */
+
+import { EventEmitter } from 'node:events';
+import type { SpeakerEmbeddingCache } from './SpeakerEmbeddingCache.js';
+import type { SlidingWindowExtractor } from './SlidingWindowExtractor.js';
+import type { SpeakerIdentified } from './types.js';
+
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
+
+/** Size of the local feature vector in dimensions. */
+const FEATURE_DIMS = 16;
+
+/** Number of octave-band buckets used by the feature extractor. */
+const OCTAVE_BANDS = 4;
+
+/** Dimensions per octave band (FEATURE_DIMS / OCTAVE_BANDS). */
+const DIMS_PER_BAND = FEATURE_DIMS / OCTAVE_BANDS;
+
+// ---------------------------------------------------------------------------
+// Main class
+// ---------------------------------------------------------------------------
+
+/**
+ * Fully offline diarization backend.
+ *
+ * Connects a {@link SlidingWindowExtractor} to a {@link SpeakerEmbeddingCache}
+ * and emits `speaker_identified` events whenever a speaker label is resolved
+ * for a new audio chunk.
+ *
+ * @example
+ * ```ts
+ * const cache = new SpeakerEmbeddingCache(0.7);
+ * const extractor = new SlidingWindowExtractor(1500, 500, 16000);
+ * const backend = new LocalDiarizationBackend(cache, extractor);
+ *
+ * backend.on('speaker_identified', ({ speakerId }) => console.log(speakerId));
+ * backend.start();
+ *
+ * microphone.on('frame', (f) => extractor.pushAudio(f));
+ * ```
+ */
+export class LocalDiarizationBackend extends EventEmitter {
+  // -------------------------------------------------------------------------
+  // Constructor
+  // -------------------------------------------------------------------------
+
+  /**
+   * @param cache - Speaker voiceprint cache used for matching and centroid
+   *   updates.
+   * @param extractor - Sliding-window extractor that supplies audio chunks.
+   */
+  constructor(
+    private readonly cache: SpeakerEmbeddingCache,
+    private readonly extractor: SlidingWindowExtractor,
+  ) {
+    super();
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API
+  // -------------------------------------------------------------------------
+
+  /**
+   * Attach the `'chunk_ready'` listener to the extractor and begin processing
+   * audio chunks.
+   *
+   * Calling `start()` more than once is safe — subsequent calls are no-ops
+   * because the same bound listener is reused.
+   */
+  start(): void {
+    this.extractor.on('chunk_ready', this._onChunkReady);
+  }
+
+  /**
+   * Detach from the extractor and stop processing.
+   */
+  stop(): void {
+    this.extractor.off('chunk_ready', this._onChunkReady);
+  }
+
+  // -------------------------------------------------------------------------
+  // Internal event handler
+  // -------------------------------------------------------------------------
+
+  /**
+   * Bound handler for `'chunk_ready'` events from the extractor.
+   *
+   * Stored as an arrow-function property so the same reference is used for
+   * both `on()` and `off()`.
+   */
+  private readonly _onChunkReady = (chunk: Float32Array): void => {
+    const embedding = this.extractSimpleEmbedding(chunk);
+    const { speakerId } = this.cache.getOrCreateSpeaker(embedding);
+
+    const payload: SpeakerIdentified = {
+      speakerId,
+      confidence: 0.7,
+      timestamp: Date.now(),
+    };
+
+    this.emit('speaker_identified', payload);
+  };
+
+  // -------------------------------------------------------------------------
+  // Feature extraction
+  // -------------------------------------------------------------------------
+
+  /**
+   * Extract a 16-dimensional voiceprint from a raw PCM chunk.
+   *
+   * Feature layout (4 octave bands × 4 features per band):
+   * - Band 0: sub-bass  (0 –  500 Hz equivalent bucket)
+   * - Band 1: bass      (500 – 2000 Hz equivalent bucket)
+   * - Band 2: mid       (2000 – 4000 Hz equivalent bucket)
+   * - Band 3: high      (4000 Hz+ equivalent bucket)
+   *
+   * For each band the following statistics are computed:
+   * - [0] RMS energy
+   * - [1] Spectral centroid (normalised to band range)
+   * - [2] Zero-crossing rate
+   * - [3] Delta energy (mean absolute frame-to-frame change)
+   *
+   * This is a placeholder implementation suitable for demos and testing.
+   * Replace with an ONNX x-vector model for production accuracy.
+   *
+   * @param chunk - Float32 PCM audio data.
+   * @returns 16-dimensional embedding vector.
+   */
+  extractSimpleEmbedding(chunk: Float32Array): Float32Array {
+    const features = new Float32Array(FEATURE_DIMS);
+    const n = chunk.length;
+
+    if (n === 0) return features;
+
+    // Split the chunk index range into OCTAVE_BANDS equal bands.
+    const bandSize = Math.ceil(n / OCTAVE_BANDS);
+
+    for (let band = 0; band < OCTAVE_BANDS; band++) {
+      const start = band * bandSize;
+      const end = Math.min(start + bandSize, n);
+      const bandLen = end - start;
+
+      if (bandLen === 0) continue;
+
+      // --- RMS energy ---
+      let sumSq = 0;
+      for (let i = start; i < end; i++) sumSq += (chunk[i]!) ** 2;
+      const rms = Math.sqrt(sumSq / bandLen);
+
+      // --- Spectral centroid (weighted mean index, normalised to [0,1]) ---
+      let weightedIdx = 0;
+      let totalAbs = 0;
+      for (let i = start; i < end; i++) {
+        const abs = Math.abs(chunk[i]!);
+        weightedIdx += (i - start) * abs;
+        totalAbs += abs;
+      }
+      const centroid = totalAbs > 0 ? weightedIdx / totalAbs / bandLen : 0;
+
+      // --- Zero-crossing rate ---
+      let zeroCrossings = 0;
+      for (let i = start + 1; i < end; i++) {
+        if ((chunk[i]! >= 0) !== (chunk[i - 1]! >= 0)) zeroCrossings++;
+      }
+      const zcr = zeroCrossings / (bandLen - 1);
+
+      // --- Delta energy (mean |diff| between consecutive samples) ---
+      let deltaSum = 0;
+      for (let i = start + 1; i < end; i++) {
+        deltaSum += Math.abs(chunk[i]! - chunk[i - 1]!);
+      }
+      const delta = deltaSum / (bandLen - 1);
+
+      // Write into the feature vector.
+      const base = band * DIMS_PER_BAND;
+      features[base + 0] = rms;
+      features[base + 1] = centroid;
+      features[base + 2] = zcr;
+      features[base + 3] = delta;
+    }
+
+    return features;
+  }
+}