@dxos/plugin-transcription 0.8.2-main.f081794 → 0.8.2-main.fbd8ed0

package/src/components/transcription.stories.tsx

@@ -18,32 +18,40 @@ import React, {
 
 import { Events, IntentPlugin, SettingsPlugin } from '@dxos/app-framework';
 import { withPluginManager } from '@dxos/app-framework/testing';
-import { MemoryQueue } from '@dxos/echo-db';
-import { create } from '@dxos/echo-schema';
+import { AIServiceEdgeClient } from '@dxos/assistant';
+import { AI_SERVICE_ENDPOINT } from '@dxos/assistant/testing';
+import { Filter, MemoryQueue } from '@dxos/echo-db';
+import { create, createQueueDxn } from '@dxos/echo-schema';
 import { type DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { ClientPlugin } from '@dxos/plugin-client';
+import { PreviewPlugin } from '@dxos/plugin-preview';
 import { SpacePlugin } from '@dxos/plugin-space';
+import { StorybookLayoutPlugin } from '@dxos/plugin-storybook-layout';
 import { ThemePlugin } from '@dxos/plugin-theme';
-import { randomQueueDxn, useQueue } from '@dxos/react-client/echo';
+import { useQueue, useSpace } from '@dxos/react-client/echo';
 import { IconButton, Toolbar } from '@dxos/react-ui';
 import { ScrollContainer } from '@dxos/react-ui-components';
 import { defaultTx } from '@dxos/react-ui-theme';
+import { DataType } from '@dxos/schema';
+import { seedTestData, Testing } from '@dxos/schema/testing';
 import { withLayout, withTheme } from '@dxos/storybook-utils';
 
 import { renderMarkdown, Transcript } from './Transcript';
-import { TestItem } from './Transcript/testings';
 import { TranscriptionPlugin } from '../TranscriptionPlugin';
+import { processTranscriptMessage } from '../entity-extraction';
 import { useAudioFile, useAudioTrack, useQueueModelAdapter, useTranscriber } from '../hooks';
-import { type BlockModel } from '../model';
+import { type SerializationModel } from '../model';
+import { TestItem } from '../testing';
 import { type TranscriberParams } from '../transcriber';
-import { TranscriptBlock } from '../types';
 
 const TranscriptionStory: FC<{
-  model: BlockModel<TranscriptBlock>;
+  model: SerializationModel<DataType.Message>;
   running: boolean;
   onRunningChange: Dispatch<SetStateAction<boolean>>;
 }> = ({ model, running, onRunningChange }) => {
+  const space = useSpace();
+
   return (
     <div className='flex flex-col w-[30rem]'>
       <Toolbar.Root>
@@ -55,30 +63,67 @@ const TranscriptionStory: FC<{
         />
       </Toolbar.Root>
       <ScrollContainer>
-        <Transcript model={model} attendableId='story' />
+        <Transcript space={space} model={model} attendableId='story' />
       </ScrollContainer>
     </div>
   );
 };
 
-const Microphone = () => {
+const aiService = new AIServiceEdgeClient({
+  endpoint: AI_SERVICE_ENDPOINT.REMOTE,
+});
+
+const Microphone = ({ entityExtraction }: { entityExtraction?: boolean }) => {
   const [running, setRunning] = useState(false);
 
   // Audio.
   const track = useAudioTrack(running);
 
   // Queue.
-  const queueDxn = useMemo(() => randomQueueDxn(), []);
-  const queue = useMemo(() => new MemoryQueue<TranscriptBlock>(queueDxn), [queueDxn]);
-  const model = useQueueModelAdapter(renderMarkdown, queue);
+  const queueDxn = useMemo(() => createQueueDxn(), []);
+  const queue = useMemo(() => new MemoryQueue<DataType.Message>(queueDxn), [queueDxn]);
+  const model = useQueueModelAdapter(renderMarkdown([]), queue);
+  const space = useSpace();
 
   // Transcriber.
   const handleSegments = useCallback<TranscriberParams['onSegments']>(
-    async (segments) => {
-      const block = create(TranscriptBlock, { segments });
-      void queue.append([block]);
+    async (blocks) => {
+      const message = create(DataType.Message, { sender: { name: 'You' }, created: new Date().toISOString(), blocks });
+
+      if (entityExtraction) {
+        if (!space) {
+          log.warn('no space');
+          return;
+        }
+        // TODO(dmaretskyi): Move to vector search index.
+        const { objects } = await space.db
+          .query(
+            Filter.or(
+              Filter.type(DataType.Person),
+              Filter.type(DataType.Organization),
+              Filter.type(Testing.DocumentType),
+            ),
+          )
+          .run();
+
+        log.info('context', { objects });
+
+        const result = await processTranscriptMessage({
+          message,
+          aiService,
+          context: {
+            objects,
+          },
+          options: {
+            fallbackToRaw: true,
+          },
+        });
+        void queue.append([result.message]);
+      } else {
+        void queue.append([message]);
+      }
     },
-    [queue],
+    [queue, space],
   );
 
   const config = useRef<TranscriberParams['config']>({
@@ -94,10 +139,8 @@ const Microphone = () => {
   const [isOpen, setIsOpen] = useState(false);
   useEffect(() => {
     void transcriber?.open().then(() => setIsOpen(true));
-    log.info('transcriber open');
     return () => {
       void transcriber?.close().then(() => setIsOpen(false));
-      log.info('transcriber close');
     };
   }, [transcriber]);
 
@@ -105,10 +148,8 @@ const Microphone = () => {
   useEffect(() => {
     if (running && transcriber?.isOpen) {
       transcriber?.startChunksRecording();
-      log.info('transcriber startChunksRecording');
     } else if (!running) {
       transcriber?.stopChunksRecording();
-      log.info('transcriber stopChunksRecording');
     }
   }, [transcriber, running, isOpen]);
 
@@ -141,12 +182,12 @@ const AudioFile = ({ queueDxn, audioUrl }: { queueDxn: DXN; audioUrl: string; tr
   }, [audio, running]);
 
   // Transcriber.
-  const queue = useQueue<TranscriptBlock>(queueDxn, { pollInterval: 500 });
-  const model = useQueueModelAdapter(renderMarkdown, queue);
+  const queue = useQueue<DataType.Message>(queueDxn, { pollInterval: 500 });
+  const model = useQueueModelAdapter(renderMarkdown([]), queue);
   const handleSegments = useCallback<TranscriberParams['onSegments']>(
-    async (segments) => {
-      const block = create(TranscriptBlock, { authorName: 'test', authorHue: 'cyan', segments });
-      queue?.append([block]);
+    async (blocks) => {
+      const message = create(DataType.Message, { sender: { name: 'You' }, created: new Date().toISOString(), blocks });
+      void queue?.append([message]);
     },
     [queue],
   );
@@ -187,14 +228,19 @@ const meta: Meta<typeof AudioFile> = {
     withPluginManager({
       plugins: [
         ThemePlugin({ tx: defaultTx }),
+        StorybookLayoutPlugin(),
         ClientPlugin({
-          types: [TestItem],
+          types: [TestItem, DataType.Person, DataType.Organization, Testing.DocumentType],
           onClientInitialized: async (_, client) => {
             await client.halo.createIdentity();
+            await client.spaces.waitUntilReady();
+            await client.spaces.default.waitUntilReady();
+            await seedTestData(client.spaces.default);
           },
         }),
         SpacePlugin(),
         SettingsPlugin(),
+        PreviewPlugin(),
         IntentPlugin(),
         TranscriptionPlugin(),
       ],
@@ -218,26 +264,30 @@ const meta: Meta<typeof AudioFile> = {
       // ],
     }),
     withTheme,
-    withLayout({
-      tooltips: true,
-      fullscreen: true,
-      classNames: 'justify-center',
-    }),
+    withLayout({ fullscreen: true, classNames: 'justify-center' }),
   ],
 };
 
 export default meta;
 
-type Story = StoryObj<typeof AudioFile>;
+export const Default: StoryObj<typeof Microphone> = {
+  render: Microphone,
+  args: {
+    entityExtraction: false,
+  },
+};
 
-export const Default: Story = {
+export const EntityExtraction: StoryObj<typeof Microphone> = {
   render: Microphone,
+  args: {
+    entityExtraction: true,
+  },
 };
 
-export const File: Story = {
+export const File: StoryObj<typeof AudioFile> = {
   render: AudioFile,
   args: {
-    queueDxn: randomQueueDxn(),
+    queueDxn: createQueueDxn(),
     // https://learnenglish.britishcouncil.org/general-english/audio-zone/living-london
     transcriptUrl: 'https://dxos.network/audio-london.txt',
     audioUrl: 'https://dxos.network/audio-london.m4a',

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

@@ -2,31 +2,96 @@
 // Copyright 2025 DXOS.org
 //
 
-import { describe, test } from 'vitest';
+import { describe, test, expect } from 'vitest';
 
 import { AIServiceEdgeClient } from '@dxos/assistant';
 import { AI_SERVICE_ENDPOINT } from '@dxos/assistant/testing';
 import { log } from '@dxos/log';
+import { createTestData } from '@dxos/schema/testing';
+import { range } from '@dxos/util';
 
-import { processTranscriptBlock } from './entity-extraction';
-import * as TestData from '../testing/test-data';
+import { processTranscriptMessage, postprocessText } from './entity-extraction';
 
 const aiService = new AIServiceEdgeClient({
   endpoint: AI_SERVICE_ENDPOINT.REMOTE,
 });
 
 describe.skip('EntityExtraction', { timeout: 180_000 }, () => {
-  test('should process a transcript block', async () => {
-    for (const block of TestData.transcriptBlocks) {
-      log.info('input', block);
-      const { block: enhancedBlock } = await processTranscriptBlock({
-        block,
+  test.skip('should process a transcript block', async () => {
+    const { transcriptMessages, documents, contacts } = await createTestData();
+
+    log.info('context', { documents, contacts });
+
+    for (const message of transcriptMessages) {
+      log.info('input', message);
+      const { message: enhancedMessage } = await processTranscriptMessage({
+        message,
+        aiService,
+        context: {
+          objects: [...documents, ...Object.values(contacts)],
+        },
+      });
+      log.info('output', enhancedMessage);
+    }
+  });
+
+  test.only('computational irreducibility', async () => {
+    const { transcriptWoflram, documents, contacts } = await createTestData();
+
+    log.info('context', { documents, contacts });
+    const message = transcriptWoflram[0];
+    log.info('input', message);
+
+    await Promise.all(
+      range(10).map(async () => {
+        const { message: enhancedMessage, timeElapsed } = await processTranscriptMessage({
+          message,
+          aiService,
+          context: {
+            objects: [...documents, ...Object.values(contacts)],
+          },
+        });
+        log.info('output', { message: enhancedMessage.blocks[0], timeElapsed });
+      }),
+    );
+  });
+
+  test('org and document linking', async () => {
+    const { transcriptJosiah, documents, contacts, organizations } = await createTestData();
+
+    log.info('context', { contacts, organizations, documents });
+
+    for (const message of transcriptJosiah) {
+      log.info('input', message);
+
+      const { message: enhancedMessage } = await processTranscriptMessage({
+        message,
         aiService,
         context: {
-          objects: [...TestData.documents, ...Object.values(TestData.contacts)],
+          objects: [...documents, ...Object.values(contacts), ...Object.values(organizations)],
         },
       });
-      log.info('output', enhancedBlock);
+      log.info('output', enhancedMessage);
     }
   });
 });
+
+describe('postprocessText', () => {
+  test('should replace quotes with DXN references', () => {
+    const quotes = {
+      references: [{ quote: 'computational irreducibility', id: '01JTG9JW11XGWJZ32AW8ET93D1' }],
+    };
+
+    expect(postprocessText('This is a computational irreducibility test.', quotes)).toBe(
+      'This is a [computational irreducibility][dxn:echo:@:01JTG9JW11XGWJZ32AW8ET93D1] test.',
+    );
+    expect(
+      postprocessText(
+        "And what I'd like to talk today about is Steven Wolfram's concept of a computational irreducibility.",
+        quotes,
+      ),
+    ).toBe(
+      "And what I'd like to talk today about is Steven Wolfram's concept of a [computational irreducibility][dxn:echo:@:01JTG9JW11XGWJZ32AW8ET93D1].",
+    );
+  });
+});

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;
 };