diff --git a/mito-ai/mito_ai/completions/prompt_builders/agent_system_message.py b/mito-ai/mito_ai/completions/prompt_builders/agent_system_message.py index 1c5e47633e..e0ecfcc6b2 100644 --- a/mito-ai/mito_ai/completions/prompt_builders/agent_system_message.py +++ b/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} + ### User Message 1: diff --git a/mito-ai/mito_ai/completions/prompt_builders/chat_system_message.py b/mito-ai/mito_ai/completions/prompt_builders/chat_system_message.py index 643446d393..9e6cb70c59 100644 --- a/mito-ai/mito_ai/completions/prompt_builders/chat_system_message.py +++ b/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} + {ACTIVE_CELL_ID_SECTION_HEADING} '7b3a9e2c-5d14-4c83-b2f9-d67891e4a5f2' diff --git a/mito-ai/mito_ai/completions/prompt_builders/prompt_constants.py b/mito-ai/mito_ai/completions/prompt_builders/prompt_constants.py index 921295f580..157f0eb1b1 100644 --- a/mito-ai/mito_ai/completions/prompt_builders/prompt_constants.py +++ b/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. diff --git a/mito-ai/mito_ai/completions/prompt_builders/utils.py b/mito-ai/mito_ai/completions/prompt_builders/utils.py index e9e2ceb059..312605a091 100644 --- a/mito-ai/mito_ai/completions/prompt_builders/utils.py +++ b/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) diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdown.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdown.tsx index 5ef6252755..689a55e538 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdown.tsx +++ b/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 = ({ isDropdownFromButton = false, onFilterChange, onClose, + notebookTracker, }) => { const [selectedIndex, setSelectedIndex] = useState(0); const [localFilterText, setLocalFilterText] = useState(filterText); @@ -107,16 +119,29 @@ const ChatDropdown: React.FC = ({ // 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 = ({ } 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 = ({ /> ); } + case 'cell': { + const uniqueKey = `cell-${option.cellNumber}`; + return ( + onSelect(option)} + /> + ); + } default: return null; } diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdownItems.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdownItems.tsx index d5b4c4ede8..63fb134ba2 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdownItems.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatDropdownItems.tsx @@ -135,4 +135,54 @@ export const RuleDropdownItem: React.FC = ({ rule, index, ) +} + +interface CellDropdownItemProps { + cellNumber: number; + cellId: string; + cellType: string; + index: number; + selectedIndex: number; + onSelect: (cellNumber: number, cellId: string) => void; +} + +export const CellDropdownItem: React.FC = ({ + 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 ( +
  • onSelect(cellNumber, cellId)} + data-testid={`chat-dropdown-item-cell-${cellNumber}`} + > + + {getShortType(cellType)} + + + Cell{cellNumber} + +
  • + ) } \ No newline at end of file diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx index c09b828d94..b8c5a8771a 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/ChatInput.tsx +++ b/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'; @@ -226,6 +227,15 @@ const ChatInput: React.FC = ({ display: option.variable.variable_name } ]); + } else if (option.type === 'cell') { + setAdditionalContext(prev => [ + ...prev, + { + type: 'cell', + value: option.cellId, + display: `Cell ${option.cellNumber}` + } + ]); } setDropdownVisible(false); @@ -272,21 +282,32 @@ const ChatInput: React.FC = ({ ...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}` + // Store the stable cell ID in additionalContext, not the @CellN format + setAdditionalContext([ + ...additionalContext, + { type: 'cell', value: option.cellId, display: `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 +327,11 @@ const ChatInput: React.FC = ({ })); }; + // 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[] = [ @@ -366,12 +392,28 @@ const ChatInput: React.FC = ({ display: 'Active Cell' }]); } + + // Remove the current notebook context item + const hasNotebookContext = additionalContext.some(context => context.type === 'notebook'); + if (hasNotebookContext) { + setAdditionalContext(prev => prev.filter(context => context.type !== 'notebook')); + } } else if (agentModeEnabled) { // Remove active cell context when in agent mode const hasActiveCellContext = additionalContext.some(context => context.type === 'active_cell'); if (hasActiveCellContext) { setAdditionalContext(prev => prev.filter(context => context.type !== 'active_cell')); } + + const hasNotebookContext = additionalContext.some(context => context.type === 'notebook'); + if (!hasNotebookContext) { + // Add a current notebook context item + setAdditionalContext(prev => [...prev, { + type: 'notebook', + value: 'Notebook', + display: 'Notebook' + }]); + } } }, [agentModeEnabled, additionalContext, activeCellCode]); @@ -407,6 +449,7 @@ const ChatInput: React.FC = ({ onRemove={() => setAdditionalContext(additionalContext.filter((_, i) => i !== index))} notebookTracker={notebookTracker} activeCellID={activeCellID} + value={context.value} /> ))} @@ -440,8 +483,9 @@ const ChatInput: React.FC = ({ 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,6 +508,7 @@ const ChatInput: React.FC = ({ isDropdownFromButton={isDropdownFromButton} onFilterChange={setDropdownFilter} onClose={handleDropdownClose} + notebookTracker={notebookTracker} /> )} @@ -471,8 +516,9 @@ const ChatInput: React.FC = ({ {isEditing &&
    diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/Citation.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/Citation.tsx index b9b764abc8..c32a1343d0 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/Citation.tsx +++ b/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 = ({ 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 = ({ 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 ( {citationIndex} diff --git a/mito-ai/src/Extensions/AiChat/ChatMessage/MarkdownBlock.tsx b/mito-ai/src/Extensions/AiChat/ChatMessage/MarkdownBlock.tsx index b01d6ea9e5..7b041ad3b0 100644 --- a/mito-ai/src/Extensions/AiChat/ChatMessage/MarkdownBlock.tsx +++ b/mito-ai/src/Extensions/AiChat/ChatMessage/MarkdownBlock.tsx @@ -3,11 +3,14 @@ * Distributed under the terms of the GNU Affero General Public License v3.0 License. */ -import React, { useEffect, useState, useRef, useCallback } from 'react'; +import React, { useEffect, useState, useRef, useCallback, useMemo } from 'react'; import { IRenderMimeRegistry, MimeModel } from '@jupyterlab/rendermime'; import { createPortal } from 'react-dom'; import { Citation, CitationProps, CitationLine } from './Citation'; import { INotebookTracker } from '@jupyterlab/notebook'; +import { scrollToCell, highlightCodeCell } from '../../../utils/notebook'; +import { useCellOrder } from '../../../hooks/useCellOrder'; +import '../../../../style/CellReference.css'; /** * React Portals in Markdown Rendering @@ -67,9 +70,26 @@ interface Citation { }; } +interface CellRef { + id: string; + cellId: string; +} + const MarkdownBlock: React.FC = ({ markdown, renderMimeRegistry, notebookTracker }) => { const [citationPortals, setCitationPortals] = useState([]); const containerRef = useRef(null); + + // Track cell order to update cell references when cells are reordered + const cellOrder = useCellOrder(notebookTracker); + + // Create a serialized version of cell order for dependency tracking + // This ensures re-renders when cells are reordered (even if count stays the same) + const cellOrderKey = useMemo(() => { + return Array.from(cellOrder.entries()) + .sort((a, b) => a[0].localeCompare(b[0])) // Sort by cellId for stable string + .map(([cellId, cellNumber]) => `${cellId}:${cellNumber}`) + .join(','); + }, [cellOrder]); // Helper function to parse line numbers or ranges const parseLineNumber = (lineStr: string): CitationLine => { @@ -99,23 +119,32 @@ const MarkdownBlock: React.FC = ({ markdown, renderMimeRegis } }; - // Extract citations from the markdown, returning the markdown with the JSON citations replaced with - // citation placeholders {{${id}}} and an array of citation objects. - const extractCitations = useCallback((text: string): { processedMarkdown: string; citations: Citation[] } => { - // Updated regex to match both single lines and line ranges: [MITO_CITATION:cell_id:line_number] or [MITO_CITATION:cell_id:start_line-end_line] + // Extract citations and cell references from the markdown, returning the markdown with + // placeholders and arrays of citation/cell reference objects. + const extractCitationsAndCellRefs = useCallback((text: string): { + processedMarkdown: string; + citations: Citation[]; + cellRefs: CellRef[]; + } => { + // Regex for citations: [MITO_CITATION:cell_id:line_number] or [MITO_CITATION:cell_id:start_line-end_line] const citationRegex = /\[MITO_CITATION:([^:]+):(\d+(?:-\d+)?)\]/g; + // Regex for cell references: [MITO_CELL_REF:cell_id] + const cellRefRegex = /\[MITO_CELL_REF:([^\]]+)\]/g; + const citations: Citation[] = []; - let counter = 0; + const cellRefs: CellRef[] = []; + let citationCounter = 0; + let cellRefCounter = 0; - // Replace each citation with a placeholder - const processedMarkdown = text.replace(citationRegex, (match, cellId, lineStr) => { + // First, replace citations with placeholders + let processedMarkdown = text.replace(citationRegex, (match, cellId, lineStr) => { try { - const id = `citation-${counter++}`; + const id = `citation-${citationCounter++}`; const line = parseLineNumber(lineStr); citations.push({ id, data: { - citation_index: counter, + citation_index: citationCounter, cell_id: cellId, line: line } @@ -127,7 +156,17 @@ const MarkdownBlock: React.FC = ({ markdown, renderMimeRegis } }); - return { processedMarkdown, citations }; + // Then, replace cell references with placeholders + processedMarkdown = processedMarkdown.replace(cellRefRegex, (match, cellId) => { + const id = `cellref-${cellRefCounter++}`; + cellRefs.push({ + id, + cellId: cellId.trim() + }); + return `{{${id}}}`; + }); + + return { processedMarkdown, citations, cellRefs }; }, []); // Uses the Jupyter markdowm MimeRenderer to render the markdown content as normal HTML @@ -150,14 +189,19 @@ const MarkdownBlock: React.FC = ({ markdown, renderMimeRegis } }, [renderMimeRegistry]); - // Replace the citation placeholders with Citation Components in the DOM - const createCitationPortals = useCallback((citations: Citation[]): React.ReactElement[] => { - if (!containerRef.current || citations.length === 0) return []; + + // Replace the citation and cell reference placeholders with components in the DOM + const createPortalsFromPlaceholders = useCallback(( + citations: Citation[], + cellRefs: CellRef[] + ): React.ReactElement[] => { + if (!containerRef.current || (citations.length === 0 && cellRefs.length === 0)) return []; const newPortals: React.ReactElement[] = []; - // Create a map of placeholder to citation for faster lookup + // Create maps for faster lookup const citationMap = new Map(citations.map(citation => [`{{${citation.id}}}`, citation])); + const cellRefMap = new Map(cellRefs.map(ref => [`{{${ref.id}}}`, ref])); // Find all text nodes that contain our placeholder like {{citation-id}}). // Since these placeholders exist within the text content (not as separate DOM elements): @@ -178,7 +222,7 @@ const MarkdownBlock: React.FC = ({ markdown, renderMimeRegis // Check if this node contains any placeholders let containsPlaceholder = false; - for (const placeholder of citationMap.keys()) { + for (const placeholder of [...citationMap.keys(), ...cellRefMap.keys()]) { if (node.nodeValue.includes(placeholder)) { containsPlaceholder = true; break; @@ -187,8 +231,8 @@ const MarkdownBlock: React.FC = ({ markdown, renderMimeRegis if (!containsPlaceholder) return; - // Create a regex to match all placeholders - const placeholderPattern = /\{\{citation-\d+\}\}/g; + // Create a regex to match all placeholders (citations and cell refs) + const placeholderPattern = /\{\{(citation|cellref)-\d+\}\}/g; const matches = [...node.nodeValue.matchAll(placeholderPattern)]; if (matches.length === 0) return; @@ -199,9 +243,6 @@ const MarkdownBlock: React.FC = ({ markdown, renderMimeRegis matches.forEach(match => { const placeholder = match[0]; - const citation = citationMap.get(placeholder); - if (!citation) return; - const startIndex = match.index!; // Add text before the placeholder @@ -211,23 +252,58 @@ const MarkdownBlock: React.FC = ({ markdown, renderMimeRegis ); } - // Create span for the citation - const span = document.createElement('span'); - span.classList.add('citation-container'); - span.dataset.citationId = citation.id; - fragment.appendChild(span); - - // Create React portal for this span - newPortals.push( - - ); + // Check if it's a citation or cell reference + const citation = citationMap.get(placeholder); + const cellRef = cellRefMap.get(placeholder); + + if (citation) { + // Create span for the citation + const span = document.createElement('span'); + span.classList.add('citation-container'); + span.dataset.citationId = citation.id; + fragment.appendChild(span); + + // Create React portal for this span + newPortals.push( + + ); + } else if (cellRef) { + // Create clickable span for cell reference + const cellNumber = cellOrder.get(cellRef.cellId); + const isMissing = cellNumber === undefined; + const displayText = isMissing ? 'Cell' : `Cell ${cellNumber}`; + + const span = document.createElement('span'); + span.className = isMissing ? 'cell-reference cell-reference-missing' : 'cell-reference'; + span.textContent = displayText; + span.title = isMissing + ? 'Cell not found (may have been deleted or is in a different notebook)' + : `Click to navigate to ${displayText}`; + + // Only add click handler if cell exists + if (!isMissing) { + span.addEventListener('click', (e) => { + e.preventDefault(); + e.stopPropagation(); + if (notebookTracker.currentWidget) { + scrollToCell(notebookTracker.currentWidget, cellRef.cellId, undefined, 'center'); + // Highlight the cell after scrolling + setTimeout(() => { + highlightCodeCell(notebookTracker, cellRef.cellId); + }, 500); + } + }); + } + + fragment.appendChild(span); + } lastIndex = startIndex + placeholder.length; }); @@ -246,24 +322,25 @@ const MarkdownBlock: React.FC = ({ markdown, renderMimeRegis }); return newPortals; - }, []); + }, [notebookTracker, cellOrder]); // Process everything in one effect, but with clear separation via helper functions + // cellOrderKey triggers re-render when notebook loads or cells are reordered (fixes race condition on refresh) useEffect(() => { const processMarkdown = async (): Promise => { - // Step 1: Extract citations and get processed markdown - const { processedMarkdown, citations } = extractCitations(markdown); + // Step 1: Extract citations and cell references, get processed markdown + const { processedMarkdown, citations, cellRefs } = extractCitationsAndCellRefs(markdown); // Step 2: Render markdown with placeholders await renderMarkdownContent(processedMarkdown); - // Step 3: Create and insert citation portals - const portals = createCitationPortals(citations); + // Step 3: Create and insert portals for citations and cell references + const portals = createPortalsFromPlaceholders(citations, cellRefs); setCitationPortals(portals); }; void processMarkdown(); - }, [markdown, extractCitations, renderMarkdownContent, createCitationPortals]); + }, [markdown, extractCitationsAndCellRefs, renderMarkdownContent, createPortalsFromPlaceholders, cellOrderKey]); return (
    diff --git a/mito-ai/src/Extensions/CellNumbering/CellNumberingPlugin.ts b/mito-ai/src/Extensions/CellNumbering/CellNumberingPlugin.ts new file mode 100644 index 0000000000..3cd64f077a --- /dev/null +++ b/mito-ai/src/Extensions/CellNumbering/CellNumberingPlugin.ts @@ -0,0 +1,89 @@ +/* + * Copyright (c) Saga Inc. + * Distributed under the terms of the GNU Affero General Public License v3.0 License. + */ + +import { + JupyterFrontEnd, + JupyterFrontEndPlugin +} from '@jupyterlab/application'; +import { INotebookTracker, NotebookPanel } from '@jupyterlab/notebook'; +import '../../../style/CellNumbering.css'; + +/** + * Updates cell numbers for all cells in a notebook. + * Uses notebook.widgets which is always in the correct order. + */ +function updateAllCellNumbers(notebookPanel: NotebookPanel): void { + const notebook = notebookPanel.content; + notebook.widgets.forEach((cell, index) => { + const header = cell.node.querySelector('.jp-Cell-header'); + if (header) { + // 1-indexed for display (Cell 1, Cell 2, etc.) + header.setAttribute('data-cell-number', String(index + 1)); + } + }); +} + +/** + * Sets up cell numbering for a notebook panel. + */ +function setupCellNumbering(notebookPanel: NotebookPanel): void { + const notebook = notebookPanel.content; + + // Store handler references so they can be disconnected on cleanup + const handleCellsChanged = (): void => { + updateAllCellNumbers(notebookPanel); + }; + + const handleActiveCellChanged = (): void => { + updateAllCellNumbers(notebookPanel); + }; + + // Update when cells are added, removed, or moved + notebook.model?.cells.changed.connect(handleCellsChanged); + + // Update when active cell changes (often happens when scrolling) + notebook.activeCellChanged.connect(handleActiveCellChanged); + + // Use MutationObserver to handle virtualization (cells being attached/detached) + const observer = new MutationObserver(() => { + updateAllCellNumbers(notebookPanel); + }); + observer.observe(notebook.node, { childList: true, subtree: true }); + + // Cleanup all handlers when notebook is disposed + notebookPanel.disposed.connect(() => { + notebook.model?.cells.changed.disconnect(handleCellsChanged); + notebook.activeCellChanged.disconnect(handleActiveCellChanged); + observer.disconnect(); + }); + + // Initial update + updateAllCellNumbers(notebookPanel); +} + +/** + * Plugin to add cell numbering to notebook cells. + */ +const CellNumberingPlugin: JupyterFrontEndPlugin = { + id: 'mito-ai:cell-numbering', + description: 'Add cell numbers to notebook cells', + autoStart: true, + requires: [INotebookTracker], + activate: (app: JupyterFrontEnd, notebookTracker: INotebookTracker): void => { + console.log('mito-ai: CellNumberingPlugin activated'); + + // Setup for all existing notebooks + notebookTracker.forEach((notebookPanel: NotebookPanel) => { + setupCellNumbering(notebookPanel); + }); + + // Setup for newly opened notebooks + notebookTracker.widgetAdded.connect((sender, notebookPanel: NotebookPanel) => { + setupCellNumbering(notebookPanel); + }); + } +}; + +export default CellNumberingPlugin; diff --git a/mito-ai/src/components/SelectedContextContainer.tsx b/mito-ai/src/components/SelectedContextContainer.tsx index a939093d8b..f466fa5a0d 100644 --- a/mito-ai/src/components/SelectedContextContainer.tsx +++ b/mito-ai/src/components/SelectedContextContainer.tsx @@ -9,7 +9,7 @@ import RuleIcon from '../icons/RuleIcon'; import CodeIcon from '../icons/CodeIcon'; import DatabaseIcon from '../icons/DatabaseIcon'; import PhotoIcon from '../icons/PhotoIcon'; -import { highlightCodeCell, getCellByID } from '../utils/notebook'; +import { highlightCodeCell, getCellByID, scrollToCell } from '../utils/notebook'; interface SelectedContextContainerProps { title: string; @@ -18,6 +18,7 @@ interface SelectedContextContainerProps { onClick?: () => void; notebookTracker?: any; activeCellID?: string; + value?: string; // The underlying value (e.g., cellId for cell type) } const SelectedContextContainer: React.FC = ({ @@ -26,7 +27,8 @@ const SelectedContextContainer: React.FC = ({ onRemove, onClick, notebookTracker, - activeCellID + activeCellID, + value }) => { const [isHovered, setIsHovered] = useState(false); @@ -42,6 +44,10 @@ const SelectedContextContainer: React.FC = ({ icon = ; } else if (type === 'active_cell') { icon = ; + } else if (type === 'notebook') { + icon = ; + } else if (type === 'cell') { + icon = ; } const handleClick = (): void => { @@ -60,6 +66,15 @@ const SelectedContextContainer: React.FC = ({ } } // If notebookTracker or activeCellID are not available, do nothing + } else if (type === 'cell' && notebookTracker && value) { + // Handle cell context click - scroll to the cell + if (notebookTracker.currentWidget) { + scrollToCell(notebookTracker.currentWidget, value, undefined, 'center'); + } + // Highlight the cell + setTimeout(() => { + highlightCodeCell(notebookTracker, value); + }, 500); } else if (onClick) { // Call the custom onClick handler for other context types onClick(); @@ -82,7 +97,7 @@ const SelectedContextContainer: React.FC = ({ }} title={isHovered ? "Remove rule" : "Selected rule"} > - {isHovered && type !== 'active_cell' ? ( + {isHovered && type !== 'active_cell' && type !== 'notebook' ? ( X ) : ( {icon} diff --git a/mito-ai/src/hooks/useCellOrder.tsx b/mito-ai/src/hooks/useCellOrder.tsx new file mode 100644 index 0000000000..30e67225d0 --- /dev/null +++ b/mito-ai/src/hooks/useCellOrder.tsx @@ -0,0 +1,78 @@ +/* + * Copyright (c) Saga Inc. + * Distributed under the terms of the GNU Affero General Public License v3.0 License. + */ + +import { useState, useEffect, useMemo } from 'react'; +import { INotebookTracker } from '@jupyterlab/notebook'; + +/** + * Hook that tracks the current cell order in the notebook. + * Returns a map of cellId → cellNumber (1-indexed). + * Updates automatically when cells are added, removed, or reordered. + * + * @param notebookTracker - The notebook tracker to monitor + * @returns A map of cellId to cellNumber (1-indexed) + */ +export const useCellOrder = (notebookTracker: INotebookTracker): Map => { + const [cellOrderKey, setCellOrderKey] = useState(0); + + // Track current widget ID to detect notebook switches + const currentWidgetId = notebookTracker.currentWidget?.id ?? null; + + // Compute the cell order mapping + const cellOrder = useMemo(() => { + const orderMap = new Map(); + const notebookPanel = notebookTracker.currentWidget; + + if (!notebookPanel) { + return orderMap; + } + + const notebook = notebookPanel.content; + notebook.widgets.forEach((cell, index) => { + // 1-indexed cell numbers for display + orderMap.set(cell.model.id, index + 1); + }); + + return orderMap; + }, [notebookTracker, cellOrderKey, currentWidgetId]); + + // Listen to cell changes to trigger re-computation + // Include currentWidgetId in dependencies so listener re-attaches when switching notebooks + useEffect(() => { + const notebookPanel = notebookTracker.currentWidget; + if (!notebookPanel) { + return; + } + + const notebook = notebookPanel.content; + + // Update when cells are added, removed, or reordered + const handleCellChange = (): void => { + setCellOrderKey(prev => prev + 1); + }; + + // Listen to cell model changes (fires when cells are added, removed, or reordered) + notebook.model?.cells.changed.connect(handleCellChange); + + return () => { + notebook.model?.cells.changed.disconnect(handleCellChange); + }; + }, [notebookTracker, currentWidgetId]); + + // Also update when the current notebook changes + useEffect(() => { + const handleNotebookChange = (): void => { + setCellOrderKey(prev => prev + 1); + }; + + notebookTracker.currentChanged.connect(handleNotebookChange); + return () => { + notebookTracker.currentChanged.disconnect(handleNotebookChange); + }; + }, [notebookTracker]); + + return cellOrder; +}; + diff --git a/mito-ai/src/index.ts b/mito-ai/src/index.ts index 30f8fa8f4f..da79224978 100644 --- a/mito-ai/src/index.ts +++ b/mito-ai/src/index.ts @@ -15,6 +15,7 @@ import SettingsManagerPlugin from './Extensions/SettingsManager/SettingsManagerP import { versionCheckPlugin } from './Extensions/VersionCheck'; import NotebookFooterPlugin from './Extensions/NotebookFooter'; import ManageAppsPlugin from "./Extensions/AppManager/ManageAppsPlugin" +import CellNumberingPlugin from './Extensions/CellNumbering/CellNumberingPlugin'; // This is the main entry point to the mito-ai extension. It must export all of the top level // extensions that we want to load. @@ -30,5 +31,6 @@ export default [ SettingsManagerPlugin, versionCheckPlugin, NotebookFooterPlugin, - ManageAppsPlugin + ManageAppsPlugin, + CellNumberingPlugin ]; diff --git a/mito-ai/src/tests/AiChat/ChatInput.test.tsx b/mito-ai/src/tests/AiChat/ChatInput.test.tsx index f355ea887b..e4d08d4d38 100644 --- a/mito-ai/src/tests/AiChat/ChatInput.test.tsx +++ b/mito-ai/src/tests/AiChat/ChatInput.test.tsx @@ -415,8 +415,8 @@ describe('ChatInput Component', () => { // After clicking, dropdown should be closed expect(screen.queryByTestId('chat-dropdown')).not.toBeInTheDocument(); - // Variable should be inserted with backticks - expect(textarea).toHaveValue('`df`'); + // Variable should be inserted with backticks and a space + expect(textarea).toHaveValue('`df` '); }); }); @@ -507,8 +507,8 @@ describe('ChatInput Component', () => { // After clicking, dropdown should be closed expect(screen.queryByTestId('chat-dropdown')).not.toBeInTheDocument(); - // Rule should be inserted with backticks - expect(textarea).toHaveValue('Data Analysis'); + // Rule should be inserted with a space after it + expect(textarea).toHaveValue('Data Analysis '); // Wait for the SelectedContextContainer to appear const selectedContextContainers = await screen.findAllByTestId('selected-context-container'); @@ -816,6 +816,9 @@ describe('ChatInput Component', () => { const selectedContextContainer = screen.getByTestId('selected-context-container'); expect(selectedContextContainer).toBeInTheDocument(); expect(within(selectedContextContainer).getByText('Active Cell')).toBeInTheDocument(); + + // Should not show the notebook context container + expect(screen.queryByText('Notebook')).not.toBeInTheDocument(); }); it('does not show active cell context in Agent mode', () => { @@ -823,9 +826,15 @@ describe('ChatInput Component', () => { // Should not show the active cell context container expect(screen.queryByText('Active Cell')).not.toBeInTheDocument(); + + // Should show the active cell context container + const activeCellContainer = screen.getByText('Notebook'); + expect(activeCellContainer).toBeInTheDocument(); - // Should not have any SelectedContextContainer - expect(screen.queryByTestId('selected-context-container')).not.toBeInTheDocument(); + // Should be inside a SelectedContextContainer + const selectedContextContainer = screen.getByTestId('selected-context-container'); + expect(selectedContextContainer).toBeInTheDocument(); + expect(within(selectedContextContainer).getByText('Notebook')).toBeInTheDocument(); }); }); }); \ No newline at end of file diff --git a/mito-ai/src/utils/cellReferences.ts b/mito-ai/src/utils/cellReferences.ts new file mode 100644 index 0000000000..5ca68d4c05 --- /dev/null +++ b/mito-ai/src/utils/cellReferences.ts @@ -0,0 +1,156 @@ +/* + * Copyright (c) Saga Inc. + * Distributed under the terms of the GNU Affero General Public License v3.0 License. + */ + +import { NotebookPanel } from '@jupyterlab/notebook'; +import { getCellIDByIndexInNotebookPanel } from './notebook'; + +/** + * Represents a cell reference in a chat message. + * Cell numbers are 1-indexed (Cell 1, Cell 2, etc.) for user display. + */ +export interface CellReference { + cellNumber: number; // 1-indexed cell number + cellId: string; // The actual cell ID used by JupyterLab + startIndex: number; // Start position in the message text + endIndex: number; // End position in the message text +} + +/** + * Parses cell references from a message text. + * Matches patterns like "@Cell 1", "@Cell 2", "@cell 5", etc. + * + * @param messageText - The message text to parse + * @param notebookPanel - The notebook panel to resolve cell IDs + * @returns Array of cell references found in the message + */ +export function parseCellReferences( + messageText: string, + notebookPanel: NotebookPanel | null +): CellReference[] { + const references: CellReference[] = []; + + if (!notebookPanel) { + return references; + } + + // Match @CellN or @cellN (case insensitive, no space) + const cellReferenceRegex = /@[Cc]ell(\d+)/g; + let match; + + while ((match = cellReferenceRegex.exec(messageText)) !== null) { + if (!match[1]) { + continue; + } + const cellNumber = parseInt(match[1], 10); + const startIndex = match.index; + const endIndex = startIndex + match[0].length; + + // Convert 1-indexed cell number to 0-indexed for lookup + const cellIndex = cellNumber - 1; + const cellId = getCellIDByIndexInNotebookPanel(notebookPanel, cellIndex); + + if (cellId) { + references.push({ + cellNumber, + cellId, + startIndex, + endIndex + }); + } + } + + return references; +} + +/** + * Converts @Cell N references in a message to [MITO_CELL_REF:cell_id] format. + * This ensures cell references are stable even if cells are reordered. + * + * @param messageText - The original message text containing @Cell N references + * @param notebookPanel - The notebook panel to resolve cell IDs + * @returns The message with @Cell N replaced by [MITO_CELL_REF:cell_id] + */ +export function convertCellReferencesToStableFormat( + messageText: string, + notebookPanel: NotebookPanel | null +): string { + if (!notebookPanel) { + return messageText; + } + + // Find all @Cell N references + const references = parseCellReferences(messageText, notebookPanel); + + if (references.length === 0) { + return messageText; + } + + // Replace from end to start to preserve indices + let processedMessage = messageText; + for (let i = references.length - 1; i >= 0; i--) { + const ref = references[i]; + if (!ref) continue; + const before = processedMessage.substring(0, ref.startIndex); + const after = processedMessage.substring(ref.endIndex); + processedMessage = before + `[MITO_CELL_REF:${ref.cellId}]` + after; + } + + return processedMessage; +} + +/** + * Gets the current cell number for a given cell ID. + * + * @param cellId - The cell ID to look up + * @param notebookPanel - The notebook panel + * @returns The 1-indexed cell number, or undefined if not found + */ +export function getCellNumberById( + cellId: string, + notebookPanel: NotebookPanel | null +): number | undefined { + if (!notebookPanel) { + return undefined; + } + + const notebook = notebookPanel.content; + const cellIndex = notebook.widgets.findIndex(cell => cell.model.id === cellId); + + if (cellIndex === -1) { + return undefined; + } + + // Return 1-indexed cell number + return cellIndex + 1; +} + +/** + * Gets all available cells with their numbers and IDs. + * Useful for populating dropdowns or autocomplete. + * + * @param notebookPanel - The notebook panel + * @returns Array of { cellNumber, cellId, cellType } objects + */ +export function getAllCellReferences( + notebookPanel: NotebookPanel | null +): Array<{ cellNumber: number; cellId: string; cellType: string }> { + if (!notebookPanel) { + return []; + } + + const notebook = notebookPanel.content; + const cells: Array<{ cellNumber: number; cellId: string; cellType: string }> = []; + + notebook.widgets.forEach((cell, index) => { + cells.push({ + cellNumber: index + 1, // 1-indexed + cellId: cell.model.id, + cellType: cell.model.type + }); + }); + + return cells; +} + diff --git a/mito-ai/style/CellNumbering.css b/mito-ai/style/CellNumbering.css new file mode 100644 index 0000000000..fff7387524 --- /dev/null +++ b/mito-ai/style/CellNumbering.css @@ -0,0 +1,50 @@ +/* + * Copyright (c) Saga Inc. + * Distributed under the terms of the GNU Affero General Public License v3.0 License. + */ + +/* Add top padding to cells to make room for the cell number label */ +.jp-Cell { + position: relative; + padding-top: 24px !important; +} + +.jp-InputArea-editor, +.jp-MarkdownOutput { + min-height: 50px; + padding-top: 20px !important; + padding-bottom: 20px !important; + background-color: transparent !important; +} + +/* Position the header to hold the label */ +.jp-Cell-header { + display: block; + height: 20px; + position: absolute; + top: 0; + left: 0; + right: 0; +} + +/* Label positioned to overlap the cell's top border (fieldset legend style) */ +.jp-Cell-header::before { + content: "Cell " attr(data-cell-number); + position: absolute; + top: 15px; /* Position so label sits on the cell border at 24px */ + left: calc(var(--jp-cell-collapser-width) + var(--jp-cell-prompt-width) + 10px); + background-color: var(--jp-layout-color0); /* Solid background to mask the border */ + padding: 2px 8px; + font-size: 11px; + font-family: var(--jp-ui-font-family); + color: var(--jp-ui-font-color2); + user-select: none; + z-index: 10; + line-height: 14px; +} + +/* Only show cell labels when the attribute is present */ +.jp-Cell-header:not([data-cell-number])::before { + content: ""; + display: none; +} diff --git a/mito-ai/style/CellReference.css b/mito-ai/style/CellReference.css new file mode 100644 index 0000000000..181f5aa8bc --- /dev/null +++ b/mito-ai/style/CellReference.css @@ -0,0 +1,34 @@ +/* + * Copyright (c) Saga Inc. + * Distributed under the terms of the GNU Affero General Public License v3.0 License. + */ + +/* Cell reference links in chat messages - styled as standard links */ +.cell-reference { + color: var(--purple-600); + text-decoration: underline; + cursor: pointer; + font-weight: 500; + padding: 1px 5px; + border-radius: 3px; + position: relative; + vertical-align: baseline; +} + +.cell-reference:hover { + background-color: var(--purple-300); + color: var(--purple-700); +} + +/* Greyed out style for missing/unresolved cell references (deleted or in different notebook) */ +.cell-reference.cell-reference-missing { + color: var(--jp-ui-font-color2); + text-decoration: none; + cursor: default; + opacity: 0.6; +} + +.cell-reference.cell-reference-missing:hover { + background-color: var(--jp-layout-color3); + color: var(--jp-ui-font-color2); +} diff --git a/mito-ai/style/Citation.css b/mito-ai/style/Citation.css index 0986e0f140..d1ae66d3aa 100644 --- a/mito-ai/style/Citation.css +++ b/mito-ai/style/Citation.css @@ -3,23 +3,34 @@ * Distributed under the terms of the GNU Affero General Public License v3.0 License. */ +/* Citation links - styled as superscript with background for visibility */ .citation-button { - background-color: var(--purple-400); color: var(--purple-700); - height: 14px; - width: 14px; - border-radius: 2px; - font-size: 12px; + background-color: var(--purple-200); + text-decoration: underline; cursor: pointer; - display: inline-flex; - align-items: center; - justify-content: center; + font-size: 11px; + font-weight: 600; + padding: 1px 5px; + border-radius: 3px; position: relative; - top: -0.5em; /* Makes it appear as superscript */ + top: -0.4em; vertical-align: baseline; - margin-right: 2px; + margin: 0 2px; } .citation-button:hover { - background-color: var(--purple-500); -} \ No newline at end of file + background-color: var(--purple-300); +} + +/* Greyed out style for missing/unresolved citations (deleted or in different notebook) */ +.citation-button.citation-missing { + color: var(--jp-ui-font-color2); + text-decoration: none; + cursor: default; + opacity: 0.6; +} + +.citation-button.citation-missing:hover { + background-color: var(--jp-layout-color3); +}