@puku-ai/sdk 4.0.0 → 4.0.2

package/dist/vendor/lib/BetaMessageStream.js

@@ -0,0 +1,662 @@
+import { STAINLESS_HELPER_METHOD_HEADER } from '../internal/stainless-helper-header.js';
+import { PukuError, APIUserAbortError } from '../error.js';
+import { isAbortError } from '../internal/errors.js';
+import { checkNever } from '../internal/utils/values.js';
+import { Stream } from '../streaming.js';
+import { maybeParseBetaMessage } from './beta-parser.js';
+import { JSON_BUF_PROPERTY, withLazyInput } from '../internal/message-stream-utils.js';
+function tracksToolInput(content) {
+    return content.type === 'tool_use' || content.type === 'server_tool_use' || content.type === 'mcp_tool_use';
+}
+export class BetaMessageStream {
+    messages = [];
+    receivedMessages = [];
+    #currentMessageSnapshot;
+    #params = null;
+    controller = new AbortController();
+    #connectedPromise;
+    #resolveConnectedPromise = () => { };
+    #rejectConnectedPromise = () => { };
+    #endPromise;
+    #resolveEndPromise = () => { };
+    #rejectEndPromise = () => { };
+    #listeners = {};
+    #ended = false;
+    #errored = false;
+    #aborted = false;
+    #catchingPromiseCreated = false;
+    #response;
+    #request_id;
+    #workspace_id;
+    #logger;
+    constructor(params, opts) {
+        this.#connectedPromise = new Promise((resolve, reject) => {
+            this.#resolveConnectedPromise = resolve;
+            this.#rejectConnectedPromise = reject;
+        });
+        this.#endPromise = new Promise((resolve, reject) => {
+            this.#resolveEndPromise = resolve;
+            this.#rejectEndPromise = reject;
+        });
+        // Don't let these promises cause unhandled rejection errors.
+        // we will manually cause an unhandled rejection error later
+        // if the user hasn't registered any error listener or called
+        // any promise-returning method.
+        this.#connectedPromise.catch(() => { });
+        this.#endPromise.catch(() => { });
+        this.#params = params;
+        this.#logger = opts?.logger ?? console;
+    }
+    get response() {
+        return this.#response;
+    }
+    get request_id() {
+        return this.#request_id;
+    }
+    get workspace_id() {
+        return this.#workspace_id;
+    }
+    /**
+     * Returns the `MessageStream` data, the raw `Response` instance and the ID of the request,
+     * returned vie the `request-id` header which is useful for debugging requests and resporting
+     * issues to PukuAI.
+     *
+     * This is the same as the `APIPromise.withResponse()` method.
+     *
+     * This method will raise an error if you created the stream using `MessageStream.fromReadableStream`
+     * as no `Response` is available.
+     */
+    async withResponse() {
+        this.#catchingPromiseCreated = true;
+        const response = await this.#connectedPromise;
+        if (!response) {
+            throw new Error('Could not resolve a `Response` object');
+        }
+        return {
+            data: this,
+            response,
+            request_id: response.headers.get('request-id'),
+            workspace_id: response.headers.get('puku-workspace-id'),
+        };
+    }
+    /**
+     * Intended for use on the frontend, consuming a stream produced with
+     * `.toReadableStream()` on the backend.
+     *
+     * Note that messages sent to the model do not appear in `.on('message')`
+     * in this context.
+     */
+    static fromReadableStream(stream) {
+        const runner = new BetaMessageStream(null);
+        runner._run(() => runner._fromReadableStream(stream));
+        return runner;
+    }
+    static createMessage(messages, params, options, { logger } = {}) {
+        const runner = new BetaMessageStream(params, { logger });
+        for (const message of params.messages) {
+            runner._addMessageParam(message);
+        }
+        runner.#params = { ...params, stream: true };
+        runner._run(() => runner._createMessage(messages, { ...params, stream: true }, { ...options, headers: { ...options?.headers, [STAINLESS_HELPER_METHOD_HEADER]: 'stream' } }));
+        return runner;
+    }
+    _run(executor) {
+        executor().then(() => {
+            this._emitFinal();
+            this._emit('end');
+        }, this.#handleError);
+    }
+    _addMessageParam(message) {
+        this.messages.push(message);
+    }
+    _addMessage(message, emit = true) {
+        this.receivedMessages.push(message);
+        if (emit) {
+            this._emit('message', message);
+        }
+    }
+    async _createMessage(messages, params, options) {
+        const signal = options?.signal;
+        let abortHandler;
+        if (signal) {
+            if (signal.aborted)
+                this.controller.abort();
+            abortHandler = this.controller.abort.bind(this.controller);
+            signal.addEventListener('abort', abortHandler);
+        }
+        try {
+            this.#beginRequest();
+            const { response, data: stream } = await messages
+                .create({ ...params, stream: true }, { ...options, signal: this.controller.signal })
+                .withResponse();
+            this._connected(response);
+            for await (const event of stream) {
+                this.#addStreamEvent(event);
+            }
+            if (stream.controller.signal?.aborted) {
+                throw new APIUserAbortError();
+            }
+            this.#endRequest();
+        }
+        finally {
+            if (signal && abortHandler) {
+                signal.removeEventListener('abort', abortHandler);
+            }
+        }
+    }
+    _connected(response) {
+        if (this.ended)
+            return;
+        this.#response = response;
+        this.#request_id = response?.headers.get('request-id');
+        this.#workspace_id = response?.headers.get('puku-workspace-id');
+        this.#resolveConnectedPromise(response);
+        this._emit('connect');
+    }
+    get ended() {
+        return this.#ended;
+    }
+    get errored() {
+        return this.#errored;
+    }
+    get aborted() {
+        return this.#aborted;
+    }
+    abort() {
+        this.controller.abort();
+    }
+    /**
+     * Adds the listener function to the end of the listeners array for the event.
+     * No checks are made to see if the listener has already been added. Multiple calls passing
+     * the same combination of event and listener will result in the listener being added, and
+     * called, multiple times.
+     * @returns this MessageStream, so that calls can be chained
+     */
+    on(event, listener) {
+        const listeners = this.#listeners[event] || (this.#listeners[event] = []);
+        listeners.push({ listener });
+        return this;
+    }
+    /**
+     * Removes the specified listener from the listener array for the event.
+     * off() will remove, at most, one instance of a listener from the listener array. If any single
+     * listener has been added multiple times to the listener array for the specified event, then
+     * off() must be called multiple times to remove each instance.
+     * @returns this MessageStream, so that calls can be chained
+     */
+    off(event, listener) {
+        const listeners = this.#listeners[event];
+        if (!listeners)
+            return this;
+        const index = listeners.findIndex((l) => l.listener === listener);
+        if (index >= 0)
+            listeners.splice(index, 1);
+        return this;
+    }
+    /**
+     * Adds a one-time listener function for the event. The next time the event is triggered,
+     * this listener is removed and then invoked.
+     * @returns this MessageStream, so that calls can be chained
+     */
+    once(event, listener) {
+        const listeners = this.#listeners[event] || (this.#listeners[event] = []);
+        listeners.push({ listener, once: true });
+        return this;
+    }
+    /**
+     * This is similar to `.once()`, but returns a Promise that resolves the next time
+     * the event is triggered, instead of calling a listener callback.
+     * @returns a Promise that resolves the next time given event is triggered,
+     * or rejects if an error is emitted.  (If you request the 'error' event,
+     * returns a promise that resolves with the error).
+     *
+     * Example:
+     *
+     *   const message = await stream.emitted('message') // rejects if the stream errors
+     */
+    emitted(event) {
+        return new Promise((resolve, reject) => {
+            this.#catchingPromiseCreated = true;
+            if (event !== 'error')
+                this.once('error', reject);
+            this.once(event, resolve);
+        });
+    }
+    async done() {
+        this.#catchingPromiseCreated = true;
+        await this.#endPromise;
+    }
+    get currentMessage() {
+        return this.#currentMessageSnapshot;
+    }
+    #getFinalMessage() {
+        if (this.receivedMessages.length === 0) {
+            throw new PukuError('stream ended without producing a Message with role=assistant');
+        }
+        return this.receivedMessages.at(-1);
+    }
+    /**
+     * @returns a promise that resolves with the the final assistant Message response,
+     * or rejects if an error occurred or the stream ended prematurely without producing a Message.
+     * If structured outputs were used, this will be a ParsedMessage with a `parsed` field.
+     */
+    async finalMessage() {
+        await this.done();
+        return this.#getFinalMessage();
+    }
+    #getFinalText() {
+        if (this.receivedMessages.length === 0) {
+            throw new PukuError('stream ended without producing a Message with role=assistant');
+        }
+        const textBlocks = this.receivedMessages
+            .at(-1)
+            .content.filter((block) => block.type === 'text')
+            .map((block) => block.text);
+        if (textBlocks.length === 0) {
+            throw new PukuError('stream ended without producing a content block with type=text');
+        }
+        return textBlocks.join(' ');
+    }
+    /**
+     * @returns a promise that resolves with the the final assistant Message's text response, concatenated
+     * together if there are more than one text blocks.
+     * Rejects if an error occurred or the stream ended prematurely without producing a Message.
+     */
+    async finalText() {
+        await this.done();
+        return this.#getFinalText();
+    }
+    #handleError = (error) => {
+        this.#errored = true;
+        if (isAbortError(error)) {
+            error = new APIUserAbortError();
+        }
+        if (error instanceof APIUserAbortError) {
+            this.#aborted = true;
+            return this._emit('abort', error);
+        }
+        if (error instanceof PukuError) {
+            return this._emit('error', error);
+        }
+        if (error instanceof Error) {
+            const pukuError = new PukuError(error.message);
+            // @ts-ignore
+            pukuError.cause = error;
+            return this._emit('error', pukuError);
+        }
+        return this._emit('error', new PukuError(String(error)));
+    };
+    _emit(event, ...args) {
+        // make sure we don't emit any MessageStreamEvents after end
+        if (this.#ended)
+            return;
+        if (event === 'end') {
+            this.#ended = true;
+            this.#resolveEndPromise();
+        }
+        const listeners = this.#listeners[event];
+        if (listeners) {
+            this.#listeners[event] = listeners.filter((l) => !l.once);
+            listeners.forEach(({ listener }) => listener(...args));
+        }
+        if (event === 'abort') {
+            const error = args[0];
+            if (!this.#catchingPromiseCreated && !listeners?.length) {
+                Promise.reject(error);
+            }
+            this.#rejectConnectedPromise(error);
+            this.#rejectEndPromise(error);
+            this._emit('end');
+            return;
+        }
+        if (event === 'error') {
+            // NOTE: _emit('error', error) should only be called from #handleError().
+            const error = args[0];
+            if (!this.#catchingPromiseCreated && !listeners?.length) {
+                // Trigger an unhandled rejection if the user hasn't registered any error handlers.
+                // If you are seeing stack traces here, make sure to handle errors via either:
+                // - runner.on('error', () => ...)
+                // - await runner.done()
+                // - await runner.final...()
+                // - etc.
+                Promise.reject(error);
+            }
+            this.#rejectConnectedPromise(error);
+            this.#rejectEndPromise(error);
+            this._emit('end');
+        }
+    }
+    _emitFinal() {
+        const finalMessage = this.receivedMessages.at(-1);
+        if (finalMessage) {
+            this._emit('finalMessage', this.#getFinalMessage());
+        }
+    }
+    #beginRequest() {
+        if (this.ended)
+            return;
+        this.#currentMessageSnapshot = undefined;
+    }
+    #addStreamEvent(event) {
+        if (this.ended)
+            return;
+        const messageSnapshot = this.#accumulateMessage(event);
+        this._emit('streamEvent', event, messageSnapshot);
+        switch (event.type) {
+            case 'content_block_delta': {
+                const content = messageSnapshot.content.at(-1);
+                switch (event.delta.type) {
+                    case 'text_delta': {
+                        if (content.type === 'text') {
+                            this._emit('text', event.delta.text, content.text || '');
+                        }
+                        break;
+                    }
+                    case 'citations_delta': {
+                        if (content.type === 'text') {
+                            this._emit('citation', event.delta.citation, content.citations ?? []);
+                        }
+                        break;
+                    }
+                    case 'input_json_delta': {
+                        if (tracksToolInput(content) && this.#listeners.inputJson?.length) {
+                            let jsonSnapshot;
+                            try {
+                                jsonSnapshot = content.input;
+                            }
+                            catch (err) {
+                                this.#handleError(this.#toolInputParseError(content, err));
+                                break;
+                            }
+                            this._emit('inputJson', event.delta.partial_json, jsonSnapshot);
+                        }
+                        break;
+                    }
+                    case 'thinking_delta': {
+                        if (content.type === 'thinking') {
+                            this._emit('thinking', event.delta.thinking, content.thinking);
+                        }
+                        break;
+                    }
+                    case 'signature_delta': {
+                        if (content.type === 'thinking') {
+                            this._emit('signature', content.signature);
+                        }
+                        break;
+                    }
+                    case 'compaction_delta': {
+                        if (content.type === 'compaction' && content.content) {
+                            this._emit('compaction', content.content);
+                        }
+                        break;
+                    }
+                    default:
+                        checkNever(event.delta);
+                }
+                break;
+            }
+            case 'message_stop': {
+                this._addMessageParam(messageSnapshot);
+                this._addMessage(maybeParseBetaMessage(messageSnapshot, this.#params, { logger: this.#logger }), true);
+                break;
+            }
+            case 'content_block_stop': {
+                this._emit('contentBlock', messageSnapshot.content.at(-1));
+                break;
+            }
+            case 'message_start': {
+                this.#currentMessageSnapshot = messageSnapshot;
+                break;
+            }
+            case 'content_block_start':
+            case 'message_delta':
+                break;
+        }
+    }
+    #endRequest() {
+        if (this.ended) {
+            throw new PukuError(`stream has ended, this shouldn't happen`);
+        }
+        const snapshot = this.#currentMessageSnapshot;
+        if (!snapshot) {
+            throw new PukuError(`request ended without sending any chunks`);
+        }
+        this.#currentMessageSnapshot = undefined;
+        return maybeParseBetaMessage(snapshot, this.#params, { logger: this.#logger });
+    }
+    async _fromReadableStream(readableStream, options) {
+        const signal = options?.signal;
+        let abortHandler;
+        if (signal) {
+            if (signal.aborted)
+                this.controller.abort();
+            abortHandler = this.controller.abort.bind(this.controller);
+            signal.addEventListener('abort', abortHandler);
+        }
+        try {
+            this.#beginRequest();
+            this._connected(null);
+            const stream = Stream.fromReadableStream(readableStream, this.controller);
+            for await (const event of stream) {
+                this.#addStreamEvent(event);
+            }
+            if (stream.controller.signal?.aborted) {
+                throw new APIUserAbortError();
+            }
+            this.#endRequest();
+        }
+        finally {
+            if (signal && abortHandler) {
+                signal.removeEventListener('abort', abortHandler);
+            }
+        }
+    }
+    /**
+     * Mutates this.#currentMessage with the current event. Handling the accumulation of multiple messages
+     * will be needed to be handled by the caller, this method will throw if you try to accumulate for multiple
+     * messages.
+     */
+    #accumulateMessage(event) {
+        let snapshot = this.#currentMessageSnapshot;
+        if (event.type === 'message_start') {
+            if (snapshot) {
+                throw new PukuError(`Unexpected event order, got ${event.type} before receiving "message_stop"`);
+            }
+            return event.message;
+        }
+        if (!snapshot) {
+            throw new PukuError(`Unexpected event order, got ${event.type} before "message_start"`);
+        }
+        switch (event.type) {
+            case 'message_stop':
+                return snapshot;
+            case 'message_delta':
+                snapshot.stop_reason = event.delta.stop_reason;
+                snapshot.stop_sequence = event.delta.stop_sequence;
+                snapshot.stop_details = event.delta.stop_details;
+                snapshot.usage.output_tokens = event.usage.output_tokens;
+                if (event.delta.container != null) {
+                    snapshot.container = event.delta.container;
+                }
+                if (event.context_management != null) {
+                    snapshot.context_management = event.context_management;
+                }
+                if (event.input_transformations != null) {
+                    snapshot.input_transformations = event.input_transformations;
+                }
+                // The remaining usage counters are cumulative whole-message totals that are
+                // omitted when they don't apply, so overwrite when present and never add.
+                if (event.usage.input_tokens != null) {
+                    snapshot.usage.input_tokens = event.usage.input_tokens;
+                }
+                if (event.usage.cache_creation_input_tokens != null) {
+                    snapshot.usage.cache_creation_input_tokens = event.usage.cache_creation_input_tokens;
+                }
+                if (event.usage.cache_read_input_tokens != null) {
+                    snapshot.usage.cache_read_input_tokens = event.usage.cache_read_input_tokens;
+                }
+                if (event.usage.server_tool_use != null) {
+                    snapshot.usage.server_tool_use = event.usage.server_tool_use;
+                }
+                if (event.usage.iterations != null) {
+                    snapshot.usage.iterations = event.usage.iterations;
+                }
+                if (event.usage.fallback_credit != null) {
+                    snapshot.usage.fallback_credit = event.usage.fallback_credit;
+                }
+                if (event.usage.output_tokens_details != null) {
+                    snapshot.usage.output_tokens_details = event.usage.output_tokens_details;
+                }
+                return snapshot;
+            case 'content_block_start':
+                snapshot.content.push(event.content_block);
+                if (event.content_block.type === 'fallback') {
+                    // the final hop's fallback block names the model that served the response —
+                    // keeps the snapshot consistent with the relabeled non-streaming message
+                    snapshot.model = event.content_block.to.model;
+                }
+                return snapshot;
+            case 'content_block_delta': {
+                const snapshotContent = snapshot.content.at(event.index);
+                switch (event.delta.type) {
+                    case 'text_delta': {
+                        if (snapshotContent?.type === 'text') {
+                            snapshot.content[event.index] = {
+                                ...snapshotContent,
+                                text: (snapshotContent.text || '') + event.delta.text,
+                            };
+                        }
+                        break;
+                    }
+                    case 'citations_delta': {
+                        if (snapshotContent?.type === 'text') {
+                            snapshot.content[event.index] = {
+                                ...snapshotContent,
+                                citations: [...(snapshotContent.citations ?? []), event.delta.citation],
+                            };
+                        }
+                        break;
+                    }
+                    case 'input_json_delta': {
+                        if (snapshotContent && tracksToolInput(snapshotContent)) {
+                            const jsonBuf = (snapshotContent[JSON_BUF_PROPERTY] || '') + event.delta.partial_json;
+                            snapshot.content[event.index] = withLazyInput(snapshotContent, jsonBuf);
+                        }
+                        break;
+                    }
+                    case 'thinking_delta': {
+                        if (snapshotContent?.type === 'thinking') {
+                            snapshot.content[event.index] = {
+                                ...snapshotContent,
+                                thinking: snapshotContent.thinking + event.delta.thinking,
+                            };
+                        }
+                        break;
+                    }
+                    case 'signature_delta': {
+                        if (snapshotContent?.type === 'thinking') {
+                            snapshot.content[event.index] = {
+                                ...snapshotContent,
+                                signature: event.delta.signature,
+                            };
+                        }
+                        break;
+                    }
+                    case 'compaction_delta': {
+                        if (snapshotContent?.type === 'compaction') {
+                            snapshot.content[event.index] = {
+                                ...snapshotContent,
+                                content: (snapshotContent.content || '') + event.delta.content,
+                                encrypted_content: event.delta.encrypted_content,
+                            };
+                        }
+                        break;
+                    }
+                    default:
+                        checkNever(event.delta);
+                }
+                return snapshot;
+            }
+            case 'content_block_stop': {
+                const snapshotContent = snapshot.content.at(event.index);
+                if (snapshotContent && tracksToolInput(snapshotContent) && JSON_BUF_PROPERTY in snapshotContent) {
+                    let input;
+                    try {
+                        input = snapshotContent.input;
+                    }
+                    catch (err) {
+                        input = {};
+                        this.#handleError(this.#toolInputParseError(snapshotContent, err));
+                    }
+                    Object.defineProperty(snapshotContent, 'input', {
+                        value: input,
+                        enumerable: true,
+                        configurable: true,
+                        writable: true,
+                    });
+                }
+                return snapshot;
+            }
+        }
+    }
+    #toolInputParseError(block, err) {
+        const jsonBuf = block[JSON_BUF_PROPERTY];
+        return new PukuError(`Unable to parse tool parameter JSON from model. Please retry your request or adjust your prompt. Error: ${err}. JSON: ${jsonBuf}`);
+    }
+    [Symbol.asyncIterator]() {
+        const pushQueue = [];
+        const readQueue = [];
+        let done = false;
+        this.on('streamEvent', (event) => {
+            const reader = readQueue.shift();
+            if (reader) {
+                reader.resolve(event);
+            }
+            else {
+                pushQueue.push(event);
+            }
+        });
+        this.on('end', () => {
+            done = true;
+            for (const reader of readQueue) {
+                reader.resolve(undefined);
+            }
+            readQueue.length = 0;
+        });
+        this.on('abort', (err) => {
+            done = true;
+            for (const reader of readQueue) {
+                reader.reject(err);
+            }
+            readQueue.length = 0;
+        });
+        this.on('error', (err) => {
+            done = true;
+            for (const reader of readQueue) {
+                reader.reject(err);
+            }
+            readQueue.length = 0;
+        });
+        return {
+            next: async () => {
+                if (!pushQueue.length) {
+                    if (done) {
+                        return { value: undefined, done: true };
+                    }
+                    return new Promise((resolve, reject) => readQueue.push({ resolve, reject })).then((chunk) => (chunk ? { value: chunk, done: false } : { value: undefined, done: true }));
+                }
+                const chunk = pushQueue.shift();
+                return { value: chunk, done: false };
+            },
+            return: async () => {
+                this.abort();
+                return { value: undefined, done: true };
+            },
+        };
+    }
+    toReadableStream() {
+        const stream = new Stream(this[Symbol.asyncIterator].bind(this), this.controller);
+        return stream.toReadableStream();
+    }
+}
+//# sourceMappingURL=BetaMessageStream.js.map