From 8cf736458a2a5967bc3c52ff7ac96e5ce7f8f7c3 Mon Sep 17 00:00:00 2001
From: "DESKTOP-RTLN3BA\\$punk" <vermarohanfinal@gmail.com>
Date: Thu, 24 Apr 2025 23:45:50 -0700
Subject: [PATCH 1/2] fix: list citations in markdown

---
 surfsense_web/components/markdown-viewer.tsx | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/surfsense_web/components/markdown-viewer.tsx b/surfsense_web/components/markdown-viewer.tsx
index 03fe7bada..3cd75434c 100644
--- a/surfsense_web/components/markdown-viewer.tsx
+++ b/surfsense_web/components/markdown-viewer.tsx
@@ -37,6 +37,13 @@ export function MarkdownViewer({ content, className, getCitationSource }: Markdo
               : children;
             return <a className="text-primary hover:underline" {...props}>{processedChildren}</a>;
           },
+          li: ({node, children, ...props}) => {
+            // Process citations within list item content
+            const processedChildren = getCitationSource 
+              ? processCitationsInReactChildren(children, getCitationSource) 
+              : children;
+            return <li {...props}>{processedChildren}</li>;
+          },
           ul: ({node, ...props}) => <ul className="list-disc pl-5 my-2" {...props} />,
           ol: ({node, ...props}) => <ol className="list-decimal pl-5 my-2" {...props} />,
           h1: ({node, children, ...props}) => {

From 91936b0b28b9553f4a13b4ed4f52d394a7053cc3 Mon Sep 17 00:00:00 2001
From: "DESKTOP-RTLN3BA\\$punk" <vermarohanfinal@gmail.com>
Date: Fri, 25 Apr 2025 00:19:45 -0700
Subject: [PATCH 2/2] fix: wrong citation mapping

---
 .../researcher/[chat_id]/page.tsx             | 117 ++++++++++++------
 1 file changed, 82 insertions(+), 35 deletions(-)

diff --git a/surfsense_web/app/dashboard/[search_space_id]/researcher/[chat_id]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/researcher/[chat_id]/page.tsx
index 8156f6d2a..b10302566 100644
--- a/surfsense_web/app/dashboard/[search_space_id]/researcher/[chat_id]/page.tsx
+++ b/surfsense_web/app/dashboard/[search_space_id]/researcher/[chat_id]/page.tsx
@@ -617,48 +617,88 @@ const ChatPage = () => {
   };
 
   // Function to get a citation source by ID
-  const getCitationSource = (citationId: number): Source | null => {
+  const getCitationSource = (citationId: number, messageIndex?: number): Source | null => {
     if (!messages || messages.length === 0) return null;
 
-    // Find the latest assistant message
-    const assistantMessages = messages.filter(msg => msg.role === 'assistant');
-    if (assistantMessages.length === 0) return null;
+    // If no specific message index is provided, use the latest assistant message
+    if (messageIndex === undefined) {
+      // Find the latest assistant message
+      const assistantMessages = messages.filter(msg => msg.role === 'assistant');
+      if (assistantMessages.length === 0) return null;
 
-    const latestAssistantMessage = assistantMessages[assistantMessages.length - 1];
-    if (!latestAssistantMessage?.annotations) return null;
+      const latestAssistantMessage = assistantMessages[assistantMessages.length - 1];
+      if (!latestAssistantMessage?.annotations) return null;
 
-    // Find all SOURCES annotations
-    const annotations = latestAssistantMessage.annotations as any[];
-    const sourcesAnnotations = annotations.filter(
-      (annotation) => annotation.type === 'SOURCES'
-    );
+      // Find all SOURCES annotations
+      const annotations = latestAssistantMessage.annotations as any[];
+      const sourcesAnnotations = annotations.filter(
+        (annotation) => annotation.type === 'SOURCES'
+      );
 
-    // Get the latest SOURCES annotation
-    if (sourcesAnnotations.length === 0) return null;
-    const latestSourcesAnnotation = sourcesAnnotations[sourcesAnnotations.length - 1];
+      // Get the latest SOURCES annotation
+      if (sourcesAnnotations.length === 0) return null;
+      const latestSourcesAnnotation = sourcesAnnotations[sourcesAnnotations.length - 1];
 
-    if (!latestSourcesAnnotation.content) return null;
+      if (!latestSourcesAnnotation.content) return null;
 
-    // Flatten all sources from all connectors
-    const allSources: Source[] = [];
-    latestSourcesAnnotation.content.forEach((connector: ConnectorSource) => {
-      if (connector.sources && Array.isArray(connector.sources)) {
-        connector.sources.forEach((source: SourceItem) => {
-          allSources.push({
-            id: source.id,
-            title: source.title,
-            description: source.description,
-            url: source.url,
-            connectorType: connector.type
+      // Flatten all sources from all connectors
+      const allSources: Source[] = [];
+      latestSourcesAnnotation.content.forEach((connector: ConnectorSource) => {
+        if (connector.sources && Array.isArray(connector.sources)) {
+          connector.sources.forEach((source: SourceItem) => {
+            allSources.push({
+              id: source.id,
+              title: source.title,
+              description: source.description,
+              url: source.url,
+              connectorType: connector.type
+            });
           });
-        });
-      }
-    });
+        }
+      });
 
-    // Find the source with the matching ID
-    const foundSource = allSources.find(source => source.id === citationId);
+      // Find the source with the matching ID
+      const foundSource = allSources.find(source => source.id === citationId);
 
-    return foundSource || null;
+      return foundSource || null;
+    } else {
+      // Use the specific message by index
+      const message = messages[messageIndex];
+      if (!message || message.role !== 'assistant' || !message.annotations) return null;
+
+      // Find all SOURCES annotations
+      const annotations = message.annotations as any[];
+      const sourcesAnnotations = annotations.filter(
+        (annotation) => annotation.type === 'SOURCES'
+      );
+
+      // Get the latest SOURCES annotation
+      if (sourcesAnnotations.length === 0) return null;
+      const latestSourcesAnnotation = sourcesAnnotations[sourcesAnnotations.length - 1];
+
+      if (!latestSourcesAnnotation.content) return null;
+
+      // Flatten all sources from all connectors
+      const allSources: Source[] = [];
+      latestSourcesAnnotation.content.forEach((connector: ConnectorSource) => {
+        if (connector.sources && Array.isArray(connector.sources)) {
+          connector.sources.forEach((source: SourceItem) => {
+            allSources.push({
+              id: source.id,
+              title: source.title,
+              description: source.description,
+              url: source.url,
+              connectorType: connector.type
+            });
+          });
+        }
+      });
+
+      // Find the source with the matching ID
+      const foundSource = allSources.find(source => source.id === citationId);
+
+      return foundSource || null;
+    }
   };
 
   return (
@@ -685,7 +725,11 @@ const ChatPage = () => {
                 <div className="flex-1">
                   <Card className="border-gray-300 dark:border-gray-700">
                     <CardContent className="p-3">
-                      <MarkdownViewer content={message.content} getCitationSource={getCitationSource} className="text-sm" />
+                      <MarkdownViewer 
+                        content={message.content}
+                        getCitationSource={(id) => getCitationSource(id, index)}
+                        className="text-sm" 
+                      />
                     </CardContent>
                   </Card>
                 </div>
@@ -901,13 +945,16 @@ const ChatPage = () => {
                               return (
                                 <MarkdownViewer
                                   content={latestAnswer.content.join('\n')}
-                                  getCitationSource={getCitationSource}
+                                  getCitationSource={(id) => getCitationSource(id, index)}
                                 />
                               );
                             }
 
                             // Fallback to the message content if no ANSWER annotation is available
-                            return <MarkdownViewer content={message.content} getCitationSource={getCitationSource} />;
+                            return <MarkdownViewer 
+                              content={message.content} 
+                              getCitationSource={(id) => getCitationSource(id, index)} 
+                            />;
                           })()}
                         </div>
                       }