@celljs/ai-core 3.2.0 → 3.4.0

package/src/common/model/model-protocol.ts

@@ -39,7 +39,9 @@ export interface ResultMetadata {}
  * object within the core domain, enhancing the understanding and management of AI model
  * responses in various applications.
  */
-export interface ResponseMetadata {}
+export interface ResponseMetadata {
+    readonly extra: Record<string, any>;
+}
 
 /**
  * This interface provides methods to access the main output of the AI model and the
@@ -182,6 +184,6 @@ export interface StreamingModel<TReq extends ModelRequest<any>, TResChunk extend
      * @param request the request object to be sent to the AI model
      * @return the streaming response from the AI model
      */
-    stream(request: TReq): Observable<TResChunk>;
+    stream(request: TReq): Promise<Observable<TResChunk>>;
 
 }

package/src/common/sse/index.ts

@@ -0,0 +1,3 @@
+export * from './line-decoder';
+export * from './sse-protocol';
+export * from './sse-decoder';

package/src/common/sse/line-decoder.ts

@@ -0,0 +1,65 @@
+import { Bytes, ByteUtil } from '@celljs/core';
+import { LineDecoder } from './sse-protocol';
+
+/**
+ * Line decoder.
+ */
+export class LineDecoderImpl implements LineDecoder {
+    private buffer: string[];
+    private trailingCR: boolean;
+
+    static readonly NEWLINE_CHARS = new Set(['\n', '\r', '\x0b', '\x0c', '\x1c', '\x1d', '\x1e', '\x85', '\u2028', '\u2029']);
+    static readonly NEWLINE_REGEXP = /\r\n|[\n\r\x0b\x0c\x1c\x1d\x1e\x85\u2028\u2029]/g;
+
+    constructor() {
+        this.buffer = [];
+        this.trailingCR = false;
+    }
+
+    decode(chunk: Bytes): string[] {
+        let text = ByteUtil.decode(chunk);
+
+        if (this.trailingCR) {
+            text = '\r' + text;
+            this.trailingCR = false;
+        }
+        if (text.endsWith('\r')) {
+            this.trailingCR = true;
+            text = text.slice(0, -1);
+        }
+
+        if (!text) {
+            return [];
+        }
+
+        const trailingNewline = LineDecoderImpl.NEWLINE_CHARS.has(text[text.length - 1] || '');
+        let lines = text.split(LineDecoderImpl.NEWLINE_REGEXP);
+
+        if (lines.length === 1 && !trailingNewline) {
+            this.buffer.push(lines[0]!);
+            return [];
+        }
+
+        if (this.buffer.length > 0) {
+            lines = [this.buffer.join('') + lines[0], ...lines.slice(1)];
+            this.buffer = [];
+        }
+
+        if (!trailingNewline) {
+            this.buffer = [lines.pop() || ''];
+        }
+
+        return lines;
+    }
+
+    flush(): string[] {
+        if (this.buffer.length === 0 && !this.trailingCR) {
+            return [];
+        }
+
+        const lines = [this.buffer.join('')];
+        this.buffer = [];
+        this.trailingCR = false;
+        return lines;
+    }
+}

package/src/common/sse/sse-decoder.ts

@@ -0,0 +1,88 @@
+import { EventDecoder, StreamEvent } from './sse-protocol';
+
+/**
+ * Server-sent event.
+ */
+export class SSEDecoder implements EventDecoder {
+    private event?: string;
+    private data: string[];
+    private chunks: string[];
+
+    constructor() {
+        this.event = undefined;
+        this.data = [];
+        this.chunks = [];
+    }
+
+    protected buildServerSentEvent(): StreamEvent<string> {
+        const sse: StreamEvent<string> = {
+            event: this.event,
+            data: this.data.join('\n'),
+            raw: this.chunks,
+        };
+
+        this.event = undefined;
+        this.data = [];
+        this.chunks = [];
+
+        return sse;
+    }
+
+    decode(line: string): StreamEvent<string> | undefined {
+        if (line.endsWith('\r')) {
+            line = line.substring(0, line.length - 1);
+        }
+
+        if (!line) {
+            if (!this.event && this.data.length === 0) {
+                return;
+            }
+
+            return this.buildServerSentEvent();
+        }
+
+        this.chunks.push(line);
+
+        if (line.startsWith(':')) {
+            return;
+        }
+
+        if (line.startsWith('{') && line.endsWith('}')) {
+            this.data.push(line);
+            return this.buildServerSentEvent();
+        }
+
+        const [fieldName, , value] = this.partition(line, ':');
+
+        const trimmedValue = value.startsWith(' ') ? value.substring(1) : value;
+
+        if (fieldName === 'event') {
+            this.event = trimmedValue;
+        } else if (fieldName === 'data') {
+            this.data.push(trimmedValue);
+        }
+
+        return;
+    }
+
+    /**
+     * Partitions a string into three parts based on a delimiter.
+     * @param str The string to partition.
+     * @param delimiter The delimiter to partition the string by.
+     * @returns A tuple containing the three parts of the partitioned string.
+     * The first part is the string before the delimiter, the second part is the delimiter itself,
+     * and the third part is the string after the delimiter.
+     * If the delimiter is not found in the string, the first part is the entire string, the second part is an empty string, and the third part is an empty string.
+     * @example
+     * partition('hello,world', ',') // ['hello', ',', 'world']
+     * partition('hello,world', '.') // ['hello,world', '', '']
+     * partition('hello,world', 'l') // ['he', 'l', 'lo,world']
+     */
+    protected partition(str: string, delimiter: string): [string, string, string] {
+        const index = str.indexOf(delimiter);
+        if (index !== -1) {
+            return [str.substring(0, index), delimiter, str.substring(index + delimiter.length)];
+        }
+        return [str, '', ''];
+    }
+}

package/src/common/sse/sse-protocol.ts

@@ -0,0 +1,25 @@
+import { Bytes } from '@celljs/core';
+
+/**
+ * Server-sent event.
+ */
+export interface StreamEvent<Data> {
+    event?: string;
+    data: Data;
+    raw: string[];
+};
+
+/**
+ * Decoder for server-sent events.
+ */
+export interface EventDecoder {
+    decode(line: string): StreamEvent<string> | undefined;
+}
+
+/**
+ * Decoder for line-based protocols.
+ */
+export interface LineDecoder {
+    decode(chunk: Bytes): string[];
+    flush(): string[];
+}

package/src/common/utils/index.ts

@@ -1 +1,2 @@
 export * from './prompt-util';
+export * from './sse-util';

package/src/common/utils/sse-util.ts

@@ -0,0 +1,115 @@
+import { Observable, Subscriber } from 'rxjs';
+import { StreamEvent, SSEDecoder, LineDecoderImpl } from '../sse';
+import { Bytes } from '@celljs/core';
+
+/**
+ * SSE utility class.
+ */
+export class SSEUtil {
+
+    static toObservable<Data>(readableStream: ReadableStream, controller?: AbortController): Observable<StreamEvent<Data>> {
+        return new Observable<StreamEvent<Data>>(subscriber => {
+            (async () => {
+                const decoder = new SSEDecoder();
+                const lineDecoder = new LineDecoderImpl();
+
+                try {
+                    const iterable = this.readableStreamAsyncIterable<Bytes>(readableStream);
+                    for await (const chunk of iterable) {
+                        const lines = lineDecoder.decode(chunk);
+                        for (const line of lines) {
+                            const sse = decoder.decode(line);
+                            this.handleSSE(sse, subscriber, controller);
+                        }
+                    }
+
+                    const remainingLines = lineDecoder.flush();
+                    for (const line of remainingLines) {
+                        const sse = decoder.decode(line);
+                        if (sse) {
+                            this.handleSSE(sse, subscriber, controller);
+                        }
+                    }
+
+                    subscriber.complete();
+                } catch (e) {
+                    if (e instanceof Error && e.name === 'AbortError') {
+                        subscriber.complete();
+                    } else {
+                        subscriber.error(e);
+                    }
+                } finally {
+                    controller?.abort();
+                }
+
+            })();
+
+        });
+
+    }
+
+    private static handleSSE<Item>(sse: StreamEvent<string> | undefined, subscriber: Subscriber<StreamEvent<Item>>, controller?: AbortController): void {
+        if (sse) {
+            if (sse.data.startsWith('[DONE]')) {
+                subscriber.complete();
+                controller?.abort();
+                return;
+            }
+
+            if (!sse.event || sse.event === 'error') {
+                let data;
+                try {
+                    data = JSON.parse(sse.data);
+                } catch (e) {
+                    console.error('Could not parse message into JSON:', sse.data);
+                    console.error('From chunk:', sse.raw);
+                    subscriber.error(e);
+                    return;
+                }
+
+                if (data && data.error) {
+                    subscriber.error(data.error);
+                    return;
+                }
+
+                subscriber.next({
+                    event: sse.event,
+                    data,
+                    raw: sse.raw
+                });
+
+                if (data.done === true) {
+                    subscriber.complete();
+                    controller?.abort();
+                }
+            }
+        }
+    }
+
+    private static readableStreamAsyncIterable<T>(stream: any): AsyncIterableIterator<T> {
+        if (stream[Symbol.asyncIterator]) {return stream; }
+
+        const reader = stream.getReader();
+        return {
+            async next() {
+                try {
+                    const result = await reader.read();
+                    if (result?.done) {reader.releaseLock(); }
+                    return result;
+                } catch (e) {
+                    reader.releaseLock();
+                    throw e;
+                }
+            },
+            async return() {
+                const cancelPromise = reader.cancel();
+                reader.releaseLock();
+                await cancelPromise;
+                return { done: true, value: undefined };
+            },
+            [Symbol.asyncIterator]() {
+                return this;
+            },
+        };
+    }
+}