frontend/jupyter/llm: fix ai cell generator context selection

Author: haraldschilly

src/packages/frontend/components/raw-prompt.tsx

@@ -1,8 +1,8 @@
 import { CSS } from "@cocalc/frontend/app-framework";
 import { useBottomScroller } from "@cocalc/frontend/app-framework/use-bottom-scroller";
-import { COLORS } from "@cocalc/util/theme";
-import StaticMarkdown from "@cocalc/frontend/editors/slate/static-markdown";
 import { Paragraph } from "@cocalc/frontend/components";
+import StaticMarkdown from "@cocalc/frontend/editors/slate/static-markdown";
+import { COLORS } from "@cocalc/util/theme";
 
 const STYLE = {
   border: "1px solid lightgrey",
@@ -21,18 +21,20 @@ interface Props {
   input: React.JSX.Element | string;
   style?: CSS;
   scrollBottom?: boolean;
+  rawText?: boolean;
 }
 
 export function RawPrompt({
   input,
   style: style0,
   scrollBottom = false,
+  rawText = false,
 }: Props) {
   const ref = useBottomScroller(scrollBottom, input);
   const style = { ...STYLE, ...style0 };
-  if (typeof input == "string") {
+  if (typeof input == "string" && !rawText) {
     // this looks so much nicer; I realize it doesn't implement scrollBottom.
-    // But just dropping the input as plain text like below just seems 
+    // But just dropping the input as plain text like below just seems
     // utterly broken!
     return <StaticMarkdown style={style} value={input} />;
   } else {

src/packages/frontend/cspell.json

@@ -3,13 +3,17 @@
   "language": "en",
   "enabled": true,
   "words": [
+    "ANOVA",
+    "ARIMA",
     "autograde",
     "autograding",
     "Bioconductor",
+    "bioinformatics",
     "CoCalc",
     "conat",
     "dplyr",
     "formatjs",
+    "geospatial",
     "ggplot",
     "ipython",
     "ipywidgets",
@@ -36,13 +40,15 @@
     "sagews",
     "scikit",
     "simplelocalize",
+    "sklearn",
     "SocketIO",
     "statsmodels",
     "syncdb",
     "syncdoc",
     "syncstring",
     "synctable",
     "synctables",
+    "tidymodels",
     "tidyverse",
     "timetravel",
     "tolerations",

src/packages/frontend/jupyter/_jupyter.sass

@@ -68,7 +68,7 @@
 .cocalc-jupyter-insert-cell
   margin-bottom: 10px
   height: 8px
-  margin-top: 8px
+  margin-top: 10px
   cursor: pointer
   text-align: center
   &:hover
@@ -111,4 +111,4 @@
 // %matplotlib widgets
 // from pylab import plot; plot([1,2],[3,4])
 cocalc-lumino-adapter
-  position: relative
+  position: relative

src/packages/frontend/jupyter/cell-buttonbar.tsx

@@ -31,7 +31,7 @@ import {
   RUN_ALL_CELLS_ABOVE_ICON,
   RUN_ALL_CELLS_BELOW_ICON,
 } from "./consts";
-import { LLMCellTool } from "./llm";
+import { LLMCellTool } from "./llm/cell-tool";
 
 interface Props {
   id: string;
@@ -193,7 +193,14 @@ export const CellButtonBar: React.FC<Props> = React.memo(
 
     function renderCodeBarLLMButtons() {
       if (!llmTools || !haveLLMCellTools || is_readonly) return;
-      return <LLMCellTool id={id} actions={actions} llmTools={llmTools} cellType={isCodeCell ? "code" : "markdown"} />;
+      return (
+        <LLMCellTool
+          id={id}
+          actions={actions}
+          llmTools={llmTools}
+          cellType={isCodeCell ? "code" : "markdown"}
+        />
+      );
     }
 
     function renderCodeBarFormatButton() {

src/packages/frontend/jupyter/insert-cell/ai-cell-generator.tsx

@@ -40,6 +40,7 @@ import { labels } from "@cocalc/frontend/i18n";
 import { JupyterActions } from "@cocalc/frontend/jupyter/browser-actions";
 import { LLMCellContextSelector } from "@cocalc/frontend/jupyter/llm/cell-context-selector";
 import { splitCells } from "@cocalc/frontend/jupyter/llm/split-cells";
+import { backtickSequence } from "@cocalc/frontend/markdown/util";
 import { LLMCostEstimation } from "@cocalc/frontend/misc/llm-cost-estimation";
 import { useProjectContext } from "@cocalc/frontend/project/context";
 import { LLMEvent } from "@cocalc/frontend/project/history/types";
@@ -207,16 +208,17 @@ export function AIGenerateCodeCell({
     const aboveCount = rangeStart < 0 ? Math.abs(rangeStart) : 0;
     const belowCount = rangeEnd > 0 ? rangeEnd : 0;
 
-    return getNonemptyCellContents({
+    const result = getNonemptyCellContents({
       actions: frameActions.current,
       id,
-      direction: "around",
-      cellCount: "all",
       cellTypes,
       lang,
       aboveCount,
       belowCount,
+      includeCurrentCellInAbove: true,
     });
+
+    return result;
   }
 
   function insertCells() {
@@ -454,6 +456,7 @@ export function AIGenerateCodeCell({
           onCellTypesChange={setCellTypes}
           currentCellId={id}
           frameActions={frameActions.current}
+          mode="insertion"
         />
       </>
     );
@@ -574,6 +577,7 @@ export function AIGenerateCodeCell({
               children: (
                 <RawPrompt
                   input={raw}
+                  rawText
                   style={{ border: "none", padding: "0", margin: "0" }}
                 />
               ),
@@ -729,11 +733,13 @@ function getInput({
     contextInfo = `Context: The new cell will be inserted with ${beforeCells} and ${afterCells} the insertion point.\n\n`;
 
     if (contextContent.before) {
-      contextInfo += `Cells BEFORE insertion point:\n<before>\n${contextContent.before}\n</before>\n\n`;
+      const beforeDelim = backtickSequence(contextContent.before);
+      contextInfo += `Cells BEFORE insertion point:\n\n${beforeDelim}\n${contextContent.before}\n${beforeDelim}\n\n`;
     }
 
     if (contextContent.after) {
-      contextInfo += `Cells AFTER insertion point:\n<after>\n${contextContent.after}\n</after>\n\n`;
+      const afterDelim = backtickSequence(contextContent.after);
+      contextInfo += `Cells AFTER insertion point:\n\n${afterDelim}\n${contextContent.after}\n${afterDelim}\n\n`;
     }
   } else {
     contextInfo =

src/packages/frontend/jupyter/insert-cell/index.tsx

@@ -12,6 +12,8 @@ buttons in the hover state (even when tacking mouse moves!),
 which is confusing.
 */
 
+// cSpell:ignore aicell
+
 import { Button, Space, Tooltip } from "antd";
 import { ReactNode } from "react";
 
@@ -96,6 +98,65 @@ export function InsertCell({
     classNames.push("cocalc-jupyter-insert-cell-below");
   }
 
+  const isActiveAIGenerator = !hide && showAICellGen === position;
+
+  function renderControls() {
+    return (
+      <div
+        className="cocalc-jupyter-insert-cell-controls"
+        style={
+          showAICellGen || alwaysShow
+            ? {
+                visibility: "visible",
+                opacity: 1,
+              }
+            : undefined
+        }
+      >
+        <Space size="large">
+          <TinyButton
+            type="code"
+            title="Insert code cell (click line)"
+            handleButtonClick={handleButtonClick}
+          >
+            <Icon name="code" /> Code
+          </TinyButton>
+          <TinyButton
+            type="markdown"
+            title="Insert text cell (shift+click line)"
+            handleButtonClick={handleButtonClick}
+          >
+            <Icon name="pen" /> Text
+          </TinyButton>
+          <TinyButton
+            type="paste"
+            title="Insert clipboard content as cell"
+            handleButtonClick={handleButtonClick}
+          >
+            <Icon name="paste" /> Paste
+          </TinyButton>
+          {showGenerateCell && llmTools && (
+            <TinyButton
+              type="aicell"
+              title="Create code based on your description (alt+click line)"
+              handleButtonClick={handleButtonClick}
+            >
+              <Space>
+                <AIAvatar
+                  backgroundColor={"transparent"}
+                  size={14}
+                  innerStyle={{ color: "default", top: "-2px" }}
+                />
+                Generate...
+              </Space>
+            </TinyButton>
+          )}
+        </Space>
+      </div>
+    );
+  }
+
+  // the AIGenerateCodeCell is only rendered if active – otherwise that dialog is rendered for each cell insertion line
   return (
     <div
       className={classNames.join(" ")}
@@ -105,66 +166,20 @@ export function InsertCell({
       }}
       onClick={showAICellGen ? undefined : handleBarClick}
     >
-      <AIGenerateCodeCell
-        setShowAICellGen={setShowAICellGen}
-        showAICellGen={!hide ? showAICellGen : null}
-        actions={actions}
-        frameActions={frameActions}
-        id={id}
-        llmTools={llmTools}
-      >
-        <div
-          className="cocalc-jupyter-insert-cell-controls"
-          style={
-            showAICellGen || alwaysShow
-              ? {
-                  visibility: "visible",
-                  opacity: 1,
-                }
-              : undefined
-          }
+      {isActiveAIGenerator ? (
+        <AIGenerateCodeCell
+          setShowAICellGen={setShowAICellGen}
+          showAICellGen={showAICellGen}
+          actions={actions}
+          frameActions={frameActions}
+          id={id}
+          llmTools={llmTools}
         >
-          <Space size="large">
-            <TinyButton
-              type="code"
-              title="Insert code cell (click line)"
-              handleButtonClick={handleButtonClick}
-            >
-              <Icon name="code" /> Code
-            </TinyButton>
-            <TinyButton
-              type="markdown"
-              title="Insert text cell (shift+click line)"
-              handleButtonClick={handleButtonClick}
-            >
-              <Icon name="pen" /> Text
-            </TinyButton>
-            <TinyButton
-              type="paste"
-              title="Insert clipboard content as cell"
-              handleButtonClick={handleButtonClick}
-            >
-              <Icon name="paste" /> Paste
-            </TinyButton>
-            {showGenerateCell && llmTools && (
-              <TinyButton
-                type="aicell"
-                title="Create code based on your description (alt+click line)"
-                handleButtonClick={handleButtonClick}
-              >
-                <Space>
-                  <AIAvatar
-                    backgroundColor={"transparent"}
-                    size={14}
-                    innerStyle={{ color: "default", top: "-2px" }}
-                  />
-                  Generate...
-                </Space>
-              </TinyButton>
-            )}
-          </Space>
-        </div>
-      </AIGenerateCodeCell>
+          {renderControls()}
+        </AIGenerateCodeCell>
+      ) : (
+        renderControls()
+      )}
     </div>
   );
 }