@dxos/assistant 0.8.3 → 0.8.4-main.1da679c

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

@@ -2,25 +2,29 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Message } from '@dxos/ai';
+// ISSUE(burdon): defineFunction
+// @ts-nocheck
+
+import { AiService } from '@dxos/ai';
 import { Obj } from '@dxos/echo';
 import { create } from '@dxos/echo-schema';
-import { AiService, defineFunction, type FunctionDefinition } from '@dxos/functions';
-import { DataType } from '@dxos/schema';
+import { type FunctionDefinition, defineFunction } from '@dxos/functions';
+import { type ContentBlock, DataType } from '@dxos/schema';
 
-import PROMPT from './instructions.tpl?raw';
-import { AISession } from '../../session';
+import { AiSession } from '../../session';
 import { ExtractionInput, ExtractionOutput } from '../extraction';
-import { insertReferences, ReferencedQuotes } from '../quotes';
+import { ReferencedQuotes, insertReferences } from '../quotes';
+
+import PROMPT from './instructions.tpl?raw';
 
-export const extractionAnthropicFn: FunctionDefinition<ExtractionInput, ExtractionOutput> = defineFunction({
+export const extractionAnthropicFunction: FunctionDefinition<ExtractionInput, ExtractionOutput> = defineFunction({
   description: 'Extract entities from the transcript message and add them to the message.',
   inputSchema: ExtractionInput,
   outputSchema: ExtractionOutput,
   handler: async ({ data: { message, objects }, context }) => {
     const startTime = performance.now();
-    const ai = context.getService(AiService);
-    const session = new AISession({ operationModel: 'configured' });
+    const ai = context.getService(AiService.AiService);
+    const session = new AiSession({ operationModel: 'configured' });
     const result = await session.runStructured(ReferencedQuotes, {
       generationOptions: {
         model: '@anthropic/claude-3-5-haiku-20241022',
@@ -28,30 +32,32 @@ export const extractionAnthropicFn: FunctionDefinition<ExtractionInput, Extracti
       client: ai.client,
       systemPrompt: PROMPT,
       history: [
-        Obj.make(Message, {
-          role: 'user',
-          content: [
+        Obj.make(DataType.Message, {
+          created: new Date().toISOString(),
+          sender: { role: 'user' },
+          blocks: [
             {
-              type: 'text',
+              _tag: 'text',
               text: `<context>${JSON.stringify(objects)}</context>`,
-            },
+            } satisfies ContentBlock.Text,
             {
-              type: 'text',
+              _tag: 'text',
               text: `<transcript>${JSON.stringify(message.blocks)}</transcript>`,
-            },
+            } satisfies ContentBlock.Text,
           ],
         }),
       ],
       artifacts: [],
       prompt: '',
       tools: [],
-    });
+    } as any); // TODO(burdon): Rewrite test.
 
     return {
+      timeElapsed: performance.now() - startTime,
       message: create(DataType.Message, {
         ...message,
-        blocks: message.blocks.map((block, i) =>
-          block.type !== 'transcription'
+        blocks: message.blocks.map((block) =>
+          block._tag !== 'transcript'
             ? block
             : {
                 ...block,
@@ -59,7 +65,6 @@ export const extractionAnthropicFn: FunctionDefinition<ExtractionInput, Extracti
               },
         ),
       }),
-      timeElapsed: performance.now() - startTime,
     };
   },
 });

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

@@ -2,19 +2,20 @@
 // Copyright 2025 DXOS.org
 //
 
-import { describe, test, beforeAll } from 'vitest';
+import { beforeAll, describe, test } from 'vitest';
 
 import { type EchoDatabase } from '@dxos/echo-db';
 import { EchoTestBuilder } from '@dxos/echo-db/testing';
 import { FunctionExecutor, ServiceContainer } from '@dxos/functions';
 import { log } from '@dxos/log';
 import { DataType } from '@dxos/schema';
-import { createTestData, Testing } from '@dxos/schema/testing';
+import { Testing, createTestData } from '@dxos/schema/testing';
 import { range } from '@dxos/util';
 
-import { extractionNerFn } from './extraction-ner-function';
 import { processTranscriptMessage } from '../extraction';
 
+import { extractionNerFunction } from './extraction-ner-function';
+
 describe.skip('NER EntityExtraction', () => {
   let builder: EchoTestBuilder;
   let db: EchoDatabase;
@@ -27,6 +28,7 @@ describe.skip('NER EntityExtraction', () => {
     contacts: Record<string, DataType.Person>;
     organizations: Record<string, DataType.Organization>;
   };
+
   const TYPES = [DataType.Organization, DataType.Person, Testing.Contact, Testing.DocumentType];
 
   beforeAll(async () => {
@@ -44,8 +46,8 @@ describe.skip('NER EntityExtraction', () => {
       contacts: Object.fromEntries(Object.entries(data.contacts).map(([key, value]) => [key, db.add(value)])),
       organizations: Object.fromEntries(Object.entries(data.organizations).map(([key, value]) => [key, db.add(value)])),
     };
-    await db.flush();
 
+    await db.flush();
     executor = new FunctionExecutor(
       new ServiceContainer().setServices({
         database: { db },
@@ -64,7 +66,7 @@ describe.skip('NER EntityExtraction', () => {
           message,
           objects: [...documents, ...Object.values(contacts)],
         },
-        function: extractionNerFn,
+        function: extractionNerFunction,
         executor,
       });
       log.info('output', enhancedMessage);
@@ -73,7 +75,6 @@ describe.skip('NER EntityExtraction', () => {
 
   test('computational irreducibility', async () => {
     const { transcriptWoflram, documents, contacts } = testData;
-
     log.info('context', { documents, contacts });
     const message = transcriptWoflram[0];
     log.info('input', message);
@@ -85,7 +86,7 @@ describe.skip('NER EntityExtraction', () => {
             message,
             objects: [...documents, ...Object.values(contacts)],
           },
-          function: extractionNerFn,
+          function: extractionNerFunction,
           executor,
         });
         log.info('output', { message: enhancedMessage.blocks[0], timeElapsed });
@@ -95,7 +96,6 @@ describe.skip('NER EntityExtraction', () => {
 
   test('org and document linking', async () => {
     const { transcriptJosiah, documents, contacts, organizations } = testData;
-
     log.info('context', { contacts, organizations, documents });
 
     for (const message of transcriptJosiah) {
@@ -106,7 +106,7 @@ describe.skip('NER EntityExtraction', () => {
           message,
           objects: [...documents, ...Object.values(contacts), ...Object.values(organizations)],
         },
-        function: extractionNerFn,
+        function: extractionNerFunction,
         executor,
       });
       log.info('output', enhancedMessage);

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

@@ -2,6 +2,9 @@
 // Copyright 2025 DXOS.org
 //
 
+// ISSUE(burdon): defineFunction
+// @ts-nocheck
+
 import { asyncTimeout } from '@dxos/async';
 import { create } from '@dxos/echo-schema';
 import { DatabaseService, defineFunction } from '@dxos/functions';
@@ -9,12 +12,14 @@ import { invariant } from '@dxos/invariant';
 import { log } from '@dxos/log';
 import { DataType } from '@dxos/schema';
 
-import { extractFullEntities } from './named-entity-recognition';
 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({
   description: 'Extracts named entities from a text',
   inputSchema: ExtractionInput,
   outputSchema: ExtractionOutput,
@@ -22,30 +27,33 @@ export const extractionNerFn = defineFunction({
     log.info('input', { message, options });
     const startTime = performance.now();
     const { db } = context.getService(DatabaseService);
+    const { extractFullEntities } = await import('./named-entity-recognition');
 
     const entitiesPromise = Promise.all(
       message.blocks.map(async (block) => {
-        invariant(block.type === 'transcription' || block.type === '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.type === 'transcription' || block.type === 'text', 'Block must have text');
-      return { ...block, text: insertReferences(block.text, quoteReferences) };
+
+    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

@@ -57,7 +57,7 @@ describe('NamedEntityRecognition', () => {
     const { transcriptWoflram, transcriptJosiah } = await createTestData();
 
     const blocks = [...transcriptWoflram, ...transcriptJosiah].flatMap((message) =>
-      message.blocks.flatMap((block) => (block.type === 'transcription' ? [block.text] : [])),
+      message.blocks.flatMap((block) => (block._tag === 'transcript' ? [block.text] : [])),
     );
 
     for (const block of blocks) {
@@ -70,7 +70,7 @@ describe('NamedEntityRecognition', () => {
     const { transcriptMessages } = await createTestData();
 
     const blocks = transcriptMessages.flatMap((message) =>
-      message.blocks.flatMap((block) => (block.type === 'transcription' ? [block.text] : [])),
+      message.blocks.flatMap((block) => (block._tag === 'transcript' ? [block.text] : [])),
     );
 
     for (const block of blocks) {

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

@@ -4,14 +4,15 @@
 
 import {
   type TokenClassificationOutput,
+  type TokenClassificationPipelineType,
   type TokenClassificationSingle,
   pipeline,
-  type TokenClassificationPipelineType,
 } from '@xenova/transformers';
 
 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/index.ts

@@ -0,0 +1,5 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './services';

package/src/functions/services.ts

@@ -0,0 +1,119 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { AiTool, type AiToolkit } from '@effect/ai';
+import { Context, Effect, Layer, Record, Schema } from 'effect';
+
+import { AiToolNotFoundError, ToolExecutionService, ToolResolverService } from '@dxos/ai';
+import { todo } from '@dxos/debug';
+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 functionDef = functions.find((fn) => fn.name === id);
+      if (!functionDef) {
+        return yield* Effect.fail(new AiToolNotFoundError(id));
+      }
+
+      return projectFunctionToTool(functionDef);
+    }),
+  });
+};
+
+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 functionDef = functions.find((fn) => fn.name === functionName);
+              if (!functionDef) {
+                return yield* Effect.fail(new AiToolNotFoundError(tool.name));
+              }
+
+              return yield* localFunctionExecutionService
+                .invokeFunction(functionDef, 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>> => {
+  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}`);
+  }
+};

package/src/index.ts

@@ -2,8 +2,10 @@
 // Copyright 2024 DXOS.org
 //
 
-export * from './extraction';
-export * from './research';
+export * from './conversation';
+export * from './diff';
+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,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,39 @@
+//
+// 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,
+  }),
+});
+
+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 new AiError.AiError({
+        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);
+    }),
+  );
+});