@dxos/assistant 0.7.4-df37807 → 0.7.5-feature-compute.4d9d99a

package/src/ai-service/types/message.ts

@@ -0,0 +1,147 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { Schema as S } from '@effect/schema';
+
+import { EchoObject, ObjectId } from '@dxos/echo-schema';
+import { SpaceId } from '@dxos/keys';
+
+// TODO(dmaretskyi): Dedupe package with dxos/edge.
+// TODO(dmaretskyi): Extract IDs to protocols.
+
+export const SpaceIdSchema: S.Schema<SpaceId, string> = S.String.pipe(S.filter(SpaceId.isValid));
+
+export const Space = S.Struct({
+  id: SpaceIdSchema,
+});
+
+export interface Space extends S.Schema.Type<typeof Space> {}
+
+/**
+ * Thread.
+ */
+export const Thread = S.Struct({
+  id: ObjectId,
+  spaceId: SpaceIdSchema,
+});
+
+export interface Thread extends S.Schema.Type<typeof Thread> {}
+
+export const MessageRole = S.String.pipe(S.filter((role) => role === 'user' || role === 'assistant'));
+export type MessageRole = S.Schema.Type<typeof MessageRole>;
+
+/**
+ * Tagged block.
+ */
+export const TagContentBlock = S.Struct({
+  type: S.Literal('tag'),
+  tag: S.String,
+  text: S.String,
+}).pipe(S.mutable);
+
+export type TagContentBlock = S.Schema.Type<typeof TagContentBlock>;
+
+/**
+ * Text.
+ */
+export const TextContentBlock = S.Struct({
+  type: S.Literal('text'),
+  text: S.String,
+}).pipe(S.mutable);
+
+export type TextContentBlock = S.Schema.Type<typeof TextContentBlock>;
+
+/**
+ * Image.
+ */
+export const ImageSource = S.Struct({
+  type: S.Literal('base64'),
+  media_type: S.Literal('image/jpeg', 'image/png', 'image/gif', 'image/webp'),
+  data: S.String,
+}).pipe(S.mutable);
+
+export type ImageSource = S.Schema.Type<typeof ImageSource>;
+
+export const ImageContentBlock = S.Struct({
+  type: S.Literal('image'),
+  id: S.optional(S.String),
+  source: S.optional(ImageSource),
+}).pipe(S.mutable);
+
+export type ImageContentBlock = S.Schema.Type<typeof ImageContentBlock>;
+
+/**
+ * Tool use.
+ */
+export const ToolUseContentBlock = S.Struct({
+  type: S.Literal('tool_use'),
+
+  /**
+   * Opaque service-defined ID of the tool invocation.
+   * Used to match the tool result block.
+   */
+  id: S.String,
+
+  /**
+   * Tool name.
+   */
+  name: S.String,
+
+  input: S.Unknown,
+
+  /**
+   * Used to accumulate the partial tool input JSON in streaming mode.
+   */
+  // TODO(burdon): Remove?
+  inputJson: S.optional(S.String),
+}).pipe(S.mutable);
+
+export const ToolResultContentBlock = S.Struct({
+  type: S.Literal('tool_result'),
+  toolUseId: S.String,
+  content: S.String,
+  isError: S.optional(S.Boolean),
+});
+
+/**
+ * Content union.
+ */
+export const MessageContentBlock = S.Union(
+  TagContentBlock,
+  TextContentBlock,
+  ImageContentBlock,
+  ToolUseContentBlock,
+  ToolResultContentBlock,
+);
+
+export type MessageContentBlock = S.Schema.Type<typeof MessageContentBlock>;
+
+/**
+ * Message.
+ */
+const MessageSchema = S.Struct({
+  id: ObjectId,
+  threadId: S.optional(ObjectId),
+  spaceId: S.optional(SpaceIdSchema),
+
+  /**
+   * ID of the message from the foreign provider.
+   */
+  // TODO(dmaretskyi): Should be in meta/keys.
+  foreignId: S.optional(S.String),
+
+  // TODO(dmaretskyi): Figure out how to deal with those.
+  // created: S.optional(S.DateFromString),
+  // updated: S.optional(S.DateFromString),
+
+  role: MessageRole,
+
+  // TODO(burdon): Rename blocks?
+  content: S.Array(MessageContentBlock).pipe(S.mutable),
+});
+
+// TODO(burdon): Reconcile with Chat/Message types.
+export const Message = MessageSchema.pipe(EchoObject('dxos.org/type/Message', '0.1.0'));
+
+export type Message = S.Schema.Type<typeof Message>;

package/src/ai-service/types/schema.ts

@@ -0,0 +1,128 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { Schema as S } from '@effect/schema';
+
+import { JsonSchemaType, ObjectId } from '@dxos/echo-schema';
+
+import { Message, type MessageContentBlock, SpaceIdSchema } from './message';
+
+export const LLMModel = S.Literal(
+  '@hf/nousresearch/hermes-2-pro-mistral-7b',
+  '@anthropic/claude-3-5-sonnet-20241022',
+  '@anthropic/claude-3-5-haiku-20241022',
+  '@ollama/llama-3-2-3b',
+  '@ollama/llama-3-1-nemotron-70b-instruct',
+  '@ollama/llama-3-1-nemotron-mini-4b-instruct',
+  '@cf/deepseek-ai/deepseek-r1-distill-qwen-32b',
+);
+export type LLMModel = S.Schema.Type<typeof LLMModel>;
+
+export const ToolTypes = Object.freeze({
+  // TODO(dmaretskyi): Not implemented yet.
+  // DatabaseQuery: 'database_query',
+  TextToImage: 'text_to_image',
+});
+
+export const LLMTool = S.Struct({
+  name: S.String,
+
+  /**
+   * If the tool is implemented by the service a type should be provided.
+   * For user-implemented tools, this field should be omitted.
+   * See {@link ToolTypes} for the list of supported types.
+   */
+  type: S.optional(S.String),
+
+  /**
+   * Required for user-implemented tools.
+   */
+  description: S.optional(S.String),
+
+  /**
+   * Input schema for the tool in the JSON Schema format.
+   * Required for user-implemented tools.
+   */
+  parameters: S.optional(JsonSchemaType),
+
+  /**
+   * Tool-specific options.
+   */
+  options: S.optional(S.Any),
+
+  /**
+   * Javascript function to execute the tool.
+   */
+  execute: S.optional(S.Any),
+});
+
+export interface LLMTool extends S.Schema.Type<typeof LLMTool> {}
+
+export const GenerateRequest = S.Struct({
+  model: LLMModel,
+
+  spaceId: S.optional(SpaceIdSchema),
+  threadId: S.optional(ObjectId),
+
+  /**
+   * History of messages to include in the context window.
+   */
+  history: S.optional(S.Array(Message)),
+
+  /**
+   * System instructions to the LLM.
+   */
+  // TODO(dmaretskyi): Should this be part of the thread?.
+  systemPrompt: S.optional(S.String),
+
+  /**
+   * Tools available for the LLM.
+   */
+  tools: S.optional(S.Array(LLMTool).pipe(S.mutable)),
+});
+
+export type GenerateRequest = S.Schema.Type<typeof GenerateRequest>;
+
+export type LLMStopReason = 'tool_use' | 'end_turn';
+
+// TODO(dmaretskyi): Effect schema.
+// TODO(dmaretskyi): Rename GenerationStreamEvent.
+export type ResultStreamEvent =
+  | {
+      // TODO(dmaretskyi): Normalize types to our schema.
+      type: 'message_start';
+      message: Omit<Message, 'spaceId' | 'threadId'>;
+    }
+  | {
+      type: 'message_delta';
+      delta: {
+        stopReason: LLMStopReason;
+      };
+    }
+  | {
+      type: 'message_stop';
+    }
+  | {
+      type: 'content_block_start';
+      index: number;
+      content: MessageContentBlock;
+    }
+  | {
+      type: 'content_block_delta';
+      index: number;
+      delta:
+        | {
+            type: 'text_delta';
+            text: string;
+          }
+        | {
+            partial_json: string;
+
+            type: 'input_json_delta';
+          };
+    }
+  | {
+      type: 'content_block_stop';
+      index: number;
+    };

package/src/ai-service/util/index.ts

@@ -0,0 +1,5 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './sse';

package/src/ai-service/{sse.ts → util/sse.ts}

@@ -3,6 +3,11 @@
 // Copied from anthropic SDK source
 //
 
+/**
+ * Server sent events (SSE) iterator.
+ * @param response
+ * @param controller
+ */
 export async function* iterSSEMessages(
   response: Response,
   controller: AbortController,

package/src/bin/chat.ts

@@ -2,11 +2,13 @@
 // Copyright 2024 DXOS.org
 //
 
+import { writeFileSync } from 'fs';
 import inquirer from 'inquirer';
 
+import { ObjectId } from '@dxos/echo-schema';
 import { SpaceId } from '@dxos/keys';
 
-import { AIServiceClientImpl, ObjectId, ToolTypes } from '../ai-service';
+import { AIServiceClientImpl, ToolTypes } from '../ai-service';
 import { runLLM, createUserMessage } from '../conversation';
 import {
   createLogger,
@@ -17,16 +19,11 @@ import {
   Org,
   Project,
   Task,
+  AI_SERVICE_ENDPOINT,
 } from '../testing';
-import { readFileSync, writeFileSync } from 'fs';
-
-// TODO(burdon): Move out of src?
-
-const ENDPOINT = 'http://localhost:8787';
-// const ENDPOINT = 'https://ai-service.dxos.workers.dev';
 
 const client = new AIServiceClientImpl({
-  endpoint: ENDPOINT,
+  endpoint: AI_SERVICE_ENDPOINT.LOCAL,
 });
 
 const dataSource = createTestData();
@@ -46,7 +43,7 @@ while (true) {
       message: 'Enter a message:',
     },
   ]);
-  await client.insertMessages([createUserMessage(spaceId, threadId, prompt.message)]);
+  await client.appendMessages([createUserMessage(spaceId, threadId, prompt.message)]);
 
   await runLLM({
     model: '@anthropic/claude-3-5-sonnet-20241022',

package/src/conversation/conversation.test.ts

@@ -5,37 +5,35 @@
 import { Schema as S } from '@effect/schema';
 import { describe, test } from 'vitest';
 
+import { ObjectId } from '@dxos/echo-schema';
 import { SpaceId } from '@dxos/keys';
 import { log } from '@dxos/log';
 
 import { runLLM, type ConversationEvent } from './conversation';
 import { createUserMessage, defineTool, LLMToolResult } from './types';
-import { AIServiceClientImpl, ObjectId } from '../ai-service';
-
-const ENDPOINT = 'http://localhost:8787';
+import { AIServiceClientImpl } from '../ai-service';
+import { AI_SERVICE_ENDPOINT } from '../testing';
 
 describe.skip('Conversation tests', () => {
   const client = new AIServiceClientImpl({
-    endpoint: ENDPOINT,
+    endpoint: AI_SERVICE_ENDPOINT.LOCAL,
   });
 
   test('hello', async ({ expect }) => {
     const spaceId = SpaceId.random();
     const threadId = ObjectId.random();
 
-    await client.insertMessages([createUserMessage(spaceId, threadId, 'Hello, how are you?')]);
     const result = await runLLM({
       model: '@anthropic/claude-3-5-sonnet-20241022',
-      spaceId,
-      threadId,
+      history: [createUserMessage(spaceId, threadId, 'Hello, how are you?')],
       tools: [],
       client,
       logger: messageLogger,
     });
-    log.info('result', { result });
+    log('result', { result });
   });
 
-  test('tool call', async ({ expect }) => {
+  test.only('tool call', async ({ expect }) => {
     const custodian = defineTool({
       name: 'custodian',
       description: 'Custodian can tell you the password if you say the magic word',
@@ -54,21 +52,19 @@ describe.skip('Conversation tests', () => {
     const spaceId = SpaceId.random();
     const threadId = ObjectId.random();
 
-    await client.insertMessages([createUserMessage(spaceId, threadId, 'What is the password? Ask the custodian.')]);
     const result = await runLLM({
       model: '@anthropic/claude-3-5-sonnet-20241022',
-      spaceId,
-      threadId,
+      history: [createUserMessage(spaceId, threadId, 'What is the password? Ask the custodian.')],
       tools: [custodian],
       client,
       logger: messageLogger,
     });
-    log.info('result', { result });
+    log('result', { result });
   });
 
   const messageLogger = (event: ConversationEvent) => {
     if (event.type === 'message') {
-      log.info('message', { message: event.message });
+      log('message', { message: event.message });
     }
   };
 });

package/src/conversation/conversation.ts

@@ -2,13 +2,13 @@
 // Copyright 2024 DXOS.org
 //
 
+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 { LLMToolResult } from './types';
+import { type LLMTool, Message, type AIServiceClient, type ResultStreamEvent, type LLMModel } from '../ai-service';
 
 export type CreateLLMConversationParams = {
   model: LLMModel;
@@ -18,11 +18,13 @@ export type CreateLLMConversationParams = {
    */
   system?: string;
 
-  spaceId: SpaceId;
-  threadId: ObjectId;
+  spaceId?: SpaceId;
+  threadId?: ObjectId;
 
   tools: LLMTool[];
 
+  history?: Message[];
+
   client: AIServiceClient;
 
   logger?: (event: ConversationEvent) => void;
@@ -37,29 +39,33 @@ 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';
@@ -67,7 +73,7 @@ export const runLLM = async (params: CreateLLMConversationParams) => {
       const toolCalls = message.content.filter((c) => c.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}`);
       }
@@ -75,7 +81,7 @@ export const runLLM = async (params: CreateLLMConversationParams) => {
       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 +97,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 +118,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 +138,92 @@ export const runLLM = async (params: CreateLLMConversationParams) => {
 
   return {
     result: conversationResult,
+    history,
   };
 };
+
+export const isToolUse = (message: Message) => message.content.at(-1)?.type === 'tool_use';
+
+type RunToolsOptions = {
+  message: Message;
+  tools: LLMTool[];
+  extensions?: LLMToolContextExtensions;
+};
+
+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: LLMToolResult;
+  try {
+    toolResult = await tool.execute(toolCall.input, { extensions });
+  } catch (error: any) {
+    log('tool error', { error });
+    toolResult = LLMToolResult.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,
+      };
+    }
+  }
+};