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

package/src/functions/entity-extraction/entity-extraction.test.ts

@@ -7,21 +7,15 @@ import * as Effect from 'effect/Effect';
 import * as Layer from 'effect/Layer';
 
 import { AiService, MemoizedAiService } from '@dxos/ai';
-import { AiServiceTestingPreset } from '@dxos/ai/testing';
+import { TestAiService } from '@dxos/ai/testing';
 import { makeToolExecutionServiceFromFunctions, makeToolResolverFromFunctions } from '@dxos/assistant';
 import { Blueprint } from '@dxos/blueprints';
 import { Obj } from '@dxos/echo';
 import { TestHelpers } from '@dxos/effect';
-import {
-  ComputeEventLogger,
-  CredentialsService,
-  DatabaseService,
-  FunctionInvocationService,
-  TracingService,
-} from '@dxos/functions';
-import { TestDatabaseLayer } from '@dxos/functions/testing';
-import { ObjectId, PublicKey } from '@dxos/keys';
-import { DataType } from '@dxos/schema';
+import { CredentialsService, DatabaseService, FunctionInvocationService, TracingService } from '@dxos/functions';
+import { FunctionInvocationServiceLayerTest, TestDatabaseLayer } from '@dxos/functions-runtime/testing';
+import { ObjectId } from '@dxos/keys';
+import { Message, Organization, Person } from '@dxos/types';
 
 import { testToolkit } from '../../blueprints/testing';
 import { ResearchGraph } from '../research';
@@ -34,17 +28,15 @@ const TestLayer = Layer.mergeAll(
   AiService.model('@anthropic/claude-opus-4-0'),
   makeToolResolverFromFunctions([], testToolkit),
   makeToolExecutionServiceFromFunctions(testToolkit, testToolkit.toLayer({}) as any),
-  ComputeEventLogger.layerFromTracing,
 ).pipe(
-  Layer.provideMerge(FunctionInvocationService.layerTest({ functions: [entityExtraction] })),
+  Layer.provideMerge(FunctionInvocationServiceLayerTest({ functions: [entityExtraction] })),
   Layer.provideMerge(
     Layer.mergeAll(
-      MemoizedAiService.layerTest().pipe(Layer.provide(AiServiceTestingPreset('direct'))),
+      TestAiService(),
       TestDatabaseLayer({
-        // TODO(dmaretskyi): define the constant space key in the test-builder.
-        spaceKey: PublicKey.from('665c420e0dec9aa36c2bedca567afb0778701920e346eaf83ab2bd3403859723'),
+        spaceKey: 'fixed',
         indexing: { vector: true },
-        types: [Blueprint.Blueprint, DataType.Message, DataType.Person, DataType.Organization, ResearchGraph],
+        types: [Blueprint.Blueprint, Message.Message, Person.Person, Organization.Organization, ResearchGraph],
       }),
       CredentialsService.configuredLayer([]),
       TracingService.layerNoop,
@@ -52,13 +44,13 @@ const TestLayer = Layer.mergeAll(
   ),
 );
 
-describe.skip('Entity extraction', () => {
+describe('Entity extraction', () => {
   it.effect(
     'call a function to generate a research report',
     Effect.fnUntraced(
       function* (_) {
         const email = yield* DatabaseService.add(
-          Obj.make(DataType.Message, {
+          Obj.make(Message.Message, {
             [Obj.Meta]: {
               tags: ['important'],
             },

package/src/functions/entity-extraction/entity-extraction.ts

@@ -11,20 +11,22 @@ import * as Schema from 'effect/Schema';
 import { AiService } from '@dxos/ai';
 import { AiSession, makeToolExecutionServiceFromFunctions, makeToolResolverFromFunctions } from '@dxos/assistant';
 import { Filter, Obj, Ref } from '@dxos/echo';
-import { DatabaseService, FunctionInvocationService, defineFunction } from '@dxos/functions';
+import { DatabaseService, defineFunction } from '@dxos/functions';
+import { FunctionInvocationServiceLayerTest } from '@dxos/functions-runtime/testing';
 import { type DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { LegacyOrganization, Message, Organization, Person } from '@dxos/types';
 import { trim } from '@dxos/util';
 
-import { contextQueueLayerFromResearchGraph, makeGraphWriterHandler, makeGraphWriterToolkit } from '../research';
+import { makeGraphWriterHandler, makeGraphWriterToolkit } from '../../crud';
+import { contextQueueLayerFromResearchGraph } from '../research';
 
 export default defineFunction({
   key: 'dxos.org/functions/entity-extraction',
   name: 'Entity Extraction',
   description: 'Extracts entities from emails and transcripts.',
   inputSchema: Schema.Struct({
-    source: DataType.Message.annotations({ description: 'Email or transcript to extract entities from.' }),
+    source: Message.Message.annotations({ description: 'Email or transcript to extract entities from.' }),
 
     // TODO(dmaretskyi): Consider making this an array of blueprints instead.
     instructions: Schema.optional(Schema.String).annotations({ description: 'Instructions extraction process.' }),
@@ -39,11 +41,11 @@ export default defineFunction({
   handler: Effect.fnUntraced(
     function* ({ data: { source, instructions } }) {
       const contact = yield* extractContact(source);
-      let organization: DataType.Organization | null = null;
+      let organization: Organization.Organization | null = null;
 
       if (contact && !contact.organization) {
         const created: DXN[] = [];
-        const GraphWriterToolkit = makeGraphWriterToolkit({ schema: [DataType.LegacyOrganization] }).pipe();
+        const GraphWriterToolkit = makeGraphWriterToolkit({ schema: [LegacyOrganization] }).pipe();
         const GraphWriterHandler = makeGraphWriterHandler(GraphWriterToolkit, {
           onAppend: (dxns) => created.push(...dxns),
         });
@@ -64,7 +66,7 @@ export default defineFunction({
         if (created.length > 1) {
           throw new Error('Multiple organizations created');
         } else if (created.length === 1) {
-          organization = yield* DatabaseService.resolve(created[0], DataType.Organization);
+          organization = yield* DatabaseService.resolve(created[0], Organization.Organization);
           Obj.getMeta(organization).tags ??= [];
           Obj.getMeta(organization).tags!.push(...(Obj.getMeta(source)?.tags ?? []));
           contact.organization = Ref.make(organization);
@@ -83,14 +85,14 @@ export default defineFunction({
       ).pipe(
         Layer.provide(
           // TODO(dmaretskyi): This should be provided by environment.
-          Layer.mergeAll(FunctionInvocationService.layerTest()),
+          Layer.mergeAll(FunctionInvocationServiceLayerTest()),
         ),
       ),
     ),
   ),
 });
 
-const extractContact = Effect.fn('extractContact')(function* (message: DataType.Message) {
+const extractContact = Effect.fn('extractContact')(function* (message: Message.Message) {
   const name = message.sender.name;
   const email = message.sender.email;
   if (!email) {
@@ -98,12 +100,12 @@ const extractContact = Effect.fn('extractContact')(function* (message: DataType.
     return undefined;
   }
 
-  const { objects: existingContacts } = yield* DatabaseService.runQuery(Filter.type(DataType.Person));
+  const { objects: existingContacts } = yield* DatabaseService.runQuery(Filter.type(Person.Person));
 
   // Check for existing contact
   // TODO(dmaretskyi): Query filter DSL - https://linear.app/dxos/issue/DX-541/filtercontains-should-work-with-partial-objects
   const existingContact = existingContacts.find((contact) =>
-    contact.emails?.some((contactEmail) => contactEmail.value === email),
+    contact.emails?.some((contactEmail: any) => contactEmail.value === email),
   );
 
   if (existingContact) {
@@ -111,7 +113,7 @@ const extractContact = Effect.fn('extractContact')(function* (message: DataType.
     return existingContact;
   }
 
-  const newContact = Obj.make(DataType.Person, {
+  const newContact = Obj.make(Person.Person, {
     [Obj.Meta]: {
       tags: Obj.getMeta(message)?.tags,
     },
@@ -131,7 +133,7 @@ const extractContact = Effect.fn('extractContact')(function* (message: DataType.
 
   log.info('extracted email domain', { emailDomain });
 
-  const { objects: existingOrganisations } = yield* DatabaseService.runQuery(Filter.type(DataType.Organization));
+  const { objects: existingOrganisations } = yield* DatabaseService.runQuery(Filter.type(Organization.Organization));
   const matchingOrg = existingOrganisations.find((org) => {
     if (org.website) {
       try {

package/src/functions/linear/linear.test.ts

@@ -14,15 +14,14 @@ import { AiServiceTestingPreset } from '@dxos/ai/testing';
 import { makeToolExecutionServiceFromFunctions, makeToolResolverFromFunctions } from '@dxos/assistant';
 import { Obj, Query } from '@dxos/echo';
 import { TestHelpers } from '@dxos/effect';
+import { CredentialsService, DatabaseService, FunctionInvocationService } from '@dxos/functions';
+import { TracingServiceExt } from '@dxos/functions-runtime';
 import {
-  ComputeEventLogger,
-  CredentialsService,
-  DatabaseService,
-  FunctionInvocationService,
-  TracingService,
-} from '@dxos/functions';
-import { TestDatabaseLayer, testStoragePath } from '@dxos/functions/testing';
-import { DataType } from '@dxos/schema';
+  FunctionInvocationServiceLayerTestMocked,
+  TestDatabaseLayer,
+  testStoragePath,
+} from '@dxos/functions-runtime/testing';
+import { Person, Project, Task } from '@dxos/types';
 
 import { LINEAR_ID_KEY, default as fetchLinearIssues } from './sync-issues';
 
@@ -30,20 +29,18 @@ const TestLayer = Layer.mergeAll(
   AiService.model('@anthropic/claude-opus-4-0'),
   makeToolResolverFromFunctions([], Toolkit.make()),
   makeToolExecutionServiceFromFunctions(Toolkit.make() as any, Layer.empty as any),
-  ComputeEventLogger.layerFromTracing,
 ).pipe(
   Layer.provideMerge(
     Layer.mergeAll(
       AiServiceTestingPreset('direct'),
       TestDatabaseLayer({
         // indexing: { vector: true },
-        types: [DataType.Task, DataType.Person, DataType.Project],
+        types: [Task.Task, Person.Person, Project.Project],
         storagePath: testStoragePath({ name: 'feed-test-13' }),
       }),
       CredentialsService.layerConfig([{ service: 'linear.app', apiKey: Config.redacted('LINEAR_API_KEY') }]),
-      FunctionInvocationService.layerTestMocked({ functions: [fetchLinearIssues] }).pipe(
-        Layer.provideMerge(ComputeEventLogger.layerFromTracing),
-        Layer.provideMerge(TracingService.layerLogInfo()),
+      FunctionInvocationServiceLayerTestMocked({ functions: [fetchLinearIssues] }).pipe(
+        Layer.provideMerge(TracingServiceExt.layerLogInfo()),
       ),
       FetchHttpClient.layer,
     ),
@@ -61,17 +58,17 @@ describe('Linear', { timeout: 600_000 }, () => {
           team: '1127c63a-6f77-4725-9229-50f6cd47321c',
         });
 
-        const { objects: persons } = yield* DatabaseService.runQuery(Query.type(DataType.Person));
+        const { objects: persons } = yield* DatabaseService.runQuery(Query.type(Person.Person));
         console.log('people', {
           count: persons.length,
           people: persons.map((_) => `(${_.id}) ${Obj.getLabel(_)} [${Obj.getKeys(_, LINEAR_ID_KEY)[0]?.id}]`),
         });
-        const { objects: projects } = yield* DatabaseService.runQuery(Query.type(DataType.Project));
+        const { objects: projects } = yield* DatabaseService.runQuery(Query.type(Project.Project));
         console.log('projects', {
           count: projects.length,
           projects: projects.map((_) => `(${_.id}) ${Obj.getLabel(_)} [${Obj.getKeys(_, LINEAR_ID_KEY)[0]?.id}]`),
         });
-        const { objects: tasks } = yield* DatabaseService.runQuery(Query.type(DataType.Task));
+        const { objects: tasks } = yield* DatabaseService.runQuery(Query.type(Task.Task));
         console.log('tasks', {
           count: tasks.length,
           tasks: tasks.map((_) => `(${_.id}) ${Obj.getLabel(_)} [${Obj.getKeys(_, LINEAR_ID_KEY)[0]?.id}]`),

package/src/functions/linear/sync-issues.ts

@@ -12,7 +12,7 @@ import * as Schema from 'effect/Schema';
 import { Filter, Obj, Query, Ref, type Type } from '@dxos/echo';
 import { DatabaseService, defineFunction, withAuthorization } from '@dxos/functions';
 import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { Person, Project, Task } from '@dxos/types';
 
 import { syncObjects } from '../../sync';
 import { graphqlRequestBody } from '../../util';
@@ -87,7 +87,7 @@ export default defineFunction({
     const client = yield* HttpClient.HttpClient.pipe(Effect.map(withAuthorization({ service: 'linear.app' })));
 
     // Get the timestamp that was previosly synced.
-    const after = yield* getLatestUpdateTimestamp(data.team, DataType.Task);
+    const after = yield* getLatestUpdateTimestamp(data.team, Task.Task);
     log.info('will fetch', { after });
 
     // Fetch the issues that have changed since the last sync.
@@ -126,8 +126,8 @@ const getLatestUpdateTimestamp: (
   );
 });
 
-const mapLinearPerson = (person: LinearPerson, { teamId }: { teamId: string }): DataType.Person =>
-  Obj.make(DataType.Person, {
+const mapLinearPerson = (person: LinearPerson, { teamId }: { teamId: string }): Person.Person =>
+  Obj.make(Person.Person, {
     [Obj.Meta]: {
       keys: [
         {
@@ -143,8 +143,8 @@ const mapLinearPerson = (person: LinearPerson, { teamId }: { teamId: string }):
     nickname: person.name,
   });
 
-const mapLinearIssue = (issue: LinearIssue, { teamId }: { teamId: string }): DataType.Task =>
-  Obj.make(DataType.Task, {
+const mapLinearIssue = (issue: LinearIssue, { teamId }: { teamId: string }): Task.Task =>
+  Obj.make(Task.Task, {
     [Obj.Meta]: {
       keys: [
         {
@@ -170,7 +170,7 @@ const mapLinearIssue = (issue: LinearIssue, { teamId }: { teamId: string }): Dat
     project: !issue.project
       ? undefined
       : Ref.make(
-          DataType.makeProject({
+          Project.make({
             [Obj.Meta]: {
               keys: [
                 {

package/src/functions/research/{create-document.ts → document-create.ts}

@@ -5,65 +5,71 @@
 import * as Effect from 'effect/Effect';
 import * as Schema from 'effect/Schema';
 
-import { Relation } from '@dxos/echo';
+import { ArtifactId } from '@dxos/assistant';
+import { Obj, Relation } from '@dxos/echo';
 import { DatabaseService, TracingService, defineFunction } from '@dxos/functions';
-import { invariant } from '@dxos/invariant';
-import { DXN, ObjectId } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { Markdown } from '@dxos/plugin-markdown/types';
-import { DataType } from '@dxos/schema';
+import { HasSubject } from '@dxos/types';
 import { trim } from '@dxos/util';
 
 export default defineFunction({
-  key: 'dxos.org/function/research/create-document',
+  key: 'dxos.org/function/research/document-create',
   name: 'Create research document',
   description: 'Creates a note summarizing the research.',
   inputSchema: Schema.Struct({
+    subject: ArtifactId.annotations({
+      description: trim`
+        ID of the object (organization, contact, etc.) for which the research was performed. 
+      `,
+    }),
     name: Schema.String.annotations({
-      description: 'Name of the note.',
+      description: 'Name of the document.',
     }),
-
     content: Schema.String.annotations({
       description: trim`
         Content of the note. 
         Supports (and are prefered) references to research objects using @ syntax and <object> tags (refer to research blueprint instructions).
       `,
     }),
-
-    // TODO(dmaretskyi): Use a specialized type for this (e.g., ArtifactId renamed as RefFromLLM).
-    target: Schema.String.annotations({
-      description: trim`
-        Id of the object (organization, contact, etc.) for which the research was performed. 
-        This must be a ulid.
-      `,
+  }),
+  outputSchema: Schema.Struct({
+    document: ArtifactId.annotations({
+      description: 'DXN of the created document.',
     }),
   }),
-  outputSchema: Schema.Struct({}), // TODO(burdon): Schema.Void?
-  handler: Effect.fnUntraced(function* ({ data: { target, name, content } }) {
-    log.info('Creating research document', { target, name, content });
+  handler: Effect.fnUntraced(function* ({ data: { subject, name, content } }) {
+    log.info('Creating research document', { subject, name, content });
 
+    // TODO(burdon): Auto flush before and after calling function?
     yield* DatabaseService.flush({ indexes: true });
     yield* TracingService.emitStatus({ message: 'Creating research document...' });
-    invariant(ObjectId.isValid(target));
 
-    const targetObj = yield* DatabaseService.resolve(DXN.fromLocalObjectId(target));
+    // TODO(burdon): Type check.
+    const target = (yield* DatabaseService.resolve(ArtifactId.toDXN(subject))) as Obj.Any;
 
-    const doc = yield* DatabaseService.add(
-      Markdown.makeDocument({
+    // Create document.
+    const object = yield* DatabaseService.add(
+      Markdown.make({
         name,
         content,
       }),
     );
+
+    // Create relation.
     yield* DatabaseService.add(
-      Relation.make(DataType.HasSubject, {
-        [Relation.Source]: doc,
-        [Relation.Target]: targetObj as any,
+      Relation.make(HasSubject.HasSubject, {
+        [Relation.Source]: object,
+        [Relation.Target]: target,
         completedAt: new Date().toISOString(),
       }),
     );
+
     yield* DatabaseService.flush({ indexes: true });
+    log.info('Created research document', { subject, object });
 
-    log.info('Created research document', { target, name, content });
-    return {};
+    return {
+      document: Obj.getDXN(object).toString(),
+    };
   }),
 });

package/src/functions/research/index.ts

@@ -2,10 +2,9 @@
 // Copyright 2025 DXOS.org
 //
 
-import { default as create$ } from './create-document';
+import { default as create$ } from './document-create';
 import { default as research$ } from './research';
 
-export * from './graph';
 export * from './research-graph';
 export * from './types';
 

package/src/functions/research/{instructions-research.tpl → research-instructions.tpl}

@@ -6,10 +6,11 @@ The Research Agent outputs results in a structured format matching the schema pr
 The Research Agent is equipped with the ability to:
 
 - Generate precise and effective search queries 
-- Request web pages by query (through a `web_search` tool)
-- Read the full content of retrieved pages
+- Request web pages by query.
 - Synthesize accurate, clear, and structured answers using reliable information from the retrieved content
+{{#if entityExtraction}}
 - Search the local database for information using a vector index (through a `local_search` tool)
+{{/if}}
 
 The Research Agent always follows these principles:
 
@@ -19,6 +20,7 @@ The Research Agent always follows these principles:
 - Transparency: The Research Agent mentions which sources were used and explains how it arrived at conclusions.
 - Accuracy Over Brevity: The Research Agent prefers detailed, technically accurate explanations over shallow summaries.
 - The Research Agent admits uncertainty rather than misleading.
+{{#if entityExtraction}}
 - The Research Agent picks the most concrete schema types for extracted information.
 - The Research Agent fills schema fields completely with information it is confident about, and omits fields it is not confident about.
 - When outputting results, the Research Agent adds extra data that fits the schema even if not directly related to the user's question.
@@ -26,6 +28,7 @@ The Research Agent always follows these principles:
 - 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 deleted, so only the data written to the graph will be preserved.
+{{/if}}
 
 The Research Agent may be asked for:
 
@@ -41,8 +44,6 @@ The Research Agent breaks it into sub-questions (if applicable).
 
 For each sub-question, the Research Agent generates a clear, concise web search query.
 
-The Research Agent uses `web_search`(query) to retrieve information.
-
 The Research Agent extracts and synthesizes relevant answers.
 
 The Research Agent's output includes:
@@ -58,23 +59,27 @@ Here's how the Research Agent operates:
 2. The Research Agent performs a web search for each topic.
 3. The Research Agent reads and analyzes results, cross references information from multiple sources, and represents conflicting information as ranges of possible values.
 
+{{#if entityExtraction}}
 4. The Research Agent searches the local database for information using a vector index that might link to the user's question.
 6. The Research Agent creates relations and references between new objects and existing database objects when related, using existing object IDs as references.
 7. The Research Agent selects the most concrete schema types for extracted information, using multiple types as needed, and prints its decision and reasoning.
 5. The Research Agent creates a clear, structured answer to the user's question.
 8. The Research Agent submits results using the specific schema.
+{{/if}}
 
+{{#if entityExtraction}}
 IMPORTANT:
-
 - The Research Agent always runs the `local_search` tool to search the local database at least once before submitting results.
 - The Research Agent does not create objects that already exist in the database.
 - Ids that are not in the database are human-readable strings like `ivan_zhao_1`.
+{{/if}}
 
 Status reporting:
 
 The Research Agent reports its status frequently using the `<status>` tags: <status>Searching for Google Founders</status>
 The Research Agent reports its status in-between each tool call and before submitting results.
 
+{{#if entityExtraction}}
 <example>
 
 Based on my research, I can now provide information about Google and it's founders.
@@ -95,4 +100,7 @@ I will use the following schema to construct new objects:
 
 <status>Formatting results</status>
 
-</example>
+</example>
+{{/if}}
+
+Last content block is the full research note -- the result of the research.