@dxos/assistant 0.8.2 → 0.8.3-main.7f5a14c

package/src/blueprint/machine.ts

@@ -6,21 +6,39 @@ 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';
 
+const SYSTEM_PROMPT = `
+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.
+`;
+
+export type BlueprintTraceStep = {
+  status: 'done' | 'bailed' | 'skipped';
+  stepId: ObjectId;
+  comment: string;
+};
+
 export type BlueprintMachineState = {
   history: Message[];
   trace: BlueprintTraceStep[];
@@ -33,14 +51,8 @@ const INITIAL_STATE: BlueprintMachineState = {
   trace: [],
 };
 
-export type BlueprintTraceStep = {
-  status: 'done' | 'bailed' | 'skipped';
-  stepId: ObjectId;
-  comment: string;
-};
-
 type ExecutionOptions = {
-  aiService: AIServiceClient;
+  aiClient: AIServiceClient;
 
   /**
    * Input to the blueprint.
@@ -48,6 +60,18 @@ type ExecutionOptions = {
   input?: unknown;
 };
 
+export type BlueprintEvent =
+  | { type: 'begin' }
+  | { type: 'end' }
+  | { type: 'step-start'; step: BlueprintStep }
+  | { type: 'step-complete'; step: BlueprintStep }
+  | { type: 'message'; message: Message }
+  | { type: 'block'; block: MessageContentBlock };
+
+export interface BlueprintLogger {
+  log(event: BlueprintEvent): void;
+}
+
 /**
  * Blueprint state machine.
  */
@@ -59,43 +83,62 @@ export class BlueprintMachine {
   public readonly message = new Event<Message>();
   public readonly block = new Event<MessageContentBlock>();
 
-  state: BlueprintMachineState = structuredClone(INITIAL_STATE);
+  // TODO(burdon): Replace events with logger?
+  private logger?: BlueprintLogger;
 
-  constructor(readonly blueprint: Blueprint) {}
+  private _state: BlueprintMachineState = structuredClone(INITIAL_STATE);
+
+  constructor(
+    readonly registry: ToolRegistry,
+    readonly blueprint: Blueprint,
+  ) {}
+
+  get state(): BlueprintMachineState {
+    return this._state;
+  }
+
+  setLogger(logger: BlueprintLogger): this {
+    this.logger = logger;
+    return this;
+  }
 
   async runToCompletion(options: ExecutionOptions): Promise<void> {
     log.info('runToCompletion', options);
 
     this.begin.emit();
+    this.logger?.log({ type: 'begin' });
+
     let firstStep = true;
-    while (this.state.state !== 'done') {
+    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)!);
+      this._state = await this._execStep(this._state, { input, aiClient: options.aiClient });
+
+      const step = this.blueprint.steps.find((step) => step.id === this._state.trace.at(-1)?.stepId)!;
+      this.stepComplete.emit(step);
+      this.logger?.log({ type: 'step-complete', step });
 
-      if (this.state.state === 'bail') {
+      if (this._state.state === 'bail') {
         throw new Error('Agent unable to follow the blueprint');
       }
     }
 
     this.end.emit();
+    this.logger?.log({ type: 'end' });
   }
 
   private async _execStep(state: BlueprintMachineState, options: ExecutionOptions): Promise<BlueprintMachineState> {
-    const prevStep = this.blueprint.steps.findIndex((step) => step.id === this.state.trace.at(-1)?.stepId);
+    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);
+    this.logger?.log({ type: 'step-start', step: nextStep });
 
     const ReportSchema = Schema.Struct({
       status: Schema.Literal('done', 'bailed', 'skipped').annotations({
@@ -109,7 +152,7 @@ export class BlueprintMachine {
 
     const report = createTool('system', {
       name: 'report',
-      description: 'This tool reports that the agent has completed the task or is unable to do so',
+      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),
     });
@@ -130,30 +173,25 @@ export class BlueprintMachine {
           }),
         ]
       : [];
-    inputMessages.forEach((message) => this.message.emit(message));
+    inputMessages.forEach((message) => {
+      this.message.emit(message);
+      this.logger?.log({ type: 'message', 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.
-        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.
-      `,
+      systemPrompt: SYSTEM_PROMPT,
       history: [...state.history, ...inputMessages],
-      tools: [...nextStep.tools, report] as ExecutableTool[], // TODO(burdon): REMOVE CAST!
+      tools: [...tools, report],
       artifacts: [],
-      client: options.aiService,
+      client: options.aiClient,
       prompt: nextStep.instructions,
     });
 
     const { messages: trimmedHistory, call: lastBlock } = popLastToolCall(messages);
-
     if (!lastBlock) {
-      // TODO(dmaretskyi): Handle this with grace
+      // TODO(dmaretskyi): Handle this with grace.
       throw new Error('Agent did not call the report tool');
     }
 

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,19 +318,21 @@ const preprocessSchema = (schema: Schema.Schema.AnyNoContext) => {
 
 export const createGraphWriterTool = ({
   db,
-  schemaTypes,
+  queue,
+  schema,
   onDone = async (x) => x,
 }: {
   db: EchoDatabase;
-  schemaTypes: Schema.Schema.AnyNoContext[];
+  queue?: Queue;
+  schema: Schema.Schema.AnyNoContext[];
   onDone?: (data: AnyEchoObject[]) => Promise<any>;
 }) => {
   return createTool('graph', {
     name: 'writer',
     description: 'Write to the local graph database',
-    schema: createExtractionSchema(schemaTypes),
+    schema: createExtractionSchema(schema),
     execute: async (input) => {
-      const data = await sanitizeObjects(schemaTypes, input as any, db);
+      const data = await sanitizeObjects(schema, input as any, db, queue);
       return ToolResult.Success(await onDone(data as AnyEchoObject[]));
     },
   });

package/src/research/research.test.ts

@@ -20,7 +20,7 @@ import { researchFn } from './research';
 const REMOTE_AI = true;
 const MOCK_SEARCH = false;
 
-describe('Research', () => {
+describe.skip('Research', () => {
   let builder: EchoTestBuilder;
   let db: EchoDatabase;
   let executor: FunctionExecutor;

package/src/research/research.ts

@@ -60,7 +60,7 @@ export const researchFn = defineFunction({
     });
     session.streamEvent.on((event) => log('stream', { event }));
 
-    const graphWriteTool = createGraphWriterTool({ db, schemaTypes: DataTypes });
+    const graphWriteTool = createGraphWriterTool({ db, schema: DataTypes });
     log.info('graphWriteTool', { schema: graphWriteTool.parameters });
 
     // TODO(dmaretskyi): Consider adding this pattern as the "Graph" output mode for the session.