@dxos/assistant 0.7.4 → 0.7.5-labs.071a3e2

package/src/ai-service/client.ts

@@ -2,41 +2,38 @@
 // Copyright 2024 DXOS.org
 //
 
-/* eslint-disable no-console */
-
 import { Schema as S } from '@effect/schema';
 
-import { Trigger } from '@dxos/async';
+import { Message, type Space, type Thread } from '@dxos/artifact';
 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 { createGenerationStream, type GenerationStream } from './stream';
+import { type GenerateRequest } from './types';
 
+// TODO(burdon): Rename.
 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>;
 }
 
+export type AIServiceClientParams = {
+  endpoint: string;
+};
+
+/**
+ * Edge GPT client.
+ */
+// TODO(burdon): Create mock.
 export class AIServiceClientImpl implements AIServiceClient {
   private readonly _endpoint: string;
+
   constructor({ endpoint }: AIServiceClientParams) {
+    invariant(endpoint, 'endpoint is required');
     this._endpoint = endpoint;
   }
 
@@ -48,13 +45,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 +61,7 @@ export class AIServiceClientImpl implements AIServiceClient {
     });
 
     if (!res.ok) {
-      throw new Error(await res.text());
+      throw new Error(`${await res.text()} [${url}]`);
     }
   }
 
@@ -72,172 +70,19 @@ export class AIServiceClientImpl implements AIServiceClient {
     throw new Error('Not implemented');
   }
 
+  /**
+   * Open message stream.
+   */
   async generate(request: GenerateRequest): Promise<GenerationStream> {
+    const controller = new AbortController();
     const response = await fetch(`${this._endpoint}/generate`, {
+      signal: controller.signal,
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
       body: JSON.stringify(request),
     });
-    invariant(response.body instanceof ReadableStream);
-
-    return GenerationStream.fromSSEResponse(
-      {
-        spaceId: request.spaceId,
-        threadId: request.threadId,
-      },
-      response,
-    );
-  }
-}
-
-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();
+    invariant(response.body instanceof ReadableStream);
+    return createGenerationStream(response, controller);
   }
 }

package/src/ai-service/index.ts

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

package/src/ai-service/parser.ts

@@ -0,0 +1,338 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Message, type MessageContentBlock } from '@dxos/artifact';
+import { Event } from '@dxos/async';
+import { createStatic } from '@dxos/echo-schema';
+import { invariant } from '@dxos/invariant';
+import { log } from '@dxos/log';
+import { isNotFalsy, safeParseJson } from '@dxos/util';
+
+import { type GenerationStream } from './stream';
+import { StreamTransform, type StreamBlock } from './transform';
+
+/**
+ * Parse mixed content of plain text, XML fragments, and JSON blocks.
+ */
+export class MixedStreamParser {
+  /**
+   * New message.
+   */
+  public message = new Event<Message>();
+
+  /**
+   * Complete block added to Message.
+   */
+  public block = new Event<MessageContentBlock>();
+
+  /**
+   * Update partial block (while streaming).
+   */
+  public update = new Event<MessageContentBlock>();
+
+  /**
+   * Current message.
+   */
+  private _message?: Message | undefined;
+
+  private _emitBlock(contentBlock: MessageContentBlock, streamBlock?: StreamBlock) {
+    const messageBlock = streamBlock ? mergeMessageBlock(contentBlock, streamBlock) : contentBlock;
+    if (messageBlock) {
+      if (messageBlock.type === 'text' && messageBlock.text.length === 0) {
+        return;
+      }
+
+      invariant(this._message);
+      this._message.content.push(messageBlock);
+      this.block.emit(messageBlock);
+    }
+  }
+
+  private _emitUpdate(contentBlock: MessageContentBlock, streamBlock: StreamBlock) {
+    const messageBlock = mergeMessageBlock(contentBlock, streamBlock);
+    if (messageBlock) {
+      messageBlock.pending = true;
+      this.update.emit(messageBlock);
+    }
+  }
+
+  /**
+   * Parse stream until end.
+   */
+  async parse(stream: GenerationStream) {
+    const transformer = new StreamTransform();
+
+    /**
+     * Current content message block.
+     */
+    let contentBlock: MessageContentBlock | undefined;
+
+    /**
+     * Current partial block used to accumulate content.
+     */
+    let streamBlock: StreamBlock | undefined;
+    const stack: StreamBlock[] = [];
+
+    for await (const event of stream) {
+      // log.info('event', { type: event.type, event });
+      switch (event.type) {
+        //
+        // Messages.
+        //
+
+        case 'message_start': {
+          if (this._message) {
+            log.warn('unexpected message_start');
+          }
+
+          this._message = createStatic(Message, { role: 'assistant', content: [] });
+          this.message.emit(this._message);
+          break;
+        }
+
+        case 'message_delta': {
+          if (!this._message) {
+            log.warn('unexpected message_delta');
+          }
+
+          break;
+        }
+
+        case 'message_stop': {
+          if (!this._message) {
+            log.warn('unexpected message_stop');
+          } else {
+            this._message = undefined;
+          }
+          break;
+        }
+
+        //
+        // Content blocks.
+        //
+
+        case 'content_block_start': {
+          if (contentBlock) {
+            log.warn('unexpected content_block_start', { content: contentBlock });
+          }
+
+          // TODO(burdon): Parse initial content (factor out from content_block_delta).
+          contentBlock = event.content;
+          break;
+        }
+
+        case 'content_block_delta': {
+          if (!contentBlock) {
+            log.warn('unexpected content_block_delta');
+            break;
+          }
+
+          switch (event.delta.type) {
+            //
+            // Text
+            //
+            case 'text_delta': {
+              const chunks = transformer.transform(event.delta.text);
+              for (const chunk of chunks) {
+                // log.info('text_delta', { chunk, current });
+
+                switch (streamBlock?.type) {
+                  //
+                  // XML Fragment.
+                  //
+                  case 'tag': {
+                    if (chunk.type === 'tag') {
+                      if (chunk.selfClosing) {
+                        streamBlock.content.push(chunk);
+                      } else if (chunk.closing) {
+                        if (stack.length > 0) {
+                          const top = stack.pop();
+                          invariant(top && top.type === 'tag');
+                          top.content.push(streamBlock);
+                          streamBlock = top;
+                        } else {
+                          this._emitBlock(contentBlock, streamBlock);
+                          streamBlock = undefined;
+                        }
+                      } else {
+                        stack.push(streamBlock);
+                        streamBlock = chunk;
+                      }
+                    } else {
+                      // Append text.
+                      if (streamBlock.content.length === 0) {
+                        streamBlock.content.push(chunk);
+                      } else {
+                        const last = streamBlock.content.at(-1);
+                        invariant(last);
+                        if (last.type === 'text') {
+                          last.content += chunk.content;
+                        } else {
+                          streamBlock.content.push(chunk);
+                        }
+                      }
+                    }
+                    break;
+                  }
+
+                  //
+                  // Text Fragment.
+                  //
+                  case 'text': {
+                    if (chunk.type === 'tag') {
+                      this._emitBlock(contentBlock, streamBlock);
+                      if (chunk.selfClosing) {
+                        this._emitBlock(contentBlock, chunk);
+                        streamBlock = undefined;
+                      } else {
+                        streamBlock = chunk;
+                      }
+                    } else {
+                      // Append text.
+                      streamBlock.content += chunk.content;
+                    }
+                    break;
+                  }
+
+                  //
+                  // No current chunk.
+                  //
+                  default: {
+                    if (chunk.type === 'tag' && chunk.selfClosing) {
+                      this._emitBlock(contentBlock, chunk);
+                    } else {
+                      streamBlock = chunk;
+                    }
+                  }
+                }
+              }
+              break;
+            }
+
+            //
+            // JSON
+            //
+            case 'input_json_delta': {
+              if (!contentBlock) {
+                log.warn('unexpected input_json_delta', { contentBlock });
+                break;
+              }
+
+              if (streamBlock?.type === 'json') {
+                streamBlock.content += event.delta.partial_json;
+              } else {
+                streamBlock = { type: 'json', content: event.delta.partial_json };
+              }
+              break;
+            }
+          }
+          break;
+        }
+
+        case 'content_block_stop': {
+          if (!contentBlock) {
+            log.warn('unexpected content_block_stop');
+            break;
+          }
+
+          this._emitBlock(contentBlock, streamBlock);
+          contentBlock = undefined;
+          streamBlock = undefined;
+          break;
+        }
+      }
+
+      if (contentBlock && streamBlock) {
+        this._emitUpdate(contentBlock, streamBlock);
+      }
+    } // for.
+  }
+}
+
+/**
+ * Convert stream block to message content block.
+ */
+export const mergeMessageBlock = (
+  contentBlock: MessageContentBlock,
+  streamBlock: StreamBlock,
+): MessageContentBlock | undefined => {
+  log('mergeMessageBlock', { contentBlock, streamBlock });
+
+  switch (streamBlock.type) {
+    //
+    // Text
+    //
+    case 'text': {
+      return { ...contentBlock, type: 'text', text: streamBlock.content };
+    }
+
+    //
+    // XML
+    //
+    case 'tag': {
+      switch (streamBlock.tag) {
+        case 'cot': {
+          const content = streamBlock.content
+            .map((block) => {
+              switch (block.type) {
+                case 'text':
+                  return block.content;
+                default:
+                  return null;
+              }
+            })
+            .filter(isNotFalsy)
+            .join('\n');
+
+          return { ...contentBlock, type: 'text', disposition: 'cot', text: content };
+        }
+
+        case 'artifact': {
+          const { attributes } = streamBlock;
+          return { type: 'json', disposition: 'artifact', json: JSON.stringify(attributes) };
+        }
+
+        case 'suggest': {
+          if (streamBlock.content.length === 1 && streamBlock.content[0].type === 'text') {
+            return {
+              type: 'json',
+              disposition: 'suggest',
+              json: JSON.stringify({ text: streamBlock.content[0].content }),
+            };
+          }
+          break;
+        }
+
+        case 'select': {
+          return {
+            type: 'json',
+            disposition: 'select',
+            json: JSON.stringify({
+              options: streamBlock.content.flatMap((content) =>
+                content.type === 'tag' && content.content.length === 1 && content.content[0].type === 'text'
+                  ? [content.content[0].content]
+                  : [],
+              ),
+            }),
+          };
+        }
+      }
+
+      break;
+    }
+
+    //
+    // JSON
+    //
+    case 'json': {
+      switch (contentBlock.type) {
+        case 'tool_use': {
+          return { ...contentBlock, input: safeParseJson(streamBlock.content) ?? {} };
+        }
+      }
+
+      return { ...contentBlock, type: 'json', disposition: 'artifact', json: streamBlock.content };
+    }
+  }
+};

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

@@ -0,0 +1,101 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { describe, it } from 'vitest';
+
+import { type Message } from '@dxos/artifact';
+import { log } from '@dxos/log';
+
+import { MixedStreamParser } from './parser';
+import { createGenerationStream } from './stream';
+import { createTestSSEStream } from './testing';
+import { type GenerationStreamEvent } from './types';
+
+// TODO(burdon): Tool request/response.
+
+const TEST_BLOCKS = [
+  [
+    // Text
+    'some text',
+    'on multiple lines.',
+
+    // XML
+    '<cot>',
+    '  1. Analyze content.',
+    '  2. Create plan.',
+    '  3. Generate artifacts.',
+    '</cot>',
+
+    // Text
+    'some more text.',
+
+    // XML
+    '<artifact id="123" />',
+
+    // XML
+    '<artifact id="456" />',
+
+    // Text
+    'and some more text.',
+
+    '<suggest>Plan me a trip to the Rio.</suggest>',
+    '<suggest>Show me steak recipes.</suggest>',
+    '<select><option>Yes</option><option>No</option></select>',
+  ].join('\n'),
+
+  // JSON
+  {
+    type: 'input',
+    value: {
+      level: 1,
+      optional: true,
+    },
+  },
+];
+
+describe('GenerationStream', () => {
+  it('should generate a stream of blocks', async ({ expect }) => {
+    const stream = createGenerationStream(new Response(createTestSSEStream(TEST_BLOCKS)));
+    const events: GenerationStreamEvent[] = [];
+    for await (const event of stream) {
+      events.push(event);
+    }
+
+    expect(events.map((event) => event.type === 'content_block_start').filter(Boolean)).to.have.length(2);
+  });
+
+  for (const splitBy of ['word'] as const) {
+    it(`should emit xml blocks (splitBy: ${splitBy})`, async ({ expect }) => {
+      const stream = createGenerationStream(new Response(createTestSSEStream(TEST_BLOCKS, { splitBy })));
+      const parser = new MixedStreamParser();
+
+      let message: Message | undefined;
+      parser.message.on((_message) => {
+        message = _message;
+      });
+
+      await parser.parse(stream);
+
+      for (const block of message?.content ?? []) {
+        log.info('block', { block });
+      }
+      expect(message?.content.map((block) => block.type)).to.deep.eq([
+        //
+        'text',
+        'text',
+        'text',
+        'json',
+        'text',
+        'json',
+        'text',
+        'json',
+        'text',
+        'json',
+        'text',
+        'json',
+        'json',
+      ]);
+    });
+  }
+});