@dxos/assistant 0.8.2-main.5885341 → 0.8.2-main.600d381

package/src/session/session.test.ts

@@ -2,35 +2,33 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Schema as S } from 'effect';
-import { inspect } from 'node:util';
+import { Schema } from 'effect';
 import { describe, test } from 'vitest';
 
-import {
-  ArtifactId,
-  defineArtifact,
-  defineTool,
-  ToolResult,
-  type Message,
-  type MessageContentBlock,
-} from '@dxos/artifact';
+import { ArtifactId, defineArtifact, defineTool, ToolResult } from '@dxos/artifact';
 import { AIServiceEdgeClient } from '@dxos/assistant';
-import { create, EchoObject, ObjectId } from '@dxos/echo-schema';
+import { Type } from '@dxos/echo';
+import { create, ObjectId } from '@dxos/echo-schema';
 import { DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
 
 import { AISession } from './session';
-import { AI_SERVICE_ENDPOINT } from '../testing';
+import { AI_SERVICE_ENDPOINT, ConsolePrinter } from '../testing';
 
 // Define a calendar event artifact schema.
-const CalendarEventSchema = S.Struct({
-  title: S.String,
-  startTime: S.String,
-  endTime: S.String,
-  description: S.String,
-}).pipe(EchoObject({ typename: 'example.com/type/CalendarEvent', version: '0.1.0' }));
+const CalendarEventSchema = Schema.Struct({
+  title: Schema.String,
+  startTime: Schema.String,
+  endTime: Schema.String,
+  description: Schema.String,
+}).pipe(
+  Type.def({
+    typename: 'example.com/type/CalendarEvent',
+    version: '0.1.0',
+  }),
+);
 
-type CalendarEvent = S.Schema.Type<typeof CalendarEventSchema>;
+type CalendarEvent = Schema.Schema.Type<typeof CalendarEventSchema>;
 
 // TODO(burdon): Flaky.
 describe.skip('AISession with Ollama', () => {
@@ -53,7 +51,7 @@ describe.skip('AISession with Ollama', () => {
         defineTool('calendar', {
           name: 'query',
           description: 'Query the calendar for events',
-          schema: S.Struct({}),
+          schema: Schema.Struct({}),
           execute: async () => {
             return ToolResult.Success(CALENDAR_EVENTS);
           },
@@ -65,13 +63,13 @@ describe.skip('AISession with Ollama', () => {
       id: 'artifact:dxos.org/plugin/table',
       name: 'Table',
       instructions: 'Use this to create and manage tables. Each table has a unique id.',
-      schema: S.Struct({}),
+      schema: Schema.Struct({}),
       tools: [
         defineTool('table', {
           name: 'create',
           description: 'Create a table',
-          schema: S.Struct({
-            data: S.Array(S.Any).annotations({ description: 'Array of data payloads to add as rows' }),
+          schema: Schema.Struct({
+            data: Schema.Array(Schema.Any).annotations({ description: 'Array of data payloads to add as rows' }),
           }),
           execute: async ({ data }) => {
             log('create table', { data });
@@ -89,12 +87,12 @@ describe.skip('AISession with Ollama', () => {
       name: 'Map',
       instructions:
         'Use this to create and manage maps. Maps source data from tables. Table id is required to create a map.',
-      schema: S.Struct({}),
+      schema: Schema.Struct({}),
       tools: [
         defineTool('map', {
           name: 'create',
           description: 'Create a map',
-          schema: S.Struct({
+          schema: Schema.Struct({
             source: ArtifactId.annotations({
               description: 'The table that will be used as the source of the map',
             }),
@@ -119,7 +117,7 @@ describe.skip('AISession with Ollama', () => {
       id: 'artifact:dxos.org/plugin/script',
       name: 'Script',
       instructions: 'Use this to create and manage scripts',
-      schema: S.Struct({}),
+      schema: Schema.Struct({}),
       tools: [],
     });
 
@@ -127,9 +125,10 @@ describe.skip('AISession with Ollama', () => {
     //   printStreamEvent(event);
     // });
 
-    session.message.on(printMessage);
-    session.userMessage.on(printMessage);
-    session.block.on(printContentBlock);
+    const printer = new ConsolePrinter();
+    session.message.on((message) => printer.printMessage(message));
+    session.userMessage.on((message) => printer.printMessage(message));
+    session.block.on((block) => printer.printContentBlock(block));
 
     // session.update.on((update) => {
     //   log('update', { update });
@@ -189,34 +188,6 @@ const CALENDAR_EVENTS: CalendarEvent[] = [
   }),
 ];
 
-const printMessage = (message: Message) => {
-  console.log(`${message.role.toUpperCase()}\n`);
-  for (const content of message.content) {
-    printContentBlock(content);
-  }
-};
-
-const printContentBlock = (content: MessageContentBlock) => {
-  switch (content.type) {
-    case 'text':
-      console.log(content.text);
-      break;
-    case 'tool_use':
-      console.log(`⚙️ [Tool Use] ${content.name} ${inspect(content.input, { depth: null, colors: true })}`);
-      break;
-    case 'tool_result': {
-      let data: any;
-      try {
-        data = JSON.parse(content.content);
-      } catch {
-        data = content.content;
-      }
-      console.log(`⚙️ [Tool Result] ${content.toolUseId} ${inspect(data, { depth: null, colors: true })}`);
-      break;
-    }
-  }
-};
-
 // const printStreamEvent = (event: GenerationStreamEvent) => {
 //   switch (event.type) {
 //     case 'message_start': {

package/src/session/session.ts

@@ -2,11 +2,12 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Option, Schema as S } from 'effect';
+import { Option, Schema } from 'effect';
 
 import {
   defineTool,
   Message,
+  structuredOutputParser,
   ToolResult,
   type ArtifactDefinition,
   type MessageContentBlock,
@@ -18,9 +19,9 @@ import { create, type ObjectId } from '@dxos/echo-schema';
 import { invariant } from '@dxos/invariant';
 import { log } from '@dxos/log';
 
+import { isToolUse, runTools } from './tools';
 import { VersionPin } from './version-pin';
 import { MixedStreamParser, type AIServiceClient, type GenerateRequest, type GenerationStream } from '../ai-service';
-import { isToolUse, runTools } from '../conversation';
 
 /**
  * Contains message history, tools, current context.
@@ -245,6 +246,15 @@ export class AISession {
     return this._pending;
   }
 
+  async runStructured<S extends Schema.Schema.AnyNoContext>(schema: S, options: SessionRunOptions) {
+    const parser = structuredOutputParser(schema);
+    const result = await this.run({
+      ...options,
+      tools: [...options.tools, parser.tool],
+    });
+    return parser.getResult(result);
+  }
+
   private async _formatUserPrompt(
     artifactDiffResolver: ArtifactDiffResolver | undefined,
     prompt: string,
@@ -287,7 +297,7 @@ export class AISession {
     return defineTool('system', {
       name: 'query_artifact_definitions',
       description: 'Query the available artifact definitions',
-      schema: S.Struct({}),
+      schema: Schema.Struct({}),
       execute: async () => {
         return ToolResult.Success({
           artifactDefinitions: artifacts.map((artifact) => ({
@@ -308,8 +318,8 @@ export class AISession {
       name: 'require_artifact_definitions',
       description:
         'Require the use of specific artifact definitions. This will allow the model to interact with artifact definitions and use their tools.',
-      schema: S.Struct({
-        artifactDefinitionIds: S.Array(S.String).annotations({
+      schema: Schema.Struct({
+        artifactDefinitionIds: Schema.Array(Schema.String).annotations({
           description: 'The ids of the artifact definitions to require',
           examples: [['artifact:dxos.org/example/Test']],
         }),
@@ -336,20 +346,20 @@ export class AISession {
       name: 'create_plan',
       description:
         'Create a plan. Make sure that each step is independent and can be executed solely on the data returned by the previous step. The steps only share the data that is specified in the plan.',
-      schema: S.Struct({
-        goal: S.String.annotations({ description: 'The goal that the plan will achieve.' }),
-        steps: S.Array(
-          S.Struct({
-            action: S.String.annotations({
+      schema: Schema.Struct({
+        goal: Schema.String.annotations({ description: 'The goal that the plan will achieve.' }),
+        steps: Schema.Array(
+          Schema.Struct({
+            action: Schema.String.annotations({
               description: 'Complete, detailed action to perform.',
             }),
-            input: S.String.annotations({
+            input: Schema.String.annotations({
               description: 'The required input data for this step. First step must not require any data.',
             }),
-            output: S.String.annotations({
+            output: Schema.String.annotations({
               description: 'The output data from this step. This is passed to the next step as input.',
             }),
-            requiredArtifactIds: S.Array(S.String).annotations({
+            requiredArtifactIds: Schema.Array(Schema.String).annotations({
               description: 'The ids of the artifacts required to perform this step.',
               examples: [['artifact:dxos.org/example/Test']],
             }),
@@ -473,7 +483,7 @@ const createBaseInstructions = ({
   ${availableArtifacts.length > 0 ? `Artifacts already in context: ${availableArtifacts.join('\n')}` : ''}
 `;
 
-export class AIServiecOverloadedError extends S.TaggedError<AIServiecOverloadedError>()(
+export class AIServiecOverloadedError extends Schema.TaggedError<AIServiecOverloadedError>()(
   'AIServiecOverloadedError',
   {},
 ) {}
@@ -483,7 +493,7 @@ const gatherObjectVersions = (messages: Message[]): Map<ObjectId, ObjectVersion>
   for (const message of messages) {
     for (const block of message.content) {
       if (block.type === 'json' && block.disposition === VersionPin.DISPOSITION) {
-        const pin = VersionPin.pipe(S.decodeOption)(JSON.parse(block.json)).pipe(Option.getOrUndefined);
+        const pin = VersionPin.pipe(Schema.decodeOption)(JSON.parse(block.json)).pipe(Option.getOrUndefined);
         if (!pin) {
           continue;
         }

package/src/session/tools.ts

@@ -0,0 +1,117 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { Message, type Tool, ToolResult } from '@dxos/artifact';
+import { create } from '@dxos/echo-schema';
+import { invariant } from '@dxos/invariant';
+import { log } from '@dxos/log';
+
+export const isToolUse = (message: Message, { onlyToolNames }: { onlyToolNames?: string[] } = {}) => {
+  const block = message.content.at(-1);
+  invariant(block);
+  return block.type === 'tool_use' && (!onlyToolNames || onlyToolNames.includes(block.name));
+};
+
+export type RunToolsOptions = {
+  message: Message;
+  tools: Tool[];
+  extensions?: ToolContextExtensions;
+};
+
+export type RunToolsResult =
+  | {
+      type: 'continue';
+      message: Message;
+    }
+  | {
+      type: 'break';
+      result: unknown;
+    };
+
+export const runTools = async ({ message, tools, extensions }: RunToolsOptions): Promise<RunToolsResult> => {
+  const toolCalls = message.content.filter((block) => block.type === 'tool_use');
+  invariant(toolCalls.length === 1);
+  const toolCall = toolCalls[0];
+  const tool = tools.find((tool) => tool.name === toolCall.name);
+  if (!tool) {
+    const resultMessage = create(Message, {
+      role: 'user',
+      content: [
+        {
+          type: 'tool_result',
+          toolUseId: toolCall.id,
+          content: `Tool not found: ${toolCall.name}`,
+          isError: true,
+        },
+      ],
+    });
+
+    // TODO(burdon): Retry count?
+    return {
+      type: 'continue',
+      message: resultMessage,
+    };
+  }
+
+  let toolResult: ToolResult;
+  try {
+    invariant(tool.execute);
+    toolResult = await tool.execute(toolCall.input, { extensions });
+  } catch (error: any) {
+    log('tool error', { error });
+    toolResult = ToolResult.Error(error.message);
+  }
+
+  switch (toolResult.kind) {
+    case 'error': {
+      log('tool error', { message: toolResult.message });
+      const resultMessage = create(Message, {
+        role: 'user',
+        content: [
+          {
+            type: 'tool_result',
+            toolUseId: toolCall.id,
+            content: toolResult.message,
+            isError: true,
+          },
+        ],
+      });
+
+      // TODO(burdon): Retry count?
+      return {
+        type: 'continue',
+        message: resultMessage,
+      };
+    }
+
+    case 'success': {
+      log('tool success', { result: toolResult.result });
+      const resultMessage = create(Message, {
+        role: 'user',
+        content: [
+          {
+            type: 'tool_result',
+            toolUseId: toolCall.id,
+            content:
+              typeof toolResult.result === 'string' ? toolResult.result : JSON.stringify(toolResult.result) ?? '',
+          },
+          ...(toolResult.extractContentBlocks ?? []),
+        ],
+      });
+
+      return {
+        type: 'continue',
+        message: resultMessage,
+      };
+    }
+
+    case 'break': {
+      log('tool break', { result: toolResult.result });
+      return {
+        type: 'break',
+        result: toolResult.result,
+      };
+    }
+  }
+};

package/src/session/version-pin.ts

@@ -18,6 +18,12 @@ const VersionPinSchema = Schema.Struct({
   // TODO(dmaretskyi): Could be opaque
   version: ObjectVersionSchema,
 });
+
+/**
+ * Used to pin a specific version of an object during the course of a conversation.
+ *
+ * This allows us to detect when the object changes during the course of the conversation.
+ */
 export interface VersionPin extends Schema.Schema.Type<typeof VersionPinSchema> {}
 
 export const VersionPin: typeof VersionPinSchema & {

package/src/testing/artifacts.test.ts

@@ -3,7 +3,7 @@
 //
 
 /*
-import { Schema as S } from '@effect/schema';
+import { Schema } from '@effect/schema';
 import { describe, test } from 'vitest';
 
 import { ObjectId, toJsonSchema } from '@dxos/echo-schema';
@@ -15,26 +15,26 @@ import { AIServiceEdgeClient, Tool } from '../ai-service';
 import { defineTool, ToolResult } from '../conversation';
 
 // TODO(dmaretskyi): Effect schema.
-const ArtifactDef = S.Struct({
-  typename: S.String,
-  description: S.String,
-  schema: S.Any,
-  actions: S.Array(Tool),
-}).pipe(S.mutable);
+const ArtifactDef = Schema.Struct({
+  typename: Schema.String,
+  description: Schema.String,
+  schema: Schema.Any,
+  actions: Schema.Array(Tool),
+}).pipe(Schema.mutable);
 
-type ArtifactDef = S.Schema.Type<typeof ArtifactDef>;
+type ArtifactDef = Schema.Schema.Type<typeof ArtifactDef>;
 
-const ListItem = S.Struct({
-  name: S.String,
+const ListItem = Schema.Struct({
+  name: Schema.String,
 });
 
-const ListSchema = S.Struct({
+const ListSchema = Schema.Struct({
   id: ObjectId,
-  title: S.String,
-  items: S.Array(ListItem).pipe(S.mutable),
-}).pipe(S.mutable);
+  title: Schema.String,
+  items: Schema.Array(ListItem).pipe(Schema.mutable),
+}).pipe(Schema.mutable);
 
-type ListSchema = S.Schema.Type<typeof ListSchema>;
+type ListSchema = Schema.Schema.Type<typeof ListSchema>;
 
 const list: ArtifactDef = {
   typename: 'example.com/type/List',
@@ -44,7 +44,7 @@ const list: ArtifactDef = {
     defineTool('list', {
       name: 'query',
       description: 'Query all lists',
-      schema: S.Struct({}),
+      schema: Schema.Struct({}),
       execute: async ({}, context) => {
         return ToolResult.Success(LISTS[id]);
       },
@@ -52,7 +52,7 @@ const list: ArtifactDef = {
     defineTool('list', {
       name: 'inspect',
       description: 'Get contents of the list',
-      schema: S.Struct({
+      schema: Schema.Struct({
         id: ObjectId.annotations({ description: 'The list to inspect' }),
       }),
       execute: async ({ id }, context) => {
@@ -62,9 +62,9 @@ const list: ArtifactDef = {
     defineTool('list', {
       name: 'create',
       description: 'Add one or more items to an existing list',
-      schema: S.Struct({
+      schema: Schema.Struct({
         id: ObjectId.annotations({ description: 'The list to add items to' }),
-        items: S.Array(ListItem).annotations({ description: 'The items to add to the list' }),
+        items: Schema.Array(ListItem).annotations({ description: 'The items to add to the list' }),
       }),
       execute: async ({ id, items }, context) => {
         LISTS[id].items.push(...items);
@@ -81,7 +81,7 @@ const artifacts = [list];
 const getArtifactDefinitions = defineTool('system', {
   name: 'getArtifactDefinitions',
   description: 'Queries the definitions and metadata of artifacts defined in the system',
-  schema: S.Any,
+  schema: Schema.Any,
   execute: async (_, context) => {
     return ToolResult.Success(
       artifacts.map((artifact) => ({
@@ -116,8 +116,8 @@ describe('Artifacts', () => {
       name: 'custodian',
       description: 'Custodian can tell you the password if you say the magic word',
       parameters: toJsonSchema(
-        S.Struct({
-          magicWord: S.String.annotations({ description: 'The magic word. Should be exactly "pretty please"' }),
+        Schema.Struct({
+          magicWord: Schema.String.annotations({ description: 'The magic word. Should be exactly "pretty please"' }),
         }),
       ),
     };

package/src/testing/index.ts

@@ -8,3 +8,4 @@ export * from './ollama';
 export * from './query-promts';
 export * from './test-data';
 export * from './test-schema';
+export * from './printer';

package/src/testing/logger.ts

@@ -4,7 +4,7 @@
 
 import type { ImageSource, MessageContentBlock } from '@dxos/artifact';
 
-import type { ConversationEvent } from '../conversation';
+import type { ConversationEvent } from '../experimental';
 
 export const createLogger = ({
   stream,

package/src/testing/printer.ts

@@ -0,0 +1,41 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { inspect } from 'node:util';
+
+import type { Message, MessageContentBlock } from '@dxos/artifact';
+
+export class ConsolePrinter {
+  printMessage = (message: Message) => {
+    // eslint-disable-next-line no-console
+    console.log(`${message.role.toUpperCase()}\n`);
+    for (const content of message.content) {
+      this.printContentBlock(content);
+    }
+  };
+
+  printContentBlock = (content: MessageContentBlock) => {
+    switch (content.type) {
+      case 'text':
+        // eslint-disable-next-line no-console
+        console.log(content.text);
+        break;
+      case 'tool_use':
+        // eslint-disable-next-line no-console
+        console.log(`⚙️ [Tool Use] ${content.name} ${inspect(content.input, { depth: null, colors: true })}`);
+        break;
+      case 'tool_result': {
+        let data: any;
+        try {
+          data = JSON.parse(content.content);
+        } catch {
+          data = content.content;
+        }
+        // eslint-disable-next-line no-console
+        console.log(`⚙️ [Tool Result] ${content.toolUseId} ${inspect(data, { depth: null, colors: true })}`);
+        break;
+      }
+    }
+  };
+}

package/src/testing/query-promts.ts

@@ -2,23 +2,23 @@
 // Copyright 2024 DXOS.org
 //
 
-import { Schema as S } from 'effect';
+import { Schema } from 'effect';
 
 import { defineTool, ToolResult } from '@dxos/artifact';
 import { toJsonSchema } from '@dxos/echo-schema';
 import { log } from '@dxos/log';
 
-import { executeQuery, formatJsonSchemaForLLM, type DataSource } from '../cypher';
+import { executeQuery, formatJsonSchemaForLLM, type DataSource } from '../experimental';
 import { trim } from '../util';
 
-export const createCypherTool = (dataSource: DataSource, schemaTypes: S.Schema.Any[] = []) =>
+export const createCypherTool = (dataSource: DataSource, schemaTypes: Schema.Schema.Any[] = []) =>
   defineTool('dxos.org/echo', {
     name: 'query',
     description:
       'Query the ECHO graph database in cypher query language. Returns data from the database.' +
       (schemaTypes.length > 0 ? `\n\n${createSystemPrompt(schemaTypes)}` : ''),
-    schema: S.Struct({
-      query: S.String.annotations({
+    schema: Schema.Struct({
+      query: Schema.String.annotations({
         description: `
           A valid cypher query string to execute.
           Query must have one MATCH clause.
@@ -55,7 +55,7 @@ export const createCypherTool = (dataSource: DataSource, schemaTypes: S.Schema.A
     },
   });
 
-export const createSystemPrompt = (schemaTypes: S.Schema.Any[]) => trim`
+export const createSystemPrompt = (schemaTypes: Schema.Schema.Any[]) => trim`
   You have access to ECHO graph database. You can query the database to get data to satisfy user prompts.
 
   Database schema is defined in pseud-cypher syntax. The schema is:

package/src/testing/query.test.ts

@@ -16,8 +16,7 @@ import { createCypherTool, createSystemPrompt } from './query-promts';
 import { createTestData, seedTestData } from './test-data';
 import { Contact, Organization, Project, Task } from './test-schema';
 import { AIServiceEdgeClient, DEFAULT_EDGE_MODEL } from '../ai-service';
-import { runLLM } from '../conversation';
-import { EchoDataSource } from '../cypher';
+import { runLLM, EchoDataSource } from '../experimental';
 
 const client = new AIServiceEdgeClient({
   endpoint: AI_SERVICE_ENDPOINT.LOCAL,

package/src/testing/test-data.ts

@@ -2,7 +2,7 @@
 // Copyright 2024 DXOS.org
 //
 
-import { type Schema as S } from 'effect';
+import { type Schema } from 'effect';
 
 import { raise } from '@dxos/debug';
 import type { EchoDatabase } from '@dxos/echo-db';
@@ -23,7 +23,7 @@ import {
   type Relationship,
   formatInferredRelationshipLabel,
   formatNodeLabel,
-} from '../cypher';
+} from '../experimental';
 import { Contact, Organization, Project, Task } from '../testing';
 
 export const seedTestData = (db: EchoDatabase) => {
@@ -254,7 +254,7 @@ export class MockDataSource implements DataSource {
   objects: Record<
     string,
     {
-      schema: S.Schema.All;
+      schema: Schema.Schema.All;
       typeDxn: string;
       id: string;
       data: unknown;
@@ -265,7 +265,7 @@ export class MockDataSource implements DataSource {
     string,
     {
       json: JsonSchemaType;
-      instance: S.Schema.All;
+      instance: Schema.Schema.All;
     }
   > = {};
 
@@ -280,7 +280,7 @@ export class MockDataSource implements DataSource {
     return this.relationships.filter((relationship) => !label || relationship.label === label);
   }
 
-  add<S extends S.Schema.All>(schema: S, data: S.Schema.Type<S>) {
+  add<S extends Schema.Schema.All>(schema: S, data: Schema.Schema.Type<S>) {
     invariant(typeof data.id === 'string', 'Data must have an id');
     const jsonSchema = toJsonSchema(schema);
     if (!this.schema[jsonSchema.$id!]) {