@dxos/assistant 0.8.2 → 0.8.3-main.672df60

package/src/blueprint/logger.ts

@@ -2,41 +2,84 @@
 // Copyright 2025 DXOS.org
 //
 
+import { type ReadonlySignal, signal } from '@preact/signals-core';
 import chalk from 'chalk';
 
 import { ConsolePrinter } from '@dxos/ai';
+import { type CleanupFn, combine } from '@dxos/async';
 
 import { type Blueprint } from './blueprint';
 import { type BlueprintMachine, type BlueprintMachineState } from './machine';
 
 /* eslint-disable no-console */
 
-// Force chalk colors on for tests
+// Force chalk colors on for tests.
 chalk.level = 2;
 
-export const setConsolePrinter = (machine: BlueprintMachine, extraLogging = false) => {
+export class Logger {
+  private _messages = signal<string[]>([]);
+
+  get messages(): ReadonlySignal<string[]> {
+    return this._messages;
+  }
+
+  clear() {
+    this._messages.value = [];
+  }
+
+  log(message: string) {
+    this._messages.value = [...this._messages.value, message];
+  }
+}
+
+export const setLogger = (machine: BlueprintMachine, logger: Logger): CleanupFn => {
+  return combine(
+    machine.begin.on(() => {
+      logger.log('Starting...');
+    }),
+    machine.stepStart.on((step) => {
+      const index = machine.blueprint.steps.indexOf(step);
+      logger.log(`Step ${index + 1} of ${machine.blueprint.steps.length}`);
+      logger.log(`Instructions: ${step.instructions}`);
+    }),
+    machine.stepComplete.on((step) => {
+      const trace = machine.state.trace.find((t) => t.stepId === step.id);
+      logger.log(`Complete: ${trace?.comment}`);
+    }),
+    machine.end.on(() => {
+      logger.log('OK');
+    }),
+  );
+};
+
+export const setConsolePrinter = (machine: BlueprintMachine, verbose = false): CleanupFn => {
   const printer = new ConsolePrinter();
-  if (extraLogging) {
+  return combine(
+    machine.begin.on(() => printTrace(machine.blueprint, machine.state)),
     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));
+    ),
+    machine.stepComplete.on(() => printTrace(machine.blueprint, machine.state)),
+    machine.end.on(() => console.log('DONE')),
+    verbose
+      ? [
+          machine.message.on((msg) => printer.printMessage(msg)),
+          machine.block.on((block) => printer.printContentBlock(block)),
+        ]
+      : [],
+  );
 };
 
-const printTrace = (blueprint: Blueprint, state: BlueprintMachineState) => {
-  console.log('\n==============================================\n');
-  console.log(chalk.bold('\nThe Blueprint:'));
+const BREAK_LINE = `\n${'_'.repeat(80)}\n`;
 
-  blueprint.steps.forEach((step, index) => {
+const printTrace = (blueprint: Blueprint, state: BlueprintMachineState) => {
+  console.group(chalk.bold('\nBlueprint'));
+  blueprint.steps.forEach((step) => {
     const traceStep = state.trace.find((t) => t.stepId === step.id);
 
-    let color = chalk.gray; // Not executed
+    let color = chalk.gray; // Not executed.
     let bullet = '○';
     if (traceStep) {
       switch (traceStep.status) {
@@ -56,11 +99,11 @@ const printTrace = (blueprint: Blueprint, state: BlueprintMachineState) => {
     }
 
     console.log(color(`\n${bullet} ${step.instructions}`));
-
     if (traceStep?.comment) {
-      console.log(chalk.white(`    ↳ ${traceStep.comment}`));
+      console.log(chalk.white(`  ↳ ${traceStep.comment}`));
     }
   });
-  console.log('\n');
-  console.log('\n==============================================\n');
+
+  console.groupEnd();
+  console.log(BREAK_LINE);
 };

package/src/blueprint/machine.ts

@@ -6,21 +6,28 @@ import { Match, Schema } from 'effect';
 
 import {
   createTool,
-  type ExecutableTool,
-  Message,
-  ToolResult,
   type AIServiceClient,
+  Message,
   type MessageContentBlock,
+  ToolResult,
   type ToolUseContentBlock,
+  type ToolRegistry,
 } 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 { isNonNullable } from '@dxos/util';
 
 import type { Blueprint, BlueprintStep } from './blueprint';
 import { AISession } from '../session';
 
+export type BlueprintTraceStep = {
+  status: 'done' | 'bailed' | 'skipped';
+  stepId: ObjectId;
+  comment: string;
+};
+
 export type BlueprintMachineState = {
   history: Message[];
   trace: BlueprintTraceStep[];
@@ -33,12 +40,6 @@ const INITIAL_STATE: BlueprintMachineState = {
   trace: [],
 };
 
-export type BlueprintTraceStep = {
-  status: 'done' | 'bailed' | 'skipped';
-  stepId: ObjectId;
-  comment: string;
-};
-
 type ExecutionOptions = {
   aiService: AIServiceClient;
 
@@ -61,7 +62,10 @@ export class BlueprintMachine {
 
   state: BlueprintMachineState = structuredClone(INITIAL_STATE);
 
-  constructor(readonly blueprint: Blueprint) {}
+  constructor(
+    readonly registry: ToolRegistry,
+    readonly blueprint: Blueprint,
+  ) {}
 
   async runToCompletion(options: ExecutionOptions): Promise<void> {
     log.info('runToCompletion', options);
@@ -132,6 +136,8 @@ export class BlueprintMachine {
       : [];
     inputMessages.forEach((message) => this.message.emit(message));
 
+    // TODO(wittjosiah): Warn if tool is not found.
+    const tools = nextStep.tools.map((tool) => this.registry.get(tool)).filter(isNonNullable);
     const messages = await session.run({
       systemPrompt: `
         You are a smart Rule-Following Agent.
@@ -144,7 +150,7 @@ export class BlueprintMachine {
         The Rule-Following Agent precisely follows the instructions.
       `,
       history: [...state.history, ...inputMessages],
-      tools: [...nextStep.tools, report] as ExecutableTool[], // TODO(burdon): REMOVE CAST!
+      tools: [...tools, report],
       artifacts: [],
       client: options.aiService,
       prompt: nextStep.instructions,

package/src/research/graph.ts

@@ -5,7 +5,8 @@
 import { Schema, identity, Option, SchemaAST } from 'effect';
 
 import { createTool, ToolResult } from '@dxos/ai';
-import { type EchoDatabase } from '@dxos/echo-db';
+import type { Obj, Relation } from '@dxos/echo';
+import { type EchoDatabase, type Queue } from '@dxos/echo-db';
 import { isEncodedReference } from '@dxos/echo-protocol';
 import {
   EntityKind,
@@ -20,14 +21,16 @@ import {
   ObjectId,
   Query,
   ReferenceAnnotationId,
+  type BaseObject,
+  RelationSourceDXNId,
+  RelationTargetDXNId,
   RelationSourceId,
   RelationTargetId,
-  type BaseObject,
 } from '@dxos/echo-schema';
 import { mapAst } from '@dxos/effect';
 import { DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
-import { deepMapValues } from '@dxos/util';
+import { deepMapValues, isNonNullable } from '@dxos/util';
 
 // TODO(burdon): Unify with the graph schema.
 export const Subgraph = Schema.Struct({
@@ -95,7 +98,7 @@ const isSchemaAddressableByDxn = (schema: Schema.Schema.AnyNoContext, dxn: DXN):
 /**
  * Perform vector search in the local database.
  */
-export const createLocalSearchTool = (db: EchoDatabase) => {
+export const createLocalSearchTool = (db: EchoDatabase, queue?: Queue) => {
   return createTool('search', {
     name: 'local_search',
     description: 'Search the local database for information using a vector index',
@@ -106,9 +109,16 @@ export const createLocalSearchTool = (db: EchoDatabase) => {
     }),
     execute: async ({ query }) => {
       const { objects } = await db.query(Query.select(Filter.text(query, { type: 'vector' }))).run();
+      const results = [...objects];
+      if (queue) {
+        const queueObjects = await queue.queryObjects();
+        // TODO(dmaretskyi): Text search on the queue.
+        results.push(...queueObjects);
+      }
+
       return ToolResult.Success(`
         <local_context>
-          ${JSON.stringify(objects, null, 2)}
+          ${JSON.stringify(results, null, 2)}
         </local_context>
         `);
     },
@@ -141,6 +151,7 @@ export const sanitizeObjects = async (
   types: Schema.Schema.AnyNoContext[],
   data: Record<string, readonly unknown[]>,
   db: EchoDatabase,
+  queue?: Queue,
 ): Promise<BaseObject[]> => {
   const entries = types
     .map(
@@ -153,8 +164,8 @@ export const sanitizeObjects = async (
     .flat();
 
   const idMap = new Map<string, string>();
-
   const existingIds = new Set<ObjectId>();
+  const enitties = new Map<ObjectId, Obj.Any | Relation.Any>();
 
   const resolveId = (id: string): DXN | undefined => {
     if (ObjectId.isValid(id)) {
@@ -210,23 +221,58 @@ export const sanitizeObjects = async (
         }
         delete data.source;
         delete data.target;
-        data[RelationSourceId] = sourceDxn;
-        data[RelationTargetId] = targetDxn;
+        data[RelationSourceDXNId] = sourceDxn;
+        data[RelationTargetDXNId] = targetDxn;
       }
 
-      return create(entry.schema, data);
+      return {
+        data,
+        schema: entry.schema,
+      };
     })
-    .filter((object) => !existingIds.has(object.id)); // TODO(dmaretskyi): This dissallows updating existing objects.
+    .filter((object) => !existingIds.has(object.data.id)); // TODO(dmaretskyi): This dissallows updating existing objects.
+
+  // TODO(dmaretskyi): Use ref resolver.
+  const { objects: dbObjects } = await db.query(Query.select(Filter.ids(...existingIds))).run();
+  const queueObjects = (await queue?.getObjectsById([...existingIds])) ?? [];
+  const objects = [...dbObjects, ...queueObjects].filter(isNonNullable);
 
-  const { objects } = await db.query(Query.select(Filter.ids(...existingIds))).run();
   // TODO(dmaretskyi): Returns everything if IDs are empty!
-  log.info('objects', { objects, existingIds });
+  log.info('objects', { dbObjects, queueObjects, existingIds });
   const missing = Array.from(existingIds).filter((id) => !objects.some((object) => object.id === id));
   if (missing.length > 0) {
     throw new Error(`Object IDs do not point to existing objects: ${missing.join(', ')}`);
   }
 
-  return res;
+  return res.flatMap(({ data, schema }) => {
+    let skip = false;
+    if (RelationSourceDXNId in data) {
+      const id = (data[RelationSourceDXNId] as DXN).asEchoDXN()?.echoId;
+      const obj = objects.find((object) => object.id === id) ?? enitties.get(id!);
+      if (obj) {
+        delete data[RelationSourceDXNId];
+        data[RelationSourceId] = obj;
+      } else {
+        skip = true;
+      }
+    }
+    if (RelationTargetDXNId in data) {
+      const id = (data[RelationTargetDXNId] as DXN).asEchoDXN()?.echoId;
+      const obj = objects.find((object) => object.id === id) ?? enitties.get(id!);
+      if (obj) {
+        delete data[RelationTargetDXNId];
+        data[RelationTargetId] = obj;
+      } else {
+        skip = true;
+      }
+    }
+    if (!skip) {
+      const obj = create(schema, data);
+      enitties.set(obj.id, obj);
+      return [obj];
+    }
+    return [];
+  });
 };
 
 const SoftRef = Schema.Struct({
@@ -272,10 +318,12 @@ const preprocessSchema = (schema: Schema.Schema.AnyNoContext) => {
 
 export const createGraphWriterTool = ({
   db,
+  queue,
   schemaTypes,
   onDone = async (x) => x,
 }: {
   db: EchoDatabase;
+  queue?: Queue;
   schemaTypes: Schema.Schema.AnyNoContext[];
   onDone?: (data: AnyEchoObject[]) => Promise<any>;
 }) => {
@@ -284,7 +332,7 @@ export const createGraphWriterTool = ({
     description: 'Write to the local graph database',
     schema: createExtractionSchema(schemaTypes),
     execute: async (input) => {
-      const data = await sanitizeObjects(schemaTypes, input as any, db);
+      const data = await sanitizeObjects(schemaTypes, input as any, db, queue);
       return ToolResult.Success(await onDone(data as AnyEchoObject[]));
     },
   });