@dxos/assistant 0.8.4-main.fd6878d → 0.8.4-main.fffef41

package/src/execution-graph.ts

@@ -0,0 +1,509 @@
+//
+// 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, Message } from '@dxos/types';
+import { isNonNullable } from '@dxos/util';
+
+const SKIP_BLOCKS: ContentBlock.Any['_tag'][] = ['text'];
+
+/**
+ * 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
+
+  constructor(private readonly _skipBlocks: ContentBlock.Any['_tag'][] = SKIP_BLOCKS) {}
+
+  /**
+   * Adds events to the graph.
+   */
+  addEvents(events: Obj.Any[]) {
+    const sortedEvents = this.sortObjectsByCreated(events);
+    for (const event of sortedEvents) {
+      if (Obj.instanceOf(Message.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: Message.Message) {
+    const messageCommits = messageToCommits(
+      message,
+      this._lastBlockId,
+      this._toolCallCommitIds,
+      this._lastCommitByBranch,
+      this._skipBlocks,
+    );
+    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];
+  }
+}
+
+/**
+ * Creates commits for all blocks in a message.
+ */
+// TODO(burdon): Pass in AiToolProvider.
+const messageToCommits = (
+  message: Message.Message,
+  lastBlockId?: string,
+  toolCallIds?: Map<string, string>,
+  lastCommitByBranch?: Map<string, string>,
+  skipBlocks?: ContentBlock.Any['_tag'][],
+): Commit[] => {
+  let previousBlockId: string | undefined = lastBlockId;
+
+  return message.blocks
+    .map((block, idx) => {
+      const branch = getMessageBranch(message);
+      const parents = getBlockParents(block, previousBlockId, message, toolCallIds, lastCommitByBranch);
+      if (skipBlocks?.includes(block._tag)) {
+        return null;
+      }
+
+      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: Message.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: Message.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: Message.Message) => {
+  return getBranchName({
+    parentMessage: message.parentMessage,
+    toolCallId: message.properties?.[MESSAGE_PROPERTY_TOOL_CALL_ID],
+  });
+};
+
+const getParentId = (message: Message.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

@@ -5,14 +5,14 @@
 import { beforeAll, describe, test } from 'vitest';
 
 import { todo } from '@dxos/debug';
-import { FunctionExecutor } from '@dxos/functions';
+import { FunctionExecutor } from '@dxos/functions-runtime';
 import { log } from '@dxos/log';
-import { createTestData } from '@dxos/schema/testing';
+import { createTestData } from '@dxos/types/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', () => {
@@ -33,7 +33,7 @@ describe.skip('LLM EntityExtraction', () => {
           message,
           objects: [...documents, ...Object.values(contacts)],
         },
-        function: extractionAnthropicFn,
+        function: extractionAnthropicFunction,
         executor,
       });
       log.info('output', enhancedMessage);
@@ -42,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);
 
@@ -54,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 });
       }),
     );
@@ -64,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,70 +2,61 @@
 // Copyright 2025 DXOS.org
 //
 
-// ISSUE(burdon): defineFunction
-// @ts-nocheck
-
-import { AiService, ToolRegistry } from '@dxos/ai';
-import { Obj } from '@dxos/echo';
-import { create } from '@dxos/echo-schema';
 import { type FunctionDefinition, defineFunction } from '@dxos/functions';
-import { type ContentBlock, DataType } from '@dxos/schema';
 
-import { AiSession } from '../../session';
 import { ExtractionInput, ExtractionOutput } from '../extraction';
-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.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: [
-        Obj.make(DataType.Message, {
-          created: new Date().toISOString(),
-          sender: { role: 'user' },
-          blocks: [
-            {
-              _tag: 'text',
-              text: `<context>${JSON.stringify(objects)}</context>`,
-            } satisfies ContentBlock.Text,
-            {
-              _tag: 'text',
-              text: `<transcript>${JSON.stringify(message.blocks)}</transcript>`,
-            } satisfies ContentBlock.Text,
-          ],
-        }),
-      ],
-      artifacts: [],
-      prompt: '',
-      tools: [],
-      toolResolver: new ToolRegistry([]),
-    } as any); // TODO(burdon): Rewrite test.
+    throw new Error('not implemented');
+    // const startTime = performance.now();
+    // const ai = context.getService(AiService.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: [
+    //     Obj.make(DataType.Message.Message, {
+    //       created: new Date().toISOString(),
+    //       sender: { role: 'user' },
+    //       blocks: [
+    //         {
+    //           _tag: 'text',
+    //           text: `<context>${JSON.stringify(objects)}</context>`,
+    //         } satisfies DataType.ContentBlock.Text,
+    //         {
+    //           _tag: 'text',
+    //           text: `<transcript>${JSON.stringify(message.blocks)}</transcript>`,
+    //         } satisfies DataType.ContentBlock.Text,
+    //       ],
+    //     }),
+    //   ],
+    //   artifacts: [],
+    //   prompt: '',
+    //   tools: [],
+    // } as any); // TODO(burdon): Rewrite test.
 
-    return {
-      message: create(DataType.Message, {
-        ...message,
-        blocks: message.blocks.map((block, i) =>
-          block._tag !== 'transcript'
-            ? block
-            : {
-                ...block,
-                text: insertReferences(block.text, result),
-              },
-        ),
-      }),
-      timeElapsed: performance.now() - startTime,
-    };
+    // return {
+    //   timeElapsed: performance.now() - startTime,
+    //   message: create(DataType.Message.Message, {
+    //     ...message,
+    //     blocks: message.blocks.map((block) =>
+    //       block._tag !== 'transcript'
+    //         ? block
+    //         : {
+    //             ...block,
+    //             text: insertReferences(block.text, result),
+    //           },
+    //     ),
+    //   }),
+    // };
   },
 });