From 0217d3c3fe998e7d96d5b74d8398af15e96dc8ed Mon Sep 17 00:00:00 2001
From: JulesBelveze <jules.belveze@hotmail.fr>
Date: Sat, 14 Dec 2024 17:02:34 +0100
Subject: [PATCH] [front/components/actions/retrieval] - feature: integrate
 icon mapping for document citations

 - Implement icon association based on document provider in citations
 - Utilize new `connectorIconMapper` to map providers to SVG icons in citations

[front/components/markdown] - refactor: update MarkdownCitation to use icon components

 - Replace the static list of citation icon types with a dynamic icon component mapper
 - Include SVG component imports for icon representation in markdown citations

[types] - refactor: expose ConnectorProviderDocumentType for usage across packages

 - Export type `ConnectorProviderDocumentType` to enable type usage in front-end components
---
 front/components/actions/retrieval/utils.ts   |  4 +-
 .../components/markdown/MarkdownCitation.tsx  | 47 +++++++++++++------
 .../src/front/assistant/actions/retrieval.ts  |  2 +-
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/front/components/actions/retrieval/utils.ts b/front/components/actions/retrieval/utils.ts
index 86597f0677d1d..a7cb8aa2c1b9d 100644
--- a/front/components/actions/retrieval/utils.ts
+++ b/front/components/actions/retrieval/utils.ts
@@ -5,14 +5,16 @@ import {
 } from "@dust-tt/types";
 
 import type { MarkdownCitation } from "@app/components/markdown/MarkdownCitation";
+import { connectorIconMapper } from "@app/components/markdown/MarkdownCitation";
 
 export function makeDocumentCitation(
   document: RetrievalDocumentType
 ): MarkdownCitation {
+  const provider = getProviderFromRetrievedDocument(document);
   return {
     href: document.sourceUrl ?? undefined,
     title: getTitleFromRetrievedDocument(document),
-    type: getProviderFromRetrievedDocument(document),
+    icon: connectorIconMapper[provider],
   };
 }
 
diff --git a/front/components/markdown/MarkdownCitation.tsx b/front/components/markdown/MarkdownCitation.tsx
index b7eae8b110e60..4fc9fc95c670f 100644
--- a/front/components/markdown/MarkdownCitation.tsx
+++ b/front/components/markdown/MarkdownCitation.tsx
@@ -1,22 +1,39 @@
-const CITATION_ICONS = [
-  "confluence",
-  "document",
-  "github",
-  "google_drive",
-  "intercom",
-  "microsoft",
-  "zendesk",
-  "notion",
-  "slack",
-  "image",
-  "snowflake",
-] as const;
+import {
+  ConfluenceLogo,
+  DocumentTextIcon,
+  DriveLogo,
+  GithubLogo,
+  ImageIcon,
+  IntercomLogo,
+  MicrosoftLogo,
+  NotionLogo,
+  SlackLogo,
+  SnowflakeLogo,
+  ZendeskLogo,
+} from "@dust-tt/sparkle";
+import type { ConnectorProviderDocumentType } from "@dust-tt/types";
+import type { SVGProps } from "react";
 
-export type CitationIconType = (typeof CITATION_ICONS)[number];
+export const connectorIconMapper: Record<
+  ConnectorProviderDocumentType | "image",
+  (props: SVGProps<SVGSVGElement>) => React.JSX.Element
+> = {
+  confluence: ConfluenceLogo,
+  document: DocumentTextIcon,
+  github: GithubLogo,
+  google_drive: DriveLogo,
+  intercom: IntercomLogo,
+  microsoft: MicrosoftLogo,
+  zendesk: ZendeskLogo,
+  notion: NotionLogo,
+  slack: SlackLogo,
+  image: ImageIcon,
+  snowflake: SnowflakeLogo,
+};
 
 export interface MarkdownCitation {
   description?: string;
   href?: string;
   title: string;
-  type: CitationIconType;
+  icon: (props: SVGProps<SVGSVGElement>) => React.JSX.Element;
 }
diff --git a/types/src/front/assistant/actions/retrieval.ts b/types/src/front/assistant/actions/retrieval.ts
index efa8e5dd6edf6..995e5d07b4299 100644
--- a/types/src/front/assistant/actions/retrieval.ts
+++ b/types/src/front/assistant/actions/retrieval.ts
@@ -90,7 +90,7 @@ export interface RetrievalDocumentType {
   timestamp: number;
 }
 
-type ConnectorProviderDocumentType =
+export type ConnectorProviderDocumentType =
   | Exclude<ConnectorProvider, "webcrawler">
   | "document";