@dxos/assistant 0.8.3-main.672df60 → 0.8.3-staging.0fa589b

package/src/blueprint/blueprint.test.ts

@@ -5,11 +5,11 @@
 import { Schema } from 'effect';
 import { describe, test } from 'vitest';
 
-import { AIServiceEdgeClient, ToolRegistry, ToolResult, createTool } from '@dxos/ai';
+import { EdgeAiServiceClient, ToolRegistry, ToolResult, createTool } from '@dxos/ai';
 import { AI_SERVICE_ENDPOINT, EXA_API_KEY } from '@dxos/ai/testing';
 import { ArtifactId } from '@dxos/artifact';
+import { Obj } from '@dxos/echo';
 import { EchoTestBuilder } from '@dxos/echo-db/testing';
-import { create } from '@dxos/echo-schema';
 import { DataType, DataTypes } from '@dxos/schema';
 
 import { BlueprintBuilder } from './blueprint';
@@ -23,7 +23,7 @@ import { createExaTool } from '../research/exa';
 
 // TODO(burdon): Don't run on CI.
 describe.skip('Blueprint', () => {
-  const aiService = new AIServiceEdgeClient({
+  const aiClient = new EdgeAiServiceClient({
     endpoint: AI_SERVICE_ENDPOINT.REMOTE,
     defaultGenerationOptions: {
       // model: '@anthropic/claude-sonnet-4-20250514',
@@ -39,10 +39,10 @@ describe.skip('Blueprint', () => {
       .step('Write a pitch deck for the product')
       .build();
 
-    const registry = new ToolRegistry([]);
-    const machine = new BlueprintMachine(registry, blueprint);
+    const tools = new ToolRegistry([]);
+    const machine = new BlueprintMachine(tools, blueprint);
     setConsolePrinter(machine, true);
-    await machine.runToCompletion({ aiService });
+    await machine.runToCompletion({ aiClient });
   });
 
   test('email bot', { timeout: 60_000 }, async () => {
@@ -96,10 +96,10 @@ describe.skip('Blueprint', () => {
       })
       .build();
 
-    const registry = new ToolRegistry([replyTool, labelTool]);
-    const machine = new BlueprintMachine(registry, blueprint);
+    const tools = new ToolRegistry([replyTool, labelTool]);
+    const machine = new BlueprintMachine(tools, blueprint);
     setConsolePrinter(machine);
-    await machine.runToCompletion({ aiService, input: TEST_EMAILS[0] });
+    await machine.runToCompletion({ aiClient, input: TEST_EMAILS[0] });
   });
 
   test.only('research', { timeout: 120_000 }, async () => {
@@ -108,14 +108,14 @@ describe.skip('Blueprint', () => {
 
     const [org1] = [
       db.add(
-        create(DataType.Organization, {
+        Obj.make(DataType.Organization, {
           name: 'Exa',
           website: 'https://exa.ai',
           description: 'An AI-powered search engine company building search infrastructure for AI agents',
         }),
       ),
       db.add(
-        create(DataType.Organization, {
+        Obj.make(DataType.Organization, {
           name: 'Cresta',
           website: 'https://cresta.ai',
           description: 'A company that builds AI agents',
@@ -125,25 +125,27 @@ describe.skip('Blueprint', () => {
 
     await db.flush({ indexes: true });
 
-    const webSearchTool = createExaTool({ apiKey: EXA_API_KEY });
-    const localSearchTool = createLocalSearchTool(db);
-    const graphWriterTool = createGraphWriterTool({ db, schemaTypes: DataTypes });
-    const registry = new ToolRegistry([webSearchTool, localSearchTool, graphWriterTool]);
+    const [exa, localSearch, graphWriter] = [
+      createExaTool({ apiKey: EXA_API_KEY }),
+      createLocalSearchTool(db),
+      createGraphWriterTool({ db, schema: DataTypes }),
+    ];
 
     const blueprint = BlueprintBuilder.create()
       .step('Research founders of the organization. Do deep research.', {
-        tools: [webSearchTool.id],
+        tools: [exa.id],
       })
       .step('Based on your research select matching entires that are already in the graph', {
-        tools: [localSearchTool.id],
+        tools: [localSearch.id],
       })
       .step('Add researched data to the graph', {
-        tools: [localSearchTool.id, graphWriterTool.id],
+        tools: [localSearch.id, graphWriter.id],
       })
       .build();
 
-    const machine = new BlueprintMachine(registry, blueprint);
+    const tools = new ToolRegistry([exa, localSearch, graphWriter]);
+    const machine = new BlueprintMachine(tools, blueprint);
     setConsolePrinter(machine, true);
-    await machine.runToCompletion({ aiService, input: org1 });
+    await machine.runToCompletion({ aiClient, input: org1 });
   });
 });

package/src/blueprint/blueprint.ts

@@ -4,26 +4,30 @@
 
 import { Schema } from 'effect';
 
-import { Obj, Type } from '@dxos/echo';
-import { ObjectId } from '@dxos/keys';
+import { Key, Obj, Type } from '@dxos/echo';
 
-export const BlueprintStepSchema = Schema.Struct({
-  id: ObjectId,
+export const BlueprintStep = Schema.Struct({
+  id: Key.ObjectId,
   instructions: Schema.String,
-  tools: Schema.Array(Schema.String),
+  tools: Schema.optional(Schema.Array(Schema.String)),
 });
-export type BlueprintStep = Schema.Schema.Type<typeof BlueprintStepSchema>;
+export interface BlueprintStep extends Schema.Schema.Type<typeof BlueprintStep> {}
 
-export const BlueprintSchema = Schema.Struct({
-  steps: Schema.Array(BlueprintStepSchema.pipe(Schema.omit('id'))),
+export const BlueprintDefinition = Schema.Struct({
+  steps: Schema.Array(BlueprintStep.pipe(Schema.omit('id'))),
 });
-export type BlueprintDef = Schema.Schema.Type<typeof BlueprintSchema>;
+export interface BlueprintDefinition extends Schema.Schema.Type<typeof BlueprintDefinition> {}
 
 export const Blueprint = Schema.Struct({
   name: Schema.optional(Schema.String),
-  steps: Schema.Array(BlueprintStepSchema),
-}).pipe(Type.Obj({ typename: 'dxos.org/type/Blueprint', version: '0.1.0' }));
-export type Blueprint = Schema.Schema.Type<typeof Blueprint>;
+  steps: Schema.Array(BlueprintStep),
+}).pipe(
+  Type.Obj({
+    typename: 'dxos.org/type/Blueprint',
+    version: '0.1.0',
+  }),
+);
+export interface Blueprint extends Schema.Schema.Type<typeof Blueprint> {}
 
 /**
  * Blueprint builder API.
@@ -36,7 +40,7 @@ export namespace BlueprintBuilder {
 
     step(instructions: string, options?: { tools?: string[] }): Builder {
       this._steps.push({
-        id: ObjectId.random(),
+        id: Key.ObjectId.random(),
         instructions,
         tools: options?.tools ?? [],
       });
@@ -57,7 +61,7 @@ export namespace BlueprintParser {
   export const create = () => new Parser();
 
   class Parser {
-    parse({ steps }: BlueprintDef): Blueprint {
+    parse({ steps }: BlueprintDefinition): Blueprint {
       const builder = BlueprintBuilder.create();
       for (const step of steps) {
         builder.step(step.instructions, {

package/src/blueprint/logger.ts

@@ -7,16 +7,34 @@ import chalk from 'chalk';
 
 import { ConsolePrinter } from '@dxos/ai';
 import { type CleanupFn, combine } from '@dxos/async';
+import { log } from '@dxos/log';
 
 import { type Blueprint } from './blueprint';
-import { type BlueprintMachine, type BlueprintMachineState } from './machine';
+import {
+  type BlueprintEvent,
+  type BlueprintLogger,
+  type BlueprintMachine,
+  type BlueprintMachineState,
+} from './machine';
 
 /* eslint-disable no-console */
 
 // Force chalk colors on for tests.
 chalk.level = 2;
 
-export class Logger {
+// TODO(burdon): Factor out.
+type ConsoleLogSignature = (message: string, arg?: any) => void;
+
+interface Logger {
+  log: ConsoleLogSignature;
+}
+
+const DEFAULT_LOGGER: Logger = { log: log.info };
+
+/**
+ * Reactive bufferedlogger.
+ */
+export class BufferedLogger implements Logger {
   private _messages = signal<string[]>([]);
 
   get messages(): ReadonlySignal<string[]> {
@@ -32,7 +50,35 @@ export class Logger {
   }
 }
 
-export const setLogger = (machine: BlueprintMachine, logger: Logger): CleanupFn => {
+// TODO(burdon): Reconcile with ConsolePrinter.
+export class BlueprintLoggerAdapter implements BlueprintLogger {
+  constructor(private readonly logger: Logger = DEFAULT_LOGGER) {}
+
+  log(event: BlueprintEvent) {
+    switch (event.type) {
+      case 'begin':
+        this.logger.log('begin', { invocationId: event.invocationId });
+        break;
+      case 'end':
+        this.logger.log('end', { invocationId: event.invocationId });
+        break;
+      case 'step-start':
+        this.logger.log('step-start', { invocationId: event.invocationId, step: event.step });
+        break;
+      case 'step-complete':
+        this.logger.log('step-complete', { invocationId: event.invocationId, step: event.step });
+        break;
+      case 'message':
+        this.logger.log('message', { invocationId: event.invocationId, message: event.message });
+        break;
+      case 'block':
+        this.logger.log('block', { invocationId: event.invocationId, block: event.block });
+        break;
+    }
+  }
+}
+
+export const setLogger = (machine: BlueprintMachine, logger: BufferedLogger): CleanupFn => {
   return combine(
     machine.begin.on(() => {
       logger.log('Starting...');

package/src/blueprint/machine.ts

@@ -6,7 +6,7 @@ import { Match, Schema } from 'effect';
 
 import {
   createTool,
-  type AIServiceClient,
+  type AiServiceClient,
   Message,
   type MessageContentBlock,
   ToolResult,
@@ -14,7 +14,7 @@ import {
   type ToolRegistry,
 } from '@dxos/ai';
 import { Event } from '@dxos/async';
-import { create } from '@dxos/echo-schema';
+import { Key, Obj } from '@dxos/echo';
 import { type ObjectId } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { isNonNullable } from '@dxos/util';
@@ -22,6 +22,17 @@ import { isNonNullable } from '@dxos/util';
 import type { Blueprint, BlueprintStep } from './blueprint';
 import { AISession } from '../session';
 
+const SYSTEM_PROMPT = `
+You are a smart Rule-Following Agent.
+Rule-Following Agent executes the command that the user sent in the last message.
+After doing the work, the Rule-Following Agent calls report tool.
+If the Rule-Following agent believes that no action is needed, it calls the report tool with "skipped" status.
+If the Rule-Following Agent is unable to perform the task, it calls the report tool with "bail" status.
+Rule-Following Agent explains the reason it is unable to perform the task before bailing.
+The Rule-Following Agent can express creativity and imagination in the way it performs the task.
+The Rule-Following Agent precisely follows the instructions.
+`;
+
 export type BlueprintTraceStep = {
   status: 'done' | 'bailed' | 'skipped';
   stepId: ObjectId;
@@ -41,7 +52,7 @@ const INITIAL_STATE: BlueprintMachineState = {
 };
 
 type ExecutionOptions = {
-  aiService: AIServiceClient;
+  aiClient: AiServiceClient;
 
   /**
    * Input to the blueprint.
@@ -49,6 +60,18 @@ type ExecutionOptions = {
   input?: unknown;
 };
 
+export type BlueprintEvent =
+  | { type: 'begin'; invocationId: string }
+  | { type: 'end'; invocationId: string }
+  | { type: 'step-start'; invocationId: string; step: BlueprintStep }
+  | { type: 'step-complete'; invocationId: string; step: BlueprintStep }
+  | { type: 'message'; invocationId: string; message: Message }
+  | { type: 'block'; invocationId: string; block: MessageContentBlock };
+
+export interface BlueprintLogger {
+  log(event: BlueprintEvent): void;
+}
+
 /**
  * Blueprint state machine.
  */
@@ -60,46 +83,67 @@ export class BlueprintMachine {
   public readonly message = new Event<Message>();
   public readonly block = new Event<MessageContentBlock>();
 
-  state: BlueprintMachineState = structuredClone(INITIAL_STATE);
+  // TODO(burdon): Replace events with logger?
+  private logger?: BlueprintLogger;
+
+  private _state: BlueprintMachineState = structuredClone(INITIAL_STATE);
 
   constructor(
     readonly registry: ToolRegistry,
     readonly blueprint: Blueprint,
   ) {}
 
+  get state(): BlueprintMachineState {
+    return this._state;
+  }
+
+  setLogger(logger: BlueprintLogger): this {
+    this.logger = logger;
+    return this;
+  }
+
   async runToCompletion(options: ExecutionOptions): Promise<void> {
-    log.info('runToCompletion', options);
+    const invocationId = Key.ObjectId.random();
+    log.info('runToCompletion', { invocationId, options });
 
     this.begin.emit();
+    this.logger?.log({ type: 'begin', invocationId });
+
     let firstStep = true;
-    while (this.state.state !== 'done') {
+    while (this._state.state !== 'done') {
       const input = firstStep ? options.input : undefined;
-
       firstStep = false;
-      this.state = await this._execStep(this.state, {
-        input,
-        aiService: options.aiService,
-      });
 
-      this.stepComplete.emit(this.blueprint.steps.find((step) => step.id === this.state.trace.at(-1)?.stepId)!);
+      this._state = await this._execStep(invocationId, this._state, { input, aiClient: options.aiClient });
+
+      const step = this.blueprint.steps.find((step) => step.id === this._state.trace.at(-1)?.stepId)!;
+      this.stepComplete.emit(step);
+      this.logger?.log({ type: 'step-complete', invocationId, step });
 
-      if (this.state.state === 'bail') {
+      if (this._state.state === 'bail') {
         throw new Error('Agent unable to follow the blueprint');
       }
     }
 
     this.end.emit();
+    this.logger?.log({ type: 'end', invocationId });
   }
 
-  private async _execStep(state: BlueprintMachineState, options: ExecutionOptions): Promise<BlueprintMachineState> {
-    const prevStep = this.blueprint.steps.findIndex((step) => step.id === this.state.trace.at(-1)?.stepId);
+  private async _execStep(
+    invocationId: string,
+    state: BlueprintMachineState,
+    options: ExecutionOptions,
+  ): Promise<BlueprintMachineState> {
+    const prevStep = this.blueprint.steps.findIndex((step) => step.id === this._state.trace.at(-1)?.stepId);
     if (prevStep === this.blueprint.steps.length - 1) {
       throw new Error('Done execution blueprint');
     }
 
     const nextStep = this.blueprint.steps[prevStep + 1];
     const onLastStep = prevStep === this.blueprint.steps.length - 2;
+
     this.stepStart.emit(nextStep);
+    this.logger?.log({ type: 'step-start', invocationId, step: nextStep });
 
     const ReportSchema = Schema.Struct({
       status: Schema.Literal('done', 'bailed', 'skipped').annotations({
@@ -113,7 +157,7 @@ export class BlueprintMachine {
 
     const report = createTool('system', {
       name: 'report',
-      description: 'This tool reports that the agent has completed the task or is unable to do so',
+      description: 'This tool reports that the agent has completed the task or is unable to do so.',
       schema: ReportSchema,
       execute: async (input) => ToolResult.Break(input),
     });
@@ -128,38 +172,31 @@ export class BlueprintMachine {
 
     const inputMessages = options.input
       ? [
-          create(Message, {
+          Obj.make(Message, {
             role: 'user',
             content: [taggedDataBlock('input', options.input)],
           }),
         ]
       : [];
-    inputMessages.forEach((message) => this.message.emit(message));
+    inputMessages.forEach((message) => {
+      this.message.emit(message);
+      this.logger?.log({ type: 'message', invocationId, message });
+    });
 
     // TODO(wittjosiah): Warn if tool is not found.
-    const tools = nextStep.tools.map((tool) => this.registry.get(tool)).filter(isNonNullable);
+    const tools = nextStep.tools?.map((tool) => this.registry.get(tool)).filter(isNonNullable) ?? [];
     const messages = await session.run({
-      systemPrompt: `
-        You are a smart Rule-Following Agent.
-        Rule-Following Agent executes the command that the user sent in the last message.
-        After doing the work, the Rule-Following Agent calls report tool.
-        If the Rule-Following agent believes that no action is needed, it calls the report tool with "skipped" status.
-        If the Rule-Following Agent is unable to perform the task, it calls the report tool with "bail" status.
-        Rule-Following Agent explains the reason it is unable to perform the task before bailing.
-        The Rule-Following Agent can express creativity and imagination in the way it performs the task.
-        The Rule-Following Agent precisely follows the instructions.
-      `,
+      systemPrompt: SYSTEM_PROMPT,
       history: [...state.history, ...inputMessages],
       tools: [...tools, report],
       artifacts: [],
-      client: options.aiService,
+      client: options.aiClient,
       prompt: nextStep.instructions,
     });
 
     const { messages: trimmedHistory, call: lastBlock } = popLastToolCall(messages);
-
     if (!lastBlock) {
-      // TODO(dmaretskyi): Handle this with grace
+      // TODO(dmaretskyi): Handle this with grace.
       throw new Error('Agent did not call the report tool');
     }
 

package/src/blueprint/test-data.ts

@@ -2,12 +2,12 @@
 // Copyright 2025 DXOS.org
 //
 
-import { create } from '@dxos/echo-schema';
+import { Obj } from '@dxos/echo';
 import { DataType } from '@dxos/schema';
 
 // TODO(burdon): Convert to queue.
 export const TEST_EMAILS: DataType.Message[] = [
-  create(DataType.Message, {
+  Obj.make(DataType.Message, {
     properties: {
       subject: 'Introduction email',
     },
@@ -22,7 +22,7 @@ export const TEST_EMAILS: DataType.Message[] = [
     ],
     created: new Date().toISOString(),
   }),
-  create(DataType.Message, {
+  Obj.make(DataType.Message, {
     properties: {
       subject: 'Question about integration',
     },
@@ -37,7 +37,7 @@ export const TEST_EMAILS: DataType.Message[] = [
     ],
     created: new Date().toISOString(),
   }),
-  create(DataType.Message, {
+  Obj.make(DataType.Message, {
     properties: {
       subject: 'MAKE MONEY FAST!!!',
     },

package/src/extraction/extraction-llm-function/extraction-llm-function.test.ts

@@ -4,7 +4,7 @@
 
 import { describe, test, beforeAll } from 'vitest';
 
-import { AIServiceEdgeClient } from '@dxos/ai';
+import { EdgeAiServiceClient } from '@dxos/ai';
 import { AI_SERVICE_ENDPOINT } from '@dxos/ai/testing';
 import { FunctionExecutor, ServiceContainer } from '@dxos/functions';
 import { log } from '@dxos/log';
@@ -21,7 +21,7 @@ describe.skip('LLM EntityExtraction', () => {
     executor = new FunctionExecutor(
       new ServiceContainer().setServices({
         ai: {
-          client: new AIServiceEdgeClient({
+          client: new EdgeAiServiceClient({
             endpoint: AI_SERVICE_ENDPOINT.REMOTE,
             defaultGenerationOptions: {
               // model: '@anthropic/claude-sonnet-4-20250514',

package/src/extraction/extraction-llm-function/extraction-llm-function.ts

@@ -3,6 +3,7 @@
 //
 
 import { Message } from '@dxos/ai';
+import { Obj } from '@dxos/echo';
 import { create } from '@dxos/echo-schema';
 import { AiService, defineFunction, type FunctionDefinition } from '@dxos/functions';
 import { DataType } from '@dxos/schema';
@@ -27,7 +28,7 @@ export const extractionAnthropicFn: FunctionDefinition<ExtractionInput, Extracti
       client: ai.client,
       systemPrompt: PROMPT,
       history: [
-        create(Message, {
+        Obj.make(Message, {
           role: 'user',
           content: [
             {

package/src/extraction/quotes.ts

@@ -4,8 +4,9 @@
 
 import { Schema } from 'effect';
 
+import { Query, Filter } from '@dxos/echo';
 import { type EchoDatabase } from '@dxos/echo-db';
-import { Filter, ObjectId, Query } from '@dxos/echo-schema';
+import { ObjectId } from '@dxos/echo-schema';
 import { DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
 

package/src/research/graph.ts

@@ -6,6 +6,7 @@ import { Schema, identity, Option, SchemaAST } from 'effect';
 
 import { createTool, ToolResult } from '@dxos/ai';
 import type { Obj, Relation } from '@dxos/echo';
+import { Query, Filter } from '@dxos/echo';
 import { type EchoDatabase, type Queue } from '@dxos/echo-db';
 import { isEncodedReference } from '@dxos/echo-protocol';
 import {
@@ -13,13 +14,10 @@ import {
   getSchemaTypename,
   getTypeAnnotation,
   getTypeIdentifierAnnotation,
-  type AnyEchoObject,
   create,
-  Filter,
   getEntityKind,
   getSchemaDXN,
   ObjectId,
-  Query,
   ReferenceAnnotationId,
   type BaseObject,
   RelationSourceDXNId,
@@ -319,21 +317,21 @@ const preprocessSchema = (schema: Schema.Schema.AnyNoContext) => {
 export const createGraphWriterTool = ({
   db,
   queue,
-  schemaTypes,
+  schema,
   onDone = async (x) => x,
 }: {
   db: EchoDatabase;
   queue?: Queue;
-  schemaTypes: Schema.Schema.AnyNoContext[];
-  onDone?: (data: AnyEchoObject[]) => Promise<any>;
+  schema: Schema.Schema.AnyNoContext[];
+  onDone?: (data: Obj.Any[]) => Promise<any>;
 }) => {
   return createTool('graph', {
     name: 'writer',
     description: 'Write to the local graph database',
-    schema: createExtractionSchema(schemaTypes),
+    schema: createExtractionSchema(schema),
     execute: async (input) => {
-      const data = await sanitizeObjects(schemaTypes, input as any, db, queue);
-      return ToolResult.Success(await onDone(data as AnyEchoObject[]));
+      const data = await sanitizeObjects(schema, input as any, db, queue);
+      return ToolResult.Success(await onDone(data as Obj.Any[]));
     },
   });
 };

package/src/research/research.test.ts

@@ -5,13 +5,13 @@
 import { inspect } from 'node:util';
 import { beforeAll, describe, test } from 'vitest';
 
-import { AIServiceEdgeClient, OllamaClient, structuredOutputParser } from '@dxos/ai';
+import { EdgeAiServiceClient, OllamaAiServiceClient, structuredOutputParser } from '@dxos/ai';
 import { AI_SERVICE_ENDPOINT, EXA_API_KEY } from '@dxos/ai/testing';
+import { Obj } from '@dxos/echo';
 import { type EchoDatabase } from '@dxos/echo-db';
 import { EchoTestBuilder } from '@dxos/echo-db/testing';
 import { getSchemaDXN } from '@dxos/echo-schema';
 import { ConfiguredCredentialsService, FunctionExecutor, ServiceContainer, TracingService } from '@dxos/functions';
-import { live } from '@dxos/live-object';
 import { DataType, DataTypes } from '@dxos/schema';
 
 import { createExtractionSchema, getSanitizedSchemaName } from './graph';
@@ -20,7 +20,7 @@ import { researchFn } from './research';
 const REMOTE_AI = true;
 const MOCK_SEARCH = false;
 
-describe('Research', () => {
+describe.skip('Research', () => {
   let builder: EchoTestBuilder;
   let db: EchoDatabase;
   let executor: FunctionExecutor;
@@ -28,21 +28,22 @@ describe('Research', () => {
   beforeAll(async () => {
     // TODO(dmaretskyi): Helper to scaffold this from a config.
     builder = await new EchoTestBuilder().open();
-    const { db: db1 } = await builder.createDatabase({ indexing: { vector: true } });
-    db = db1;
+
+    db = (await builder.createDatabase({ indexing: { vector: true } })).db;
     db.graph.schemaRegistry.addSchema(DataTypes);
+
     executor = new FunctionExecutor(
       new ServiceContainer().setServices({
         ai: {
           client: REMOTE_AI
-            ? new AIServiceEdgeClient({
+            ? new EdgeAiServiceClient({
                 endpoint: AI_SERVICE_ENDPOINT.REMOTE,
                 defaultGenerationOptions: {
                   // model: '@anthropic/claude-sonnet-4-20250514',
                   model: '@anthropic/claude-3-5-sonnet-20241022',
                 },
               })
-            : new OllamaClient({
+            : new OllamaAiServiceClient({
                 overrides: {
                   model: 'llama3.1:8b',
                 },
@@ -57,7 +58,7 @@ describe('Research', () => {
 
   test('should generate a research report', { timeout: 300_000 }, async () => {
     db.add(
-      live(DataType.Organization, {
+      Obj.make(DataType.Organization, {
         name: 'Notion',
         website: 'https://www.notion.com',
       }),

package/src/research/research.ts

@@ -41,7 +41,7 @@ export const researchFn = defineFunction({
     const ai = context.getService(AiService);
     const credentials = context.getService(CredentialsService);
     const { db } = context.getService(DatabaseService);
-    // const queues = context.getService(QueuesService);
+    // const queues = context.getService(QueueService);
     const tracing = context.getService(TracingService);
 
     tracing.write(create(AgentStatusReport, { message: 'Researching...' }));
@@ -60,7 +60,7 @@ export const researchFn = defineFunction({
     });
     session.streamEvent.on((event) => log('stream', { event }));
 
-    const graphWriteTool = createGraphWriterTool({ db, schemaTypes: DataTypes });
+    const graphWriteTool = createGraphWriterTool({ db, schema: DataTypes });
     log.info('graphWriteTool', { schema: graphWriteTool.parameters });
 
     // TODO(dmaretskyi): Consider adding this pattern as the "Graph" output mode for the session.