frontend/llm-history: only use one dstream to safe resources

haraldschilly · haraldschilly · commit f7bef671cae6 · 2025-08-22T14:29:07.000+02:00
diff --git a/src/packages/frontend/frame-editors/llm/use-llm-history.ts b/src/packages/frontend/frame-editors/llm/use-llm-history.ts
@@ -7,16 +7,13 @@
 // To debug LLM history in the browser console:
 c = cc.client.conat_client
 // Get the shared LLM history streams
-generalStream = await c.dstream({account_id: cc.client.account_id, name: 'llm-history-general'})
-formulaStream = await c.dstream({account_id: cc.client.account_id, name: 'llm-history-formula'})
-// View general prompts
-console.log('General LLM prompts:', generalStream.getAll())
-// View formula prompts
-console.log('Formula prompts:', formulaStream.getAll())
+stream = await c.dstream({account_id: cc.client.account_id, name: 'llm-history'})
+// View prompts
+console.log('LLM prompts:', stream.getAll())
 // Add a prompt to general
-generalStream.push("New prompt")
+stream.push("New prompt")
 // Listen to changes
-generalStream.on('change', (prompt) => console.log('New general prompt:', prompt))
+stream.on('change', (prompt) => console.log('New prompt:', prompt))
 */
 
 import { useState } from "react";
@@ -32,16 +29,20 @@ import { reuseInFlight } from "@cocalc/util/reuse-in-flight";
 const MAX_PROMPTS_NUM = 1000;
 const MAX_PROMPTS_BYTES = 1024 * 1024;
 
-export type LLMHistoryType = "general" | "formula";
+export type LLMHistoryType = "general" | "formula" | "generate";
 
-// Cache for dstream instances per type
-const streamCache = new Map<LLMHistoryType, DStream<string>>();
+interface LLMHistoryEntry {
+  type: LLMHistoryType;
+  prompt: string;
+}
+
+// Single cache for the shared dstream
+let streamCache: DStream<LLMHistoryEntry> | null = null;
 
-// Get or create dstream for a specific history type
-const getDStream = reuseInFlight(async (type: LLMHistoryType) => {
-  const cachedStream = streamCache.get(type);
-  if (cachedStream) {
-    return cachedStream;
+// Get or create the single shared dstream
+const getDStream = reuseInFlight(async () => {
+  if (streamCache) {
+    return streamCache;
   }
 
   try {
@@ -53,17 +54,17 @@ const getDStream = reuseInFlight(async (type: LLMHistoryType) => {
     });
 
     const account_id = store.get_account_id();
-    const stream = await webapp_client.conat_client.dstream<string>({
+    const stream = await webapp_client.conat_client.dstream<LLMHistoryEntry>({
       account_id,
-      name: `${CONAT_LLM_HISTORY_KEY}-${type}`,
+      name: CONAT_LLM_HISTORY_KEY,
       config: {
         discard_policy: "old",
         max_msgs: MAX_PROMPTS_NUM,
         max_bytes: MAX_PROMPTS_BYTES,
       },
     });
 
-    streamCache.set(type, stream);
+    streamCache = stream;
     return stream;
   } catch (err) {
     console.warn(`dstream LLM history initialization error -- ${err}`);
@@ -75,22 +76,31 @@ const getDStream = reuseInFlight(async (type: LLMHistoryType) => {
 export function useLLMHistory(type: LLMHistoryType = "general") {
   const [prompts, setPrompts] = useState<string[]>([]);
 
+  // Filter prompts by type and extract just the prompt strings (newest first)
+  function filterPromptsByType(entries: LLMHistoryEntry[]): string[] {
+    return entries
+      .filter((entry) => entry.type === type)
+      .map((entry) => entry.prompt)
+      .reverse();
+  }
+
   // Initialize dstream and set up listeners
   useAsyncEffect(async () => {
     try {
-      const stream = await getDStream(type);
-
-      // Load existing prompts from stream (newest first)
-      const allPrompts = stream.getAll().reverse();
-      setPrompts(allPrompts);
+      const stream = await getDStream();
+      const allEntries = stream.getAll();
+      setPrompts(filterPromptsByType(allEntries));
 
       // Listen for new prompts being added
-      const handleChange = (newPrompt: string) => {
-        setPrompts((prev) => {
-          // Remove duplicate if exists, then add to front
-          const filtered = prev.filter((p) => p !== newPrompt);
-          return [newPrompt, ...filtered];
-        });
+      const handleChange = (newEntry: LLMHistoryEntry) => {
+        // Only update if the new entry matches our type
+        if (newEntry.type === type) {
+          setPrompts((prev) => {
+            // Remove duplicate if exists, then add to front
+            const filtered = prev.filter((p) => p !== newEntry.prompt);
+            return [newEntry.prompt, ...filtered];
+          });
+        }
       };
 
       stream.on("change", handleChange);
@@ -105,35 +115,32 @@ export function useLLMHistory(type: LLMHistoryType = "general") {
   }, [type]);
 
   async function addPrompt(prompt: string) {
-    if (!prompt.trim()) {
-      console.warn("Empty prompt provided");
+    const trimmedPrompt = prompt.trim();
+
+    if (!trimmedPrompt) {
+      console.warn("use-llm-history: ignoring empty prompt");
       return;
     }
 
     try {
-      const stream = await getDStream(type);
-      const trimmedPrompt = prompt.trim();
+      const stream = await getDStream();
 
-      // Add prompt to stream - this will trigger change event
-      stream.push(trimmedPrompt);
+      // Create entry object with type and prompt
+      const entry: LLMHistoryEntry = {
+        type,
+        prompt: trimmedPrompt,
+      };
 
-      // Clean up old prompts if we exceed MAX_PROMPTS
-      const currentLength = stream.length;
-      if (currentLength > MAX_PROMPTS_NUM) {
-        // Note: dstream doesn't have a built-in way to remove old entries
-        // but we limit the display to MAX_PROMPTS in the UI
-        console.warn(
-          `LLM history has ${currentLength} entries, exceeding MAX_PROMPTS=${MAX_PROMPTS_NUM}`,
-        );
-      }
+      // Add entry to stream - this will trigger a change event
+      stream.push(entry);
     } catch (err) {
       console.warn(`Error adding prompt to LLM history -- ${err}`);
     }
   }
 
   async function clearHistory() {
     try {
-      const stream = await getDStream(type);
+      const stream = await getDStream();
 
       // Clear local state immediately
       setPrompts([]);
@@ -142,14 +149,14 @@ export function useLLMHistory(type: LLMHistoryType = "general") {
       await stream.delete();
 
       // Remove from cache so a new stream will be created
-      streamCache.delete(type);
+      streamCache = null;
     } catch (err) {
       console.warn(`Error clearing LLM history -- ${err}`);
       // Reload prompts on error
       try {
-        const stream = await getDStream(type);
-        const allPrompts = stream.getAll().slice(-MAX_PROMPTS_NUM).reverse();
-        setPrompts(allPrompts);
+        const stream = await getDStream();
+        const allEntries = stream.getAll();
+        setPrompts(filterPromptsByType(allEntries));
       } catch (reloadErr) {
         console.warn(
           `Error reloading prompts after clear failure -- ${reloadErr}`,
diff --git a/src/packages/frontend/project/page/home-page/ai-generate-document.tsx b/src/packages/frontend/project/page/home-page/ai-generate-document.tsx
@@ -25,6 +25,7 @@ import { delay } from "awaiting";
 import { debounce, isEmpty, throttle } from "lodash";
 import { useEffect, useRef, useState } from "react";
 import { FormattedMessage, useIntl } from "react-intl";
+
 import { useLanguageModelSetting } from "@cocalc/frontend/account/useLanguageModelSetting";
 import {
   CSS,
@@ -51,10 +52,12 @@ import { file_options } from "@cocalc/frontend/editor-tmp";
 import { Actions as CodeEditorActions } from "@cocalc/frontend/frame-editors/code-editor/actions";
 import { JupyterEditorActions } from "@cocalc/frontend/frame-editors/jupyter-editor/actions";
 import { Actions as LatexActions } from "@cocalc/frontend/frame-editors/latex-editor/actions";
+import { LLMHistorySelector } from "@cocalc/frontend/frame-editors/llm/llm-history-selector";
 import { LLMQueryDropdownButton } from "@cocalc/frontend/frame-editors/llm/llm-query-dropdown";
 import LLMSelector, {
   modelToName,
 } from "@cocalc/frontend/frame-editors/llm/llm-selector";
+import { useLLMHistory } from "@cocalc/frontend/frame-editors/llm/use-llm-history";
 import { Actions as RmdActions } from "@cocalc/frontend/frame-editors/rmd-editor/actions";
 import { dialogs, labels } from "@cocalc/frontend/i18n";
 import getKernelSpec from "@cocalc/frontend/jupyter/kernelspecs";
@@ -156,6 +159,7 @@ function AIGenerateDocument({
   const [paperSize, setPaperSize] = useState<string | null>(null);
   // User's description of document they want to generate.
   const [prompt, setPrompt] = useState<string>("");
+  const { prompts: historyPrompts, addPrompt } = useLLMHistory("generate");
   const [querying, setQuerying] = useState<boolean>(false);
   const [saving, setSaving] = useState<boolean>(false);
   const [error, setError] = useState<string>("");
@@ -362,6 +366,9 @@ function AIGenerateDocument({
     };
     projectActions?.log(event);
 
+    // Add prompt to history
+    addPrompt(prompt);
+
     const what = intl.formatMessage(
       {
         id: "project.page.ai-generate-document.create_document.what",
@@ -578,7 +585,7 @@ function AIGenerateDocument({
       if (cancel.current) {
         // we abort this
         llmStream.removeAllListeners();
-        // singal "finalization"
+        // signal "finalization"
         processTokens(answer, true);
         return;
       }
@@ -645,7 +652,7 @@ function AIGenerateDocument({
         case "ipynb":
         case "ipynb-sagemath":
           return spec != null
-            ? (JUPYTER[spec.language?.toLowerCase()] ?? [])
+            ? JUPYTER[spec.language?.toLowerCase()] ?? []
             : [];
         default:
           return DOCUMENT[ext];
@@ -836,21 +843,28 @@ function AIGenerateDocument({
           />
         </Paragraph>
         <Paragraph>
-          <Input.TextArea
-            ref={promptRef}
-            allowClear
-            autoSize={{ minRows: 3, maxRows: 6 }}
-            maxLength={3000}
-            placeholder={placeholder}
-            value={prompt}
-            disabled={querying}
-            onChange={({ target: { value } }) => setPrompt(value)}
-            onPressEnter={(e) => {
-              if (e.shiftKey) {
-                generate();
-              }
-            }}
-          />
+          <Space.Compact style={{ width: "100%" }}>
+            <Input.TextArea
+              ref={promptRef}
+              allowClear
+              autoSize={{ minRows: 3, maxRows: 6 }}
+              maxLength={3000}
+              placeholder={placeholder}
+              value={prompt}
+              disabled={querying}
+              onChange={({ target: { value } }) => setPrompt(value)}
+              onPressEnter={(e) => {
+                if (e.shiftKey) {
+                  generate();
+                }
+              }}
+            />
+            <LLMHistorySelector
+              prompts={historyPrompts}
+              onSelect={setPrompt}
+              disabled={querying}
+            />
+          </Space.Compact>
         </Paragraph>
         {!error ? renderExamples() : undefined}
         {!error ? renderPaperSize() : undefined}
