diff --git a/connectors/src/connectors/slack/bot.ts b/connectors/src/connectors/slack/bot.ts index 70370487fca2..58041c479913 100644 --- a/connectors/src/connectors/slack/bot.ts +++ b/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, ""); } diff --git a/connectors/src/lib/dust_api.ts b/connectors/src/lib/dust_api.ts index 9e28bb83ec2d..105367a1528a 100644 --- a/connectors/src/lib/dust_api.ts +++ b/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;