@dxos/assistant 0.8.4-main.67995b8 → 0.8.4-main.a4bbb77

package/src/extraction/extraction-ner-function/extraction-ner-function.ts

@@ -2,7 +2,7 @@
 // Copyright 2025 DXOS.org
 //
 
-// TODO(burdon): Fix.
+// ISSUE(burdon): defineFunction
 // @ts-nocheck
 
 import { asyncTimeout } from '@dxos/async';
@@ -15,8 +15,13 @@ import { DataType } from '@dxos/schema';
 import { ExtractionInput, ExtractionOutput } from '../extraction';
 import { findQuotes, insertReferences } from '../quotes';
 
+/**
+ * Named Entity Recognition model
+ */
 // TODO(mykola): Make it use vector index to query objects and not pass objects as input.
-export const extractionNerFn = defineFunction({
+export const extractionNerFunction = defineFunction({
+  key: 'dxos.org/function/extraction/extract-entities',
+  name: 'Extract Entities',
   description: 'Extracts named entities from a text',
   inputSchema: ExtractionInput,
   outputSchema: ExtractionOutput,
@@ -28,27 +33,29 @@ export const extractionNerFn = defineFunction({
 
     const entitiesPromise = Promise.all(
       message.blocks.map(async (block) => {
-        invariant(block._tag === 'transcript' || block._tag === 'text', 'Block must have text');
+        invariant(block._tag === 'transcript' || block._tag === 'text');
         return extractFullEntities(block.text);
       }),
     ).then((entities) => entities.flat());
     const entities = options?.timeout ? await asyncTimeout(entitiesPromise, options.timeout) : await entitiesPromise;
+
     log.info('entities', { entities });
     const quoteReferences = await findQuotes(
       entities.map((entity) => entity.word),
       db,
     );
-    const blocksWithReferences = message.blocks.map((block, i) => {
-      invariant(block._tag === 'transcript' || block._tag === 'text', 'Block must have text');
+
+    const blocksWithReferences = message.blocks.map((block) => {
+      invariant(block._tag === 'transcript' || block._tag === 'text');
       return { ...block, text: insertReferences(block._tag, quoteReferences) };
     });
+
     const messageWithReferences = create(DataType.Message, {
       ...message,
       blocks: blocksWithReferences,
     });
 
     log.info('output', { messageWithReferences });
-
     return {
       message: messageWithReferences,
       timeElapsed: performance.now() - startTime,

package/src/extraction/extraction-ner-function/named-entity-recognition.test.ts

@@ -10,7 +10,8 @@ import { createTestData } from '@dxos/schema/testing';
 
 import { combineNerTokens, createTokenGroups, extractFullEntities, getNer } from './named-entity-recognition';
 
-describe('NamedEntityRecognition', () => {
+// TODO(mykola): This should not run on CI.
+describe.skip('NamedEntityRecognition', () => {
   let ner: TokenClassificationPipelineType;
 
   beforeAll(async () => {

package/src/extraction/extraction-ner-function/named-entity-recognition.ts

@@ -12,6 +12,7 @@ import {
 import { log } from '@dxos/log';
 
 let _ner: Promise<TokenClassificationPipelineType>;
+
 /**
  * Named Entity Recognition pipeline.
  * Initializes the pipeline on first call.
@@ -24,6 +25,7 @@ export const getNer = () => {
       return ner;
     });
   }
+
   return _ner;
 };
 

package/src/extraction/quotes.test.ts

@@ -2,7 +2,7 @@
 // Copyright 2025 DXOS.org
 //
 
-import { describe, test, expect } from 'vitest';
+import { describe, expect, test } from 'vitest';
 
 import { insertReferences } from './quotes';
 
@@ -13,7 +13,7 @@ describe('insertReferences', () => {
     };
 
     expect(insertReferences('This is a computational irreducibility test.', quotes)).toBe(
-      'This is a [computational irreducibility][dxn:echo:@:01JTG9JW11XGWJZ32AW8ET93D1] test.',
+      'This is a [computational irreducibility](dxn:echo:@:01JTG9JW11XGWJZ32AW8ET93D1) test.',
     );
     expect(
       insertReferences(
@@ -21,7 +21,7 @@ describe('insertReferences', () => {
         quotes,
       ),
     ).toBe(
-      "And what I'd like to talk today about is Steven Wolfram's concept of a [computational irreducibility][dxn:echo:@:01JTG9JW11XGWJZ32AW8ET93D1].",
+      "And what I'd like to talk today about is Steven Wolfram's concept of a [computational irreducibility](dxn:echo:@:01JTG9JW11XGWJZ32AW8ET93D1).",
     );
   });
 });

package/src/extraction/quotes.ts

@@ -4,7 +4,7 @@
 
 import { Schema } from 'effect';
 
-import { Query, Filter } from '@dxos/echo';
+import { Filter, Query } from '@dxos/echo';
 import { type EchoDatabase } from '@dxos/echo-db';
 import { ObjectId } from '@dxos/echo-schema';
 import { DXN } from '@dxos/keys';
@@ -51,7 +51,7 @@ export const insertReferences = (text: string, quotes: ReferencedQuotes) => {
 
     // Use a case-insensitive regular expression to replace the quote.
     const regex = new RegExp(quote.quote.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'gi');
-    text = text.replace(regex, `[${quote.quote}][${DXN.fromLocalObjectId(quote.id).toString()}]`);
+    text = text.replace(regex, `[${quote.quote}](${DXN.fromLocalObjectId(quote.id).toString()})`);
   }
 
   return text;

package/src/functions/services.ts

@@ -2,79 +2,127 @@
 // Copyright 2025 DXOS.org
 //
 
-import { AiTool } from '@effect/ai';
-import { Context, Effect, Layer, Match, Predicate, Record, Schema } from 'effect';
+import { Tool, type Toolkit } from '@effect/ai';
+import { Context, Effect, Layer, Record, Schema } from 'effect';
 
 import { AiToolNotFoundError, ToolExecutionService, ToolResolverService } from '@dxos/ai';
-import { LocalFunctionExecutionService, type FunctionDefinition } from '@dxos/functions';
+import { todo } from '@dxos/debug';
+import { Query } from '@dxos/echo';
+import { DatabaseService, FunctionDefinition, FunctionInvocationService, FunctionType } from '@dxos/functions';
 import { invariant } from '@dxos/invariant';
 
+/**
+ * Constructs a `ToolResolverService` whose `resolve(id)` looks up tools in the following order:
+ *  1. Toolkit: return an existing tool from the provided `toolkit` if one is present under `id`.
+ *  2. Functions in DB: query `FunctionType` by `key=id`; if found, deserialize and project to a tool
+ *     (propagating any `deployedFunctionId` discovered in the object's metadata).
+ *  3. Functions passed in: fall back to a matching `FunctionDefinition` from the `functions` array.
+ *
+ * If none of the above yield a match, the effect fails with `AiToolNotFoundError`.
+ *
+ * Requires `DatabaseService` in the environment.
+ */
 export const makeToolResolverFromFunctions = (
   functions: FunctionDefinition<any, any>[],
-): Layer.Layer<ToolResolverService> => {
-  return Layer.succeed(ToolResolverService, {
-    resolve: Effect.fn('resolveTool')(function* (id) {
-      const fn = functions.find((f) => f.name === id);
-      if (!fn) {
-        return yield* Effect.fail(new AiToolNotFoundError(id));
-      }
-
-      return projectFunctionToTool(fn);
+  toolkit: Toolkit.Toolkit<any>,
+): Layer.Layer<ToolResolverService, never, DatabaseService> => {
+  return Layer.effect(
+    ToolResolverService,
+    Effect.gen(function* () {
+      const dbService = yield* DatabaseService;
+      return {
+        resolve: (id): Effect.Effect<Tool.Any, AiToolNotFoundError> =>
+          Effect.gen(function* () {
+            const tool = toolkit.tools[id];
+            if (tool) {
+              return tool;
+            }
+
+            const {
+              objects: [dbFunction],
+            } = yield* DatabaseService.runQuery(Query.type(FunctionType, { key: id }));
+
+            const functionDef = dbFunction
+              ? FunctionDefinition.deserialize(dbFunction)
+              : functions.find((fn) => fn.key === id);
+
+            if (!functionDef) {
+              return yield* Effect.fail(new AiToolNotFoundError(id));
+            }
+
+            return projectFunctionToTool(functionDef);
+          }).pipe(Effect.provideService(DatabaseService, dbService)),
+      } satisfies Context.Tag.Service<ToolResolverService>;
     }),
-  });
+  );
 };
 
 export const makeToolExecutionServiceFromFunctions = (
-  functions: FunctionDefinition<any, any>[],
-): Layer.Layer<ToolExecutionService, never, LocalFunctionExecutionService> => {
+  toolkit: Toolkit.Toolkit<any>,
+  handlersLayer: Layer.Layer<Tool.Handler<any>, never, never>,
+): Layer.Layer<ToolExecutionService, never, FunctionInvocationService> => {
   return Layer.effect(
     ToolExecutionService,
     Effect.gen(function* () {
-      const localFunctionExecutionService = yield* LocalFunctionExecutionService;
+      const toolkitHandler = yield* toolkit.pipe(Effect.provide(handlersLayer));
+      invariant(isHandlerLike(toolkitHandler));
+      const functionInvocationService = yield* FunctionInvocationService;
+
       return {
         handlersFor: (toolkit) => {
-          const makeHandler = (tool: AiTool.Any): ((params: unknown) => Effect.Effect<unknown, any, any>) => {
-            return Effect.fn('toolFunctionHandler')(function* (input: any) {
-              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));
+          const makeHandler = (tool: Tool.Any): ((params: unknown) => Effect.Effect<unknown, any, any>) => {
+            return Effect.fn(`toolFunctionHandler ${tool.name}`)(function* (input: any) {
+              if (toolkitHandler.tools[tool.name]) {
+                if (Tool.isProviderDefined(tool)) {
+                  throw new Error('Attempted to call a provider-defined tool');
+                }
+                // TODO(wittjosiah): Everything is `never` here.
+                return yield* (toolkitHandler.handle as any)(tool.name, input);
               }
 
-              return yield* localFunctionExecutionService
-                .invokeFunction(fnDef, input)
+              const { definition: functionDef } = Context.get(FunctionToolAnnotation)(tool.annotations as any);
+
+              return yield* functionInvocationService
+                .invokeFunction(functionDef, input as any)
                 .pipe(Effect.catchAllDefect((defect) => Effect.fail(defect)));
             });
           };
 
-          return toolkit.of(Record.map(toolkit.tools, (tool, name) => makeHandler(tool)) as any) as any;
+          return toolkit.of(
+            Record.map(toolkit.tools, (tool, _name) => (Tool.isUserDefined(tool) ? makeHandler(tool) : null)) as any,
+          ) as any;
         },
       };
     }),
   );
 };
 
-/**
- * Tools that are projected from functions have this annotation.
- */
-class FunctionToolAnnotation extends Context.Tag('@dxos/assisatnt/FunctionToolAnnotation')<
+class FunctionToolAnnotation extends Context.Tag('@dxos/assistant/FunctionToolAnnotation')<
   FunctionToolAnnotation,
-  {
-    functionName: string;
-  }
+  { definition: FunctionDefinition<any, any> }
 >() {}
 
-const toolCache = new WeakMap<FunctionDefinition<any, any>, AiTool.Any>();
-const projectFunctionToTool = (fn: FunctionDefinition<any, any>): AiTool.Any => {
+const toolCache = new WeakMap<FunctionDefinition<any, any>, Tool.Any>();
+
+/**
+ * Projects a `FunctionDefinition` into an `AiTool`.
+ *
+ * @param fn The function definition to project into a tool.
+ * @param meta Optional projection metadata.
+ * @param meta.deployedFunctionId Backend deployment ID used for remote invocation when present. This is the
+ *        EDGE service's function deployment identifier (not the ECHO object ID/DXN and not `FunctionDefinition.key`).
+ */
+const projectFunctionToTool = (fn: FunctionDefinition<any, any>): Tool.Any => {
   if (toolCache.has(fn)) {
     return toolCache.get(fn)!;
   }
-  const tool = AiTool.make(makeToolName(fn.name), {
+
+  const tool = Tool.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 });
+  }).annotate(FunctionToolAnnotation, { definition: fn });
   toolCache.set(fn, tool);
   return tool;
 };
@@ -83,17 +131,23 @@ const projectFunctionToTool = (fn: FunctionDefinition<any, any>): AiTool.Any =>
  * @returns Tool name produced from function name by escaping invalid characters.
  */
 const makeToolName = (name: string) => {
-  const toolName = name.replace(/[^a-zA-Z0-9]/g, '');
+  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,
-  );
+  switch (schema.ast._tag) {
+    case 'TypeLiteral':
+      return Object.fromEntries(schema.ast.propertySignatures.map((prop) => [prop.name, Schema.make(prop.type)]));
+    case 'VoidKeyword':
+      return {};
+    default:
+      return todo(`Unsupported schema AST: ${schema.ast._tag}`);
+  }
+};
+
+const isHandlerLike = (value: unknown): value is Toolkit.WithHandler<Record<string, Tool.Any>> => {
+  return typeof (value as any).tools === 'object' && typeof (value as any).handle === 'function';
 };

package/src/index.ts

@@ -2,10 +2,11 @@
 // Copyright 2024 DXOS.org
 //
 
-export * from './context';
 export * from './conversation';
+export * from './diff';
+export * from './errors';
+export * from './execution-graph';
 export * from './functions';
-export * from './research';
 export * from './session';
 export * from './templates';
 export * from './types';

package/src/session/artifact-diff.ts

@@ -0,0 +1,41 @@
+//
+// 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(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/error-handling.ts

@@ -0,0 +1,40 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { AiError } from '@effect/ai';
+import { HttpClientError } from '@effect/platform';
+import { Effect, Schema } from 'effect';
+
+import { AiModelError } from '../errors';
+
+const AnthropicErrorResponse = Schema.Struct({
+  type: Schema.Literal('error'),
+  error: Schema.Struct({
+    type: Schema.String,
+    message: Schema.String,
+  }),
+});
+
+// TODO(dmaretskyi): Needs rework
+export const mapAiError = (err: AiError.AiError): Effect.Effect<AiError.AiError> =>
+  Effect.gen(function* () {
+    const cause = err.cause;
+    if (HttpClientError.isHttpClientError(cause) && cause.reason === 'StatusCode') {
+      const body = yield* cause.response.json.pipe(
+        Effect.flatMap(Schema.decodeUnknown(AnthropicErrorResponse, { exact: false })),
+      );
+
+      const parsedCause = new AiModelError({
+        message: body.error.message,
+        context: { model: 'anthropic', type: body.error.type },
+      });
+      return AiError.UnknownError.make({
+        description: body.error.message,
+        module: err.module,
+        method: err.method,
+        cause: parsedCause,
+      });
+    }
+    return err;
+  }).pipe(Effect.catchAll(() => Effect.succeed(err)));

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 AiSessionRunParams } from './session';
+
+/**
+ * Formats the system prompt.
+ */
+// TODO(burdon): Move to AiPreprocessor.
+export const formatSystemPrompt = ({
+  system,
+  blueprints = [],
+  objects = [],
+}: Pick<AiSessionRunParams<any>, 'system' | 'blueprints' | 'objects'>) =>
+  Effect.gen(function* () {
+    const blueprintDefs = yield* pipe(
+      blueprints,
+      Effect.forEach((blueprint) => Effect.succeed(blueprint.instructions)),
+      Effect.flatMap(Effect.forEach((template) => DatabaseService.loadOption(template.source))),
+      Effect.map(Array.filter(Option.isSome)),
+      Effect.map(
+        Array.map(
+          (template) => trim`
+            <blueprint>
+              ${Template.process(template.value.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<AiSessionRunParams<any>, 'prompt' | 'history'>) =>
+  Effect.gen(function* () {
+    const blocks: 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({ message: 'Artifact diff resolution error' }),
+      });
+
+      log('version', { artifactDiff, versions });
+      for (const [id, { version }] of [...artifactDiff.entries()]) {
+        if (ObjectVersion.equals(version, versions.get(id)!)) {
+          artifactDiff.delete(id);
+          continue;
+        }
+
+        blocks.push({ _tag: 'anchor', objectId: id, version });
+      }
+
+      if (artifactDiff.size > 0) {
+        blocks.push(createArtifactUpdateBlock(artifactDiff));
+      }
+    }
+
+    return Obj.make(DataType.Message, {
+      created: new Date().toISOString(),
+      sender: { role: 'user' },
+      blocks: [...blocks, { _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,8 @@
 // Copyright 2025 DXOS.org
 //
 
+export * from './artifact-diff';
+export * from './observer';
 export * from './session';
+export * from './toolkit';
 export * from './version-pin';