@dxos/assistant 0.8.0 → 0.8.1-main.013e445

package/dist/types/src/testing/query-promts.d.ts

@@ -1,4 +1,4 @@
-import { Schema as S } from '@effect/schema';
+import { Schema as S } from 'effect';
 import { type DataSource } from '../cypher';
 export declare const createCypherTool: (dataSource: DataSource, schemaTypes?: S.Schema.Any[]) => import("@dxos/artifact").Tool;
 export declare const createSystemPrompt: (schemaTypes: S.Schema.Any[]) => string;

package/dist/types/src/testing/query-promts.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"query-promts.d.ts","sourceRoot":"","sources":["../../../../src/testing/query-promts.ts"],"names":[],"mappings":"AAIA,OAAO,EAAE,MAAM,IAAI,CAAC,EAAE,MAAM,gBAAgB,CAAC;AAM7C,OAAO,EAAwC,KAAK,UAAU,EAAE,MAAM,WAAW,CAAC;AAGlF,eAAO,MAAM,gBAAgB,eAAgB,UAAU,gBAAe,CAAC,CAAC,MAAM,CAAC,GAAG,EAAE,kCA0ChF,CAAC;AAEL,eAAO,MAAM,kBAAkB,gBAAiB,CAAC,CAAC,MAAM,CAAC,GAAG,EAAE,WAiB7D,CAAC"}
+{"version":3,"file":"query-promts.d.ts","sourceRoot":"","sources":["../../../../src/testing/query-promts.ts"],"names":[],"mappings":"AAIA,OAAO,EAAE,MAAM,IAAI,CAAC,EAAE,MAAM,QAAQ,CAAC;AAMrC,OAAO,EAAwC,KAAK,UAAU,EAAE,MAAM,WAAW,CAAC;AAGlF,eAAO,MAAM,gBAAgB,eAAgB,UAAU,gBAAe,CAAC,CAAC,MAAM,CAAC,GAAG,EAAE,kCA0ChF,CAAC;AAEL,eAAO,MAAM,kBAAkB,gBAAiB,CAAC,CAAC,MAAM,CAAC,GAAG,EAAE,WAiB7D,CAAC"}

package/dist/types/src/testing/test-data.d.ts

@@ -1,4 +1,4 @@
-import { type Schema as S } from '@effect/schema';
+import { type Schema as S } from 'effect';
 import type { EchoDatabase } from '@dxos/echo-db';
 import { type JsonSchemaType } from '@dxos/echo-schema';
 import { type DataSource, type Node, type Relationship } from '../cypher';

package/dist/types/src/testing/test-data.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"test-data.d.ts","sourceRoot":"","sources":["../../../../src/testing/test-data.ts"],"names":[],"mappings":"AAIA,OAAO,EAAE,KAAK,MAAM,IAAI,CAAC,EAAE,MAAM,gBAAgB,CAAC;AAGlD,OAAO,KAAK,EAAE,YAAY,EAAE,MAAM,eAAe,CAAC;AAClD,OAAO,EAAmD,KAAK,cAAc,EAAY,MAAM,mBAAmB,CAAC;AAInH,OAAO,EACL,KAAK,UAAU,EACf,KAAK,IAAI,EACT,KAAK,YAAY,EAGlB,MAAM,WAAW,CAAC;AAGnB,eAAO,MAAM,YAAY,OAAQ,YAAY,SAoH5C,CAAC;AAGF,eAAO,MAAM,cAAc,QAAO,cAuGjC,CAAC;AAEF,qBAAa,cAAe,YAAW,UAAU;IAC/C,OAAO,EAAE,MAAM,CACb,MAAM,EACN;QACE,MAAM,EAAE,CAAC,CAAC,MAAM,CAAC,GAAG,CAAC;QACrB,OAAO,EAAE,MAAM,CAAC;QAChB,EAAE,EAAE,MAAM,CAAC;QACX,IAAI,EAAE,OAAO,CAAC;KACf,CACF,CAAM;IAEP,MAAM,EAAE,MAAM,CACZ,MAAM,EACN;QACE,IAAI,EAAE,cAAc,CAAC;QACrB,QAAQ,EAAE,CAAC,CAAC,MAAM,CAAC,GAAG,CAAC;KACxB,CACF,CAAM;IAEP,KAAK,EAAE,IAAI,EAAE,CAAM;IACnB,aAAa,EAAE,YAAY,EAAE,CAAM;IAE7B,QAAQ,CAAC,EAAE,KAAK,EAAE,EAAE;QAAE,KAAK,CAAC,EAAE,MAAM,CAAA;KAAE;IAItC,gBAAgB,CAAC,EAAE,KAAK,EAAE,EAAE;QAAE,KAAK,CAAC,EAAE,MAAM,CAAA;KAAE;IAIpD,GAAG,CAAC,CAAC,SAAS,CAAC,CAAC,MAAM,CAAC,GAAG,EAAE,MAAM,EAAE,CAAC,EAAE,IAAI,EAAE,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,CAAC,CAAC;IAsB7D,YAAY;CA0Eb"}
+{"version":3,"file":"test-data.d.ts","sourceRoot":"","sources":["../../../../src/testing/test-data.ts"],"names":[],"mappings":"AAIA,OAAO,EAAE,KAAK,MAAM,IAAI,CAAC,EAAE,MAAM,QAAQ,CAAC;AAG1C,OAAO,KAAK,EAAE,YAAY,EAAE,MAAM,eAAe,CAAC;AAClD,OAAO,EAAmD,KAAK,cAAc,EAAY,MAAM,mBAAmB,CAAC;AAInH,OAAO,EACL,KAAK,UAAU,EACf,KAAK,IAAI,EACT,KAAK,YAAY,EAGlB,MAAM,WAAW,CAAC;AAGnB,eAAO,MAAM,YAAY,OAAQ,YAAY,SAoH5C,CAAC;AAGF,eAAO,MAAM,cAAc,QAAO,cAuGjC,CAAC;AAEF,qBAAa,cAAe,YAAW,UAAU;IAC/C,OAAO,EAAE,MAAM,CACb,MAAM,EACN;QACE,MAAM,EAAE,CAAC,CAAC,MAAM,CAAC,GAAG,CAAC;QACrB,OAAO,EAAE,MAAM,CAAC;QAChB,EAAE,EAAE,MAAM,CAAC;QACX,IAAI,EAAE,OAAO,CAAC;KACf,CACF,CAAM;IAEP,MAAM,EAAE,MAAM,CACZ,MAAM,EACN;QACE,IAAI,EAAE,cAAc,CAAC;QACrB,QAAQ,EAAE,CAAC,CAAC,MAAM,CAAC,GAAG,CAAC;KACxB,CACF,CAAM;IAEP,KAAK,EAAE,IAAI,EAAE,CAAM;IACnB,aAAa,EAAE,YAAY,EAAE,CAAM;IAE7B,QAAQ,CAAC,EAAE,KAAK,EAAE,EAAE;QAAE,KAAK,CAAC,EAAE,MAAM,CAAA;KAAE;IAItC,gBAAgB,CAAC,EAAE,KAAK,EAAE,EAAE;QAAE,KAAK,CAAC,EAAE,MAAM,CAAA;KAAE;IAIpD,GAAG,CAAC,CAAC,SAAS,CAAC,CAAC,MAAM,CAAC,GAAG,EAAE,MAAM,EAAE,CAAC,EAAE,IAAI,EAAE,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,CAAC,CAAC;IAsB7D,YAAY;CA0Eb"}

package/dist/types/src/testing/test-schema.d.ts

@@ -1,4 +1,4 @@
-import { Schema as S } from '@effect/schema';
+import { Schema as S } from 'effect';
 export declare const ObjectId: S.SchemaClass<string, string, never>;
 export declare const Contact: S.mutable<S.Struct<{
     id: S.SchemaClass<string, string, never>;

package/dist/types/src/testing/test-schema.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"test-schema.d.ts","sourceRoot":"","sources":["../../../../src/testing/test-schema.ts"],"names":[],"mappings":"AAIA,OAAO,EAAE,MAAM,IAAI,CAAC,EAAE,MAAM,gBAAgB,CAAC;AAQ7C,eAAO,MAAM,QAAQ,sCAAyE,CAAC;AAE/F,eAAO,MAAM,OAAO;;;;GAahB,CAAC;AACL,MAAM,WAAW,OAAQ,SAAQ,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,OAAO,CAAC;CAAG;AAEjE,eAAO,MAAM,OAAO;;;;GAahB,CAAC;AACL,MAAM,WAAW,OAAQ,SAAQ,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,OAAO,CAAC;CAAG;AAEjE,eAAO,MAAM,IAAI;;;;;;;;;;;;;;GAeb,CAAC;AACL,MAAM,WAAW,IAAK,SAAQ,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,IAAI,CAAC;CAAG;AAE3D,eAAO,MAAM,GAAG;;;;;;;;;;;;;GAcZ,CAAC;AACL,MAAM,WAAW,GAAI,SAAQ,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,GAAG,CAAC;CAAG"}
+{"version":3,"file":"test-schema.d.ts","sourceRoot":"","sources":["../../../../src/testing/test-schema.ts"],"names":[],"mappings":"AAIA,OAAO,EAAE,MAAM,IAAI,CAAC,EAAE,MAAM,QAAQ,CAAC;AAQrC,eAAO,MAAM,QAAQ,sCAAyE,CAAC;AAE/F,eAAO,MAAM,OAAO;;;;GAahB,CAAC;AACL,MAAM,WAAW,OAAQ,SAAQ,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,OAAO,CAAC;CAAG;AAEjE,eAAO,MAAM,OAAO;;;;GAahB,CAAC;AACL,MAAM,WAAW,OAAQ,SAAQ,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,OAAO,CAAC;CAAG;AAEjE,eAAO,MAAM,IAAI;;;;;;;;;;;;;;GAeb,CAAC;AACL,MAAM,WAAW,IAAK,SAAQ,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,IAAI,CAAC;CAAG;AAE3D,eAAO,MAAM,GAAG;;;;;;;;;;;;;GAcZ,CAAC;AACL,MAAM,WAAW,GAAI,SAAQ,CAAC,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,GAAG,CAAC;CAAG"}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dxos/assistant",
-  "version": "0.8.0",
+  "version": "0.8.1-main.013e445",
   "description": "Natural language interface to DXOS",
   "homepage": "https://dxos.org",
   "bugs": "https://github.com/dxos/dxos/issues",
@@ -37,27 +37,30 @@
     "src"
   ],
   "dependencies": {
-    "@effect/schema": "^0.75.5",
     "parjs": "^1.3.9",
     "parsimmon": "^1.18.1",
     "ulidx": "^2.3.0",
-    "@dxos/debug": "0.8.0",
-    "@dxos/context": "0.8.0",
-    "@dxos/artifact": "0.8.0",
-    "@dxos/echo-db": "0.8.0",
-    "@dxos/async": "0.8.0",
-    "@dxos/invariant": "0.8.0",
-    "@dxos/echo-schema": "0.8.0",
-    "@dxos/keys": "0.8.0",
-    "@dxos/live-object": "0.8.0",
-    "@dxos/log": "0.8.0",
-    "@dxos/node-std": "0.8.0",
-    "@dxos/util": "0.8.0"
+    "@dxos/artifact": "0.8.1-main.013e445",
+    "@dxos/async": "0.8.1-main.013e445",
+    "@dxos/echo-db": "0.8.1-main.013e445",
+    "@dxos/echo-schema": "0.8.1-main.013e445",
+    "@dxos/context": "0.8.1-main.013e445",
+    "@dxos/keys": "0.8.1-main.013e445",
+    "@dxos/live-object": "0.8.1-main.013e445",
+    "@dxos/debug": "0.8.1-main.013e445",
+    "@dxos/log": "0.8.1-main.013e445",
+    "@dxos/node-std": "0.8.1-main.013e445",
+    "@dxos/util": "0.8.1-main.013e445",
+    "@dxos/invariant": "0.8.1-main.013e445"
   },
   "devDependencies": {
     "@types/parsimmon": "^1.10.9",
+    "effect": "3.13.3",
     "inquirer": "^12.0.0"
   },
+  "peerDependencies": {
+    "effect": "^3.13.3"
+  },
   "publishConfig": {
     "access": "public"
   },

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

@@ -2,25 +2,27 @@
 // Copyright 2024 DXOS.org
 //
 
-import { Schema as S } from '@effect/schema';
+import { Schema as S } from 'effect';
 import { test, describe } from 'vitest';
 
-import { type Message, type Tool } from '@dxos/artifact';
-import { toJsonSchema, ObjectId } from '@dxos/echo-schema';
+import { defineTool, Message, ToolResult, type Tool } from '@dxos/artifact';
+import { toJsonSchema, ObjectId, createStatic } from '@dxos/echo-schema';
 import { invariant } from '@dxos/invariant';
 import { SpaceId } from '@dxos/keys';
 import { log } from '@dxos/log';
 
-import { AIServiceClientImpl } from './client';
-import { DEFAULT_LLM_MODEL } from './defs';
+import { DEFAULT_EDGE_MODEL } from './defs';
+import { AIServiceEdgeClient } from './edge-client';
+import { OllamaClient } from './ollama-client';
+import { MixedStreamParser } from './parser';
 import { ToolTypes } from './types';
-import { AI_SERVICE_ENDPOINT } from '../testing';
+import { createTestOllamaClient, AI_SERVICE_ENDPOINT } from '../testing';
 
 // log.config({ filter: 'debug' });
 
 describe.skip('AI Service Client', () => {
   test('client generation', async () => {
-    const client = new AIServiceClientImpl({
+    const client = new AIServiceEdgeClient({
       endpoint: AI_SERVICE_ENDPOINT.LOCAL,
     });
 
@@ -38,7 +40,7 @@ describe.skip('AI Service Client', () => {
     ]);
 
     const stream = await client.exec({
-      model: DEFAULT_LLM_MODEL,
+      model: DEFAULT_EDGE_MODEL,
       spaceId,
       threadId,
       systemPrompt: 'You are a poet',
@@ -54,7 +56,7 @@ describe.skip('AI Service Client', () => {
   });
 
   test('tool calls', async () => {
-    const client = new AIServiceClientImpl({
+    const client = new AIServiceEdgeClient({
       endpoint: AI_SERVICE_ENDPOINT.LOCAL,
     });
 
@@ -83,7 +85,7 @@ describe.skip('AI Service Client', () => {
 
     {
       const stream1 = await client.exec({
-        model: DEFAULT_LLM_MODEL,
+        model: DEFAULT_EDGE_MODEL,
         spaceId,
         threadId,
         systemPrompt: 'You are a helpful assistant.',
@@ -117,7 +119,7 @@ describe.skip('AI Service Client', () => {
 
     {
       const stream2 = await client.exec({
-        model: DEFAULT_LLM_MODEL,
+        model: DEFAULT_EDGE_MODEL,
         spaceId,
         threadId,
         systemPrompt: 'You are a helpful assistant.',
@@ -137,7 +139,7 @@ describe.skip('AI Service Client', () => {
   });
 
   test.skip('image generation', async () => {
-    const client = new AIServiceClientImpl({
+    const client = new AIServiceEdgeClient({
       endpoint: AI_SERVICE_ENDPOINT.LOCAL,
     });
 
@@ -155,7 +157,7 @@ describe.skip('AI Service Client', () => {
     ]);
 
     const stream = await client.exec({
-      model: DEFAULT_LLM_MODEL,
+      model: DEFAULT_EDGE_MODEL,
       spaceId,
       threadId,
       tools: [
@@ -176,3 +178,88 @@ describe.skip('AI Service Client', () => {
     log('full message', { message: await stream.complete() });
   });
 });
+
+describe.skip('Ollama Client', () => {
+  test('basic', async (ctx) => {
+    const isRunning = await OllamaClient.isRunning();
+    if (!isRunning) {
+      ctx.skip();
+    }
+
+    const client = createTestOllamaClient();
+    const parser = new MixedStreamParser();
+
+    const messages = await parser.parse(
+      await client.exec({
+        prompt: createStatic(Message, {
+          role: 'user',
+          content: [{ type: 'text', text: 'Hello, world!' }],
+        }),
+      }),
+    );
+
+    log.info('messages', { messages });
+  });
+
+  test('tool calls', async (ctx) => {
+    const isRunning = await OllamaClient.isRunning();
+    if (!isRunning) {
+      ctx.skip();
+    }
+
+    const client = createTestOllamaClient({
+      tools: [
+        defineTool('test', {
+          name: 'encrypt',
+          description: 'Encrypt a message',
+          schema: S.Struct({
+            message: S.String.annotations({ description: 'The message to encrypt' }),
+          }),
+          execute: async ({ message }) => ToolResult.Success(message.split('').reverse().join('')),
+        }),
+      ],
+    });
+    const parser = new MixedStreamParser();
+    parser.streamEvent.on((event) => {
+      // log.info('event', { event });
+    });
+
+    const messages = await parser.parse(
+      await client.exec({
+        prompt: createStatic(Message, {
+          role: 'user',
+          content: [{ type: 'text', text: 'What is the encrypted message for "Hello, world!"' }],
+        }),
+      }),
+    );
+
+    log.info('messages', { messages });
+  });
+
+  test('text-to-image', async (ctx) => {
+    const isRunning = await OllamaClient.isRunning();
+    if (!isRunning) {
+      ctx.skip();
+    }
+
+    const client = createTestOllamaClient();
+    const parser = new MixedStreamParser();
+
+    const messages = await parser.parse(
+      await client.exec({
+        prompt: createStatic(Message, {
+          role: 'user',
+          content: [{ type: 'text', text: 'Generate an image of a cat' }],
+        }),
+        tools: [
+          {
+            name: 'text-to-image',
+            type: ToolTypes.TextToImage,
+          },
+        ],
+      }),
+    );
+
+    log.info('messages', { messages });
+  });
+});

package/src/ai-service/defs.ts

@@ -5,9 +5,9 @@
 // TODO(burdon): Share with Edge.
 // TODO(burdon): Move to service config retrieved by HTTP?
 
-export const DEFAULT_LLM_MODEL = '@anthropic/claude-3-5-sonnet-20241022';
+export const DEFAULT_EDGE_MODEL = '@anthropic/claude-3-5-sonnet-20241022';
 
-export const DEFAULT_LLM_MODELS = [
+export const DEFAULT_EDGE_MODELS = [
   // AI Gateway.
   // https://developers.cloudflare.com/ai-gateway
   '@anthropic/claude-3-5-haiku-20241022',
@@ -21,4 +21,19 @@ export const DEFAULT_LLM_MODELS = [
   '@ollama/llama-3-1-nemotron-70b-instruct',
   '@ollama/llama-3-1-nemotron-mini-4b-instruct',
   '@ollama/llama-3-2-3b',
-];
+] as const;
+
+// TODO(burdon): Config.
+export const DEFAULT_OLLAMA_ENDPOINT = 'http://localhost:11434';
+
+export const DEFAULT_OLLAMA_MODEL = 'llama3.2:1b';
+
+/**
+ * https://ollama.com/library
+ */
+export const DEFAULT_OLLAMA_MODELS = [
+  //
+  'llama3.2:1b',
+  'llama3:70b',
+  'deepseek-r1:latest',
+] as const;

package/src/ai-service/{client.ts → edge-client.ts}

@@ -2,39 +2,32 @@
 // Copyright 2024 DXOS.org
 //
 
-import { Schema as S } from '@effect/schema';
+import { Schema as S } from 'effect';
 
 import { Message } from '@dxos/artifact';
-import { invariant } from '@dxos/invariant';
+import { assertArgument, invariant } from '@dxos/invariant';
 import { log } from '@dxos/log';
 
-import { createGenerationStream, type GenerationStream } from './stream';
+import type { AIServiceClient, GenerationStream } from './service';
+import { createGenerationStream } from './stream';
 import { type GenerateRequest } from './types';
 
-export interface AIService {
-  exec(request: GenerateRequest): Promise<GenerationStream>;
-}
-
-export interface AIServiceClient extends AIService {
-  appendMessages(messages: Message[]): Promise<void>;
-}
-
-export type AIServiceClientParams = {
+export type AIServiceEdgeClientOptions = {
   endpoint: string;
 };
 
-/**
- * Edge GPT client.
- */
-// TODO(burdon): Create mock.
-export class AIServiceClientImpl implements AIServiceClient {
+export class AIServiceEdgeClient implements AIServiceClient {
   private readonly _endpoint: string;
 
-  constructor({ endpoint }: AIServiceClientParams) {
+  constructor({ endpoint }: AIServiceEdgeClientOptions) {
     invariant(endpoint, 'endpoint is required');
     this._endpoint = endpoint;
   }
 
+  /**
+   * @deprecated
+   */
+  // TODO(burdon): Remove.
   async appendMessages(messages: Message[]): Promise<void> {
     const url = `${this._endpoint}/message`;
     const res = await fetch(url, {
@@ -54,19 +47,18 @@ export class AIServiceClientImpl implements AIServiceClient {
    * Process request and open message stream.
    */
   async exec(request: GenerateRequest): Promise<GenerationStream> {
-    // TODO(dmaretskyi): Errors if tools are not provided.
-    request = {
-      tools: request.tools ?? [],
-      ...request,
-    };
-
-    log.info('requesting', { endpoint: this._endpoint });
+    assertArgument(typeof request.model === 'string', 'model is required');
+    log('requesting', { endpoint: this._endpoint });
     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),
+      body: JSON.stringify({
+        // TODO(dmaretskyi): Errors if tools are not provided.
+        tools: request.tools ?? [],
+        ...request,
+      }),
     });
 
     invariant(response.body instanceof ReadableStream);

package/src/ai-service/index.ts

@@ -2,9 +2,12 @@
 // Copyright 2024 DXOS.org
 //
 
-export * from './client';
 export * from './defs';
+export * from './edge-client';
+export * from './message-collector';
+export * from './ollama-client';
 export * from './parser';
+export * from './service';
 export * from './stream';
 export * from './transform';
 export * from './types';

package/src/ai-service/message-collector.ts

@@ -0,0 +1,84 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import type { Message } from '@dxos/artifact';
+
+import type { GenerationStreamEvent } from './types';
+
+export class MessageCollector {
+  private _messages: Message[] = [];
+  private _pendingMessage?: Message = undefined;
+
+  public get messages() {
+    return this._messages;
+  }
+
+  public get pendingMessage() {
+    return this._pendingMessage;
+  }
+
+  push(event: GenerationStreamEvent) {
+    switch (event.type) {
+      case 'message_start':
+        this._pendingMessage = event.message;
+        break;
+
+      case 'message_stop':
+        if (this._pendingMessage) {
+          this._messages.push(this._pendingMessage);
+          this._pendingMessage = undefined;
+        }
+        break;
+
+      case 'content_block_start':
+        if (this._pendingMessage) {
+          // Ensure the content array exists
+          this._pendingMessage.content = this._pendingMessage.content || [];
+          // Add the new content block at the specified index
+          this._pendingMessage.content[event.index] = event.content;
+        }
+        break;
+
+      case 'content_block_delta':
+        if (this._pendingMessage && this._pendingMessage.content) {
+          const block = this._pendingMessage.content[event.index];
+          if (block && event.delta.type === 'text_delta' && block.type === 'text') {
+            block.text = (block.text || '') + event.delta.text;
+          }
+        }
+        break;
+
+      case 'content_block_stop':
+        // Mark the block as no longer pending if needed
+        if (this._pendingMessage && this._pendingMessage.content) {
+          const block = this._pendingMessage.content[event.index];
+          if (block) {
+            block.pending = false;
+          }
+        }
+        break;
+
+      case 'message_delta':
+        // Apply any delta updates to the message
+        if (this._pendingMessage && event.delta) {
+          Object.assign(this._pendingMessage, event.delta);
+        }
+        break;
+    }
+  }
+
+  pushMessage(message: Message) {
+    this._messages.push(message);
+  }
+}
+
+export function* emitMessageAsEvents(message: Message): Iterable<GenerationStreamEvent> {
+  yield {
+    type: 'message_start',
+    message,
+  };
+  yield {
+    type: 'message_stop',
+  };
+}