@dxos/assistant-toolkit 0.8.4-main.ae835ea → 0.8.4-main.e8ec1fe

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

@@ -9,7 +9,7 @@ import * as Effect from 'effect/Effect';
 import * as Layer from 'effect/Layer';
 
 import { AiService, ConsolePrinter, MemoizedAiService } from '@dxos/ai';
-import { AiServiceTestingPreset, EXA_API_KEY } from '@dxos/ai/testing';
+import { TestAiService } from '@dxos/ai/testing';
 import {
   AiConversation,
   type ContextBinding,
@@ -18,25 +18,26 @@ import {
   makeToolResolverFromFunctions,
 } from '@dxos/assistant';
 import { Blueprint } from '@dxos/blueprints';
-import { Obj, Ref, Type } from '@dxos/echo';
+import { Filter, Obj, Query, Ref } from '@dxos/echo';
 import { TestHelpers, acquireReleaseResource } from '@dxos/effect';
 import {
-  ComputeEventLogger,
   CredentialsService,
   DatabaseService,
   FunctionInvocationService,
   QueueService,
   TracingService,
 } from '@dxos/functions';
-import { TestDatabaseLayer } from '@dxos/functions/testing';
+import { FunctionInvocationServiceLayerTest, TestDatabaseLayer } from '@dxos/functions-runtime/testing';
+import { invariant } from '@dxos/invariant';
 import { ObjectId } from '@dxos/keys';
-import { DataType } from '@dxos/schema';
+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 createDocument from './create-document';
-import { createExtractionSchema, getSanitizedSchemaName } from './graph';
+import { default as createDocument } from './document-create';
 import { default as research } from './research';
 import { ResearchGraph, queryResearchGraph } from './research-graph';
 import { ResearchDataTypes } from './types';
@@ -45,41 +46,45 @@ ObjectId.dangerouslyDisableRandomness();
 
 const TestLayer = Layer.mergeAll(
   AiService.model('@anthropic/claude-opus-4-0'),
-  makeToolResolverFromFunctions([research, createDocument], testToolkit),
+  makeToolResolverFromFunctions(
+    [research, createDocument, MarkdownFunction.create, MarkdownFunction.open, MarkdownFunction.update],
+    testToolkit,
+  ),
   makeToolExecutionServiceFromFunctions(testToolkit, testToolkit.toLayer({}) as any),
-  ComputeEventLogger.layerFromTracing,
 ).pipe(
-  Layer.provideMerge(FunctionInvocationService.layerTest({ functions: [research, createDocument] })),
+  Layer.provideMerge(
+    FunctionInvocationServiceLayerTest({
+      functions: [research, createDocument, MarkdownFunction.create, MarkdownFunction.open, MarkdownFunction.update],
+    }),
+  ),
   Layer.provideMerge(
     Layer.mergeAll(
-      MemoizedAiService.layerTest().pipe(Layer.provide(AiServiceTestingPreset('direct'))),
+      TestAiService(),
       TestDatabaseLayer({
+        spaceKey: 'fixed',
         indexing: { vector: true },
-        types: [...ResearchDataTypes, ResearchGraph, Blueprint.Blueprint],
+        types: [...ResearchDataTypes, ResearchGraph, Blueprint.Blueprint, Markdown.Document, HasSubject.HasSubject],
       }),
-      CredentialsService.configuredLayer([{ service: 'exa.ai', apiKey: EXA_API_KEY }]),
+      CredentialsService.configuredLayer([]),
       TracingService.layerNoop,
     ),
   ),
 );
 
-// TODO(dmaretskyi): Out-of-memory.
-describe.skip('Research', () => {
+describe('Research', () => {
   it.effect(
     'call a function to generate a research report',
     Effect.fnUntraced(
       function* (_) {
         yield* DatabaseService.add(
-          Obj.make(DataType.Organization, {
-            name: 'Airbnb',
-            website: 'https://www.airbnb.com/',
+          Obj.make(Organization.Organization, {
+            name: 'BlueYard',
+            website: 'https://blueyard.com',
           }),
         );
         yield* DatabaseService.flush({ indexes: true });
-
         const result = yield* FunctionInvocationService.invokeFunction(research, {
-          query: 'Founders and investors of airbnb.',
-          mockSearch: false,
+          query: 'Founders and portfolio of BlueYard.',
         });
 
         console.log(inspect(result, { depth: null, colors: true }));
@@ -87,154 +92,88 @@ describe.skip('Research', () => {
 
         yield* DatabaseService.flush({ indexes: true });
         const researchGraph = yield* queryResearchGraph();
-        const data = yield* DatabaseService.load(researchGraph!.queue).pipe(
-          Effect.flatMap((queue) => Effect.promise(() => queue.queryObjects())),
-        );
-        console.log(inspect(data, { depth: null, colors: true }));
+        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 : undefined,
+    MemoizedAiService.isGenerationEnabled() ? 240_000 : 30_000,
   );
 
   it.scoped(
-    'research blueprint',
+    'create and update research report',
     Effect.fnUntraced(
       function* (_) {
-        const queue = yield* QueueService.createQueue<DataType.Message | ContextBinding>();
+        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));
 
-        const org = Obj.make(DataType.Organization, { name: 'Airbnb', website: 'https://www.airbnb.com/' });
-        yield* DatabaseService.add(org);
         yield* DatabaseService.flush({ indexes: true });
-
-        const blueprint = yield* DatabaseService.add(Obj.clone(ResearchBlueprint));
-        yield* Effect.promise(() => conversation.context.bind({ blueprints: [Ref.make(blueprint)] }));
+        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: `Research airbnb founders.`,
+          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)),
+          });
+        }
 
-        const researchGraph = yield* queryResearchGraph();
-        const data = yield* DatabaseService.load(researchGraph!.queue).pipe(
-          Effect.flatMap((queue) => Effect.promise(() => queue.queryObjects())),
-        );
-        console.log(inspect(data, { depth: null, colors: true }));
+        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 : undefined,
+    MemoizedAiService.isGenerationEnabled() ? 240_000 : 30_000,
   );
 });
-
-describe('misc', () => {
-  it('createExtractionSchema', () => {
-    const _schema = createExtractionSchema(ResearchDataTypes);
-    // log.info('schema', { schema });
-  });
-
-  it('getSanitizedSchemaName', () => {
-    const _names = ResearchDataTypes.map(getSanitizedSchemaName);
-    // log.info('names', { names }) ;
-  });
-
-  it('getTypeAnnotation', () => {
-    for (const schema of ResearchDataTypes) {
-      const _dxn = Type.getDXN(schema);
-      // log.info('dxn', { schema, dxn });
-    }
-  });
-
-  it.skip('sanitizeObjects', () => {
-    const _TEST_DATA = {
-      objects_dxos_org_type_Project: [
-        {
-          name: 'Reor',
-          description:
-            'A private & local AI personal knowledge management app focused on high entropy thinkers. Features include Q&A chat with full context of notes, automatic connection of ideas, semantic search, and local-first architecture with LLMs and vector databases.',
-          image: 'https://reorhomepage-2-cwy0zagzg-reor-team.vercel.app/opengraph-image.jpg',
-          id: 'reor-project-01',
-        },
-        {
-          name: 'Personal AI',
-          description:
-            'A platform for creating personal AI models with unlimited memory capabilities. Features include memory stacking, data uploads, and personalized language models that can be trained on individual knowledge bases.',
-          image: 'https://cdn.prod.website-files.com/5ff65c460ce39f5ec5681c6a/663d12aab1b425e1ad40d3a6_Memory-min.jpg',
-          id: 'personal-ai-project-01',
-        },
-        {
-          name: 'IKI AI',
-          description:
-            'An AI-native workspace for research, strategy, and creative work. Offers features like AI summarization, content analysis, and team knowledge sharing capabilities.',
-          image: 'https://framerusercontent.com/assets/cI6Uo7x4q0W3uxzOt2preXjv6aE.jpg',
-          id: 'iki-ai-project-01',
-        },
-        {
-          name: 'Forethink',
-          description:
-            'A privacy-first, on-device AI system that proactively links current activities to relevant knowledge. Features real-time learning across teams and local data processing.',
-          image: 'https://framerusercontent.com/assets/XDhN6tICnlElB134JWcWXOuXc.jpeg',
-          id: 'forethink-project-01',
-        },
-        {
-          name: 'Amurex',
-          description:
-            'An invisible AI companion that handles tasks, remembers important information, and supports workflow. Includes features for meeting automation, email categorization, and cross-platform search.',
-          image: 'https://amurex.ai/ogimages/og_amurex.jpg',
-          id: 'amurex-project-01',
-        },
-      ],
-      objects_dxos_org_type_Organization: [
-        {
-          name: 'Reor',
-          description:
-            'A technology company focused on developing local-first AI solutions for personal knowledge management.',
-          status: 'active',
-          website: 'https://www.reorproject.org',
-          id: 'reor-org-01',
-        },
-        {
-          name: 'Personal.ai',
-          description: 'A company specializing in personal AI development and memory management solutions.',
-          status: 'active',
-          website: 'https://www.personal.ai',
-          id: 'personal-ai-org-01',
-        },
-        {
-          name: 'IKI AI',
-          description: 'A startup backed by 500 Global, developing AI-native workspaces for knowledge workers.',
-          status: 'active',
-          website: 'https://iki.ai',
-          id: 'iki-ai-org-01',
-        },
-        {
-          name: 'Forethink.ai',
-          description:
-            'A company focused on privacy-first AI solutions for knowledge management and team collaboration.',
-          status: 'active',
-          website: 'https://www.forethink.ai',
-          id: 'forethink-org-01',
-        },
-      ],
-      objects_dxos_org_type_Text: [
-        {
-          content:
-            'Current trends in AI-powered Personal Knowledge Management (PKM) systems include: 1) Shift from document-based to graph-based systems for better knowledge connection, 2) Integration of large language models for enhanced content understanding and retrieval, 3) Focus on privacy-first and local-first architectures, 4) Emphasis on automatic knowledge connection and discovery, 5) Implementation of semantic search capabilities.',
-          id: 'pkm-trends-text-01',
-        },
-        {
-          content:
-            'Key features emerging across modern PKM tools include: automated content categorization, semantic analysis and linking, personalized AI models that learn from user interaction, integration with existing workflows, and robust privacy protection. These tools are increasingly focusing on proactive knowledge surfacing rather than passive storage.',
-          id: 'pkm-features-text-01',
-        },
-      ],
-    };
-
-    // const data = await sanitizeObjects(TYPES, TEST_DATA, db);
-    // log.info('data', { data });
-  });
-});

package/src/functions/research/research.ts

@@ -6,8 +6,11 @@ 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 {
@@ -17,17 +20,18 @@ import {
   makeToolExecutionServiceFromFunctions,
   makeToolResolverFromFunctions,
 } from '@dxos/assistant';
+import { Template } from '@dxos/blueprints';
 import { type DXN, Obj } from '@dxos/echo';
-import { DatabaseService, FunctionInvocationService, TracingService, defineFunction } from '@dxos/functions';
-import { DataType } from '@dxos/schema';
+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 { 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 { contextQueueLayerFromResearchGraph } from './research-graph';
+import PROMPT from './research-instructions.tpl?raw';
 import { ResearchDataTypes } from './types';
 
 /**
@@ -37,22 +41,21 @@ export default defineFunction({
   key: 'dxos.org/function/research',
   name: 'Research',
   description: trim`
-    Research the web for information. 
+    Search the web to research information about the given subject.
     Inserts structured data into the research graph. 
-    Will return research summary and the objects created.
+    Creates a research summary and returns the objects created.
   `,
   inputSchema: Schema.Struct({
     query: Schema.String.annotations({
       description: trim`
-        The query to search for.
-        If doing research on an object, load it first and pass it as a JSON string.
+        The search query.
+        If doing research on an object then load it first and pass it as JSON.
       `,
     }),
 
-    researchInstructions: Schema.optional(Schema.String).annotations({
+    instructions: Schema.optional(Schema.String).annotations({
       description: trim`
         The instructions for the research agent. 
-        E.g., preference on fast responses or in-depth analysis, number of web searcher or the objects created.
       `,
     }),
 
@@ -61,20 +64,28 @@ export default defineFunction({
       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({
-    note: Schema.optional(Schema.String).annotations({
-      description: 'A note from the research agent.',
+    document: Schema.optional(Schema.String).annotations({
+      description: 'The generated research document.',
     }),
     objects: Schema.Array(Schema.Unknown).annotations({
-      description: 'The structured objects created as a result of the research.',
+      description: 'Structured objects created during the research process.',
     }),
   }),
   handler: Effect.fnUntraced(
-    function* ({ data: { query, mockSearch, researchInstructions } }) {
+    function* ({ data: { query, instructions, mockSearch = false, entityExtraction = false } }) {
       if (mockSearch) {
         const mockPerson = yield* DatabaseService.add(
-          Obj.make(DataType.Person, {
+          Obj.make(Person.Person, {
             preferredName: 'John Doe',
             emails: [{ value: 'john.doe@example.com' }],
             phoneNumbers: [{ value: '123-456-7890' }],
@@ -82,9 +93,8 @@ export default defineFunction({
         );
 
         return {
-          note: trim`
-            The research run in test-mode and was mocked. 
-            Proceed as usual.
+          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)],
@@ -92,49 +102,47 @@ export default defineFunction({
       }
 
       yield* DatabaseService.flush({ indexes: true });
-      yield* TracingService.emitStatus({ message: 'Researching...' });
+      yield* TracingService.emitStatus({ message: 'Starting research...' });
 
-      const objectDXNs: DXN[] = [];
-      const GraphWriterToolkit = makeGraphWriterToolkit({ schema: ResearchDataTypes });
-      const GraphWriterHandler = makeGraphWriterHandler(GraphWriterToolkit, {
-        onAppend: (dxns) => objectDXNs.push(...dxns),
-      });
       const NativeWebSearch = Toolkit.make(AnthropicTool.WebSearch_20250305({}));
 
-      const toolkit = yield* createToolkit({
-        toolkit: Toolkit.merge(LocalSearchToolkit, GraphWriterToolkit, NativeWebSearch),
-        // toolIds: [mockSearch ? ToolId.make(exaMockFunction.key) : ToolId.make(exaFunction.key)],
-      }).pipe(
-        Effect.provide(
-          Layer.mergeAll(
-            //
-            GraphWriterHandler,
-            LocalSearchHandler,
-          ).pipe(Layer.provide(contextQueueLayerFromResearchGraph)),
-        ),
-      );
+      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:
-          PROMPT +
-          (researchInstructions
-            ? '\n\n' + `<research_instructions>${researchInstructions}</research_instructions>`
-            : ''),
-        toolkit,
+        system: join(
+          Template.process(PROMPT, { entityExtraction }),
+          instructions && `<instructions>${instructions}</instructions>`,
+        ),
+        toolkit: finishedToolkit,
         observer: GenerationObserver.fromPrinter(new ConsolePrinter({ tag: 'research' })),
       });
-      const note = result
-        .at(-1)
-        ?.blocks.filter((block) => block._tag === 'text')
-        .at(-1)?.text;
+
       const objects = yield* Effect.forEach(objectDXNs, (dxn) => DatabaseService.resolve(dxn)).pipe(
         Effect.map(Array.map((obj) => Obj.toJSON(obj))),
       );
 
       return {
-        note,
+        document: extractLastTextBlock(result),
         objects,
       };
     },
@@ -147,9 +155,36 @@ export default defineFunction({
       ).pipe(
         Layer.provide(
           // TODO(dmaretskyi): This should be provided by environment.
-          Layer.mergeAll(FunctionInvocationService.layerTestMocked({ functions: [exaFunction, exaMockFunction] })),
+          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

@@ -2,23 +2,25 @@
 // Copyright 2025 DXOS.org
 //
 
-import { DataType } from '@dxos/schema';
+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 = [
+export const ResearchDataTypes: (Type.Obj.Any | Type.Relation.Any)[] = [
   // Objects
-  DataType.Event,
-  DataType.LegacyOrganization,
-  DataType.LegacyPerson,
-  DataType.Project,
-  DataType.Task,
-  DataType.Text,
+  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.
-  // DataType.Employer,
-  DataType.HasRelationship,
-  DataType.HasConnection,
+  // 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/index.ts

@@ -6,3 +6,5 @@ export * from './blueprints';
 export * from './functions';
 export * from './plugins';
 export * from './sync';
+export * from './crud';
+export * from './toolkits';