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

package/src/ai-service/parser.ts

@@ -0,0 +1,183 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export type Block = TextBlock | TagBlock | CodeBlock;
+
+export type TextBlock = {
+  type: 'text';
+  content: string;
+};
+
+export type TagBlock = {
+  type: 'tag';
+  name: string;
+  content: string;
+  attributes: Record<string, string>;
+  selfClosing: boolean;
+};
+
+export type CodeBlock = {
+  type: 'code';
+  language: string;
+  content: string;
+};
+
+/**
+ * Permissive streaming parser for XML content and plain text.
+ */
+export class StreamingParser {
+  private buffer = '';
+
+  constructor(private readonly onBlock: (block: Block) => void) {}
+
+  /**
+   * Write a chunk of text to the parser.
+   */
+  write(chunk: string) {
+    this.buffer += chunk;
+    this.processBuffer();
+  }
+
+  /**
+   * Signal that no more input will be written.
+   */
+  end() {
+    if (this.buffer.length > 0) {
+      this.emitTextBlock(this.buffer);
+      this.buffer = '';
+    }
+  }
+
+  private processBuffer() {
+    while (this.buffer.length > 0) {
+      // Try to parse code block first.
+      if (this.tryParseCodeBlock()) {
+        continue;
+      }
+
+      // If no tag start found, keep buffering.
+      const tagStartIndex = this.buffer.indexOf('<');
+      if (tagStartIndex === -1) {
+        return;
+      }
+
+      // Emit any text before the tag.
+      if (tagStartIndex > 0) {
+        this.emitTextBlock(this.buffer.slice(0, tagStartIndex));
+        this.buffer = this.buffer.slice(tagStartIndex);
+        continue;
+      }
+
+      // If incomplete tag, keep buffering.
+      const tagEndIndex = this.buffer.indexOf('>', tagStartIndex);
+      if (tagEndIndex === -1) {
+        return;
+      }
+
+      // Check if this is a self-closing tag.
+      const isSelfClosing = this.buffer[tagEndIndex - 1] === '/';
+      const actualTagEndIndex = isSelfClosing ? tagEndIndex - 1 : tagEndIndex;
+
+      // Parse the tag name and attributes.
+      const tagContent = this.buffer.slice(tagStartIndex + 1, actualTagEndIndex).trim();
+      const [tagName, ...attrParts] = tagContent.split(/\s+/);
+      const attributes = this.parseAttributes(attrParts.join(' '));
+
+      if (isSelfClosing) {
+        this.emitTagBlock(tagName, '', attributes, true);
+        this.buffer = this.buffer.slice(tagEndIndex + 1);
+        continue;
+      }
+
+      // If no closing tag yet, keep buffering.
+      const closingTagIndex = this.buffer.indexOf('</' + tagName + '>', tagEndIndex);
+      if (closingTagIndex === -1) {
+        return;
+      }
+
+      // Complete tag.
+      const content = this.buffer.slice(tagEndIndex + 1, closingTagIndex);
+      this.emitTagBlock(tagName, content, attributes, false);
+      this.buffer = this.buffer.slice(closingTagIndex + tagName.length + 3); // +3 for '</>'
+    }
+  }
+
+  private tryParseCodeBlock(): boolean {
+    const codeBlockStart = this.buffer.indexOf('```');
+    if (codeBlockStart === -1) {
+      return false;
+    }
+
+    // If there's text before the code block, emit it first.
+    if (codeBlockStart > 0) {
+      this.emitTextBlock(this.buffer.slice(0, codeBlockStart));
+      this.buffer = this.buffer.slice(codeBlockStart);
+      return true;
+    }
+
+    // Look for the closing fence.
+    const nextNewline = this.buffer.indexOf('\n', codeBlockStart);
+    if (nextNewline === -1) {
+      return false; // Incomplete, keep buffering.
+    }
+
+    const codeBlockEnd = this.buffer.indexOf('\n```', nextNewline);
+    if (codeBlockEnd === -1) {
+      return false; // No closing fence yet, keep buffering.
+    }
+
+    // Extract language (if specified) and content.
+    const languageLine = this.buffer.slice(3, nextNewline).trim();
+    const content = this.buffer.slice(nextNewline + 1, codeBlockEnd);
+
+    this.emitCodeBlock(languageLine, content);
+    this.buffer = this.buffer.slice(codeBlockEnd + 4); // +4 for '\n```'
+    return true;
+  }
+
+  private parseAttributes(attrString: string): Record<string, string> {
+    const attributes: Record<string, string> = {};
+    if (!attrString) {
+      return attributes;
+    }
+
+    // Match attribute patterns like: name="value" or name='value' or name=value.
+    const attrPattern = /(\w+)\s*=\s*(?:"([^"]*)"|'([^']*)'|(\S+))/g;
+    let match;
+
+    while ((match = attrPattern.exec(attrString)) !== null) {
+      const [, name, doubleQuoted, singleQuoted, unquoted] = match;
+      attributes[name] = doubleQuoted ?? singleQuoted ?? unquoted;
+    }
+
+    return attributes;
+  }
+
+  private emitTextBlock(text: string) {
+    if (text.trim().length > 0) {
+      this.onBlock({
+        type: 'text',
+        content: text,
+      });
+    }
+  }
+
+  private emitTagBlock(name: string, content: string, attributes: Record<string, string>, selfClosing: boolean) {
+    this.onBlock({
+      type: 'tag',
+      name,
+      content: content.trim(),
+      attributes,
+      selfClosing,
+    });
+  }
+
+  private emitCodeBlock(language: string, content: string) {
+    this.onBlock({
+      type: 'code',
+      language,
+      content: content.trim(),
+    });
+  }
+}

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

@@ -0,0 +1,106 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { describe, it } from 'vitest';
+
+import { type Block, StreamingParser } from './parser';
+import { GenerationStream } from './stream';
+/** @ts-ignore */
+import TEXT from './stream.test.txt?raw';
+
+type Part = { event: string; data: any };
+
+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();
+  });
+
+  it('should parse mixed content.', ({ expect }) => {
+    const blocks: Block[] = [];
+    const parser = new StreamingParser((block) => {
+      blocks.push(block);
+    });
+
+    const words = TEXT.split(' ');
+    for (const word of words) {
+      parser.write(word + ' ');
+    }
+    parser.end();
+
+    expect(blocks).to.have.length(7);
+    expect(blocks.map((block) => block.type)).to.deep.eq(['text', 'tag', 'text', 'tag', 'text', 'code', 'code']);
+    expect(blocks.map((block) => block.type === 'tag' && block.name).filter(Boolean)).to.deep.eq(['cot', 'artifact']);
+  });
+});
+
+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();
+    },
+  });
+};

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

@@ -0,0 +1,33 @@
+{
+  "test": 100
+}
+
+<cot>
+1. Planning step.
+2. Research step.
+3. Output step.
+</cot>
+
+This is some text.
+
+This is a tool call.
+
+{
+  "test": 100
+}
+
+<artifact id="100" />
+
+This is some code.
+
+```json
+{
+  "test": 100
+}
+```
+
+```ts
+const Test = () => {
+  return <div>Hello</div>;
+}
+```

package/src/ai-service/stream.ts

@@ -0,0 +1,182 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { type Message } from '@dxos/artifact';
+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 ResultStreamEvent } from './defs';
+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/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,13 +2,17 @@
 // Copyright 2024 DXOS.org
 //
 
+import { writeFileSync } from 'fs';
 import inquirer from 'inquirer';
 
+import { createUserMessage } from '@dxos/artifact';
+import { ObjectId } from '@dxos/echo-schema';
 import { SpaceId } from '@dxos/keys';
 
-import { AIServiceClientImpl, ObjectId, ToolTypes } from '../ai-service';
-import { runLLM, createUserMessage } from '../conversation';
+import { AIServiceClientImpl, ToolTypes } from '../ai-service';
+import { runLLM } from '../conversation';
 import {
+  AI_SERVICE_ENDPOINT,
   createLogger,
   createCypherTool,
   createSystemPrompt,
@@ -18,15 +22,9 @@ import {
   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,
+  endpoint: AI_SERVICE_ENDPOINT.LOCAL,
 });
 
 const dataSource = createTestData();
@@ -46,7 +44,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 { createUserMessage, defineTool, ToolResult } from '@dxos/artifact';
+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',
@@ -44,9 +42,9 @@ describe.skip('Conversation tests', () => {
       }),
       execute: async ({ magicWord }) => {
         if (magicWord === 'pretty please') {
-          return LLMToolResult.Success('The password is: "The sky is gray"');
+          return ToolResult.Success('The password is: "The sky is gray"');
         } else {
-          return LLMToolResult.Error('Wrong magic word');
+          return ToolResult.Error('Wrong 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 });
     }
   };
 });