@dxos/assistant 0.8.2-main.fbd8ed0 → 0.8.2-staging.4d6ad0f

package/src/extraction/extraction-llm-function/extraction-llm-function.ts

@@ -0,0 +1,64 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Message } from '@dxos/ai';
+import { create } from '@dxos/echo-schema';
+import { AiService, defineFunction, type FunctionDefinition } from '@dxos/functions';
+import { DataType } from '@dxos/schema';
+
+import PROMPT from './instructions.tpl?raw';
+import { AISession } from '../../session';
+import { ExtractionInput, ExtractionOutput } from '../extraction';
+import { insertReferences, ReferencedQuotes } from '../quotes';
+
+export const extractionAnthropicFn: FunctionDefinition<ExtractionInput, ExtractionOutput> = defineFunction({
+  description: 'Extract entities from the transcript message and add them to the message.',
+  inputSchema: ExtractionInput,
+  outputSchema: ExtractionOutput,
+  handler: async ({ data: { message, objects }, context }) => {
+    const startTime = performance.now();
+    const ai = context.getService(AiService);
+    const session = new AISession({ operationModel: 'configured' });
+    const result = await session.runStructured(ReferencedQuotes, {
+      generationOptions: {
+        model: '@anthropic/claude-3-5-haiku-20241022',
+      },
+      client: ai.client,
+      systemPrompt: PROMPT,
+      history: [
+        create(Message, {
+          role: 'user',
+          content: [
+            {
+              type: 'text',
+              text: `<context>${JSON.stringify(objects)}</context>`,
+            },
+            {
+              type: 'text',
+              text: `<transcript>${JSON.stringify(message.blocks)}</transcript>`,
+            },
+          ],
+        }),
+      ],
+      artifacts: [],
+      prompt: '',
+      tools: [],
+    });
+
+    return {
+      message: create(DataType.Message, {
+        ...message,
+        blocks: message.blocks.map((block, i) =>
+          block.type !== 'transcription'
+            ? block
+            : {
+                ...block,
+                text: insertReferences(block.text, result),
+              },
+        ),
+      }),
+      timeElapsed: performance.now() - startTime,
+    };
+  },
+});

package/src/extraction/extraction-llm-function/index.ts

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

package/src/extraction/extraction-llm-function/instructions.tpl

@@ -0,0 +1,28 @@
+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:
+You have been provided with a tool that specifies output format
+
+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/extraction/extraction-ner-function/extraction-ner-function.test.ts

@@ -0,0 +1,115 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { describe, test, beforeAll } from 'vitest';
+
+import { type EchoDatabase } from '@dxos/echo-db';
+import { EchoTestBuilder } from '@dxos/echo-db/testing';
+import { FunctionExecutor, ServiceContainer } from '@dxos/functions';
+import { log } from '@dxos/log';
+import { DataType } from '@dxos/schema';
+import { createTestData, Testing } from '@dxos/schema/testing';
+import { range } from '@dxos/util';
+
+import { extractionNerFn } from './extraction-ner-function';
+import { processTranscriptMessage } from '../extraction';
+
+describe.skip('NER EntityExtraction', () => {
+  let builder: EchoTestBuilder;
+  let db: EchoDatabase;
+  let executor: FunctionExecutor;
+  let testData: {
+    transcriptJosiah: DataType.Message[];
+    transcriptWoflram: DataType.Message[];
+    transcriptMessages: DataType.Message[];
+    documents: Testing.DocumentType[];
+    contacts: Record<string, DataType.Person>;
+    organizations: Record<string, DataType.Organization>;
+  };
+  const TYPES = [DataType.Organization, DataType.Person, Testing.Contact, Testing.DocumentType];
+
+  beforeAll(async () => {
+    // TODO(dmaretskyi): Helper to scaffold this from a config.
+    builder = await new EchoTestBuilder().open();
+    const { db: db1 } = await builder.createDatabase({ indexing: { vector: true } });
+    db = db1;
+    db.graph.schemaRegistry.addSchema(TYPES);
+    const data = createTestData();
+    testData = {
+      transcriptJosiah: data.transcriptJosiah,
+      transcriptWoflram: data.transcriptWoflram,
+      transcriptMessages: data.transcriptMessages,
+      documents: data.documents.map((document) => db.add(document)),
+      contacts: Object.fromEntries(Object.entries(data.contacts).map(([key, value]) => [key, db.add(value)])),
+      organizations: Object.fromEntries(Object.entries(data.organizations).map(([key, value]) => [key, db.add(value)])),
+    };
+    await db.flush();
+
+    executor = new FunctionExecutor(
+      new ServiceContainer().setServices({
+        database: { db },
+      }),
+    );
+  }, 30_000);
+
+  test('should process a transcript block', async () => {
+    const { transcriptMessages, documents, contacts } = testData;
+    log.info('context', { documents, contacts });
+
+    for (const message of transcriptMessages) {
+      log.info('input', message);
+      const { message: enhancedMessage } = await processTranscriptMessage({
+        input: {
+          message,
+          objects: [...documents, ...Object.values(contacts)],
+        },
+        function: extractionNerFn,
+        executor,
+      });
+      log.info('output', enhancedMessage);
+    }
+  });
+
+  test('computational irreducibility', async () => {
+    const { transcriptWoflram, documents, contacts } = testData;
+
+    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({
+          input: {
+            message,
+            objects: [...documents, ...Object.values(contacts)],
+          },
+          function: extractionNerFn,
+          executor,
+        });
+        log.info('output', { message: enhancedMessage.blocks[0], timeElapsed });
+      }),
+    );
+  });
+
+  test('org and document linking', async () => {
+    const { transcriptJosiah, documents, contacts, organizations } = testData;
+
+    log.info('context', { contacts, organizations, documents });
+
+    for (const message of transcriptJosiah) {
+      log.info('input', message);
+
+      const { message: enhancedMessage } = await processTranscriptMessage({
+        input: {
+          message,
+          objects: [...documents, ...Object.values(contacts), ...Object.values(organizations)],
+        },
+        function: extractionNerFn,
+        executor,
+      });
+      log.info('output', enhancedMessage);
+    }
+  });
+});

package/src/extraction/extraction-ner-function/extraction-ner-function.ts

@@ -0,0 +1,54 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { asyncTimeout } from '@dxos/async';
+import { create } from '@dxos/echo-schema';
+import { DatabaseService, defineFunction } from '@dxos/functions';
+import { invariant } from '@dxos/invariant';
+import { log } from '@dxos/log';
+import { DataType } from '@dxos/schema';
+
+import { extractFullEntities } from './named-entity-recognition';
+import { ExtractionInput, ExtractionOutput } from '../extraction';
+import { findQuotes, insertReferences } from '../quotes';
+
+// TODO(mykola): Make it use vector index to query objects and not pass objects as input.
+export const extractionNerFn = defineFunction({
+  description: 'Extracts named entities from a text',
+  inputSchema: ExtractionInput,
+  outputSchema: ExtractionOutput,
+  handler: async ({ data: { message, options }, context }) => {
+    log.info('input', { message, options });
+    const startTime = performance.now();
+    const { db } = context.getService(DatabaseService);
+
+    const entitiesPromise = Promise.all(
+      message.blocks.map(async (block) => {
+        invariant(block.type === 'transcription' || block.type === 'text', 'Block must have text');
+        return extractFullEntities(block.text);
+      }),
+    ).then((entities) => entities.flat());
+    const entities = options?.timeout ? await asyncTimeout(entitiesPromise, options.timeout) : await entitiesPromise;
+    log.info('entities', { entities });
+    const quoteReferences = await findQuotes(
+      entities.map((entity) => entity.word),
+      db,
+    );
+    const blocksWithReferences = message.blocks.map((block, i) => {
+      invariant(block.type === 'transcription' || block.type === 'text', 'Block must have text');
+      return { ...block, text: insertReferences(block.text, quoteReferences) };
+    });
+    const messageWithReferences = create(DataType.Message, {
+      ...message,
+      blocks: blocksWithReferences,
+    });
+
+    log.info('output', { messageWithReferences });
+
+    return {
+      message: messageWithReferences,
+      timeElapsed: performance.now() - startTime,
+    };
+  },
+});

package/src/extraction/extraction-ner-function/index.ts

@@ -0,0 +1,6 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './extraction-ner-function';
+export * from './named-entity-recognition';

package/src/extraction/extraction-ner-function/named-entity-recognition.test.ts

@@ -0,0 +1,94 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { type TokenClassificationOutput, type TokenClassificationPipelineType } from '@xenova/transformers';
+import { beforeAll, describe, expect, test } from 'vitest';
+
+import { log } from '@dxos/log';
+import { createTestData } from '@dxos/schema/testing';
+
+import { combineNerTokens, createTokenGroups, extractFullEntities, getNer } from './named-entity-recognition';
+
+describe('NamedEntityRecognition', () => {
+  let ner: TokenClassificationPipelineType;
+
+  beforeAll(async () => {
+    // Initialize the NER pipeline
+    ner = await getNer();
+  }, 30_000);
+
+  test('should identify named entities in text', async () => {
+    const text = 'Elon Musk is the CEO of SpaceX and Tesla.';
+    const result = (await ner(text)) as TokenClassificationOutput;
+
+    // Verify the structure of the result
+    expect(Array.isArray(result)).toBe(true);
+    expect(result.length).toBeGreaterThan(0);
+
+    log.info('ner', { result });
+
+    const combinedTokens = createTokenGroups(result).map(combineNerTokens);
+    log.info('combined', { text, result: combinedTokens });
+
+    // Check for specific entities
+    const entities = combinedTokens.map((item) => ({
+      text: item.word,
+      entity: item.entity,
+      score: item.score,
+    }));
+
+    log.info('entities', { entities });
+
+    // Verify Elon Musk is recognized as a person
+    const elonMusk = entities.find((e) => e.text.includes('Elon') || e.text.includes('Musk'));
+    expect(elonMusk).toBeDefined();
+    expect(elonMusk?.entity).toMatch(/PER/);
+
+    // Verify SpaceX and Tesla are recognized as organizations
+    const companies = entities.filter((e: any) => e.text === 'SpaceX' || e.text === 'Tesla');
+    expect(companies.length).toBe(2);
+    companies.forEach((company: any) => {
+      expect(company.entity).toMatch(/ORG/);
+    });
+  });
+
+  test('own data', async () => {
+    const { transcriptWoflram, transcriptJosiah } = await createTestData();
+
+    const blocks = [...transcriptWoflram, ...transcriptJosiah].flatMap((message) =>
+      message.blocks.flatMap((block) => (block.type === 'transcription' ? [block.text] : [])),
+    );
+
+    for (const block of blocks) {
+      const entities = await extractFullEntities(block);
+      log.info('ner', { text: block, result: entities });
+    }
+  });
+
+  test('more own data', async () => {
+    const { transcriptMessages } = await createTestData();
+
+    const blocks = transcriptMessages.flatMap((message) =>
+      message.blocks.flatMap((block) => (block.type === 'transcription' ? [block.text] : [])),
+    );
+
+    for (const block of blocks) {
+      const entities = await extractFullEntities(block);
+      log.info('ner', { text: block, result: entities });
+    }
+  });
+
+  test('should handle empty text', async () => {
+    const result = await ner('');
+    expect(Array.isArray(result)).toBe(true);
+    expect(result.length).toBe(0);
+  });
+
+  test('should handle text with no named entities', async () => {
+    const text = 'The weather is nice today.';
+    const result = await ner(text);
+    expect(Array.isArray(result)).toBe(true);
+    expect(result.length).toBe(0);
+  });
+});

package/src/extraction/extraction-ner-function/named-entity-recognition.ts

@@ -0,0 +1,97 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import {
+  type TokenClassificationOutput,
+  type TokenClassificationSingle,
+  pipeline,
+  type TokenClassificationPipelineType,
+} from '@xenova/transformers';
+
+import { log } from '@dxos/log';
+
+let _ner: Promise<TokenClassificationPipelineType>;
+/**
+ * Named Entity Recognition pipeline.
+ * Initializes the pipeline on first call.
+ * @returns The singleton promise that resolves to a token classification pipeline.
+ */
+export const getNer = () => {
+  if (!_ner) {
+    _ner = pipeline('ner', 'Xenova/bert-base-NER').then((ner) => {
+      log.info('NER model is ready');
+      return ner;
+    });
+  }
+  return _ner;
+};
+
+/**
+ * Extracts entities from text using the named entity recognition pipeline and does simple grouping of tokens.
+ */
+export const extractFullEntities = async (text: string): Promise<TokenClassificationSingle[]> => {
+  const ner = await getNer();
+  const entities = (await ner(text)) as TokenClassificationOutput;
+  const tokenGroups = createTokenGroups(entities);
+  const combinedEntities = tokenGroups.map(combineNerTokens);
+  return combinedEntities;
+};
+
+export const createTokenGroups = (tokens: TokenClassificationSingle[]): TokenClassificationSingle[][] => {
+  const tokenGroups: TokenClassificationSingle[][] = [];
+
+  const canBeCombined = (prevToken: TokenClassificationSingle, currentToken: TokenClassificationSingle) => {
+    const [currentState, currentEntity] = currentToken.entity.split('-'); // B, I, O
+    const [prevState, prevEntity] = prevToken.entity.split('-'); // B, I, O
+    if (prevEntity !== currentEntity) {
+      return false;
+    }
+    if (prevState === 'O') {
+      return false;
+    }
+    if (currentState === 'I') {
+      return true;
+    }
+    if ((prevState === 'B' || prevState === 'I') && currentState === 'B') {
+      return currentToken.word.startsWith('#');
+    }
+    return false;
+  };
+
+  for (const token of tokens) {
+    if (tokenGroups.length === 0) {
+      tokenGroups.push([token]);
+      continue;
+    }
+    const lastTokenGroup = tokenGroups.at(-1)!;
+    if (canBeCombined(lastTokenGroup.at(-1)!, token)) {
+      lastTokenGroup.push(token);
+      continue;
+    }
+    tokenGroups.push([token]);
+  }
+
+  return tokenGroups;
+};
+
+export const combineNerTokens = (group: TokenClassificationSingle[]): TokenClassificationSingle => {
+  const combinedToken: TokenClassificationSingle = {
+    entity: group.at(0)!.entity.split('-')[1],
+    score: group.reduce((acc, token) => acc + token.score, 0) / group.length,
+    index: group.at(0)!.index,
+    word: group.at(0)!.word,
+    start: group.at(0)!.start,
+    end: group.at(-1)!.end,
+  };
+
+  for (const token of group.slice(1)) {
+    if (token.word.startsWith('#')) {
+      combinedToken.word += token.word.replace(/^#+/, '');
+    } else {
+      combinedToken.word += ' ' + token.word;
+    }
+  }
+
+  return combinedToken;
+};

package/src/extraction/extraction.ts

@@ -0,0 +1,75 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Schema } from 'effect';
+
+import { asyncTimeout } from '@dxos/async';
+import { Expando } from '@dxos/echo-schema';
+import { type FunctionDefinition, type FunctionExecutor } from '@dxos/functions';
+import { log } from '@dxos/log';
+import { DataType } from '@dxos/schema';
+
+export const ExtractionInput = Schema.Struct({
+  message: DataType.Message,
+  objects: Schema.optional(Schema.Array(Expando)),
+  options: Schema.optional(
+    Schema.Struct({
+      timeout: Schema.optional(Schema.Number),
+      fallbackToRaw: Schema.optional(Schema.Boolean),
+    }),
+  ),
+});
+export interface ExtractionInput extends Schema.Schema.Type<typeof ExtractionInput> {}
+
+export const ExtractionOutput = Schema.Struct({
+  message: DataType.Message,
+  timeElapsed: Schema.Number,
+});
+export interface ExtractionOutput extends Schema.Schema.Type<typeof ExtractionOutput> {}
+
+export type ProcessTranscriptMessageParams = {
+  input: ExtractionInput;
+  function: FunctionDefinition<ExtractionInput, ExtractionOutput>;
+  executor: FunctionExecutor;
+
+  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;
+  };
+};
+
+export type ExtractionFunction = FunctionDefinition<ExtractionInput, ExtractionOutput>;
+
+/**
+ * Extract entities from the transcript message and add them to the message.
+ */
+// TODO(dmaretskyi): Move context to a vector search index.
+export const processTranscriptMessage = async (params: ProcessTranscriptMessageParams): Promise<ExtractionOutput> => {
+  try {
+    if (params.options?.timeout && params.options.timeout > 0) {
+      return await asyncTimeout(params.executor.invoke(params.function, params.input), params.options.timeout);
+    } else {
+      return await params.executor.invoke(params.function, params.input);
+    }
+  } catch (error) {
+    if (params.options?.fallbackToRaw) {
+      log.warn('failed to process message, falling back to raw text', { error });
+      return {
+        message: params.input.message,
+        timeElapsed: 0,
+      };
+    } else {
+      throw error;
+    }
+  }
+};

package/src/extraction/index.ts

@@ -0,0 +1,8 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './extraction';
+export * from './extraction-llm-function';
+export * from './extraction-ner-function';
+export * from './quotes';

package/src/extraction/quotes.test.ts

@@ -0,0 +1,27 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { describe, test, expect } from 'vitest';
+
+import { insertReferences } from './quotes';
+
+describe('insertReferences', () => {
+  test('should replace quotes with DXN references', () => {
+    const quotes = {
+      references: [{ quote: 'computational irreducibility', id: '01JTG9JW11XGWJZ32AW8ET93D1' }],
+    };
+
+    expect(insertReferences('This is a computational irreducibility test.', quotes)).toBe(
+      'This is a [computational irreducibility][dxn:echo:@:01JTG9JW11XGWJZ32AW8ET93D1] test.',
+    );
+    expect(
+      insertReferences(
+        "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/extraction/quotes.ts

@@ -0,0 +1,57 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Schema } from 'effect';
+
+import { type EchoDatabase } from '@dxos/echo-db';
+import { Filter, ObjectId, Query } from '@dxos/echo-schema';
+import { DXN } from '@dxos/keys';
+import { log } from '@dxos/log';
+
+export const ReferencedQuotes = Schema.Struct({
+  references: Schema.Array(
+    Schema.Struct({
+      id: Schema.String,
+      quote: Schema.String, // TODO(burdon): Quote?
+    }),
+  )
+    .pipe(Schema.mutable)
+    .annotations({
+      // TODO(burdon): Does this description make sense?
+      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.
+    `,
+    }),
+});
+export type ReferencedQuotes = Schema.Schema.Type<typeof ReferencedQuotes>;
+
+export const findQuotes = async (quotes: string[], db: EchoDatabase): Promise<ReferencedQuotes> => {
+  const references = await Promise.all(
+    quotes.map(async (quote) => {
+      const { objects } = await db.query(Query.select(Filter.text(quote, { type: 'vector' }))).run();
+      log.info('found quotes in vector index', { objects, quote });
+      return objects.length > 0 ? [{ id: objects[0].id.toString(), quote }] : [];
+    }),
+  );
+  return { references: references.flat() };
+};
+
+/**
+ * Finds and replaces all quotes with DXNs references.
+ */
+// TODO(dmaretskyi): Lookup and verifiy ids from provided context.
+export const insertReferences = (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/index.ts

@@ -2,8 +2,8 @@
 // Copyright 2024 DXOS.org
 //
 
-export * from './ai-service';
+export * from './extraction';
+export * from './research';
 export * from './session';
-
-// TODO(dmaretskyi): Separate entrypoint for experimental features.
-export * from './experimental';
+export * from './types';
+export * from './blueprint';