ai 2.2.1 → 2.2.2

package/dist/index.mjs

@@ -28,28 +28,36 @@ function createEventStreamTransformer(customParser) {
     }
   });
 }
-function createCallbacksTransformer(callbacks) {
+function createCallbacksTransformer(cb) {
   const textEncoder = new TextEncoder();
   let aggregatedResponse = "";
-  const { onStart, onToken, onCompletion } = callbacks || {};
+  const callbacks = cb || {};
   return new TransformStream({
     async start() {
-      if (onStart)
-        await onStart();
+      if (callbacks.onStart)
+        await callbacks.onStart();
     },
     async transform(message, controller) {
       controller.enqueue(textEncoder.encode(message));
-      if (onToken)
-        await onToken(message);
-      if (onCompletion)
+      if (callbacks.onToken)
+        await callbacks.onToken(message);
+      if (callbacks.onCompletion)
         aggregatedResponse += message;
     },
     async flush() {
-      if (onCompletion)
-        await onCompletion(aggregatedResponse);
+      const isOpenAICallbacks = isOfTypeOpenAIStreamCallbacks(callbacks);
+      if (callbacks.onCompletion) {
+        await callbacks.onCompletion(aggregatedResponse);
+      }
+      if (callbacks.onFinal && !isOpenAICallbacks) {
+        await callbacks.onFinal(aggregatedResponse);
+      }
     }
   });
 }
+function isOfTypeOpenAIStreamCallbacks(callbacks) {
+  return "experimental_onFunctionCall" in callbacks;
+}
 function trimStartOfStreamHelper() {
   let isStreamStart = true;
   return (text) => {
@@ -109,6 +117,140 @@ function readableFromAsyncIterable(iterable) {
   });
 }
 
+// shared/utils.ts
+import { customAlphabet } from "nanoid/non-secure";
+var nanoid = customAlphabet(
+  "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz",
+  7
+);
+function createChunkDecoder(complex) {
+  const decoder = new TextDecoder();
+  if (!complex) {
+    return function(chunk) {
+      if (!chunk)
+        return "";
+      return decoder.decode(chunk, { stream: true });
+    };
+  }
+  return function(chunk) {
+    const decoded = decoder.decode(chunk, { stream: true }).split("\n");
+    return decoded.map(getStreamStringTypeAndValue).filter(Boolean);
+  };
+}
+var StreamStringPrefixes = {
+  text: 0,
+  function_call: 1,
+  data: 2
+  // user_err: 3?
+};
+var isStreamStringEqualToType = (type, value) => value.startsWith(`${StreamStringPrefixes[type]}:`) && value.endsWith("\n");
+var getStreamString = (type, value) => `${StreamStringPrefixes[type]}:${JSON.stringify(value)}
+`;
+var getStreamStringTypeAndValue = (line) => {
+  const firstSeperatorIndex = line.indexOf(":");
+  const prefix = line.slice(0, firstSeperatorIndex);
+  const type = Object.keys(StreamStringPrefixes).find(
+    (key) => StreamStringPrefixes[key] === Number(prefix)
+  );
+  const val = line.slice(firstSeperatorIndex + 1);
+  let parsedVal = val;
+  if (!val) {
+    return { type, value: "" };
+  }
+  try {
+    parsedVal = JSON.parse(val);
+  } catch (e) {
+    console.error("Failed to parse JSON value:", val);
+  }
+  return { type, value: parsedVal };
+};
+var COMPLEX_HEADER = "X-Experimental-Stream-Data";
+
+// streams/stream-data.ts
+var experimental_StreamData = class {
+  constructor() {
+    this.encoder = new TextEncoder();
+    this.controller = null;
+    // closing the stream is synchronous, but we want to return a promise
+    // in case we're doing async work
+    this.isClosedPromise = null;
+    this.isClosedPromiseResolver = void 0;
+    this.isClosed = false;
+    // array to store appended data
+    this.data = [];
+    this.isClosedPromise = new Promise((resolve) => {
+      this.isClosedPromiseResolver = resolve;
+    });
+    const self = this;
+    this.stream = new TransformStream({
+      start: async (controller) => {
+        self.controller = controller;
+      },
+      transform: async (chunk, controller) => {
+        controller.enqueue(chunk);
+        if (self.data.length > 0) {
+          const encodedData = self.encoder.encode(
+            getStreamString("data", JSON.stringify(self.data))
+          );
+          self.data = [];
+          controller.enqueue(encodedData);
+        }
+      },
+      async flush(controller) {
+        const warningTimeout = process.env.NODE_ENV === "development" ? setTimeout(() => {
+          console.warn(
+            "The data stream is hanging. Did you forget to close it with `data.close()`?"
+          );
+        }, 3e3) : null;
+        await self.isClosedPromise;
+        if (warningTimeout !== null) {
+          clearTimeout(warningTimeout);
+        }
+        if (self.data.length) {
+          const encodedData = self.encoder.encode(
+            getStreamString("data", JSON.stringify(self.data))
+          );
+          controller.enqueue(encodedData);
+        }
+      }
+    });
+  }
+  async close() {
+    var _a;
+    if (this.isClosed) {
+      throw new Error("Data Stream has already been closed.");
+    }
+    if (!this.controller) {
+      throw new Error("Stream controller is not initialized.");
+    }
+    (_a = this.isClosedPromiseResolver) == null ? void 0 : _a.call(this);
+    this.isClosed = true;
+  }
+  append(value) {
+    if (this.isClosed) {
+      throw new Error("Data Stream has already been closed.");
+    }
+    this.data.push(value);
+  }
+};
+function createStreamDataTransformer(experimental_streamData) {
+  if (!experimental_streamData) {
+    return new TransformStream({
+      transform: async (chunk, controller) => {
+        controller.enqueue(chunk);
+      }
+    });
+  }
+  const encoder = new TextEncoder();
+  const decoder = new TextDecoder();
+  return new TransformStream({
+    transform: async (chunk, controller) => {
+      const message = decoder.decode(chunk);
+      controller.enqueue(encoder.encode(getStreamString("text", message)));
+    }
+  });
+}
+
 // streams/openai-stream.ts
 function parseOpenAIStream() {
   const extract = chunkToText();
@@ -152,27 +294,49 @@ function OpenAIStream(res, callbacks) {
   let stream;
   if (Symbol.asyncIterator in res) {
     stream = readableFromAsyncIterable(streamable(res)).pipeThrough(
-      createCallbacksTransformer(cb)
+      createCallbacksTransformer(
+        (cb == null ? void 0 : cb.experimental_onFunctionCall) ? {
+          ...cb,
+          onFinal: void 0
+        } : {
+          ...cb
+        }
+      )
     );
   } else {
-    stream = AIStream(res, parseOpenAIStream(), cb);
+    stream = AIStream(
+      res,
+      parseOpenAIStream(),
+      (cb == null ? void 0 : cb.experimental_onFunctionCall) ? {
+        ...cb,
+        onFinal: void 0
+      } : {
+        ...cb
+      }
+    );
   }
   if (cb && cb.experimental_onFunctionCall) {
     const functionCallTransformer = createFunctionCallTransformer(cb);
     return stream.pipeThrough(functionCallTransformer);
   } else {
-    return stream;
+    return stream.pipeThrough(
+      createStreamDataTransformer(cb == null ? void 0 : cb.experimental_streamData)
+    );
   }
 }
 function createFunctionCallTransformer(callbacks) {
   const textEncoder = new TextEncoder();
   let isFirstChunk = true;
   let aggregatedResponse = "";
+  let aggregatedFinalCompletionResponse = "";
   let isFunctionStreamingIn = false;
   let functionCallMessages = callbacks[__internal__OpenAIFnMessagesSymbol] || [];
+  const isComplexMode = callbacks == null ? void 0 : callbacks.experimental_streamData;
+  const decode = createChunkDecoder();
   return new TransformStream({
     async transform(chunk, controller) {
-      const message = new TextDecoder().decode(chunk);
+      const message = decode(chunk);
+      aggregatedFinalCompletionResponse += message;
       const shouldHandleAsFunction = isFirstChunk && message.startsWith('{"function_call":');
       if (shouldHandleAsFunction) {
         isFunctionStreamingIn = true;
@@ -181,64 +345,80 @@ function createFunctionCallTransformer(callbacks) {
         return;
       }
       if (!isFunctionStreamingIn) {
-        controller.enqueue(chunk);
+        controller.enqueue(
+          isComplexMode ? textEncoder.encode(getStreamString("text", message)) : chunk
+        );
         return;
       } else {
         aggregatedResponse += message;
       }
     },
     async flush(controller) {
-      const isEndOfFunction = !isFirstChunk && callbacks.experimental_onFunctionCall && isFunctionStreamingIn;
-      if (isEndOfFunction && callbacks.experimental_onFunctionCall) {
-        isFunctionStreamingIn = false;
-        const payload = JSON.parse(aggregatedResponse);
-        const argumentsPayload = JSON.parse(payload.function_call.arguments);
-        let newFunctionCallMessages = [...functionCallMessages];
-        const functionResponse = await callbacks.experimental_onFunctionCall(
-          {
-            name: payload.function_call.name,
-            arguments: argumentsPayload
-          },
-          (result) => {
-            newFunctionCallMessages = [
-              ...functionCallMessages,
-              {
-                role: "assistant",
-                content: "",
-                function_call: payload.function_call
-              },
-              {
-                role: "function",
-                name: payload.function_call.name,
-                content: JSON.stringify(result)
-              }
-            ];
-            return newFunctionCallMessages;
+      try {
+        const isEndOfFunction = !isFirstChunk && callbacks.experimental_onFunctionCall && isFunctionStreamingIn;
+        if (isEndOfFunction && callbacks.experimental_onFunctionCall) {
+          isFunctionStreamingIn = false;
+          const payload = JSON.parse(aggregatedResponse);
+          const argumentsPayload = JSON.parse(payload.function_call.arguments);
+          let newFunctionCallMessages = [
+            ...functionCallMessages
+          ];
+          const functionResponse = await callbacks.experimental_onFunctionCall(
+            {
+              name: payload.function_call.name,
+              arguments: argumentsPayload
+            },
+            (result) => {
+              newFunctionCallMessages = [
+                ...functionCallMessages,
+                {
+                  role: "assistant",
+                  content: "",
+                  function_call: payload.function_call
+                },
+                {
+                  role: "function",
+                  name: payload.function_call.name,
+                  content: JSON.stringify(result)
+                }
+              ];
+              return newFunctionCallMessages;
+            }
+          );
+          if (!functionResponse) {
+            controller.enqueue(
+              textEncoder.encode(
+                isComplexMode ? getStreamString("function_call", aggregatedResponse) : aggregatedResponse
+              )
+            );
+            return;
+          } else if (typeof functionResponse === "string") {
+            controller.enqueue(
+              isComplexMode ? textEncoder.encode(getStreamString("text", functionResponse)) : textEncoder.encode(functionResponse)
+            );
+            return;
           }
-        );
-        if (!functionResponse) {
-          controller.enqueue(textEncoder.encode(aggregatedResponse));
-          return;
-        } else if (typeof functionResponse === "string") {
-          controller.enqueue(textEncoder.encode(functionResponse));
-          return;
-        }
-        const filteredCallbacks = {
-          ...callbacks,
-          onStart: void 0,
-          onCompletion: void 0
-        };
-        const openAIStream = OpenAIStream(functionResponse, {
-          ...filteredCallbacks,
-          [__internal__OpenAIFnMessagesSymbol]: newFunctionCallMessages
-        });
-        const reader = openAIStream.getReader();
-        while (true) {
-          const { done, value } = await reader.read();
-          if (done) {
-            break;
+          const filteredCallbacks = {
+            ...callbacks,
+            onStart: void 0
+          };
+          callbacks.onFinal = void 0;
+          const openAIStream = OpenAIStream(functionResponse, {
+            ...filteredCallbacks,
+            [__internal__OpenAIFnMessagesSymbol]: newFunctionCallMessages
+          });
+          const reader = openAIStream.getReader();
+          while (true) {
+            const { done, value } = await reader.read();
+            if (done) {
+              break;
+            }
+            controller.enqueue(value);
           }
-          controller.enqueue(value);
+        }
+      } finally {
+        if (callbacks.onFinal && aggregatedFinalCompletionResponse) {
+          await callbacks.onFinal(aggregatedFinalCompletionResponse);
         }
       }
     }
@@ -247,12 +427,17 @@ function createFunctionCallTransformer(callbacks) {
 
 // streams/streaming-text-response.ts
 var StreamingTextResponse = class extends Response {
-  constructor(res, init) {
-    super(res, {
+  constructor(res, init, data) {
+    let processedStream = res;
+    if (data) {
+      processedStream = res.pipeThrough(data.stream);
+    }
+    super(processedStream, {
       ...init,
       status: 200,
       headers: {
         "Content-Type": "text/plain; charset=utf-8",
+        [COMPLEX_HEADER]: data ? "true" : "false",
         ...init == null ? void 0 : init.headers
       }
     });
@@ -304,7 +489,9 @@ function createParser2(res) {
   });
 }
 function HuggingFaceStream(res, callbacks) {
-  return createParser2(res).pipeThrough(createCallbacksTransformer(callbacks));
+  return createParser2(res).pipeThrough(createCallbacksTransformer(callbacks)).pipeThrough(
+    createStreamDataTransformer(callbacks == null ? void 0 : callbacks.experimental_streamData)
+  );
 }
 
 // streams/cohere-stream.ts
@@ -351,7 +538,9 @@ function createParser3(res) {
   });
 }
 function CohereStream(reader, callbacks) {
-  return createParser3(reader).pipeThrough(createCallbacksTransformer(callbacks));
+  return createParser3(reader).pipeThrough(createCallbacksTransformer(callbacks)).pipeThrough(
+    createStreamDataTransformer(callbacks == null ? void 0 : callbacks.experimental_streamData)
+  );
 }
 
 // streams/anthropic-stream.ts
@@ -383,11 +572,11 @@ async function* streamable2(stream) {
 }
 function AnthropicStream(res, cb) {
   if (Symbol.asyncIterator in res) {
-    return readableFromAsyncIterable(streamable2(res)).pipeThrough(
-      createCallbacksTransformer(cb)
-    );
+    return readableFromAsyncIterable(streamable2(res)).pipeThrough(createCallbacksTransformer(cb)).pipeThrough(createStreamDataTransformer(cb == null ? void 0 : cb.experimental_streamData));
   } else {
-    return AIStream(res, parseAnthropicStream(), cb);
+    return AIStream(res, parseAnthropicStream(), cb).pipeThrough(
+      createStreamDataTransformer(cb == null ? void 0 : cb.experimental_streamData)
+    );
   }
 }
 
@@ -412,7 +601,9 @@ function LangChainStream(callbacks) {
     }
   };
   return {
-    stream: stream.readable.pipeThrough(createCallbacksTransformer(callbacks)),
+    stream: stream.readable.pipeThrough(createCallbacksTransformer(callbacks)).pipeThrough(
+      createStreamDataTransformer(callbacks == null ? void 0 : callbacks.experimental_streamData)
+    ),
     handlers: {
       handleLLMNewToken: async (token) => {
         await writer.ready;
@@ -465,35 +656,29 @@ async function ReplicateStream(res, cb) {
       Accept: "text/event-stream"
     }
   });
-  return AIStream(eventStream, void 0, cb);
-}
-
-// shared/utils.ts
-import { customAlphabet } from "nanoid/non-secure";
-var nanoid = customAlphabet(
-  "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz",
-  7
-);
-function createChunkDecoder() {
-  const decoder = new TextDecoder();
-  return function(chunk) {
-    if (!chunk)
-      return "";
-    return decoder.decode(chunk, { stream: true });
-  };
+  return AIStream(eventStream, void 0, cb).pipeThrough(
+    createStreamDataTransformer(cb == null ? void 0 : cb.experimental_streamData)
+  );
 }
 export {
   AIStream,
   AnthropicStream,
+  COMPLEX_HEADER,
   CohereStream,
   HuggingFaceStream,
   LangChainStream,
   OpenAIStream,
   ReplicateStream,
+  StreamStringPrefixes,
   StreamingTextResponse,
   createCallbacksTransformer,
   createChunkDecoder,
   createEventStreamTransformer,
+  createStreamDataTransformer,
+  experimental_StreamData,
+  getStreamString,
+  getStreamStringTypeAndValue,
+  isStreamStringEqualToType,
   nanoid,
   readableFromAsyncIterable,
   streamToResponse,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai",
-  "version": "2.2.1",
+  "version": "2.2.2",
   "license": "Apache-2.0",
   "sideEffects": false,
   "main": "./dist/index.js",

package/react/dist/index.d.ts

@@ -203,6 +203,8 @@ type UseChatHelpers = {
     metadata?: Object;
     /** Whether the API request is in progress */
     isLoading: boolean;
+    /** Additional data added on the server via StreamData */
+    data?: any;
 };
 declare function useChat({ api, id, initialMessages, initialInput, sendExtraMessageFields, experimental_onFunctionCall, onResponse, onFinish, onError, credentials, headers, body }?: UseChatOptions): UseChatHelpers;