@dxos/assistant 0.8.0 → 0.8.1-main.a06ad20

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

@@ -2,17 +2,19 @@
 // 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 { AIServiceEdgeClient } from './edge-client';
+import { OllamaClient } from './ollama-client';
+import { MixedStreamParser } from './parser';
 import { ToolTypes } from './types';
 import { AI_SERVICE_ENDPOINT } from '../testing';
 
@@ -20,7 +22,7 @@ import { AI_SERVICE_ENDPOINT } from '../testing';
 
 describe.skip('AI Service Client', () => {
   test('client generation', async () => {
-    const client = new AIServiceClientImpl({
+    const client = new AIServiceEdgeClient({
       endpoint: AI_SERVICE_ENDPOINT.LOCAL,
     });
 
@@ -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,
     });
 
@@ -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,
     });
 
@@ -176,3 +178,88 @@ describe.skip('AI Service Client', () => {
     log('full message', { message: await stream.complete() });
   });
 });
+
+describe('Ollama Client', () => {
+  test('basic', async (ctx) => {
+    const isRunning = await OllamaClient.isRunning();
+    if (!isRunning) {
+      ctx.skip();
+    }
+
+    const client = OllamaClient.createClient();
+    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 = OllamaClient.createClient({
+      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 = OllamaClient.createClient();
+    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/{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 { 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,17 @@ 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 });
     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',
+  };
+}