@dxos/assistant 0.8.4-main.c4373fc → 0.8.4-main.c85a9c8dae

package/src/conversation/conversation.ts

@@ -2,16 +2,28 @@
 // Copyright 2025 DXOS.org
 //
 
+import type * as Tool from '@effect/ai/Tool';
 import type * as Toolkit from '@effect/ai/Toolkit';
+import type { Registry } from '@effect-atom/atom-react';
 import * as Array from 'effect/Array';
+import * as Context from 'effect/Context';
 import * as Effect from 'effect/Effect';
-import * as Option from 'effect/Option';
-
+import * as Either from 'effect/Either';
+import { pipe } from 'effect/Function';
+import * as Layer from 'effect/Layer';
+
+import { type ToolExecutionService, type ToolResolverService } from '@dxos/ai';
+import { type GenericToolkit } from '@dxos/ai';
+import { type Blueprint } from '@dxos/blueprints';
+import { Resource } from '@dxos/context';
 import { Obj } from '@dxos/echo';
 import { type Queue } from '@dxos/echo-db';
-import { DatabaseService } from '@dxos/functions';
+import { acquireReleaseResource } from '@dxos/effect';
+import { QueueService } from '@dxos/functions';
+import { invariant } from '@dxos/invariant';
 import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { McpToolkit } from '@dxos/mcp-client';
+import { Message } from '@dxos/types';
 
 import {
   AiSession,
@@ -23,99 +35,213 @@ import {
 
 import { AiContextBinder, AiContextService, type ContextBinding } from './context';
 
-export interface AiConversationRunParams {
+export interface AiConversationRunProps {
   prompt: string;
   system?: string;
   observer?: GenerationObserver;
 }
 
 export type AiConversationOptions = {
-  queue: Queue<DataType.Message | ContextBinding>;
+  queue: Queue<Message.Message | ContextBinding>;
+  toolkit?: Toolkit.Any;
+  registry?: Registry.Registry;
 };
 
+// TODO(dmaretskyi): Take from model characteristics. opus has 200k max tokens.
+/**
+ * Summarization threshold in tokens.
+ */
+const SUMMARY_THRESHOLD = 80_000;
+
 /**
  * Durable conversation state (initiated by users and agents) backed by a Queue.
  * Executes tools based on AI responses and supports cancellation of in-progress requests.
  */
-export class AiConversation {
-  /**
-   * Message and binding queue.
-   */
-  private readonly _queue: Queue<DataType.Message | ContextBinding>;
-
-  /**
-   * Blueprints bound to the conversation.
-   */
-  private readonly _context: AiContextBinder;
-
+export class AiConversation extends Resource {
   /**
-   * Toolkit from the current session request.
+   * Blueprints and objects bound to the conversation.
    */
-  private _toolkit: Toolkit.WithHandler<any> | undefined;
+  private readonly _binder: AiContextBinder;
+  private readonly _queue: Queue<Message.Message | ContextBinding>;
+  private readonly _toolkit?: Toolkit.Any;
 
   public constructor(options: AiConversationOptions) {
+    super();
     this._queue = options.queue;
-    this._context = new AiContextBinder(this._queue);
+    this._toolkit = options.toolkit;
+    invariant(this._queue);
+    this._binder = new AiContextBinder({ queue: this._queue, registry: options.registry });
+  }
+
+  protected override async _open(): Promise<void> {
+    await this._binder.open(this._ctx);
+  }
+
+  public get queue() {
+    return this._queue;
   }
 
   public get context() {
-    return this._context;
+    return this._binder;
   }
 
   public get toolkit() {
     return this._toolkit;
   }
 
-  public async getHistory(): Promise<DataType.Message[]> {
-    const queueItems = await this._queue.queryObjects();
-    return queueItems.filter(Obj.instanceOf(DataType.Message));
+  public async getHistory(): Promise<Message.Message[]> {
+    const queueItems = await this._queue.queryObjects(); // TODO(burdon): Update.
+    return queueItems.filter(Obj.instanceOf(Message.Message));
+  }
+
+  getTools(): Effect.Effect<Record<string, Tool.Any>, never, ToolExecutionService | ToolResolverService> {
+    return Effect.gen(this, function* () {
+      const blueprints = this.context.getBlueprints();
+      const tookit = yield* createToolkit({ toolkit: this._toolkit, blueprints });
+      return tookit.tools;
+    }).pipe(Effect.orDie);
   }
 
   /**
    * Creates a new cancelable request effect.
    */
-  createRequest(
-    params: AiConversationRunParams,
-  ): Effect.Effect<DataType.Message[], AiSessionRunError, AiSessionRunRequirements> {
+  public createRequest(
+    params: AiConversationRunProps,
+  ): Effect.Effect<Message.Message[], AiSessionRunError, AiSessionRunRequirements> {
     return Effect.gen(this, function* () {
-      const session = new AiSession();
       const history = yield* Effect.promise(() => this.getHistory());
 
-      // Get context objects.
-      const context = yield* Effect.promise(() => this.context.query());
-      const blueprints = yield* Effect.forEach(context.blueprints.values(), DatabaseService.loadOption).pipe(
-        Effect.map(Array.filter(Option.isSome)),
-        Effect.map(Array.map((option) => option.value)),
-      );
-      const objects = yield* Effect.forEach(context.objects.values(), DatabaseService.loadOption).pipe(
-        Effect.map(Array.filter(Option.isSome)),
-        Effect.map(Array.map((option) => option.value)),
-      );
-
       // Create toolkit.
-      const toolkit = yield* createToolkit({ blueprints });
-      this._toolkit = toolkit;
+      const blueprints = this.context.getBlueprints();
+
+      const mcps = yield* connectMcpServers(blueprints);
+
+      const toolkit = yield* createToolkit({
+        toolkit: this._toolkit,
+        blueprints,
+        genericToolkits: mcps,
+      });
+
+      // Context objects.
+      const objects = this.context.getObjects();
 
-      const start = Date.now();
       log('run', {
         history: history.length,
         blueprints: blueprints.length,
+        tools: Object.keys(toolkit.tools).length,
         objects: objects.length,
-        tools: this._toolkit?.tools.length ?? 0,
       });
 
       // Process request.
-      const messages = yield* session.run({ history, blueprints, objects, toolkit, ...params }).pipe(
-        Effect.provideService(AiContextService, {
-          binder: this.context,
-        }),
-      );
-
-      log('result', { messages: messages.length, duration: Date.now() - start });
+      const session = new AiSession({
+        summarizationThreshold: SUMMARY_THRESHOLD,
+      });
+      const messages = yield* session
+        .run({
+          history,
+          blueprints,
+          toolkit,
+          objects,
+          ...params,
+          onOutput: (message) => Effect.promise(() => this._queue.append([message])),
+        })
+        .pipe(
+          Effect.provide(
+            Layer.mergeAll(
+              Layer.succeed(AiContextService, {
+                binder: this.context,
+              }),
+              Layer.succeed(AiConversationService, this),
+            ),
+          ),
+        );
+
+      log('result', {
+        messages: messages.length,
+        duration: session.duration,
+        toolCalls: session.toolCalls,
+      });
 
-      // Append to queue.
-      yield* Effect.promise(() => this._queue.append(messages));
       return messages;
-    }).pipe(Effect.withSpan('AiConversation.request'));
+    });
   }
 }
+
+/**
+ * Gives access to the ai conversation.
+ */
+export class AiConversationService extends Context.Tag('@dxos/assistant/AiConversationService')<
+  AiConversationService,
+  AiConversation
+>() {
+  /**
+   * Create a new conversation layer from options.
+   */
+  static layer = (options: AiConversationOptions): Layer.Layer<AiConversationService | AiContextService> =>
+    aiContextFromConversation.pipe(
+      Layer.provideMerge(
+        Layer.scoped(
+          AiConversationService,
+          Effect.gen(function* () {
+            const conversation = yield* acquireReleaseResource(() => new AiConversation(options));
+            return conversation;
+          }),
+        ),
+      ),
+    );
+
+  /**
+   * Create a new conversation with a new queue.
+   */
+  static layerNewQueue = (
+    options?: Omit<AiConversationOptions, 'queue'>,
+  ): Layer.Layer<AiConversationService | AiContextService, never, QueueService> =>
+    Layer.unwrapScoped(
+      Effect.gen(function* () {
+        return AiConversationService.layer({
+          ...options,
+          queue: yield* QueueService.createQueue<Message.Message | ContextBinding>(),
+        });
+      }),
+    );
+
+  /**
+   * Run a prompt in the current conversation.
+   */
+  static run = (
+    params: AiConversationRunProps,
+  ): Effect.Effect<Message.Message[], AiSessionRunError, AiSessionRunRequirements | AiConversationService> =>
+    Effect.gen(function* () {
+      const conversation = yield* AiConversationService;
+      return yield* conversation.createRequest(params);
+    });
+}
+
+const aiContextFromConversation = Layer.effect(
+  AiContextService,
+  Effect.gen(function* () {
+    const conversation = yield* AiConversationService;
+    return {
+      binder: conversation.context,
+    };
+  }),
+);
+
+const connectMcpServers = (
+  blueprints: readonly Blueprint.Blueprint[],
+): Effect.Effect<GenericToolkit.GenericToolkit[]> =>
+  pipe(
+    blueprints,
+    Array.flatMap((_) => _.mcpServers ?? []),
+    Effect.forEach(({ url, protocol }) =>
+      McpToolkit.make({ url, kind: protocol }).pipe(
+        // NOTE: Type-inference fails here without explicit void return.
+        Effect.tap((toolkit): void =>
+          log.info('Connected to MCP server', { url, tools: Object.keys(toolkit.toolkit.tools).length }),
+        ),
+        Effect.tapDefect((error) => Effect.sync(() => log.warn('Failed to connect to MCP server', { error }))),
+        Effect.either,
+      ),
+    ),
+    Effect.map(Array.filterMap((_) => Either.getRight(_))),
+  );

package/src/diff.test.ts

@@ -2,20 +2,27 @@
 // Copyright 2025 DXOS.org
 //
 
-import { describe, expect, it } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
 import { createDocAccessor } from '@dxos/echo-db';
 import { EchoTestBuilder } from '@dxos/echo-db/testing';
-import { DataType } from '@dxos/schema';
+import { Text } from '@dxos/schema';
 import { trim } from '@dxos/util';
 
 import { applyDiffs } from './diff';
 
 describe('diff', () => {
+  let builder: EchoTestBuilder;
+  beforeEach(async () => {
+    builder = await new EchoTestBuilder().open();
+  });
+
+  afterEach(async () => {
+    await builder.close();
+  });
+
   it('should apply diffs', async () => {
-    const builder = new EchoTestBuilder();
-    const { db, graph } = await builder.createDatabase();
-    graph.schemaRegistry.addSchema([DataType.Text]);
+    const { db } = await builder.createDatabase({ types: [Text.Text] });
 
     const document = trim`
       # Hello World
@@ -24,7 +31,7 @@ describe('diff', () => {
       But not this one.
     `;
 
-    const text = db.add(DataType.makeText(document));
+    const text = db.add(Text.make(document));
     const accessor = createDocAccessor(text, ['content']);
     const result = applyDiffs(accessor, [
       `- There'z a typo in this sentence.`,
@@ -40,4 +47,28 @@ describe('diff', () => {
       But not this one.
     `);
   });
+
+  it('should apply diffs that append content at the end', async () => {
+    const { db } = await builder.createDatabase({ types: [Text.Text] });
+
+    const document = '# Shopping list';
+
+    const text = db.add(Text.make(document));
+    const accessor = createDocAccessor(text, ['content']);
+    const result = applyDiffs(accessor, ['- # Shopping list', '+ # Shopping list\n- Milk']);
+
+    expect(result).toBe('# Shopping list\n- Milk');
+  });
+
+  it('should append text to an empty document', async () => {
+    const builder = new EchoTestBuilder();
+    const { db, graph } = await builder.createDatabase();
+    await graph.schemaRegistry.register([Text.Text]);
+
+    const text = db.add(Text.make(''));
+    const accessor = createDocAccessor(text, ['content']);
+    const result = applyDiffs(accessor, ['+ # Shopping list']);
+
+    expect(result).toBe('# Shopping list');
+  });
 });

package/src/diff.ts

@@ -11,6 +11,11 @@ import { isNonNullable } from '@dxos/util';
 // Relevant info:
 // https://github.com/Minahil-official/GPT-4.1-Promting-Guide?tab=readme-ov-file#-apply-patch-tool-format-v4a-diff
 
+// TODO(burdon): Anthropic only.
+//  https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/text-editor-tool#example-str-replace-command
+//  AI_TOOL_NOT_FOUND: str_replace_based_edit_tool
+// 'str_replace_based_edit_tool',
+
 // TOOD(burdon): Support append to end (no "-").
 
 export type Diff = {
@@ -42,9 +47,10 @@ export const reduceDiffs = (diffs: readonly string[]): Diff[] => {
         }
         return acc;
       }
+      default: {
+        throw new Error(`Invalid diff format: each line must start with "- " or "+ ", got: ${JSON.stringify(diff)}`);
+      }
     }
-
-    return acc;
   }, []);
 };
 
@@ -60,7 +66,11 @@ export const computeDiffsWithCursors = <T>(accessor: DocAccessor<T>, diffs: read
     .filter(isNonNullable);
 };
 
-export const applyDiffs = <T>(accessor: DocAccessor<T>, diffs: readonly string[]): string => {
+export const applyDiffs = <T>(
+  accessor: DocAccessor<T>,
+  diffs: readonly string[],
+  { errorOnNotFound = false }: { errorOnNotFound?: boolean } = {},
+): string => {
   for (const diff of reduceDiffs(diffs)) {
     accessor.handle.change((doc: Doc<T>) => {
       const text = DocAccessor.getValue<string>(accessor);
@@ -70,7 +80,11 @@ export const applyDiffs = <T>(accessor: DocAccessor<T>, diffs: readonly string[]
         // NOTE: This only replaces the first match.
         A.splice(doc, accessor.path as A.Prop[], idx, diff.match.length, diff.replace);
       } else {
-        log.warn('diff not found', diff);
+        if (errorOnNotFound) {
+          throw new Error(`Diff not found: ${diff.match}`);
+        } else {
+          log.warn('diff not found', diff);
+        }
       }
     });
   }

package/src/errors.ts

@@ -7,9 +7,9 @@ import { BaseError } from '@dxos/errors';
 /**
  * Generic error from AI model.
  */
-export class AiModelError extends BaseError.extend('AI_MODEL_ERROR', 'AI model error') {}
+export class AiModelError extends BaseError.extend('AiModelError', 'AI model error') {}
 
 /**
  * Generic error for AI agent execution.
  */
-export class AiAssistantError extends BaseError.extend('AI_ASSISTANT_ERROR', 'AI assistant error') {}
+export class AiAssistantError extends BaseError.extend('AiAssistantError', 'AI assistant error') {}