@dxos/assistant 0.8.2-staging.7ac8446 → 0.8.2

package/src/blueprint/blueprint.test.ts

@@ -0,0 +1,142 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Schema } from 'effect';
+import { describe, test } from 'vitest';
+
+import { AIServiceEdgeClient, ToolResult, createTool } from '@dxos/ai';
+import { AI_SERVICE_ENDPOINT, EXA_API_KEY } from '@dxos/ai/testing';
+import { ArtifactId } from '@dxos/artifact';
+import { EchoTestBuilder } from '@dxos/echo-db/testing';
+import { create } from '@dxos/echo-schema';
+import { DataType, DataTypes } from '@dxos/schema';
+
+import { BlueprintBuilder } from './blueprint';
+import { setConsolePrinter } from './logger';
+import { BlueprintMachine } from './machine';
+import { TEST_EMAILS } from './test-data';
+import { createGraphWriterTool, createLocalSearchTool } from '../research';
+import { createExaTool } from '../research/exa';
+
+// TODO(burdon): Conslidate with existing artifact definition and create JSON DSL.
+
+// TODO(burdon): Don't run on CI.
+describe.skip('Blueprint', () => {
+  const aiService = new AIServiceEdgeClient({
+    endpoint: AI_SERVICE_ENDPOINT.REMOTE,
+    defaultGenerationOptions: {
+      // model: '@anthropic/claude-sonnet-4-20250514',
+      model: '@anthropic/claude-3-5-sonnet-20241022',
+    },
+  });
+
+  test('follows a simple blueprint', { timeout: 60_000 }, async () => {
+    const blueprint = BlueprintBuilder.create()
+      .step('Generate an idea for a new product. Do not use any external tools for this.')
+      .step('Write a short description of the product.')
+      .step('Run a market research to see if the product is viable. Do not use any external tools for this.')
+      .step('Write a pitch deck for the product')
+      .build();
+
+    const machine = new BlueprintMachine(blueprint);
+    setConsolePrinter(machine, true);
+    await machine.runToCompletion({ aiService });
+  });
+
+  test('email bot', { timeout: 60_000 }, async () => {
+    const replyTool = createTool('email', {
+      name: 'reply',
+      description: 'Reply to the email',
+      schema: Schema.Struct({
+        toEmail: ArtifactId,
+        subject: Schema.String.annotations({
+          description: 'The subject of the reply',
+        }),
+        body: Schema.String.annotations({
+          description: 'The body of the reply',
+        }),
+      }),
+      execute: async (params) => {
+        console.log('reply', params);
+        return ToolResult.Success('Sent!');
+      },
+    });
+
+    const labelTool = createTool('email', {
+      name: 'label',
+      description: 'Apply a label to the email',
+      schema: Schema.Struct({
+        toEmail: ArtifactId,
+        label: Schema.String.annotations({
+          description: 'The label to apply to the email',
+        }),
+      }),
+      execute: async (params) => {
+        return ToolResult.Success('Labeled!');
+      },
+    });
+
+    const blueprint = BlueprintBuilder.create()
+      .step(
+        'Determine if the email is introduction, question, or spam. Bail if email does not fit into one of these categories.',
+      )
+      .step('If the email is spam, label it as spam and do not respond.', {
+        tools: [labelTool],
+      })
+      .step(
+        'If the email is an introduction, respond with a short introduction of yourself and ask for more information.',
+        {
+          tools: [replyTool],
+        },
+      )
+      .step('If the email is a question, respond with a short answer and ask for more information.', {
+        tools: [replyTool],
+      })
+      .build();
+
+    const machine = new BlueprintMachine(blueprint);
+    setConsolePrinter(machine);
+    await machine.runToCompletion({ aiService, input: TEST_EMAILS[0] });
+  });
+
+  test.only('research', { timeout: 120_000 }, async () => {
+    const builder = await new EchoTestBuilder().open();
+    const { db } = await builder.createDatabase({ indexing: { vector: true }, types: DataTypes });
+
+    const [org1] = [
+      db.add(
+        create(DataType.Organization, {
+          name: 'Exa',
+          website: 'https://exa.ai',
+          description: 'An AI-powered search engine company building search infrastructure for AI agents',
+        }),
+      ),
+      db.add(
+        create(DataType.Organization, {
+          name: 'Cresta',
+          website: 'https://cresta.ai',
+          description: 'A company that builds AI agents',
+        }),
+      ),
+    ];
+
+    await db.flush({ indexes: true });
+
+    const blueprint = BlueprintBuilder.create()
+      .step('Research founders of the organization. Do deep research.', {
+        tools: [createExaTool({ apiKey: EXA_API_KEY })],
+      })
+      .step('Based on your research select matching entires that are already in the graph', {
+        tools: [createLocalSearchTool(db)],
+      })
+      .step('Add researched data to the graph', {
+        tools: [createLocalSearchTool(db), createGraphWriterTool({ db, schemaTypes: DataTypes })],
+      })
+      .build();
+
+    const machine = new BlueprintMachine(blueprint);
+    setConsolePrinter(machine, true);
+    await machine.runToCompletion({ aiService, input: org1 });
+  });
+});

package/src/blueprint/blueprint.ts

@@ -0,0 +1,95 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Schema } from 'effect';
+
+import { type ExecutableTool, Tool } from '@dxos/ai';
+import { raise } from '@dxos/debug';
+import { ObjectId } from '@dxos/keys';
+
+export const BlueprintStep = Schema.Struct({
+  id: Schema.String,
+  instructions: Schema.String,
+  // TODO(burdon): ExecutableTool can't be serialized.
+  tools: Schema.Array(Tool).pipe(Schema.mutable),
+});
+export interface BlueprintStep extends Schema.Schema.Type<typeof BlueprintStep> {}
+
+export const Blueprint = Schema.Struct({
+  steps: Schema.Array(BlueprintStep).pipe(Schema.mutable),
+});
+export interface Blueprint extends Schema.Schema.Type<typeof Blueprint> {}
+
+/**
+ * Blueprint builder API.
+ */
+export namespace BlueprintBuilder {
+  export const create = () => new Builder();
+
+  class Builder {
+    private readonly _steps: BlueprintStep[] = [];
+
+    step(instructions: string, options?: { tools?: ExecutableTool[] }): Builder {
+      this._steps.push({
+        id: ObjectId.random(),
+        instructions,
+        tools: options?.tools ?? [],
+      });
+
+      return this;
+    }
+
+    build(): Blueprint {
+      return {
+        steps: this._steps,
+      };
+    }
+  }
+}
+
+/**
+ * Blueprint parser API.
+ */
+export namespace BlueprintParser {
+  export type Step = {
+    instructions: string;
+    // TODO(burdon): Tool DXN? Additional metadata?
+    tools?: string[];
+  };
+
+  export type DSL = {
+    steps: Step[];
+  };
+
+  export const create = (tools: ExecutableTool[] = []) => new Parser(tools);
+
+  class Parser {
+    constructor(private readonly _tools: ExecutableTool[]) {}
+
+    toJSON() {
+      return {
+        tools: this._tools.map((tool) => ({
+          name: tool.name,
+          namespace: tool.namespace,
+          type: tool.type,
+        })),
+      };
+    }
+
+    parse({ steps }: DSL): Blueprint {
+      const builder = BlueprintBuilder.create();
+
+      for (const step of steps) {
+        builder.step(step.instructions, {
+          // TODO(burdon): Tool resolution is duplicated in the session and ollama-client.
+          tools: step.tools?.map(
+            (tool) => this._tools.find(({ id }) => id === tool) ?? raise(new Error(`Tool not found: ${tool}`)),
+          ),
+        });
+      }
+
+      return builder.build();
+    }
+  }
+}

package/src/blueprint/index.ts

@@ -0,0 +1,7 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './blueprint';
+export * from './machine';
+export * from './logger';

package/src/blueprint/logger.ts

@@ -0,0 +1,66 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import chalk from 'chalk';
+
+import { ConsolePrinter } from '@dxos/ai';
+
+import { type Blueprint } from './blueprint';
+import { type BlueprintMachine, type BlueprintMachineState } from './machine';
+
+/* eslint-disable no-console */
+
+// Force chalk colors on for tests
+chalk.level = 2;
+
+export const setConsolePrinter = (machine: BlueprintMachine, extraLogging = false) => {
+  const printer = new ConsolePrinter();
+  if (extraLogging) {
+    machine.stepStart.on((step) =>
+      console.log(
+        `\n${chalk.magenta(`${chalk.bold(`STEP ${machine.blueprint.steps.indexOf(step) + 1} of ${machine.blueprint.steps.length}:`)} ${step.instructions}`)}\n`,
+      ),
+    );
+    machine.message.on((msg) => printer.printMessage(msg));
+    machine.block.on((block) => printer.printContentBlock(block));
+  }
+  machine.begin.on(() => printTrace(machine.blueprint, machine.state));
+  machine.stepComplete.on(() => printTrace(machine.blueprint, machine.state));
+};
+
+const printTrace = (blueprint: Blueprint, state: BlueprintMachineState) => {
+  console.log('\n==============================================\n');
+  console.log(chalk.bold('\nThe Blueprint:'));
+
+  blueprint.steps.forEach((step, index) => {
+    const traceStep = state.trace.find((t) => t.stepId === step.id);
+
+    let color = chalk.gray; // Not executed
+    let bullet = '○';
+    if (traceStep) {
+      switch (traceStep.status) {
+        case 'done':
+          color = chalk.green;
+          bullet = '✓';
+          break;
+        case 'skipped':
+          color = chalk.blue;
+          bullet = '↓';
+          break;
+        case 'bailed':
+          color = chalk.red;
+          bullet = '✗';
+          break;
+      }
+    }
+
+    console.log(color(`\n${bullet} ${step.instructions}`));
+
+    if (traceStep?.comment) {
+      console.log(chalk.white(`    ↳ ${traceStep.comment}`));
+    }
+  });
+  console.log('\n');
+  console.log('\n==============================================\n');
+};

package/src/blueprint/machine.ts

@@ -0,0 +1,208 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Match, Schema } from 'effect';
+
+import {
+  createTool,
+  type ExecutableTool,
+  Message,
+  ToolResult,
+  type AIServiceClient,
+  type MessageContentBlock,
+  type ToolUseContentBlock,
+} from '@dxos/ai';
+import { Event } from '@dxos/async';
+import { create } from '@dxos/echo-schema';
+import { type ObjectId } from '@dxos/keys';
+import { log } from '@dxos/log';
+
+import type { Blueprint, BlueprintStep } from './blueprint';
+import { AISession } from '../session';
+
+export type BlueprintMachineState = {
+  history: Message[];
+  trace: BlueprintTraceStep[];
+  state: 'working' | 'bail' | 'done';
+};
+
+const INITIAL_STATE: BlueprintMachineState = {
+  state: 'working',
+  history: [],
+  trace: [],
+};
+
+export type BlueprintTraceStep = {
+  status: 'done' | 'bailed' | 'skipped';
+  stepId: ObjectId;
+  comment: string;
+};
+
+type ExecutionOptions = {
+  aiService: AIServiceClient;
+
+  /**
+   * Input to the blueprint.
+   */
+  input?: unknown;
+};
+
+/**
+ * Blueprint state machine.
+ */
+export class BlueprintMachine {
+  public readonly begin = new Event<void>();
+  public readonly end = new Event<void>();
+  public readonly stepStart = new Event<BlueprintStep>();
+  public readonly stepComplete = new Event<BlueprintStep>();
+  public readonly message = new Event<Message>();
+  public readonly block = new Event<MessageContentBlock>();
+
+  state: BlueprintMachineState = structuredClone(INITIAL_STATE);
+
+  constructor(readonly blueprint: Blueprint) {}
+
+  async runToCompletion(options: ExecutionOptions): Promise<void> {
+    log.info('runToCompletion', options);
+
+    this.begin.emit();
+    let firstStep = true;
+    while (this.state.state !== 'done') {
+      const input = firstStep ? options.input : undefined;
+
+      firstStep = false;
+      this.state = await this._execStep(this.state, {
+        input,
+        aiService: options.aiService,
+      });
+
+      this.stepComplete.emit(this.blueprint.steps.find((step) => step.id === this.state.trace.at(-1)?.stepId)!);
+
+      if (this.state.state === 'bail') {
+        throw new Error('Agent unable to follow the blueprint');
+      }
+    }
+
+    this.end.emit();
+  }
+
+  private async _execStep(state: BlueprintMachineState, options: ExecutionOptions): Promise<BlueprintMachineState> {
+    const prevStep = this.blueprint.steps.findIndex((step) => step.id === this.state.trace.at(-1)?.stepId);
+    if (prevStep === this.blueprint.steps.length - 1) {
+      throw new Error('Done execution blueprint');
+    }
+
+    const nextStep = this.blueprint.steps[prevStep + 1];
+    const onLastStep = prevStep === this.blueprint.steps.length - 2;
+    this.stepStart.emit(nextStep);
+
+    const ReportSchema = Schema.Struct({
+      status: Schema.Literal('done', 'bailed', 'skipped').annotations({
+        description: 'The status of the task completion',
+      }),
+      comment: Schema.String.annotations({
+        description:
+          'A comment about the task completion. If you bailed, you must explain why in detail (min couple of sentences).',
+      }),
+    });
+
+    const report = createTool('system', {
+      name: 'report',
+      description: 'This tool reports that the agent has completed the task or is unable to do so',
+      schema: ReportSchema,
+      execute: async (input) => ToolResult.Break(input),
+    });
+
+    const session = new AISession({
+      operationModel: 'configured',
+    });
+
+    session.userMessage.pipeInto(this.message);
+    session.message.pipeInto(this.message);
+    session.block.pipeInto(this.block);
+
+    const inputMessages = options.input
+      ? [
+          create(Message, {
+            role: 'user',
+            content: [taggedDataBlock('input', options.input)],
+          }),
+        ]
+      : [];
+    inputMessages.forEach((message) => this.message.emit(message));
+
+    const messages = await session.run({
+      systemPrompt: `
+        You are a smart Rule-Following Agent.
+        Rule-Following Agent executes the command that the user sent in the last message.
+        After doing the work, the Rule-Following Agent calls report tool.
+        If the Rule-Following agent believes that no action is needed, it calls the report tool with "skipped" status.
+        If the Rule-Following Agent is unable to perform the task, it calls the report tool with "bail" status.
+        Rule-Following Agent explains the reason it is unable to perform the task before bailing.
+        The Rule-Following Agent can express creativity and imagination in the way it performs the task.
+        The Rule-Following Agent precisely follows the instructions.
+      `,
+      history: [...state.history, ...inputMessages],
+      tools: [...nextStep.tools, report] as ExecutableTool[], // TODO(burdon): REMOVE CAST!
+      artifacts: [],
+      client: options.aiService,
+      prompt: nextStep.instructions,
+    });
+
+    const { messages: trimmedHistory, call: lastBlock } = popLastToolCall(messages);
+
+    if (!lastBlock) {
+      // TODO(dmaretskyi): Handle this with grace
+      throw new Error('Agent did not call the report tool');
+    }
+
+    const { status, comment } = ReportSchema.pipe(Schema.decodeUnknownSync)(lastBlock.input);
+
+    return {
+      history: [...state.history, ...inputMessages, ...trimmedHistory],
+      trace: [
+        ...state.trace,
+        {
+          stepId: nextStep.id,
+          status,
+          comment,
+        },
+      ],
+      state: Match.value(status).pipe(
+        Match.withReturnType<BlueprintMachineState['state']>(),
+        Match.when('done', () => (onLastStep ? 'done' : 'working')),
+        Match.when('skipped', () => (onLastStep ? 'done' : 'working')),
+        Match.when('bailed', () => 'bail'),
+        Match.exhaustive,
+      ),
+    };
+  }
+}
+
+const taggedDataBlock = (tag: string, data: unknown): MessageContentBlock => {
+  return {
+    type: 'text',
+    text: `<${tag}>\n${JSON.stringify(data, null, 2)}\n</${tag}>`,
+  } satisfies MessageContentBlock;
+};
+
+const popLastToolCall = (messages: Message[]): { messages: Message[]; call?: ToolUseContentBlock } => {
+  const lastBlock = messages.at(-1)?.content.at(-1);
+  if (lastBlock?.type !== 'tool_use') {
+    return { messages };
+  }
+
+  const trimmedHistory = [
+    ...messages.slice(0, -1),
+    {
+      ...messages.at(-1)!,
+      content: [...messages.at(-1)!.content.slice(0, -1)],
+    } satisfies Message,
+  ];
+
+  return {
+    messages: trimmedHistory,
+    call: lastBlock,
+  };
+};

package/src/blueprint/test-data.ts

@@ -0,0 +1,55 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { create } from '@dxos/echo-schema';
+import { DataType } from '@dxos/schema';
+
+// TODO(burdon): Convert to queue.
+export const TEST_EMAILS: DataType.Message[] = [
+  create(DataType.Message, {
+    properties: {
+      subject: 'Introduction email',
+    },
+    sender: {
+      email: 'alice@example.com',
+    },
+    blocks: [
+      {
+        type: 'text',
+        text: "Hi there! I came across your work and would love to connect. I'm working on some interesting projects in the decentralized space and think there could be good collaboration opportunities.",
+      },
+    ],
+    created: new Date().toISOString(),
+  }),
+  create(DataType.Message, {
+    properties: {
+      subject: 'Question about integration',
+    },
+    sender: {
+      email: 'bob@example.com',
+    },
+    blocks: [
+      {
+        type: 'text',
+        text: "Hello, I was wondering if your system supports integration with external APIs? We have some custom services we'd like to connect. Could you explain the process?",
+      },
+    ],
+    created: new Date().toISOString(),
+  }),
+  create(DataType.Message, {
+    properties: {
+      subject: 'MAKE MONEY FAST!!!',
+    },
+    sender: {
+      email: 'spammer123@sketchy.com',
+    },
+    blocks: [
+      {
+        type: 'text',
+        text: 'CONGRATULATIONS!!! You have been selected to receive $10 MILLION dollars! Just send us your bank details and social security number to claim your prize NOW!!!',
+      },
+    ],
+    created: new Date().toISOString(),
+  }),
+];

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

@@ -0,0 +1,95 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { describe, test, beforeAll } from 'vitest';
+
+import { AIServiceEdgeClient } from '@dxos/ai';
+import { AI_SERVICE_ENDPOINT } from '@dxos/ai/testing';
+import { FunctionExecutor, ServiceContainer } from '@dxos/functions';
+import { log } from '@dxos/log';
+import { createTestData } from '@dxos/schema/testing';
+import { range } from '@dxos/util';
+
+import { extractionAnthropicFn } from './extraction-llm-function';
+import { processTranscriptMessage } from '../extraction';
+
+describe.skip('LLM EntityExtraction', () => {
+  let executor: FunctionExecutor;
+
+  beforeAll(async () => {
+    executor = new FunctionExecutor(
+      new ServiceContainer().setServices({
+        ai: {
+          client: new AIServiceEdgeClient({
+            endpoint: AI_SERVICE_ENDPOINT.REMOTE,
+            defaultGenerationOptions: {
+              // model: '@anthropic/claude-sonnet-4-20250514',
+              model: '@anthropic/claude-3-5-haiku-20241022',
+            },
+          }),
+        },
+      }),
+    );
+  });
+
+  test('should process a transcript block', async () => {
+    const { transcriptMessages, documents, contacts } = createTestData();
+    log.info('context', { documents, contacts });
+
+    for (const message of transcriptMessages.slice(0, 1)) {
+      log.info('input', message);
+      const { message: enhancedMessage } = await processTranscriptMessage({
+        input: {
+          message,
+          objects: [...documents, ...Object.values(contacts)],
+        },
+        function: extractionAnthropicFn,
+        executor,
+      });
+      log.info('output', enhancedMessage);
+    }
+  });
+
+  test('computational irreducibility', async () => {
+    const { transcriptWoflram, documents, contacts } = 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({
+          input: {
+            message,
+            objects: [...documents, ...Object.values(contacts)],
+          },
+          function: extractionAnthropicFn,
+          executor,
+        });
+        log.info('output', { message: enhancedMessage.blocks[0], timeElapsed });
+      }),
+    );
+  });
+
+  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,
+        executor,
+      });
+      log.info('output', enhancedMessage);
+    }
+  });
+});