@dxos/assistant 0.7.4-df37807

package/src/ai-service/sse.ts

@@ -0,0 +1,336 @@
+//
+// Copyright 2024 DXOS.org
+// Copied from anthropic SDK source
+//
+
+export async function* iterSSEMessages(
+  response: Response,
+  controller: AbortController,
+): AsyncGenerator<ServerSentEvent, void, unknown> {
+  if (!response.body) {
+    controller.abort();
+    throw new Error('Attempted to iterate over a response with no body');
+  }
+
+  const sseDecoder = new SSEDecoder();
+  const lineDecoder = new LineDecoder();
+
+  const iter = readableStreamAsyncIterable<Bytes>(response.body);
+  for await (const sseChunk of iterSSEChunks(iter)) {
+    for (const line of lineDecoder.decode(sseChunk)) {
+      const sse = sseDecoder.decode(line);
+      if (sse) {
+        yield sse;
+      }
+    }
+  }
+
+  for (const line of lineDecoder.flush()) {
+    const sse = sseDecoder.decode(line);
+    if (sse) {
+      yield sse;
+    }
+  }
+}
+
+export type ServerSentEvent = {
+  event: string | null;
+  data: string;
+  raw: string[];
+};
+
+/**
+ * Given an async iterable iterator, iterates over it and yields full
+ * SSE chunks, i.e. yields when a double new-line is encountered.
+ */
+async function* iterSSEChunks(iterator: AsyncIterableIterator<Bytes>): AsyncGenerator<Uint8Array> {
+  let data = new Uint8Array();
+
+  for await (const chunk of iterator) {
+    if (chunk == null) {
+      continue;
+    }
+
+    const binaryChunk =
+      chunk instanceof ArrayBuffer
+        ? new Uint8Array(chunk)
+        : typeof chunk === 'string'
+          ? new TextEncoder().encode(chunk)
+          : chunk;
+
+    const newData = new Uint8Array(data.length + binaryChunk.length);
+    newData.set(data);
+    newData.set(binaryChunk, data.length);
+    data = newData;
+
+    let patternIndex;
+    while ((patternIndex = findDoubleNewlineIndex(data)) !== -1) {
+      yield data.slice(0, patternIndex);
+      data = data.slice(patternIndex);
+    }
+  }
+
+  if (data.length > 0) {
+    yield data;
+  }
+}
+
+const findDoubleNewlineIndex = (buffer: Uint8Array): number => {
+  // This function searches the buffer for the end patterns (\r\r, \n\n, \r\n\r\n)
+  // and returns the index right after the first occurrence of any pattern,
+  // or -1 if none of the patterns are found.
+  const newline = 0x0a; // \n
+  const carriage = 0x0d; // \r
+
+  for (let i = 0; i < buffer.length - 2; i++) {
+    if (buffer[i] === newline && buffer[i + 1] === newline) {
+      // \n\n
+      return i + 2;
+    }
+    if (buffer[i] === carriage && buffer[i + 1] === carriage) {
+      // \r\r
+      return i + 2;
+    }
+    if (
+      buffer[i] === carriage &&
+      buffer[i + 1] === newline &&
+      i + 3 < buffer.length &&
+      buffer[i + 2] === carriage &&
+      buffer[i + 3] === newline
+    ) {
+      // \r\n\r\n
+      return i + 4;
+    }
+  }
+
+  return -1;
+};
+
+class SSEDecoder {
+  private data: string[];
+  private event: string | null;
+  private chunks: string[];
+
+  constructor() {
+    this.event = null;
+    this.data = [];
+    this.chunks = [];
+  }
+
+  decode(line: string) {
+    if (line.endsWith('\r')) {
+      line = line.substring(0, line.length - 1);
+    }
+
+    if (!line) {
+      // empty line and we didn't previously encounter any messages
+      if (!this.event && !this.data.length) {
+        return null;
+      }
+
+      const sse: ServerSentEvent = {
+        event: this.event,
+        data: this.data.join('\n'),
+        raw: this.chunks,
+      };
+
+      this.event = null;
+      this.data = [];
+      this.chunks = [];
+
+      return sse;
+    }
+
+    this.chunks.push(line);
+
+    if (line.startsWith(':')) {
+      return null;
+    }
+
+    let [fieldname, _, value] = partition(line, ':');
+
+    if (value.startsWith(' ')) {
+      value = value.substring(1);
+    }
+
+    if (fieldname === 'event') {
+      this.event = value;
+    } else if (fieldname === 'data') {
+      this.data.push(value);
+    }
+
+    return null;
+  }
+}
+
+/** This is an internal helper function that's just used for testing */
+export const _decodeChunks = (chunks: string[]): string[] => {
+  const decoder = new LineDecoder();
+  const lines: string[] = [];
+  for (const chunk of chunks) {
+    lines.push(...decoder.decode(chunk));
+  }
+
+  return lines;
+};
+
+const partition = (str: string, delimiter: string): [string, string, string] => {
+  const index = str.indexOf(delimiter);
+  if (index !== -1) {
+    return [str.substring(0, index), delimiter, str.substring(index + delimiter.length)];
+  }
+
+  return [str, '', ''];
+};
+
+/**
+ * Most browsers don't yet have async iterable support for ReadableStream,
+ * and Node has a very different way of reading bytes from its "ReadableStream".
+ *
+ * This polyfill was pulled from https://github.com/MattiasBuelens/web-streams-polyfill/pull/122#issuecomment-1627354490
+ */
+export function readableStreamAsyncIterable<T>(stream: any): AsyncIterableIterator<T> {
+  if (stream[Symbol.asyncIterator]) {
+    return stream;
+  }
+
+  const reader = stream.getReader();
+  return {
+    next: async () => {
+      try {
+        const result = await reader.read();
+        if (result?.done) {
+          reader.releaseLock();
+        } // release lock when stream becomes closed
+        return result;
+      } catch (e) {
+        reader.releaseLock(); // release lock when stream becomes errored
+        throw e;
+      }
+    },
+    return: async () => {
+      const cancelPromise = reader.cancel();
+      reader.releaseLock();
+      await cancelPromise;
+      return { done: true, value: undefined };
+    },
+    [Symbol.asyncIterator]() {
+      return this;
+    },
+  };
+}
+
+export type Bytes = string | ArrayBuffer | Uint8Array | Buffer | null | undefined;
+
+/**
+ * A re-implementation of httpx's `LineDecoder` in Python that handles incrementally
+ * reading lines from text.
+ *
+ * https://github.com/encode/httpx/blob/920333ea98118e9cf617f246905d7b202510941c/httpx/_decoders.py#L258
+ */
+export class LineDecoder {
+  // prettier-ignore
+  static NEWLINE_CHARS = new Set(['\n', '\r']);
+  static NEWLINE_REGEXP = /\r\n|[\n\r]/g;
+
+  buffer: string[];
+  trailingCR: boolean;
+  textDecoder: any; // TextDecoder found in browsers; not typed to avoid pulling in either "dom" or "node" types.
+
+  constructor() {
+    this.buffer = [];
+    this.trailingCR = false;
+  }
+
+  decode(chunk: Bytes): string[] {
+    let text = this.decodeText(chunk);
+
+    if (this.trailingCR) {
+      text = '\r' + text;
+      this.trailingCR = false;
+    }
+    if (text.endsWith('\r')) {
+      this.trailingCR = true;
+      text = text.slice(0, -1);
+    }
+
+    if (!text) {
+      return [];
+    }
+
+    const trailingNewline = LineDecoder.NEWLINE_CHARS.has(text[text.length - 1] || '');
+    let lines = text.split(LineDecoder.NEWLINE_REGEXP);
+
+    // if there is a trailing new line then the last entry will be an empty
+    // string which we don't care about
+    if (trailingNewline) {
+      lines.pop();
+    }
+
+    if (lines.length === 1 && !trailingNewline) {
+      this.buffer.push(lines[0]!);
+      return [];
+    }
+
+    if (this.buffer.length > 0) {
+      lines = [this.buffer.join('') + lines[0], ...lines.slice(1)];
+      this.buffer = [];
+    }
+
+    if (!trailingNewline) {
+      this.buffer = [lines.pop() || ''];
+    }
+
+    return lines;
+  }
+
+  decodeText(bytes: Bytes): string {
+    if (bytes == null) {
+      return '';
+    }
+    if (typeof bytes === 'string') {
+      return bytes;
+    }
+
+    // Node:
+    if (typeof Buffer !== 'undefined') {
+      if (bytes instanceof Buffer) {
+        return bytes.toString();
+      }
+      if (bytes instanceof Uint8Array) {
+        return Buffer.from(bytes).toString();
+      }
+
+      throw new Error(
+        `Unexpected: received non-Uint8Array (${bytes.constructor.name}) stream chunk in an environment with a global "Buffer" defined, which this library assumes to be Node. Please report this error.`,
+      );
+    }
+
+    // Browser
+    if (typeof TextDecoder !== 'undefined') {
+      if (bytes instanceof Uint8Array || bytes instanceof ArrayBuffer) {
+        this.textDecoder ??= new TextDecoder('utf8');
+        return this.textDecoder.decode(bytes);
+      }
+
+      throw new Error(
+        `Unexpected: received non-Uint8Array/ArrayBuffer (${
+          (bytes as any).constructor.name
+        }) in a web platform. Please report this error.`,
+      );
+    }
+
+    throw new Error('Unexpected: neither Buffer nor TextDecoder are available as globals. Please report this error.');
+  }
+
+  flush(): string[] {
+    if (!this.buffer.length && !this.trailingCR) {
+      return [];
+    }
+
+    const lines = [this.buffer.join('')];
+    this.buffer = [];
+    this.trailingCR = false;
+    return lines;
+  }
+}

package/src/bin/chat.ts

@@ -0,0 +1,79 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import inquirer from 'inquirer';
+
+import { SpaceId } from '@dxos/keys';
+
+import { AIServiceClientImpl, ObjectId, ToolTypes } from '../ai-service';
+import { runLLM, createUserMessage } from '../conversation';
+import {
+  createLogger,
+  createCypherTool,
+  createSystemPrompt,
+  createTestData,
+  Contact,
+  Org,
+  Project,
+  Task,
+} 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,
+});
+
+const dataSource = createTestData();
+
+const cypherTool = createCypherTool(dataSource);
+
+const schemaTypes = [Org, Project, Task, Contact];
+
+const spaceId = SpaceId.random();
+const threadId = ObjectId.random();
+
+while (true) {
+  const prompt = await inquirer.prompt([
+    {
+      type: 'input',
+      name: 'message',
+      message: 'Enter a message:',
+    },
+  ]);
+  await client.insertMessages([createUserMessage(spaceId, threadId, prompt.message)]);
+
+  await runLLM({
+    model: '@anthropic/claude-3-5-sonnet-20241022',
+    spaceId,
+    threadId,
+    system: createSystemPrompt(schemaTypes),
+    tools: [
+      cypherTool,
+      {
+        name: 'text-to-image',
+        type: ToolTypes.TextToImage,
+      },
+    ],
+    client,
+    logger: createLogger({
+      stream: true,
+      filter: (e) => {
+        return true;
+      },
+      onImage: (img) => {
+        const path = `/tmp/image-${img.id}.jpeg`;
+        writeFileSync(path, Buffer.from(img.source.data, 'base64'));
+        console.log(`Saved image to ${path}`);
+        // Print image in iTerm using ANSI escape sequence
+        const imageData = img.source.data;
+        process.stdout.write('\x1b]1337;File=inline=1:' + imageData + '\x07');
+      },
+    }),
+  });
+}

package/src/conversation/conversation.test.ts

@@ -0,0 +1,74 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { Schema as S } from '@effect/schema';
+import { describe, test } from 'vitest';
+
+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';
+
+describe.skip('Conversation tests', () => {
+  const client = new AIServiceClientImpl({
+    endpoint: ENDPOINT,
+  });
+
+  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,
+      tools: [],
+      client,
+      logger: messageLogger,
+    });
+    log.info('result', { result });
+  });
+
+  test('tool call', async ({ expect }) => {
+    const custodian = defineTool({
+      name: 'custodian',
+      description: 'Custodian can tell you the password if you say the magic word',
+      schema: S.Struct({
+        magicWord: S.String.annotations({ description: 'The magic word. Should be exactly "pretty please"' }),
+      }),
+      execute: async ({ magicWord }) => {
+        if (magicWord === 'pretty please') {
+          return LLMToolResult.Success('The password is: "The sky is gray"');
+        } else {
+          return LLMToolResult.Error('Wrong magic word');
+        }
+      },
+    });
+
+    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,
+      tools: [custodian],
+      client,
+      logger: messageLogger,
+    });
+    log.info('result', { result });
+  });
+
+  const messageLogger = (event: ConversationEvent) => {
+    if (event.type === 'message') {
+      log.info('message', { message: event.message });
+    }
+  };
+});

package/src/conversation/conversation.ts

@@ -0,0 +1,134 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+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';
+
+export type CreateLLMConversationParams = {
+  model: LLMModel;
+
+  /**
+   * System prompt that specifies instructions for the LLM.
+   */
+  system?: string;
+
+  spaceId: SpaceId;
+  threadId: ObjectId;
+
+  tools: LLMTool[];
+
+  client: AIServiceClient;
+
+  logger?: (event: ConversationEvent) => void;
+};
+
+export type ConversationEvent =
+  | {
+      type: 'message';
+      message: Message;
+    }
+  | ResultStreamEvent;
+
+export const runLLM = async (params: CreateLLMConversationParams) => {
+  let conversationResult: any = null;
+
+  const generate = async () => {
+    log('llm generate', { tools: params.tools });
+    const beginTs = Date.now();
+    const result = await params.client.generate({
+      model: params.model,
+      spaceId: params.spaceId,
+      threadId: params.threadId,
+      systemPrompt: params.system,
+      tools: params.tools as any,
+    });
+
+    for await (const event of result) {
+      params.logger?.(event);
+    }
+    const messages = await result.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') })),
+    );
+
+    const isToolUse = message.content.at(-1)?.type === 'tool_use';
+    if (isToolUse) {
+      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);
+      if (!tool) {
+        throw new Error(`Tool not found: ${toolCall.name}`);
+      }
+
+      const toolResult = await tool.execute(toolCall.input, {});
+      switch (toolResult.kind) {
+        case 'error': {
+          log('tool error', { message: toolResult.message });
+          const resultMessage: Message = {
+            id: ObjectId.random(),
+            spaceId: params.spaceId,
+            threadId: params.threadId,
+            role: 'user',
+            content: [
+              {
+                type: 'tool_result',
+                toolUseId: toolCall.id,
+                content: toolResult.message,
+                isError: true,
+              },
+            ],
+          };
+          params.logger?.({ type: 'message', message: resultMessage });
+          await params.client.insertMessages([resultMessage]);
+
+          return true;
+        }
+        case 'success': {
+          log('tool success', { result: toolResult.result });
+          const resultMessage: Message = {
+            id: ObjectId.random(),
+            spaceId: params.spaceId,
+            threadId: params.threadId,
+            role: 'user',
+            content: [
+              {
+                type: 'tool_result',
+                toolUseId: toolCall.id,
+                content: JSON.stringify(toolResult.result),
+              },
+            ],
+          };
+          params.logger?.({ type: 'message', message: resultMessage });
+          await params.client.insertMessages([resultMessage]);
+
+          return true;
+        }
+        case 'break': {
+          log('tool break', { result: toolResult.result });
+          conversationResult = toolResult.result;
+          return false;
+        }
+      }
+    }
+    return false;
+  };
+
+  // eslint-disable-next-line no-empty
+  while (await generate()) {}
+
+  return {
+    result: conversationResult,
+  };
+};

package/src/conversation/index.ts

@@ -0,0 +1,6 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+export * from './conversation';
+export * from './types';