@dxos/assistant 0.8.4-main.e098934 → 0.8.4-main.e8ec1fe

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

@@ -2,61 +2,55 @@
 // 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';
-import { invariant } from '@dxos/invariant';
-import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { defineFunction } from '@dxos/functions';
 
 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 extractionNerFunction = defineFunction({
+  key: 'dxos.org/function/extraction/extract-entities',
+  name: 'Extract Entities',
   description: 'Extracts named entities from a text',
   inputSchema: ExtractionInput,
   outputSchema: ExtractionOutput,
   handler: async ({ data: { message, options }, context }) => {
-    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._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) => {
-      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,
-    };
+    throw new Error('not implemented');
+    // 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._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) => {
+    //   invariant(block._tag === 'transcript' || block._tag === 'text');
+    //   return { ...block, text: insertReferences(block._tag, quoteReferences) };
+    // });
+
+    // const messageWithReferences = create(Message.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

@@ -6,11 +6,13 @@ import { type TokenClassificationOutput, type TokenClassificationPipelineType }
 import { beforeAll, describe, expect, test } from 'vitest';
 
 import { log } from '@dxos/log';
-import { createTestData } from '@dxos/schema/testing';
+import { type ContentBlock, type Message } from '@dxos/types';
+import { createTestData } from '@dxos/types/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 () => {
@@ -56,8 +58,10 @@ describe('NamedEntityRecognition', () => {
   test('own data', async () => {
     const { transcriptWoflram, transcriptJosiah } = await createTestData();
 
-    const blocks = [...transcriptWoflram, ...transcriptJosiah].flatMap((message) =>
-      message.blocks.flatMap((block) => (block._tag === 'transcript' ? [block.text] : [])),
+    const blocks = [...transcriptWoflram, ...transcriptJosiah].flatMap((message: Message.Message) =>
+      message.blocks.flatMap((block: ContentBlock.Any) =>
+        block._tag === 'transcript' ? [(block as ContentBlock.Transcript).text] : [],
+      ),
     );
 
     for (const block of blocks) {
@@ -69,8 +73,10 @@ describe('NamedEntityRecognition', () => {
   test('more own data', async () => {
     const { transcriptMessages } = await createTestData();
 
-    const blocks = transcriptMessages.flatMap((message) =>
-      message.blocks.flatMap((block) => (block._tag === 'transcript' ? [block.text] : [])),
+    const blocks = transcriptMessages.flatMap((message: Message.Message) =>
+      message.blocks.flatMap((block: ContentBlock.Any) =>
+        block._tag === 'transcript' ? [(block as ContentBlock.Transcript).text] : [],
+      ),
     );
 
     for (const block of blocks) {

package/src/extraction/extraction.ts

@@ -2,16 +2,17 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Schema } from 'effect';
+import * as Schema from 'effect/Schema';
 
 import { asyncTimeout } from '@dxos/async';
-import { Expando } from '@dxos/echo-schema';
-import { type FunctionDefinition, type FunctionExecutor } from '@dxos/functions';
+import { Expando } from '@dxos/echo/internal';
+import { type FunctionDefinition } from '@dxos/functions';
+import { type FunctionExecutor } from '@dxos/functions-runtime';
 import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { Message } from '@dxos/types';
 
 export const ExtractionInput = Schema.Struct({
-  message: DataType.Message,
+  message: Message.Message,
   objects: Schema.optional(Schema.Array(Expando)),
   options: Schema.optional(
     Schema.Struct({
@@ -23,7 +24,7 @@ export const ExtractionInput = Schema.Struct({
 export interface ExtractionInput extends Schema.Schema.Type<typeof ExtractionInput> {}
 
 export const ExtractionOutput = Schema.Struct({
-  message: DataType.Message,
+  message: Message.Message,
   timeElapsed: Schema.Number,
 });
 export interface ExtractionOutput extends Schema.Schema.Type<typeof ExtractionOutput> {}

package/src/extraction/quotes.ts

@@ -2,11 +2,11 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Schema } from 'effect';
+import * as Schema from 'effect/Schema';
 
 import { Filter, Query } from '@dxos/echo';
+import { ObjectId } from '@dxos/echo/internal';
 import { type EchoDatabase } from '@dxos/echo-db';
-import { ObjectId } from '@dxos/echo-schema';
 import { DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
 

package/src/functions/services.ts

@@ -2,97 +2,136 @@
 // Copyright 2025 DXOS.org
 //
 
-import { AiTool, type AiToolkit } from '@effect/ai';
-import { Context, Effect, Layer, Record, Schema } from 'effect';
+import * as Tool from '@effect/ai/Tool';
+import type * as Toolkit from '@effect/ai/Toolkit';
+import * as Context from 'effect/Context';
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+import * as Record from 'effect/Record';
+import * as Schema from 'effect/Schema';
 
 import { AiToolNotFoundError, ToolExecutionService, ToolResolverService } from '@dxos/ai';
 import { todo } from '@dxos/debug';
-import { type FunctionDefinition, LocalFunctionExecutionService } from '@dxos/functions';
+import { Query } from '@dxos/echo';
+import { DatabaseService } from '@dxos/echo-db';
+import { Function, FunctionDefinition, FunctionInvocationService } 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>[],
-  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);
+  functions: FunctionDefinition.Any[],
+  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(Function.Function, { 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>[],
-  toolkit: AiToolkit.AiToolkit<AiTool.Any>,
-  handlersLayer: Layer.Layer<AiTool.ToHandler<AiTool.AiTool<any>>, never, never>,
-): 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 toolkitHandler = yield* toolkit.pipe(Effect.provide(handlersLayer));
-      const localFunctionExecutionService = yield* LocalFunctionExecutionService;
+      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) {
-              if (toolkitHandler.tools.find((t: AiTool.Any) => t.name === 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);
+                const { result } = yield* (toolkitHandler.handle as any)(tool.name, input);
+                return result;
               }
 
-              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));
-              }
+              const { definition: functionDef } = Context.get(FunctionToolAnnotation)(tool.annotations as any);
 
-              return yield* localFunctionExecutionService
-                .invokeFunction(functionDef, input)
+              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 }
 >() {}
 
-const toolCache = new WeakMap<FunctionDefinition<any, any>, AiTool.Any>();
+const toolCache = new WeakMap<FunctionDefinition.Any, Tool.Any>();
 
-const projectFunctionToTool = (fn: FunctionDefinition<any, any>): AiTool.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): Tool.Any => {
   if (toolCache.has(fn)) {
     return toolCache.get(fn)!;
   }
 
-  const tool = AiTool.make(makeToolName(fn.name), {
+  // TODO(burdon): Use and map function.key?
+  // TODO(dmaretskyi): Use function key instead.
+  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;
 };
@@ -101,8 +140,8 @@ 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, '_');
-  invariant(toolName.match(/^[a-zA-Z_][a-zA-Z0-9-_]*$/));
+  const toolName = name.toLowerCase().replace(/[^a-zA-Z0-9]/g, '-');
+  invariant(toolName.match(/^[a-z_][a-z0-9-_]*$/));
   return toolName;
 };
 
@@ -117,3 +156,7 @@ const createStructFieldsFromSchema = (schema: Schema.Schema<any, any>): Record<s
       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

@@ -4,8 +4,9 @@
 
 export * from './conversation';
 export * from './diff';
+export * from './errors';
+export * from './execution-graph';
 export * from './functions';
 export * from './session';
 export * from './templates';
 export * from './types';
-export * from './errors';

package/src/session/GenericToolkit.ts

@@ -0,0 +1,130 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import type * as Tool from '@effect/ai/Tool';
+import * as Toolkit from '@effect/ai/Toolkit';
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+import * as Pipeable from 'effect/Pipeable';
+import type * as Schema from 'effect/Schema';
+import type * as SchemaAST from 'effect/SchemaAST';
+
+/**
+ * Unique identifier for generic toolkit instances.
+ */
+export const TypeId = '~@dxos/assistant/GenericToolkit';
+
+/**
+ * Type-level representation of the generic toolkit identifier.
+ */
+export type TypeId = typeof TypeId;
+
+/**
+ * Type-safe way to define toolkits where we don't know specific types of tools,
+ * and we want to bundle definition and handlers together.
+ *
+ * Usefull for plugin-based systems, where tools are dynamically added.
+ *
+ * @param TR requirements type to the tool invocation.
+ * @param E failure type to the handlers layer.
+ * @param R requirements type to the handlers layer.
+ */
+export interface GenericToolkit<TR = never, E = never, R = never> extends Pipeable.Pipeable {
+  readonly [TypeId]: TypeId;
+
+  /**
+   * Toolkit definition.
+   */
+  readonly toolkit: Toolkit.Toolkit<GenericTools<TR>>;
+
+  /**
+   * Handlers layer.
+   */
+  readonly layer: Layer.Layer<unknown, E, R>;
+
+  /**
+   * Handlers effect.
+   */
+  readonly handlers: Effect.Effect<Toolkit.WithHandler<GenericTools<TR>>, E, R>;
+}
+
+/**
+ * Generic constraint.
+ *
+ * NOTE: Only use in place of `T extends PortableToolkit.Any`. Not suitable for standalone use.
+ */
+export interface Any {
+  readonly [TypeId]: TypeId;
+  readonly toolkit: Toolkit.Toolkit<any>;
+  readonly layer: Layer.Layer<unknown, any, any>;
+}
+
+export type InvocationRequirements<T extends Any> = T extends GenericToolkit<infer TR, infer _E, infer _R> ? TR : never;
+export type Failure<T extends Any> = T extends GenericToolkit<infer _TR, infer E, infer _R> ? E : never;
+export type Requirements<T extends Any> = T extends GenericToolkit<infer _TR, infer _E, infer R> ? R : never;
+
+/**
+ * Creates a portable toolkit from toolkit definition and handlers layer.
+ */
+export const make = <Tools extends Record<string, Tool.Any>, E, R>(
+  toolkit: Toolkit.Toolkit<Tools>,
+  layer: Layer.Layer<Tool.HandlersFor<Tools>, E, R>,
+): GenericToolkit<E, R> => {
+  return {
+    [TypeId]: TypeId,
+    toolkit: toolkit as any,
+    layer: layer as any,
+    handlers: toolkit.pipe(Effect.provide(layer)) as any,
+    pipe() {
+      // eslint-disable-next-line prefer-rest-params
+      return Pipeable.pipeArguments(this, arguments);
+    },
+  };
+};
+
+/**
+ * Merges multiple portable toolkits into a single portable toolkit.
+ */
+export const merge = <const Toolkits extends ReadonlyArray<Any>>(
+  /**
+   * The toolkits to merge together.
+   */
+  ...toolkits: Toolkits
+): GenericToolkit<
+  InvocationRequirements<Toolkits[number]>,
+  Failure<Toolkits[number]>,
+  Requirements<Toolkits[number]>
+> => {
+  return make(
+    Toolkit.merge(...toolkits.map((t) => t.toolkit)),
+    Layer.mergeAll(...(toolkits.map((t) => t.layer) as any)),
+  ) as any;
+};
+
+/**
+ * Generic tools type where we cannot statically enumerate individual tools.
+ */
+export type GenericTools<R = never> = Record<
+  string,
+  Tool.Tool<
+    string,
+    {
+      readonly parameters: AnyStructSchemaNoContext;
+      readonly success: Schema.Schema.AnyNoContext;
+      readonly failure: typeof Schema.Never;
+    },
+    R
+  >
+>;
+
+export interface AnyStructSchemaNoContext extends Pipeable.Pipeable {
+  readonly [Schema.TypeId]: any;
+  readonly make: any;
+  readonly Type: any;
+  readonly Encoded: any;
+  readonly Context: never;
+  readonly ast: SchemaAST.AST;
+  readonly fields: Schema.Struct.Fields;
+  readonly annotations: any;
+}

package/src/session/artifact-diff.ts

@@ -2,10 +2,10 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Context } from 'effect';
+import * as Context from 'effect/Context';
 
+import { type ObjectId } from '@dxos/echo/internal';
 import { type ObjectVersion } from '@dxos/echo-db';
-import { type ObjectId } from '@dxos/echo-schema';
 
 /**
  * Resolves artifact ids to their versions.

package/src/session/error-handling.ts

@@ -2,9 +2,10 @@
 // Copyright 2025 DXOS.org
 //
 
-import { AiError } from '@effect/ai';
-import { HttpClientError } from '@effect/platform';
-import { Effect, Schema } from 'effect';
+import * as AiError from '@effect/ai/AiError';
+import * as HttpClientError from '@effect/platform/HttpClientError';
+import * as Effect from 'effect/Effect';
+import * as Schema from 'effect/Schema';
 
 import { AiModelError } from '../errors';
 
@@ -16,6 +17,7 @@ const AnthropicErrorResponse = Schema.Struct({
   }),
 });
 
+// TODO(dmaretskyi): Needs rework
 export const mapAiError = (err: AiError.AiError): Effect.Effect<AiError.AiError> =>
   Effect.gen(function* () {
     const cause = err.cause;
@@ -28,7 +30,7 @@ export const mapAiError = (err: AiError.AiError): Effect.Effect<AiError.AiError>
         message: body.error.message,
         context: { model: 'anthropic', type: body.error.type },
       });
-      return new AiError.AiError({
+      return AiError.UnknownError.make({
         description: body.error.message,
         module: err.module,
         method: err.method,

package/src/session/format.test.ts

@@ -3,13 +3,14 @@
 //
 
 import { beforeAll, describe, it } from '@effect/vitest';
-import { Effect } from 'effect';
+import * as Effect from 'effect/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 { Text } from '@dxos/schema';
+import { Organization } from '@dxos/types';
 import { trim } from '@dxos/util';
 
 import { createSystemPrompt } from '../templates/system';
@@ -21,7 +22,7 @@ describe('format', () => {
   beforeAll(async () => {
     const builder = await new EchoTestBuilder().open();
     ({ db } = await builder.createDatabase({
-      types: [DataType.Text, DataType.Organization, Blueprint.Blueprint],
+      types: [Text.Text, Organization.Organization, Blueprint.Blueprint],
     }));
   });
 
@@ -29,7 +30,7 @@ describe('format', () => {
     'should format',
     Effect.fn(function* () {
       const object = db.add(
-        Obj.make(DataType.Organization, {
+        Obj.make(Organization.Organization, {
           name: 'Test',
           website: 'https://www.test.com',
         }),