@framers/agentos-ext-diarization 0.2.0

package/src/ProviderDiarizationBackend.ts

@@ -0,0 +1,119 @@
+/**
+ * @file ProviderDiarizationBackend.ts
+ * @description Thin adapter that extracts speaker labels from STT provider
+ * transcript events containing word-level speaker annotations.
+ *
+ * Many cloud STT providers (Deepgram, AssemblyAI, …) support native
+ * diarization that produces a `speaker` label on each word in the transcript.
+ * This backend consumes those events and re-emits them as
+ * {@link DiarizedSegment} objects without any additional compute.
+ *
+ * @module diarization/ProviderDiarizationBackend
+ */
+
+import type { TranscriptEvent, DiarizedSegment, TranscriptWord } from './types.js';
+
+// ---------------------------------------------------------------------------
+// Main class
+// ---------------------------------------------------------------------------
+
+/**
+ * Extracts speaker labels from word-level STT results.
+ *
+ * Stateless — each call to {@link pushTranscriptEvent} is independent.  The
+ * caller is responsible for any session-level bookkeeping.
+ *
+ * @example
+ * ```ts
+ * const backend = new ProviderDiarizationBackend();
+ * sttSession.on('transcript', (event) => {
+ *   const segment = backend.pushTranscriptEvent(event);
+ *   if (segment) console.log(segment.speakerId, segment.text);
+ * });
+ * ```
+ */
+export class ProviderDiarizationBackend {
+  // -------------------------------------------------------------------------
+  // Public API
+  // -------------------------------------------------------------------------
+
+  /**
+   * Process a transcript event and extract the dominant speaker label.
+   *
+   * The dominant speaker is determined by majority vote across the word-level
+   * `speaker` annotations.  If no words carry speaker labels, `null` is
+   * returned.
+   *
+   * For final transcripts (`isFinal: true`) the returned segment covers the
+   * full utterance.  For interim transcripts the segment is a best-effort
+   * approximation based on the current words.
+   *
+   * @param event - Transcript event from an STT provider session.
+   * @returns A {@link DiarizedSegment} if speaker information is available,
+   *   or `null` if the event lacks word-level speaker labels.
+   */
+  pushTranscriptEvent(event: TranscriptEvent): DiarizedSegment | null {
+    const words = event.words;
+
+    if (!words || words.length === 0) {
+      return null;
+    }
+
+    // Collect all words that have a speaker label.
+    const labelledWords = words.filter((w) => w.speaker !== undefined && w.speaker !== '');
+
+    if (labelledWords.length === 0) {
+      return null;
+    }
+
+    // Majority-vote for the dominant speaker in this transcript event.
+    const speakerId = this.dominantSpeaker(labelledWords);
+
+    // Derive segment boundaries from word timestamps (seconds → milliseconds).
+    const firstWord = labelledWords[0]!;
+    const lastWord = labelledWords[labelledWords.length - 1]!;
+
+    // Average per-word confidence (default 1 if not present).
+    const avgConfidence =
+      labelledWords.reduce((sum, w) => sum + (w.confidence ?? 1), 0) / labelledWords.length;
+
+    return {
+      speakerId,
+      text: event.text,
+      startMs: Math.round(firstWord.start * 1000),
+      endMs: Math.round(lastWord.end * 1000),
+      confidence: avgConfidence,
+    };
+  }
+
+  // -------------------------------------------------------------------------
+  // Internal helpers
+  // -------------------------------------------------------------------------
+
+  /**
+   * Determine the most-frequently-occurring speaker label among a set of words.
+   *
+   * @param words - Words that all have a defined `speaker` property.
+   * @returns The dominant `speaker` label.
+   */
+  private dominantSpeaker(words: TranscriptWord[]): string {
+    const counts = new Map<string, number>();
+
+    for (const word of words) {
+      const label = word.speaker!;
+      counts.set(label, (counts.get(label) ?? 0) + 1);
+    }
+
+    let best = words[0]!.speaker!;
+    let bestCount = 0;
+
+    for (const [label, count] of counts) {
+      if (count > bestCount) {
+        bestCount = count;
+        best = label;
+      }
+    }
+
+    return best;
+  }
+}

package/src/SlidingWindowExtractor.ts

@@ -0,0 +1,174 @@
+/**
+ * @file SlidingWindowExtractor.ts
+ * @description Ring buffer that accumulates audio frames into fixed-size chunks
+ * with configurable overlap, tuned for diarization (1.5 s chunks, 0.5 s overlap).
+ *
+ * Mirrors the design of `SlidingWindowBuffer` from the Whisper STT pack but
+ * operates on {@link AudioFrame} objects rather than raw sample arrays, and
+ * uses millisecond-based chunk/overlap sizes for a friendlier API.
+ *
+ * @module diarization/SlidingWindowExtractor
+ */
+
+import { EventEmitter } from 'node:events';
+import type { AudioFrame } from './types.js';
+
+// ---------------------------------------------------------------------------
+// Events interface
+// ---------------------------------------------------------------------------
+
+/** Event map for {@link SlidingWindowExtractor}. */
+export interface SlidingWindowExtractorEvents {
+  /** Emitted when a complete chunk of audio samples is ready. */
+  chunk_ready: [chunk: Float32Array];
+}
+
+// ---------------------------------------------------------------------------
+// Main class
+// ---------------------------------------------------------------------------
+
+/**
+ * Sliding-window audio extractor for the diarization pipeline.
+ *
+ * Accepts {@link AudioFrame} objects pushed by the voice pipeline, accumulates
+ * samples in an internal ring buffer, and emits `'chunk_ready'` events when a
+ * full chunk is available.  The last `overlapSamples` of each chunk are kept
+ * as context at the head of the next chunk.
+ *
+ * @example
+ * ```ts
+ * const extractor = new SlidingWindowExtractor(1500, 500, 16000);
+ * extractor.on('chunk_ready', (chunk) => processChunk(chunk));
+ * microphone.on('frame', (f) => extractor.pushAudio(f));
+ * ```
+ */
+export class SlidingWindowExtractor extends EventEmitter {
+  // -------------------------------------------------------------------------
+  // Private state
+  // -------------------------------------------------------------------------
+
+  /** Internal sample store.  Sized to {@link chunkSizeSamples}. */
+  private buffer: Float32Array;
+
+  /** Current write position within {@link buffer}. */
+  private writePos = 0;
+
+  // -------------------------------------------------------------------------
+  // Computed dimensions (from ms → samples)
+  // -------------------------------------------------------------------------
+
+  /** Number of samples per emitted chunk. */
+  private readonly chunkSizeSamples: number;
+
+  /** Number of overlap samples carried forward to the next chunk. */
+  private readonly overlapSamples: number;
+
+  // -------------------------------------------------------------------------
+  // Constructor
+  // -------------------------------------------------------------------------
+
+  /**
+   * @param chunkSizeMs - Duration of each emitted chunk in milliseconds.
+   *   @defaultValue 1500
+   * @param overlapMs - Overlap between consecutive chunks in milliseconds.
+   *   Must be less than `chunkSizeMs`.
+   *   @defaultValue 500
+   * @param sampleRate - Expected audio sample rate in Hz.
+   *   @defaultValue 16000
+   */
+  constructor(
+    private readonly chunkSizeMs: number = 1500,
+    private readonly overlapMs: number = 500,
+    private readonly sampleRate: number = 16_000,
+  ) {
+    super();
+
+    this.chunkSizeSamples = Math.round((chunkSizeMs / 1000) * sampleRate);
+    this.overlapSamples = Math.round((overlapMs / 1000) * sampleRate);
+
+    if (this.overlapSamples >= this.chunkSizeSamples) {
+      throw new RangeError(
+        `overlapMs (${overlapMs}) must result in fewer samples than chunkSizeMs (${chunkSizeMs})`,
+      );
+    }
+
+    this.buffer = new Float32Array(this.chunkSizeSamples);
+  }
+
+  // -------------------------------------------------------------------------
+  // Public API
+  // -------------------------------------------------------------------------
+
+  /**
+   * Push an audio frame into the extractor.
+   *
+   * Samples are appended to the internal ring buffer.  When the buffer fills
+   * to `chunkSizeSamples`, a `'chunk_ready'` event is emitted and the overlap
+   * region is copied to the head of the buffer for the next chunk.
+   *
+   * @param frame - Audio frame to append.
+   */
+  pushAudio(frame: AudioFrame): void {
+    this.pushSamples(frame.samples);
+  }
+
+  /**
+   * Flush any remaining samples as a final partial chunk.
+   *
+   * If the buffer is empty this is a no-op.
+   */
+  flush(): void {
+    if (this.writePos === 0) return;
+    this.emit('chunk_ready', this.buffer.slice(0, this.writePos));
+    this.reset();
+  }
+
+  /**
+   * Clear the buffer without emitting a chunk.
+   */
+  reset(): void {
+    this.buffer = new Float32Array(this.chunkSizeSamples);
+    this.writePos = 0;
+  }
+
+  // -------------------------------------------------------------------------
+  // Accessors
+  // -------------------------------------------------------------------------
+
+  /** Number of samples currently held in the buffer. */
+  get bufferedSamples(): number {
+    return this.writePos;
+  }
+
+  // -------------------------------------------------------------------------
+  // Internal helpers
+  // -------------------------------------------------------------------------
+
+  /**
+   * Append raw PCM samples and emit chunks as they fill.
+   *
+   * @param samples - Float32 PCM data to append.
+   */
+  private pushSamples(samples: Float32Array): void {
+    let srcOffset = 0;
+
+    while (srcOffset < samples.length) {
+      const spaceLeft = this.chunkSizeSamples - this.writePos;
+      const copyCount = Math.min(spaceLeft, samples.length - srcOffset);
+
+      this.buffer.set(samples.subarray(srcOffset, srcOffset + copyCount), this.writePos);
+      this.writePos += copyCount;
+      srcOffset += copyCount;
+
+      if (this.writePos >= this.chunkSizeSamples) {
+        // Emit a copy — the buffer is reused for the next chunk.
+        this.emit('chunk_ready', this.buffer.slice());
+
+        // Copy the overlap region to the start of the buffer.
+        const overlapStart = this.chunkSizeSamples - this.overlapSamples;
+        this.buffer.copyWithin(0, overlapStart, this.chunkSizeSamples);
+        this.writePos = this.overlapSamples;
+      }
+    }
+  }
+}

package/src/SpeakerEmbeddingCache.ts

@@ -0,0 +1,228 @@
+/**
+ * @file SpeakerEmbeddingCache.ts
+ * @description Stores and matches speaker voiceprints using cosine similarity.
+ *
+ * Each speaker is represented as a running centroid — the mean of all
+ * embeddings observed so far.  New audio embeddings are matched against
+ * existing centroids; if the best match exceeds the similarity threshold the
+ * frame is attributed to that speaker and the centroid is updated.  Otherwise
+ * a new `Speaker_N` identity is created.
+ *
+ * @module diarization/SpeakerEmbeddingCache
+ */
+
+// ---------------------------------------------------------------------------
+// Internal types
+// ---------------------------------------------------------------------------
+
+/** Running centroid entry stored per speaker. */
+interface CentroidEntry {
+  /** Current mean embedding vector. */
+  embedding: Float32Array;
+  /** Number of frames that have contributed to this centroid. */
+  count: number;
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Compute the cosine similarity between two embedding vectors.
+ *
+ * Returns a value in the range `[-1, 1]`.  A small epsilon is added to the
+ * denominator to prevent division-by-zero on silent/zero frames.
+ *
+ * @param a - First embedding vector.
+ * @param b - Second embedding vector.
+ * @returns Cosine similarity.
+ */
+export function cosineSimilarity(a: Float32Array, b: Float32Array): number {
+  let dot = 0;
+  let magA = 0;
+  let magB = 0;
+
+  for (let i = 0; i < a.length; i++) {
+    dot += a[i]! * b[i]!;
+    magA += (a[i]!) ** 2;
+    magB += (b[i]!) ** 2;
+  }
+
+  return dot / (Math.sqrt(magA) * Math.sqrt(magB) + 1e-8);
+}
+
+// ---------------------------------------------------------------------------
+// Main class
+// ---------------------------------------------------------------------------
+
+/**
+ * In-memory cache of speaker voiceprint centroids.
+ *
+ * Centroids are updated with an online running-average as new embeddings
+ * arrive, so memory usage is O(speakers × embedding dimensions) regardless of
+ * the number of frames observed.
+ *
+ * @example
+ * ```ts
+ * const cache = new SpeakerEmbeddingCache(0.7);
+ * const { speakerId, isNew } = cache.getOrCreateSpeaker(embedding);
+ * ```
+ */
+export class SpeakerEmbeddingCache {
+  // -------------------------------------------------------------------------
+  // Private state
+  // -------------------------------------------------------------------------
+
+  /**
+   * Centroid map: `speakerId → { embedding, count }`.
+   *
+   * Access via {@link centroids} for testing/introspection.
+   */
+  private readonly _centroids = new Map<string, CentroidEntry>();
+
+  // -------------------------------------------------------------------------
+  // Constructor
+  // -------------------------------------------------------------------------
+
+  /**
+   * @param threshold - Minimum cosine similarity required to attribute a new
+   *   embedding to an existing speaker.  Embeddings below this threshold
+   *   create a new speaker identity.
+   *   @defaultValue 0.7
+   */
+  constructor(private readonly threshold: number = 0.7) {}
+
+  // -------------------------------------------------------------------------
+  // Public API
+  // -------------------------------------------------------------------------
+
+  /**
+   * Update the running centroid for `speakerId` with a new embedding.
+   *
+   * The centroid is recalculated as an online running average:
+   * `new_mean = old_mean + (x - old_mean) / (count + 1)`.
+   *
+   * @param speakerId - Speaker identity to update.
+   * @param embedding - New embedding vector to incorporate.
+   */
+  addEmbedding(speakerId: string, embedding: Float32Array): void {
+    const existing = this._centroids.get(speakerId);
+
+    if (!existing) {
+      // First frame — store a copy of the embedding directly.
+      this._centroids.set(speakerId, {
+        embedding: embedding.slice(),
+        count: 1,
+      });
+      return;
+    }
+
+    // Online running average: mean_n = mean_{n-1} + (x - mean_{n-1}) / n
+    const newCount = existing.count + 1;
+    const newEmbedding = new Float32Array(existing.embedding.length);
+
+    for (let i = 0; i < existing.embedding.length; i++) {
+      newEmbedding[i] = existing.embedding[i]! + (embedding[i]! - existing.embedding[i]!) / newCount;
+    }
+
+    this._centroids.set(speakerId, { embedding: newEmbedding, count: newCount });
+  }
+
+  /**
+   * Find the existing speaker whose centroid is most similar to `embedding`.
+   *
+   * Returns `null` if no speaker exceeds the similarity threshold or the cache
+   * is empty.
+   *
+   * @param embedding - Query embedding vector.
+   * @returns Best match above threshold, or `null`.
+   */
+  findClosestSpeaker(embedding: Float32Array): { speakerId: string; similarity: number } | null {
+    let bestId: string | null = null;
+    let bestSim = -Infinity;
+
+    for (const [id, entry] of this._centroids) {
+      const sim = cosineSimilarity(embedding, entry.embedding);
+      if (sim > bestSim) {
+        bestSim = sim;
+        bestId = id;
+      }
+    }
+
+    if (bestId === null || bestSim < this.threshold) {
+      return null;
+    }
+
+    return { speakerId: bestId, similarity: bestSim };
+  }
+
+  /**
+   * Look up or create a speaker identity for `embedding`.
+   *
+   * If a match above the threshold is found, the centroid is updated and the
+   * existing ID is returned.  Otherwise a new `Speaker_N` identity is
+   * registered.
+   *
+   * @param embedding - Audio embedding to resolve.
+   * @returns Speaker ID and whether a new identity was created.
+   */
+  getOrCreateSpeaker(embedding: Float32Array): { speakerId: string; isNew: boolean } {
+    const match = this.findClosestSpeaker(embedding);
+
+    if (match) {
+      this.addEmbedding(match.speakerId, embedding);
+      return { speakerId: match.speakerId, isNew: false };
+    }
+
+    const newId = `Speaker_${this._centroids.size}`;
+    this.addEmbedding(newId, embedding);
+    return { speakerId: newId, isNew: true };
+  }
+
+  /**
+   * Pre-register a known speaker voiceprint, overwriting any existing centroid
+   * for this ID.
+   *
+   * Use this to seed the cache with enrolment recordings before a session
+   * begins so that known participants are matched by name rather than the
+   * auto-generated `Speaker_N` label.
+   *
+   * @param id - Human-readable speaker identifier.
+   * @param voiceprint - Reference embedding for this speaker.
+   */
+  enrollSpeaker(id: string, voiceprint: Float32Array): void {
+    this._centroids.set(id, { embedding: voiceprint.slice(), count: 1 });
+  }
+
+  // -------------------------------------------------------------------------
+  // Accessors (useful for testing and ClusteringStrategy)
+  // -------------------------------------------------------------------------
+
+  /**
+   * Read-only view of the internal centroid map.
+   *
+   * Returns a snapshot — modifications to the returned map do not affect the
+   * cache.
+   */
+  get centroids(): Map<string, Float32Array> {
+    const result = new Map<string, Float32Array>();
+    for (const [id, entry] of this._centroids) {
+      result.set(id, entry.embedding.slice());
+    }
+    return result;
+  }
+
+  /** Number of distinct speaker identities currently tracked. */
+  get size(): number {
+    return this._centroids.size;
+  }
+
+  /**
+   * Remove all centroids from the cache.
+   *
+   * Useful for resetting state between test cases.
+   */
+  clear(): void {
+    this._centroids.clear();
+  }
+}

package/src/index.ts

@@ -0,0 +1,105 @@
+/**
+ * @file index.ts
+ * @description Pack factory for the Diarization extension pack.
+ *
+ * ### Usage (direct)
+ * ```ts
+ * import { createDiarizationEngine } from '@framers/agentos-ext-diarization';
+ *
+ * const engine = createDiarizationEngine();
+ * const session = engine.startSession({ backend: 'local' });
+ * session.on('speaker_identified', ({ speakerId }) => console.log(speakerId));
+ * ```
+ *
+ * ### Usage (manifest-driven)
+ * ```json
+ * { "packs": [{ "module": "@framers/agentos-ext-diarization" }] }
+ * ```
+ *
+ * @module diarization
+ */
+
+import { DiarizationEngine } from './DiarizationEngine.js';
+
+// ---------------------------------------------------------------------------
+// Local types (avoid hard runtime dep on @framers/agentos)
+// ---------------------------------------------------------------------------
+
+/** Subset of ExtensionDescriptor required by this pack. */
+interface ExtensionDescriptor {
+  id: string;
+  kind: string;
+  payload: unknown;
+  enableByDefault?: boolean;
+  metadata?: Record<string, unknown>;
+}
+
+/** Subset of ExtensionPack required by this pack. */
+interface ExtensionPack {
+  id: string;
+  descriptors: ExtensionDescriptor[];
+}
+
+/** Subset of ExtensionPackContext required by this pack. */
+interface ExtensionPackContext {
+  getSecret?: (id: string) => string | undefined;
+  options?: Record<string, unknown>;
+}
+
+/** Kind constant matching packages/agentos/src/extensions/types.ts. */
+const EXTENSION_KIND_DIARIZATION = 'diarization-provider';
+
+// ---------------------------------------------------------------------------
+// Factory
+// ---------------------------------------------------------------------------
+
+/**
+ * Create a standalone {@link DiarizationEngine} instance.
+ *
+ * Use this when composing the engine programmatically outside of the
+ * AgentOS extension system.
+ *
+ * @returns A configured {@link DiarizationEngine}.
+ */
+export function createDiarizationEngine(): DiarizationEngine {
+  return new DiarizationEngine();
+}
+
+/**
+ * AgentOS manifest factory function.
+ *
+ * Returns an {@link ExtensionPack} containing a single `diarization-provider`
+ * descriptor backed by {@link DiarizationEngine}.
+ *
+ * @param _context - Pack context supplied by the extension manager.
+ * @returns A fully configured {@link ExtensionPack}.
+ */
+export function createExtensionPack(_context: ExtensionPackContext): ExtensionPack {
+  const engine = new DiarizationEngine();
+
+  return {
+    id: 'diarization',
+    descriptors: [
+      {
+        id: 'diarization-engine',
+        kind: EXTENSION_KIND_DIARIZATION,
+        payload: engine,
+        enableByDefault: true,
+        metadata: { providerId: 'diarization-local' },
+      },
+    ],
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Re-exports
+// ---------------------------------------------------------------------------
+
+export * from './types.js';
+export { DiarizationEngine } from './DiarizationEngine.js';
+export { DiarizationSession } from './DiarizationSession.js';
+export { SpeakerEmbeddingCache, cosineSimilarity } from './SpeakerEmbeddingCache.js';
+export { ClusteringStrategy } from './ClusteringStrategy.js';
+export { SlidingWindowExtractor } from './SlidingWindowExtractor.js';
+export { ProviderDiarizationBackend } from './ProviderDiarizationBackend.js';
+export { LocalDiarizationBackend } from './LocalDiarizationBackend.js';