feat(UI): Enhance model statistics tracking and citation handling in chat components

.github/copilot-instructions.md

@@ -150,3 +150,4 @@ When working on this codebase, you might need to:
   - `/langchain-ai/langchainjs` for LangChain
   - `/langchain-ai/langgraph` for LangGraph
   - `/quantizor/markdown-to-jsx` for Markdown to JSX conversion
+  - `/context7/headlessui_com` for Headless UI components

src/app/api/chat/route.ts

@@ -56,6 +56,11 @@ type Body = {
 type ModelStats = {
   modelName: string;
   responseTime?: number;
+  usage?: {
+    input_tokens: number;
+    output_tokens: number;
+    total_tokens: number;
+  };
 };
 
 const handleEmitterEvents = async (

src/components/ChatWindow.tsx

@@ -16,6 +16,11 @@ import NextError from 'next/error';
 export type ModelStats = {
   modelName: string;
   responseTime?: number;
+  usage?: {
+    input_tokens: number;
+    output_tokens: number;
+    total_tokens: number;
+  };
 };
 
 export type AgentActionEvent = {

src/components/CitationLink.tsx

@@ -0,0 +1,55 @@
+import { Document } from '@langchain/core/documents';
+import { useState } from 'react';
+import MessageSource from './MessageSource';
+
+interface CitationLinkProps {
+  number: string;
+  source?: Document;
+  url?: string;
+}
+
+const CitationLink = ({ number, source, url }: CitationLinkProps) => {
+  const [showTooltip, setShowTooltip] = useState(false);
+  const linkContent = (
+    <a
+      href={url}
+      target="_blank"
+      rel="noopener noreferrer"
+      className="bg-light-secondary dark:bg-dark-secondary px-1 rounded ml-1 no-underline text-xs text-black/70 dark:text-white/70 relative hover:bg-light-200 dark:hover:bg-dark-200 transition-colors duration-200"
+    >
+      {number}
+    </a>
+  );
+
+  // If we have source data, wrap with tooltip
+  if (source) {
+    return (
+      <div className="relative inline-block">
+        <div
+          onMouseEnter={() => setShowTooltip(true)}
+          onMouseLeave={() => setShowTooltip(false)}
+        >
+          {linkContent}
+        </div>
+
+        {showTooltip && (
+          <div className="absolute z-50 bottom-full mb-2 left-1/2 transform -translate-x-1/2 animate-in fade-in-0 duration-150">
+            <div className="bg-light-primary dark:bg-dark-primary border rounded-lg border-light-200 dark:border-dark-200 shadow-lg w-96">
+              <MessageSource
+                source={source}
+                className="shadow-none border-none bg-transparent hover:bg-transparent dark:hover:bg-transparent cursor-pointer"
+              />
+            </div>
+            {/* Tooltip arrow */}
+            <div className="absolute top-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-l-4 border-r-4 border-t-4 border-transparent border-t-light-200 dark:border-t-dark-200"></div>
+          </div>
+        )}
+      </div>
+    );
+  }
+
+  // Otherwise, just return the plain link
+  return linkContent;
+};
+
+export default CitationLink;

src/components/MarkdownRenderer.tsx

@@ -2,7 +2,14 @@
 'use client';
 
 import { cn } from '@/lib/utils';
-import { CheckCheck, Copy as CopyIcon, Search, FileText, Globe, Settings } from 'lucide-react';
+import {
+  CheckCheck,
+  Copy as CopyIcon,
+  Search,
+  FileText,
+  Globe,
+  Settings,
+} from 'lucide-react';
 import Markdown, { MarkdownToJSX } from 'markdown-to-jsx';
 import { useState } from 'react';
 import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
@@ -12,6 +19,8 @@ import {
 } from 'react-syntax-highlighter/dist/cjs/styles/prism';
 import { useTheme } from 'next-themes';
 import ThinkBox from './ThinkBox';
+import { Document } from '@langchain/core/documents';
+import CitationLink from './CitationLink';
 
 // Helper functions for think overlay
 const extractThinkContent = (content: string): string | null => {
@@ -51,6 +60,7 @@ interface MarkdownRendererProps {
     thinkBoxId: string,
     expanded: boolean,
   ) => void;
+  sources?: Document[];
 }
 
 // Custom ToolCall component for markdown
@@ -162,8 +172,8 @@ const ThinkTagProcessor = ({
   onToggle?: (thinkBoxId: string, expanded: boolean) => void;
 }) => {
   return (
-    <ThinkBox 
+    <ThinkBox
-      content={children} 
+      content={children}
       expanded={isExpanded}
       onToggle={() => {
         if (id && onToggle) {
@@ -172,7 +182,8 @@ const ThinkTagProcessor = ({
       }}
     />
   );
-};const CodeBlock = ({
+};
+const CodeBlock = ({
   className,
   children,
 }: {
@@ -248,6 +259,7 @@ const MarkdownRenderer = ({
   messageId,
   expandedThinkBoxes,
   onThinkBoxToggle,
+  sources,
 }: MarkdownRendererProps) => {
   // Preprocess content to add stable IDs to think tags
   const processedContent = addThinkBoxIds(content);
@@ -265,7 +277,7 @@ const MarkdownRenderer = ({
   };
 
   // Determine what content to render based on showThinking parameter
-  const contentToRender = showThinking 
+  const contentToRender = showThinking
     ? processedContent
     : removeThinkTags(processedContent);
   // Markdown formatting options
@@ -317,9 +329,28 @@ const MarkdownRenderer = ({
         component: ({ children }) => children,
       },
       a: {
-        component: (props) => (
+        component: (props) => {
-          <a {...props} target="_blank" rel="noopener noreferrer" />
+          // Check if this is a citation link with data-citation attribute
-        ),
+          const citationNumber = props['data-citation'];
+
+          if (sources && citationNumber) {
+            const number = parseInt(citationNumber);
+            const source = sources[number - 1];
+
+            if (source) {
+              return (
+                <CitationLink
+                  number={number.toString()}
+                  source={source}
+                  url={props.href}
+                />
+              );
+            }
+          }
+
+          // Default link behavior
+          return <a {...props} target="_blank" rel="noopener noreferrer" />;
+        },
       },
       // Prevent rendering of certain HTML elements for security
       iframe: () => null, // Don't render iframes

src/components/MessageActions/ModelInfo.tsx

@@ -48,21 +48,21 @@ const ModelInfoButton: React.FC<ModelInfoButtonProps> = ({ modelStats }) => {
       {showPopover && (
         <div
           ref={popoverRef}
-          className="absolute z-10 left-6 top-0 w-64 rounded-md shadow-lg bg-white dark:bg-dark-secondary border border-light-200 dark:border-dark-200"
+          className="absolute z-10 left-6 top-0 w-72 rounded-md shadow-lg bg-white dark:bg-dark-secondary border border-light-200 dark:border-dark-200"
         >
           <div className="py-2 px-3">
             <h4 className="text-sm font-medium mb-2 text-black dark:text-white">
               Model Information
             </h4>
             <div className="space-y-1 text-xs">
-              <div className="flex justify-between">
+              <div className="flex space-x-2">
                 <span className="text-black/70 dark:text-white/70">Model:</span>
                 <span className="text-black dark:text-white font-medium">
                   {modelName}
                 </span>
               </div>
               {modelStats?.responseTime && (
-                <div className="flex justify-between">
+                <div className="flex space-x-2">
                   <span className="text-black/70 dark:text-white/70">
                     Response time:
                   </span>
@@ -71,6 +71,34 @@ const ModelInfoButton: React.FC<ModelInfoButtonProps> = ({ modelStats }) => {
                   </span>
                 </div>
               )}
+              {modelStats?.usage && (
+                <>
+                  <div className="flex space-x-2">
+                    <span className="text-black/70 dark:text-white/70">
+                      Input tokens:
+                    </span>
+                    <span className="text-black dark:text-white font-medium">
+                      {modelStats.usage.input_tokens.toLocaleString()}
+                    </span>
+                  </div>
+                  <div className="flex space-x-2">
+                    <span className="text-black/70 dark:text-white/70">
+                      Output tokens:
+                    </span>
+                    <span className="text-black dark:text-white font-medium">
+                      {modelStats.usage.output_tokens.toLocaleString()}
+                    </span>
+                  </div>
+                  <div className="flex space-x-2">
+                    <span className="text-black/70 dark:text-white/70">
+                      Total tokens:
+                    </span>
+                    <span className="text-black dark:text-white font-medium">
+                      {modelStats.usage.total_tokens.toLocaleString()}
+                    </span>
+                  </div>
+                </>
+              )}
             </div>
           </div>
         </div>

src/components/MessageSource.tsx

@@ -0,0 +1,111 @@
+/* eslint-disable @next/next/no-img-element */
+import { Document } from '@langchain/core/documents';
+import { File, Zap, Microscope, FileText, Sparkles } from 'lucide-react';
+
+interface MessageSourceProps {
+  source: Document;
+  index?: number;
+  style?: React.CSSProperties;
+  className?: string;
+}
+
+const MessageSource = ({
+  source,
+  index,
+  style,
+  className,
+}: MessageSourceProps) => {
+  return (
+    <a
+      className={`bg-light-100 hover:bg-light-200 dark:bg-dark-100 dark:hover:bg-dark-200 transition duration-200 rounded-lg p-4 flex flex-row no-underline space-x-3 font-medium ${className || ''}`}
+      href={source.metadata.url}
+      target="_blank"
+      style={style}
+    >
+      {/* Left side: Favicon/Icon and source number */}
+      <div className="flex flex-col items-center space-y-2 flex-shrink-0">
+        {source.metadata.url === 'File' ? (
+          <div className="bg-dark-200 hover:bg-dark-100 transition duration-200 flex items-center justify-center w-8 h-8 rounded-full">
+            <File size={16} className="text-white/70" />
+          </div>
+        ) : (
+          <img
+            src={`https://s2.googleusercontent.com/s2/favicons?domain_url=${source.metadata.url}`}
+            width={28}
+            height={28}
+            alt="favicon"
+            className="rounded-lg h-7 w-7"
+          />
+        )}
+        <div className="flex flex-row items-center space-x-1 text-black/50 dark:text-white/50 text-xs">
+          {typeof index === 'number' && (
+            <span className="font-semibold">{index + 1}</span>
+          )}
+          {/* Processing type indicator */}
+          {source.metadata.processingType === 'preview-only' && (
+            <span title="Partial content analyzed" className="inline-flex">
+              <Zap size={12} className="text-black/40 dark:text-white/40" />
+            </span>
+          )}
+          {source.metadata.processingType === 'full-content' && (
+            <span title="Full content analyzed" className="inline-flex">
+              <Microscope
+                size={12}
+                className="text-black/40 dark:text-white/40"
+              />
+            </span>
+          )}
+          {source.metadata.processingType === 'url-direct-content' && (
+            <span title="Direct URL content" className="inline-flex">
+              <FileText
+                size={12}
+                className="text-black/40 dark:text-white/40"
+              />
+            </span>
+          )}
+          {source.metadata.processingType === 'url-content-extraction' && (
+            <span title="Summarized URL content" className="inline-flex">
+              <Sparkles
+                size={12}
+                className="text-black/40 dark:text-white/40"
+              />
+            </span>
+          )}
+        </div>
+      </div>
+
+      {/* Right side: Content */}
+      <div className="flex-1 flex flex-col space-y-2">
+        {/* Title */}
+        <h3 className="dark:text-white text-sm font-semibold leading-tight">
+          {source.metadata.title}
+        </h3>
+
+        {/* URL */}
+        <p className="text-xs text-black/50 dark:text-white/50">
+          {source.metadata.url.replace(/.+\/\/|www.|\..+/g, '')}
+        </p>
+
+        {/* Preview content */}
+        <p
+          className="text-xs text-black/70 dark:text-white/70 leading-relaxed overflow-hidden"
+          style={{
+            display: '-webkit-box',
+            WebkitLineClamp: 3,
+            WebkitBoxOrient: 'vertical',
+          }}
+        >
+          {/* Use snippet for preview-only content, otherwise use pageContent */}
+          {source.metadata.processingType === 'preview-only' &&
+          source.metadata.snippet
+            ? source.metadata.snippet
+            : source.pageContent?.length > 250
+              ? source.pageContent.slice(0, 250) + '...'
+              : source.pageContent || 'No preview available'}
+        </p>
+      </div>
+    </a>
+  );
+};
+
+export default MessageSource;

src/components/MessageSources.tsx

@@ -1,94 +1,12 @@
 /* eslint-disable @next/next/no-img-element */
 import { Document } from '@langchain/core/documents';
-import { File, Zap, Microscope, FileText, Sparkles } from 'lucide-react';
+import MessageSource from './MessageSource';
 
 const MessageSources = ({ sources }: { sources: Document[] }) => {
   return (
     <div className="flex flex-col space-y-3">
       {sources.map((source, i) => (
-        <a
+        <MessageSource key={i} source={source} index={i} />
-          className="bg-light-100 hover:bg-light-200 dark:bg-dark-100 dark:hover:bg-dark-200 transition duration-200 rounded-lg p-4 flex flex-row space-x-3 font-medium"
-          key={i}
-          href={source.metadata.url}
-          target="_blank"
-        >
-          {/* Left side: Favicon/Icon and source number */}
-          <div className="flex flex-col items-center space-y-2 flex-shrink-0">
-            {source.metadata.url === 'File' ? (
-              <div className="bg-dark-200 hover:bg-dark-100 transition duration-200 flex items-center justify-center w-8 h-8 rounded-full">
-                <File size={16} className="text-white/70" />
-              </div>
-            ) : (
-              <img
-                src={`https://s2.googleusercontent.com/s2/favicons?domain_url=${source.metadata.url}`}
-                width={28}
-                height={28}
-                alt="favicon"
-                className="rounded-lg h-7 w-7"
-              />
-            )}
-            <div className="flex flex-row items-center space-x-1 text-black/50 dark:text-white/50 text-xs">
-              <span className="font-semibold">{i + 1}</span>
-              {/* Processing type indicator */}
-              {source.metadata.processingType === 'preview-only' && (
-                <span title="Partial content analyzed" className="inline-flex">
-                  <Zap
-                    size={12}
-                    className="text-black/40 dark:text-white/40"
-                  />
-                </span>
-              )}
-              {source.metadata.processingType === 'full-content' && (
-                <span title="Full content analyzed" className="inline-flex">
-                  <Microscope
-                    size={12}
-                    className="text-black/40 dark:text-white/40"
-                  />
-                </span>
-              )}
-              {source.metadata.processingType === 'url-direct-content' && (
-                <span title="Direct URL content" className="inline-flex">
-                  <FileText
-                    size={12}
-                    className="text-black/40 dark:text-white/40"
-                  />
-                </span>
-              )}
-              {source.metadata.processingType === 'url-content-extraction' && (
-                <span title="Summarized URL content" className="inline-flex">
-                  <Sparkles
-                    size={12}
-                    className="text-black/40 dark:text-white/40"
-                  />
-                </span>
-              )}
-            </div>
-          </div>
-
-          {/* Right side: Content */}
-          <div className="flex-1 flex flex-col space-y-2">
-            {/* Title */}
-            <h3 className="dark:text-white text-sm font-semibold leading-tight">
-              {source.metadata.title}
-            </h3>
-            
-            {/* URL */}
-            <p className="text-xs text-black/50 dark:text-white/50">
-              {source.metadata.url.replace(/.+\/\/|www.|\..+/g, '')}
-            </p>
-
-            {/* Preview content */}
-            <p className="text-xs text-black/70 dark:text-white/70 leading-relaxed overflow-hidden" style={{ display: '-webkit-box', WebkitLineClamp: 3, WebkitBoxOrient: 'vertical' }}>
-              {/* Use snippet for preview-only content, otherwise use pageContent */}
-              {source.metadata.processingType === 'preview-only' && source.metadata.snippet
-                ? source.metadata.snippet
-                : source.pageContent?.length > 250 
-                  ? source.pageContent.slice(0, 250) + '...'
-                  : source.pageContent || 'No preview available'
-              }
-            </p>
-          </div>
-        </a>
       ))}
     </div>
   );

src/components/MessageTabs.tsx

@@ -138,7 +138,7 @@ const MessageTabs = ({
                 const url = source?.metadata?.url;
 
                 if (url) {
-                  return `<a href="${url}" target="_blank" className="bg-light-secondary dark:bg-dark-secondary px-1 rounded ml-1 no-underline text-xs text-black/70 dark:text-white/70 relative">${numStr}</a>`;
+                  return `<a href="${url}" target="_blank" data-citation="${number}" className="bg-light-secondary dark:bg-dark-secondary px-1 rounded ml-1 no-underline text-xs text-black/70 dark:text-white/70 relative hover:bg-light-200 dark:hover:bg-dark-200 transition-colors duration-200">${numStr}</a>`;
                 } else {
                   return `[${numStr}]`;
                 }
@@ -279,14 +279,16 @@ const MessageTabs = ({
         {/* Answer Tab */}
         {activeTab === 'text' && (
           <div className="flex flex-col space-y-4 animate-fadeIn">
-            <MarkdownRenderer 
+            <MarkdownRenderer
-              content={parsedMessage} 
+              content={parsedMessage}
-              className="px-4" 
+              className="px-4"
               messageId={message.messageId}
               expandedThinkBoxes={message.expandedThinkBoxes}
               onThinkBoxToggle={onThinkBoxToggle}
               showThinking={true}
-            />            {loading && isLast ? null : (
+              sources={message.sources}
+            />{' '}
+            {loading && isLast ? null : (
               <div className="flex flex-row items-center justify-between w-full text-black dark:text-white px-4 py-4">
                 <div className="flex flex-row items-center space-x-1">
                   <Rewrite rewrite={rewrite} messageId={message.messageId} />
@@ -315,7 +317,6 @@ const MessageTabs = ({
                 </div>
               </div>
             )}
-
             {isLast && message.role === 'assistant' && !loading && (
               <>
                 <div className="border-t border-light-secondary dark:border-dark-secondary px-4 pt-4 mt-4">

src/components/dashboard/WidgetConfigModal.tsx

@@ -419,8 +419,13 @@ const WidgetConfigModal = ({
                       </h4>
                       <div className="flex items-center gap-2">
                         <div className="flex items-center gap-2">
-                          <Brain size={16} className="text-gray-600 dark:text-gray-400" />
+                          <Brain
-                          <span className="text-sm text-gray-700 dark:text-gray-300">Thinking</span>
+                            size={16}
+                            className="text-gray-600 dark:text-gray-400"
+                          />
+                          <span className="text-sm text-gray-700 dark:text-gray-300">
+                            Thinking
+                          </span>
                           <Switch
                             checked={showThinking}
                             onChange={setShowThinking}

src/lib/search/simplifiedAgent.ts

@@ -20,6 +20,43 @@ import { formatDateForLLM } from '../utils';
 import { getModelName } from '../utils/modelUtils';
 import { removeThinkingBlocks } from '../utils/contentUtils';
 
+/**
+ * Normalize usage metadata from different LLM providers
+ */
+function normalizeUsageMetadata(usageData: any): {
+  input_tokens: number;
+  output_tokens: number;
+  total_tokens: number;
+} {
+  if (!usageData) return { input_tokens: 0, output_tokens: 0, total_tokens: 0 };
+
+  // Handle different provider formats
+  const inputTokens =
+    usageData.input_tokens ||
+    usageData.prompt_tokens ||
+    usageData.promptTokens ||
+    usageData.usedTokens ||
+    0;
+
+  const outputTokens =
+    usageData.output_tokens ||
+    usageData.completion_tokens ||
+    usageData.completionTokens ||
+    0;
+
+  const totalTokens =
+    usageData.total_tokens ||
+    usageData.totalTokens ||
+    usageData.usedTokens ||
+    inputTokens + outputTokens;
+
+  return {
+    input_tokens: inputTokens,
+    output_tokens: outputTokens,
+    total_tokens: totalTokens,
+  };
+}
+
 /**
  * Simplified Agent using createReactAgent
  *
@@ -487,6 +524,11 @@ Use all available tools strategically to provide comprehensive, well-researched,
       let finalResult: any = null;
       let collectedDocuments: any[] = [];
       let currentResponseBuffer = '';
+      let totalUsage = {
+        input_tokens: 0,
+        output_tokens: 0,
+        total_tokens: 0,
+      };
 
       // Process the event stream
       for await (const event of eventStream) {
@@ -540,6 +582,31 @@ Use all available tools strategically to provide comprehensive, well-researched,
         // Handle streaming tool calls (for thought messages)
         if (event.event === 'on_chat_model_end' && event.data.output) {
           const output = event.data.output;
+
+          // Collect token usage from chat model end events
+          if (output.usage_metadata) {
+            const normalized = normalizeUsageMetadata(output.usage_metadata);
+            totalUsage.input_tokens += normalized.input_tokens;
+            totalUsage.output_tokens += normalized.output_tokens;
+            totalUsage.total_tokens += normalized.total_tokens;
+            console.log(
+              'SimplifiedAgent: Collected usage from usage_metadata:',
+              normalized,
+            );
+          } else if (output.response_metadata?.usage) {
+            // Fallback to response_metadata for different model providers
+            const normalized = normalizeUsageMetadata(
+              output.response_metadata.usage,
+            );
+            totalUsage.input_tokens += normalized.input_tokens;
+            totalUsage.output_tokens += normalized.output_tokens;
+            totalUsage.total_tokens += normalized.total_tokens;
+            console.log(
+              'SimplifiedAgent: Collected usage from response_metadata:',
+              normalized,
+            );
+          }
+
           if (
             output._getType() === 'ai' &&
             output.tool_calls &&
@@ -589,6 +656,25 @@ Use all available tools strategically to provide comprehensive, well-researched,
           }
         }
 
+        // Handle LLM end events for token usage tracking
+        if (event.event === 'on_llm_end' && event.data.output) {
+          const output = event.data.output;
+
+          // Collect token usage from LLM end events
+          if (output.llmOutput?.tokenUsage) {
+            const normalized = normalizeUsageMetadata(
+              output.llmOutput.tokenUsage,
+            );
+            totalUsage.input_tokens += normalized.input_tokens;
+            totalUsage.output_tokens += normalized.output_tokens;
+            totalUsage.total_tokens += normalized.total_tokens;
+            console.log(
+              'SimplifiedAgent: Collected usage from llmOutput:',
+              normalized,
+            );
+          }
+        }
+
         // Handle token-level streaming for the final response
         if (event.event === 'on_chat_model_stream' && event.data.chunk) {
           const chunk = event.data.chunk;
@@ -676,11 +762,15 @@ Use all available tools strategically to provide comprehensive, well-researched,
 
       // Emit model stats and end signal after streaming is complete
       const modelName = getModelName(this.llm);
+      console.log('SimplifiedAgent: Total usage collected:', totalUsage);
       this.emitter.emit(
         'stats',
         JSON.stringify({
           type: 'modelStats',
-          data: { modelName },
+          data: {
+            modelName,
+            usage: totalUsage,
+          },
         }),
       );