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/mito_ai/completions/prompt_builders/utils.py

@@ -39,6 +39,7 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
     selected_files = [context["value"] for context in additional_context if context.get("type") == "file"]
     selected_db_connections = [context["value"] for context in additional_context if context.get("type") == "db"]
     selected_images = [context["value"] for context in additional_context if context.get("type", "").startswith("image/")]
+    selected_cells = [context["value"] for context in additional_context if context.get("type") == "cell"]
 
     # STEP 2: Create a list of strings (instructions) for each context type
     context_parts = []
@@ -66,6 +67,12 @@ def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]])
             "The following images have been selected by the user to be used in the task:\n"
             + "\n".join(selected_images)
         )
+
+    if len(selected_cells) > 0:
+        context_parts.append(
+            "The following cells have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_cells)
+        )
 
     # STEP 3: Combine into a single string
     return "\n\n".join(context_parts)

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>
+    )
 }