@dxos/assistant 0.8.3 → 0.8.4-main.2e9d522

package/src/research/exa.ts

@@ -2,78 +2,75 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Schema } from 'effect';
+import { AiTool, AiToolkit } from '@effect/ai';
+import { Effect, Schema } from 'effect';
 import Exa from 'exa-js';
 
-import { createTool, ToolResult } from '@dxos/ai';
+import { CredentialsService } from '@dxos/functions';
 import { log } from '@dxos/log';
 
 import { SEARCH_RESULTS } from '../testing';
 
-type CreateExaToolOptions = {
-  apiKey: string;
-};
-
-export const createExaTool = ({ apiKey }: CreateExaToolOptions) => {
-  const exa = new Exa(apiKey);
-
-  return createTool('search', {
-    name: 'web_search',
+export class ExaToolkit extends AiToolkit.make(
+  AiTool.make('exa_search', {
     description: 'Search the web for information',
-    schema: Schema.Struct({
+    parameters: {
       query: Schema.String.annotations({ description: 'The query to search for' }),
-    }),
-    execute: async ({ query }) => {
-      const context = await exa.searchAndContents(query, {
+    },
+    success: Schema.Unknown,
+    failure: Schema.Never,
+  }).addRequirement<CredentialsService>(),
+) {}
+
+export const LiveExaHandler = ExaToolkit.toLayer({
+  exa_search: Effect.fn(function* ({ query }) {
+    const credential = yield* CredentialsService.getCredential({ service: 'exa.ai' });
+    const exa = new Exa(credential.apiKey);
+
+    const context = yield* Effect.promise(async () =>
+      exa.searchAndContents(query, {
         type: 'auto',
         text: {
           maxCharacters: 3_000,
         },
         livecrawl: 'always',
-      });
-
-      log.info('exa search', { query, costDollars: context.costDollars });
-
-      return ToolResult.Success(context);
-    },
-  });
-};
-
-export const createMockExaTool = () => {
-  return createTool('search', {
-    name: 'web_search',
-    description: 'Search the web for information',
-    schema: Schema.Struct({
-      query: Schema.String.annotations({ description: 'The query to search for' }),
-    }),
-    execute: async ({ query }) => {
-      // Find result with closest matching autoprompt using weighted Levenshtein distance
-      const result = SEARCH_RESULTS.reduce(
-        (closest, current) => {
-          if (!current.autopromptString) {
-            return closest;
-          }
-          if (!closest) {
-            return current;
-          }
-
-          // Calculate Levenshtein distance
-          const dist1 = levenshteinDistance(query, current.autopromptString);
-          const dist2 = levenshteinDistance(query, closest.autopromptString || '');
-
-          // Weight by length of the longer string to normalize
-          const weight1 = dist1 / Math.max(query.length, current.autopromptString.length);
-          const weight2 = dist2 / Math.max(query.length, closest.autopromptString?.length || 0);
-
-          return weight1 < weight2 ? current : closest;
-        },
-        null as (typeof SEARCH_RESULTS)[0] | null,
-      );
-
-      return ToolResult.Success(result);
-    },
-  });
-};
+      }),
+    );
+
+    log.info('exa search', { query, costDollars: context.costDollars });
+
+    return context;
+  }),
+});
+
+export const MockExaHandler = ExaToolkit.toLayer({
+  exa_search: Effect.fn(function* ({ query }) {
+    // Find result with closest matching autoprompt using weighted Levenshtein distance
+    const result = SEARCH_RESULTS.reduce(
+      (closest, current) => {
+        if (!current.autopromptString) {
+          return closest;
+        }
+        if (!closest) {
+          return current;
+        }
+
+        // Calculate Levenshtein distance
+        const dist1 = levenshteinDistance(query, current.autopromptString);
+        const dist2 = levenshteinDistance(query, closest.autopromptString || '');
+
+        // Weight by length of the longer string to normalize
+        const weight1 = dist1 / Math.max(query.length, current.autopromptString.length);
+        const weight2 = dist2 / Math.max(query.length, closest.autopromptString?.length || 0);
+
+        return weight1 < weight2 ? current : closest;
+      },
+      null as (typeof SEARCH_RESULTS)[0] | null,
+    );
+
+    return result;
+  }),
+});
 
 const levenshteinDistance = (str1: string, str2: string): number => {
   const m = str1.length;

package/src/research/graph.test.ts

@@ -2,14 +2,63 @@
 // Copyright 2025 DXOS.org
 //
 
-import { describe } from 'vitest';
+import { AnthropicClient } from '@effect/ai-anthropic';
+import { FetchHttpClient } from '@effect/platform';
+import { describe, it } from '@effect/vitest';
+import { Config, Effect, Layer, pipe } from 'effect';
+
+import { AiService, AiServiceRouter, ToolExecutionService, ToolResolverService } from '@dxos/ai';
+import { tapHttpErrors } from '@dxos/ai/testing';
+import { TestHelpers } from '@dxos/effect';
+import { DatabaseService } from '@dxos/functions';
+import { log } from '@dxos/log';
+import { DataType } from '@dxos/schema';
+
+import { makeGraphWriterHandler, makeGraphWriterToolkit } from './graph';
+import { AiSession } from '../session';
 
 // import { type EchoTestBuilder } from '@dxos/echo-db/testing';
 
+const TestLayer = pipe(
+  AiService.model('@anthropic/claude-3-5-sonnet-20241022'),
+  Layer.provideMerge(DatabaseService.notAvailable),
+  Layer.provideMerge(ToolResolverService.layerEmpty),
+  Layer.provideMerge(ToolExecutionService.layerEmpty),
+  Layer.provide(AiServiceRouter.AiServiceRouter),
+  Layer.provide(
+    AnthropicClient.layerConfig({
+      apiKey: Config.redacted('ANTHROPIC_API_KEY'),
+      transformClient: tapHttpErrors,
+    }),
+  ),
+  Layer.provide(FetchHttpClient.layer),
+);
+
 describe('graph', () => {
   // let builder: EchoTestBuilder;
   // test('findRelatedSchema', async () => {
   //   const db = await createEchoDatabase();
   //   const relatedSchemas = await findRelatedSchema(db, Schema.Struct({}));
   // });
+
+  it.effect.skip(
+    'calculator',
+    Effect.fn(
+      function* ({ expect }) {
+        const graphWriteToolkit = makeGraphWriterToolkit({ schema: [DataType.Project] });
+
+        const session = new AiSession();
+        const response = yield* session
+          .run({
+            prompt: 'What is 10 + 20?',
+            history: [],
+            toolkit: graphWriteToolkit,
+          })
+          .pipe(Effect.provide(makeGraphWriterHandler(graphWriteToolkit)));
+        log.info('response', { response });
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.runIf(process.env.ANTHROPIC_API_KEY),
+    ),
+  );
 });

package/src/research/graph.ts

@@ -2,30 +2,31 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Schema, identity, Option, SchemaAST } from 'effect';
+import { AiTool, AiToolkit } from '@effect/ai';
+import { Context, Effect, identity, Option, Schema, SchemaAST } from 'effect';
 
-import { createTool, ToolResult } from '@dxos/ai';
 import type { Obj, Relation } from '@dxos/echo';
-import { Query, Filter } from '@dxos/echo';
+import { Filter, Query } from '@dxos/echo';
 import { type EchoDatabase, type Queue } from '@dxos/echo-db';
 import { isEncodedReference } from '@dxos/echo-protocol';
 import {
+  create,
   EntityKind,
+  getEntityKind,
+  getSchemaDXN,
   getSchemaTypename,
   getTypeAnnotation,
   getTypeIdentifierAnnotation,
-  create,
-  getEntityKind,
-  getSchemaDXN,
   ObjectId,
   ReferenceAnnotationId,
-  type BaseObject,
   RelationSourceDXNId,
-  RelationTargetDXNId,
   RelationSourceId,
+  RelationTargetDXNId,
   RelationTargetId,
+  type BaseObject,
 } from '@dxos/echo-schema';
 import { mapAst } from '@dxos/effect';
+import { ContextQueueService, DatabaseService } from '@dxos/functions';
 import { DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { deepMapValues, isNonNullable } from '@dxos/util';
@@ -96,33 +97,80 @@ const isSchemaAddressableByDxn = (schema: Schema.Schema.AnyNoContext, dxn: DXN):
 /**
  * Perform vector search in the local database.
  */
-export const createLocalSearchTool = (db: EchoDatabase, queue?: Queue) => {
-  return createTool('search', {
-    name: 'local_search',
+// TODO(dmaretskyi): Rename `GraphReadToolkit`.
+export class LocalSearchToolkit extends AiToolkit.make(
+  AiTool.make('search_local_search', {
     description: 'Search the local database for information using a vector index',
-    schema: Schema.Struct({
+    parameters: {
       query: Schema.String.annotations({
         description: 'The query to search for. Could be a question or a topic or a set of keywords.',
       }),
-    }),
-    execute: async ({ query }) => {
-      const { objects } = await db.query(Query.select(Filter.text(query, { type: 'vector' }))).run();
-      const results = [...objects];
-      if (queue) {
-        const queueObjects = await queue.queryObjects();
-        // TODO(dmaretskyi): Text search on the queue.
-        results.push(...queueObjects);
-      }
-
-      return ToolResult.Success(`
-        <local_context>
-          ${JSON.stringify(results, null, 2)}
-        </local_context>
-        `);
     },
-  });
+    success: Schema.Unknown,
+    failure: Schema.Never,
+  }).addRequirement<DatabaseService>(),
+) {}
+
+export const LocalSearchHandler = LocalSearchToolkit.toLayer({
+  search_local_search: Effect.fn(function* ({ query }) {
+    const { objects } = yield* DatabaseService.runQuery(Query.select(Filter.text(query, { type: 'vector' })));
+    const results = [...objects];
+
+    const contextQueue = yield* Effect.serviceOption(ContextQueueService);
+    if (Option.isSome(contextQueue)) {
+      const queueObjects = yield* Effect.promise(() => contextQueue.value.contextQueue.queryObjects());
+      // TODO(dmaretskyi): Text search on the queue.
+      results.push(...queueObjects);
+    }
+
+    return `
+      <local_context>
+        ${JSON.stringify(results, null, 2)}
+      </local_context>
+      `;
+  }),
+});
+
+/**
+ * Attached as an annotation to the writer tool.
+ */
+class GraphWriterSchema extends Context.Tag('GraphWriterSchema')<
+  GraphWriterSchema,
+  { schema: Schema.Schema.AnyNoContext[] }
+>() {}
+
+/**
+ * Forms typed objects that can be written to the graph database.
+ */
+export const makeGraphWriterToolkit = ({ schema }: { schema: Schema.Schema.AnyNoContext[] }) => {
+  return AiToolkit.make(
+    AiTool.make('graph_writer', {
+      description: 'Write to the local graph database',
+      parameters: createExtractionSchema(schema).fields,
+      success: Schema.Unknown,
+      failure: Schema.Never,
+    })
+      .addRequirement<DatabaseService>()
+      .annotateContext(Context.make(GraphWriterSchema, { schema })),
+  );
 };
 
+export const makeGraphWriterHandler = (toolkit: ReturnType<typeof makeGraphWriterToolkit>) => {
+  const { schema } = Context.get(
+    toolkit.tools.graph_writer.annotations as Context.Context<GraphWriterSchema>,
+    GraphWriterSchema,
+  );
+  return toolkit.toLayer({
+    graph_writer: Effect.fn(function* (input) {
+      const { db } = yield* DatabaseService;
+      const contextQueue = yield* Effect.serviceOption(ContextQueueService);
+      const data = yield* Effect.promise(() =>
+        sanitizeObjects(schema, input as any, db, contextQueue.pipe(Option.getOrUndefined)?.contextQueue),
+      );
+      return data as Obj.Any[];
+    }),
+  });
+};
 /**
  * Create a schema for structured data extraction.
  */
@@ -282,15 +330,21 @@ const SoftRef = Schema.Struct({
 const preprocessSchema = (schema: Schema.Schema.AnyNoContext) => {
   const isRelationSchema = getEntityKind(schema) === 'relation';
 
-  const go = (ast: SchemaAST.AST): SchemaAST.AST => {
+  const go = (ast: SchemaAST.AST, visited = new Set<SchemaAST.AST>()): SchemaAST.AST => {
+    if (visited.has(ast)) {
+      // Already visited this node, prevent infinite recursion.
+      return ast;
+    }
+    visited.add(ast);
+
     if (SchemaAST.getAnnotation(ast, ReferenceAnnotationId).pipe(Option.isSome)) {
       return SoftRef.ast;
     }
 
-    return mapAst(ast, go);
+    return mapAst(ast, (child) => go(child, visited));
   };
 
-  return Schema.make<any, any, never>(mapAst(schema.ast, go)).pipe(
+  return Schema.make<any, any, never>(mapAst(schema.ast, (ast) => go(ast))).pipe(
     Schema.omit('id'),
     Schema.extend(
       Schema.Struct({
@@ -313,25 +367,3 @@ const preprocessSchema = (schema: Schema.Schema.AnyNoContext) => {
       : identity<Schema.Schema.AnyNoContext>,
   );
 };
-
-export const createGraphWriterTool = ({
-  db,
-  queue,
-  schema,
-  onDone = async (x) => x,
-}: {
-  db: EchoDatabase;
-  queue?: Queue;
-  schema: Schema.Schema.AnyNoContext[];
-  onDone?: (data: Obj.Any[]) => Promise<any>;
-}) => {
-  return createTool('graph', {
-    name: 'writer',
-    description: 'Write to the local graph database',
-    schema: createExtractionSchema(schema),
-    execute: async (input) => {
-      const data = await sanitizeObjects(schema, input as any, db, queue);
-      return ToolResult.Success(await onDone(data as Obj.Any[]));
-    },
-  });
-};

package/src/research/instructions-research.tpl

@@ -25,6 +25,7 @@ The Research Agent always follows these principles:
 - The Research Agent creates relations and references between new objects found and what's already in the database.
 - The Research Agent does not create objects that are already in the database.
 - The Research Agent re-uses existing object IDs as references when enriching existing objects.
+- The Research Agent ALWAYS calls the `graph_writer` at the end to save the data. This conversation will be delted, so only the data written to the graph will be preserved.
 
 The Research Agent may be asked for:
 

package/src/research/research.test.ts

@@ -2,11 +2,14 @@
 // Copyright 2025 DXOS.org
 //
 
+import { AnthropicClient } from '@effect/ai-anthropic';
+import { FetchHttpClient } from '@effect/platform';
+import { Config, Layer, ManagedRuntime } from 'effect';
 import { inspect } from 'node:util';
-import { beforeAll, describe, test } from 'vitest';
+import { afterAll, beforeAll, describe, test } from 'vitest';
 
-import { EdgeAiServiceClient, OllamaAiServiceClient, structuredOutputParser } from '@dxos/ai';
-import { AI_SERVICE_ENDPOINT, EXA_API_KEY } from '@dxos/ai/testing';
+import { AiService, AiServiceRouter, structuredOutputParser } from '@dxos/ai';
+import { EXA_API_KEY, tapHttpErrors } from '@dxos/ai/testing';
 import { Obj } from '@dxos/echo';
 import { type EchoDatabase } from '@dxos/echo-db';
 import { EchoTestBuilder } from '@dxos/echo-db/testing';
@@ -17,38 +20,35 @@ import { DataType, DataTypes } from '@dxos/schema';
 import { createExtractionSchema, getSanitizedSchemaName } from './graph';
 import { researchFn } from './research';
 
-const REMOTE_AI = true;
 const MOCK_SEARCH = false;
 
+const AnthropicLayer = AnthropicClient.layerConfig({
+  apiKey: Config.redacted('ANTHROPIC_API_KEY'),
+  transformClient: tapHttpErrors,
+});
+
+const AiServiceLayer = AiServiceRouter.AiServiceRouter.pipe(
+  Layer.provide(AnthropicLayer),
+  Layer.provide(FetchHttpClient.layer),
+);
+
 describe.skip('Research', () => {
+  let runtime: ManagedRuntime.ManagedRuntime<AiService, any>;
   let builder: EchoTestBuilder;
   let db: EchoDatabase;
   let executor: FunctionExecutor;
 
   beforeAll(async () => {
+    runtime = ManagedRuntime.make(AiServiceLayer);
+
     // TODO(dmaretskyi): Helper to scaffold this from a config.
     builder = await new EchoTestBuilder().open();
 
     db = (await builder.createDatabase({ indexing: { vector: true } })).db;
-    db.graph.schemaRegistry.addSchema(DataTypes);
 
     executor = new FunctionExecutor(
       new ServiceContainer().setServices({
-        ai: {
-          client: REMOTE_AI
-            ? new EdgeAiServiceClient({
-                endpoint: AI_SERVICE_ENDPOINT.REMOTE,
-                defaultGenerationOptions: {
-                  // model: '@anthropic/claude-sonnet-4-20250514',
-                  model: '@anthropic/claude-3-5-sonnet-20241022',
-                },
-              })
-            : new OllamaAiServiceClient({
-                overrides: {
-                  model: 'llama3.1:8b',
-                },
-              }),
-        },
+        ai: await runtime.runPromise(AiService),
         credentials: new ConfiguredCredentialsService([{ service: 'exa.ai', apiKey: EXA_API_KEY }]),
         database: { db },
         tracing: TracingService.console,
@@ -56,13 +56,13 @@ describe.skip('Research', () => {
     );
   });
 
-  test('should generate a research report', { timeout: 300_000 }, async () => {
-    db.add(
-      Obj.make(DataType.Organization, {
-        name: 'Notion',
-        website: 'https://www.notion.com',
-      }),
-    );
+  afterAll(async () => {
+    await runtime.dispose();
+  });
+
+  // TODO(dmaretskyi): Refactor using effect.
+  test.only('should generate a research report', { timeout: 300_000 }, async () => {
+    db.add(Obj.make(DataType.Organization, { name: 'Notion', website: 'https://www.notion.com' }));
     await db.flush({ indexes: true });
 
     const result = await executor.invoke(researchFn, {

package/src/research/research.ts

@@ -2,19 +2,22 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Schema } from 'effect';
+// TODO(burdon): Fix.
+// @ts-nocheck
 
-import { AgentStatusReport, ConsolePrinter } from '@dxos/ai';
-import { create } from '@dxos/echo-schema';
-import { AiService, CredentialsService, DatabaseService, defineFunction, TracingService } from '@dxos/functions';
+import { AiToolkit } from '@effect/ai';
+import { Effect, Layer, Schema } from 'effect';
+
+import { AiService, ConsolePrinter } from '@dxos/ai';
+import { defineFunction, TracingService } from '@dxos/functions';
 import { log } from '@dxos/log';
 import { DataTypes } from '@dxos/schema';
 
-import { createExaTool, createMockExaTool } from './exa';
-import { createGraphWriterTool } from './graph';
+import { ExaToolkit, LiveExaHandler, MockExaHandler } from './exa';
+import { LocalSearchHandler, LocalSearchToolkit, makeGraphWriterHandler, makeGraphWriterToolkit } from './graph';
 // TODO(dmaretskyi): Vite build bug with instruction files with the same filename getting mixed-up
 import PROMPT from './instructions-research.tpl?raw';
-import { AISession } from '../session';
+import { AiSession } from '../session';
 
 /**
  * Exec external service and return the results as a Subgraph.
@@ -34,56 +37,51 @@ export const researchFn = defineFunction({
     }),
   }),
   outputSchema: Schema.Struct({
-    // result: Schema.String,
-    result: Schema.Any,
+    result: Schema.Unknown,
   }),
-  handler: async ({ data: { query, mockSearch }, context }) => {
-    const ai = context.getService(AiService);
-    const credentials = context.getService(CredentialsService);
-    const { db } = context.getService(DatabaseService);
-    // const queues = context.getService(QueueService);
-    const tracing = context.getService(TracingService);
-
-    tracing.write(create(AgentStatusReport, { message: 'Researching...' }));
-
-    const exaCredential = await credentials.getCredential({ service: 'exa.ai' });
-    const searchTool = mockSearch ? createMockExaTool() : createExaTool({ apiKey: exaCredential.apiKey! });
+  handler: Effect.fnUntraced(
+    function* ({ data: { query, mockSearch } }) {
+      yield* TracingService.emitStatus({ message: 'Researching...' });
 
-    const printer = new ConsolePrinter();
-    const session = new AISession({ operationModel: 'configured' });
-    session.message.on((message) => printer.printMessage(message));
-    session.userMessage.on((message) => printer.printMessage(message));
-    session.block.on((block) => printer.printContentBlock(block));
-    session.statusReport.on((status) => {
-      log.info('[agent] status', { status });
-      tracing.write(status);
-    });
-    session.streamEvent.on((event) => log('stream', { event }));
+      const graphWriteToolkit = makeGraphWriterToolkit({ schema: DataTypes });
+      const toolkit = yield* AiToolkit.merge(ExaToolkit, LocalSearchToolkit, graphWriteToolkit).pipe(
+        Effect.provide(
+          Layer.mergeAll(
+            mockSearch ? MockExaHandler : LiveExaHandler,
+            LocalSearchHandler,
+            makeGraphWriterHandler(graphWriteToolkit),
+          ),
+        ),
+      );
 
-    const graphWriteTool = createGraphWriterTool({ db, schema: DataTypes });
-    log.info('graphWriteTool', { schema: graphWriteTool.parameters });
+      const printer = new ConsolePrinter();
+      const session = new AiSession();
+      session.message.on((message) => printer.printMessage(message));
+      session.userMessage.on((message) => printer.printMessage(message));
+      session.block.on((block) => printer.printContentBlock(block));
+      session.streamEvent.on((event) => log('stream', { event }));
 
-    // TODO(dmaretskyi): Consider adding this pattern as the "Graph" output mode for the session.
-    const result = await session.run({
-      client: ai.client,
-      systemPrompt: PROMPT,
-      artifacts: [],
-      tools: [searchTool, graphWriteTool],
-      history: [],
-      prompt: query,
-    });
+      // TODO(dmaretskyi): Consider adding this pattern as the "Graph" output mode for the session.
+      const result = yield* session.run({
+        prompt: query,
+        history: [],
+        systemPrompt: PROMPT,
+        toolkit,
+      });
 
-    return {
-      result: result as any,
-    };
+      return {
+        result,
+      };
 
-    // queues.contextQueue!.append(data);
+      // queues.contextQueue!.append(data);
 
-    // return {
-    //   result: `
-    //   The research results are placed in the following objects:
-    //     ${data.map((object, id) => `[obj_${id}][dxn:echo:@:${object.id}]`).join('\n')}
-    //   `,
-    // };
-  },
+      // return {
+      //   result: `
+      //   The research results are placed in the following objects:
+      //     ${data.map((object, id) => `[obj_${id}][dxn:echo:@:${object.id}]`).join('\n')}
+      //   `,
+      // };
+    },
+    Effect.provide(AiService.model('@anthropic/claude-3-5-sonnet-20241022')),
+  ),
 });

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",
+    },
+  },
+]
+`;