@dxos/plugin-transcription 0.8.1 → 0.8.2-main.2f9c567

package/src/components/transcription.stories.tsx

@@ -16,84 +16,148 @@ import React, {
   useState,
 } from 'react';
 
-import { createStatic } from '@dxos/echo-schema';
+import { Events, IntentPlugin, SettingsPlugin } from '@dxos/app-framework';
+import { withPluginManager } from '@dxos/app-framework/testing';
+import { AIServiceEdgeClient } from '@dxos/assistant';
+import { AI_SERVICE_ENDPOINT } from '@dxos/assistant/testing';
+import { Filter, MemoryQueue } from '@dxos/echo-db';
+import { create } from '@dxos/echo-schema';
 import { type DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
-import { Config } from '@dxos/react-client';
-import { useQueue } from '@dxos/react-client/echo';
-import { withClientProvider } from '@dxos/react-client/testing';
+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 { createQueueDxn, 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 { Transcript } from './Transcript';
-import { useAudioFile, useAudioTrack, useTranscriber } from '../hooks';
+import { renderMarkdown, Transcript } from './Transcript';
+import { TranscriptionPlugin } from '../TranscriptionPlugin';
+import { processTranscriptMessage } from '../entity-extraction';
+import { useAudioFile, useAudioTrack, useQueueModelAdapter, useTranscriber } from '../hooks';
+import { type SerializationModel } from '../model';
+import { TestItem } from '../testing';
 import { type TranscriberParams } from '../transcriber';
-import { TranscriptBlock } from '../types';
-import { randomQueueDxn } from '../util';
 
 const TranscriptionStory: FC<{
-  playing: boolean;
-  setPlaying: Dispatch<SetStateAction<boolean>>;
-  blocks?: TranscriptBlock[];
-}> = ({ playing, setPlaying, blocks }) => {
+  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>
         <IconButton
           iconOnly
-          icon={playing ? 'ph--pause--regular' : 'ph--play--regular'}
-          label={playing ? 'Pause' : 'Play'}
-          onClick={() => setPlaying((playing) => !playing)}
+          icon={running ? 'ph--pause--regular' : 'ph--play--regular'}
+          label={running ? 'Pause' : 'Play'}
+          onClick={() => onRunningChange((running) => !running)}
         />
       </Toolbar.Root>
       <ScrollContainer>
-        <Transcript blocks={blocks} attendableId='story' />
+        <Transcript space={space} model={model} attendableId='story' />
       </ScrollContainer>
     </div>
   );
 };
 
-const Microphone = () => {
-  const [playing, setPlaying] = useState(false);
+const aiService = new AIServiceEdgeClient({
+  endpoint: AI_SERVICE_ENDPOINT.REMOTE,
+});
+
+const Microphone = ({ entityExtraction }: { entityExtraction?: boolean }) => {
+  const [running, setRunning] = useState(false);
 
   // Audio.
-  const track = useAudioTrack(playing);
+  const track = useAudioTrack(running);
 
   // Queue.
-  const queueDxn = useMemo(() => randomQueueDxn(), []);
-  const queue = useQueue<TranscriptBlock>(queueDxn, { pollInterval: 500 });
+  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 = createStatic(TranscriptBlock, { segments });
-      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.schema(DataType.Person),
+              Filter.schema(DataType.Organization),
+              Filter.schema(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 transcriber = useTranscriber({ audioStreamTrack: track, onSegments: handleSegments });
+
+  const config = useRef<TranscriberParams['config']>({
+    transcribeAfterChunksAmount: 20,
+    prefixBufferChunksAmount: 5,
+  });
+
+  const transcriber = useTranscriber({
+    audioStreamTrack: track,
+    onSegments: handleSegments,
+    config: config.current,
+  });
+  const [isOpen, setIsOpen] = useState(false);
   useEffect(() => {
-    void transcriber?.open();
+    void transcriber?.open().then(() => setIsOpen(true));
     return () => {
-      void transcriber?.close();
+      void transcriber?.close().then(() => setIsOpen(false));
     };
   }, [transcriber]);
 
   // Manage transcription state.
   useEffect(() => {
-    if (playing && transcriber?.isOpen) {
-      void transcriber?.startChunksRecording();
-    } else if (!playing) {
-      void transcriber?.stopChunksRecording();
+    if (running && transcriber?.isOpen) {
+      transcriber?.startChunksRecording();
+    } else if (!running) {
+      transcriber?.stopChunksRecording();
     }
-  }, [transcriber, playing, transcriber?.isOpen]);
+  }, [transcriber, running, isOpen]);
 
-  return <TranscriptionStory playing={playing} setPlaying={setPlaying} blocks={queue?.items} />;
+  return <TranscriptionStory model={model} running={running} onRunningChange={setRunning} />;
 };
 
 const AudioFile = ({ queueDxn, audioUrl }: { queueDxn: DXN; audioUrl: string; transcriptUrl: string }) => {
-  const [playing, setPlaying] = useState(false);
+  const [running, setRunning] = useState(false);
 
   // Audio.
   const audioElement = useRef<HTMLAudioElement>(null);
@@ -110,19 +174,20 @@ const AudioFile = ({ queueDxn, audioUrl }: { queueDxn: DXN; audioUrl: string; tr
       return;
     }
 
-    if (playing) {
+    if (running) {
       void audio.play();
     } else {
       void audio.pause();
     }
-  }, [audio, playing]);
+  }, [audio, running]);
 
   // Transcriber.
-  const queue = useQueue<TranscriptBlock>(queueDxn, { pollInterval: 500 });
+  const queue = useQueue<DataType.Message>(queueDxn, { pollInterval: 500 });
+  const model = useQueueModelAdapter(renderMarkdown([]), queue);
   const handleSegments = useCallback<TranscriberParams['onSegments']>(
-    async (segments) => {
-      const block = createStatic(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],
   );
@@ -133,15 +198,15 @@ const AudioFile = ({ queueDxn, audioUrl }: { queueDxn: DXN; audioUrl: string; tr
   });
 
   useEffect(() => {
-    if (transcriber && playing) {
-      void transcriber.open();
-    } else if (!playing) {
+    if (transcriber && running) {
+      void transcriber.open().then(manageChunkRecording);
+    } else if (!running) {
       transcriber?.stopChunksRecording();
       void transcriber?.close();
     }
-  }, [transcriber, playing]);
+  }, [transcriber, running]);
 
-  useEffect(() => {
+  const manageChunkRecording = () => {
     if (track?.readyState === 'live' && transcriber?.isOpen) {
       log.info('starting transcription');
       transcriber.startChunksRecording();
@@ -149,24 +214,54 @@ const AudioFile = ({ queueDxn, audioUrl }: { queueDxn: DXN; audioUrl: string; tr
       log.info('stopping transcription');
       transcriber.stopChunksRecording();
     }
+  };
+  useEffect(() => {
+    manageChunkRecording();
   }, [transcriber, track?.readyState, transcriber?.isOpen]);
 
-  return <TranscriptionStory playing={playing} setPlaying={setPlaying} blocks={queue?.items} />;
+  return <TranscriptionStory model={model} running={running} onRunningChange={setRunning} />;
 };
 
 const meta: Meta<typeof AudioFile> = {
   title: 'plugins/plugin-transcription/transcription',
   decorators: [
-    withClientProvider({
-      config: new Config({
-        runtime: {
-          client: { edgeFeatures: { signaling: true } },
-          services: {
-            edge: { url: 'https://edge.dxos.workers.dev/' },
-            iceProviders: [{ urls: 'https://edge.dxos.workers.dev/ice' }],
+    withPluginManager({
+      plugins: [
+        ThemePlugin({ tx: defaultTx }),
+        StorybookLayoutPlugin(),
+        ClientPlugin({
+          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(),
+      ],
+      fireEvents: [Events.SetupAppGraph],
+      // capabilities: [
+      //   contributes(
+      //     Capabilities.ReactSurface,
+      //     createSurface({
+      //       id: 'preview-test',
+      //       role: 'preview',
+      //       component: ({ data }) => {
+      //         const schema = getSchema(data);
+      //         if (!schema) {
+      //           return null;
+      //         }
+
+      //         return <Form schema={schema} values={data} />;
+      //       },
+      //     }),
+      //   ),
+      // ],
     }),
     withTheme,
     withLayout({
@@ -179,16 +274,24 @@ const meta: Meta<typeof AudioFile> = {
 
 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

@@ -0,0 +1,97 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+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 { processTranscriptMessage, postprocessText } from './entity-extraction';
+
+const aiService = new AIServiceEdgeClient({
+  endpoint: AI_SERVICE_ENDPOINT.REMOTE,
+});
+
+describe.skip('EntityExtraction', { timeout: 180_000 }, () => {
+  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: [...documents, ...Object.values(contacts), ...Object.values(organizations)],
+        },
+      });
+      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

@@ -0,0 +1,149 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Schema } from 'effect';
+
+import { createTemplate, Message, structuredOutputParser } from '@dxos/artifact';
+import type { AIServiceClient } from '@dxos/assistant';
+import { MixedStreamParser } from '@dxos/assistant';
+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';
+
+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 ProcessTranscriptMessageResult = {
+  message: DataType.Message;
+  timeElapsed: number;
+};
+
+/**
+ * Process Handlebars template.
+ */
+const createSystemPrompt = (): string => {
+  const template = createTemplate(SYSTEM_PROMPT);
+  return template({});
+};
+
+const ReferencedQuotes = Schema.Struct({
+  references: Schema.Array(
+    Schema.Struct({
+      quote: Schema.String,
+      id: Schema.String,
+    }),
+  ).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],
+          }),
+        ),
+      );
+      log.info('entity extraction result', { refs: result.references });
+
+      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 quotes with DXNs references.
+ */
+// TODO(dmaretskyi): Lookup and verifiy ids from provided context.
+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/index.ts

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

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

@@ -0,0 +1,38 @@
+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/entity-extraction/types.d.ts

@@ -0,0 +1,8 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+declare module '*.tpl?raw' {
+  const content: string;
+  export default content;
+}

package/src/hooks/index.ts

@@ -4,5 +4,6 @@
 
 export * from './useAudioFile';
 export * from './useAudioTrack';
+export * from './useQueueModelAdapter';
 export * from './useTranscriber';
 export * from './useVoiceInput';