Slackbot handles mentions

connectors/src/connectors/slack/bot.ts

@@ -1,7 +1,9 @@
 import {
+  AgentActionType,
   AgentGenerationSuccessEvent,
   AgentMessageType,
   DustAPI,
+  RetrievalDocumentType,
 } from "@connectors/lib/dust_api";
 import {
   Connector,
@@ -238,6 +240,7 @@ async function botAnswerMessage(
   }
 
   let fullAnswer = "";
+  let action: AgentActionType | null = null;
   let lastSentDate = new Date();
   for await (const event of streamRes.value.eventStream) {
     switch (event.type) {
@@ -264,22 +267,25 @@ async function botAnswerMessage(
         lastSentDate = new Date();
         await slackClient.chat.update({
           channel: slackChannel,
-          text: fullAnswer,
+          text: _processCiteMention(fullAnswer, action),
           ts: mainMessage.ts as string,
           thread_ts: slackMessageTs,
         });
         break;
       }
+      case "agent_action_success": {
+        action = event.action;
+        break;
+      }
       case "agent_generation_success": {
-        const finalAnswer = `${_removeCiteMention(
-          event.text
-        )}\n\n <${DUST_API}/w/${connector.workspaceId}/assistant/${
-          conversation.sId
-        }|Continue this conversation on Dust>`;
+        fullAnswer = event.text;
+
+        let processedText = _processCiteMention(fullAnswer, action);
+        processedText += `\n\n <${DUST_API}/w/${connector.workspaceId}/assistant/${conversation.sId}|Continue this conversation on Dust>`;
 
         await slackClient.chat.update({
           channel: slackChannel,
-          text: finalAnswer,
+          text: processedText,
           ts: mainMessage.ts as string,
           thread_ts: slackMessageTs,
         });
@@ -293,10 +299,39 @@ async function botAnswerMessage(
   return new Err(new Error("Failed to get the final answer from Dust"));
 }
 
-/*
- * Temp > until I have a PR to properly handle mentions
- */
-function _removeCiteMention(message: string) {
+function _processCiteMention(
+  content: string,
+  action: AgentActionType | null
+): string {
+  const references: { [key: string]: RetrievalDocumentType } = {};
+
+  if (action && action.type === "retrieval_action" && action.documents) {
+    action.documents.forEach((d) => {
+      references[d.reference] = d;
+    });
+  }
+
+  if (references) {
+    let counter = 0;
+    return content.replace(/:cite\[[a-zA-Z0-9,]+\]/g, (match) => {
+      const keys = match.slice(6, -1).split(","); // slice off ":cite[" and "]" then split by comma
+      return keys
+        .map((key) => {
+          const ref = references[key];
+          if (ref && ref.sourceUrl) {
+            counter++;
+            return `<${ref.sourceUrl}|${counter}>`;
+          }
+          return "";
+        })
+        .join("");
+    });
+  }
+
+  return _removeCiteMention(content);
+}
+
+function _removeCiteMention(message: string): string {
   const regex = /:cite\[[a-zA-Z0-9,]+\]/g;
   return message.replace(regex, "");
 }

connectors/src/lib/dust_api.ts

@@ -280,7 +280,6 @@ export type AgentMessageType = {
   visibility: MessageVisibility;
   version: number;
   parentMessageId: string | null;
-
   // configuration: AgentConfigurationType;
   status: AgentMessageStatus;
   action: AgentActionType | null;