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

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

@@ -0,0 +1,149 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { inspect } from 'node:util';
+
+import { describe, it } from '@effect/vitest';
+import * as Effect from 'effect/Effect';
+
+import { ConsolePrinter } from '@dxos/ai';
+import { MemoizedAiService } from '@dxos/ai/testing';
+import { AiConversation, type ContextBinding, GenerationObserver } from '@dxos/assistant';
+import { AssistantTestLayer } from '@dxos/assistant/testing';
+import { Blueprint } from '@dxos/blueprints';
+import { Database, Filter, Obj, Query, Ref } from '@dxos/echo';
+import { acquireReleaseResource } from '@dxos/effect';
+import { TestHelpers } from '@dxos/effect/testing';
+import { FunctionInvocationService, QueueService } from '@dxos/functions';
+import { invariant } from '@dxos/invariant';
+import { ObjectId } from '@dxos/keys';
+import { MarkdownBlueprint } from '@dxos/plugin-markdown/blueprints';
+import { Markdown } from '@dxos/plugin-markdown/types';
+import { HasSubject, type Message, Organization } from '@dxos/types';
+
+import ResearchBlueprint from '../blueprint';
+import { ResearchDataTypes, ResearchGraph } from '../types';
+
+import { default as createDocument } from './document-create';
+import { default as research } from './research';
+
+ObjectId.dangerouslyDisableRandomness();
+
+const TestLayer = AssistantTestLayer({
+  functions: [research, createDocument, ...MarkdownBlueprint.functions],
+  types: [
+    ...ResearchDataTypes,
+    ResearchGraph.ResearchGraph,
+    Blueprint.Blueprint,
+    Markdown.Document,
+    HasSubject.HasSubject,
+  ],
+});
+
+describe('Research', () => {
+  it.effect(
+    'call a function to generate a research report',
+    Effect.fnUntraced(
+      function* (_) {
+        yield* Database.add(
+          Obj.make(Organization.Organization, {
+            name: 'BlueYard',
+            website: 'https://blueyard.com',
+          }),
+        );
+        yield* Database.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* Database.flush({ indexes: true });
+        const researchGraph = yield* ResearchGraph.query();
+        if (researchGraph) {
+          const data = yield* Database.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* Database.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* Database.flush({ indexes: true });
+        const researchBlueprint = yield* Database.add(Obj.clone(ResearchBlueprint.make()));
+        const markdownBlueprint = yield* Database.add(Obj.clone(MarkdownBlueprint.make()));
+        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 docs = yield* Database.runQuery(
+            Query.select(Filter.id(organization.id)).targetOf(HasSubject.HasSubject).source(),
+          );
+          if (docs.length !== 1) {
+            throw new Error(`Expected 1 research document; got ${docs.length}: ${docs.map((_) => (_ as any).name)}`);
+          }
+
+          const doc = docs[0];
+          invariant(Obj.instanceOf(Markdown.Document, doc));
+          console.log({
+            name: doc.name,
+            content: yield* Database.load(doc.content).pipe(Effect.map((_) => _.content)),
+          });
+        }
+
+        yield* conversation.createRequest({
+          observer,
+          prompt: 'Add a section about their portfolio.',
+        });
+        {
+          const docs = yield* Database.runQuery(
+            Query.select(Filter.id(organization.id)).targetOf(HasSubject.HasSubject).source(),
+          );
+          if (docs.length !== 1) {
+            throw new Error(`Expected 1 research document; got ${docs.length}: ${docs.map((_) => (_ as any).name)}`);
+          }
+
+          const doc = docs[0];
+          invariant(Obj.instanceOf(Markdown.Document, doc));
+          console.log({
+            name: doc.name,
+            content: yield* Database.load(doc.content).pipe(Effect.map((_) => _.content)),
+          });
+        }
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+    MemoizedAiService.isGenerationEnabled() ? 240_000 : 30_000,
+  );
+});

package/src/blueprints/research/functions/research.ts

@@ -0,0 +1,175 @@
+//
+// 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 { GenericToolkit } from '@dxos/ai';
+import { AiSession, GenerationObserver, ToolExecutionServices, createToolkit } from '@dxos/assistant';
+import { Template } from '@dxos/blueprints';
+import { type DXN, Entity, Obj } from '@dxos/echo';
+import { Database } from '@dxos/echo';
+import { TracingService, defineFunction } from '@dxos/functions';
+import { type Message, Person } from '@dxos/types';
+import { trim } from '@dxos/util';
+
+import { LocalSearchHandler, LocalSearchToolkit, makeGraphWriterHandler, makeGraphWriterToolkit } from '../../../crud';
+import { ResearchGraph } from '../types';
+import { ResearchDataTypes } from '../types';
+
+import PROMPT from './research-instructions.tpl?raw';
+
+/**
+ * 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* Database.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* Database.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(ResearchGraph.contextQueueLayer),
+        ) as any;
+      }
+
+      const finishedToolkit = yield* createToolkit({ toolkit }).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) => Database.resolve(dxn)).pipe(
+        Effect.map(Array.map((obj) => Entity.toJSON(obj))),
+      );
+
+      return {
+        document: extractLastTextBlock(result),
+        objects,
+      };
+    },
+    Effect.provide(
+      AiService.model('@anthropic/claude-sonnet-4-0').pipe(
+        Layer.merge(ToolExecutionServices),
+        Layer.provide(GenericToolkit.providerEmpty),
+      ),
+    ),
+  ),
+});
+
+// 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 === 'stats'),
+        Array.takeWhile((_: any) => _._tag === 'text'),
+        Array.reverse,
+        Array.map((_: any) => _.text),
+        Array.reduce('', String.concat),
+      ),
+    ),
+    Option.getOrElse(() => ''),
+  );
+};

package/src/blueprints/research/index.ts

@@ -2,6 +2,5 @@
 // Copyright 2025 DXOS.org
 //
 
-import blueprint from './research-blueprint';
-
-export default blueprint;
+export { default as ResearchBlueprint } from './blueprint';
+export * from './types';

package/src/blueprints/research/types/ResearchGraph.ts

@@ -0,0 +1,49 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+import * as Schema from 'effect/Schema';
+
+import { Database, Obj, Query, Ref, Type } from '@dxos/echo';
+import { SystemTypeAnnotation } from '@dxos/echo/internal';
+import { Queue } from '@dxos/echo-db';
+import { ContextQueueService, QueueService } from '@dxos/functions';
+
+/**
+ * Container for a set of ephemeral research results.
+ */
+export const ResearchGraph = Schema.Struct({
+  queue: Ref.Ref(Queue),
+}).pipe(
+  Type.object({
+    typename: 'dxos.org/type/ResearchGraph',
+    version: '0.1.0',
+  }),
+  SystemTypeAnnotation.set(true),
+);
+
+export interface ResearchGraph extends Schema.Schema.Type<typeof ResearchGraph> {}
+
+export const create: () => Effect.Effect<ResearchGraph, never, Database.Service | QueueService> = Effect.fn(
+  'createResearchGraph',
+)(function* () {
+  const queue = yield* QueueService.createQueue();
+  return yield* Database.add(Obj.make(ResearchGraph, { queue: Ref.fromDXN(queue.dxn) }));
+});
+
+export const query: () => Effect.Effect<ResearchGraph | undefined, never, Database.Service> = Effect.fn(
+  'queryResearchGraph',
+)(function* () {
+  const objects = yield* Database.runQuery(Query.type(ResearchGraph));
+  return objects.at(0);
+});
+
+export const contextQueueLayer = Layer.unwrapEffect(
+  Effect.gen(function* () {
+    const researchGraph = (yield* query()) ?? (yield* create());
+    const researchQueue = yield* Database.load(researchGraph.queue);
+    return ContextQueueService.layer(researchQueue);
+  }),
+);

package/src/blueprints/research/types/index.ts

@@ -0,0 +1,32 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import { type Type } from '@dxos/echo';
+import { Text } from '@dxos/schema';
+import { Event, HasConnection, HasRelationship, LegacyOrganization, LegacyPerson, Pipeline, Task } from '@dxos/types';
+
+export * as ResearchGraph from './ResearchGraph';
+
+/**
+ * Data types for research.
+ */
+// TODO(burdon): This should not be hardcoded.
+export const ResearchDataTypes: Type.AnyEntity[] = [
+  // Objects
+  Event.Event,
+  Task.Task,
+  Text.Text,
+
+  // TODO(wittjosiah): Until views (e.g., Table) support relations this needs to be expressed via organization ref.
+  // Employer.Employer,
+  // Organization.Organization,
+  LegacyPerson,
+  LegacyOrganization,
+
+  Pipeline.Pipeline,
+
+  // Relations
+  HasRelationship.HasRelationship,
+  HasConnection.HasConnection,
+];

package/src/blueprints/testing.ts

@@ -2,17 +2,36 @@
 // Copyright 2025 DXOS.org
 //
 
-import * as Toolkit from '@effect/ai/Toolkit';
+import * as Arr from 'effect/Array';
 import * as Effect from 'effect/Effect';
+import { pipe } from 'effect/Function';
 
 import { ConsolePrinter } from '@dxos/ai';
-import { type AiConversation, type AiConversationRunParams, GenerationObserver } from '@dxos/assistant';
+import {
+  AiContextService,
+  type AiConversation,
+  type AiConversationRunProps,
+  GenerationObserver,
+} from '@dxos/assistant';
+import type { Blueprint } from '@dxos/blueprints';
+import { Database, Ref } from '@dxos/echo';
+import { type FunctionDefinition } from '@dxos/functions';
 import { log } from '@dxos/log';
 
-export type TestStep = Pick<AiConversationRunParams, 'prompt' | 'system'> & {
+export type TestStep = Pick<AiConversationRunProps, 'prompt' | 'system'> & {
   test?: () => Promise<void>;
 };
 
+/**
+ * Blueprint definition type for testing.
+ * Mirrors AppCapabilities.BlueprintDefinition to avoid circular dependency with app-toolkit.
+ */
+export interface BlueprintDefinition {
+  key: string;
+  functions: FunctionDefinition.Any[];
+  make: () => Blueprint.Blueprint;
+}
+
 /**
  * Runs the prompt steps, calling the test function after each step.
  */
@@ -30,5 +49,17 @@ export const runSteps = Effect.fn(function* (conversation: AiConversation, steps
   }
 });
 
-// TODO(wittjosiah): Don't cast.
-export const testToolkit = Toolkit.empty as Toolkit.Toolkit<any>;
+/**
+ * Binds blueprints from the blueprint definitions.
+ */
+// TODO(dmaretskyi): Potentially the agent will auto-bind the blueprints.
+export const addBlueprints = Effect.fnUntraced(function* (blueprints: BlueprintDefinition[]) {
+  yield* AiContextService.bindContext({
+    blueprints: yield* pipe(
+      blueprints,
+      Arr.map((blueprint) => blueprint.make()),
+      Effect.forEach(Database.add),
+      Effect.map(Arr.map(Ref.make)),
+    ),
+  });
+});

package/src/blueprints/websearch/blueprint.ts

@@ -0,0 +1,32 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { ToolId } from '@dxos/ai';
+import { type AppCapabilities } from '@dxos/app-toolkit';
+import { Blueprint } from '@dxos/blueprints';
+import { Ref } from '@dxos/echo';
+import { Text } from '@dxos/schema';
+
+const BLUEPRINT_KEY = 'dxos.org/blueprint/web-search';
+
+const functions: AppCapabilities.BlueprintDefinition['functions'] = [];
+
+const make = () =>
+  Blueprint.make({
+    key: BLUEPRINT_KEY,
+    name: 'Web Search',
+    description: 'Search the web.',
+    instructions: {
+      source: Ref.make(Text.make()),
+    },
+    tools: [ToolId.make('AnthropicWebSearch')],
+  });
+
+const blueprint: AppCapabilities.BlueprintDefinition = {
+  key: BLUEPRINT_KEY,
+  functions,
+  make,
+};
+
+export default blueprint;

package/src/blueprints/websearch/index.ts

@@ -2,7 +2,5 @@
 // Copyright 2025 DXOS.org
 //
 
-import blueprint from './websearch-blueprint';
-
-export { WebSearchToolkit } from './websearch-toolkit';
-export default blueprint;
+export { default as WebSearchBlueprint } from './blueprint';
+export { WebSearchToolkit } from './toolkit';

package/src/crud/graph.test.ts

@@ -0,0 +1,46 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { describe, it } from '@effect/vitest';
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+
+import { AiSession } from '@dxos/assistant';
+import { AssistantTestLayer } from '@dxos/assistant/testing';
+import { TestHelpers } from '@dxos/effect/testing';
+import { log } from '@dxos/log';
+import { Pipeline } from '@dxos/types';
+
+import { makeGraphWriterHandler, makeGraphWriterToolkit } from './graph';
+
+const TestLayer = AssistantTestLayer({});
+
+describe('graph', () => {
+  // let builder: EchoTestBuilder;
+  // test('findRelatedSchema', async () => {
+  //   const db = await createEchoDatabase();
+  //   const relatedSchemas = await findRelatedSchema(db, Schema.Struct({}));
+  // });
+
+  const Toolkit = makeGraphWriterToolkit({ schema: [Pipeline.Pipeline] });
+  const ToolkitLayer = makeGraphWriterHandler(Toolkit);
+
+  it.effect.skip(
+    'calculator',
+    Effect.fn(
+      function* (_) {
+        const session = new AiSession();
+        const toolkit = yield* Toolkit;
+        const response = yield* session.run({
+          toolkit,
+          prompt: 'What is 10 + 20?',
+        });
+
+        log.info('response', { response });
+      },
+      Effect.provide(Layer.mergeAll(TestLayer, ToolkitLayer)),
+      TestHelpers.provideTestContext,
+    ),
+  );
+});