From 2dc7e64dd7df0fed2e5594d339fa8a89b94a2fa4 Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Thu, 18 Jul 2024 16:15:19 -0700
Subject: [PATCH] fix internet search icons / text + assistants tab (#1862)

---
 web/src/app/assistants/SidebarWrapper.tsx |  1 +
 web/src/app/assistants/mine/page.tsx      |  3 ---
 web/src/app/chat/message/Messages.tsx     | 20 +++++++++++---------
 web/src/lib/sources.ts                    | 19 ++++++++++++++++++-
 4 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/web/src/app/assistants/SidebarWrapper.tsx b/web/src/app/assistants/SidebarWrapper.tsx
index 87081e07f5..c94cce370b 100644
--- a/web/src/app/assistants/SidebarWrapper.tsx
+++ b/web/src/app/assistants/SidebarWrapper.tsx
@@ -143,6 +143,7 @@ export default function SidebarWrapper<T extends object>({
           <div className="mt-4 mx-auto">{content(contentProps)}</div>
         </div>
       </div>
+      <FixedLogo />
     </div>
   );
 }
diff --git a/web/src/app/assistants/mine/page.tsx b/web/src/app/assistants/mine/page.tsx
index 748217b5d6..839299726c 100644
--- a/web/src/app/assistants/mine/page.tsx
+++ b/web/src/app/assistants/mine/page.tsx
@@ -67,9 +67,6 @@ export default async function GalleryPage({
           user={user}
           assistants={assistants}
         />
-
-        {/* Temporary - fixed logo */}
-        <FixedLogo />
       </ChatProvider>
     </>
   );
diff --git a/web/src/app/chat/message/Messages.tsx b/web/src/app/chat/message/Messages.tsx
index a435c97233..b01de17f38 100644
--- a/web/src/app/chat/message/Messages.tsx
+++ b/web/src/app/chat/message/Messages.tsx
@@ -57,6 +57,8 @@ import {
 import { ValidSources } from "@/lib/types";
 import { Tooltip } from "@/components/tooltip/Tooltip";
 import { useMouseTracking } from "./hooks";
+import { InternetSearchIcon } from "@/components/InternetSearchIcon";
+import { getTitleFromDocument } from "@/lib/sources";
 
 const TOOLS_WITH_CUSTOM_HANDLING = [
   SEARCH_TOOL_NAME,
@@ -443,17 +445,17 @@ export const AIMessage = ({
                                   >
                                     <Citation link={doc.link} index={ind + 1} />
                                     <p className="shrink truncate ellipsis break-all ">
-                                      {
-                                        doc.document_id.split("/")[
-                                          doc.document_id.split("/").length - 1
-                                        ]
-                                      }
+                                      {getTitleFromDocument(doc)}
                                     </p>
                                     <div className="ml-auto flex-none">
-                                      <SourceIcon
-                                        sourceType={doc.source_type}
-                                        iconSize={18}
-                                      />
+                                      {doc.is_internet ? (
+                                        <InternetSearchIcon url={doc.link} />
+                                      ) : (
+                                        <SourceIcon
+                                          sourceType={doc.source_type}
+                                          iconSize={18}
+                                        />
+                                      )}
                                     </div>
                                   </a>
                                   <div className="flex overscroll-x-scroll mt-.5">
diff --git a/web/src/lib/sources.ts b/web/src/lib/sources.ts
index 66b99ca516..89891ef789 100644
--- a/web/src/lib/sources.ts
+++ b/web/src/lib/sources.ts
@@ -38,7 +38,11 @@ import {
   ColorSlackIcon,
 } from "@/components/icons/icons";
 import { ValidSources } from "./types";
-import { SourceCategory, SourceMetadata } from "./search/interfaces";
+import {
+  DanswerDocument,
+  SourceCategory,
+  SourceMetadata,
+} from "./search/interfaces";
 import { Persona } from "@/app/admin/assistants/interfaces";
 import { FaAccessibleIcon, FaSlack } from "react-icons/fa";
 
@@ -290,3 +294,16 @@ export function getSourcesForPersona(persona: Persona): ValidSources[] {
   });
   return personaSources;
 }
+
+function stripTrailingSlash(str: string) {
+  if (str.substr(-1) === "/") {
+    return str.substr(0, str.length - 1);
+  }
+  return str;
+}
+
+export const getTitleFromDocument = (document: DanswerDocument) => {
+  return stripTrailingSlash(document.document_id).split("/")[
+    stripTrailingSlash(document.document_id).split("/").length - 1
+  ];
+};