@dxos/assistant-toolkit 0.8.4-main.bc674ce → 0.8.4-main.c85a9c8dae

package/src/functions/entity-extraction/entity-extraction.test.ts

@@ -4,53 +4,34 @@
 
 import { describe, expect, it } from '@effect/vitest';
 import * as Effect from 'effect/Effect';
-import * as Layer from 'effect/Layer';
 
-import { AiService } from '@dxos/ai';
-import { MemoizedAiService, TestAiService } from '@dxos/ai/testing';
-import { makeToolExecutionServiceFromFunctions, makeToolResolverFromFunctions } from '@dxos/assistant';
+import { MemoizedAiService } from '@dxos/ai/testing';
+import { AssistantTestLayer } from '@dxos/assistant/testing';
 import { Blueprint } from '@dxos/blueprints';
 import { Obj } from '@dxos/echo';
 import { Database } from '@dxos/echo';
 import { TestHelpers } from '@dxos/effect/testing';
-import { CredentialsService, FunctionInvocationService, TracingService } from '@dxos/functions';
-import { FunctionInvocationServiceLayerTest, TestDatabaseLayer } from '@dxos/functions-runtime/testing';
+import { FunctionInvocationService } from '@dxos/functions';
 import { ObjectId } from '@dxos/keys';
 import { Message, Organization, Person } from '@dxos/types';
 
-import { testToolkit } from '../../blueprints/testing';
-import { ResearchGraph } from '../research';
+import { ResearchGraph } from '../../blueprints';
 
 import { default as entityExtraction } from './entity-extraction';
 
 ObjectId.dangerouslyDisableRandomness();
 
-const TestLayer = Layer.mergeAll(
-  AiService.model('@anthropic/claude-opus-4-0'),
-  makeToolResolverFromFunctions([], testToolkit),
-  makeToolExecutionServiceFromFunctions(testToolkit, testToolkit.toLayer({}) as any),
-).pipe(
-  Layer.provideMerge(FunctionInvocationServiceLayerTest({ functions: [entityExtraction] })),
-  Layer.provideMerge(
-    Layer.mergeAll(
-      TestAiService(),
-      TestDatabaseLayer({
-        spaceKey: 'fixed',
-        indexing: { vector: true },
-        types: [Blueprint.Blueprint, Message.Message, Person.Person, Organization.Organization, ResearchGraph],
-      }),
-      CredentialsService.configuredLayer([]),
-      TracingService.layerNoop,
-    ),
-  ),
-);
+const TestLayer = AssistantTestLayer({
+  functions: [entityExtraction],
+  types: [Blueprint.Blueprint, Message.Message, Person.Person, Organization.Organization, ResearchGraph.ResearchGraph],
+});
 
 describe('Entity extraction', () => {
   it.effect(
-    'call a function to generate a research report',
+    'calls a function to generate a research report',
     Effect.fnUntraced(
       function* (_) {
-        const email = yield* Database.Service.add(
+        const email = yield* Database.add(
           Obj.make(Message.Message, {
             [Obj.Meta]: {
               tags: ['important'],
@@ -76,7 +57,7 @@ describe('Entity extraction', () => {
             ],
           }),
         );
-        yield* Database.Service.flush({ indexes: true });
+        yield* Database.flush({ indexes: true });
         const result = yield* FunctionInvocationService.invokeFunction(entityExtraction, {
           source: email,
         });

package/src/functions/entity-extraction/entity-extraction.ts

@@ -2,16 +2,15 @@
 // Copyright 2025 DXOS.org
 //
 
-import * as Toolkit from '@effect/ai/Toolkit';
 import * as Effect from 'effect/Effect';
 import * as Layer from 'effect/Layer';
 import * as Predicate from 'effect/Predicate';
 import * as Schema from 'effect/Schema';
 
 import { AiService } from '@dxos/ai';
-import { AiSession, makeToolExecutionServiceFromFunctions, makeToolResolverFromFunctions } from '@dxos/assistant';
-import { Filter, Obj, Ref, Type } from '@dxos/echo';
-import { Database } from '@dxos/echo';
+import { GenericToolkit } from '@dxos/ai';
+import { AiSession, ToolExecutionServices } from '@dxos/assistant';
+import { Database, Filter, Obj, Ref } from '@dxos/echo';
 import { defineFunction } from '@dxos/functions';
 import { FunctionInvocationServiceLayerTest } from '@dxos/functions-runtime/testing';
 import { type DXN } from '@dxos/keys';
@@ -19,8 +18,8 @@ import { log } from '@dxos/log';
 import { type Actor, LegacyOrganization, Message, Organization, Person } from '@dxos/types';
 import { trim } from '@dxos/util';
 
+import { ResearchGraph } from '../../blueprints';
 import { makeGraphWriterHandler, makeGraphWriterToolkit } from '../../crud';
-import { contextQueueLayerFromResearchGraph } from '../research';
 
 export default defineFunction({
   key: 'dxos.org/functions/entity-extraction',
@@ -38,7 +37,7 @@ export default defineFunction({
   }),
   outputSchema: Schema.Struct({
     entities: Schema.optional(
-      Schema.Array(Type.Obj).annotations({
+      Schema.Array(Obj.Unknown).annotations({
         description: 'Extracted entities.',
       }),
     ),
@@ -58,7 +57,7 @@ export default defineFunction({
           onAppend: (dxns) => created.push(...dxns),
         });
         const toolkit = yield* GraphWriterToolkit.pipe(
-          Effect.provide(GraphWriterHandler.pipe(Layer.provide(contextQueueLayerFromResearchGraph))),
+          Effect.provide(GraphWriterHandler.pipe(Layer.provide(ResearchGraph.contextQueueLayer))),
         );
 
         yield* new AiSession().run({
@@ -74,7 +73,7 @@ export default defineFunction({
         if (created.length > 1) {
           throw new Error('Multiple organizations created');
         } else if (created.length === 1) {
-          organization = yield* Database.Service.resolve(created[0], Organization.Organization);
+          organization = yield* Database.resolve(created[0], Organization.Organization);
           Obj.change(organization, (org) => {
             const meta = Obj.getMeta(org);
             meta.tags ??= [];
@@ -93,12 +92,11 @@ export default defineFunction({
     Effect.provide(
       Layer.mergeAll(
         AiService.model('@anthropic/claude-sonnet-4-0'), // TODO(dmaretskyi): Extract.
-        makeToolResolverFromFunctions([], Toolkit.make()),
-        makeToolExecutionServiceFromFunctions(Toolkit.make() as any, Layer.empty as any),
+        ToolExecutionServices,
       ).pipe(
         Layer.provide(
           // TODO(dmaretskyi): This should be provided by environment.
-          Layer.mergeAll(FunctionInvocationServiceLayerTest()),
+          Layer.mergeAll(GenericToolkit.providerEmpty, FunctionInvocationServiceLayerTest()),
         ),
       ),
     ),
@@ -113,7 +111,7 @@ const extractContact = Effect.fn('extractContact')(function* (actor: Actor.Actor
     return undefined;
   }
 
-  const existingContacts = yield* Database.Service.runQuery(Filter.type(Person.Person));
+  const existingContacts = yield* Database.runQuery(Filter.type(Person.Person));
 
   // Check for existing contact
   // TODO(dmaretskyi): Query filter DSL - https://linear.app/dxos/issue/DX-541/filtercontains-should-work-with-partial-objects
@@ -130,7 +128,7 @@ const extractContact = Effect.fn('extractContact')(function* (actor: Actor.Actor
     ...(tags ? { [Obj.Meta]: { tags: [...tags] } } : {}),
     emails: [{ value: email }],
   });
-  yield* Database.Service.add(newContact);
+  yield* Database.add(newContact);
 
   if (name) {
     Obj.change(newContact, (c) => {
@@ -146,7 +144,7 @@ const extractContact = Effect.fn('extractContact')(function* (actor: Actor.Actor
 
   log.info('extracted email domain', { emailDomain });
 
-  const existingOrganisations = yield* Database.Service.runQuery(Filter.type(Organization.Organization));
+  const existingOrganisations = yield* Database.runQuery(Filter.type(Organization.Organization));
   const matchingOrg = existingOrganisations.find((org) => {
     if (org.website) {
       try {

package/src/functions/entity-extraction/index.ts

@@ -2,8 +2,8 @@
 // Copyright 2025 DXOS.org
 //
 
-import { default as extract$ } from './entity-extraction';
+import Extract from './entity-extraction';
 
-export namespace EntityExtraction {
-  export const extract = extract$;
-}
+export const EntityExtractionFunctions = {
+  Extract,
+};

package/src/functions/exa/exa.ts

@@ -9,8 +9,8 @@ import Exa from 'exa-js';
 import { CredentialsService, defineFunction } from '@dxos/functions';
 
 export default defineFunction({
-  key: 'dxos.org/function/exa',
-  name: 'Exa',
+  key: 'dxos.org/function/exa-search',
+  name: 'Exa Search',
   description: 'Search the web for information',
   inputSchema: Schema.Struct({
     query: Schema.String.annotations({

package/src/functions/exa/index.ts

@@ -2,5 +2,10 @@
 // Copyright 2025 DXOS.org
 //
 
-export { default as exaFunction } from './exa';
-export { default as exaMockFunction } from './mock';
+import Search from './exa';
+import Mock from './mock';
+
+export const ExaFunctions = {
+  Search,
+  Mock,
+};

package/src/functions/index.ts

@@ -3,9 +3,5 @@
 //
 
 export * from './agent';
-export * from './discord';
-export * from './document';
 export * from './entity-extraction';
-export * from './linear';
-export * from './research';
-export * from './tasks';
+export * from './exa';

package/src/index.ts

@@ -3,7 +3,8 @@
 //
 
 export * from './blueprints';
+export * from './crud';
 export * from './functions';
 export * from './sync';
-export * from './crud';
 export * from './toolkits';
+export * from './types';

package/src/sync/sync.ts

@@ -42,7 +42,7 @@ export const syncObjects: (
 
       const schema = Obj.getSchema(obj) ?? failedInvariant('No schema.');
       const foreignId = Obj.getKeys(obj, foreignKeyId)[0]?.id ?? failedInvariant('No foreign key.');
-      const [existing] = yield* Database.Service.runQuery(
+      const [existing] = yield* Database.runQuery(
         Query.select(Filter.foreignKeys(schema, [{ source: foreignKeyId, id: foreignId }])),
       );
       log('sync object', {
@@ -51,7 +51,7 @@ export const syncObjects: (
         existing: existing ? Obj.getDXN(existing) : undefined,
       });
       if (!existing) {
-        yield* Database.Service.add(obj);
+        yield* Database.add(obj);
         return obj;
       } else {
         copyObjectData(existing, obj);

package/src/testing/index.ts

@@ -2,4 +2,5 @@
 // Copyright 2025 DXOS.org
 //
 
+export * from '../blueprints/testing';
 export * from './plugins';

package/src/testing/plugins.tsx

@@ -5,8 +5,10 @@
 import * as Schema from 'effect/Schema';
 import React from 'react';
 
-import { Capability, Common } from '@dxos/app-framework';
+import { Capabilities, Capability } from '@dxos/app-framework';
+import { Surface } from '@dxos/app-framework/ui';
 import { Format, type Obj, Type } from '@dxos/echo';
+import { Card } from '@dxos/react-ui';
 import { JsonFilter } from '@dxos/react-ui-syntax-highlighter';
 
 export const MapSchema = Schema.Struct({
@@ -38,31 +40,33 @@ const isImage = (data: any): data is any => false;
 
 export const capabilities: Capability.Any[] = [
   Capability.contributes(
-    Common.Capability.ReactSurface,
-    Common.createSurface({
+    Capabilities.ReactSurface,
+    Surface.create({
       id: 'plugin-image',
-      role: 'card--extrinsic',
+      role: 'card--content',
       filter: (data: any): data is any => isImage(data.value),
       component: ({ data }) => (
-        <img
-          className='grow object-cover'
-          src={`data:image/jpeg;base64,${data.value.source.data}`}
-          alt={data.value.prompt ?? `Generated image [id=${data.value.id}]`}
-        />
+        <Card.Content>
+          <img
+            className='grow object-cover'
+            src={`data:image/jpeg;base64,${data.value.source.data}`}
+            alt={data.value.prompt ?? `Generated image [id=${data.value.id}]`}
+          />
+        </Card.Content>
       ),
     }),
   ),
-
-  //
-  // Default
-  //
   Capability.contributes(
-    Common.Capability.ReactSurface,
-    Common.createSurface({
+    Capabilities.ReactSurface,
+    Surface.create({
       id: 'plugin-default',
-      role: 'card--extrinsic',
+      role: 'card--content',
       position: 'fallback',
-      component: ({ data }) => <JsonFilter data={data} />,
+      component: ({ data }) => (
+        <Card.Content>
+          <JsonFilter data={data} />
+        </Card.Content>
+      ),
     }),
   ),
 ];

package/src/toolkits/WebToolkit.ts

@@ -0,0 +1,33 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Tool from '@effect/ai/Tool';
+import * as Toolkit from '@effect/ai/Toolkit';
+import * as AnthropicTool from '@effect/ai-anthropic/AnthropicTool';
+import * as Effect from 'effect/Effect';
+import * as Schema from 'effect/Schema';
+
+// TODO(dmaretskyi): Testing only.
+export const WebToolkit = Toolkit.make(
+  AnthropicTool.WebSearch_20250305({}).pipe(
+    // TODO(dmaretskyi): Effect bug -- provider-defined tools don't support annotations.
+    // ToolFormatter.assign({
+    //   debugFormatResult: (result) =>
+    //     Array.isArray(result) ? result.map(Struct.pick('title', 'type', 'url', 'page_age')) : result,
+    // }),
+  ),
+  Tool.make('WebFetch', {
+    parameters: {
+      url: Schema.String,
+    },
+    success: Schema.String,
+  }),
+);
+
+export const layer = WebToolkit.toLayer({
+  WebFetch: Effect.fnUntraced(function* ({ url }) {
+    const response = yield* Effect.promise(() => fetch(url).then((response) => response.text()));
+    return response;
+  }),
+});

package/src/toolkits/index.ts

@@ -2,5 +2,4 @@
 // Copyright 2025 DXOS.org
 //
 
-export * as AssistantToolkit from './AssistantToolkit';
-export * as SystemToolkit from './SystemToolkit';
+export * as WebToolkit from './WebToolkit';

package/src/types/Chat.ts

@@ -0,0 +1,49 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import * as Schema from 'effect/Schema';
+
+import { Annotation, Obj, Ref, Type } from '@dxos/echo';
+import { FormInputAnnotation, LabelAnnotation } from '@dxos/echo/internal';
+import { Queue } from '@dxos/echo-db';
+
+/**
+ * AI chat.
+ */
+export const Chat = Schema.Struct({
+  name: Schema.String.pipe(Schema.optional),
+  queue: Ref.Ref(Queue).pipe(FormInputAnnotation.set(false)),
+  // TODO(dmaretskyi): Eventually this and the message queue will be the same.
+  traceQueue: Ref.Ref(Queue).pipe(FormInputAnnotation.set(false), Schema.optional),
+}).pipe(
+  Type.object({
+    typename: 'dxos.org/type/assistant/Chat',
+    version: '0.2.0',
+  }),
+  LabelAnnotation.set(['name']),
+  Annotation.IconAnnotation.set({
+    icon: 'ph--atom--regular',
+    hue: 'blue',
+  }),
+);
+
+export interface Chat extends Schema.Schema.Type<typeof Chat> {}
+
+export const make = (props: Obj.MakeProps<typeof Chat>) => Obj.make(Chat, props);
+
+/**
+ * Relation between a Chat and companion objects (e.g., artifacts).
+ */
+export const CompanionTo = Schema.Struct({
+  id: Obj.ID,
+}).pipe(
+  Type.relation({
+    typename: 'dxos.org/relation/assistant/CompanionTo',
+    version: '0.1.0',
+    source: Chat,
+    target: Obj.Unknown,
+  }),
+);
+
+export interface CompanionTo extends Schema.Schema.Type<typeof CompanionTo> {}

package/src/types/Plan.ts

@@ -0,0 +1,107 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import * as Schema from 'effect/Schema';
+
+import { Obj, Ref, Type } from '@dxos/echo';
+
+import * as Chat from './Chat';
+
+export const TaskId = Schema.String.pipe(Schema.brand('@dxos/assistant-toolkit/TaskId'));
+export type TaskId = Schema.Schema.Type<typeof TaskId>;
+
+export const Task = Schema.Struct({
+  /**
+   * Short task ID unique within the plan.
+   */
+  id: TaskId,
+
+  title: Schema.String.annotations({
+    description: 'Task title and description.',
+  }),
+
+  status: Schema.Literal('todo', 'in-progress', 'done'),
+
+  /**
+   * Parent task ID.
+   */
+  parent: Schema.optional(TaskId).annotations({
+    description: 'Parent task ID.',
+  }),
+
+  /**
+   * Chat object that this task is associated with.
+   */
+  chat: Schema.optional(Ref.Ref(Chat.Chat)),
+});
+export interface Task extends Schema.Schema.Type<typeof Task> {}
+
+/**
+ * Hierarchical collection of tasks for humans and agents to track progress.
+ */
+export const Plan = Schema.Struct({
+  tasks: Schema.Array(Task),
+}).pipe(
+  Type.object({
+    typename: 'dxos.org/type/Plan',
+    version: '0.1.0',
+  }),
+);
+export interface Plan extends Schema.Schema.Type<typeof Plan> {}
+
+export const generateTaskId = (plan: Plan): TaskId => {
+  const existingIds = plan.tasks
+    .map((task) => {
+      const [num, letter] = task.id.split('-');
+      return parseInt(num);
+    })
+    .filter((_) => !isNaN(_));
+  let newId;
+  if (existingIds.length === 0) {
+    newId = 1;
+  } else {
+    newId = Math.max(...existingIds) + 1;
+  }
+  // Add random suffix to avoid collisions.
+  return TaskId.make(`${newId}-${crypto.randomUUID().slice(0, 2)}`);
+};
+
+/**
+ * Add new tasks to a plan, generating IDs for new tasks.
+ * Use inside an `Obj.change` callback.
+ */
+export const addTasks = (
+  plan: Obj.Mutable<Plan>,
+  tasks: (Pick<Task, 'title'> & Partial<Pick<Task, 'status' | 'parent' | 'chat'>>)[],
+) => {
+  for (const task of tasks) {
+    const taskId = generateTaskId(plan);
+    plan.tasks.push({ id: taskId, ...task, status: task.status ?? 'todo' });
+  }
+  return plan;
+};
+
+interface MakePlanProps {
+  tasks: {
+    title: string;
+    status?: 'todo' | 'in-progress' | 'done';
+  }[];
+}
+
+export const makePlan = (props: MakePlanProps): Plan => {
+  const plan = Obj.make(Plan, { tasks: [] });
+  Obj.change(plan, (plan) => {
+    addTasks(plan, props.tasks);
+  });
+  return plan;
+};
+
+/**
+ * Formats plan to markdown format.
+ */
+export const formatPlan = (plan: Plan): string => {
+  return plan.tasks
+    .map((task) => `- **${task.status?.toLocaleUpperCase()}**: ${task.title ?? 'No title'} <!-- id=${task.id} -->`)
+    .join('\n');
+};