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

package/src/ai-service/client.ts

@@ -2,40 +2,35 @@
 // Copyright 2024 DXOS.org
 //
 
-/* eslint-disable no-console */
-
 import { Schema as S } from '@effect/schema';
 
-import { Trigger } from '@dxos/async';
 import { invariant } from '@dxos/invariant';
-import type { SpaceId } from '@dxos/keys';
-
-import {
-  Message,
-  type GenerateRequest,
-  type ObjectId,
-  type ResultStreamEvent,
-  type Space,
-  type Thread,
-} from './schema';
-import { iterSSEMessages } from './sse';
-import { log } from '@dxos/log';
+import { type SpaceId } from '@dxos/keys';
 
-export type AIServiceClientParams = {
-  endpoint: string;
-};
+import { GenerationStream } from './stream';
+import { Message, type GenerateRequest, type Space, type Thread } from './types';
 
 export interface AIServiceClient {
   getSpace(spaceId: SpaceId): Promise<Space>;
   getThread(spaceId: SpaceId, threadId: string): Promise<Thread>;
-  getMessagesInThread(spaceId: SpaceId, threadId: string): Promise<Message[]>;
-  insertMessages(messages: Message[]): Promise<void>;
+  getMessages(spaceId: SpaceId, threadId: string): Promise<Message[]>;
+  appendMessages(messages: Message[]): Promise<void>;
   updateMessage(spaceId: SpaceId, threadId: string, messageId: string, message: Message): Promise<void>;
   generate(request: GenerateRequest): Promise<GenerationStream>;
 }
 
+// TODO(burdon): Create mock.
+
+export type AIServiceClientParams = {
+  endpoint: string;
+};
+
+/**
+ * Edge GPT client.
+ */
 export class AIServiceClientImpl implements AIServiceClient {
   private readonly _endpoint: string;
+
   constructor({ endpoint }: AIServiceClientParams) {
     this._endpoint = endpoint;
   }
@@ -48,13 +43,14 @@ export class AIServiceClientImpl implements AIServiceClient {
     throw new Error('Not implemented');
   }
 
-  async getMessagesInThread(spaceId: SpaceId, threadId: string): Promise<Message[]> {
+  async getMessages(spaceId: SpaceId, threadId: string): Promise<Message[]> {
     const res = await fetch(`${this._endpoint}/space/${spaceId}/thread/${threadId}/message`);
     return S.decodePromise(S.Array(Message).pipe(S.mutable))((await res.json()).data.results);
   }
 
-  async insertMessages(messages: Message[]): Promise<void> {
-    const res = await fetch(`${this._endpoint}/message`, {
+  async appendMessages(messages: Message[]): Promise<void> {
+    const url = `${this._endpoint}/message`;
+    const res = await fetch(url, {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
       body: JSON.stringify({
@@ -63,7 +59,7 @@ export class AIServiceClientImpl implements AIServiceClient {
     });
 
     if (!res.ok) {
-      throw new Error(await res.text());
+      throw new Error(`${await res.text()} [${url}]`);
     }
   }
 
@@ -72,172 +68,28 @@ export class AIServiceClientImpl implements AIServiceClient {
     throw new Error('Not implemented');
   }
 
+  /**
+   * Open message stream.
+   */
   async generate(request: GenerateRequest): Promise<GenerationStream> {
+    const controller = new AbortController();
+    const signal = controller.signal;
+
     const response = await fetch(`${this._endpoint}/generate`, {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
       body: JSON.stringify(request),
+      signal,
     });
-    invariant(response.body instanceof ReadableStream);
 
+    invariant(response.body instanceof ReadableStream);
     return GenerationStream.fromSSEResponse(
       {
         spaceId: request.spaceId,
         threadId: request.threadId,
       },
       response,
+      controller,
     );
   }
 }
-
-type GenerationParams = {
-  spaceId: SpaceId;
-  threadId: ObjectId;
-};
-
-export class GenerationStream implements AsyncIterable<ResultStreamEvent> {
-  static fromSSEResponse(params: GenerationParams, response: Response) {
-    const controller = new AbortController();
-    const iterator = async function* () {
-      for await (const sse of iterSSEMessages(response, controller)) {
-        if (sse.event === 'completion') {
-          try {
-            yield JSON.parse(sse.data);
-          } catch (e) {
-            console.error('Could not parse message into JSON:', sse.data);
-            console.error('From chunk:', sse.raw);
-            throw e;
-          }
-        }
-
-        if (
-          sse.event === 'message_start' ||
-          sse.event === 'message_delta' ||
-          sse.event === 'message_stop' ||
-          sse.event === 'content_block_start' ||
-          sse.event === 'content_block_delta' ||
-          sse.event === 'content_block_stop'
-        ) {
-          try {
-            yield JSON.parse(sse.data);
-          } catch (e) {
-            console.error('Could not parse message into JSON:', sse.data);
-            console.error('From chunk:', sse.raw);
-            throw e;
-          }
-        }
-
-        if (sse.event === 'ping') {
-          continue;
-        }
-
-        if (sse.event === 'error') {
-          throw new Error(`Message generation error: ${sse.data}`);
-        }
-      }
-    };
-    return new GenerationStream(params, iterator);
-  }
-
-  private _previousMessages: Message[] = [];
-  private _accumulatedMessage?: Message = undefined;
-  private _done = new Trigger();
-  private _iterator?: AsyncIterator<ResultStreamEvent> = undefined;
-
-  constructor(
-    private readonly _params: GenerationParams,
-    private readonly _getIterator: () => AsyncIterableIterator<ResultStreamEvent>,
-  ) {}
-
-  [Symbol.asyncIterator](): AsyncIterator<ResultStreamEvent> {
-    return this._createIterator();
-  }
-
-  cancel() {
-    throw new Error('Not implemented');
-  }
-
-  get accumulatedMessages(): Message[] {
-    // TODO(dmaretskyi): Support multiple accumulated messages.
-    return [...this._previousMessages, ...(this._accumulatedMessage == null ? [] : [this._accumulatedMessage])];
-  }
-
-  /**
-   * Returns the complete message list generated by the AI service.
-   */
-  async complete(): Promise<Message[]> {
-    await this._done.wait();
-    return this.accumulatedMessages;
-  }
-
-  private _createIterator(): AsyncIterator<ResultStreamEvent> {
-    const self = this;
-    return (this._iterator ??= (async function* () {
-      for await (const event of self._getIterator()) {
-        self._processEvent(event);
-        yield event;
-      }
-    })());
-  }
-
-  private _processEvent(event: ResultStreamEvent) {
-    // TODO(dmaretskyi): Support multiple message, e.g. for service-defined tools.
-    switch (event.type) {
-      case 'message_start': {
-        if (this._accumulatedMessage) {
-          this._previousMessages.push(this._accumulatedMessage);
-          this._accumulatedMessage = undefined;
-        }
-        this._accumulatedMessage = {
-          ...event.message,
-          threadId: this._params.threadId,
-          spaceId: this._params.spaceId,
-        };
-        break;
-      }
-      case 'message_delta': {
-        if (this._accumulatedMessage === undefined) {
-          throw new Error('Received message delta without a message start');
-        }
-        break;
-      }
-      case 'message_stop': {
-        if (this._accumulatedMessage === undefined) {
-          throw new Error('Received message stop without a message start');
-        }
-
-        break;
-      }
-      case 'content_block_start': {
-        invariant(this._accumulatedMessage);
-        this._accumulatedMessage.content.push(event.content);
-        break;
-      }
-      case 'content_block_delta': {
-        invariant(this._accumulatedMessage);
-        const snapshotContent = this._accumulatedMessage.content.at(event.index);
-        if (snapshotContent?.type === 'text' && event.delta.type === 'text_delta') {
-          snapshotContent.text += event.delta.text;
-        } else if (snapshotContent?.type === 'tool_use' && event.delta.type === 'input_json_delta') {
-          snapshotContent.inputJson ??= '';
-          snapshotContent.inputJson += event.delta.partial_json;
-
-          // TODO(dmaretskyi): Partial parsing.
-          // if (jsonBuf) {
-          //   snapshotContent.input = partialParse(jsonBuf);
-          // }
-        }
-        break;
-      }
-      case 'content_block_stop': {
-        invariant(this._accumulatedMessage);
-        const lastBlock = this._accumulatedMessage.content.at(-1);
-        if (lastBlock?.type === 'tool_use') {
-          lastBlock.input = JSON.parse(lastBlock.inputJson ?? '');
-        }
-        break;
-      }
-    }
-    this._done.wake();
-  }
-}

package/src/ai-service/index.ts

@@ -3,4 +3,5 @@
 //
 
 export * from './client';
-export * from './schema';
+export * from './stream';
+export * from './types';

package/src/ai-service/stream.test.ts

@@ -0,0 +1,186 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { SaxesParser } from 'saxes';
+import { describe, it } from 'vitest';
+
+import { ObjectId } from '@dxos/echo-schema';
+import { log } from '@dxos/log';
+
+import { GenerationStream } from './stream';
+import { type Message } from './types';
+
+type Part = { event: string; data: any };
+
+export const createTestStream = (blocks: string[]): ReadableStream => {
+  const encoder = new TextEncoder();
+  return new ReadableStream({
+    start: (controller) => {
+      const push = (message: Part) => {
+        controller.enqueue(encoder.encode(`event: ${message.event}\ndata: ${JSON.stringify(message.data)}\n\n`));
+      };
+
+      push({
+        event: 'message_start',
+        data: { type: 'message_start', message: { id: '123', role: 'assistant', content: [] } },
+      });
+
+      for (const block of blocks) {
+        push({
+          event: 'content_block_start',
+          data: { type: 'content_block_start', content: { type: 'text', text: '' } },
+        });
+
+        let index = 0;
+        for (const word of block.split(' ')) {
+          push({
+            event: 'content_block_delta',
+            data: { type: 'content_block_delta', index, delta: { type: 'text_delta', text: word + ' ' } },
+          });
+
+          index += word.length + 1;
+        }
+
+        push({
+          event: 'content_block_stop',
+          data: { type: 'content_block_stop' },
+        });
+      }
+
+      push({
+        event: 'message_stop',
+        data: { type: 'message_stop' },
+      });
+
+      controller.close();
+    },
+  });
+};
+
+describe('GenerationStream', () => {
+  it('should stream messages', async ({ expect }) => {
+    const input = [['<cot>', '1. Ready.', '2. Aim.', '3. Fire.', '</cot>'].join('\n'), 'Hello world!'];
+    const output = [];
+
+    let block: string[] = [];
+    const stream = GenerationStream.fromSSEResponse({}, new Response(createTestStream(input)));
+    for await (const event of stream) {
+      switch (event.type) {
+        case 'content_block_start': {
+          block = [(event.content as any).text];
+          break;
+        }
+
+        case 'content_block_delta': {
+          block.push((event.delta as any).text);
+          break;
+        }
+
+        case 'content_block_stop': {
+          output.push(block.join('').trim());
+          break;
+        }
+      }
+    }
+
+    expect(output.length).to.equal(input.length);
+    expect(output).to.deep.equal(input);
+    await stream.complete();
+  });
+
+  /**
+   * https://astexplorer.net
+   */
+  it('should parse the stream', async ({ expect }) => {
+    const input = [
+      // Text
+      `
+      Hello world!
+      `,
+
+      // COT
+      `
+      <cot>
+      1. Prepare the ship.
+      2. Arm the torpedoes.
+      3. Fire lasers.
+      </cot>
+      `,
+
+      // Artifact
+      // <?xml version="1.0" encoding="UTF-8"?>
+      `
+      <artifact id="123" />
+      <choice id="123">
+        <artifact id="123" />
+      </choice>
+      `,
+
+      // JSON
+      `
+      <json>
+      <![CDATA[
+      {
+        "type": "text",
+        "text": "<test>OK</test>"
+      }
+      ]]>
+      </json>
+      `,
+
+      // Code/JSX
+      `,
+      <code>
+      <![CDATA[
+      const Test = () => {
+        return <div>Hello world!</div>;
+      }
+      ]]>
+      </code>
+      `,
+    ];
+
+    // TODO(burdon): Delimiter for JSON blocks.
+
+    const output: Message = { id: ObjectId.random(), role: 'assistant', content: [] };
+
+    /**
+     * https://www.npmjs.com/package/saxex
+     */
+    const parser = new SaxesParser({
+      fragment: true,
+    });
+
+    parser.on('error', (error) => {
+      log.error('error', { error });
+    });
+    parser.on('opentag', (node) => {
+      log.info('opentag', { node });
+    });
+    parser.on('closetag', (node) => {
+      log.info('closetag', { node });
+    });
+    parser.on('cdata', (text) => {
+      log.info('cdata', { text });
+    });
+    parser.on('text', (text) => {
+      log.info('text', { text });
+    });
+    parser.on('end', () => {
+      log.info('end');
+    });
+
+    const stream = GenerationStream.fromSSEResponse({}, new Response(createTestStream(input)));
+    for await (const event of stream) {
+      if (event.type === 'content_block_delta') {
+        // log.info('delta', { delta: (event.delta as any).text });
+        parser.write((event.delta as any).text);
+      }
+    }
+
+    await stream.complete();
+
+    expect(output.content).to.deep.equal([]);
+  });
+});

package/src/ai-service/stream.ts

@@ -0,0 +1,181 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { Trigger } from '@dxos/async';
+import { type ObjectId } from '@dxos/echo-schema';
+import { invariant } from '@dxos/invariant';
+import { type SpaceId } from '@dxos/keys';
+import { log } from '@dxos/log';
+
+import { type Message, type ResultStreamEvent } from './types';
+import { iterSSEMessages } from './util';
+
+export type GenerationParams = {
+  spaceId?: SpaceId;
+  threadId?: ObjectId;
+};
+
+export class GenerationStream implements AsyncIterable<ResultStreamEvent> {
+  /**
+   * Creates a stream from an SSE response.
+   */
+  static fromSSEResponse(params: GenerationParams, response: Response, controller = new AbortController()) {
+    const iterator = async function* () {
+      for await (const sse of iterSSEMessages(response, controller)) {
+        if (sse.event === 'completion') {
+          try {
+            yield JSON.parse(sse.data);
+          } catch (err) {
+            log.error('could not parse message into JSON:', { data: sse.data, raw: sse.raw });
+            throw err;
+          }
+        }
+
+        if (
+          sse.event === 'message_start' ||
+          sse.event === 'message_delta' ||
+          sse.event === 'message_stop' ||
+          sse.event === 'content_block_start' ||
+          sse.event === 'content_block_delta' ||
+          sse.event === 'content_block_stop'
+        ) {
+          try {
+            yield JSON.parse(sse.data);
+          } catch (err) {
+            log.error('could not parse message into JSON:', { data: sse.data, raw: sse.raw });
+            throw err;
+          }
+        }
+
+        if (sse.event === 'ping') {
+          continue;
+        }
+
+        if (sse.event === 'error') {
+          throw new Error(`Message generation error: ${sse.data}`);
+        }
+      }
+    };
+
+    return new GenerationStream(controller, iterator, params);
+  }
+
+  private readonly _done = new Trigger();
+  private _previousMessages: Message[] = [];
+  private _accumulatedMessage?: Message = undefined;
+  private _iterator?: AsyncIterator<ResultStreamEvent> = undefined;
+
+  constructor(
+    private readonly _controller: AbortController,
+    private readonly _getIterator: () => AsyncIterableIterator<ResultStreamEvent>,
+    private readonly _params: GenerationParams,
+  ) {}
+
+  [Symbol.asyncIterator](): AsyncIterator<ResultStreamEvent> {
+    return this._createIterator();
+  }
+
+  get accumulatedMessages(): Message[] {
+    // TODO(dmaretskyi): Support multiple accumulated messages.
+    return [...this._previousMessages, ...(this._accumulatedMessage == null ? [] : [this._accumulatedMessage])];
+  }
+
+  /**
+   * Aborts the stream.
+   */
+  abort() {
+    this._controller.abort();
+  }
+
+  /**
+   * Returns the complete message list generated by the AI service.
+   */
+  async complete(): Promise<Message[]> {
+    await this._done.wait();
+    return this.accumulatedMessages;
+  }
+
+  private _createIterator(): AsyncIterator<ResultStreamEvent> {
+    const self = this;
+    return (this._iterator ??= (() => {
+      const generator = async function* (this: GenerationStream) {
+        try {
+          for await (const event of self._getIterator()) {
+            self._processEvent(event);
+            yield event;
+          }
+          this._done.wake();
+        } catch (err: any) {
+          this._done.throw(err);
+        }
+      };
+      return generator.call(this);
+    })());
+  }
+
+  private _processEvent(event: ResultStreamEvent) {
+    // TODO(dmaretskyi): Support multiple message, e.g., for service-defined tools.
+    switch (event.type) {
+      case 'message_start': {
+        if (this._accumulatedMessage) {
+          this._previousMessages.push(this._accumulatedMessage);
+          this._accumulatedMessage = undefined;
+        }
+        this._accumulatedMessage = {
+          ...event.message,
+          threadId: this._params.threadId,
+          spaceId: this._params.spaceId,
+        };
+        break;
+      }
+
+      case 'message_delta': {
+        if (this._accumulatedMessage === undefined) {
+          throw new Error('Received message delta without a message start');
+        }
+        break;
+      }
+
+      case 'message_stop': {
+        if (this._accumulatedMessage === undefined) {
+          throw new Error('Received message stop without a message start');
+        }
+        break;
+      }
+
+      case 'content_block_start': {
+        invariant(this._accumulatedMessage);
+        this._accumulatedMessage.content.push(event.content);
+        break;
+      }
+
+      case 'content_block_delta': {
+        invariant(this._accumulatedMessage);
+        const snapshotContent = this._accumulatedMessage.content.at(event.index);
+        if (snapshotContent?.type === 'text' && event.delta.type === 'text_delta') {
+          snapshotContent.text += event.delta.text;
+        } else if (snapshotContent?.type === 'tool_use' && event.delta.type === 'input_json_delta') {
+          snapshotContent.inputJson ??= '';
+          snapshotContent.inputJson += event.delta.partial_json;
+
+          // TODO(dmaretskyi): Partial parsing.
+          // if (jsonBuf) {
+          //   snapshotContent.input = partialParse(jsonBuf);
+          // }
+        }
+        break;
+      }
+
+      case 'content_block_stop': {
+        invariant(this._accumulatedMessage);
+        const lastBlock = this._accumulatedMessage.content.at(-1);
+        if (lastBlock?.type === 'tool_use') {
+          lastBlock.input =
+            lastBlock.inputJson == null || lastBlock.inputJson === '' ? {} : JSON.parse(lastBlock.inputJson ?? '');
+        }
+        break;
+      }
+    }
+  }
+}

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

@@ -0,0 +1,6 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './message';
+export * from './schema';