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

2025-08-04 00:41:31 -06:00 · 2025-08-04 00:41:31 -06:00 · 71120c997a
commit 71120c997a
parent 3e238303b0
11 changed files with 354 additions and 104 deletions
--- a/.github/copilot-instructions.md
+++ b/.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
--- a/src/app/api/chat/route.ts
+++ b/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 (
--- a/src/components/ChatWindow.tsx
+++ b/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 = {
--- a/src/components/CitationLink.tsx
+++ b/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;
--- a/src/components/MarkdownRenderer.tsx
+++ b/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 
-      content={children} 
+    <ThinkBox
+      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) => (
-          <a {...props} target="_blank" rel="noopener noreferrer" />
-        ),
+        component: (props) => {
+          // 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
--- a/src/components/MessageActions/ModelInfo.tsx
+++ b/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>
--- a/src/components/MessageSource.tsx
+++ b/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;
--- a/src/components/MessageSources.tsx
+++ b/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
-          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>
+        <MessageSource key={i} source={source} index={i} />
      ))}
    </div>
  );
--- a/src/components/MessageTabs.tsx
+++ b/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 
-              content={parsedMessage} 
-              className="px-4" 
+            <MarkdownRenderer
+              content={parsedMessage}
+              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">
--- a/src/components/dashboard/WidgetConfigModal.tsx
+++ b/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" />
-                          <span className="text-sm text-gray-700 dark:text-gray-300">Thinking</span>
+                          <Brain
+                            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}
--- a/src/lib/search/simplifiedAgent.ts
+++ b/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,
+          },
        }),
      );