@dxos/assistant 0.8.2-main.f11618f → 0.8.2-main.fbd8ed0

package/src/experimental/research/research.test.ts

@@ -0,0 +1,263 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Option, Schema, SchemaAST } from 'effect';
+import { describe, test } from 'vitest';
+
+import { structuredOutputParser } from '@dxos/artifact';
+import { isEncodedReference } from '@dxos/echo-protocol';
+import { create, getSchemaDXN, ObjectId, ReferenceAnnotationId } from '@dxos/echo-schema';
+import { mapAst } from '@dxos/effect';
+import { log } from '@dxos/log';
+import { DataType } from '@dxos/schema';
+import { deepMapValues } from '@dxos/util';
+
+import { createExaTool } from './exa';
+import INSTRUCTIONS from './instructions.tpl?raw';
+import { AIServiceEdgeClient, OllamaClient } from '../../ai-service';
+import { AISession } from '../../session';
+import { AI_SERVICE_ENDPOINT, ConsolePrinter } from '../../testing';
+
+const EXA_API_KEY = '9c7e17ff-0c85-4cd5-827a-8b489f139e03';
+const REMOTE_AI = true;
+
+const aiService = REMOTE_AI
+  ? new AIServiceEdgeClient({
+      endpoint: AI_SERVICE_ENDPOINT.REMOTE,
+    })
+  : new OllamaClient({
+      overrides: {
+        model: 'llama3.1:8b',
+      },
+    });
+
+const TYPES = [DataType.Event, DataType.Organization, DataType.Person, DataType.Project, DataType.Task, DataType.Text];
+
+describe.skip('Research', () => {
+  test('should generate a research report', { timeout: 1000000 }, async () => {
+    const searchTool = createExaTool({ apiKey: EXA_API_KEY });
+
+    const session = new AISession({ operationModel: 'configured' });
+
+    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));
+
+    // TODO(dmaretskyi): Consider adding this pattern as the "Graph" output mode for the session.
+    const outputSchema = createExtractionSchema(TYPES);
+    const result = await session.runStructured(outputSchema, {
+      client: aiService,
+      systemPrompt: INSTRUCTIONS,
+      artifacts: [],
+      tools: [searchTool],
+      generationOptions: {
+        model: '@anthropic/claude-3-5-sonnet-20241022',
+      },
+      history: [],
+      prompt: 'Find projects that are in the space of AI and personal knowledge management',
+    });
+    const data = sanitizeObjects(TYPES, result as any);
+
+    log.info('result', { data });
+  });
+});
+
+describe.skip('misc', () => {
+  test('createExtractionSchema', () => {
+    const schema = createExtractionSchema(TYPES);
+    log.info('schema', { schema });
+  });
+
+  test('extract schema json schema', () => {
+    const schema = createExtractionSchema(TYPES);
+    const parser = structuredOutputParser(schema);
+    log.info('schema', { json: parser.tool.parameters });
+  });
+
+  test('getSanitizedSchemaName', () => {
+    const names = TYPES.map(getSanitizedSchemaName);
+    log.info('names', { names });
+  });
+
+  test('getTypeAnnotation', () => {
+    for (const schema of TYPES) {
+      const dxn = getSchemaDXN(schema);
+      log.info('dxn', { schema, dxn });
+    }
+  });
+
+  test('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 = sanitizeObjects(TYPES, TEST_DATA);
+    log.info('data', { data });
+  });
+});
+
+const createExtractionSchema = (types: Schema.Schema.AnyNoContext[]) => {
+  return Schema.Struct({
+    ...Object.fromEntries(
+      types.map(preprocessSchema).map((schema, index) => [
+        `objects_${getSanitizedSchemaName(types[index])}`,
+        Schema.optional(Schema.Array(schema)).annotations({
+          description: `The objects to answer the query of type: ${getSchemaDXN(types[index])?.asTypeDXN()!.type}`,
+        }),
+      ]),
+    ),
+  });
+};
+
+const getSanitizedSchemaName = (schema: Schema.Schema.AnyNoContext) => {
+  return getSchemaDXN(schema)!
+    .asTypeDXN()!
+    .type.replaceAll(/[^a-zA-Z0-9]+/g, '_');
+};
+
+const sanitizeObjects = (types: Schema.Schema.AnyNoContext[], data: Record<string, readonly unknown[]>) => {
+  const entries = types
+    .map(
+      (type) =>
+        data[`objects_${getSanitizedSchemaName(type)}`]?.map((object: any) => ({
+          data: object,
+          schema: type,
+        })) ?? [],
+    )
+    .flat();
+
+  const idMap = new Map<string, string>();
+
+  return entries
+    .map((entry) => {
+      idMap.set(entry.data.id, ObjectId.random());
+      entry.data.id = idMap.get(entry.data.id);
+      return entry;
+    })
+    .map((entry) => {
+      const data = deepMapValues(entry.data, (value, recurse) => {
+        if (isEncodedReference(value)) {
+          const ref = value['/'];
+          if (idMap.has(ref)) {
+            // TODO(dmaretskyi): Whats the best way to represent a local url.
+            return { '/': `dxn:echo:@:${idMap.get(ref)}` };
+          } else {
+            // Search URIs?
+            return { '/': `search:?q=${encodeURIComponent(ref)}` };
+          }
+        }
+
+        return recurse(value);
+      });
+
+      return create(entry.schema, data);
+    });
+};
+
+const SoftRef = Schema.Struct({
+  '/': Schema.String,
+}).annotations({
+  description: 'Reference to another object.',
+});
+
+const preprocessSchema = (schema: Schema.Schema.AnyNoContext) => {
+  const go = (ast: SchemaAST.AST): SchemaAST.AST => {
+    if (SchemaAST.getAnnotation(ast, ReferenceAnnotationId).pipe(Option.isSome)) {
+      return SoftRef.ast;
+    }
+
+    return mapAst(ast, go);
+  };
+
+  const processed = Schema.make<any, any, never>(mapAst(schema.ast, go));
+  return Schema.extend(
+    processed.pipe(Schema.omit('id')),
+    Schema.Struct({
+      id: Schema.String.annotations({
+        description: 'The id of this object. Come up with a unique id based on your judgement.',
+      }),
+    }),
+  );
+};

package/src/index.ts

@@ -3,6 +3,7 @@
 //
 
 export * from './ai-service';
-export * from './cypher';
-export * from './conversation';
 export * from './session';
+
+// TODO(dmaretskyi): Separate entrypoint for experimental features.
+export * from './experimental';

package/src/session/index.ts

@@ -4,3 +4,4 @@
 
 export * from './session';
 export * from './version-pin';
+export * from './tools';

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.
@@ -230,7 +231,7 @@ export class AISession {
     } catch (err) {
       log.catch(err);
       if (err instanceof Error && err.message.includes('Overloaded')) {
-        error = new AIServiecOverloadedError('AI service overloaded', { cause: err });
+        error = new AIServiecOverloadedError('AI service overloaded', { cause: err } as any); // TODO(dmaretskyi): Schema errors cant have cause
       } else {
         error = new Error('AI service error', { cause: err });
       }
@@ -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 & {