Merge pull request #1364 from mito-ds/spike-at-mention-vars

Implemented @ Mention Variables

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

@@ -0,0 +1,118 @@
+import React, { useState, useEffect } from 'react';
+import { classNames } from '../../../utils/classNames';
+import { ExpandedVariable } from './ChatInput';
+
+interface ChatDropdownProps {
+    options: ExpandedVariable[];
+    onSelect: (variableName: string, parentDf?: string) => void;
+    filterText: string;
+    maxDropdownItems?: number;
+    position?: 'above' | 'below';
+}
+
+const ChatDropdown: React.FC<ChatDropdownProps> = ({
+    options,
+    onSelect,
+    filterText,
+    maxDropdownItems = 10,
+    position = 'below'
+}) => {
+    const [selectedIndex, setSelectedIndex] = useState(0);
+
+    const filteredOptions = options.filter((variable) =>
+        variable.variable_name.toLowerCase().includes(filterText.toLowerCase()) &&
+        variable.type !== "<class 'module'>"
+    ).slice(0, maxDropdownItems);
+
+    useEffect(() => {
+        setSelectedIndex(0);
+    }, [options, filterText]);
+
+    const handleKeyDown = (event: KeyboardEvent) => {
+        switch (event.key) {
+            case 'ArrowDown' || 'Down':
+                event.preventDefault();
+                setSelectedIndex((prev) =>
+                    prev < filteredOptions.length - 1 ? prev + 1 : 0
+                );
+                break;
+            case 'ArrowUp' || 'Up':
+                event.preventDefault();
+                setSelectedIndex((prev) =>
+                    prev > 0 ? prev - 1 : filteredOptions.length - 1
+                );
+                break;
+            case 'Enter' || 'Return':
+                event.preventDefault();
+                if (filteredOptions[selectedIndex]) {
+                    onSelect(filteredOptions[selectedIndex].variable_name, filteredOptions[selectedIndex].parent_df);
+                }
+                break;
+        }
+    };
+
+    useEffect(() => {
+        document.addEventListener('keydown', handleKeyDown);
+        return () => document.removeEventListener('keydown', handleKeyDown);
+    }, [filteredOptions, selectedIndex]);
+
+    const getShortType = (type: string) => {
+        return type.includes("DataFrame") ? "df"
+            : type.includes("<class '") ? type.split("'")[1]
+                : type;
+    }
+
+    return (
+        <div className={`chat-dropdown ${position}`}>
+            <ul className="chat-dropdown-list" >
+                {filteredOptions.length === 0 && (
+                    <li className="chat-dropdown-item">No variables found</li>
+                )}
+
+                {filteredOptions.map((option, index) => {
+                    const uniqueKey = option.parent_df
+                        ? `${option.parent_df}.${option.variable_name}`
+                        : option.variable_name;
+
+                    return (
+                        <li
+                            key={uniqueKey}
+                            className={classNames("chat-dropdown-item", { selected: index === selectedIndex })}
+                            onClick={() => onSelect(option.variable_name, option.parent_df)}
+                        >
+                            <span className="chat-dropdown-item-type"
+                                style={{
+                                    color: getShortType(option.type) === 'df' ? 'blue'
+                                        : getShortType(option.type) === 'col' ? 'orange'
+                                            : "green"
+                                }}
+                                title={getShortType(option.type)}
+                            >
+                                {getShortType(option.type)}
+                            </span>
+                            <span
+                                className="chat-dropdown-item-name"
+                                title={option.variable_name}
+                                ref={(el) => {
+                                    // Show full text on hover if the text is too long
+                                    if (el) {
+                                        el.title = el.scrollWidth > el.clientWidth ? option.variable_name : '';
+                                    }
+                                }}
+                            >
+                                {option.variable_name}
+                            </span>
+                            {option.parent_df && (
+                                <span className="chat-dropdown-item-parent-df">
+                                    {option.parent_df}
+                                </span>
+                            )}
+                        </li>
+                    );
+                })}
+            </ul>
+        </div>
+    );
+};
+
+export default ChatDropdown; 

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

@@ -1,22 +1,35 @@
-import React, { useEffect } from 'react';
+import React, { useState, useEffect } from 'react';
 import { classNames } from '../../../utils/classNames';
+import { IVariableManager } from '../../VariableManager/VariableManagerPlugin';
+import ChatDropdown from './ChatDropdown';
+import { Variable } from '../../VariableManager/VariableInspector';
 
 interface ChatInputProps {
     initialContent: string;
     placeholder: string;
     onSave: (content: string) => void;
-    onCancel?: () => void ;
+    onCancel?: () => void;
     isEditing: boolean;
+    variableManager?: IVariableManager;
+}
+
+export interface ExpandedVariable extends Variable {
+    parent_df?: string;
 }
 
 const ChatInput: React.FC<ChatInputProps> = ({
     initialContent,
     placeholder,
     onSave,
     onCancel,
-    isEditing
+    isEditing,
+    variableManager
 }) => {
-    const [input, setInput] = React.useState(initialContent);
+    const [input, setInput] = useState(initialContent);
+    const [expandedVariables, setExpandedVariables] = useState<ExpandedVariable[]>([]);
+    const [isDropdownVisible, setDropdownVisible] = useState(false);
+    const [dropdownFilter, setDropdownFilter] = useState('');
+    const [showDropdownAbove, setShowDropdownAbove] = useState(false);
     const textAreaRef = React.useRef<HTMLTextAreaElement>(null);
 
     // TextAreas cannot automatically adjust their height based on the content that they contain, 
@@ -33,19 +46,122 @@ const ChatInput: React.FC<ChatInputProps> = ({
         textarea.style.height = `${minHeight}px`;
     };
 
+    const handleInputChange = (event: React.ChangeEvent<HTMLTextAreaElement>) => {
+        const value = event.target.value;
+        setInput(value);
+
+        const cursorPosition = event.target.selectionStart;
+        const textBeforeCursor = value.slice(0, cursorPosition);
+        const words = textBeforeCursor.split(/\s+/);
+        const currentWord = words[words.length - 1];
+
+        if (currentWord.startsWith("@")) {
+            const query = currentWord.slice(1);
+            setDropdownFilter(query);
+            setDropdownVisible(true);
+        } else {
+            setDropdownVisible(false);
+            setDropdownFilter('');
+        }
+    };
+
+    const handleOptionSelect = (variableName: string, parentDf?: string) => {
+        const textarea = textAreaRef.current;
+        if (!textarea) return;
+
+        const cursorPosition = textarea.selectionStart;
+        const textBeforeCursor = input.slice(0, cursorPosition);
+        const atIndex = textBeforeCursor.lastIndexOf("@");
+        const textAfterCursor = input.slice(cursorPosition);
+
+        let variableNameWithBackticks: string;
+        if (!parentDf) {
+            variableNameWithBackticks = `\`${variableName}\``
+        } else {
+            variableNameWithBackticks = `\`${parentDf}['${variableName}']\``
+        }
+
+        const newValue =
+            input.slice(0, atIndex) +
+            variableNameWithBackticks +
+            textAfterCursor;
+        setInput(newValue);
+
+        setDropdownVisible(false);
+
+        // After updating the input value, set the cursor position after the inserted variable name
+        // We use setTimeout to ensure this happens after React's state update
+        setTimeout(() => {
+            if (textarea) {
+                const newCursorPosition = atIndex + variableNameWithBackticks.length;
+                textarea.focus();
+                textarea.setSelectionRange(newCursorPosition, newCursorPosition);
+            }
+        }, 0);
+    };
+
     useEffect(() => {
         adjustHeight();
     }, [textAreaRef?.current?.value]);
 
+    // Update the expandedVariables arr when the variable manager changes
+    useEffect(() => {
+        const expandedVariables: ExpandedVariable[] = [
+            // Add base variables (excluding DataFrames)
+            ...(variableManager?.variables.filter(variable => variable.type !== "pd.DataFrame") || []),
+            // Add DataFrames
+            ...(variableManager?.variables.filter((variable) => variable.type === "pd.DataFrame") || []),
+            // Add series with parent DataFrame references
+            ...(variableManager?.variables
+                .filter((variable) => variable.type === "pd.DataFrame")
+                .flatMap((df) =>
+                    Object.entries(df.value).map(([seriesName, details]) => ({
+                        variable_name: seriesName,
+                        type: "col",
+                        value: "replace_me",
+                        parent_df: df.variable_name,
+                    }))
+                ) || [])
+        ];
+        setExpandedVariables(expandedVariables);
+    }, [variableManager?.variables]);
+
+    const calculateDropdownPosition = () => {
+        if (!textAreaRef.current) return;
+
+        const textarea = textAreaRef.current;
+        const textareaRect = textarea.getBoundingClientRect();
+        const windowHeight = window.innerHeight;
+        const spaceBelow = windowHeight - textareaRect.bottom;
+
+        // If space below is less than 200px (typical dropdown height), show above
+        setShowDropdownAbove(spaceBelow < 200);
+    };
+
+    useEffect(() => {
+        if (isDropdownVisible) {
+            calculateDropdownPosition();
+        }
+    }, [isDropdownVisible]);
+
     return (
-        <>
+        <div style={{ position: 'relative' }}>
             <textarea
                 ref={textAreaRef}
                 className={classNames("message", "message-user", 'chat-input')}
                 placeholder={placeholder}
                 value={input}
-                onChange={(e) => { setInput(e.target.value) }}
+                onChange={handleInputChange}
                 onKeyDown={(e) => {
+                    // If dropdown is visible, only handle escape to close it
+                    if (isDropdownVisible) {
+                        if (e.key === 'Escape') {
+                            e.preventDefault();
+                            setDropdownVisible(false);
+                        }
+                        return;
+                    }
+
                     // Enter key sends the message, but we still want to allow 
                     // shift + enter to add a new line.
                     if (e.key === 'Enter' && !e.shiftKey) {
@@ -62,13 +178,21 @@ const ChatInput: React.FC<ChatInputProps> = ({
                     }
                 }}
             />
-            {isEditing && 
+            {isEditing &&
                 <div className="message-edit-buttons">
                     <button onClick={() => onSave(input)}>Save</button>
                     <button onClick={onCancel}>Cancel</button>
                 </div>
             }
-        </>
+            {isDropdownVisible && (
+                <ChatDropdown
+                    options={expandedVariables}
+                    onSelect={handleOptionSelect}
+                    filterText={dropdownFilter}
+                    position={showDropdownAbove ? "above" : "below"}
+                />
+            )}
+        </div>
     )
 };
 

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

@@ -11,6 +11,7 @@ import { OperatingSystem } from '../../../utils/user';
 import { UnifiedDiffLine } from '../../../utils/codeDiff';
 import PencilIcon from '../../../icons/Pencil';
 import ChatInput from './ChatInput';
+import { IVariableManager } from '../../VariableManager/VariableManagerPlugin';
 
 interface IChatMessageProps {
     message: OpenAI.Chat.ChatCompletionMessageParam
@@ -25,6 +26,7 @@ interface IChatMessageProps {
     acceptAICode: () => void
     rejectAICode: () => void
     onUpdateMessage: (messageIndex: number, newContent: string) => void
+    variableManager?: IVariableManager
 }
 
 const ChatMessage: React.FC<IChatMessageProps> = ({
@@ -40,6 +42,7 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
     acceptAICode,
     rejectAICode,
     onUpdateMessage,
+    variableManager
 }): JSX.Element | null => {
     const [isEditing, setIsEditing] = useState(false);
 
@@ -74,6 +77,7 @@ const ChatMessage: React.FC<IChatMessageProps> = ({
                     onSave={handleSave}
                     onCancel={handleCancel}
                     isEditing={isEditing}
+                    variableManager={variableManager}
                 />
             </div>
         );

mito-ai/src/Extensions/AiChat/ChatTaskpane.tsx

@@ -379,6 +379,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
                             acceptAICode={acceptAICode}
                             rejectAICode={rejectAICode}
                             onUpdateMessage={handleUpdateMessage}
+                            variableManager={variableManager}
                         />
                     )
                 }).filter(message => message !== null)}
@@ -394,6 +395,7 @@ const ChatTaskpane: React.FC<IChatTaskpaneProps> = ({
                 onSave={sendChatInputMessage}
                 onCancel={undefined}
                 isEditing={false}
+                variableManager={variableManager}
             />
         </div>
     );