Cell references

mito-ai/mito_ai/completions/prompt_builders/agent_system_message.py

@@ -3,6 +3,7 @@
 
 from mito_ai.completions.prompt_builders.prompt_constants import (
     CITATION_RULES,
+    CELL_REFERENCE_RULES,
     FILES_SECTION_HEADING,
     JUPYTER_NOTEBOOK_SECTION_HEADING,
     VARIABLES_SECTION_HEADING,
@@ -291,7 +292,7 @@ def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
 
 Important information:
 1. The message is a short summary of the ALL the work that you've completed on this task. It should not just refer to the final message. It could be something like "I've completed the sales strategy analysis by exploring key relationships in the data and summarizing creating a report with three recommendations to boost sales.""
-2. The message should include citations for any insights that you shared with the user.
+2. The message should include citations for any insights that you shared with the user and cell references for whenever you refer to specific cells that you've updated or created.
 3. The next_steps is an optional list of 2 or 3 suggested follow-up tasks or analyses that the user might want to perform next. These should be concise, actionable suggestions that build on the work you've just completed. For example: ["Export the cleaned data to CSV", "Analyze revenue per customer", "Convert notebook into an app"].
 4. The next_steps should be as relevant to the user's actual task as possible. Try your best not to make generic suggestions like "Analyze the data" or "Visualize the results". For example, if the user just asked you to calculate LTV of their customers, you might suggest the following next steps: ["Graph key LTV drivers: churn and average transaction value", "Visualize LTV per age group"].
 5. If you are not sure what the user might want to do next, err on the side of suggesting next steps instead of making an assumption and using more CELL_UPDATES.
@@ -345,6 +346,9 @@ def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
 ==== 
 {CITATION_RULES}
 
+====
+{CELL_REFERENCE_RULES}
+
 <Citation Example>
 
 ### User Message 1:

mito-ai/mito_ai/completions/prompt_builders/chat_system_message.py

@@ -3,7 +3,8 @@
 
 from mito_ai.completions.prompt_builders.prompt_constants import (
     CHAT_CODE_FORMATTING_RULES,
-    CITATION_RULES, 
+    CITATION_RULES,
+    CELL_REFERENCE_RULES,
     ACTIVE_CELL_ID_SECTION_HEADING, 
     CODE_SECTION_HEADING,
     get_database_rules
@@ -28,6 +29,9 @@ def create_chat_system_message_prompt() -> str:
 ====
 {CITATION_RULES}
 
+====
+{CELL_REFERENCE_RULES}
+
 <Example 1> 
 {ACTIVE_CELL_ID_SECTION_HEADING}
 '7b3a9e2c-5d14-4c83-b2f9-d67891e4a5f2'

mito-ai/mito_ai/completions/prompt_builders/prompt_constants.py

@@ -46,6 +46,28 @@
 8. Do not include the citation in the code block as a comment. ONLY include the citation in the message field of your response.
 """
 
+CELL_REFERENCE_RULES = """RULES FOR REFERENCING CELLS
+
+When referring to specific cells in the notebook in your messages, use cell references so the user can easily navigate to the cell you're talking about. The user sees cells numbered as "Cell 1", "Cell 2", etc., but internally cells are identified by their unique IDs.
+
+To reference a cell, use this format inline in your message:
+[MITO_CELL_REF:cell_id]
+
+This will be displayed to the user as a clickable "Cell N" link that navigates to the referenced cell.
+
+Cell Reference Rules:
+
+1. Use cell references when discussing specific cells you've created or modified (e.g., "I've added the data cleaning code in [MITO_CELL_REF:abc123]").
+2. Use cell references when referring to cells the user mentioned or that contain relevant context.
+3. The cell_id must be an actual cell ID from the notebook - do not make up IDs.
+4. Place the reference inline where it makes sense in your message, similar to how you would write "Cell 3" in natural language.
+5. Do not use cell references in code - only in the message field of your responses.
+6. Cell references are different from citations. Use citations for specific line-level insights; use cell references for general cell-level navigation.
+
+Example:
+"I've loaded the sales data in [MITO_CELL_REF:c68fdf19-db8c-46dd-926f-d90ad35bb3bc] and will now calculate the monthly totals."
+"""
+
 def get_active_cell_output_str(has_active_cell_output: bool) -> str:
     """
     Used to tell the AI about the output of the active code cell. 

mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdown.tsx

@@ -6,7 +6,9 @@
 import React, { useState, useEffect, useRef } from 'react';
 import { ExpandedVariable } from './ChatInput';
 import { getDatabaseConnections, getRules } from '../../../restAPI/RestAPI';
-import { VariableDropdownItem, FileDropdownItem, RuleDropdownItem } from './ChatDropdownItems';
+import { VariableDropdownItem, FileDropdownItem, RuleDropdownItem, CellDropdownItem } from './ChatDropdownItems';
+import { INotebookTracker } from '@jupyterlab/notebook';
+import { getAllCellReferences } from '../../../utils/cellReferences';
 
 interface ChatDropdownProps {
     options: ExpandedVariable[];
@@ -16,6 +18,7 @@ interface ChatDropdownProps {
     isDropdownFromButton?: boolean;
     onFilterChange?: (filterText: string) => void;
     onClose?: () => void;
+    notebookTracker?: INotebookTracker;
 }
 
 interface ChatDropdownVariableOption {
@@ -38,11 +41,19 @@ interface ChatDropdownFileOption {
     file: ExpandedVariable;
 }
 
+interface ChatDropdownCellOption {
+    type: 'cell'
+    cellNumber: number;
+    cellId: string;
+    cellType: string;
+}
+
 export type ChatDropdownOption =
     | ChatDropdownVariableOption
     | ChatDropdownRuleOption
     | ChatDropdownFileOption
-    | ChatDropdownDatabaseOption;
+    | ChatDropdownDatabaseOption
+    | ChatDropdownCellOption;
 
 const priortizeByType = (options: ChatDropdownOption[], maxPerType: number): ChatDropdownOption[] => {
     /* 
@@ -76,6 +87,7 @@ const ChatDropdown: React.FC<ChatDropdownProps> = ({
     isDropdownFromButton = false,
     onFilterChange,
     onClose,
+    notebookTracker,
 }) => {
     const [selectedIndex, setSelectedIndex] = useState(0);
     const [localFilterText, setLocalFilterText] = useState(filterText);
@@ -107,16 +119,29 @@ const ChatDropdown: React.FC<ChatDropdownProps> = ({
     // Use local filter text when search input is shown, otherwise use prop
     const effectiveFilterText = isDropdownFromButton ? localFilterText : filterText;
 
+    // Get cell references if notebook tracker is available
+    const cellReferences = notebookTracker?.currentWidget 
+        ? getAllCellReferences(notebookTracker.currentWidget)
+        : [];
+
     // Create a list of all options with the format 
     // ['type': 'variable', "expandedVariable": variable]
     // ['type': 'rule', "rule": rule]
     // ['type': 'file', "file": file]
+    // ['type': 'cell', "cellNumber": number, "cellId": string]
     const allOptions: ChatDropdownOption[] = [
         // Rules first
         ...rules.map((rule): ChatDropdownRuleOption => ({
             type: 'rule',
             rule: rule
         })),
+        // Cells second (when user types @Cell or @cell)
+        ...cellReferences.map((cell): ChatDropdownCellOption => ({
+            type: 'cell',
+            cellNumber: cell.cellNumber,
+            cellId: cell.cellId,
+            cellType: cell.cellType
+        })),
         // Files second
         ...options
             .filter(variable => variable.file_name) // Only files
@@ -167,6 +192,12 @@ const ChatDropdown: React.FC<ChatDropdownProps> = ({
         } else if (option.type === 'db') {
             return option.variable.variable_name.toLowerCase().includes(effectiveFilterText.toLowerCase()) ||
                 option.variable.value.toLowerCase().includes(effectiveFilterText.toLowerCase());
+        } else if (option.type === 'cell') {
+            // Match "CellN" (no space)
+            const cellText = `cell${option.cellNumber}`;
+            const numberText = String(option.cellNumber);
+            return cellText.includes(effectiveFilterText.toLowerCase()) ||
+                numberText.includes(effectiveFilterText.toLowerCase());
         } else {
             return option.rule.toLowerCase().includes(effectiveFilterText.toLowerCase());
         }
@@ -331,6 +362,20 @@ const ChatDropdown: React.FC<ChatDropdownProps> = ({
                                 />
                             );
                         }
+                        case 'cell': {
+                            const uniqueKey = `cell-${option.cellNumber}`;
+                            return (
+                                <CellDropdownItem
+                                    key={uniqueKey}
+                                    cellNumber={option.cellNumber}
+                                    cellId={option.cellId}
+                                    cellType={option.cellType}
+                                    index={index}
+                                    selectedIndex={selectedIndex}
+                                    onSelect={() => onSelect(option)}
+                                />
+                            );
+                        }
                         default:
                             return null;
                     }

mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdownItems.tsx

@@ -135,4 +135,54 @@ export const RuleDropdownItem: React.FC<RuleDropdownItemProps> = ({ rule, index,
             </span>
         </li>
     )
+}
+
+interface CellDropdownItemProps {
+    cellNumber: number;
+    cellId: string;
+    cellType: string;
+    index: number;
+    selectedIndex: number;
+    onSelect: (cellNumber: number, cellId: string) => void;
+}
+
+export const CellDropdownItem: React.FC<CellDropdownItemProps> = ({ 
+    cellNumber, 
+    cellId, 
+    cellType, 
+    index, 
+    selectedIndex, 
+    onSelect 
+}) => {
+    const getShortType = (type: string): string => {
+        if (type === 'code') {
+            return 'code';
+        } else if (type === 'markdown') {
+            return 'md';
+        } else {
+            return 'raw';
+        }
+    };
+
+    return (
+        <li
+            className={classNames("chat-dropdown-item", { selected: index === selectedIndex })}
+            onClick={() => onSelect(cellNumber, cellId)}
+            data-testid={`chat-dropdown-item-cell-${cellNumber}`}
+        >
+            <span className="chat-dropdown-item-type"
+                title={cellType}
+                data-testid={`chat-dropdown-item-type-cell-${cellNumber}`}
+            >
+                {getShortType(cellType)}
+            </span>
+            <span
+                className="chat-dropdown-item-name"
+                title={`Cell ${cellNumber}`}
+                data-testid={`chat-dropdown-item-name-cell-${cellNumber}`}
+            >
+                Cell{cellNumber}
+            </span>
+        </li>
+    )
 } 

mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx

@@ -10,6 +10,7 @@ import ChatDropdown from './ChatDropdown';
 import { Variable } from '../../ContextManager/VariableInspector';
 import { getActiveCellID, getActiveCellCode } from '../../../utils/notebook';
 import { INotebookTracker } from '@jupyterlab/notebook';
+import { convertCellReferencesToStableFormat } from '../../../utils/cellReferences';
 import '../../../../style/ChatInput.css';
 import '../../../../style/ChatDropdown.css';
 import { useDebouncedFunction } from '../../../hooks/useDebouncedFunction';
@@ -272,21 +273,27 @@ const ChatInput: React.FC<ChatInputProps> = ({
                 ...additionalContext,
                 { type: 'db', value: option.variable.value, display: option.variable.variable_name }
             ]);
+        } else if (option.type === 'cell') {
+            // For cells, add them as @CellN mentions (no space for easier filtering)
+            contextChatRepresentation = `@Cell${option.cellNumber}`
         }
 
+        // Add a space after the selected item so user can continue typing
+        const contextChatRepresentationWithSpace = contextChatRepresentation + ' ';
+
         const newValue =
             input.slice(0, atIndex) +
-            contextChatRepresentation +
+            contextChatRepresentationWithSpace +
             textAfterCursor;
         setInput(newValue);
 
         setDropdownVisible(false);
 
-        // After updating the input value, set the cursor position after the inserted variable name
+        // After updating the input value, set the cursor position after the inserted item and space
         // We use setTimeout to ensure this happens after React's state update
         setTimeout(() => {
             if (textarea) {
-                const newCursorPosition = atIndex + contextChatRepresentation.length;
+                const newCursorPosition = atIndex + contextChatRepresentationWithSpace.length;
                 textarea.focus();
                 textarea.setSelectionRange(newCursorPosition, newCursorPosition);
             }
@@ -306,6 +313,11 @@ const ChatInput: React.FC<ChatInputProps> = ({
         }));
     };
 
+    // Convert @Cell N references to [MITO_CELL_REF:cell_id] format before submitting
+    const processMessageForSubmission = (messageText: string): string => {
+        return convertCellReferencesToStableFormat(messageText, notebookTracker.currentWidget);
+    };
+
     const getExpandedVarialbes = (): ExpandedVariable[] => {
         const activeNotebookContext = contextManager?.getActiveNotebookContext();
         const expandedVariables: ExpandedVariable[] = [
@@ -440,8 +452,9 @@ const ChatInput: React.FC<ChatInputProps> = ({
                         if (e.key === 'Enter' && !e.shiftKey) {
                             e.preventDefault();
                             adjustHeight(true)
+                            const processedMessage = processMessageForSubmission(input);
                             const additionalContextWithoutDisplayNames = getAdditionContextWithoutDisplayNames();
-                            handleSubmitUserMessage(input, messageIndex, additionalContextWithoutDisplayNames);
+                            handleSubmitUserMessage(processedMessage, messageIndex, additionalContextWithoutDisplayNames);
 
                             // Reset
                             setInput('')
@@ -464,15 +477,17 @@ const ChatInput: React.FC<ChatInputProps> = ({
                         isDropdownFromButton={isDropdownFromButton}
                         onFilterChange={setDropdownFilter}
                         onClose={handleDropdownClose}
+                        notebookTracker={notebookTracker}
                     />
                 )}
             </div>
 
             {isEditing &&
                 <div className="message-edit-buttons">
                     <button onClick={() => {
+                        const processedMessage = processMessageForSubmission(input);
                         const additionalContextWithoutDisplayNames = getAdditionContextWithoutDisplayNames();
-                        handleSubmitUserMessage(input, messageIndex, additionalContextWithoutDisplayNames);
+                        handleSubmitUserMessage(processedMessage, messageIndex, additionalContextWithoutDisplayNames);
                     }}>Save</button>
                     <button onClick={onCancel}>Cancel</button>
                 </div>

mito-ai/src/Extensions/AiChat/ChatMessage/Citation.tsx

@@ -6,6 +6,7 @@
 import React from 'react';
 import { scrollToAndHighlightCell } from '../../../utils/notebook';
 import { INotebookTracker } from '@jupyterlab/notebook';
+import { getCellNumberById } from '../../../utils/cellReferences';
 import '../../../../style/Citation.css';
 
 // Citation line can be either a single line number or a range of lines
@@ -31,7 +32,13 @@ const getLineDisplayText = (line: CitationLine): string => {
 // Citation button component
 export const Citation: React.FC<CitationProps> = ({ citationIndex, cellId, line, notebookTracker }): JSX.Element => {
 
+  // Check if the cell exists in the current notebook
+  const cellNumber = getCellNumberById(cellId, notebookTracker.currentWidget);
+  const isMissing = cellNumber === undefined;
+
   const handleClick = (): void => {
+    if (isMissing) return;
+
     const lineStart = typeof line === 'number' ? line : line.start;
     // In order to support old citations that have just one line, we 
     // we set the end line to the start line if only a single line number is provided.
@@ -41,11 +48,16 @@ export const Citation: React.FC<CitationProps> = ({ citationIndex, cellId, line,
     scrollToAndHighlightCell(notebookTracker.currentWidget, cellId, lineStart, lineEnd);
   };
 
+  const className = isMissing ? 'citation-button citation-missing' : 'citation-button';
+  const title = isMissing 
+    ? 'Cell not found (may have been deleted or is in a different notebook)'
+    : getLineDisplayText(line);
+
   return (
     <span
-      className="citation-button"
+      className={className}
       onClick={handleClick}
-      title={getLineDisplayText(line)}
+      title={title}
     >
       {citationIndex}
     </span>