@dxos/assistant 0.8.3 → 0.8.4-main.28f8d3d

package/src/functions/services.ts

@@ -0,0 +1,114 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { AiTool, type AiToolkit } from '@effect/ai';
+import { Context, Effect, Layer, Match, Predicate, Record, Schema } from 'effect';
+
+import { AiToolNotFoundError, ToolExecutionService, ToolResolverService } from '@dxos/ai';
+import { type FunctionDefinition, LocalFunctionExecutionService } from '@dxos/functions';
+import { invariant } from '@dxos/invariant';
+
+export const makeToolResolverFromFunctions = (
+  functions: FunctionDefinition<any, any>[],
+  toolkit: AiToolkit.Any,
+): Layer.Layer<ToolResolverService> => {
+  return Layer.succeed(ToolResolverService, {
+    resolve: Effect.fn('resolveTool')(function* (id) {
+      const tool = toolkit.tools[id];
+      if (tool) {
+        return tool;
+      }
+
+      const fn = functions.find((f) => f.name === id);
+      if (!fn) {
+        return yield* Effect.fail(new AiToolNotFoundError(id));
+      }
+
+      return projectFunctionToTool(fn);
+    }),
+  });
+};
+
+export const makeToolExecutionServiceFromFunctions = (
+  functions: FunctionDefinition<any, any>[],
+  toolkit: AiToolkit.AiToolkit<AiTool.Any>,
+  handlersLayer: Layer.Layer<AiTool.ToHandler<AiTool.AiTool<any>>, never, never>,
+): Layer.Layer<ToolExecutionService, never, LocalFunctionExecutionService> => {
+  return Layer.effect(
+    ToolExecutionService,
+    Effect.gen(function* () {
+      const toolkitHandler = yield* toolkit.pipe(Effect.provide(handlersLayer));
+
+      const localFunctionExecutionService = yield* LocalFunctionExecutionService;
+      return {
+        handlersFor: (toolkit) => {
+          const makeHandler = (tool: AiTool.Any): ((params: unknown) => Effect.Effect<unknown, any, any>) => {
+            return Effect.fn('toolFunctionHandler')(function* (input: any) {
+              if (toolkitHandler.tools.find((t: AiTool.Any) => t.name === tool.name)) {
+                // TODO(wittjosiah): Everything is `never` here.
+                return yield* (toolkitHandler.handle as any)(tool.name, input);
+              }
+
+              const { functionName } = Context.get(FunctionToolAnnotation)(tool.annotations as any);
+              const fnDef = functions.find((f) => f.name === functionName);
+              if (!fnDef) {
+                return yield* Effect.fail(new AiToolNotFoundError(tool.name));
+              }
+
+              return yield* localFunctionExecutionService
+                .invokeFunction(fnDef, input)
+                .pipe(Effect.catchAllDefect((defect) => Effect.fail(defect)));
+            });
+          };
+
+          return toolkit.of(Record.map(toolkit.tools, (tool, _name) => makeHandler(tool)) as any) as any;
+        },
+      };
+    }),
+  );
+};
+
+/**
+ * Tools that are projected from functions have this annotation.
+ */
+class FunctionToolAnnotation extends Context.Tag('@dxos/assisatnt/FunctionToolAnnotation')<
+  FunctionToolAnnotation,
+  {
+    functionName: string;
+  }
+>() {}
+
+const toolCache = new WeakMap<FunctionDefinition<any, any>, AiTool.Any>();
+const projectFunctionToTool = (fn: FunctionDefinition<any, any>): AiTool.Any => {
+  if (toolCache.has(fn)) {
+    return toolCache.get(fn)!;
+  }
+  const tool = AiTool.make(makeToolName(fn.name), {
+    description: fn.description,
+    parameters: createStructFieldsFromSchema(fn.inputSchema),
+    // TODO(dmaretskyi): Include output schema.
+    failure: Schema.Any, // TODO(dmaretskyi): Better type for the failure?
+  }).annotate(FunctionToolAnnotation, { functionName: fn.name });
+  toolCache.set(fn, tool);
+  return tool;
+};
+
+/**
+ * @returns Tool name produced from function name by escaping invalid characters.
+ */
+const makeToolName = (name: string) => {
+  const toolName = name.replace(/[^a-zA-Z0-9]/g, '_');
+  invariant(toolName.match(/^[a-zA-Z_][a-zA-Z0-9-_]*$/));
+  return toolName;
+};
+
+// TODO(dmaretskyi): Factor out.
+const createStructFieldsFromSchema = (schema: Schema.Schema<any, any>): Record<string, Schema.Schema<any, any>> => {
+  return Match.value(schema.ast).pipe(
+    Match.when(Predicate.isTagged('TypeLiteral'), (ast) => {
+      return Object.fromEntries(ast.propertySignatures.map((prop) => [prop.name, Schema.make(prop.type)]));
+    }),
+    Match.orElseAbsurd,
+  );
+};

package/src/index.ts

@@ -2,8 +2,9 @@
 // Copyright 2024 DXOS.org
 //
 
-export * from './extraction';
-export * from './research';
+export * from './conversation';
+export * from './functions';
 export * from './session';
+export * from './templates';
 export * from './types';
-export * from './blueprint';
+export * from './errors';

package/src/session/__snapshots__/session.test.ts.snap

@@ -0,0 +1,32 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`AiSession > no tools 1`] = `
+[
+  {
+    "blocks": ReactiveArray [
+      {
+        "_tag": "text",
+        "text": "Hello world!",
+      },
+    ],
+    "created": "Sun, 20 Jul 2025 10:54:08 GMT",
+    "id": "01K0KQGFTFV3GBKGEWZW67K8HK",
+    "sender": {
+      "role": "user",
+    },
+  },
+  {
+    "blocks": ReactiveArray [
+      {
+        "_tag": "text",
+        "text": "Hi there! How can I help you today?",
+      },
+    ],
+    "created": "2025-07-20T10:54:10.801Z",
+    "id": "01K0KQGHVHJ37D4YG45WG090WV",
+    "sender": {
+      "role": "assistant",
+    },
+  },
+]
+`;

package/src/session/artifact-diff.ts

@@ -0,0 +1,37 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Context } from 'effect';
+
+import { type ObjectVersion } from '@dxos/echo-db';
+import { type ObjectId } from '@dxos/echo-schema';
+
+/**
+ * Resolves artifact ids to their versions.
+ * Used to give the model a sense of the changes to the artifacts made by users during the conversation.
+ * The artifacts versions are pinned in the history, and whenever the artifact changes in-between assistant's steps,
+ * a diff is inserted into the conversation.
+ *
+ * Can be optionally provided to the session run call.
+ */
+// TODO(burdon): Move.
+// TODO(dmaretskyi): Convert to Context.Reference
+export class ArtifactDiffResolver extends Context.Tag('@dxos/assistant/ArtifactDiffResolver')<
+  ArtifactDiffResolver,
+  ArtifactDiffResolver.Service
+>() {}
+
+export namespace ArtifactDiffResolver {
+  export type Service = {
+    resolve: (artifacts: { id: ObjectId; lastVersion: ObjectVersion }[]) => Promise<
+      Map<
+        ObjectId,
+        {
+          version: ObjectVersion;
+          diff?: string;
+        }
+      >
+    >;
+  };
+}

package/src/session/format.test.ts

@@ -0,0 +1,60 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { beforeAll, describe, it } from '@effect/vitest';
+import { Effect } from 'effect';
+
+import { Blueprint, Template } from '@dxos/blueprints';
+import { Obj } from '@dxos/echo';
+import { type EchoDatabase } from '@dxos/echo-db';
+import { EchoTestBuilder } from '@dxos/echo-db/testing';
+import { DataType } from '@dxos/schema';
+import { trim } from '@dxos/util';
+
+import { createSystemPrompt } from '../templates/system';
+
+import { formatSystemPrompt } from './format';
+
+describe('format', () => {
+  let db: EchoDatabase;
+  beforeAll(async () => {
+    const builder = await new EchoTestBuilder().open();
+    ({ db } = await builder.createDatabase({
+      types: [DataType.Text, DataType.Organization, Blueprint.Blueprint],
+    }));
+  });
+
+  it.effect(
+    'should format',
+    Effect.fn(function* () {
+      const object = db.add(
+        Obj.make(DataType.Organization, {
+          name: 'Test',
+          website: 'https://www.test.com',
+        }),
+      );
+
+      const blueprint = db.add(
+        Blueprint.make({
+          key: 'example.com/blueprint/test',
+          name: 'Test',
+          instructions: Template.make({
+            source: trim`
+              Test
+              This is the test blueprint.
+            `,
+          }),
+        }),
+      );
+
+      const output = yield* formatSystemPrompt({
+        system: createSystemPrompt({}),
+        blueprints: [blueprint, blueprint],
+        objects: [object, object],
+      });
+
+      console.log(output);
+    }),
+  );
+});

package/src/session/format.ts

@@ -0,0 +1,139 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Array, Effect, Option, pipe } from 'effect';
+
+import { Template } from '@dxos/blueprints';
+import { Obj } from '@dxos/echo';
+import { ObjectVersion } from '@dxos/echo-db';
+import { type ObjectId } from '@dxos/echo-schema';
+import { DatabaseService } from '@dxos/functions';
+import { log } from '@dxos/log';
+import { type ContentBlock, DataType } from '@dxos/schema';
+import { trim } from '@dxos/util';
+
+import { AiAssistantError } from '../errors';
+
+import { ArtifactDiffResolver } from './artifact-diff';
+import { type SessionRunParams } from './session';
+
+/**
+ * Formats the system prompt.
+ */
+// TODO(burdon): Move to AiPreprocessor.
+export const formatSystemPrompt = ({
+  system,
+  blueprints = [],
+  objects = [],
+}: Pick<SessionRunParams<any>, 'system' | 'blueprints' | 'objects'>) =>
+  Effect.gen(function* () {
+    // TOOD(burdon): Should process templates.
+    const blueprintDefs = yield* pipe(
+      blueprints,
+      Effect.forEach((blueprint) => Effect.succeed(blueprint.instructions)),
+      Effect.flatMap(Effect.forEach((template) => DatabaseService.load(template.source))),
+      Effect.map(
+        Array.map(
+          (template) => trim`
+            <blueprint>
+              ${Template.process(template.content)}
+            </blueprint>
+          `,
+        ),
+      ),
+      Effect.map((blueprints) =>
+        blueprints.length > 0 ? ['## Blueprints Definitions', ...blueprints].join('\n\n') : undefined,
+      ),
+    );
+
+    const objectDefs = yield* pipe(
+      objects,
+      Effect.forEach((object) =>
+        Effect.succeed(trim`
+          <object>
+            <dxn>${Obj.getDXN(object)}</dxn>
+            <typename>${Obj.getTypename(object)}</typename>
+          </object>
+        `),
+      ),
+      Effect.map((objects) => (objects.length > 0 ? ['## Context Objects', ...objects].join('\n\n') : undefined)),
+    );
+
+    return yield* pipe(
+      Effect.succeed([system, blueprintDefs, objectDefs].filter((def): def is string => def !== undefined)),
+      Effect.map((parts) => parts.join('\n\n')),
+    );
+  });
+
+/**
+ * Formats the user prompt.
+ */
+// TODO(burdon): Move to AiPreprocessor.
+// TODO(burdon): Convert util below to `Effect.fn` (to preserve stack info)
+export const formatUserPrompt = ({ prompt, history = [] }: Pick<SessionRunParams<any>, 'prompt' | 'history'>) =>
+  Effect.gen(function* () {
+    const prelude: ContentBlock.Any[] = [];
+
+    // TODO(dmaretskyi): Evaluate other approaches as `serviceOption` isn't represented in the type system.
+    const artifactDiffResolver = yield* Effect.serviceOption(ArtifactDiffResolver);
+    if (Option.isSome(artifactDiffResolver)) {
+      const versions = gatherObjectVersions(history);
+      const artifactDiff = yield* Effect.tryPromise({
+        try: () =>
+          artifactDiffResolver.value.resolve(
+            [...versions.entries()].map(([id, version]) => ({ id, lastVersion: version })),
+          ),
+        catch: AiAssistantError.wrap('Artifact diff resolution error'),
+      });
+
+      log.info('version', { artifactDiff, versions });
+      for (const [id, { version }] of [...artifactDiff.entries()]) {
+        if (ObjectVersion.equals(version, versions.get(id)!)) {
+          artifactDiff.delete(id);
+          continue;
+        }
+
+        prelude.push({ _tag: 'anchor', objectId: id, version });
+      }
+
+      if (artifactDiff.size > 0) {
+        prelude.push(createArtifactUpdateBlock(artifactDiff));
+      }
+    }
+
+    return Obj.make(DataType.Message, {
+      created: new Date().toISOString(),
+      sender: { role: 'user' },
+      blocks: [...prelude, { _tag: 'text', text: prompt }],
+    });
+  });
+
+const gatherObjectVersions = (messages: DataType.Message[]): Map<ObjectId, ObjectVersion> => {
+  const artifactIds = new Map<ObjectId, ObjectVersion>();
+  for (const message of messages) {
+    for (const block of message.blocks) {
+      if (block._tag === 'anchor') {
+        artifactIds.set(block.objectId, block.version as ObjectVersion);
+      }
+    }
+  }
+
+  return artifactIds;
+};
+
+const createArtifactUpdateBlock = (
+  artifactDiff: Map<ObjectId, { version: ObjectVersion; diff?: string }>,
+): ContentBlock.Any => {
+  return {
+    _tag: 'text',
+    // TODO(dmaretskyi): Does this need to be a special content-block?
+    disposition: 'artifact-update',
+    text: trim`
+      The following artifacts have been updated since the last message:
+      ${[...artifactDiff.entries()]
+        .map(([id, { diff }]) => `<changed-artifact id="${id}">${diff ? `\n${diff}` : ''}</changed-artifact>`)
+        .join('\n')}
+    `,
+  };
+};

package/src/session/index.ts

@@ -2,5 +2,6 @@
 // Copyright 2025 DXOS.org
 //
 
+export * from './artifact-diff';
 export * from './session';
 export * from './version-pin';

package/src/session/session.test.ts

@@ -2,18 +2,16 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Schema } from 'effect';
-import { describe, test } from 'vitest';
-
-import { EdgeAiServiceClient, ConsolePrinter, createTool, ToolResult } from '@dxos/ai';
-import { AI_SERVICE_ENDPOINT } from '@dxos/ai/testing';
-import { ArtifactId, defineArtifact } from '@dxos/artifact';
-import { Type, Obj } from '@dxos/echo';
-import { ObjectId } from '@dxos/echo-schema';
-import { DXN } from '@dxos/keys';
+import { AiTool, AiToolkit } from '@effect/ai';
+import { describe, it } from '@effect/vitest';
+import { Effect, Layer, Schema } from 'effect';
+
+import { AiService, ToolExecutionService, ToolResolverService } from '@dxos/ai';
+import { AiServiceTestingPreset } from '@dxos/ai/testing';
+import { Obj, Type } from '@dxos/echo';
 import { log } from '@dxos/log';
 
-import { AISession } from './session';
+import { AiSession } from './session';
 
 // Define a calendar event artifact schema.
 const CalendarEventSchema = Schema.Struct({
@@ -30,15 +28,90 @@ const CalendarEventSchema = Schema.Struct({
 
 type CalendarEvent = Schema.Schema.Type<typeof CalendarEventSchema>;
 
-// TODO(burdon): Flaky.
-describe.skip('AISession with Ollama', () => {
-  test('create calendar itinerary', { timeout: 60_000 }, async () => {
-    const aiClient = new EdgeAiServiceClient({ endpoint: AI_SERVICE_ENDPOINT.REMOTE });
-    // const aiClient = new OllamaAiServiceClient({
-    //   overrides: { model: 'llama3.1:8b' },
-    // });
-    const session = new AISession({ operationModel: 'configured' });
+class TestToolkit extends AiToolkit.make(
+  AiTool.make('Calculator', {
+    description: 'Basic calculator tool',
+    parameters: {
+      input: Schema.String.annotations({
+        description: 'The calculation to perform.',
+      }),
+    },
+    success: Schema.Struct({
+      result: Schema.Number,
+    }),
+    failure: Schema.Never,
+  }),
+) {}
+
+// Tool handlers.
+const toolkitLayer = TestToolkit.toLayer({
+  Calculator: ({ input }) =>
+    Effect.gen(function* () {
+      const result = (() => {
+        // Restrict to basic arithmetic operations for safety.
+        const sanitizedInput = input.replace(/[^0-9+\-*/().\s]/g, '');
+        log.info('calculate', { sanitizedInput });
+
+        // eslint-disable-next-line @typescript-eslint/no-implied-eval
+        return Function(`"use strict"; return (${sanitizedInput})`)();
+      })();
+
+      return { result };
+    }),
+});
+
+describe.runIf(process.env.DX_RUN_SLOW_TESTS)('AiSession', () => {
+  it.effect(
+    'no tools',
+    Effect.fn(
+      function* ({ expect: _ }) {
+        const session = new AiSession({ operationModel: 'configured' });
+        const response = yield* session.run({
+          prompt: 'Hello world!',
+          history: [],
+        });
+        log.info('response', { response });
+      },
+      Effect.provide(
+        AiService.model('@anthropic/claude-3-5-sonnet-20241022').pipe(
+          Layer.provideMerge(ToolResolverService.layerEmpty),
+          Layer.provideMerge(ToolExecutionService.layerEmpty),
+          Layer.provideMerge(AiServiceTestingPreset('direct')),
+        ),
+      ),
+    ),
+  );
 
+  it.effect(
+    'calculator',
+    Effect.fn(
+      function* ({ expect: _ }) {
+        const session = new AiSession({ operationModel: 'configured' });
+        const response = yield* session.run({
+          prompt: 'What is 10 + 20?',
+          history: [],
+          toolkit: TestToolkit,
+        });
+        log.info('response', { response });
+      },
+      Effect.provide(
+        Layer.mergeAll(
+          toolkitLayer,
+          ToolResolverService.layerEmpty,
+          ToolExecutionService.layerEmpty,
+          AiService.model('@anthropic/claude-3-5-sonnet-20241022').pipe(
+            Layer.provideMerge(AiServiceTestingPreset('direct')),
+          ),
+        ),
+      ),
+    ),
+  );
+
+  // TODO(dmaretskyi): Revive test.
+  /*
+  it.skip('create calendar itinerary', { timeout: 60_000 }, async () => {
+    // overrides: { model: 'llama3.1:8b' },
+    const session = new AiSession({ operationModel: 'configured' });
     const objects = new Set<string>();
 
     // Define calendar artifact.
@@ -145,6 +218,7 @@ describe.skip('AISession with Ollama', () => {
         model: '@anthropic/claude-3-5-haiku-20241022',
       },
       prompt: 'create a table and map for a travel itinerary based on events in my calendar',
+      toolResolver: new ToolRegistry([]),
     });
 
     log('result', {
@@ -152,10 +226,12 @@ describe.skip('AISession with Ollama', () => {
       finalMessage: response.at(-1),
     });
   });
+  */
 });
 
 // Travel to rome, florence, livorno, siena, madrid for conferences
-const CALENDAR_EVENTS: CalendarEvent[] = [
+
+const _CALENDAR_EVENTS: CalendarEvent[] = [
   Obj.make(CalendarEventSchema, {
     title: 'Exploring Ancient Ruins in Rome',
     startTime: '2024-01-01T10:00:00Z',