@dxos/assistant-toolkit 0.8.4-main.2c6827d

package/src/functions/research/research.test.ts

@@ -0,0 +1,179 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { inspect } from 'node:util';
+
+import { describe, it } from '@effect/vitest';
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+
+import { AiService, ConsolePrinter, MemoizedAiService } from '@dxos/ai';
+import { TestAiService } from '@dxos/ai/testing';
+import {
+  AiConversation,
+  type ContextBinding,
+  GenerationObserver,
+  makeToolExecutionServiceFromFunctions,
+  makeToolResolverFromFunctions,
+} from '@dxos/assistant';
+import { Blueprint } from '@dxos/blueprints';
+import { Filter, Obj, Query, Ref } from '@dxos/echo';
+import { TestHelpers, acquireReleaseResource } from '@dxos/effect';
+import {
+  CredentialsService,
+  DatabaseService,
+  FunctionInvocationService,
+  QueueService,
+  TracingService,
+} from '@dxos/functions';
+import { FunctionInvocationServiceLayerTest, TestDatabaseLayer } from '@dxos/functions-runtime/testing';
+import { invariant } from '@dxos/invariant';
+import { ObjectId } from '@dxos/keys';
+import { MarkdownBlueprint, MarkdownFunction } from '@dxos/plugin-markdown/toolkit';
+import { Markdown } from '@dxos/plugin-markdown/types';
+import { HasSubject, type Message, Organization } from '@dxos/types';
+
+import { ResearchBlueprint } from '../../blueprints';
+import { testToolkit } from '../../blueprints/testing';
+
+import { default as createDocument } from './document-create';
+import { default as research } from './research';
+import { ResearchGraph, queryResearchGraph } from './research-graph';
+import { ResearchDataTypes } from './types';
+
+ObjectId.dangerouslyDisableRandomness();
+
+const TestLayer = Layer.mergeAll(
+  AiService.model('@anthropic/claude-opus-4-0'),
+  makeToolResolverFromFunctions(
+    [research, createDocument, MarkdownFunction.create, MarkdownFunction.open, MarkdownFunction.update],
+    testToolkit,
+  ),
+  makeToolExecutionServiceFromFunctions(testToolkit, testToolkit.toLayer({}) as any),
+).pipe(
+  Layer.provideMerge(
+    FunctionInvocationServiceLayerTest({
+      functions: [research, createDocument, MarkdownFunction.create, MarkdownFunction.open, MarkdownFunction.update],
+    }),
+  ),
+  Layer.provideMerge(
+    Layer.mergeAll(
+      TestAiService(),
+      TestDatabaseLayer({
+        spaceKey: 'fixed',
+        indexing: { vector: true },
+        types: [...ResearchDataTypes, ResearchGraph, Blueprint.Blueprint, Markdown.Document, HasSubject.HasSubject],
+      }),
+      CredentialsService.configuredLayer([]),
+      TracingService.layerNoop,
+    ),
+  ),
+);
+
+describe('Research', () => {
+  it.effect(
+    'call a function to generate a research report',
+    Effect.fnUntraced(
+      function* (_) {
+        yield* DatabaseService.add(
+          Obj.make(Organization.Organization, {
+            name: 'BlueYard',
+            website: 'https://blueyard.com',
+          }),
+        );
+        yield* DatabaseService.flush({ indexes: true });
+        const result = yield* FunctionInvocationService.invokeFunction(research, {
+          query: 'Founders and portfolio of BlueYard.',
+        });
+
+        console.log(inspect(result, { depth: null, colors: true }));
+        console.log(JSON.stringify(result, null, 2));
+
+        yield* DatabaseService.flush({ indexes: true });
+        const researchGraph = yield* queryResearchGraph();
+        if (researchGraph) {
+          const data = yield* DatabaseService.load(researchGraph.queue).pipe(
+            Effect.flatMap((queue) => Effect.promise(() => queue.queryObjects())),
+          );
+          console.log(inspect(data, { depth: null, colors: true }));
+        }
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+    MemoizedAiService.isGenerationEnabled() ? 240_000 : 30_000,
+  );
+
+  it.scoped(
+    'create and update research report',
+    Effect.fnUntraced(
+      function* (_) {
+        const organization = yield* DatabaseService.add(
+          Obj.make(Organization.Organization, {
+            name: 'BlueYard',
+            website: 'https://blueyard.com',
+          }),
+        );
+
+        const queue = yield* QueueService.createQueue<Message.Message | ContextBinding>();
+        const conversation = yield* acquireReleaseResource(() => new AiConversation(queue));
+
+        yield* DatabaseService.flush({ indexes: true });
+        const researchBlueprint = yield* DatabaseService.add(Obj.clone(ResearchBlueprint));
+        const markdownBlueprint = yield* DatabaseService.add(Obj.clone(MarkdownBlueprint));
+        yield* Effect.promise(() =>
+          conversation.context.bind({
+            blueprints: [Ref.make(researchBlueprint), Ref.make(markdownBlueprint)],
+            objects: [Ref.make(organization)],
+          }),
+        );
+
+        const observer = GenerationObserver.fromPrinter(new ConsolePrinter());
+
+        yield* conversation.createRequest({
+          observer,
+          prompt: `Create a research summary about ${organization.name}.`,
+        });
+        {
+          const { objects: docs } = yield* DatabaseService.runQuery(
+            Query.select(Filter.ids(organization.id)).targetOf(HasSubject.HasSubject).source(),
+          );
+          if (docs.length !== 1) {
+            throw new Error(`Expected 1 research document; got ${docs.length}: ${docs.map((_) => _.name)}`);
+          }
+
+          const doc = docs[0];
+          invariant(Obj.instanceOf(Markdown.Document, doc));
+          console.log({
+            name: doc.name,
+            content: yield* DatabaseService.load(doc.content).pipe(Effect.map((_) => _.content)),
+          });
+        }
+
+        yield* conversation.createRequest({
+          observer,
+          prompt: 'Add a section about their portfolio.',
+        });
+        {
+          const { objects: docs } = yield* DatabaseService.runQuery(
+            Query.select(Filter.ids(organization.id)).targetOf(HasSubject.HasSubject).source(),
+          );
+          if (docs.length !== 1) {
+            throw new Error(`Expected 1 research document; got ${docs.length}: ${docs.map((_) => _.name)}`);
+          }
+
+          const doc = docs[0];
+          invariant(Obj.instanceOf(Markdown.Document, doc));
+          console.log({
+            name: doc.name,
+            content: yield* DatabaseService.load(doc.content).pipe(Effect.map((_) => _.content)),
+          });
+        }
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+    MemoizedAiService.isGenerationEnabled() ? 240_000 : 30_000,
+  );
+});

package/src/functions/research/research.ts

@@ -0,0 +1,190 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Toolkit from '@effect/ai/Toolkit';
+import * as AnthropicTool from '@effect/ai-anthropic/AnthropicTool';
+import * as Array from 'effect/Array';
+import * as Effect from 'effect/Effect';
+import * as Function from 'effect/Function';
+import * as Layer from 'effect/Layer';
+import * as Option from 'effect/Option';
+import * as Schema from 'effect/Schema';
+import * as String from 'effect/String';
+
+import { AiService, ConsolePrinter } from '@dxos/ai';
+import {
+  AiSession,
+  GenerationObserver,
+  createToolkit,
+  makeToolExecutionServiceFromFunctions,
+  makeToolResolverFromFunctions,
+} from '@dxos/assistant';
+import { Template } from '@dxos/blueprints';
+import { type DXN, Obj } from '@dxos/echo';
+import { DatabaseService, TracingService, defineFunction } from '@dxos/functions';
+import { FunctionInvocationServiceLayerTestMocked } from '@dxos/functions-runtime/testing';
+import { type Message, Person } from '@dxos/types';
+import { trim } from '@dxos/util';
+
+import { LocalSearchHandler, LocalSearchToolkit, makeGraphWriterHandler, makeGraphWriterToolkit } from '../../crud';
+import { exaFunction, exaMockFunction } from '../exa';
+
+import { contextQueueLayerFromResearchGraph } from './research-graph';
+import PROMPT from './research-instructions.tpl?raw';
+import { ResearchDataTypes } from './types';
+
+/**
+ * Exec external service and return the results as a Subgraph.
+ */
+export default defineFunction({
+  key: 'dxos.org/function/research',
+  name: 'Research',
+  description: trim`
+    Search the web to research information about the given subject.
+    Inserts structured data into the research graph. 
+    Creates a research summary and returns the objects created.
+  `,
+  inputSchema: Schema.Struct({
+    query: Schema.String.annotations({
+      description: trim`
+        The search query.
+        If doing research on an object then load it first and pass it as JSON.
+      `,
+    }),
+
+    instructions: Schema.optional(Schema.String).annotations({
+      description: trim`
+        The instructions for the research agent. 
+      `,
+    }),
+
+    // TOOD(burdon): Move to context.
+    mockSearch: Schema.optional(Schema.Boolean).annotations({
+      description: 'Whether to use the mock search tool.',
+      default: false,
+    }),
+
+    entityExtraction: Schema.optional(Schema.Boolean).annotations({
+      description: trim`
+        Whether to extract structured entities from the research. 
+        Experimental feature only enable if user explicitly requests it.
+      `,
+      default: false,
+    }),
+  }),
+  outputSchema: Schema.Struct({
+    document: Schema.optional(Schema.String).annotations({
+      description: 'The generated research document.',
+    }),
+    objects: Schema.Array(Schema.Unknown).annotations({
+      description: 'Structured objects created during the research process.',
+    }),
+  }),
+  handler: Effect.fnUntraced(
+    function* ({ data: { query, instructions, mockSearch = false, entityExtraction = false } }) {
+      if (mockSearch) {
+        const mockPerson = yield* DatabaseService.add(
+          Obj.make(Person.Person, {
+            preferredName: 'John Doe',
+            emails: [{ value: 'john.doe@example.com' }],
+            phoneNumbers: [{ value: '123-456-7890' }],
+          }),
+        );
+
+        return {
+          document: trim`
+            The research ran in test-mode and was mocked. Proceed as usual.
+            We reference John Doe to test reference: ${Obj.getDXN(mockPerson)}
+          `,
+          objects: [Obj.toJSON(mockPerson)],
+        };
+      }
+
+      yield* DatabaseService.flush({ indexes: true });
+      yield* TracingService.emitStatus({ message: 'Starting research...' });
+
+      const NativeWebSearch = Toolkit.make(AnthropicTool.WebSearch_20250305({}));
+
+      let toolkit: Toolkit.Any = NativeWebSearch;
+      let handlers: Layer.Layer<any, any> = Layer.empty as any;
+
+      const objectDXNs: DXN[] = [];
+      if (entityExtraction) {
+        const GraphWriterToolkit = makeGraphWriterToolkit({ schema: ResearchDataTypes });
+        const GraphWriterHandler = makeGraphWriterHandler(GraphWriterToolkit, {
+          onAppend: (dxns) => objectDXNs.push(...dxns),
+        });
+
+        toolkit = Toolkit.merge(toolkit, LocalSearchToolkit, GraphWriterToolkit);
+        handlers = Layer.mergeAll(handlers, LocalSearchHandler, GraphWriterHandler).pipe(
+          Layer.provide(contextQueueLayerFromResearchGraph),
+        ) as any;
+      }
+
+      const finishedToolkit = yield* createToolkit({
+        toolkit: toolkit as any,
+      }).pipe(Effect.provide(handlers));
+
+      const session = new AiSession();
+      const result = yield* session.run({
+        prompt: query,
+        system: join(
+          Template.process(PROMPT, { entityExtraction }),
+          instructions && `<instructions>${instructions}</instructions>`,
+        ),
+        toolkit: finishedToolkit,
+        observer: GenerationObserver.fromPrinter(new ConsolePrinter({ tag: 'research' })),
+      });
+
+      const objects = yield* Effect.forEach(objectDXNs, (dxn) => DatabaseService.resolve(dxn)).pipe(
+        Effect.map(Array.map((obj) => Obj.toJSON(obj))),
+      );
+
+      return {
+        document: extractLastTextBlock(result),
+        objects,
+      };
+    },
+    Effect.provide(
+      Layer.mergeAll(
+        AiService.model('@anthropic/claude-sonnet-4-0'),
+        // TODO(dmaretskyi): Extract.
+        makeToolResolverFromFunctions([exaFunction, exaMockFunction], Toolkit.make()),
+        makeToolExecutionServiceFromFunctions(Toolkit.make() as any, Layer.empty as any),
+      ).pipe(
+        Layer.provide(
+          // TODO(dmaretskyi): This should be provided by environment.
+          Layer.mergeAll(FunctionInvocationServiceLayerTestMocked({ functions: [exaFunction, exaMockFunction] })),
+        ),
+      ),
+    ),
+  ),
+});
+
+// TODO(burdon): Factor out.
+const join = (...strings: (string | undefined)[]) => strings.filter(Boolean).join('\n\n');
+
+/**
+ * Extracts the last text block from the result.
+ * Skips citations.
+ */
+// TODO(burdon): Factor out.
+const extractLastTextBlock = (result: Message.Message[]) => {
+  return Function.pipe(
+    result,
+    Array.last,
+    Option.map(
+      Function.flow(
+        (_: Message.Message) => _.blocks,
+        Array.reverse,
+        Array.dropWhile((_: any) => _._tag === 'summary'),
+        Array.takeWhile((_: any) => _._tag === 'text'),
+        Array.reverse,
+        Array.map((_: any) => _.text),
+        Array.reduce('', String.concat),
+      ),
+    ),
+    Option.getOrElse(() => ''),
+  );
+};

package/src/functions/research/types.ts

@@ -0,0 +1,26 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { type Type } from '@dxos/echo';
+import { Text } from '@dxos/schema';
+import { Event, HasConnection, HasRelationship, LegacyOrganization, LegacyPerson, Project, Task } from '@dxos/types';
+
+/**
+ * Data types for research.
+ */
+export const ResearchDataTypes: (Type.Obj.Any | Type.Relation.Any)[] = [
+  // Objects
+  Event.Event,
+  LegacyOrganization,
+  Task.Task,
+  Text.Text,
+
+  // Relations
+  // TODO(wittjosiah): Until views (e.g., Table) support relations this needs to be expressed via organization ref.
+  // Employer.Employer,
+  LegacyPerson,
+  Project.Project,
+  HasRelationship.HasRelationship,
+  HasConnection.HasConnection,
+];

package/src/functions/tasks/index.ts

@@ -0,0 +1,11 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { default as read$ } from './read';
+import { default as update$ } from './update';
+
+export namespace Tasks {
+  export const read = read$;
+  export const update = update$;
+}

package/src/functions/tasks/read.ts

@@ -0,0 +1,34 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Effect from 'effect/Effect';
+import * as Schema from 'effect/Schema';
+
+import { ArtifactId } from '@dxos/assistant';
+import { DatabaseService, defineFunction } from '@dxos/functions';
+import { Markdown } from '@dxos/plugin-markdown/types';
+
+export default defineFunction({
+  key: 'dxos.org/function/markdown/read-tasks',
+  name: 'Read',
+  description: 'Read markdown tasks.',
+  inputSchema: Schema.Struct({
+    id: ArtifactId.annotations({
+      description: 'The ID of the document to read.',
+    }),
+  }),
+  outputSchema: Schema.Struct({
+    content: Schema.String,
+  }),
+  handler: Effect.fn(function* ({ data: { id } }) {
+    const doc = yield* DatabaseService.resolve(ArtifactId.toDXN(id), Markdown.Document);
+
+    // Return content with line numbers prefixed.
+    const { content } = yield* DatabaseService.load(doc.content);
+    const lines = content.split('\n');
+    const len = String(lines.length).length;
+    const numbered = lines.map((line, i) => `${String(i + 1).padStart(len, ' ')}. ${line}`).join('\n');
+    return { content: numbered };
+  }),
+});

package/src/functions/tasks/task-list.test.ts

@@ -0,0 +1,99 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { describe, expect, it } from 'vitest';
+
+import { MarkdownTasks, type TaskOperation } from './task-list';
+
+describe('MarkdownTasks', () => {
+  it('should initialize with content', () => {
+    const content = '- [ ] Task 1\n- [x] Task 2';
+    const manager = new MarkdownTasks(content);
+    expect(manager.getRawContent()).toBe(content);
+  });
+
+  it('should provide numbered content', () => {
+    const manager = new MarkdownTasks('- [ ] Task 1\n  - [ ] Subtask');
+    const numbered = manager.getNumberedContent();
+    expect(numbered).toBe('  1→- [ ] Task 1\n  2→  - [ ] Subtask');
+  });
+
+  it('should insert tasks with indentation', () => {
+    const manager = new MarkdownTasks('- [ ] Main task');
+    manager.insertTask(2, 'Subtask', false, 2);
+    expect(manager.getRawContent()).toBe('- [ ] Main task\n  - [ ] Subtask');
+  });
+
+  it('should insert completed tasks', () => {
+    const manager = new MarkdownTasks('- [ ] Task 1');
+    manager.insertTask(2, 'Completed task', true, 0);
+    expect(manager.getRawContent()).toBe('- [ ] Task 1\n- [x] Completed task');
+  });
+
+  it('should delete tasks', () => {
+    const manager = new MarkdownTasks('- [ ] Task 1\n- [ ] Task 2');
+    const success = manager.deleteTask(2);
+    expect(success).toBe(true);
+    expect(manager.getRawContent()).toBe('- [ ] Task 1');
+  });
+
+  it('should update task text', () => {
+    const manager = new MarkdownTasks('- [ ] Old task');
+    const success = manager.updateTaskText(1, 'New task description');
+    expect(success).toBe(true);
+    expect(manager.getRawContent()).toBe('- [ ] New task description');
+  });
+
+  it('should toggle task completion', () => {
+    const manager = new MarkdownTasks('- [ ] Task 1\n- [x] Task 2');
+
+    // Complete first task
+    manager.toggleTaskCompletion(1, true);
+    expect(manager.getRawContent()).toBe('- [x] Task 1\n- [x] Task 2');
+
+    // Uncomplete second task
+    manager.toggleTaskCompletion(2, false);
+    expect(manager.getRawContent()).toBe('- [x] Task 1\n- [ ] Task 2');
+  });
+
+  it('should change task indentation', () => {
+    const manager = new MarkdownTasks('- [ ] Task 1');
+    const success = manager.setTaskIndent(1, 4);
+    expect(success).toBe(true);
+    expect(manager.getRawContent()).toBe('    - [ ] Task 1');
+  });
+
+  it('should apply multiple operations', () => {
+    const manager = new MarkdownTasks('- [ ] Task 1');
+
+    const operations: TaskOperation[] = [
+      { type: 'insertTask', lineNumber: 2, text: 'New task', indent: 2 },
+      { type: 'toggleTaskCompletion', lineNumber: 1, completed: true },
+      { type: 'insertTask', lineNumber: 3, text: 'Another task', indent: 0 },
+    ];
+
+    manager.applyOperations(operations);
+    const result = manager.getRawContent();
+    expect(result).toBe('- [x] Task 1\n  - [ ] New task\n- [ ] Another task');
+  });
+
+  it('should handle hierarchical task structure', () => {
+    const manager = new MarkdownTasks();
+
+    // Build a hierarchical structure
+    manager.insertTask(1, 'Main task', false, 0);
+    manager.insertTask(2, 'Subtask 1', false, 2);
+    manager.insertTask(3, 'Sub-subtask', false, 4);
+    manager.insertTask(4, 'Subtask 2', false, 2);
+
+    const expected = '- [ ] Main task\n  - [ ] Subtask 1\n    - [ ] Sub-subtask\n  - [ ] Subtask 2';
+    expect(manager.getRawContent()).toBe(expected);
+  });
+
+  it('should handle line insertions at the end', () => {
+    const manager = new MarkdownTasks('- [ ] Task 1');
+    manager.insertTask(999, 'End task', false, 0);
+    expect(manager.getRawContent()).toBe('- [ ] Task 1\n- [ ] End task');
+  });
+});