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

package/src/ai-service/stream.ts

@@ -0,0 +1,116 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { Trigger } from '@dxos/async';
+import { log } from '@dxos/log';
+
+import { type GenerationStreamEvent } from './types';
+import { iterSSEMessages } from './util';
+
+/**
+ * Creates a stream from an SSE response.
+ */
+export const createGenerationStream = (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);
+};
+
+/**
+ * Server-Sent Events (SSE) stream from the AI service.
+ * https://html.spec.whatwg.org/multipage/server-sent-events.html#server-sent-events
+ */
+export class GenerationStream implements AsyncIterable<GenerationStreamEvent> {
+  /**
+   * Trigger event when the stream is done.
+   */
+  private readonly _done = new Trigger();
+
+  /**
+   * Iterator over the stream.
+   */
+  private _iterator?: AsyncIterator<GenerationStreamEvent> = undefined;
+
+  constructor(
+    private readonly _controller: AbortController,
+    private readonly _getIterator: () => AsyncIterableIterator<GenerationStreamEvent>,
+  ) {}
+
+  /**
+   * Returns an iterator over the stream.
+   */
+  [Symbol.asyncIterator](): AsyncIterator<GenerationStreamEvent> {
+    return this._createIterator();
+  }
+
+  /**
+   * Aborts the stream.
+   */
+  abort() {
+    this._controller.abort();
+  }
+
+  /**
+   * Returns the complete message list generated by the AI service.
+   */
+  async complete() {
+    await this._done.wait();
+  }
+
+  /**
+   * Creates an iterator over the stream.
+   */
+  private _createIterator(): AsyncIterator<GenerationStreamEvent> {
+    const self = this;
+    return (this._iterator ??= (() => {
+      const generator = async function* (this: GenerationStream) {
+        try {
+          for await (const event of self._getIterator()) {
+            yield event;
+          }
+          this._done.wake();
+        } catch (err: any) {
+          this._done.throw(err);
+        }
+      };
+
+      return generator.call(this);
+    })());
+  }
+}

package/src/ai-service/testing.ts

@@ -0,0 +1,93 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { ObjectId } from '@dxos/echo-schema';
+
+import { type GenerationStreamEvent } from './types';
+
+/**
+ * Mock server-side events (SSE) stream.
+ */
+export const createTestSSEStream = (
+  pars: (string | object)[],
+  { splitBy = 'word' }: { splitBy?: 'word' | 'character' } = {},
+): ReadableStream => {
+  const encoder = new TextEncoder();
+
+  return new ReadableStream({
+    start: (controller) => {
+      const push = (event: GenerationStreamEvent) => {
+        controller.enqueue(encoder.encode(`event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`));
+      };
+
+      push({
+        type: 'message_start',
+        message: {
+          id: ObjectId.random(),
+          role: 'assistant',
+          content: [],
+        },
+      });
+
+      for (const block of pars) {
+        let index = 0;
+        if (typeof block === 'string') {
+          push({
+            type: 'content_block_start',
+            index,
+            content: {
+              type: 'text',
+              text: '',
+            },
+          });
+
+          for (const word of block.split(splitBy === 'word' ? /([ \t\n]+)/ : '')) {
+            push({
+              type: 'content_block_delta',
+              index,
+              delta: {
+                type: 'text_delta',
+                text: word,
+              },
+            });
+            index += word.length + 1;
+          }
+        } else {
+          push({
+            type: 'content_block_start',
+            index,
+            content: {
+              type: 'json',
+              json: '',
+            },
+          });
+
+          const text = JSON.stringify(block, null, 2);
+          for (const line of text.split('\n')) {
+            push({
+              type: 'content_block_delta',
+              index,
+              delta: {
+                type: 'input_json_delta',
+                partial_json: line,
+              },
+            });
+            index += line.length + 1;
+          }
+        }
+
+        push({
+          type: 'content_block_stop',
+          index,
+        });
+      }
+
+      push({
+        type: 'message_stop',
+      });
+
+      controller.close();
+    },
+  });
+};

package/src/ai-service/transform.ts

@@ -0,0 +1,121 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as P from 'parsimmon';
+
+/**
+ * Parsed block.
+ */
+export type StreamBlock =
+  | {
+      type: 'tag';
+      tag: string;
+      attributes?: Record<string, string>;
+      content: StreamBlock[];
+      closing?: boolean;
+      selfClosing?: boolean;
+    }
+  | { type: 'text'; content: string }
+  | { type: 'json'; disposition?: string; content: string };
+
+const textChunk: P.Parser<StreamBlock> = P.regexp(/[^<]+/).map((content) => ({
+  type: 'text',
+  content,
+}));
+
+const attribute: P.Parser<[string, string]> = P.seq(
+  P.regexp(/[a-zA-Z_][a-zA-Z0-9_-]*/),
+  P.string('=').trim(P.optWhitespace),
+  P.alt(
+    P.regexp(/"[^"]*"/).map((s) => s.slice(1, -1)),
+    P.regexp(/'[^']*'/).map((s) => s.slice(1, -1)),
+  ),
+).map(([name, _, value]) => [name, value]);
+
+const tagName = P.regexp(/[a-zA-Z_][a-zA-Z0-9_-]*/);
+
+const selfClosingTag: P.Parser<StreamBlock> = P.seqMap(
+  //
+  P.string('<'),
+  tagName,
+  P.optWhitespace.then(attribute).many(),
+  P.optWhitespace.then(P.string('/>')),
+  P.optWhitespace,
+  (_, tag, attributes, __) => ({
+    type: 'tag',
+    tag,
+    attributes: Object.fromEntries(attributes),
+    content: [],
+    selfClosing: true,
+  }),
+);
+
+const openTag: P.Parser<StreamBlock> = P.seqMap(
+  //
+  P.string('<'),
+  tagName,
+  P.optWhitespace.then(attribute).many(),
+  P.string('>'),
+  P.regexp(/[^<]*/),
+  (_, tag, attributes, __, content) => ({
+    type: 'tag',
+    tag,
+    attributes: Object.fromEntries(attributes),
+    content: content.length ? [{ type: 'text', content }] : [],
+  }),
+);
+
+const closeTag: P.Parser<StreamBlock> = P.seqMap(
+  //
+  P.string('</'),
+  tagName,
+  P.string('>'),
+  P.optWhitespace,
+  (_, tag) => ({
+    type: 'tag',
+    tag,
+    content: [],
+    closing: true,
+  }),
+);
+
+const mixedChunk: P.Parser<StreamBlock> = P.alt(selfClosingTag, openTag, closeTag, textChunk);
+
+/**
+ * Permissive streaming transformer that processes mixed content (plain text and XML fragments) from the AI service.
+ */
+export class StreamTransform {
+  private _buffer = '';
+
+  transform(chunk: string): StreamBlock[] {
+    this._buffer += chunk;
+
+    const results: StreamBlock[] = [];
+    const parser = mixedChunk.many().skip(P.optWhitespace);
+    const result = parser.parse(this._buffer);
+    if (result.status) {
+      for (const chunk of result.value) {
+        results.push(chunk);
+      }
+
+      this._buffer = '';
+    }
+
+    return results;
+  }
+
+  // TODO(burdon): Process remaining buffer?
+  flush(): StreamBlock[] {
+    const remaining = this._buffer;
+    this._buffer = '';
+    return remaining.length > 0
+      ? [
+          {
+            type: 'text',
+            content: remaining,
+          },
+        ]
+      : [];
+  }
+}

package/src/ai-service/types.ts

@@ -0,0 +1,107 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import { Schema as S } from '@effect/schema';
+
+import { Tool, Message, type MessageContentBlock, SpaceIdSchema } from '@dxos/artifact';
+import { ObjectId } from '@dxos/echo-schema';
+
+export const createArtifactElement = (id: ObjectId) => `<artifact id=${id} />`;
+
+export const LLMModel = S.Literal(
+  '@anthropic/claude-3-5-haiku-20241022',
+  '@anthropic/claude-3-5-sonnet-20241022',
+  '@cf/deepseek-ai/deepseek-r1-distill-qwen-32b',
+  '@hf/nousresearch/hermes-2-pro-mistral-7b',
+  '@ollama/llama-3-2-3b',
+  '@ollama/llama-3-1-nemotron-70b-instruct',
+  '@ollama/llama-3-1-nemotron-mini-4b-instruct',
+);
+
+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',
+});
+
+/**
+ * Client GPT request.
+ */
+export const GenerateRequest = S.Struct({
+  spaceId: S.optional(SpaceIdSchema),
+  threadId: S.optional(ObjectId),
+
+  // TODO(burdon): Make optional with default.
+  model: LLMModel,
+
+  /**
+   * Current request.
+   */
+  // TODO(burdon): Remove from history.
+  // prompt: Message,
+
+  /**
+   * System instructions to the LLM.
+   */
+  systemPrompt: S.optional(S.String),
+
+  /**
+   * History of messages to include in the context window.
+   */
+  history: S.optional(S.Array(Message)),
+
+  /**
+   * Tools available for the LLM.
+   */
+  tools: S.optional(S.Array(Tool).pipe(S.mutable)),
+});
+
+export type GenerateRequest = S.Schema.Type<typeof GenerateRequest>;
+
+/**
+ * Server-Sent Events (SSE) stream from the AI service.
+ * https://docs.anthropic.com/en/api/streaming
+ * https://platform.openai.com/docs/api-reference/streaming
+ * https://html.spec.whatwg.org/multipage/server-sent-events.html#server-sent-events
+ */
+// TODO(dmaretskyi): Effect schema.
+export type GenerationStreamEvent =
+  | {
+      // TODO(dmaretskyi): Normalize types to our schema.
+      type: 'message_start';
+      message: Omit<Message, 'spaceId' | 'threadId'>;
+    }
+  | {
+      type: 'message_delta';
+      delta: {
+        stopReason: 'tool_use' | 'end_turn';
+      };
+    }
+  | {
+      type: 'message_stop';
+    }
+  | {
+      type: 'content_block_start';
+      index: number;
+      content: MessageContentBlock;
+    }
+  | {
+      type: 'content_block_delta';
+      index: number;
+      delta:
+        | {
+            type: 'text_delta';
+            text: string;
+          }
+        | {
+            type: 'input_json_delta';
+            partial_json: string;
+          };
+    }
+  | {
+      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,13 +2,18 @@
 // 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 { log } from '@dxos/log';
 
-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 +23,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 +45,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',
@@ -69,7 +68,7 @@ while (true) {
       onImage: (img) => {
         const path = `/tmp/image-${img.id}.jpeg`;
         writeFileSync(path, Buffer.from(img.source.data, 'base64'));
-        console.log(`Saved image to ${path}`);
+        log.info('Saved image', { 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

@@ -5,34 +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';
 
+// TODO(burdon): Local live LLM test.
 describe.skip('Conversation tests', () => {
   const client = new AIServiceClientImpl({
-    endpoint: ENDPOINT,
+    endpoint: AI_SERVICE_ENDPOINT.LOCAL,
   });
 
-  test('hello', async ({ expect }) => {
+  test('basic', 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 });
+    expect(result.history.length).to.equal(2);
   });
 
   test('tool call', async ({ expect }) => {
@@ -44,9 +45,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 +55,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 });
     }
   };
 });