feat: TextToDiagram v2 (#10530)

* feat: introducing TextToDiagram v2 feature * fix: eslint issue * debug mermaid bundle size * tests: covering the utils * fix: import mock chunks dynamically to shrink the bundle size * fix: removing replay feature * fix: removing unused prop * fix: bumping workbox cache limit * snapshots + yarn.lock * bump mermaid-to-excalidraw@2 and split into its own chunk * bump node@20 * css tweaks * move files around & rewrite stream chunk schema * random naming & file structure refactor + some tweaks * fix preview theme * support custom warning renderer * label and css fix * fix and rwrite 429 handling * fix label --------- Co-authored-by: dwelle <5153846+dwelle@users.noreply.github.com>
2026-01-15 19:15:41 +01:00
parent 6ebf52279d
commit a0b98a944f
45 changed files with 5226 additions and 942 deletions
@@ -0,0 +1,252 @@
+import { useRef } from "react";
+import { parseMermaidToExcalidraw } from "@excalidraw/mermaid-to-excalidraw";
+import { isFiniteNumber } from "@excalidraw/math";
+
+import { useAtom } from "../../../editor-jotai";
+
+import { trackEvent } from "../../../analytics";
+import { t } from "../../../i18n";
+
+import { errorAtom, rateLimitsAtom, chatHistoryAtom } from "../TTDContext";
+import { useChatAgent } from "../Chat";
+
+import {
+  addMessages,
+  getLastAssistantMessage,
+  getMessagesForLLM,
+  removeLastAssistantMessage,
+  updateAssistantContent,
+} from "../utils/chat";
+
+import type { TTTDDialog } from "../types";
+
+const MIN_PROMPT_LENGTH = 3;
+const MAX_PROMPT_LENGTH = 10000;
+
+export const useTextGeneration = ({
+  onTextSubmit,
+}: {
+  onTextSubmit: (
+    props: TTTDDialog.OnTextSubmitProps,
+  ) => Promise<TTTDDialog.OnTextSubmitRetValue>;
+}) => {
+  const [, setError] = useAtom(errorAtom);
+  const [rateLimits, setRateLimits] = useAtom(rateLimitsAtom);
+  const [chatHistory, setChatHistory] = useAtom(chatHistoryAtom);
+
+  const { addUserMessage, addAssistantMessage, setAssistantError } =
+    useChatAgent();
+
+  const streamingAbortControllerRef = useRef<AbortController | null>(null);
+
+  const validatePrompt = (prompt: string): boolean => {
+    if (
+      prompt.length > MAX_PROMPT_LENGTH ||
+      prompt.length < MIN_PROMPT_LENGTH ||
+      rateLimits?.rateLimitRemaining === 0
+    ) {
+      if (prompt.length < MIN_PROMPT_LENGTH) {
+        setError(
+          new Error(
+            t("chat.errors.promptTooShort", { min: MIN_PROMPT_LENGTH }),
+          ),
+        );
+      }
+      if (prompt.length > MAX_PROMPT_LENGTH) {
+        setError(
+          new Error(t("chat.errors.promptTooLong", { max: MAX_PROMPT_LENGTH })),
+        );
+      }
+
+      return false;
+    }
+    return true;
+  };
+
+  const getReadableErrorMsg = (msg: string) => {
+    try {
+      const content = JSON.parse(msg);
+
+      const innerMessages = JSON.parse(content.message);
+
+      return innerMessages
+        .map((oneMsg: { message: string }) => oneMsg.message)
+        .join("\n");
+    } catch (err) {
+      return msg;
+    }
+  };
+
+  const handleError = (error: Error, errorType: "parse" | "network") => {
+    if (errorType === "parse") {
+      trackEvent("ai", "mermaid parse failed", "ttd");
+    }
+
+    const msg = getReadableErrorMsg(error.message);
+
+    setAssistantError(msg, errorType);
+    setError(error);
+  };
+
+  const onGenerate = async (
+    promptWithContext: string,
+    isRepairFlow = false,
+  ) => {
+    if (!validatePrompt(promptWithContext)) {
+      return;
+    }
+
+    if (streamingAbortControllerRef.current) {
+      streamingAbortControllerRef.current.abort();
+    }
+
+    setError(null);
+
+    const abortController = new AbortController();
+    streamingAbortControllerRef.current = abortController;
+
+    if (!isRepairFlow) {
+      addUserMessage(promptWithContext);
+      addAssistantMessage();
+    } else {
+      const lastAsisstantMessage = getLastAssistantMessage(chatHistory);
+
+      if (lastAsisstantMessage?.errorType === "network") {
+        setChatHistory((prev) =>
+          updateAssistantContent(prev, {
+            isGenerating: true,
+            error: undefined,
+            errorType: undefined,
+            errorDetails: undefined,
+          }),
+        );
+      }
+    }
+
+    try {
+      trackEvent("ai", "generate", "ttd");
+
+      const previousMessages = getMessagesForLLM(chatHistory);
+
+      const { generatedResponse, error, rateLimit, rateLimitRemaining } =
+        await onTextSubmit({
+          messages: [
+            ...previousMessages.slice(-3),
+            { role: "user", content: promptWithContext },
+          ],
+          onStreamCreated: () => {
+            if (isRepairFlow) {
+              setChatHistory((prev) =>
+                updateAssistantContent(prev, {
+                  content: "",
+                  error: "",
+                  isGenerating: true,
+                }),
+              );
+            }
+          },
+          onChunk: (chunk: string) => {
+            setChatHistory((prev) => {
+              const lastAssistantMessage = getLastAssistantMessage(prev);
+              return updateAssistantContent(prev, {
+                content: lastAssistantMessage.content + chunk,
+              });
+            });
+          },
+          signal: abortController.signal,
+        });
+
+      setChatHistory((prev) =>
+        updateAssistantContent(prev, {
+          isGenerating: false,
+        }),
+      );
+
+      if (isFiniteNumber(rateLimit) && isFiniteNumber(rateLimitRemaining)) {
+        setRateLimits({ rateLimit, rateLimitRemaining });
+      }
+
+      if (error) {
+        const isAborted =
+          error.name === "AbortError" ||
+          error.message === "Aborted" ||
+          abortController.signal.aborted;
+
+        // do nothing if the request was aborted by the user
+        if (isAborted) {
+          return;
+        }
+
+        if (error.status === 429) {
+          setChatHistory((prev) => {
+            const chatHistory = removeLastAssistantMessage(prev);
+
+            return {
+              ...chatHistory,
+              messages: chatHistory.messages.filter(
+                (msg) =>
+                  msg.type !== "warning" ||
+                  msg.warningType === "rateLimitExceeded" ||
+                  msg.warningType === "messageLimitExceeded",
+              ),
+            };
+          });
+
+          setChatHistory((chatHistory) => {
+            return addMessages(chatHistory, [
+              {
+                type: "warning",
+                warningType:
+                  rateLimitRemaining === 0
+                    ? "messageLimitExceeded"
+                    : "rateLimitExceeded",
+              },
+            ]);
+          });
+
+          return;
+        } else if (rateLimitRemaining === 0) {
+          setChatHistory((chatHistory) => {
+            chatHistory = {
+              ...chatHistory,
+              messages: chatHistory.messages.filter(
+                (msg) =>
+                  msg.type !== "warning" ||
+                  msg.warningType === "rateLimitExceeded" ||
+                  msg.warningType === "messageLimitExceeded",
+              ),
+            };
+            return addMessages(chatHistory, [
+              {
+                type: "warning",
+                warningType: "messageLimitExceeded",
+              },
+            ]);
+          });
+        }
+
+        handleError(error as Error, "network");
+        return;
+      }
+
+      await parseMermaidToExcalidraw(generatedResponse ?? "");
+
+      trackEvent("ai", "mermaid parse success", "ttd");
+    } catch (error: unknown) {
+      handleError(error as Error, "parse");
+    } finally {
+      streamingAbortControllerRef.current = null;
+    }
+  };
+
+  const handleAbort = () => {
+    if (streamingAbortControllerRef.current) {
+      streamingAbortControllerRef.current.abort();
+    }
+  };
+
+  return {
+    onGenerate,
+    handleAbort,
+  };
+};