From 870a2fca67db1d9c8b9c19fb0b1d49911fc3cb3e Mon Sep 17 00:00:00 2001
From: JulesBelveze <jules.belveze@hotmail.fr>
Date: Wed, 18 Dec 2024 21:01:17 +0100
Subject: [PATCH] [extension] - refactor: streamline icon imports in
 AgentMessage component

 - Removed individual icon imports for various services and replaced with a generic document icon
 - Implemented a citation icon map to dynamically resolve icons based on document provider
 - Simplified makeDocumentCitation and makeWebsearchResultsCitation by integrating icon components directly
---
 .../components/conversation/AgentMessage.tsx  | 24 ++++++-------------
 1 file changed, 7 insertions(+), 17 deletions(-)

diff --git a/extension/app/src/components/conversation/AgentMessage.tsx b/extension/app/src/components/conversation/AgentMessage.tsx
index 4c7c552e830cc..8eae148a54c5a 100644
--- a/extension/app/src/components/conversation/AgentMessage.tsx
+++ b/extension/app/src/components/conversation/AgentMessage.tsx
@@ -23,25 +23,12 @@ import {
   removeNulls,
 } from "@dust-tt/client";
 import type { ConversationMessageSizeType } from "@dust-tt/sparkle";
-import {
-  ConfluenceLogo,
-  DocumentTextIcon,
-  DriveLogo,
-  GithubLogo,
-  ImageIcon,
-  IntercomLogo,
-  MicrosoftLogo,
-  NotionLogo,
-  SlackLogo,
-  SnowflakeLogo,
-  ZendeskLogo,
-} from "@dust-tt/sparkle";
+import { DocumentTextStrokeIcon } from "@dust-tt/sparkle";
 import {
   Citation,
   CitationIcons,
   CitationIndex,
   CitationTitle,
-  Icon,
 } from "@dust-tt/sparkle";
 import {
   ArrowPathIcon,
@@ -59,12 +46,13 @@ import {
 } from "@dust-tt/sparkle";
 import { AgentMessageActions } from "@extension/components/conversation/AgentMessageActions";
 import { GenerationContext } from "@extension/components/conversation/GenerationContextProvider";
-import type { MarkdownCitation } from "@extension/components/conversation/MarkdownCitation";
 import {
   CitationsContext,
   CiteBlock,
   getCiteDirective,
 } from "@extension/components/markdown/CiteBlock";
+import type { MarkdownCitation } from "@extension/components/markdown/MarkdownCitation";
+import { citationIconMap } from "@extension/components/markdown/MarkdownCitation";
 import {
   MentionBlock,
   mentionDirective,
@@ -103,10 +91,12 @@ export function visualizationDirective() {
 export function makeDocumentCitation(
   document: RetrievalDocumentPublicType
 ): MarkdownCitation {
+  const IconComponent =
+    citationIconMap[getProviderFromRetrievedDocument(document)];
   return {
     href: document.sourceUrl ?? undefined,
     title: getTitleFromRetrievedDocument(document),
-    type: getProviderFromRetrievedDocument(document),
+    icon: <IconComponent />,
   };
 }
 
@@ -117,7 +107,7 @@ export function makeWebsearchResultsCitation(
     description: result.snippet,
     href: result.link,
     title: result.title,
-    type: "document" as const,
+    icon: <DocumentTextStrokeIcon />,
   };
 }