@dxos/assistant 0.8.4-main.3a94e84 → 0.8.4-main.406dc2a

package/src/execution-graph.ts

@@ -0,0 +1,500 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { AgentStatus } from '@dxos/ai';
+import { Obj, type Ref } from '@dxos/echo';
+import { MESSAGE_PROPERTY_TOOL_CALL_ID } from '@dxos/functions';
+import { type ObjectId } from '@dxos/keys';
+import { LogLevel } from '@dxos/log';
+import { ContentBlock, DataType } from '@dxos/schema';
+import { isNonNullable } from '@dxos/util';
+
+/**
+ * Mercurial-style Commit.
+ */
+// TODO(wittjosiah): Reconcile with @dxos/react-ui-components.
+export type Commit = {
+  id: string;
+  parents: string[];
+  branch: string;
+  icon?: string;
+  level?: LogLevel;
+  message: string;
+  timestamp?: Date;
+  tags?: string[];
+};
+
+// TODO(burdon): Add colors?
+enum IconType {
+  // General status.
+  SUCCESS = 'ph--check-circle--regular',
+  WARNING = 'ph--warning-circle--regular',
+  ERROR = 'ph--x-circle--regular',
+  FLAG = 'ph--flag--regular',
+  ROCKET = 'ph--rocket--regular',
+  TIMER = 'ph--timer--regular',
+
+  // Interactions.
+  USER = 'ph--user--regular',
+  USER_INTERACTION = 'ph--user-sound--regular',
+  AGENT = 'ph--robot--regular',
+  THINKING = 'ph--brain--regular',
+  LINK = 'ph--link--regular',
+  TOOL = 'ph--wrench--regular',
+}
+
+/**
+ * Models the execution graph based on a stream of events.
+ */
+export class ExecutionGraph {
+  // TODO(dmaretskyi): Evolve the internal state to support chats, functions, circuit workflows, etc.
+  private _commits: Commit[] = [];
+  private _branchNames = new Set<string>();
+  private _lastBlockId: string | undefined;
+
+  // Tool call tracking.
+  private _toolCallCommitIds = new Map<string, string>(); // toolCallId -> commitId
+  private _lastCommitByBranch = new Map<string, string>(); // branch -> last commitId
+  private _pendingToolResults = new Map<string, string>(); // toolCallId -> toolResultCommitId
+
+  /**
+   * Adds events to the graph.
+   */
+  addEvents(events: Obj.Any[]) {
+    const sortedEvents = this.sortObjectsByCreated(events);
+    for (const event of sortedEvents) {
+      if (Obj.instanceOf(DataType.Message, event)) {
+        this._processMessage(event);
+      } else if (Obj.instanceOf(AgentStatus, event)) {
+        this._processAgentStatus(event);
+      }
+    }
+  }
+
+  /**
+   * Processes a message event and creates commits for its blocks.
+   */
+  private _processMessage(message: DataType.Message) {
+    const messageCommits = messageToCommits(
+      message,
+      this._lastBlockId,
+      this._toolCallCommitIds,
+      this._lastCommitByBranch,
+    );
+    this._commits.push(...messageCommits);
+
+    // Track branches.
+    messageCommits.forEach((commit) => this._branchNames.add(commit.branch));
+
+    // Track tool calls and update pending tool results.
+    this._trackToolCalls(messageCommits);
+
+    // Update last block ID for sequential chaining.
+    if (messageCommits.length > 0) {
+      this._lastBlockId = messageCommits[messageCommits.length - 1].id;
+    }
+  }
+
+  /**
+   * Processes an AgentStatus event.
+   */
+  private _processAgentStatus(event: AgentStatus) {
+    const branch = this._getToolCallBranch(event.parentMessage, event.toolCallId);
+    this._branchNames.add(branch);
+
+    const agentStatusCommit: Commit = {
+      id: event.id,
+      branch,
+      message: event.message,
+      icon: IconType.FLAG,
+      parents: this._getAgentStatusParents(event),
+      timestamp: new Date(event.created),
+    };
+
+    this._commits.push(agentStatusCommit);
+
+    // Update last block ID for sequential chaining.
+    this._lastBlockId = event.id;
+
+    // Update pending tool results to point to this AgentStatus.
+    if (event.toolCallId) {
+      this._lastCommitByBranch.set(branch, event.id);
+      this._updatePendingToolResults(event.toolCallId, event.id);
+    }
+  }
+
+  /**
+   * Stable sort: objects with 'created' field are sorted by it, others remain in place.
+   */
+  private sortObjectsByCreated(objects: any[]) {
+    const sortedObjects = objects.slice();
+    // Find indices of objects with 'created'
+    const createdObjects: { idx: number; obj: any }[] = [];
+    sortedObjects.forEach((obj, idx) => {
+      if (obj && typeof obj.created === 'string') {
+        createdObjects.push({ idx, obj });
+      }
+    });
+    // Sort only the objects with 'created' by their date, stable
+    createdObjects.sort((a, b) => {
+      const aDate = new Date(a.obj.created).getTime();
+      const bDate = new Date(b.obj.created).getTime();
+      return aDate - bDate;
+    });
+    // Place sorted createdObjects back into their original positions
+    let createdIdx = 0;
+    return sortedObjects.map((obj, _idx) => {
+      if (obj && typeof obj.created === 'string') {
+        return createdObjects[createdIdx++].obj;
+      }
+      return obj;
+    });
+  }
+
+  /**
+   * Returns the current state of the graph.
+   */
+  getGraph(lastRequest = false): { branches: string[]; commits: Commit[] } {
+    const idx = lastRequest ? this._commits.findLastIndex((c) => c.tags?.includes('user')) : -1;
+    const commits = idx === -1 ? this._commits : this._commits.slice(idx);
+
+    return {
+      branches: Array.from(this._branchNames),
+      commits,
+    };
+  }
+
+  /**
+   * Tracks tool calls and handles pending tool result updates.
+   */
+  private _trackToolCalls(commits: Commit[]) {
+    commits.forEach((commit) => {
+      if (this._isToolCallCommit(commit)) {
+        const toolCallId = this._extractToolCallId(commit.id);
+        if (toolCallId) {
+          this._toolCallCommitIds.set(toolCallId, commit.id);
+        }
+      }
+
+      if (this._isToolResultCommit(commit)) {
+        const toolCallId = this._extractToolCallId(commit.id);
+        if (toolCallId) {
+          this._pendingToolResults.set(toolCallId, commit.id);
+          this._handleLateToolResultUpdate(toolCallId, commit.id);
+        }
+      }
+
+      this._lastCommitByBranch.set(commit.branch, commit.id);
+    });
+  }
+
+  /**
+   * Updates pending tool results to point to an AgentStatus event.
+   */
+  private _updatePendingToolResults(toolCallId: string, agentStatusId: string) {
+    const toolResultId = this._pendingToolResults.get(toolCallId);
+    if (!toolResultId) return;
+
+    const toolResultCommit = this._commits.find((c) => c.id === toolResultId);
+    if (!toolResultCommit) return;
+
+    const toolCallCommitId = this._toolCallCommitIds.get(toolCallId);
+    if (!toolCallCommitId) return;
+
+    // Replace tool call commit with AgentStatus commit.
+    const toolCallIndex = toolResultCommit.parents.indexOf(toolCallCommitId);
+    if (toolCallIndex !== -1) {
+      toolResultCommit.parents[toolCallIndex] = agentStatusId;
+    }
+
+    this._pendingToolResults.delete(toolCallId);
+  }
+
+  /**
+   * Handles late updates when AgentStatus is processed before tool result.
+   */
+  private _handleLateToolResultUpdate(toolCallId: string, toolResultId: string) {
+    const toolCallCommitId = this._toolCallCommitIds.get(toolCallId);
+    if (!toolCallCommitId) return;
+
+    const messageId = toolCallCommitId.split('_toolCall_')[0];
+    const toolCallBranch = `${messageId}_${toolCallId}`;
+    const agentStatusCommitId = this._lastCommitByBranch.get(toolCallBranch);
+
+    if (agentStatusCommitId && agentStatusCommitId !== toolCallCommitId) {
+      const toolResultCommit = this._commits.find((c) => c.id === toolResultId);
+      if (toolResultCommit) {
+        const toolCallIndex = toolResultCommit.parents.indexOf(toolCallCommitId);
+        if (toolCallIndex !== -1) {
+          toolResultCommit.parents[toolCallIndex] = agentStatusCommitId;
+        }
+      }
+      this._pendingToolResults.delete(toolCallId);
+    }
+  }
+
+  /**
+   * Gets the branch name for a tool call.
+   */
+  private _getToolCallBranch(parentMessage?: ObjectId, toolCallId?: string): string {
+    if (toolCallId) {
+      return toolCallId;
+    }
+    return parentMessage || 'main';
+  }
+
+  /**
+   * Gets the parents for an AgentStatus commit.
+   */
+  private _getAgentStatusParents(event: AgentStatus): string[] {
+    if (event.parentMessage && event.toolCallId) {
+      return [getToolCallId(event.parentMessage, event.toolCallId)];
+    }
+    return [];
+  }
+
+  /**
+   * Checks if a commit is a tool call commit.
+   */
+  private _isToolCallCommit(commit: Commit): boolean {
+    return commit.message.includes('Calling tool');
+  }
+
+  /**
+   * Checks if a commit is a tool result commit.
+   */
+  private _isToolResultCommit(commit: Commit): boolean {
+    return commit.message.includes('Tool call succeeded') || commit.message.includes('Tool error:');
+  }
+
+  /**
+   * Extracts tool call ID from a commit ID.
+   */
+  private _extractToolCallId(commitId: string): string | undefined {
+    const toolCallMatch = commitId.match(/_toolCall_(.+)$/);
+    const toolResultMatch = commitId.match(/_toolResult_(.+)$/);
+    return toolCallMatch?.[1] || toolResultMatch?.[1];
+  }
+}
+
+// TODO(burdon): Pass in AiToolProvider.
+/**
+ * Creates commits for all blocks in a message.
+ */
+const messageToCommits = (
+  message: DataType.Message,
+  lastBlockId?: string,
+  toolCallIds?: Map<string, string>,
+  lastCommitByBranch?: Map<string, string>,
+): Commit[] => {
+  let previousBlockId: string | undefined = lastBlockId;
+
+  return message.blocks
+    .map((block, idx) => {
+      const branch = getMessageBranch(message);
+      const parents = getBlockParents(block, previousBlockId, message, toolCallIds, lastCommitByBranch);
+
+      const commit = createBlockCommit(block, message, branch, parents, idx);
+      if (commit) {
+        previousBlockId = commit.id;
+      }
+
+      return commit;
+    })
+    .filter(isNonNullable);
+};
+
+/**
+ * Determines the parent commits for a block based on its type and context.
+ */
+const getBlockParents = (
+  block: ContentBlock.Any,
+  previousBlockId: string | undefined,
+  message: DataType.Message,
+  toolCallIds?: Map<string, string>,
+  lastCommitByBranch?: Map<string, string>,
+): string[] => {
+  const parents: string[] = [];
+
+  if (block._tag === 'toolResult') {
+    // Tool results have two parents: previous block and last block from tool call branch
+    if (previousBlockId) {
+      parents.push(previousBlockId);
+    }
+
+    const toolCallCommitId = toolCallIds?.get(block.toolCallId);
+    if (toolCallCommitId) {
+      const toolCallBranch = getToolCallBranchFromCommitId(toolCallCommitId, block.toolCallId);
+      const lastBlockFromToolCallBranch = lastCommitByBranch?.get(toolCallBranch);
+
+      if (lastBlockFromToolCallBranch) {
+        parents.push(lastBlockFromToolCallBranch);
+      } else {
+        // Fallback to tool call commit if no AgentStatus exists yet.
+        parents.push(toolCallCommitId);
+      }
+    }
+  } else {
+    // For other blocks, use previous block as parent (or tool call parent for first block).
+    if (previousBlockId) {
+      parents.push(previousBlockId);
+    } else {
+      const toolCallParent = getParentId(message);
+      if (toolCallParent) {
+        parents.push(toolCallParent);
+      }
+    }
+  }
+
+  return parents;
+};
+
+/**
+ * Creates a commit for a specific block.
+ */
+const createBlockCommit = (
+  block: ContentBlock.Any,
+  message: DataType.Message,
+  branch: string,
+  parents: string[],
+  idx: number,
+): Commit | null => {
+  const timestamp = new Date(message.created);
+  switch (block._tag) {
+    case 'text':
+      if (!block.text.trim().length) return null;
+      return {
+        id: getGenericBlockId(message.id, idx),
+        branch,
+        parents,
+        timestamp,
+        ...(message.sender.role === 'user'
+          ? {
+              icon: IconType.USER,
+              tags: ['user'],
+              message: 'Processing request...',
+            }
+          : {
+              icon: IconType.AGENT,
+              message: `Response (${block.text.split(' ').length} words)`,
+            }),
+      } satisfies Commit;
+
+    case 'toolCall':
+      return {
+        id: getToolCallId(message.id, block.toolCallId),
+        branch,
+        parents,
+        timestamp,
+        icon: IconType.TOOL,
+        level: LogLevel.INFO,
+        message: `Calling tool (${block.name})`,
+      } satisfies Commit;
+
+    case 'toolResult':
+      return {
+        id: getToolResultId(message.id, block.toolCallId),
+        branch,
+        parents,
+        timestamp,
+        icon: block.error ? IconType.ERROR : IconType.SUCCESS,
+        level: block.error ? LogLevel.ERROR : LogLevel.INFO,
+        message: block.error ? 'Tool error: ' + block.error : 'Tool call succeeded',
+      } satisfies Commit;
+
+    case 'summary':
+      return {
+        id: getGenericBlockId(message.id, idx),
+        branch,
+        parents,
+        timestamp,
+        icon: IconType.ROCKET,
+        level: LogLevel.INFO,
+        message: ContentBlock.createSummaryMessage(block),
+      } satisfies Commit;
+
+    case 'status':
+      return {
+        id: getGenericBlockId(message.id, idx),
+        branch,
+        parents,
+        timestamp,
+        message: block.statusText,
+        level: LogLevel.INFO,
+        icon: IconType.FLAG,
+      } satisfies Commit;
+
+    case 'reasoning':
+      return {
+        id: getGenericBlockId(message.id, idx),
+        branch,
+        parents,
+        timestamp,
+        message: block.reasoningText ?? 'Thinking...',
+        icon: IconType.THINKING,
+      } satisfies Commit;
+
+    case 'reference':
+      return {
+        id: getGenericBlockId(message.id, idx),
+        branch,
+        parents,
+        timestamp,
+        icon: IconType.LINK,
+        message: stringifyRef(block.reference),
+      } satisfies Commit;
+
+    default:
+      return null;
+  }
+};
+
+/**
+ * Gets the tool call branch name from a tool call commit ID.
+ */
+const getToolCallBranchFromCommitId = (toolCallCommitId: string, toolCallId: string): string => {
+  const messageId = toolCallCommitId.split('_toolCall_')[0];
+  return `${messageId}_${toolCallId}`;
+};
+
+const getToolCallId = (messageId: ObjectId, toolCallId: string) => `${messageId}_toolCall_${toolCallId}`;
+const getToolResultId = (messageId: ObjectId, toolCallId: string) => `${messageId}_toolResult_${toolCallId}`;
+const getGenericBlockId = (messageId: ObjectId, idx: number) => `${messageId}_block_${idx}`;
+
+const getBranchName = (options: { parentMessage?: ObjectId; toolCallId?: string }) => {
+  if (options.toolCallId) {
+    return options.toolCallId;
+  } else if (options.parentMessage) {
+    return options.parentMessage;
+  } else {
+    return 'main';
+  }
+};
+
+const getMessageBranch = (message: DataType.Message) => {
+  return getBranchName({
+    parentMessage: message.parentMessage,
+    toolCallId: message.properties?.[MESSAGE_PROPERTY_TOOL_CALL_ID],
+  });
+};
+
+const getParentId = (message: DataType.Message) => {
+  if (message.parentMessage && message.properties?.[MESSAGE_PROPERTY_TOOL_CALL_ID]) {
+    return getToolCallId(message.parentMessage, message.properties[MESSAGE_PROPERTY_TOOL_CALL_ID]);
+  } else {
+    return undefined;
+  }
+};
+
+const stringifyRef = (ref: Ref.Any) => {
+  if (ref.target) {
+    return stringifyObject(ref.target);
+  }
+
+  return ref.dxn.asEchoDXN()?.echoId ?? ref.dxn.asQueueDXN()?.objectId ?? '';
+};
+
+const stringifyObject = (obj: Obj.Any) => {
+  return Obj.getLabel(obj) ?? Obj.getTypename(obj) ?? obj.id;
+};

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

@@ -11,7 +11,8 @@ import { createTestData } from '@dxos/schema/testing';
 import { range } from '@dxos/util';
 
 import { processTranscriptMessage } from '../extraction';
-import { extractionAnthropicFn } from './extraction-llm-function';
+
+import { extractionAnthropicFunction } from './extraction-llm-function';
 
 // TODO(burdon): Rewrite test.
 describe.skip('LLM EntityExtraction', () => {
@@ -32,7 +33,7 @@ describe.skip('LLM EntityExtraction', () => {
           message,
           objects: [...documents, ...Object.values(contacts)],
         },
-        function: extractionAnthropicFn,
+        function: extractionAnthropicFunction,
         executor,
       });
       log.info('output', enhancedMessage);
@@ -41,8 +42,8 @@ describe.skip('LLM EntityExtraction', () => {
 
   test('computational irreducibility', async () => {
     const { transcriptWoflram, documents, contacts } = createTestData();
-
     log.info('context', { documents, contacts });
+
     const message = transcriptWoflram[0];
     log.info('input', message);
 
@@ -53,9 +54,10 @@ describe.skip('LLM EntityExtraction', () => {
             message,
             objects: [...documents, ...Object.values(contacts)],
           },
-          function: extractionAnthropicFn,
+          function: extractionAnthropicFunction,
           executor,
         });
+
         log.info('output', { message: enhancedMessage.blocks[0], timeElapsed });
       }),
     );
@@ -63,20 +65,19 @@ describe.skip('LLM EntityExtraction', () => {
 
   test('org and document linking', async () => {
     const { transcriptJosiah, documents, contacts, organizations } = createTestData();
-
     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: extractionAnthropicFn,
+        function: extractionAnthropicFunction,
         executor,
       });
+
       log.info('output', enhancedMessage);
     }
   });

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

@@ -2,27 +2,30 @@
 // Copyright 2025 DXOS.org
 //
 
-// TODO(burdon): Fix.
+// ISSUE(burdon): defineFunction
 // @ts-nocheck
 
-import { ToolRegistry, AiService } from '@dxos/ai';
+import { AiService } from '@dxos/ai';
 import { Obj } from '@dxos/echo';
-import { create } from '@dxos/echo-schema';
-import { defineFunction, type FunctionDefinition } from '@dxos/functions';
+import { create } from '@dxos/echo/internal';
+import { type FunctionDefinition, defineFunction } from '@dxos/functions';
 import { type ContentBlock, DataType } from '@dxos/schema';
 
-import PROMPT from './instructions.tpl?raw';
 import { AiSession } from '../../session';
 import { ExtractionInput, ExtractionOutput } from '../extraction';
-import { insertReferences, ReferencedQuotes } from '../quotes';
+import { ReferencedQuotes, insertReferences } from '../quotes';
+
+import PROMPT from './instructions.tpl?raw';
 
-export const extractionAnthropicFn: FunctionDefinition<ExtractionInput, ExtractionOutput> = defineFunction({
+export const extractionAnthropicFunction: FunctionDefinition<ExtractionInput, ExtractionOutput> = defineFunction({
+  key: 'dxos.org/function/extraction/extract-entities',
+  name: 'Extract Entities',
   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 ai = context.getService(AiService.AiService);
     const session = new AiSession({ operationModel: 'configured' });
     const result = await session.runStructured(ReferencedQuotes, {
       generationOptions: {
@@ -49,13 +52,13 @@ export const extractionAnthropicFn: FunctionDefinition<ExtractionInput, Extracti
       artifacts: [],
       prompt: '',
       tools: [],
-      toolResolver: new ToolRegistry([]),
     } as any); // TODO(burdon): Rewrite test.
 
     return {
+      timeElapsed: performance.now() - startTime,
       message: create(DataType.Message, {
         ...message,
-        blocks: message.blocks.map((block, i) =>
+        blocks: message.blocks.map((block) =>
           block._tag !== 'transcript'
             ? block
             : {
@@ -64,7 +67,6 @@ export const extractionAnthropicFn: FunctionDefinition<ExtractionInput, Extracti
               },
         ),
       }),
-      timeElapsed: performance.now() - startTime,
     };
   },
 });

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

@@ -2,19 +2,20 @@
 // Copyright 2025 DXOS.org
 //
 
-import { describe, test, beforeAll } from 'vitest';
+import { beforeAll, describe, test } 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 { Testing, createTestData } from '@dxos/schema/testing';
 import { range } from '@dxos/util';
 
-import { extractionNerFn } from './extraction-ner-function';
 import { processTranscriptMessage } from '../extraction';
 
+import { extractionNerFunction } from './extraction-ner-function';
+
 describe.skip('NER EntityExtraction', () => {
   let builder: EchoTestBuilder;
   let db: EchoDatabase;
@@ -27,6 +28,7 @@ describe.skip('NER EntityExtraction', () => {
     contacts: Record<string, DataType.Person>;
     organizations: Record<string, DataType.Organization>;
   };
+
   const TYPES = [DataType.Organization, DataType.Person, Testing.Contact, Testing.DocumentType];
 
   beforeAll(async () => {
@@ -44,8 +46,8 @@ describe.skip('NER EntityExtraction', () => {
       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();
 
+    await db.flush();
     executor = new FunctionExecutor(
       new ServiceContainer().setServices({
         database: { db },
@@ -64,7 +66,7 @@ describe.skip('NER EntityExtraction', () => {
           message,
           objects: [...documents, ...Object.values(contacts)],
         },
-        function: extractionNerFn,
+        function: extractionNerFunction,
         executor,
       });
       log.info('output', enhancedMessage);
@@ -73,7 +75,6 @@ describe.skip('NER EntityExtraction', () => {
 
   test('computational irreducibility', async () => {
     const { transcriptWoflram, documents, contacts } = testData;
-
     log.info('context', { documents, contacts });
     const message = transcriptWoflram[0];
     log.info('input', message);
@@ -85,7 +86,7 @@ describe.skip('NER EntityExtraction', () => {
             message,
             objects: [...documents, ...Object.values(contacts)],
           },
-          function: extractionNerFn,
+          function: extractionNerFunction,
           executor,
         });
         log.info('output', { message: enhancedMessage.blocks[0], timeElapsed });
@@ -95,7 +96,6 @@ describe.skip('NER EntityExtraction', () => {
 
   test('org and document linking', async () => {
     const { transcriptJosiah, documents, contacts, organizations } = testData;
-
     log.info('context', { contacts, organizations, documents });
 
     for (const message of transcriptJosiah) {
@@ -106,7 +106,7 @@ describe.skip('NER EntityExtraction', () => {
           message,
           objects: [...documents, ...Object.values(contacts), ...Object.values(organizations)],
         },
-        function: extractionNerFn,
+        function: extractionNerFunction,
         executor,
       });
       log.info('output', enhancedMessage);