refactor: Rework LLM client types and implementations (#2417)

Refactor the LLM client, responses shape and web client to support
streaming and regular json request in a single helper.

apps/builder/app/builder/features/ai/copywriter.tsx

@@ -1,6 +1,10 @@
 import { z } from "zod";
 import untruncateJson from "untruncate-json";
-import { copywriter, requestStream } from "@webstudio-is/ai";
+import {
+  copywriter,
+  request,
+  type StreamingTextResponseType,
+} from "@webstudio-is/ai";
 import { restAiCopy } from "~/shared/router-utils";
 import {
   Box,
@@ -45,24 +49,6 @@ const patchTextInstance = (textInstance: copywriter.TextInstance) => {
   });
 };
 
-const onChunk = (completion: string) => {
-  try {
-    const jsonResponse = z
-      .array(copywriter.TextInstanceSchema)
-      .parse(JSON.parse(untruncateJson(completion)));
-
-    const currenTextInstance = jsonResponse.pop();
-
-    if (currenTextInstance === undefined) {
-      return;
-    }
-
-    patchTextInstance(currenTextInstance);
-  } catch {
-    /**/
-  }
-};
-
 const $textInstances = computed(
   [instancesStore, selectedInstanceStore],
   (instances, selectedInstance) => {
@@ -120,7 +106,7 @@ export const Copywriter = () => {
 
               abort.current = new AbortController();
               setIsLoading(true);
-              requestStream(
+              request<StreamingTextResponseType>(
                 [
                   restAiCopy(),
                   {
@@ -134,12 +120,28 @@ export const Copywriter = () => {
                   },
                 ],
                 {
-                  onChunk,
+                  onChunk: (id, { decoded }) => {
+                    try {
+                      const jsonResponse = z
+                        .array(copywriter.TextInstanceSchema)
+                        .parse(JSON.parse(untruncateJson(decoded)));
+
+                      const currenTextInstance = jsonResponse.pop();
+
+                      if (currenTextInstance === undefined) {
+                        return;
+                      }
+
+                      patchTextInstance(currenTextInstance);
+                    } catch {
+                      /**/
+                    }
+                  },
                 }
               ).then((result) => {
                 abort.current = null;
-                if (typeof result !== "string") {
-                  alert("Error " + result.type);
+                if (result.success === false) {
+                  alert(result.data.message);
                 }
                 setIsLoading(false);
               });

apps/builder/app/routes/rest.ai.copy.ts

@@ -1,6 +1,7 @@
 import type { ActionArgs } from "@remix-run/node";
 import {
   copywriter,
+  createErrorResponse,
   createGptModel,
   type GPTModelMessageFormat,
 } from "@webstudio-is/ai";
@@ -17,19 +18,24 @@ const RequestSchema = copywriter.ContextSchema.extend({
 export const action = async function action({ request }: ActionArgs) {
   if (isFeatureEnabled("aiCopy") === false) {
     return {
-      success: false,
-      type: "featureDisabled",
-      status: 503,
-      message: "The feature is not available",
+      id: "ai",
+      ...createErrorResponse({
+        error: "featureDisabled",
+        status: 503,
+        message: "The feature is not available",
+        debug: "aiCopy feature disabled",
+      }),
     };
   }
 
   if (env.OPENAI_KEY === undefined) {
     return {
-      success: false,
-      type: "invalidApiKey",
-      status: 401,
-      message: "",
+      id: "ai",
+      ...createErrorResponse({
+        error: "ai.invalidApiKey",
+        status: 401,
+        debug: "Invalid OpenAI API key",
+      }),
     };
   }
 
@@ -38,10 +44,12 @@ export const action = async function action({ request }: ActionArgs) {
     env.OPENAI_ORG.startsWith("org-") === false
   ) {
     return {
-      success: false,
-      type: "invalidOrg",
-      status: 401,
-      message: "",
+      id: "ai",
+      ...createErrorResponse({
+        error: "ai.invalidOrg",
+        status: 401,
+        debug: "Invalid OpenAI API organization",
+      }),
     };
   }
 
@@ -49,18 +57,15 @@ export const action = async function action({ request }: ActionArgs) {
 
   if (parsed.success === false) {
     return {
-      success: false,
-      type: "invalidRequest",
-      status: 400,
-      message:
-        process.env.NODE_ENV === "development"
-          ? parsed.error.errors
-          : "Invalid request data",
+      id: "ai",
+      ...createErrorResponse({
+        error: "ai.invalidRequest",
+        status: 401,
+        debug: "Invalid request data",
+      }),
     };
   }
 
-  // @todo add rate limiting
-
   const { projectId, prompt, textInstances } = parsed.data;
 
   // Permissions check
@@ -72,10 +77,13 @@ export const action = async function action({ request }: ActionArgs) {
 
   if (canEdit === false) {
     return {
-      success: false,
-      type: "unauthorized",
-      status: 401,
-      message: "You don't have edit access to this project",
+      id: "ai",
+      ...createErrorResponse({
+        error: "unauthorized",
+        status: 401,
+        message: "You don't have edit access to this project",
+        debug: "Unauthorized access attempt",
+      }),
     };
   }
   // End of Permissions check
@@ -89,11 +97,17 @@ export const action = async function action({ request }: ActionArgs) {
 
   const chain = copywriter.createChain<GPTModelMessageFormat>();
 
-  return chain({
+  const response = await chain({
     model,
     context: {
       prompt,
       textInstances,
     },
   });
+
+  if (response.success) {
+    return response.data;
+  }
+
+  return response;
 };

packages/ai/src/chains/copywriter/index.ts

@@ -1,13 +1,11 @@
 import { z } from "zod";
-import type {
-  Model as BaseModel,
-  ModelMessage,
-  ChainStream,
-} from "../../types";
+import type { Model as BaseModel, ModelMessage, Chain } from "../../types";
 import { formatPrompt } from "../../utils/format-prompt";
 import { prompt as promptSystemTemplate } from "./__generated__/copy.system.prompt";
 import { prompt as promptUserTemplate } from "./__generated__/copy.user.prompt";
 import type { Instance, Instances } from "@webstudio-is/sdk";
+import { createErrorResponse } from "../../utils/create-error-response";
+import type { StreamingTextResponseType } from "../../utils/streaming-text-response";
 
 /**
  * Copywriter chain.
@@ -40,51 +38,71 @@ export type Context = z.infer<typeof ContextSchema>;
 export const ResponseSchema = z.array(TextInstanceSchema);
 export type Response = z.infer<typeof ResponseSchema>;
 
-export const createChain = <ModelMessageFormat>(): ChainStream<
+export const createChain = <ModelMessageFormat>(): Chain<
   BaseModel<ModelMessageFormat>,
-  Context
+  Context,
+  StreamingTextResponseType
 > =>
   async function chain({ model, context }) {
     const { prompt, textInstances } = context;
 
+    const id = "copywriter";
+
+    const llmMessages: ModelMessage[] = [
+      ["system", promptSystemTemplate],
+      [
+        "user",
+        formatPrompt(
+          {
+            prompt,
+            text_nodes: JSON.stringify(textInstances),
+          },
+          promptUserTemplate
+        ),
+      ],
+    ];
+
     if (textInstances.length === 0) {
+      const message = "No text nodes found for the instance";
       return {
-        success: false,
-        type: "genericError",
-        status: 404,
-        message: "No text nodes found for the instance",
+        id,
+        ...createErrorResponse({
+          status: 404,
+          error: "ai.copywriter.textNodesNotFound",
+          message,
+          debug: message,
+        }),
+        llmMessages,
       };
     }
 
     try {
       z.array(TextInstanceSchema).parse(textInstances);
     } catch (error) {
+      const message = "Invalid nodes list";
       return {
-        success: false,
-        type: "parseError",
-        status: 500,
-        message: "Invalid text nodes list",
+        id,
+        ...createErrorResponse({
+          status: 404,
+          error: "ai.parseError",
+          message,
+          debug: message,
+        }),
+        llmMessages,
       };
     }
 
-    const systemMessage: ModelMessage = ["system", promptSystemTemplate];
+    const messages = model.generateMessages(llmMessages);
 
-    const userMessage: ModelMessage = [
-      "user",
-      formatPrompt(
-        {
-          prompt,
-          text_nodes: JSON.stringify(textInstances),
-        },
-        promptUserTemplate
-      ),
-    ];
-
-    const messages = model.generateMessages([systemMessage, userMessage]);
-
-    return model.requestStream({
+    const response = await model.completionStream({
+      id,
       messages,
     });
+
+    return {
+      ...response,
+      llmMessages,
+    };
   };
 
 export const collectTextInstances = ({

packages/ai/src/index.ts

@@ -11,7 +11,8 @@ export {
 
 /** Utils */
 export * from "./utils/request";
+export * from "./utils/create-error-response";
+export * from "./utils/streaming-text-response";
 
 /** Types */
 export * from "./types";
-export * from "./utils/request";

packages/ai/src/models/gpt.ts

@@ -1,15 +1,13 @@
 import OpenAI from "openai";
-import {
-  OpenAIStream,
-  StreamingTextResponse as _StreamingTextResponse,
-} from "ai";
+import { OpenAIStream } from "ai";
 import type {
   Model as BaseModel,
-  ErrorResponse,
+  ModelCompletion,
+  ModelCompletionStream,
   ModelGenerateMessages,
-  ModelRequest,
-  ModelRequestStream,
 } from "../types";
+import { createErrorResponse } from "../utils/create-error-response";
+import { StreamingTextResponse } from "../utils/streaming-text-response";
 
 export type Model = BaseModel<ModelMessageFormat>;
 export type ModelMessageFormat = OpenAI.Chat.Completions.ChatCompletionMessage;
@@ -25,8 +23,8 @@ export type ModelConfig = {
 export const create = function createModel(config: ModelConfig): Model {
   return {
     generateMessages,
-    request: createRequest(config),
-    requestStream: createRequestStream(config),
+    completion: createCompletion(config),
+    completionStream: createCompletionStream(config),
   };
 };
 
@@ -36,96 +34,62 @@ export const generateMessages: ModelGenerateMessages<ModelMessageFormat> = (
   return messages.map(([role, content]) => ({ role, content }));
 };
 
-const createRequest = (config: ModelConfig): ModelRequest<ModelMessageFormat> =>
-  async function request({ messages }) {
-    const response = await fetch(
-      config.endpoint ?? "https://api.openai.com/v1/chat/completions",
-      {
-        method: "POST",
-        headers: {
-          "Content-Type": "application/json",
-          Accept: "application/json",
-          Authorization: `Bearer ${config.apiKey}`,
-          "OpenAI-Organization": config.organization,
-        },
-        body: JSON.stringify({
-          model: config.model ?? "gpt-3.5-turbo",
-          messages,
-          temperature: config.temperature,
-        }),
-      }
-    );
-
-    if (response.ok === false) {
-      return errorResponse(response);
-    }
-
-    const completion: OpenAI.Chat.Completions.ChatCompletion =
-      await response.json();
-
-    return {
-      success: true,
-      choices: completion.choices.map(
-        (choice) => choice?.message?.content || ""
-      ),
-      tokens: {
-        prompt: completion.usage?.prompt_tokens ?? 0,
-        completion: completion.usage?.completion_tokens ?? 0,
-      },
-    };
-  };
-
-export const errorResponse = (response: Response): ErrorResponse => {
-  let errorType = "ai.genericError";
-
-  if ([401, 429, 500, 503].includes(response.status)) {
-    errorType =
-      OpenAIErrors[response.status as 401 | 429 | 500 | 503][
-        response.statusText
-      ] ?? "ai.genericError";
-  }
-
-  return {
-    success: false,
-    type: errorType,
-    status: response.status,
-    message: response.statusText,
-  };
-};
-
-const OpenAIErrors: Record<401 | 429 | 500 | 503, { [key in string]: string }> =
-  {
-    401: {
-      "Invalid Authentication": "ai.invalidAuth",
-      "Incorrect API key provided": "ai.invalidApiKey",
-      "You must be a member of an organization to use the API": "ai.invalidOrg",
-    },
-    429: {
-      "Rate limit reached for requests": "ai.rateLimit",
-
-      "You exceeded your current quota, please check your plan and billing details":
-        "ai.quotaExceed",
-    },
-    500: {
-      "The server had an error while processing your request":
-        "ai.genericError",
-    },
-    503: {
-      "The engine is currently overloaded, please try again later":
-        "ai.overloaded",
-    },
-  };
-
-const createRequestStream = (
+export const createCompletion = (
   config: ModelConfig
-): ModelRequestStream<ModelMessageFormat> =>
-  async function requestStream({ messages }) {
+): ModelCompletion<ModelMessageFormat> =>
+  async function completion({
+    id,
+    messages,
+  }: {
+    id: string;
+    messages: ModelMessageFormat[];
+  }) {
     try {
       const openai = new OpenAI({
         apiKey: config.apiKey,
         organization: config.organization,
       });
 
+      const completion = await openai.chat.completions.create({
+        model: config.model ?? "gpt-3.5-turbo",
+        temperature: config.temperature,
+        messages,
+      });
+
+      return {
+        id,
+        type: "json",
+        success: true,
+        tokens: {
+          prompt: completion.usage?.prompt_tokens || 0,
+          completion: completion.usage?.completion_tokens || 0,
+        },
+        data: {
+          choices: completion.choices.map(
+            (choice) => choice?.message?.content || ""
+          ),
+        },
+      } as const;
+    } catch (error) {
+      return errorToResponse(id, error);
+    }
+  };
+
+export const createCompletionStream = (
+  config: ModelConfig
+): ModelCompletionStream<ModelMessageFormat> =>
+  async function completeStream({
+    id,
+    messages,
+  }: {
+    id: string;
+    messages: ModelMessageFormat[];
+  }) {
+    try {
+      const openai = new OpenAI({
+        apiKey: config.apiKey,
+        organization: config.organization,
+      });
       // Use polyfilled TransformStream because in Webstudio Builder
       // globalThis.fetch is overriden to @remix-run/web-fetch.
       //
@@ -149,62 +113,51 @@ const createRequestStream = (
       });
 
       const stream = OpenAIStream(response);
-      return new StreamingTextResponse(stream);
+      return {
+        id,
+        type: "stream",
+        success: true,
+        data: new StreamingTextResponse(stream),
+        tokens: {
+          prompt: -1,
+          completion: -1,
+        },
+      } as const;
     } catch (error) {
-      let response = new Response(null, {
-        status: 500,
-        statusText:
-          process.env.NODE_ENV === "development" &&
-          error != null &&
-          typeof error === "object" &&
-          "message" in error &&
-          typeof error.message === "string"
-            ? error.message
-            : "",
-      });
-
-      if (error instanceof OpenAI.APIError) {
-        const { status, message } = error;
-        response = new Response(null, {
-          status,
-          statusText: message,
-        });
-      }
-
-      return errorResponse(response);
+      return errorToResponse(id, error);
     }
   };
 
-// vercel/ai's StreamingTextResponse does not include request.headers.raw()
-// which @vercel/remix uses when deployed on vercel.
-// Therefore we use a custom one.
-export class StreamingTextResponse extends _StreamingTextResponse {
-  constructor(res: ReadableStream, init?: ResponseInit) {
-    super(res, init);
-    this.getRequestHeaders();
-  }
+const errorToResponse = (id: string, error: unknown) => {
+  let status = 500;
+  let debug =
+    error != null &&
+    typeof error === "object" &&
+    "message" in error &&
+    typeof error.message === "string"
+      ? error.message
+      : "";
 
-  getRequestHeaders() {
-    return addRawHeaders(this.headers);
-  }
-}
-
-const addRawHeaders = function addRawHeaders(headers: Headers) {
-  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-  // @ts-ignore
-  headers.raw = function () {
-    const rawHeaders: { [k in string]: string[] } = {};
-    const headerEntries = headers.entries();
-    for (const [key, value] of headerEntries) {
-      const headerKey = key.toLowerCase();
-      // eslint-disable-next-line no-prototype-builtins
-      if (rawHeaders.hasOwnProperty(headerKey)) {
-        rawHeaders[headerKey].push(value);
-      } else {
-        rawHeaders[headerKey] = [value];
-      }
+  if (error instanceof OpenAI.APIError) {
+    if (typeof error.status === "number") {
+      status = error.status;
     }
-    return rawHeaders;
-  };
-  return headers;
+    debug += `\n ${error.message}`;
+  }
+
+  return {
+    id,
+    ...createErrorResponse({
+      status,
+      error: getErrorType(error),
+      message: "Something went wrong",
+      debug,
+    }),
+  } as const;
+};
+const getErrorType = (error: unknown) => {
+  if (error instanceof OpenAI.APIError) {
+    return `ai.${error.name}`;
+  }
+  return `ai.unknownError`;
 };

packages/ai/src/types.ts

@@ -1,3 +1,40 @@
+import type { StreamingTextResponseType } from "./utils/streaming-text-response";
+
+/**
+ * Generic Response types used both by Models and Chains.
+ */
+
+// @todo Convert the response types to Zod
+// so that responses can be parsed and validated on the client.
+
+export type Tokens = {
+  prompt: number;
+  completion: number;
+};
+
+export type SuccessResponse<ResponseData> = {
+  type: ResponseData extends StreamingTextResponseType ? "stream" : "json";
+  success: true;
+  tokens: Tokens;
+  data: ResponseData;
+};
+
+export type ErrorResponse = {
+  type: "json";
+  success: false;
+  tokens: Tokens;
+  data: {
+    status: number;
+    error: string;
+    message: string;
+    debug?: string;
+  };
+};
+
+type Response<ResponseData> = {
+  id: string;
+} & (SuccessResponse<ResponseData> | ErrorResponse);
+
 /**
  * Models types.
  *
@@ -6,41 +43,27 @@
  */
 
 export type ModelMessage = ["system" | "user" | "assistant", string];
-export type ModelMessages = Array<ModelMessage>;
 
 export type Model<ModelMessageFormat> = {
   // Turns ModelMessages into a model-specific messages format.
   generateMessages: ModelGenerateMessages<ModelMessageFormat>;
-  // Regular request.
-  request: ModelRequest<ModelMessageFormat>;
-  // Streaming.
-  requestStream: ModelRequestStream<ModelMessageFormat>;
+  completion: ModelCompletion<ModelMessageFormat>;
+  completionStream: ModelCompletionStream<ModelMessageFormat>;
 };
 
-export type ModelGenerateMessages<ModelMessage> = (
-  messages: ModelMessages
-) => ModelMessage[];
+export type ModelGenerateMessages<ModelMessageFormat> = (
+  messages: ModelMessage[]
+) => ModelMessageFormat[];
 
-export type ModelRequestSuccess = {
-  success: true;
-  choices: string[];
-  tokens: {
-    prompt: number;
-    completion: number;
-  };
-};
-
-export type ModelRequest<ModelMessageFormat> = ({
-  messages,
-}: {
+export type ModelCompletion<ModelMessageFormat> = (args: {
+  id: string;
   messages: ReturnType<ModelGenerateMessages<ModelMessageFormat>>;
-}) => Promise<ModelRequestSuccess | ErrorResponse>;
+}) => Promise<Response<{ choices: string[] }>>;
 
-export type ModelRequestStream<ModelMessageFormat> = ({
-  messages,
-}: {
+export type ModelCompletionStream<ModelMessageFormat> = (args: {
+  id: string;
   messages: ReturnType<ModelGenerateMessages<ModelMessageFormat>>;
-}) => Promise<StreamingTextResponse | ErrorResponse>;
+}) => Promise<Response<StreamingTextResponseType>>;
 
 /**
  * Chains types.
@@ -54,34 +77,7 @@ export type ModelRequestStream<ModelMessageFormat> = ({
  * zod types must have a Schema suffix. For example ResponseSchema.
  */
 
-import type { StreamingTextResponse } from "ai";
-
-/* Used on any kind of error response */
-export type ErrorResponse = {
-  success: false;
-  type: "generic_error" | string;
-  status: number;
-  message: string;
-};
-
-export type SuccessResponse<ResponseData> = {
-  success: true;
-  data: ResponseData;
-};
-
-export type ChainResponseSuccess<ResponseData> = {
-  success: true;
-  tokens: number;
-  data: ResponseData;
-  llmMessages: ModelMessages[];
-};
-
 export type Chain<Model, Context, ResponseData> = (args: {
   model: Model;
   context: Context;
-}) => Promise<ChainResponseSuccess<ResponseData> | ErrorResponse>;
-
-export type ChainStream<Model, Context> = (request: {
-  model: Model;
-  context: Context;
-}) => Promise<StreamingTextResponse | ErrorResponse>;
+}) => Promise<Response<ResponseData> & { llmMessages: ModelMessage[] }>;

packages/ai/src/utils/create-error-response.ts

@@ -0,0 +1,28 @@
+import type { ErrorResponse, Tokens } from "../types";
+
+export const createErrorResponse = ({
+  status,
+  error,
+  message,
+  debug,
+  tokens,
+}: {
+  status?: number;
+  error?: string;
+  message?: string;
+  debug?: string;
+  tokens?: Tokens;
+}): ErrorResponse => ({
+  type: "json",
+  success: false,
+  tokens: tokens || {
+    prompt: -1,
+    completion: -1,
+  },
+  data: {
+    status: status || 500,
+    error: error || "ai.unknownError",
+    message: message || "Something went wrong",
+    debug: debug || undefined,
+  },
+});

packages/ai/src/utils/request.ts

@@ -1,123 +1,94 @@
 import { createChunkDecoder } from "ai";
 import type { ErrorResponse, SuccessResponse } from "../types";
+import { createErrorResponse } from "./create-error-response";
+import { StreamingTextResponse } from "./streaming-text-response";
 
 type RequestOptions = {
-  retry?: number;
+  onChunk?: (
+    operationId: string,
+    data: { decoded: string; value: Uint8Array | undefined; done: boolean }
+  ) => void;
 };
 
-export const request = function request<ResponseData>(
+export const request = <ResponseData>(
   fetchArgs: Parameters<typeof fetch>,
-  options: RequestOptions = { retry: 0 }
-): Promise<SuccessResponse<ResponseData> | ErrorResponse> {
-  return fetch(...fetchArgs)
-    .then((res) => {
-      if (res.ok === false) {
-        return {
-          success: false,
-          type: "generic_error",
-          status: res.status,
-          message: res.statusText,
-        } as ErrorResponse;
-      }
-      return res.json();
-    })
-    .catch((error) => {
-      const signal = fetchArgs[1]?.signal;
-      const isAborted = signal?.aborted === true;
-
-      if (
-        isAborted === false &&
-        typeof options.retry === "number" &&
-        options.retry > 0
-      ) {
-        return requestStream(fetchArgs, {
-          ...options,
-          retry: options.retry - 1,
-        });
-      }
-      return {
-        success: false,
-        type: isAborted ? "aborted" : "generic_error",
-        status: 500,
-        message: "",
-      } as ErrorResponse;
-    });
-};
-
-type RequestStreamOptions = {
-  onChunk?: (completion: string) => void;
-  retry?: number;
-};
-
-export const requestStream = async function requestStream(
-  fetchArgs: Parameters<typeof fetch>,
-  options: RequestStreamOptions = {
-    retry: 0,
-  }
-): Promise<string | ErrorResponse> {
+  { onChunk }: RequestOptions
+) => {
   const signal = fetchArgs[1]?.signal;
   return fetch(...fetchArgs)
     .then(async (response) => {
-      if (response.ok === false || response.body == null) {
+      if (response.status !== 200) {
         return {
-          success: false,
-          type: "generic_error",
-          status: response.status,
-          message: response.body
-            ? response.statusText
-            : "The response is empty",
-        } as ErrorResponse;
-      }
-      if (response.headers.get("Content-Type")?.includes("application/json")) {
-        const json: ErrorResponse = await response.json();
-        return json;
+          id: "",
+          ...createErrorResponse({
+            status: response.status,
+            error: "ai.unhandledError",
+            debug: response.statusText,
+          }),
+        } as const;
       }
 
-      let completion = "";
-      const reader = response.body.getReader();
-      const decoder = createChunkDecoder();
+      const isStream =
+        (response.headers.get("content-type") || "").startsWith("text/plain") &&
+        response.body instanceof ReadableStream;
 
-      // eslint-disable-next-line no-constant-condition
-      while (true) {
-        const { done, value } = await reader.read();
+      if (isStream) {
+        // @todo Add delimiter to text streaming response to extract the operation id.
+        const operationId = "copywriter";
 
-        if (done) {
-          break;
+        let completion = "";
+        const reader = response.body.getReader();
+        const decoder = createChunkDecoder();
+
+        // eslint-disable-next-line no-constant-condition
+        while (true) {
+          const { done, value } = await reader.read();
+
+          if (done) {
+            break;
+          }
+
+          completion += decoder(value);
+
+          if (typeof onChunk === "function") {
+            onChunk(operationId, {
+              decoded: completion,
+              value,
+              done,
+            });
+          }
+
+          if (signal?.aborted === true) {
+            reader.cancel();
+            break;
+          }
         }
 
-        completion += decoder(value);
-
-        if (typeof options.onChunk === "function") {
-          options.onChunk(completion);
-        }
-
-        if (signal?.aborted === true) {
-          reader.cancel();
-          break;
-        }
+        return {
+          id: operationId,
+          type: "stream",
+          success: true,
+          data: new StreamingTextResponse(
+            new Blob([completion], { type: "text/plain" }).stream()
+          ),
+          tokens: { prompt: -1, completion: -1 },
+        } as const;
       }
 
-      return completion;
+      // @todo Convert the response types to Zod
+      // so that responses can be parsed and validated on the client.
+      return (await response.json()) as
+        | SuccessResponse<ResponseData>
+        | ErrorResponse;
     })
     .catch((error) => {
-      const isAborted = signal?.aborted === true;
-
-      if (
-        isAborted === false &&
-        typeof options.retry === "number" &&
-        options.retry > 0
-      ) {
-        return requestStream(fetchArgs, {
-          ...options,
-          retry: options.retry - 1,
-        });
-      }
-
       return {
-        success: false,
-        type: isAborted ? "aborted" : "generic_error",
-        status: 500,
-        message: "",
-      } as ErrorResponse;
+        id: "",
+        ...createErrorResponse({
+          status: 500,
+          error: "ai.unhandledError",
+          debug: error.message + "\n" + error.stack,
+        }),
+      } as const;
     });
 };

packages/ai/src/utils/streaming-text-response.ts

@@ -0,0 +1,37 @@
+import { StreamingTextResponse as _StreamingTextResponse } from "ai";
+
+export type StreamingTextResponseType = StreamingTextResponse;
+
+// vercel/ai's StreamingTextResponse does not include request.headers.raw()
+// which @vercel/remix uses when deployed on vercel.
+// Therefore we use a custom one.
+export class StreamingTextResponse extends _StreamingTextResponse {
+  constructor(res: ReadableStream, init?: ResponseInit) {
+    super(res, init);
+    this.getRequestHeaders();
+  }
+
+  getRequestHeaders() {
+    return addRawHeaders(this.headers);
+  }
+}
+
+const addRawHeaders = function addRawHeaders(headers: Headers) {
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  headers.raw = function () {
+    const rawHeaders: { [k in string]: string[] } = {};
+    const headerEntries = headers.entries();
+    for (const [key, value] of headerEntries) {
+      const headerKey = key.toLowerCase();
+      // eslint-disable-next-line no-prototype-builtins
+      if (rawHeaders.hasOwnProperty(headerKey)) {
+        rawHeaders[headerKey].push(value);
+      } else {
+        rawHeaders[headerKey] = [value];
+      }
+    }
+    return rawHeaders;
+  };
+  return headers;
+};