@dxos/assistant 0.7.4 → 0.7.5-labs.5f04cf6

package/src/conversation/conversation.ts

@@ -2,13 +2,13 @@
 // Copyright 2024 DXOS.org
 //
 
+import { type Tool, ToolResult, Message } from '@dxos/artifact';
+import { createStatic, ObjectId } from '@dxos/echo-schema';
 import { invariant } from '@dxos/invariant';
 import { type SpaceId } from '@dxos/keys';
 import { log } from '@dxos/log';
 
-import { type LLMToolDefinition } from './types';
-import { ObjectId, type LLMTool } from '../ai-service';
-import { type AIServiceClient, type Message, type ResultStreamEvent, type LLMModel } from '../ai-service';
+import { type AIServiceClient, type ResultStreamEvent, type LLMModel } from '../ai-service';
 
 export type CreateLLMConversationParams = {
   model: LLMModel;
@@ -18,11 +18,11 @@ export type CreateLLMConversationParams = {
    */
   system?: string;
 
-  spaceId: SpaceId;
-  threadId: ObjectId;
-
-  tools: LLMTool[];
+  spaceId?: SpaceId;
+  threadId?: ObjectId;
 
+  tools: Tool[];
+  history?: Message[];
   client: AIServiceClient;
 
   logger?: (event: ConversationEvent) => void;
@@ -37,45 +37,50 @@ export type ConversationEvent =
 
 export const runLLM = async (params: CreateLLMConversationParams) => {
   let conversationResult: any = null;
+  const history = params.history ?? [];
 
   const generate = async () => {
     log('llm generate', { tools: params.tools });
     const beginTs = Date.now();
-    const result = await params.client.generate({
+    const stream = await params.client.generate({
       model: params.model,
       spaceId: params.spaceId,
       threadId: params.threadId,
+      history,
       systemPrompt: params.system,
       tools: params.tools as any,
     });
 
-    for await (const event of result) {
+    for await (const event of stream) {
       params.logger?.(event);
     }
-    const messages = await result.complete();
+    const messages = await stream.complete();
     const message = messages.at(-1);
 
     log('llm result', { time: Date.now() - beginTs, message });
     invariant(message);
     params.logger?.({ type: 'message', message });
-    await params.client.insertMessages(
-      messages.map((msg) => ({ ...msg, content: msg.content.filter((c) => c.type !== 'image') })),
+    history.push(
+      ...messages.map(
+        (message): Message => ({ ...message, content: message.content.filter((block) => block.type !== 'image') }),
+      ),
     );
 
     const isToolUse = message.content.at(-1)?.type === 'tool_use';
     if (isToolUse) {
-      const toolCalls = message.content.filter((c) => c.type === 'tool_use');
+      const toolCalls = message.content.filter((block) => block.type === 'tool_use');
       invariant(toolCalls.length === 1);
       const toolCall = toolCalls[0];
-      const tool = params.tools.find((t) => t.name === toolCall.name);
+      const tool = params.tools.find((tool) => tool.name === toolCall.name);
       if (!tool) {
         throw new Error(`Tool not found: ${toolCall.name}`);
       }
 
-      const toolResult = await tool.execute(toolCall.input, {});
+      invariant(tool.execute);
+      const toolResult = await tool.execute(toolCall.input);
       switch (toolResult.kind) {
         case 'error': {
-          log('tool error', { message: toolResult.message });
+          log.warn('tool error', { message: toolResult.message });
           const resultMessage: Message = {
             id: ObjectId.random(),
             spaceId: params.spaceId,
@@ -91,10 +96,11 @@ export const runLLM = async (params: CreateLLMConversationParams) => {
             ],
           };
           params.logger?.({ type: 'message', message: resultMessage });
-          await params.client.insertMessages([resultMessage]);
+          history.push(resultMessage);
 
           return true;
         }
+
         case 'success': {
           log('tool success', { result: toolResult.result });
           const resultMessage: Message = {
@@ -111,10 +117,11 @@ export const runLLM = async (params: CreateLLMConversationParams) => {
             ],
           };
           params.logger?.({ type: 'message', message: resultMessage });
-          await params.client.insertMessages([resultMessage]);
+          history.push(resultMessage);
 
           return true;
         }
+
         case 'break': {
           log('tool break', { result: toolResult.result });
           conversationResult = toolResult.result;
@@ -130,5 +137,93 @@ export const runLLM = async (params: CreateLLMConversationParams) => {
 
   return {
     result: conversationResult,
+    history,
   };
 };
+
+export const isToolUse = (message: Message) => message.content.at(-1)?.type === 'tool_use';
+
+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) {
+    throw new Error(`Tool not found: ${toolCall.name}`);
+  }
+
+  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 = createStatic(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 = createStatic(Message, {
+        role: 'user',
+        content: [
+          {
+            type: 'tool_result',
+            toolUseId: toolCall.id,
+            content: typeof toolResult.result === 'string' ? toolResult.result : JSON.stringify(toolResult.result),
+          },
+        ],
+      });
+
+      return {
+        type: 'continue',
+        message: resultMessage,
+      };
+    }
+
+    case 'break': {
+      log('tool break', { result: toolResult.result });
+      return {
+        type: 'break',
+        result: toolResult.result,
+      };
+    }
+  }
+};

package/src/conversation/index.ts

@@ -3,4 +3,3 @@
 //
 
 export * from './conversation';
-export * from './types';

package/src/conversation/types.ts

@@ -4,13 +4,11 @@
 
 import { Schema as S } from '@effect/schema';
 
-import { toJsonSchema } from '@dxos/echo-schema';
-import { JsonSchemaType } from '@dxos/echo-schema';
+import { type Message, type MessageContentBlock } from '@dxos/artifact';
+import { toJsonSchema, JsonSchemaType, ObjectId } from '@dxos/echo-schema';
 import type { SpaceId } from '@dxos/keys';
 import { log } from '@dxos/log';
 
-import { ObjectId, type Message } from '../ai-service';
-
 export const createUserMessage = (spaceId: SpaceId, threadId: ObjectId, text: string): Message => ({
   id: ObjectId.random(),
   spaceId,
@@ -19,46 +17,80 @@ export const createUserMessage = (spaceId: SpaceId, threadId: ObjectId, text: st
   content: [{ type: 'text', text }],
 });
 
-export const LLMToolDefinition = S.Struct({
+/**
+ * @deprecated Use {@link LLMTool} instead.
+ */
+export const ToolDefinition = S.Struct({
   name: S.String,
   description: S.String,
   parameters: JsonSchemaType,
   execute: S.Any,
 });
 
-export type LLMToolDefinition = {
+/**
+ * @deprecated Use {@link LLMTool} instead.
+ */
+export type ToolDefinition = {
   name: string;
   description: string;
   parameters: JsonSchemaType;
-  execute: (params: unknown, context: ToolExecutionContext) => Promise<LLMToolResult>;
+  execute: (params: unknown, context: ToolExecutionContext) => Promise<ToolResult>;
 };
 
-export type ToolExecutionContext = {};
+declare global {
+  /**
+   * Extensions to the tool execution context.
+   *
+   * Different modules can extend this definition to add their own properties.
+   *
+   * @example
+   * ```ts
+   * declare global {
+   *   interface ToolContextExtensions {
+   *     chess: {
+   *       board: string;
+   *     };
+   *   }
+   * }
+   * ```
+   *
+   * @see https://www.typescriptlang.org/docs/handbook/declaration-merging.html#global-augmentation
+   */
+  interface ToolContextExtensions {}
+}
+
+export type ToolExecutionContext = {
+  extensions: ToolContextExtensions;
+};
 
-export type LLMToolResult =
-  | { kind: 'success'; result: unknown }
+export type ToolResult =
+  | { kind: 'success'; result: unknown; extractContentBlocks?: MessageContentBlock[] }
   | { kind: 'error'; message: string }
   | { kind: 'break'; result: unknown };
 
-export const LLMToolResult = Object.freeze({
+export const ToolResult = Object.freeze({
   /**
    * The tool execution was successful.
    * Gives the result back to the LLM.
    */
-  Success: (result: unknown): LLMToolResult => ({ kind: 'success', result }),
+  Success: (result: unknown, extractContentBlocks?: MessageContentBlock[]): ToolResult => ({
+    kind: 'success',
+    result,
+    extractContentBlocks,
+  }),
 
-  Error: (message: string): LLMToolResult => ({ kind: 'error', message }),
+  Error: (message: string): ToolResult => ({ kind: 'error', message }),
   /**
    * Stop the conversation and return the result.
    */
-  Break: (result: unknown): LLMToolResult => ({ kind: 'break', result }),
+  Break: (result: unknown): ToolResult => ({ kind: 'break', result }),
 });
 
 export type DefineToolParams<Params extends S.Schema.AnyNoContext> = {
   name: string;
   description: string;
   schema: Params;
-  execute: (params: S.Schema.Type<Params>, context: ToolExecutionContext) => Promise<LLMToolResult>;
+  execute: (params: S.Schema.Type<Params>, context: ToolExecutionContext) => Promise<ToolResult>;
 };
 
 export const defineTool = <Params extends S.Schema.AnyNoContext>({
@@ -66,7 +98,7 @@ export const defineTool = <Params extends S.Schema.AnyNoContext>({
   description,
   schema,
   execute,
-}: DefineToolParams<Params>): LLMToolDefinition => {
+}: DefineToolParams<Params>): ToolDefinition => {
   return {
     name,
     description,
@@ -78,9 +110,16 @@ export const defineTool = <Params extends S.Schema.AnyNoContext>({
   };
 };
 
+/**
+ * Adapts schems to be able to pass to AI providers.
+ */
 const toFunctionParameterSchema = (schema: S.Schema.All) => {
   const jsonSchema = toJsonSchema(schema);
   log('tool schema', { jsonSchema });
+  delete jsonSchema.anyOf;
+  delete jsonSchema.$id;
+  jsonSchema.type = 'object';
+  jsonSchema.properties ??= {};
   for (const key in jsonSchema.properties) {
     if (typeof jsonSchema.properties![key].description !== 'string') {
       throw new Error(`Property missing description: ${key}`);

package/src/cypher/combinators.test.ts

@@ -6,7 +6,7 @@
 /* eslint-disable @typescript-eslint/consistent-type-imports */
 
 import { createRequire } from 'node:module';
-import type { Parjser } from 'parjs';
+import { type Parjser } from 'parjs';
 import { describe, test } from 'vitest';
 
 import { log } from '@dxos/log';
@@ -58,13 +58,13 @@ describe('Cypher Combinators', () => {
     const parser = or(node)(relation);
     const input = '[a:b]';
     const result = parser.parse(input);
-    log.info('', { result });
+    log('result', { result });
   });
 
   test('properties', ({ expect }) => {
     const input = "{name: 'John'}";
     const result = properties.parse(input);
-    log.info('', { result });
+    log('result', { result });
     expect(result).toMatchObject({ isOk: true });
   });
 
@@ -92,7 +92,7 @@ describe('Cypher Combinators', () => {
     const cases = ['RETURN n, m'];
     for (const input of cases) {
       const result = returnClause.parse(input);
-      log.info('', { result });
+      log('result', { result });
       expect(result).toMatchObject({ isOk: true });
     }
   });

package/src/cypher/combinators.ts

@@ -21,7 +21,6 @@ import type {
   StringLiteral,
   WhereClause,
 } from './ast';
-
 import { Parjser, parjs, parjsCombinators } from './parjs';
 
 // Helper parsers

package/src/cypher/execution-plan.test.ts

@@ -9,10 +9,12 @@ import { log } from '@dxos/log';
 import { createExecutionPlan } from './execution-plan';
 import { parseCypherQuery } from './parser';
 
+// log.config({ filter: 'debug' });
+
 describe('Execution Plan', () => {
   test('trivial', ({ expect }) => {
     const plan = createExecutionPlan(parseCypherQuery('MATCH (n:Person) RETURN n'));
-    log.info('', { plan });
+    log('', { plan });
     expect(plan.steps.map((step) => step.type)).toEqual(['NodeScan', 'ProduceResults']);
   });
 
@@ -22,9 +24,9 @@ describe('Execution Plan', () => {
         "MATCH (org:Org {name: 'DXOS'})-[:ORG_EMPLOYEES]->(c:Contact)<-[:TASK_ASSIGNEE]-(:Task)-[:TASK_PROJECT]->(p:Project {name: 'Composer'}) RETURN c.name",
       ),
     );
-    log.info('', { plan });
+    log('', { plan });
     for (const step of plan.steps) {
-      log.info('', { step });
+      log('step', { step });
     }
     expect(plan.steps.map((step) => step.type)).toEqual([
       'NodeScan',

package/src/cypher/parjs-shims.d.ts

@@ -1,3 +1,7 @@
+//
+// Copyright 2025 DXOS.org
+//
+
 declare module '#parjs' {
   import * as parjs from 'parjs';
   import * as parjsCombinators from 'parjs/combinators';

package/src/cypher/parjs.ts

@@ -3,8 +3,8 @@
 //
 
 // Workaround for https://github.com/vitest-dev/vitest/issues/7174.
-import { parjsCombinators, parjs } from '#parjs';
-
 import type { Parjser } from 'parjs';
 
+import { parjsCombinators, parjs } from '#parjs';
+
 export { type Parjser, parjsCombinators, parjs };

package/src/cypher/query-executor.test.ts

@@ -9,7 +9,7 @@ import { log } from '@dxos/log';
 import { executeQuery } from './query-executor';
 import { createTestData } from '../testing';
 
-test('run', async ({ expect }) => {
+test.skip('run', async ({ expect }) => {
   const dataSource = createTestData();
   // for (const node of dataSource.nodes) {
   //   log.info('node', { id: node.id, label: node.label, name: node.properties.name });

package/src/cypher/schema.ts

@@ -4,7 +4,6 @@
 
 import { raise } from '@dxos/debug';
 import { JSON_SCHEMA_ECHO_REF_ID, type JsonSchemaType } from '@dxos/echo-schema';
-import { DXN } from '@dxos/keys';
 
 import { trim } from '../util';
 

package/src/index.ts

@@ -4,4 +4,4 @@
 
 export * from './ai-service';
 export * from './cypher';
-export * from './conversation';
+export * from './conversation';

package/src/testing/artifacts.test.ts

@@ -0,0 +1,182 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+/*
+import { Schema as S } from '@effect/schema';
+import { describe, test } from 'vitest';
+
+import { ObjectId, toJsonSchema } from '@dxos/echo-schema';
+import { invariant } from '@dxos/invariant';
+import { SpaceId } from '@dxos/keys';
+import { log } from '@dxos/log';
+
+import { AIServiceClientImpl, 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);
+
+type ArtifactDef = S.Schema.Type<typeof ArtifactDef>;
+
+const ListItem = S.Struct({
+  name: S.String,
+});
+
+const ListSchema = S.Struct({
+  id: ObjectId,
+  title: S.String,
+  items: S.Array(ListItem).pipe(S.mutable),
+}).pipe(S.mutable);
+
+type ListSchema = S.Schema.Type<typeof ListSchema>;
+
+const list: ArtifactDef = {
+  typename: 'example.com/type/List',
+  description: 'Ordered set of text items.',
+  schema: ListSchema,
+  actions: [
+    defineTool({
+      name: 'list.query',
+      description: 'Query all lists',
+      schema: S.Struct({}),
+      execute: async ({}, context) => {
+        return ToolResult.Success(LISTS[id]);
+      },
+    }),
+    defineTool({
+      name: 'list.inspect',
+      description: 'Get contents of the list',
+      schema: S.Struct({
+        id: ObjectId.annotations({ description: 'The list to inspect' }),
+      }),
+      execute: async ({ id }, context) => {
+        return ToolResult.Success(LISTS[id]);
+      },
+    }),
+    defineTool({
+      name: 'list.add',
+      description: 'Add one or more items to an existing list',
+      schema: S.Struct({
+        id: ObjectId.annotations({ description: 'The list to add items to' }),
+        items: S.Array(ListItem).annotations({ description: 'The items to add to the list' }),
+      }),
+      execute: async ({ id, items }, context) => {
+        LISTS[id].items.push(...items);
+        return ToolResult.Success(LISTS[id]);
+      },
+    }),
+  ],
+};
+
+const LISTS: Record<ObjectId, ListSchema> = {};
+
+const artifacts = [list];
+
+const getArtifactDefinitions = defineTool({
+  name: 'getArtifactDefinitions',
+  description: 'Queries the definitions and metadata of artifacts defined in the system',
+  schema: S.Any,
+  execute: async (_, context) => {
+    return ToolResult.Success(
+      artifacts.map((artifact) => ({
+        typename: artifact.typename,
+        description: artifact.description,
+        schema: toJsonSchema(artifact.schema),
+        actions: artifact.actions.map((action) => ({
+          name: action.name,
+          description: action.description,
+        })),
+      })),
+    );
+  },
+});
+
+const ARTIFACT_INSTRUCTIONS = `
+You have ability to interact with artifacts.
+Artifacts are persistent structured objects that have actions associated with them.
+
+Before replying to the user you should first query the artifact definitions and metadata.
+If there's
+
+`;
+
+describe('Artifacts', () => {
+  test('shopping list', async () => {
+    const client = new AIServiceClientImpl({
+      endpoint: ENDPOINT,
+    });
+
+    const custodian: Tool = {
+      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"' }),
+        }),
+      ),
+    };
+
+    const spaceId = SpaceId.random();
+    const threadId = ObjectId.random();
+
+    await client.appendMessages([
+      {
+        id: ObjectId.random(),
+        spaceId,
+        threadId,
+        role: 'user',
+        content: [{ type: 'text', text: 'What is the password? Ask the custodian' }],
+      },
+    ]);
+
+    const stream = await client.generate({
+      model: '@anthropic/claude-3-5-haiku-20241022',
+      spaceId,
+      threadId,
+      systemPrompt: 'You are a helpful assistant.',
+      tools: [custodian],
+    });
+    for await (const event of stream) {
+      log.info('event', event);
+    }
+    const [message] = await stream.complete();
+    log.info('full message', {
+      message,
+    });
+    await client.appendMessages([message]);
+
+    const toolUse = message.content.find(({ type }) => type === 'tool_use')!;
+    invariant(toolUse.type === 'tool_use');
+    await client.appendMessages([
+      {
+        id: ObjectId.random(),
+        spaceId,
+        threadId,
+        role: 'user',
+        content: [{ type: 'tool_result', toolUseId: toolUse.id, content: 'password="The sky is gray"' }],
+      },
+    ]);
+
+    const stream2 = await client.generate({
+      model: '@anthropic/claude-3-5-haiku-20241022',
+      spaceId,
+      threadId,
+      systemPrompt: 'You are a helpful assistant.',
+      tools: [custodian],
+    });
+    for await (const event of stream2) {
+      log.info('event', event);
+    }
+    const [message2] = await stream2.complete();
+    log.info('full message', {
+      message: message2,
+    });
+  });
+});
+*/

package/src/testing/data.ts

@@ -0,0 +1,13 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export const DATA_MESSAGES = [
+  `
+  <cot>
+  ${Array.from({ length: 5 })
+    .map((_, i) => `${i + 1}. Step`)
+    .join('\n')}
+  </cot>
+  `,
+];

package/src/testing/defs.ts

@@ -0,0 +1,8 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export const AI_SERVICE_ENDPOINT = {
+  LOCAL: 'http://localhost:8788',
+  REMOTE: 'https://ai-service.dxos.workers.dev',
+};

package/src/testing/index.ts

@@ -2,6 +2,7 @@
 // Copyright 2024 DXOS.org
 //
 
+export * from './defs';
 export * from './logger';
 export * from './query-promts';
 export * from './test-data';