@dxos/assistant 0.8.4-main.f9ba587 → 0.8.4-main.fcfe5033a5

package/src/conversation/session.ts

@@ -0,0 +1,300 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+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 Either from 'effect/Either';
+import { pipe } from 'effect/Function';
+import * as Layer from 'effect/Layer';
+import * as Record from 'effect/Record';
+import * as Runtime from 'effect/Runtime';
+
+import { type OpaqueToolkit, type ToolExecutionService, type ToolResolverService } from '@dxos/ai';
+import { type Blueprint } from '@dxos/blueprints';
+import { Resource } from '@dxos/context';
+import { Database, Feed, Filter, Obj } from '@dxos/echo';
+import { acquireReleaseResource } from '@dxos/effect';
+import { invariant } from '@dxos/invariant';
+import { log } from '@dxos/log';
+import { McpToolkit } from '@dxos/mcp-client';
+import { Operation, type OperationRegistry } from '@dxos/operation';
+import { Message } from '@dxos/types';
+
+import { ToolExecutionServices } from '../functions';
+import {
+  AiRequest,
+  type AiRequestRunError,
+  type AiRequestRunRequirements,
+  type GenerationObserver,
+  createToolkit,
+  formatSystemPrompt,
+} from '../session';
+import { AiContextBinder, AiContextService } from './context';
+
+export interface McpServerConfig {
+  url: string;
+  protocol: 'sse' | 'http';
+  apiKey?: string;
+}
+
+export interface AiSessionRunProps<R = never> {
+  prompt: string;
+  system?: string;
+  observer?: GenerationObserver;
+  toolkit?: OpaqueToolkit.OpaqueToolkit<R>;
+
+  /**
+   * Space-level MCP servers to connect alongside blueprint-defined ones.
+   */
+  mcpServers?: readonly McpServerConfig[];
+}
+
+export type AiSessionOptions = {
+  feed: Feed.Feed;
+  runtime: Runtime.Runtime<Feed.FeedService>;
+  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 Feed.
+ * Executes tools based on AI responses and supports cancellation of in-progress requests.
+ */
+export class AiSession extends Resource {
+  /**
+   * Blueprints and objects bound to the conversation.
+   */
+  private readonly _binder: AiContextBinder;
+  private readonly _feed: Feed.Feed;
+  private readonly _runtime: Runtime.Runtime<Feed.FeedService>;
+
+  public constructor(options: AiSessionOptions) {
+    super();
+    this._feed = options.feed;
+    this._runtime = options.runtime;
+    invariant(this._feed);
+    invariant(this._runtime);
+    this._binder = new AiContextBinder({ feed: this._feed, runtime: this._runtime, registry: options.registry });
+  }
+
+  protected override async _open(): Promise<void> {
+    await this._binder.open(this._ctx);
+  }
+
+  public get feed() {
+    return this._feed;
+  }
+
+  public get context() {
+    return this._binder;
+  }
+
+  public async getHistory(): Promise<Message.Message[]> {
+    const queryResult = await Runtime.runPromise(this._runtime)(Feed.query(this._feed, Filter.type(Message.Message)));
+    const items = await queryResult.run();
+    return items.filter(Obj.instanceOf(Message.Message));
+  }
+
+  getTools(): Effect.Effect<
+    Record<string, import('@effect/ai/Tool').Any>,
+    never,
+    ToolExecutionService | ToolResolverService
+  > {
+    return Effect.gen(this, function* () {
+      const toolkit = yield* createToolkit({ blueprints: this.context.getBlueprints() });
+      return toolkit.toolkit.tools;
+    }).pipe(Effect.orDie);
+  }
+
+  /**
+   * Not provided by default, since users might want to override them.
+   */
+  makeToolExecutionServices(): Layer.Layer<
+    ToolExecutionService | ToolResolverService,
+    never,
+    OpaqueToolkit.OpaqueToolkitProvider | Operation.Service | OperationRegistry.Service
+  > {
+    return ToolExecutionServices.pipe(
+      Layer.provide(Operation.withInvocationOptions({ conversation: Obj.getDXN(this._feed).toString() })),
+    );
+  }
+  /**
+   * Creates a new cancelable request effect.
+   */
+  public createRequest<R = never>(
+    params: AiSessionRunProps<R>,
+  ): Effect.Effect<Message.Message[], AiRequestRunError, AiRequestRunRequirements | R> {
+    return Effect.gen(this, function* () {
+      const history = yield* Effect.promise(() => this.getHistory());
+      const blueprints = this.context.getBlueprints();
+      const objects = this.context.getObjects();
+
+      log('run', {
+        history: history.length,
+        blueprints: blueprints.length,
+        objects: objects.length,
+      });
+
+      const request = new AiRequest({
+        summarizationThreshold: SUMMARY_THRESHOLD,
+        observer: params.observer,
+        onOutput: (message) =>
+          Effect.promise(() => Runtime.runPromise(this._runtime)(Feed.append(this._feed, [message]))),
+      });
+
+      yield* request.begin({
+        history,
+        blueprints,
+        objects,
+        prompt: params.prompt,
+        system: params.system,
+      });
+
+      // Turn loop: recompute toolkit and system prompt between turns to pick up dynamically enabled blueprints.
+      do {
+        yield* Effect.promise(() => this.context.sync());
+        const currentBlueprints = this.context.getBlueprints();
+        const mcps = yield* connectMcpServers(currentBlueprints, params.mcpServers);
+        const toolkit = yield* createToolkit({
+          toolkit: params.toolkit,
+          blueprints: currentBlueprints,
+          opaqueToolkits: mcps,
+        });
+
+        log('toolkit', { tools: Record.keys(toolkit.toolkit.tools) });
+        const system = yield* formatSystemPrompt({
+          system: params.system,
+          blueprints: currentBlueprints,
+          objects: this.context.getObjects(),
+        }).pipe(Effect.orDie);
+
+        const { done } = yield* request.runAgentTurn({ system, toolkit });
+        if (done) {
+          break;
+        }
+
+        yield* request.runTools({ toolkit });
+      } while (true);
+
+      log('result', {
+        messages: request.pending.length,
+        duration: request.duration,
+        toolCalls: request.toolCalls,
+      });
+
+      return [...request.pending];
+    }).pipe(
+      Effect.provide(
+        Layer.mergeAll(
+          Layer.succeed(AiContextService, {
+            binder: this.context,
+          }),
+          Layer.succeed(AiSessionService, this),
+          Operation.withInvocationOptions({ conversation: Obj.getDXN(this._feed).toString() }),
+        ),
+      ),
+      Effect.withSpan('AiSession.createRequest'),
+    );
+  }
+}
+
+/**
+ * Gives access to the ai session.
+ */
+export class AiSessionService extends Context.Tag('@dxos/assistant/AiSessionService')<AiSessionService, AiSession>() {
+  /**
+   * Create a new session layer from options.
+   */
+  static layer = (options: AiSessionOptions): Layer.Layer<AiSessionService | AiContextService> =>
+    aiContextFromSession.pipe(
+      Layer.provideMerge(
+        Layer.scoped(
+          AiSessionService,
+          Effect.gen(function* () {
+            const session = yield* acquireReleaseResource(() => new AiSession(options));
+            return session;
+          }),
+        ),
+      ),
+    );
+
+  /**
+   * Create a new session with a new feed.
+   */
+  static layerNewFeed = (
+    options?: Omit<AiSessionOptions, 'feed' | 'runtime'>,
+  ): Layer.Layer<AiSessionService | AiContextService, never, Database.Service | Feed.FeedService> =>
+    Layer.unwrapScoped(
+      Effect.gen(function* () {
+        const feed = Feed.make();
+        yield* Database.add(feed);
+        const runtime = yield* Effect.runtime<Feed.FeedService>();
+        return AiSessionService.layer({
+          ...options,
+          feed,
+          runtime,
+        });
+      }),
+    );
+
+  /**
+   * Run a prompt in the current session.
+   */
+  static run = <R = never>(
+    params: AiSessionRunProps<R>,
+  ): Effect.Effect<Message.Message[], AiRequestRunError, AiRequestRunRequirements | AiSessionService | R> =>
+    Effect.gen(function* () {
+      const session = yield* AiSessionService;
+      return yield* session.createRequest(params);
+    });
+}
+
+const aiContextFromSession = Layer.effect(
+  AiContextService,
+  Effect.gen(function* () {
+    const session = yield* AiSessionService;
+    return {
+      binder: session.context,
+    };
+  }),
+);
+
+const connectMcpServers = (
+  blueprints: readonly Blueprint.Blueprint[],
+  spaceMcpServers: readonly McpServerConfig[] = [],
+): Effect.Effect<OpaqueToolkit.OpaqueToolkit[]> => {
+  const blueprintServers: McpToolkit.McpToolkitOptions[] = pipe(
+    blueprints,
+    Array.flatMap((_) => _.mcpServers ?? []),
+    Array.map(({ url, protocol }) => ({ url, kind: protocol })),
+  );
+  const spaceServers: McpToolkit.McpToolkitOptions[] = spaceMcpServers.map(({ url, protocol, apiKey }) => ({
+    url,
+    kind: protocol,
+    apiKey,
+  }));
+  const allServers = [...blueprintServers, ...spaceServers];
+
+  return pipe(
+    allServers,
+    Effect.forEach((options) =>
+      McpToolkit.make(options).pipe(
+        // NOTE: Type-inference fails here without explicit void return.
+        Effect.tap((toolkit): void =>
+          log.info('Connected to MCP server', { url: options.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

@@ -0,0 +1,74 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import { createDocAccessor } from '@dxos/echo-db';
+import { EchoTestBuilder } from '@dxos/echo-db/testing';
+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 { db } = await builder.createDatabase({ types: [Text.Text] });
+
+    const document = trim`
+      # Hello World
+      There'z a typo in this sentence.
+      And another oon here.
+      But not this one.
+    `;
+
+    const text = db.add(Text.make({ content: document }));
+    const accessor = createDocAccessor(text, ['content']);
+    const result = applyDiffs(accessor, [
+      `- There'z a typo in this sentence.`,
+      `+ There is a typo in this sentence.`,
+      `- And another oon here.`,
+      `+ And another one here.`,
+    ]);
+
+    expect(result).toBe(trim`
+      # Hello World
+      There is a typo in this sentence.
+      And another one here.
+      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({ content: 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({ content: '' }));
+    const accessor = createDocAccessor(text, ['content']);
+    const result = applyDiffs(accessor, ['+ # Shopping list']);
+
+    expect(result).toBe('# Shopping list');
+  });
+});

package/src/diff.ts

@@ -0,0 +1,93 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { next as A, type Doc } from '@automerge/automerge';
+
+import { DocAccessor, toCursorRange } from '@dxos/echo-db';
+import { log } from '@dxos/log';
+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 = {
+  match: string;
+  replace: string;
+};
+
+/**
+ * Convert a list of diffs into a list of operations.
+ */
+export const reduceDiffs = (diffs: readonly string[]): Diff[] => {
+  return diffs.reduce<Diff[]>((acc, diff) => {
+    const current = acc[acc.length - 1];
+    const [op, ...rest] = diff.split(' ');
+    switch (op) {
+      case '-': {
+        if (current && !current.match) {
+          current.match += rest.join(' ');
+        } else {
+          acc.push({ match: rest.join(' '), replace: '' });
+        }
+        return acc;
+      }
+      case '+': {
+        if (current && !current.replace) {
+          current.replace += rest.join(' ');
+        } else {
+          acc.push({ match: '', replace: rest.join(' ') });
+        }
+        return acc;
+      }
+      default: {
+        throw new Error(`Invalid diff format: each line must start with "- " or "+ ", got: ${JSON.stringify(diff)}`);
+      }
+    }
+  }, []);
+};
+
+export const computeDiffsWithCursors = <T>(accessor: DocAccessor<T>, diffs: readonly string[]) => {
+  return reduceDiffs(diffs)
+    .map((diff) => {
+      const text = DocAccessor.getValue<string>(accessor);
+      const idx = text.indexOf(diff.match);
+      if (idx !== -1) {
+        return { cursor: toCursorRange(accessor, idx, idx + diff.match.length), text: diff.replace };
+      }
+    })
+    .filter(isNonNullable);
+};
+
+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);
+      const idx = text.indexOf(diff.match);
+      if (idx !== -1) {
+        // TODO(burdon): Replace smallest substring.
+        // NOTE: This only replaces the first match.
+        A.splice(doc, accessor.path as A.Prop[], idx, diff.match.length, diff.replace);
+      } else {
+        if (errorOnNotFound) {
+          throw new Error(`Diff not found: ${diff.match}`);
+        } else {
+          log.warn('diff not found', diff);
+        }
+      }
+    });
+  }
+
+  return DocAccessor.getValue<string>(accessor);
+};

package/src/errors.ts

@@ -0,0 +1,15 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { BaseError } from '@dxos/errors';
+
+/**
+ * Generic error from AI model.
+ */
+export class AiModelError extends BaseError.extend('AiModelError', 'AI model error') {}
+
+/**
+ * Generic error for AI agent execution.
+ */
+export class AiAssistantError extends BaseError.extend('AiAssistantError', 'AI assistant error') {}