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

package/src/entity-extraction/entity-extraction.ts

@@ -7,23 +7,39 @@ import { Schema } from 'effect';
 import { createTemplate, Message, structuredOutputParser } from '@dxos/artifact';
 import type { AIServiceClient } from '@dxos/assistant';
 import { MixedStreamParser } from '@dxos/assistant';
-import { raise } from '@dxos/debug';
-import { create } from '@dxos/echo-schema';
-import type { BaseEchoObject } from '@dxos/echo-schema';
+import { asyncTimeout } from '@dxos/async';
+import { type BaseEchoObject, create, ObjectId } from '@dxos/echo-schema';
+import { DXN } from '@dxos/keys';
+import { log } from '@dxos/log';
+import { DataType } from '@dxos/schema';
 
 import SYSTEM_PROMPT from './system-prompt.tpl?raw';
-import type { TranscriptBlock } from '../types';
 
-type ProcessTranscriptBlockParams = {
-  block: TranscriptBlock;
+type ProcessTranscriptMessageParams = {
+  message: DataType.Message;
   aiService: AIServiceClient;
   context: {
     objects?: BaseEchoObject[];
   };
+
+  options?: {
+    /**
+     * Timeout for the entity extraction.
+     */
+    timeout?: number;
+
+    /**
+     * Fallback to raw text if the entity extraction fails.
+     * Otherwise the function will throw an error.
+     * @default false
+     */
+    fallbackToRaw?: boolean;
+  };
 };
 
-type ProcessTranscriptBlockResult = {
-  block: TranscriptBlock;
+type ProcessTranscriptMessageResult = {
+  message: DataType.Message;
+  timeElapsed: number;
 };
 
 /**
@@ -34,61 +50,100 @@ const createSystemPrompt = (): string => {
   return template({});
 };
 
-export const processTranscriptBlock = async (
-  params: ProcessTranscriptBlockParams,
-): Promise<ProcessTranscriptBlockResult> => {
-  // TODO(dmaretskyi): Move context to a vector search index.
-  const systemPrompt = `
-    ${createSystemPrompt()}
-    Context:
-    ${JSON.stringify(params.context.objects)}
-  `;
-
-  const outputParser = structuredOutputParser(
+const ReferencedQuotes = Schema.Struct({
+  references: Schema.Array(
     Schema.Struct({
-      segments: Schema.Array(Schema.String).annotations({
-        description: 'The enhanced text of the transcript segments, keep the order and structure exactly as is',
-      }),
+      quote: Schema.String,
+      id: Schema.String,
     }),
-  );
-  const result = outputParser.getResult(
-    await new MixedStreamParser().parse(
-      await params.aiService.execStream({
-        model: '@anthropic/claude-3-5-haiku-20241022',
-        systemPrompt,
-        history: [
-          create(Message, {
-            role: 'user',
-            content: [
-              {
-                type: 'text',
-                text: JSON.stringify(params.block),
-              },
+  ).annotations({
+    description:
+      'The references to the context objects that are mentioned in the transcript. quote should match the original transcript text exactly, while id is the id of the context object.',
+  }),
+});
+interface ReferencedQuotes extends Schema.Schema.Type<typeof ReferencedQuotes> {}
+
+// TODO(dmaretskyi): Move context to a vector search index.
+export const processTranscriptMessage = async (
+  params: ProcessTranscriptMessageParams,
+): Promise<ProcessTranscriptMessageResult> => {
+  try {
+    const outputParser = structuredOutputParser(ReferencedQuotes);
+    const runParser = async (): Promise<ProcessTranscriptMessageResult> => {
+      const startTime = performance.now();
+      const result = outputParser.getResult(
+        await new MixedStreamParser().parse(
+          await params.aiService.execStream({
+            model: '@anthropic/claude-3-5-haiku-20241022',
+            systemPrompt: createSystemPrompt(),
+            history: [
+              create(Message, {
+                role: 'user',
+                content: [
+                  {
+                    type: 'text',
+                    text: `<context>${JSON.stringify(params.context.objects)}</context>`,
+                  },
+                  {
+                    type: 'text',
+                    text: `<transcript>${JSON.stringify(params.message.blocks)}</transcript>`,
+                  },
+                ],
+              }),
             ],
+            tools: [outputParser.tool],
           }),
-        ],
-        tools: [outputParser.tool],
-      }),
-    ),
-  );
+        ),
+      );
+      log.info('entity extraction result', { refs: result.references });
 
-  return {
-    block: {
-      ...params.block,
-      segments: params.block.segments.map((segment, i) => ({
-        ...segment,
-        text: postprocessText(result?.segments[i] ?? raise(new Error('failed to process email'))),
-      })),
-    },
-  };
+      return {
+        message: create(DataType.Message, {
+          ...params.message,
+          blocks: params.message.blocks.map((block, i) =>
+            block.type !== 'transcription'
+              ? block
+              : {
+                  ...block,
+                  text: postprocessText(block.text, result),
+                },
+          ),
+        }),
+        timeElapsed: performance.now() - startTime,
+      };
+    };
+
+    if (params.options?.timeout && params.options.timeout > 0) {
+      return await asyncTimeout(runParser(), params.options.timeout);
+    } else {
+      return await runParser();
+    }
+  } catch (error) {
+    if (params.options?.fallbackToRaw) {
+      log.warn('failed to process transcript message, falling back to raw text', { error });
+      return {
+        message: params.message,
+        timeElapsed: 0,
+      };
+    } else {
+      throw error;
+    }
+  }
 };
 
 /**
- * Finds and replaces all inline references with DXNs references.
+ * Finds and replaces all quotes with DXNs references.
  */
 // TODO(dmaretskyi): Lookup and verifiy ids from provided context.
-const postprocessText = (text: string) => {
-  return text.replace(/\n/g, ' ').replace(/\[([^\]]+)\]\[([A-Z0-9]+)\]/g, (match, name, id) => {
-    return `[${name}][dxn:echo:@:${id}]`;
-  });
+export const postprocessText = (text: string, quotes: ReferencedQuotes) => {
+  for (const quote of quotes.references) {
+    if (!ObjectId.isValid(quote.id)) {
+      continue;
+    }
+
+    // Use a case-insensitive regular expression to replace the quote
+    const regex = new RegExp(quote.quote.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'gi');
+    text = text.replace(regex, `[${quote.quote}][${DXN.fromLocalObjectId(quote.id).toString()}]`);
+  }
+  return text;
 };

package/src/entity-extraction/system-prompt.tpl

@@ -1,9 +1,38 @@
-  You are a helpful assistant that processes transcripts.
-  Your goal is to enhance the transcript with entity references.
-  Insert inline references to the entities in the transcript.
-  Keep the transcript structure and text as is.
-  Call the submit_result tool to submit your result.
-
-  The inline reference syntax is as follows:
-   - [<optional name>][<ID>] or just [<ID>]
-   - Example: [Earnings Report][01JT0JP9AX0XKGZX4MV4B69VT6]
+You are an AI assistant specialized in processing and enhancing transcripts with contextual information.
+Your task is to analyze a given transcript and extract references to the context objects.
+
+The context and transcript are provided in the user message in `context` and `transcript` XML tags.
+
+Here are your instructions:
+
+1. Read the transcript and context information provided above.
+
+2. Process each segment of the transcript separately.
+
+3. For each segment:
+   a. Analyze the content and identify relevant information from the context.
+   b. Extract a number of references to the context objects that are mentioned in the segment.
+   c. Output the references in the format `{ "quote": "earnings report", "id": "01JT0JP9AX0XKGZX4MV4B69VT6" }`.
+   d. If there are no references, just output an empty array.
+
+4. Maintain the exact order and structure of the original transcript segments.
+
+Output Format:
+The final output should be a JSON array of references.
+
+Example output structure:
+```json
+{
+  "references": [
+    { "quote": "earnings report", "id": "01JT0JP9AX0XKGZX4MV4B69VT6" },
+    { "quote": "John", "id": "01JT0VM0VN3ZKN7GDBM8CV9BA7" }
+  ]
+}
+```
+
+Remember:
+- Only quote the text that is in the transcript.
+- Analyze each segment carefully before processing it.
+- Only reference objects that appear in the context and exactly by their id, If the object is not in the context, do not insert a reference. 
+
+Now, process the transcript and context provided below according to these instructions.

package/src/hooks/useQueueModelAdapter.ts

@@ -2,29 +2,62 @@
 // Copyright 2025 DXOS.org
 //
 
-import { useMemo, useEffect } from 'react';
+import { useMemo, useEffect, useState } from 'react';
 
 import { type Queue } from '@dxos/react-client/echo';
 
-import { type Block, BlockModel, type BlockRenderer } from '../model';
+import { type Chunk, type ChunkRenderer, SerializationModel } from '../model';
 
 /**
  * Model adapter for a queue.
  */
-export const useQueueModelAdapter = <T extends Block>(
-  renderer: BlockRenderer<T>,
+export const useQueueModelAdapter = <T extends Chunk>(
+  renderer: ChunkRenderer<T>,
   queue: Queue<T> | undefined,
-  initialBlocks: T[] = [],
-): BlockModel<T> => {
-  const model = useMemo(() => new BlockModel<T>(renderer, initialBlocks), [queue]);
+  initialChunks: T[] = [],
+): SerializationModel<T> => {
+  const model = useMemo(() => new SerializationModel<T>(renderer, initialChunks), [queue]);
+  const [loaded, setLoaded] = useState(false);
+
+  // Set initial blocks.
+  useEffect(() => {
+    if (!queue) {
+      return;
+    }
+
+    const update = () => {
+      for (const block of queue?.items ?? []) {
+        model.appendChunk(block);
+      }
+
+      setLoaded(true);
+    };
+
+    // TODO(burdon): This is a hack to ensure the queue is loaded.
+    let i: NodeJS.Timeout | undefined;
+    if (queue.isLoading) {
+      i = setInterval(() => {
+        if (!queue?.isLoading) {
+          clearInterval(i);
+          update();
+        }
+      }, 1_000);
+    } else {
+      update();
+    }
+
+    return () => clearTimeout(i);
+  }, [model, queue]);
+
+  // TODO(burdon): Can we listen for queue events?
   useEffect(() => {
-    if (!queue?.items.length) {
+    if (!loaded || !queue || model.chunks.length === queue.items.length) {
       return;
     }
 
-    const block = queue.items[queue.items.length - 1];
-    model.appendBlock(block);
-  }, [model, queue?.items.length]);
+    const chunk = queue.items[queue.items.length - 1];
+    model.appendChunk(chunk);
+  }, [model, loaded, queue?.items.length]);
 
   return model;
 };

package/src/model/model.test.ts

@@ -5,71 +5,101 @@
 import { EditorView } from '@codemirror/view';
 import { describe, test } from 'vitest';
 
-import { BlockModel, DocumentAdapter, type BlockRenderer } from './model';
-import { type TranscriptBlock } from '../types';
-
-const blockToMarkdown: BlockRenderer<TranscriptBlock> = (block: TranscriptBlock, debug = true): string[] => {
-  return [`###### ${block.authorName}`, ...block.segments.map((segment) => segment.text), ''];
+import { type DataType } from '@dxos/schema';
+
+import { SerializationModel, DocumentAdapter, type ChunkRenderer } from './model';
+
+const blockToMarkdown: ChunkRenderer<DataType.Message> = (
+  message: DataType.Message,
+  index: number,
+  debug = true,
+): string[] => {
+  return [
+    `###### ${message.sender.name}`,
+    ...message.blocks.filter((block) => block.type === 'transcription').map((block) => block.text),
+    '',
+  ];
 };
 
 const createDate = () => new Date().toISOString();
 
-describe('BlockModel', () => {
+describe('SerializationModel', () => {
   test('basic', ({ expect }) => {
-    const model = new BlockModel<TranscriptBlock>(blockToMarkdown);
-    expect(model.blocks.length).to.eq(0);
+    const model = new SerializationModel<DataType.Message>(blockToMarkdown);
+    expect(model.chunks.length).to.eq(0);
     expect(model.doc.toString()).to.eq('');
 
-    // Create block.
-    const block = {
+    // Create message.
+    const message: DataType.Message = {
       id: '1',
-      authorName: 'Alice',
-      segments: [{ started: createDate(), text: 'Hello world!' }],
+      created: createDate(),
+      sender: { name: 'Alice' },
+      blocks: [
+        {
+          type: 'transcription',
+          started: createDate(),
+          text: 'Hello world!',
+        },
+      ],
     };
-    model.appendBlock(block);
+    model.appendChunk(message);
     {
       const text = model.doc.toString();
       expect(text).to.eq('###### Alice\nHello world!\n\n');
-      expect(model.blocks.length).to.eq(1);
+      expect(model.chunks.length).to.eq(1);
     }
 
-    // Update block.
-    block.segments.push({ started: new Date().toISOString(), text: 'Hello again!' });
-    model.updateBlock(block);
+    // Update message.
+    message.blocks.push({ type: 'transcription', started: createDate(), text: 'Hello again!' });
+    model.updateChunk(message);
     {
       const text = model.doc.toString();
       expect(text).to.eq('###### Alice\nHello world!\nHello again!\n\n');
-      expect(model.blocks.length).to.eq(1);
+      expect(model.chunks.length).to.eq(1);
     }
   });
 
   test('sync - append', ({ expect }) => {
     const view = new EditorView({ extensions: [], doc: '' });
-    const model = new BlockModel<TranscriptBlock>(blockToMarkdown);
+    const model = new SerializationModel<DataType.Message>(blockToMarkdown);
     const adapter = new DocumentAdapter(view);
     expect(adapter.lineCount()).to.eq(0);
     expect(view.state.doc.toString()).to.eq('');
 
-    // Append block.
+    // Append message.
     {
-      const block = {
+      const message: DataType.Message = {
         id: '1',
-        authorName: 'Alice',
-        segments: [{ started: createDate(), text: 'Hello world!' }],
+        created: createDate(),
+        sender: { name: 'Alice' },
+        blocks: [
+          {
+            type: 'transcription',
+            started: createDate(),
+            text: 'Hello world!',
+          },
+        ],
       };
-      model.appendBlock(block);
+      model.appendChunk(message);
       model.sync(adapter);
       expect(view.state.doc.toString()).to.eq('###### Alice\nHello world!\n\n');
     }
 
-    // Append block.
+    // Append message.
     {
-      const block = {
+      const message: DataType.Message = {
         id: '2',
-        authorName: 'Bob',
-        segments: [{ started: createDate(), text: 'Hello world!' }],
+        created: createDate(),
+        sender: { name: 'Bob' },
+        blocks: [
+          {
+            type: 'transcription',
+            started: createDate(),
+            text: 'Hello world!',
+          },
+        ],
       };
-      model.appendBlock(block);
+      model.appendChunk(message);
       model.sync(adapter);
       expect(view.state.doc.toString()).to.eq('###### Alice\nHello world!\n\n###### Bob\nHello world!\n\n');
     }
@@ -77,40 +107,56 @@ describe('BlockModel', () => {
 
   test('sync - append, update, delete', ({ expect }) => {
     const view = new EditorView({ extensions: [], doc: '' });
-    const model = new BlockModel<TranscriptBlock>(blockToMarkdown);
+    const model = new SerializationModel<DataType.Message>(blockToMarkdown);
     const adapter = new DocumentAdapter(view);
     expect(adapter.lineCount()).to.eq(0);
 
-    // Append block.
+    // Append message.
     {
-      const block = { id: '1', authorName: 'Alice', segments: [{ started: createDate(), text: 'Hello world!' }] };
-      model.appendBlock(block);
+      const message: DataType.Message = {
+        id: '1',
+        created: createDate(),
+        sender: { name: 'Alice' },
+        blocks: [
+          {
+            type: 'transcription',
+            started: createDate(),
+            text: 'Hello world!',
+          },
+        ],
+      };
+      model.appendChunk(message);
       model.sync(adapter);
       expect(view.state.doc.toString()).to.eq('###### Alice\nHello world!\n\n');
 
-      // Update block (add segment).
-      block.segments.push({ started: createDate(), text: 'Hello again!' });
-      model.updateBlock(block);
+      // Update message (add block).
+      message.blocks.push({ type: 'transcription', started: createDate(), text: 'Hello again!' });
+      model.updateChunk(message);
       model.sync(adapter);
       expect(view.state.doc.toString()).to.eq('###### Alice\nHello world!\nHello again!\n\n');
 
-      // Update block (remove segment).
-      block.segments.splice(0, 1);
-      model.updateBlock(block);
+      // Update message (remove block).
+      message.blocks.splice(0, 1);
+      model.updateChunk(message);
       model.sync(adapter);
       expect(view.state.doc.toString()).to.eq('###### Alice\nHello again!\n\n');
 
-      expect(model.getBlockAtLine(1)).to.eq(block);
+      expect(model.getChunkAtLine(1)).to.eq(message);
     }
 
-    // Append block.
+    // Append message.
     {
-      const block = { id: '2', authorName: 'Bob', segments: [{ started: createDate(), text: 'Hello again!' }] };
-      model.appendBlock(block);
+      const message: DataType.Message = {
+        id: '2',
+        created: createDate(),
+        sender: { name: 'Bob' },
+        blocks: [{ type: 'transcription', started: createDate(), text: 'Hello again!' }],
+      };
+      model.appendChunk(message);
       model.sync(adapter);
       expect(view.state.doc.toString()).to.eq('###### Alice\nHello again!\n\n###### Bob\nHello again!\n\n');
 
-      expect(model.getBlockAtLine(4)).to.eq(block);
+      expect(model.getChunkAtLine(4)).to.eq(message);
     }
 
     // Delete block.