frontend/jupyter/llm: generalize cell content context for llm-tool and ai-cell-generator

Author: haraldschilly

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

@@ -7,7 +7,6 @@ import {
   Dropdown,
   Flex,
   Input,
-  InputNumber,
   Popover,
   Space,
   Switch,
@@ -39,6 +38,7 @@ import LLMSelector, {
 } from "@cocalc/frontend/frame-editors/llm/llm-selector";
 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 { LLMCostEstimation } from "@cocalc/frontend/misc/llm-cost-estimation";
 import { useProjectContext } from "@cocalc/frontend/project/context";
@@ -60,12 +60,13 @@ import {
 } from "@cocalc/util/misc";
 import { COLORS } from "@cocalc/util/theme";
 import NBViewer from "../nbviewer/nbviewer";
-import { getPreviousNonemptyCellContents } from "../util/cell-content";
+import {
+  CellContextContent,
+  getNonemptyCellContents,
+} from "../util/cell-content";
 import { Position } from "./types";
 import { insertCell } from "./util";
 
-type PrevCells = "none" | number | "all above";
-
 type Cell = { cell_type: "markdown" | "code"; source: string[] };
 type Cells = Cell[];
 
@@ -127,8 +128,8 @@ export function AIGenerateCodeCell({
   const [model, setModel] = useLanguageModelSetting(project_id);
   const [prompt, setPrompt] = useState<string>("");
   const [cellTypes, setCellTypes] = useState<"code" | "all">("code");
-  const [includePreviousCells, setIncludePreviousCells] =
-    useState<PrevCells>(2);
+  // Context for the new selector component - default to 2 previous cells, 0 after
+  const [contextRange, setContextRange] = useState<[number, number]>([-2, 0]);
   const [error, setError] = useState<string>();
   const [preview, setPreview] = useState<Cells | null>(null);
   const [attribute, setAttribute] = useState<boolean>(false);
@@ -141,7 +142,7 @@ export function AIGenerateCodeCell({
 
   const open = showAICellGen != null;
 
-  const prevCodeContents = getPrevCodeContents();
+  const contextContent = getContextContents();
 
   const inputPrompt = getInput({
     frameActions,
@@ -150,7 +151,8 @@ export function AIGenerateCodeCell({
     kernel_name,
     position: showAICellGen,
     model,
-    prevCodeContents,
+    contextContent,
+    contextRange,
   });
 
   const { input } = inputPrompt;
@@ -193,16 +195,22 @@ export function AIGenerateCodeCell({
     }
   }, [preview, open]);
 
-  function getPrevCodeContents(): string {
-    if (includePreviousCells === 0 || showAICellGen == null) return "";
-    return getPreviousNonemptyCellContents(
-      frameActions.current,
+  function getContextContents(): CellContextContent {
+    const prevCount = -contextRange[0]; // contextRange[0] is negative, so -(-2) = 2
+    const nextCount = contextRange[1]; // contextRange[1] is positive for cells after
+
+    if (prevCount === 0 && nextCount === 0) return {};
+
+    return getNonemptyCellContents({
+      actions: frameActions.current,
       id,
-      showAICellGen,
-      includePreviousCells,
+      direction: "around",
+      cellCount: "all", // Use "all" for around direction
       cellTypes,
       lang,
-    );
+      aboveCount: prevCount,
+      belowCount: nextCount,
+    });
   }
 
   function insertCells() {
@@ -273,8 +281,9 @@ export function AIGenerateCodeCell({
   }
 
   async function queryLanguageModel({
-    prevCodeContents,
-    includePreviousCells,
+    contextContent,
+  }: {
+    contextContent: CellContextContent;
   }) {
     if (!prompt.trim()) return;
 
@@ -285,7 +294,8 @@ export function AIGenerateCodeCell({
       model,
       position: showAICellGen,
       prompt,
-      prevCodeContents,
+      contextContent,
+      contextRange,
     });
 
     if (!input) {
@@ -300,7 +310,7 @@ export function AIGenerateCodeCell({
         tag,
         type: "generate",
         model,
-        prev: includePreviousCells,
+        contextRange,
       });
 
       const stream = await webapp_client.openai_client.queryStream({
@@ -364,16 +374,15 @@ export function AIGenerateCodeCell({
     }
   }
 
-  function doQuery(prevCodeContents: string) {
+  function doQuery(contextContent: CellContextContent) {
     cancel.current = false;
     setError("");
     setQuerying(true);
 
     if (showAICellGen == null) return;
 
     queryLanguageModel({
-      prevCodeContents,
-      includePreviousCells,
+      contextContent,
     });
 
     // we also log this
@@ -427,74 +436,20 @@ export function AIGenerateCodeCell({
   }
 
   function renderContext() {
-    const cellStr = `${cellTypes === "code" ? "code " : ""} cell`;
     return (
       <>
         <Divider orientation="left">
           <Text>Context</Text>
         </Divider>
-        <Paragraph>
-          <Flex dir="horizontal" gap="10px" align="center" justify="center">
-            <Flex flex={1}>
-              <div>
-                Include{" "}
-                {typeof includePreviousCells === "number" ? (
-                  <>
-                    previous{" "}
-                    <InputNumber
-                      min={0}
-                      max={10}
-                      size={"small"}
-                      value={includePreviousCells}
-                      onChange={(value) => setIncludePreviousCells(value ?? 1)}
-                    />{" "}
-                    {plural(
-                      includePreviousCells,
-                      `${cellStr}.`,
-                      `${cellStr}s.`,
-                    )}
-                  </>
-                ) : includePreviousCells === "all above" ? (
-                  `all previous ${cellStr}s`
-                ) : (
-                  `no ${cellStr}s`
-                )}
-              </div>
-            </Flex>
-            <Flex flex={0}>
-              {["none", 1, 2, 3, 5, 10, "all above"].map((i: PrevCells) => {
-                const c = getRandomColor(`${i}`);
-                return (
-                  <Tag
-                    key={i}
-                    color={c}
-                    style={{ cursor: "pointer" }}
-                    onClick={() => setIncludePreviousCells(i)}
-                  >
-                    {i}
-                  </Tag>
-                );
-              })}
-            </Flex>
-          </Flex>
-        </Paragraph>
-        <Paragraph>
-          <Flex align="center" gap="10px">
-            <Flex flex={0}>
-              <Switch
-                defaultChecked={cellTypes === "all"}
-                onChange={(val) => setCellTypes(val ? "all" : "code")}
-                unCheckedChildren={"Code cells"}
-                checkedChildren={"All Cells"}
-              />
-            </Flex>
-            <Flex flex={1}>
-              <Text type="secondary">
-                Include only code cells, or all types of cells.
-              </Text>
-            </Flex>
-          </Flex>
-        </Paragraph>
+        <LLMCellContextSelector
+          contextRange={contextRange}
+          onContextRangeChange={setContextRange}
+          cellTypes={cellTypes}
+          onCellTypesChange={setCellTypes}
+          currentCellId={id}
+          frameActions={frameActions.current}
+          mode="insert-position"
+        />
       </>
     );
   }
@@ -643,7 +598,7 @@ export function AIGenerateCodeCell({
               if (!e.shiftKey) return;
               e.preventDefault(); // prevent the default action
               e.stopPropagation(); // stop event propagation
-              doQuery(prevCodeContents);
+              doQuery(contextContent);
             }}
             autoSize={{ minRows: 2, maxRows: 6 }}
           />
@@ -657,7 +612,7 @@ export function AIGenerateCodeCell({
             <LLMQueryDropdownButton
               disabled={!prompt.trim()}
               loading={querying}
-              onClick={() => doQuery(prevCodeContents)}
+              onClick={() => doQuery(contextContent)}
               llmTools={llmTools}
               task="Generate using"
             />
@@ -724,9 +679,10 @@ interface GetInputProps {
   model: LanguageModel;
   position: Position;
   prompt: string;
-  prevCodeContents: string;
+  contextContent: CellContextContent;
   lang: string;
   kernel_name: string;
+  contextRange: [number, number];
 }
 
 function getInputPrompt(prompt: string): string {
@@ -736,9 +692,10 @@ function getInputPrompt(prompt: string): string {
 function getInput({
   frameActions,
   prompt,
-  prevCodeContents,
+  contextContent,
   lang,
   kernel_name,
+  contextRange,
 }: GetInputProps): {
   input: string;
   system: string;
@@ -753,9 +710,30 @@ function getInput({
     );
     return { input: "", system: "", history: [] };
   }
-  const prevCode = prevCodeContents
-    ? `The context after which to insert the cells is:\n\n<context>\n${prevCodeContents}\n\</context>\n\n`
-    : "";
+
+  const prevCount = -contextRange[0]; // cells before insertion point
+  const afterCount = contextRange[1]; // cells after insertion point
+
+  let contextInfo = "";
+
+  if (contextContent.before || contextContent.after) {
+    const beforeCells =
+      prevCount > 0 ? `${prevCount} cells before` : "no cells before";
+    const afterCells =
+      afterCount > 0 ? `${afterCount} cells after` : "no cells after";
+    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`;
+    }
+
+    if (contextContent.after) {
+      contextInfo += `Cells AFTER insertion point:\n<after>\n${contextContent.after}\n</after>\n\n`;
+    }
+  } else {
+    contextInfo =
+      "Context: The new cell will be inserted at the beginning or end of the notebook.\n\n";
+  }
 
   const history: Message[] = [
     { role: "user", content: getInputPrompt("Show the value of foo.") },
@@ -766,8 +744,8 @@ function getInput({
   ];
 
   return {
-    input: `${prevCode}${getInputPrompt(prompt)}`,
+    input: `${contextInfo}${getInputPrompt(prompt)}`,
     history,
-    system: `Create one or more code cells in a Jupyter Notebook.\n\nKernel: "${kernel_name}".\n\nProgramming language: "${lang}".\n\nEach code cell must be wrapped in triple backticks. Do not say what the output will be. Be brief.`,
+    system: `Create one or more code cells in a Jupyter Notebook.\n\nKernel: "${kernel_name}".\n\nProgramming language: "${lang}".\n\nThe new cell(s) will be inserted at a specific position in the notebook. Pay attention to the context provided - cells marked as BEFORE come before the insertion point, cells marked as AFTER come after the insertion point.\n\nEach code cell must be wrapped in triple backticks. Do not say what the output will be. Be brief.`,
   };
 }