Feature/reference entities (#324)

* Add support for mentions

* Add logs to debug

* Improve styling

* Fix rendering using simple spans

* handle hover colors

* highlight in textarea

* Fix inline renderer

* Fix debounced mention rendering in chat content

* Improve performance for large tables/ models

* Fix spaces in table or model names

* Fix special characters problem

* Show schema and collection names

---------

Co-authored-by: nuwandavek <vivekaithal44@gmail.com>

Author: ppsreejith

web/src/components/common/ChatContent.tsx

@@ -1,11 +1,13 @@
-import React from 'react';
+import React, { useMemo } from 'react';
 import { useSelector } from 'react-redux';
 import { ChatMessageContent } from '../../state/chat/reducer'
 import { Markdown } from './Markdown';
 import { processModelToUIText } from '../../helpers/utils';
 import { getApp } from '../../helpers/app';
 import { RootState } from '../../state/store';
 import { getOrigin, getParsedIframeInfo } from '../../helpers/origin';
+import { createMentionItems, convertMentionsToDisplay } from '../../helpers/mentionUtils';
+import { MetabaseContext } from 'apps/types';
 
 
 const useAppStore = getApp().useStore()
@@ -15,19 +17,30 @@ export const ChatContent: React.FC<{content: ChatMessageContent, messageIndex?:
   messageIndex,
   role
 }) => {
-  const url = useAppStore((state) => state.toolContext)?.url || ''
+  const toolContext: MetabaseContext = useAppStore((state) => state.toolContext)
+  const url = toolContext?.url || ''
   const origin = url ? new URL(url).origin : '';
-  const pageType = useAppStore((state) => state.toolContext)?.pageType || ''
+  const pageType = toolContext?.pageType || ''
   const embedConfigs = useSelector((state: RootState) => state.configs.embed);
 
+  // Create mention items for parsing storage format mentions
+  const mentionItems = useMemo(() => {
+    if (!toolContext?.dbInfo) return []
+    const tables = toolContext.dbInfo.tables || []
+    const models = toolContext.dbInfo.models || []
+    return createMentionItems(tables, models)
+  }, [toolContext?.dbInfo]);
+  
   if (content.type == 'DEFAULT') {
-    const contentText = ((pageType === 'dashboard' || pageType === 'unknown') && role === 'assistant') ? `${content.text} {{MX_LAST_QUERY_URL}}` : content.text;
+    const baseContentText = ((pageType === 'dashboard' || pageType === 'unknown') && role === 'assistant') ? `${content.text} {{MX_LAST_QUERY_URL}}` : content.text;
+    // Convert storage format mentions (@{type:table,id:123}) to special code syntax ([mention:table:table_name])
+    const contentTextWithMentionTags = convertMentionsToDisplay(baseContentText, mentionItems);
     return (
       <div>
         {content.images.map(image => (
           <img src={image.url} key={image.url} />
         ))}
-        <Markdown content={processModelToUIText(contentText, origin, embedConfigs)} messageIndex={messageIndex} />
+        <Markdown content={processModelToUIText(contentTextWithMentionTags, origin, embedConfigs)} messageIndex={messageIndex} />
       </div>
     )
   } else {

web/src/components/common/ChatInputArea.tsx

@@ -5,19 +5,20 @@ import {
   Text,
   Badge
 } from '@chakra-ui/react'
-import React, { forwardRef, useState, useEffect, useCallback } from 'react'
+import React, { forwardRef, useState, useEffect, useCallback, useMemo } from 'react'
 import { useDispatch, useSelector } from 'react-redux'
 import { RootState } from '../../state/store'
 import RunTaskButton from './RunTaskButton'
 import AbortTaskButton from './AbortTaskButton'
-import AutosizeTextarea from './AutosizeTextarea'
+import { MentionTextarea } from './MentionTextarea'
 import { abortPlan } from '../../state/chat/reducer'
 import { setInstructions as setTaskInstructions } from '../../state/thumbnails/reducer'
 import { setConfirmChanges } from '../../state/settings/reducer'
 import { configs } from '../../constants'
 import _ from 'lodash'
 import { MetabaseContext } from 'apps/types'
 import { getApp } from '../../helpers/app'
+import { createMentionItems, convertMentionsToStorage } from '../../helpers/mentionUtils'
 
 interface ChatInputAreaProps {
   isRecording: boolean
@@ -48,6 +49,14 @@ const ChatInputArea = forwardRef<HTMLTextAreaElement, ChatInputAreaProps>(
       dispatch(setConfirmChanges(value))
     }
 
+    // Create mention items from tables and models
+    const mentionItems = useMemo(() => {
+      if (!toolContext.dbInfo) return []
+      const tables = toolContext.dbInfo.tables || []
+      const models = toolContext.dbInfo.models || []
+      return createMentionItems(tables, models)
+    }, [toolContext.dbInfo])
+
     // Sync with redux instructions when they change
     useEffect(() => {
       setInstructions(reduxInstructions)
@@ -67,15 +76,17 @@ const ChatInputArea = forwardRef<HTMLTextAreaElement, ChatInputAreaProps>(
     const onKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
       if (e.key === 'Enter' && !e.shiftKey) {
         e.preventDefault()
-        runTask(instructions)
+        // Convert mentions to storage format before sending
+        const instructionsWithStorageMentions = convertMentionsToStorage(instructions, mentionItems)
+        runTask(instructionsWithStorageMentions)
       }
     }
 
     return (
       <>
         {appEnabledStatus.inputBox && (
           <Stack aria-label="chat-input-area" position={"relative"}>
-            <AutosizeTextarea
+            <MentionTextarea
               ref={ref}
               autoFocus
               aria-label='chat-input'
@@ -84,6 +95,7 @@ const ChatInputArea = forwardRef<HTMLTextAreaElement, ChatInputAreaProps>(
               onChange={(e) => setInstructions(e.target.value)}
               onKeyDown={onKeyDown}
               style={{ width: '100%', height: "100%" }}
+              mentionItems={mentionItems}
             />
             <HStack aria-label="chat-controls" position={"absolute"} bottom={0} width={"100%"} p={2}>
               <HStack justify={"space-between"} width={"100%"}>
@@ -120,7 +132,10 @@ const ChatInputArea = forwardRef<HTMLTextAreaElement, ChatInputAreaProps>(
                   {taskInProgress ? (
                     <AbortTaskButton abortTask={() => dispatch(abortPlan())} disabled={!taskInProgress}/>
                   ) : (
-                    <RunTaskButton runTask={() => runTask(instructions)} disabled={taskInProgress} />
+                    <RunTaskButton runTask={() => {
+                      const instructionsWithStorageMentions = convertMentionsToStorage(instructions, mentionItems)
+                      runTask(instructionsWithStorageMentions)
+                    }} disabled={taskInProgress} />
                   )}
                 </HStack>
               </HStack>

web/src/components/common/Markdown.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect, useState } from 'react';
+import React, { useEffect, useState, useMemo } from 'react';
 import MarkdownComponent from 'react-markdown'
 import remarkGfm from 'remark-gfm'
 import './ChatContent.css'
@@ -10,7 +10,7 @@ import { renderString } from '../../helpers/templatize'
 import { getOrigin, getParsedIframeInfo } from '../../helpers/origin'
 import { getApp } from '../../helpers/app'
 import { getAllTemplateTagsInQuery, replaceLLMFriendlyIdentifiersInSqlWithModels } from 'apps'
-import type { MetabaseModel } from 'apps/types'
+import type { MetabaseModel, MetabaseContext } from 'apps/types'
 import { type EmbedConfigs } from '../../state/configs/reducer'
 import { Badge } from "@chakra-ui/react";
 import { CodeBlock } from './CodeBlock';
@@ -19,6 +19,7 @@ import { BsBarChartFill } from "react-icons/bs";
 import { dispatch } from '../../state/dispatch';
 import { updateIsDevToolsOpen  } from '../../state/settings/reducer';
 import { setMinusxMode } from '../../app/rpc';
+import { createMentionItems, MentionItem } from '../../helpers/mentionUtils';
 
 
 function LinkRenderer(props: any) {
@@ -83,10 +84,60 @@ function ModifiedCode(props: any) {
         const text = props.children?.toString() || '';
 
         if (text.startsWith('[badge]')) {
-        return <Badge color={"minusxGreen.600"} aria-label='mx-badge'>{text.replace('[badge]', '')}</Badge>;
+            return <Badge color={"minusxGreen.600"} aria-label='mx-badge'>{text.replace('[badge]', '')}</Badge>;
         }
         if (text.startsWith('[badge_mx]')) {
-        return <><br></br><Badge aria-label='mx-badge' borderLeftColor={"minusxGreen.600"} borderLeft={"2px solid"} color={"minusxGreen.600"} fontSize={"sm"} mt={2}>{text.replace('[badge_mx]', '')}</Badge><br></br></>;
+            return <><br></br><Badge aria-label='mx-badge' borderLeftColor={"minusxGreen.600"} borderLeft={"2px solid"} color={"minusxGreen.600"} fontSize={"sm"} mt={2}>{text.replace('[badge_mx]', '')}</Badge><br></br></>;
+        }
+        if (text.startsWith('[mention:table:')) {
+          const tableName = text.replace('[mention:table:', '').replace(']', '');
+          const mentionItem = (props as any).mentionItems?.find((item: MentionItem) => 
+            item.name === tableName && item.type === 'table'
+          );
+          
+          let tooltipText = '';
+          if (mentionItem) {
+            tooltipText = `Name: ${mentionItem.originalName}`;
+            if (mentionItem.schema) tooltipText += ` | Schema: ${mentionItem.schema}`;
+          }
+          
+          return (
+            <Tooltip label={tooltipText} placement="top" hasArrow>
+              <span style={{ color: '#3182ce', fontWeight: 500 }}>
+                @{tableName}
+              </span>
+            </Tooltip>
+          );
+        }
+        if (text.startsWith('[mention:model:')) {
+          const modelName = text.replace('[mention:model:', '').replace(']', '');
+          const mentionItem = (props as any).mentionItems?.find((item: MentionItem) => 
+            item.name === modelName && item.type === 'model'
+          );
+          
+          let tooltipText = '';
+          if (mentionItem) {
+            tooltipText = `Name: ${mentionItem.originalName}`;
+            if (mentionItem.collection) tooltipText += ` | Collection: ${mentionItem.collection}`;
+          }
+          
+          return (
+            <Tooltip label={tooltipText} placement="top" hasArrow>
+              <span style={{ color: '#805ad5', fontWeight: 500 }}>
+                @{modelName}
+              </span>
+            </Tooltip>
+          );
+        }
+        if (text.startsWith('[mention:missing:')) {
+          const parts = text.replace('[mention:missing:', '').replace(']', '').split(':');
+          const type = parts[0];
+          const id = parts[1];
+          return (
+            <span style={{ color: '#718096' }}>
+              @[{type}:{id}]
+            </span>
+          );
         }
     }
 
@@ -551,6 +602,90 @@ function extractLastQueryFromMessages(messages: any[], currentMessageIndex: numb
 
 const useAppStore = getApp().useStore();
 
+// Component to detect @ mentions in text and render with colors
+function MentionAwareText({ children }: { children: React.ReactNode }) {
+  const toolContext: MetabaseContext = useAppStore((state) => state.toolContext)
+  
+  // Create mention items for determining colors
+  const mentionItems = useMemo(() => {
+    if (!toolContext?.dbInfo) return []
+    const tables = toolContext.dbInfo.tables || []
+    const models = toolContext.dbInfo.models || []
+    return createMentionItems(tables, models)
+  }, [toolContext?.dbInfo])
+
+  // Create lookup map for determining mention types
+  const mentionMap = useMemo(() => {
+    const map = new Map<string, 'table' | 'model'>()
+    mentionItems.forEach((item: MentionItem) => {
+      map.set(item.name.toLowerCase(), item.type)
+    })
+    return map
+  }, [mentionItems])
+
+  // Only process string children
+  if (typeof children !== 'string') {
+    return <>{children}</>
+  }
+
+  const renderTextWithColors = (text: string) => {
+    // Regex to find @ mentions (@ followed by word characters)
+    const mentionRegex = /@(\w+)/g
+    const parts: (string | JSX.Element)[] = []
+    let lastIndex = 0
+    let match
+
+    while ((match = mentionRegex.exec(text)) !== null) {
+      // Add text before the mention
+      if (match.index > lastIndex) {
+        parts.push(text.slice(lastIndex, match.index))
+      }
+
+      const mentionName = match[1]
+      const mentionType = mentionMap.get(mentionName.toLowerCase())
+      
+      if (mentionType) {
+        // Render as colored mention
+        const color = mentionType === 'table' ? '#3182ce' : '#805ad5' // Blue for tables, purple for models
+        parts.push(
+          <span
+            key={`mention-${match.index}`}
+            style={{ color, fontWeight: 500 }}
+          >
+            @{mentionName}
+          </span>
+        )
+      } else {
+        // Check if it looks like a missing reference pattern @[type:id]
+        if (text.slice(match.index).match(/^@\[[^:]+:\w+\]/)) {
+          parts.push(
+            <span
+              key={`mention-missing-${match.index}`}
+              style={{ color: '#718096' }}
+            >
+              {match[0]}
+            </span>
+          )
+        } else {
+          // Regular @ mention, not in our database
+          parts.push(match[0])
+        }
+      }
+
+      lastIndex = match.index + match[0].length
+    }
+
+    // Add remaining text
+    if (lastIndex < text.length) {
+      parts.push(text.slice(lastIndex))
+    }
+
+    return parts.length > 1 ? <>{parts}</> : text
+  }
+
+  return <>{renderTextWithColors(children)}</>
+}
+
 export function Markdown({content, messageIndex}: {content: string, messageIndex?: number}) {
   const currentThread = useSelector((state: RootState) => 
     state.chat.threads[state.chat.activeThread]
@@ -566,6 +701,14 @@ export function Markdown({content, messageIndex}: {content: string, messageIndex
   const mxModels = useSelector((state: RootState) => state.cache.mxModels);
 
   // Process template variables like {{MX_LAST_QUERY_URL}}
+  // Create mention items from toolContext
+  const mentionItems = useMemo(() => {
+    if (!toolContext?.dbInfo) return []
+    const tables = toolContext.dbInfo.tables || []
+    const models = toolContext.dbInfo.models || []
+    return createMentionItems(tables, models)
+  }, [toolContext?.dbInfo])
+
   const processedContent = React.useMemo(() => {
     if (content.includes('{{MX_LAST_QUERY_URL}}')) {
       try {
@@ -594,7 +737,38 @@ export function Markdown({content, messageIndex}: {content: string, messageIndex
     return content;
   }, [content, currentThread?.messages, toolContext?.dbId, messageIndex, settings, mxModels]);
 
+  // Create a wrapped ModifiedCode component that has access to mentionItems
+  const ModifiedCodeWithMentions = (props: any) => (
+    <ModifiedCode {...props} mentionItems={mentionItems} />
+  )
+
   return (
-    <MarkdownComponent remarkPlugins={[remarkGfm]} className={"markdown"} components={{ a: LinkRenderer, p: ModifiedParagraph, ul: ModifiedUL, ol: ModifiedOL, img: ImageComponent, pre: ModifiedPre, blockquote: ModifiedBlockquote, code: ModifiedCode, hr: HorizontalLine, h1: ModifiedH1, h2: ModifiedH2, h3: ModifiedH3, h4: ModifiedH4, table: ModifiedTable, thead: ModifiedThead, tbody: ModifiedTbody, tr: ModifiedTr, th: ModifiedTh, td: ModifiedTd }}>{processedContent}</MarkdownComponent>
+    <MarkdownComponent 
+      remarkPlugins={[remarkGfm]} 
+      className={"markdown"} 
+      components={{ 
+        a: LinkRenderer, 
+        p: ModifiedParagraph, 
+        ul: ModifiedUL, 
+        ol: ModifiedOL, 
+        img: ImageComponent, 
+        pre: ModifiedPre, 
+        blockquote: ModifiedBlockquote, 
+        code: ModifiedCodeWithMentions, 
+        hr: HorizontalLine, 
+        h1: ModifiedH1, 
+        h2: ModifiedH2, 
+        h3: ModifiedH3, 
+        h4: ModifiedH4, 
+        table: ModifiedTable, 
+        thead: ModifiedThead, 
+        tbody: ModifiedTbody, 
+        tr: ModifiedTr, 
+        th: ModifiedTh, 
+        td: ModifiedTd,
+      }}
+    >
+      {processedContent}
+    </MarkdownComponent>
   )
 }