@dxos/plugin-transcription 0.8.2-main.12df754 → 0.8.2-main.30e4dbb

package/src/model/model.ts

@@ -9,17 +9,17 @@ import { Event } from '@dxos/async';
 import { invariant } from '@dxos/invariant';
 import { log } from '@dxos/log';
 
-export type Block = { id: string };
+export type Chunk = { id: string };
 
 /**
  * Covert block to a set of markdown lines.
  */
-export type BlockRenderer<T extends Block> = (block: T, debug?: boolean) => string[];
+export type ChunkRenderer<T extends Chunk> = (chunk: T, index: number, debug?: boolean) => string[];
 
 /**
  * Abstract representation of the document model.
  */
-export interface BlockDocument {
+export interface ChunkDocument {
   lineCount(): number;
   replaceLines(from: number, remove: number, insert: string[]): void;
 }
@@ -27,7 +27,7 @@ export interface BlockDocument {
 /**
  * Codemirror document adapter.
  */
-export class DocumentAdapter implements BlockDocument {
+export class DocumentAdapter implements ChunkDocument {
   constructor(private readonly _view: EditorView) {}
 
   /** The document must always have at least one line. */
@@ -54,16 +54,16 @@ export class DocumentAdapter implements BlockDocument {
   }
 }
 
-type BlockChange<T extends Block> =
+type ChunkChange<T extends Chunk> =
   | {
       type: 'append';
       index: number;
-      block: T;
+      chunk: T;
     }
   | {
       type: 'update';
       index: number;
-      block: T;
+      chunk: T;
     }
   | {
       type: 'delete';
@@ -72,98 +72,98 @@ type BlockChange<T extends Block> =
     };
 
 /**
- * Each Block can be converted into a set of lines.
- * Blocks can be added, updated or deleted from the Queue.
+ * Each Chunk can be converted into a set of lines.
+ * Chunks can be added, updated or deleted from the Queue.
  * A Document is an immutable projection that represents a flat list of lines corresponding to the Queue.
  *
  * Ideally we would implement a custom virtual Text model for the View, but this currently isn't possible in Codemirror.
- * Instead this model tracks changes and syncs them with the BlockDocument.
+ * Instead this model tracks changes and syncs them with the ChunkDocument.
  */
-export class BlockModel<T extends Block> {
-  /** Ordered set of blocks. */
-  private readonly _blocks: T[] = [];
+export class SerializationModel<T extends Chunk> {
+  /** Ordered set of chunks. */
+  private readonly _chunks: T[] = [];
 
-  /** Map of block IDs to their current line counts. */
-  private readonly _blockLineCounts = new Map<string, number>();
+  /** Map of chunk IDs to their current line counts. */
+  private readonly _chunkLineCounts = new Map<string, number>();
 
-  /** Map of line numbers to block IDs. */
-  private readonly _lineToBlock = new Map<number, T>();
+  /** Map of line numbers to chunk IDs. */
+  private readonly _lineToChunk = new Map<number, T>();
 
-  /** Track block changes since last sync. */
-  private readonly _changes: Array<BlockChange<T>> = [];
+  /** Track chunk changes since last sync. */
+  private readonly _changes: Array<ChunkChange<T>> = [];
 
   /** Emits when the document is updated. */
   public readonly update = new Event<void>();
 
   constructor(
-    private readonly _renderer: BlockRenderer<T>,
-    initialBlocks: T[] = [],
+    private readonly _renderer: ChunkRenderer<T>,
+    initialChunks: T[] = [],
   ) {
-    initialBlocks.forEach((block) => {
-      this.appendBlock(block);
+    initialChunks.forEach((chunk) => {
+      this.appendChunk(chunk);
     });
   }
 
   toJSON() {
     return {
-      blocks: this._blocks.length,
+      chunks: this._chunks.length,
       changes: this._changes.length,
     };
   }
 
   get doc() {
-    return Text.of([...this._blocks.flatMap((block) => this._renderer(block)), '']);
+    return Text.of([...this._chunks.flatMap((chunk, index) => this._renderer(chunk, index)), '']);
   }
 
-  get blocks() {
-    return this._blocks;
+  get chunks() {
+    return this._chunks;
   }
 
   /**
    * Get the block ID at the given line number.
    */
-  getBlockAtLine(line: number): T | undefined {
-    return this._lineToBlock.get(line);
+  getChunkAtLine(line: number): T | undefined {
+    return this._lineToChunk.get(line);
   }
 
   reset(): this {
-    this._blocks.length = 0;
-    this._blockLineCounts.clear();
-    this._lineToBlock.clear();
+    this._chunks.length = 0;
+    this._chunkLineCounts.clear();
+    this._lineToChunk.clear();
     this.update.emit();
     return this;
   }
 
-  appendBlock(block: T): this {
-    const index = this._blocks.length;
-    this._blocks.push(block);
-    this._changes.push({ type: 'append', index, block });
+  appendChunk(chunk: T): this {
+    const index = this._chunks.length;
+    this._chunks.push(chunk);
+    this._changes.push({ type: 'append', index, chunk });
     this.update.emit();
     return this;
   }
 
-  updateBlock(block: T): this {
-    const index = this._blocks.findIndex((b) => b.id === block.id);
+  updateChunk(chunk: T): this {
+    const index = this._chunks.findIndex((c) => c.id === chunk.id);
     invariant(index !== -1);
-    this._blocks[index] = block;
-    this._changes.push({ type: 'update', index, block });
+    this._chunks[index] = chunk;
+    this._changes.push({ type: 'update', index, chunk });
     this.update.emit();
     return this;
   }
 
   deleteBlock(id: string): this {
-    const index = this._blocks.findIndex((b) => b.id === id);
+    const index = this._chunks.findIndex((c) => c.id === id);
     invariant(index !== -1);
-    this._blocks.splice(index, 1);
+    this._chunks.splice(index, 1);
     this._changes.push({ type: 'delete', index, id });
     this.update.emit();
     return this;
   }
 
   /**
-   * Syncs the tracked changes with the BlockDocument.
+   * Syncs the tracked changes with the ChunkDocument.
    */
-  sync(document: BlockDocument): this {
+  sync(document: ChunkDocument): this {
     log('sync', { changes: this._changes });
 
     // Process each change in order.
@@ -171,7 +171,7 @@ export class BlockModel<T extends Block> {
       switch (change.type) {
         case 'append': {
           // For appends, convert block to lines and add at the end.
-          const lines = this._renderer(change.block);
+          const lines = this._renderer(change.chunk, change.index);
           const lastLine = document.lineCount();
           document.replaceLines(lastLine + 1, 0, lines);
           break;
@@ -181,10 +181,10 @@ export class BlockModel<T extends Block> {
           // For updates, find the block's line range and replace with new lines.
           let lineStart = 1;
           for (let i = 0; i < change.index; i++) {
-            lineStart += this._blockLineCounts.get(this._blocks[i].id)!;
+            lineStart += this._chunkLineCounts.get(this._chunks[i].id)!;
           }
-          const previousLineCount = this._blockLineCounts.get(change.block.id)!;
-          const newLines = this._renderer(change.block);
+          const previousLineCount = this._chunkLineCounts.get(change.chunk.id)!;
+          const newLines = this._renderer(change.chunk, change.index);
           document.replaceLines(lineStart, previousLineCount, newLines);
           break;
         }
@@ -193,9 +193,9 @@ export class BlockModel<T extends Block> {
           // For deletes, find the block's line range and remove those lines.
           let lineStart = 1;
           for (let i = 0; i < change.index; i++) {
-            lineStart += this._blockLineCounts.get(this._blocks[i].id)!;
+            lineStart += this._chunkLineCounts.get(this._chunks[i].id)!;
           }
-          const previousLineCount = this._blockLineCounts.get(change.id)!;
+          const previousLineCount = this._chunkLineCounts.get(change.id)!;
           document.replaceLines(lineStart, previousLineCount, []);
           break;
         }
@@ -203,22 +203,22 @@ export class BlockModel<T extends Block> {
     }
 
     // Clear the line-to-block mapping.
-    this._lineToBlock.clear();
+    this._lineToChunk.clear();
 
     // Update all block line counts and line-to-block mapping after processing changes.
     let currentLine = 1;
-    for (const block of this._blocks) {
-      const lineCount = this._renderer(block).length;
-      this._blockLineCounts.set(block.id, lineCount);
-      this._lineToBlock.set(currentLine, block);
+    for (const [index, chunk] of this._chunks.entries()) {
+      const lineCount = this._renderer(chunk, index).length;
+      this._chunkLineCounts.set(chunk.id, lineCount);
+      this._lineToChunk.set(currentLine, chunk);
       currentLine += lineCount;
     }
 
     // Clean up deleted blocks.
-    const blockIds = new Set(this._blocks.map((block) => block.id));
-    for (const id of this._blockLineCounts.keys()) {
-      if (!blockIds.has(id)) {
-        this._blockLineCounts.delete(id);
+    const chunkIds = new Set(this._chunks.map((chunk) => chunk.id));
+    for (const id of this._chunkLineCounts.keys()) {
+      if (!chunkIds.has(id)) {
+        this._chunkLineCounts.delete(id);
       }
     }
 

package/src/testing/index.ts

@@ -0,0 +1,5 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './testing';

package/src/testing/testing.ts

@@ -0,0 +1,221 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Schema, SchemaAST } from 'effect';
+import { useEffect, useMemo, useState } from 'react';
+
+import { AIServiceEdgeClient } from '@dxos/assistant';
+import { AI_SERVICE_ENDPOINT } from '@dxos/assistant/testing';
+import { scheduleTaskInterval } from '@dxos/async';
+import { createQueueDxn, Filter, type Queue } from '@dxos/client/echo';
+import { Context } from '@dxos/context';
+import { Type } from '@dxos/echo';
+import { create, ObjectId } from '@dxos/echo-schema';
+import { IdentityDid } from '@dxos/keys';
+import { log } from '@dxos/log';
+import { faker } from '@dxos/random';
+import { live, makeRef, useQueue, type Space } from '@dxos/react-client/echo';
+import { DataType } from '@dxos/schema';
+import { Testing, seedTestData } from '@dxos/schema/testing';
+
+import { processTranscriptMessage } from '../entity-extraction';
+
+// TODO(burdon): Reconcile with plugin-markdown. Move to @dxos/schema/testing.
+
+export const TestItem = Schema.Struct({
+  title: Schema.String.annotations({
+    [SchemaAST.TitleAnnotationId]: 'Title',
+    [SchemaAST.DescriptionAnnotationId]: 'Product title',
+  }),
+  description: Schema.String.annotations({
+    [SchemaAST.TitleAnnotationId]: 'Description',
+    [SchemaAST.DescriptionAnnotationId]: 'Product description',
+  }),
+}).pipe(
+  Type.def({
+    typename: 'dxos.org/type/Test',
+    version: '0.1.0',
+  }),
+);
+
+// TODO(wittjosiah): Make builder generic and reuse for all message types.
+abstract class AbstractMessageBuilder {
+  abstract createMessage(numSegments?: number): Promise<DataType.Message>;
+}
+
+/**
+ * Generator of transcript messages.
+ */
+export class MessageBuilder extends AbstractMessageBuilder {
+  static readonly singleton = new MessageBuilder();
+
+  users = Array.from({ length: 5 }, () => ({
+    identityDid: IdentityDid.random().toString(),
+    name: faker.person.fullName(),
+    email: faker.internet.email(),
+  }));
+
+  start = new Date(Date.now() - 24 * 60 * 60 * 10_000);
+
+  constructor(private readonly _space?: Space) {
+    super();
+  }
+
+  override async createMessage(numSegments = 1): Promise<DataType.Message> {
+    return {
+      id: ObjectId.random().toString(),
+      created: this.next().toISOString(),
+      sender: faker.helpers.arrayElement(this.users),
+      blocks: Array.from({ length: numSegments }).map(() => this.createBlock()),
+    };
+  }
+
+  createBlock(): DataType.MessageBlock.Transcription {
+    let text = faker.lorem.paragraph();
+    if (this._space) {
+      const label = faker.commerce.productName();
+      const obj = this._space.db.add(live(TestItem, { title: label, description: faker.lorem.paragraph() }));
+      const dxn = makeRef(obj).dxn.toString();
+      const words = text.split(' ');
+      words.splice(Math.floor(Math.random() * words.length), 0, `[${label}][${dxn}]`);
+      text = words.join(' ');
+    }
+
+    return {
+      type: 'transcription',
+      started: this.next().toISOString(),
+      text,
+    };
+  }
+
+  next() {
+    this.start = new Date(this.start.getTime() + Math.random() * 10_000);
+    return this.start;
+  }
+}
+
+// TODO(burdon): Reconcile with BlockBuilder.
+class EntityExtractionMessageBuilder extends AbstractMessageBuilder {
+  aiService = new AIServiceEdgeClient({
+    endpoint: AI_SERVICE_ENDPOINT.REMOTE,
+  });
+
+  space: Space | undefined;
+  currentMessage: number = 0;
+  transcriptMessages: DataType.Message[] = [];
+
+  async connect(space: Space) {
+    this.space = space;
+    const { transcriptMessages } = await seedTestData(space);
+    this.transcriptMessages = transcriptMessages;
+  }
+
+  override async createMessage(): Promise<DataType.Message> {
+    if (!this.space) {
+      throw new Error('Space not connected');
+    }
+
+    const { objects } = await this.space.db
+      .query(
+        Filter.or(
+          Filter.schema(DataType.Person),
+          Filter.schema(DataType.Organization),
+          Filter.schema(Testing.DocumentType),
+        ),
+      )
+      .run();
+
+    log.info('context', { objects });
+    const message = this.transcriptMessages[this.currentMessage];
+    this.currentMessage++;
+    this.currentMessage = this.currentMessage % this.transcriptMessages.length;
+
+    const { message: enhancedMessage } = await processTranscriptMessage({
+      message,
+      aiService: this.aiService,
+      context: {
+        objects,
+      },
+    });
+
+    return enhancedMessage;
+  }
+}
+
+type UseTestTranscriptionQueue = (
+  space: Space | undefined,
+  queueId?: ObjectId,
+  running?: boolean,
+  interval?: number,
+) => Queue<DataType.Message> | undefined;
+
+/**
+ * Test transcriptionqueue.
+ */
+export const useTestTranscriptionQueue: UseTestTranscriptionQueue = (
+  space: Space | undefined,
+  queueId?: ObjectId,
+  running = true,
+  interval = 1_000,
+) => {
+  const queueDxn = useMemo(() => (space ? createQueueDxn(space.id, queueId) : undefined), [space, queueId]);
+  const queue = useQueue<DataType.Message>(queueDxn);
+  const builder = useMemo(() => new MessageBuilder(space), [space]);
+
+  useEffect(() => {
+    if (!queue || !running) {
+      return;
+    }
+
+    const i = setInterval(() => {
+      void builder.createMessage(Math.ceil(Math.random() * 3)).then((message) => {
+        queue.append([create(DataType.Message, message)]);
+      });
+    }, interval);
+    return () => clearInterval(i);
+  }, [queue, running, interval]);
+
+  return queue;
+};
+
+/**
+ * Test transcription queue.
+ */
+// TODO(burdon): Reconcile with useTestTranscriptionQueue.
+export const useTestTranscriptionQueueWithEntityExtraction: UseTestTranscriptionQueue = (
+  space: Space | undefined,
+  queueId?: ObjectId,
+  running = true,
+  interval = 1_000,
+) => {
+  const queueDxn = useMemo(() => (space ? createQueueDxn(space.id, queueId) : undefined), [space, queueId]);
+  const queue = useQueue<DataType.Message>(queueDxn);
+  const [builder] = useState(() => new EntityExtractionMessageBuilder());
+
+  useEffect(() => {
+    if (!queue || !running) {
+      return;
+    }
+
+    if (space) {
+      void builder.connect(space);
+    }
+
+    const ctx = new Context();
+    scheduleTaskInterval(
+      ctx,
+      async () => {
+        const message = await builder.createMessage();
+        queue.append([create(DataType.Message, message)]);
+      },
+      interval,
+    );
+
+    return () => {
+      void ctx.dispose();
+    };
+  }, [space, queue, running, interval]);
+
+  return queue;
+};

package/src/transcriber/transcriber.browser.test.ts

@@ -9,11 +9,11 @@ import { WaveFile } from 'wavefile';
 
 import { Trigger } from '@dxos/async';
 import { log } from '@dxos/log';
+import { type DataType } from '@dxos/schema';
 import { openAndClose } from '@dxos/test-utils';
 import { trace, TRACE_PROCESSOR } from '@dxos/tracing';
 
 import { type AudioChunk, type AudioRecorder, Transcriber } from '../transcriber';
-import { type TranscriptSegment } from '../types';
 import { mergeFloat64Arrays } from '../util';
 
 // TODO(burdon): ReferenceError: Worker is not defined (Can only run with web workers).
@@ -93,7 +93,7 @@ describe.skip('Transcriber', () => {
   });
 
   test.skip('transcription of audio recording', { timeout: 10_000 }, async () => {
-    const trigger = new Trigger<TranscriptSegment[]>({ autoReset: true });
+    const trigger = new Trigger<DataType.MessageBlock.Transcription[]>({ autoReset: true });
     const recorder = new MockAudioRecorder({
       buffer: await readFile('test.wav'),
       chunkDuration: 3_000,
@@ -127,7 +127,7 @@ describe.skip('Transcriber', () => {
   });
 
   test.skip('transcription of audio recording with overlapping chunks', { timeout: 20_000 }, async () => {
-    const trigger = new Trigger<TranscriptSegment[]>({ autoReset: true });
+    const trigger = new Trigger<DataType.MessageBlock.Transcription[]>({ autoReset: true });
     const recorder = new MockAudioRecorder({
       buffer: await readFile('test.wav'),
       chunkDuration: 3_000,

package/src/transcriber/transcriber.ts

@@ -7,10 +7,11 @@ import { WaveFile } from 'wavefile';
 import { DeferredTask, Trigger } from '@dxos/async';
 import { type Context, LifecycleState, Resource } from '@dxos/context';
 import { log } from '@dxos/log';
+import { type DataType } from '@dxos/schema';
 import { trace } from '@dxos/tracing';
 
 import { type AudioRecorder, type AudioChunk } from './audio-recorder';
-import { TRANSCRIPTION_URL, type TranscriptSegment } from '../types';
+import { TRANSCRIPTION_URL } from '../types';
 import { mergeFloat64Arrays } from '../util';
 
 type WhisperWord = {
@@ -64,7 +65,7 @@ export type TranscribeConfig = {
 export type TranscriberParams = {
   config: TranscribeConfig;
   recorder: AudioRecorder;
-  onSegments: (segments: TranscriptSegment[]) => Promise<void>;
+  onSegments: (segments: DataType.MessageBlock.Transcription[]) => Promise<void>;
 };
 
 /**
@@ -93,6 +94,7 @@ export class Transcriber extends Resource {
   }
 
   protected override async _open(ctx: Context) {
+    log.info('opening');
     this._recorder.setOnChunk((chunk) => this._saveAudioChunk(chunk));
     await this._recorder.start();
     this._transcribeTask = new DeferredTask(ctx, async () => this._transcribe());
@@ -100,15 +102,18 @@ export class Transcriber extends Resource {
   }
 
   protected override async _close() {
+    log.info('closing');
     this._recording = false;
     this._transcribeTask = undefined;
     await this._recorder.stop();
   }
 
   startChunksRecording() {
+    log.info('starting');
     if (this._lifecycleState !== LifecycleState.OPEN) {
       return;
     }
+
     this._recording = true;
   }
 
@@ -116,8 +121,10 @@ export class Transcriber extends Resource {
     if (this._lifecycleState !== LifecycleState.OPEN || !this._recording) {
       return;
     }
+
     this._recording = false;
     this._transcribeTask?.schedule();
+    log.info('stopped');
   }
 
   private _saveAudioChunk(chunk: AudioChunk) {
@@ -194,24 +201,25 @@ export class Transcriber extends Resource {
 
     if (!response.ok) {
       this._audioChunks = [];
-      throw new Error('Failed to transcribe');
+      throw new Error(`Transcription failed: ${response.statusText}`);
     }
 
-    const {
-      segments,
-    }: {
+    const { segments } = (await response.json()) as {
       segments: WhisperSegment[];
-    } = await response.json();
+    };
 
     log.info('transcription response', {
-      segments,
+      segments: segments.length,
       string: segments.map((segments) => segments.text).join(' '),
     });
 
     return segments;
   }
 
-  private _alignSegments(segments: WhisperSegment[], originalChunks: AudioChunk[]): TranscriptSegment[] {
+  private _alignSegments(
+    segments: WhisperSegment[],
+    originalChunks: AudioChunk[],
+  ): DataType.MessageBlock.Transcription[] {
     // Absolute zero for all relative timestamps in the segments.
     const zeroTimestamp = originalChunks.at(0)!.timestamp;
 
@@ -240,8 +248,9 @@ export class Transcriber extends Resource {
 
     // Add absolute timestamp to each segment.
     return filteredSegments.map((segment) => ({
+      type: 'transcription',
       started: new Date(zeroTimestamp + segment.start * 1_000).toISOString(),
-      text: segment.text,
+      text: segment.text.trim(),
     }));
   }
 }