feat: allow to download model generated files other than images (#725)

* feat: allow to download model generated files other than images

* refactor: move markdown image fix into citations fix function

* add file name

Author: tomtobac

src/interfaces/assistants_web/src/app/(main)/(chat)/Chat.tsx

@@ -10,7 +10,7 @@ import { useCitationsStore, useConversationStore, useParamsStore } from '@/store
 import { OutputFiles } from '@/stores/slices/citationsSlice';
 import {
   createStartEndKey,
-  fixCitationsLeadingMarkdown,
+  fixInlineCitationsForMarkdown,
   mapHistoryToMessages,
   parsePythonInterpreterToolFields,
 } from '@/utils';
@@ -86,7 +86,7 @@ const Chat: React.FC<{ agentId?: string; conversationId?: string }> = ({
           }
         }
       });
-      fixCitationsLeadingMarkdown(message.citations, message.text)?.forEach((citation) => {
+      fixInlineCitationsForMarkdown(message.citations, message.text)?.forEach((citation) => {
         const startEndKey = createStartEndKey(citation.start ?? 0, citation.end ?? 0);
         const documents = citation.document_ids?.map((id) => documentsMap[id]) ?? [];
         addCitation(message.generation_id ?? '', startEndKey, documents);

src/interfaces/assistants_web/src/components/Markdown/Markdown.tsx

@@ -11,6 +11,7 @@ import remarkMath from 'remark-math';
 import { PluggableList } from 'unified';
 
 import {
+  Anchor,
   Iframe,
   P,
   Pre,
@@ -111,6 +112,7 @@ export const Markdown = ({
       references: References,
       // @ts-ignore
       iframe: Iframe,
+      a: Anchor,
       ...customComponents,
     }),
     [customComponents]

src/interfaces/assistants_web/src/components/Markdown/tags/Anchor.tsx

@@ -0,0 +1,30 @@
+'use client';
+
+import type { Component, ExtraProps } from 'hast-util-to-jsx-runtime/lib/components';
+import { ComponentPropsWithoutRef } from 'react';
+
+import { Icon } from '@/components/UI';
+import { useOutputFiles } from '@/stores';
+
+export const Anchor: Component<ComponentPropsWithoutRef<'a'> & ExtraProps> = ({ children }) => {
+  const { outputFiles } = useOutputFiles();
+
+  if (typeof children === 'string') {
+    const snakeCaseUrl = children.replace(/\s/g, '_').toLowerCase();
+    const outputFile = Object.entries(outputFiles).find(([key]) =>
+      key.startsWith(snakeCaseUrl)
+    )?.[1];
+    const downloadUrl = outputFile?.downloadUrl;
+
+    if (downloadUrl) {
+      return (
+        <a href={downloadUrl} download={snakeCaseUrl} className="flex items-center gap-1">
+          {children}
+          <Icon name="download" />
+        </a>
+      );
+    }
+  }
+
+  return <a dir="auto">{children}</a>;
+};

src/interfaces/assistants_web/src/components/Markdown/tags/index.ts

@@ -1,3 +1,4 @@
+export * from './Anchor';
 export * from './Iframe';
 export * from './P';
 export * from './Pre';

src/interfaces/assistants_web/src/hooks/use-chat.ts

@@ -42,7 +42,7 @@ import {
 } from '@/types/message';
 import {
   createStartEndKey,
-  fixCitationsLeadingMarkdown,
+  fixInlineCitationsForMarkdown,
   fixMarkdownImagesInText,
   isGroundingOn,
   parsePythonInterpreterToolFields,
@@ -384,7 +384,7 @@ export const useChat = (config?: { onSend?: (msg: string) => void }) => {
             case StreamEvent.CITATION_GENERATION: {
               const data = eventData.data;
               const newCitations = [...(data?.citations ?? [])];
-              const fixedCitations = fixCitationsLeadingMarkdown(newCitations, botResponse);
+              const fixedCitations = fixInlineCitationsForMarkdown(newCitations, botResponse);
               citations.push(...fixedCitations);
               citations.sort((a, b) => (a.start ?? 0) - (b.start ?? 0));
               saveCitations(generationId, fixedCitations, documentsMap);

src/interfaces/assistants_web/src/stores/index.ts

@@ -78,4 +78,13 @@ export const useParamsStore = () => {
   );
 };
 
+export const useOutputFiles = () => {
+  return useStore(
+    (state) => ({
+      outputFiles: state.citations.outputFiles,
+    }),
+    shallow
+  );
+};
+
 export { useSettingsStore } from '@/stores/persistedStore';

src/interfaces/assistants_web/src/stores/slices/citationsSlice.ts

@@ -1,6 +1,7 @@
 import { StateCreator } from 'zustand';
 
 import { Document } from '@/cohere-client';
+import { mapExtensionToMimeType } from '@/utils';
 
 import { StoreState } from '..';
 
@@ -22,7 +23,9 @@ interface SearchResults {
   [documentId: string]: Record<string, any>;
 }
 
-export type OutputFiles = { [name: string]: { name: string; data: string; documentId?: string } };
+export type OutputFiles = {
+  [name: string]: { name: string; data: string; documentId?: string; downloadUrl?: string };
+};
 
 type State = {
   citationReferences: CitationReferences;
@@ -80,6 +83,18 @@ export const createCitationsSlice: StateCreator<StoreState, [], [], CitationsSto
     }));
   },
   saveOutputFiles(outputFiles) {
+    for (const [name, file] of Object.entries(outputFiles)) {
+      if (file.downloadUrl) {
+        continue;
+      }
+      const data = file.data;
+      const fileExtension = file.name.split('.').pop() || '.txt';
+      const mimeType = mapExtensionToMimeType(fileExtension);
+      const blob = new Blob([data], { type: mimeType });
+      const url = URL.createObjectURL(blob);
+      outputFiles[name] = { ...file, downloadUrl: url };
+    }
+
     set((state) => ({
       citations: {
         ...state.citations,

src/interfaces/assistants_web/src/utils/citations.test.ts

@@ -1,7 +1,11 @@
 import { describe, expect, test } from 'vitest';
 
 import { Citation } from '@/cohere-client';
-import { fixCitationsLeadingMarkdown, replaceTextWithCitations } from '@/utils';
+import {
+  fixInlineCitationsForMarkdown,
+  isReferenceBetweenSpecialTags,
+  replaceTextWithCitations,
+} from '@/utils';
 
 describe('replaceTextWithCitations', () => {
   test('should replace text with citations', () => {
@@ -44,46 +48,6 @@ describe('replaceTextWithCitations', () => {
     );
   });
 
-  test('should avoid to break markdown images', () => {
-    const citations: Citation[] = [
-      {
-        start: 0,
-        end: 26,
-        text: '! [test](https://test.com)',
-        document_ids: ['12345'],
-      },
-    ];
-    const text = '![test](https://test.com)';
-    const generationId = '12345';
-    const result = replaceTextWithCitations(text, citations, generationId);
-    expect(result).toBe(
-      ':cite[![test](https://test.com)]{generationId="12345" start="0" end="26"}'
-    );
-  });
-
-  test('should handle extra space when fixing markdown images', () => {
-    const citations: Citation[] = [
-      {
-        start: 5,
-        end: 31,
-        text: '! [test](https://test.com)',
-        document_ids: ['12345'],
-      },
-      {
-        start: 32,
-        end: 39,
-        text: 'Kadabra',
-        document_ids: ['44444'],
-      },
-    ];
-    const text = 'Abra ![test](https://test.com) Kadabra';
-    const generationId = '12345';
-    const result = replaceTextWithCitations(text, citations, generationId);
-    expect(result).toBe(
-      'Abra :cite[![test](https://test.com)]{generationId="12345" start="5" end="31"} :cite[Kadabra]{generationId="12345" start="32" end="39"}'
-    );
-  });
-
   test('should allow citations as markdown elements', () => {
     const citations: Citation[] = [
       {
@@ -100,7 +64,7 @@ describe('replaceTextWithCitations', () => {
   });
 });
 
-describe('fixCitationsLeadingMarkdown', () => {
+describe('fixInlineCitationsForMarkdown', () => {
   test('should fix leading markdown citations breaking markdown', () => {
     const citations: Citation[] = [
       {
@@ -142,7 +106,7 @@ describe('fixCitationsLeadingMarkdown', () => {
     ];
     const text =
       'The `ENVIRONMENT DIVISION` should be included after the `IDENTIFICATION DIVISION`. The `CONFIGURATION SECTION` and `SPECIAL-NAMES` should be included within the `ENVIRONMENT DIVISION`.';
-    const result = fixCitationsLeadingMarkdown(citations, text);
+    const result = fixInlineCitationsForMarkdown(citations, text);
 
     expect(result).toStrictEqual([
       {
@@ -183,4 +147,122 @@ describe('fixCitationsLeadingMarkdown', () => {
       },
     ]);
   });
+
+  test('should push markdown downloable links to the end', () => {
+    const citations: Citation[] = [
+      {
+        text: '[link](https://www.google.com)',
+        start: 10,
+        end: 39,
+        document_ids: ['111111'],
+      },
+      {
+        text: "[link]('.file_path.csv')",
+        start: 64,
+        end: 87,
+        document_ids: ['111111'],
+      },
+    ];
+    const text =
+      "This is a [link](https://www.google.com) and this is downloable [link]('file_path.csv')";
+    const result = fixInlineCitationsForMarkdown(citations, text);
+    expect(result).toStrictEqual([
+      {
+        text: '[link](https://www.google.com)',
+        end: 39,
+        start: 88,
+        document_ids: ['111111'],
+      },
+      {
+        text: "[link]('.file_path.csv')",
+        end: 87,
+        start: 88,
+        document_ids: ['111111'],
+      },
+    ]);
+  });
+
+  test('should remove extra blank space on markdown images', () => {
+    const citations: Citation[] = [
+      {
+        start: 0,
+        end: 25,
+        text: '! [test](https://test.com)',
+        document_ids: ['12345'],
+      },
+      {
+        start: 26,
+        end: 33,
+        text: 'Kadabra',
+        document_ids: ['44444'],
+      },
+    ];
+    const text = '![test](https://test.com) Kadabra';
+    const result = fixInlineCitationsForMarkdown(citations, text);
+    expect(result).toStrictEqual([
+      {
+        start: 34,
+        end: 25,
+        text: '![test](https://test.com)',
+        document_ids: ['12345'],
+      },
+      {
+        start: 25,
+        end: 32,
+        text: 'Kadabra',
+        document_ids: ['44444'],
+      },
+    ]);
+  });
+});
+
+describe('isReferenceBetweenSpecialTags', () => {
+  test('should return true if the citation is between <iframe> tags', () => {
+    const matchRegex = /<iframe>.*<\/iframe>/;
+    const text = '<iframe> This is a citation </iframe>';
+    const citation: Citation = {
+      start: 19,
+      end: 27,
+      text: 'citation',
+      document_ids: ['12345'],
+    };
+    const result = isReferenceBetweenSpecialTags(matchRegex, text, citation.start);
+    expect(result).toBe(true);
+  });
+  test('should return false if the citation is not between <iframe> tags', () => {
+    const matchRegex = /<iframe>.*<\/iframe>/;
+    const text = 'This is a citation <iframe> another test citaiton </iframe>';
+    const citation: Citation = {
+      start: 10,
+      end: 18,
+      text: 'citation',
+      document_ids: ['12345'],
+    };
+    const result = isReferenceBetweenSpecialTags(matchRegex, text, citation.start);
+    expect(result).toBe(false);
+  });
+  test('should return true if the citation is between ``` tags', () => {
+    const matchRegex = /```[\s\S]*?```/g;
+    const text = '``` This is a citation ```';
+    const citation: Citation = {
+      start: 14,
+      end: 22,
+      text: 'citation',
+      document_ids: ['12345'],
+    };
+    const result = isReferenceBetweenSpecialTags(matchRegex, text, citation.start);
+    expect(result).toBe(true);
+  });
+  test('should return false if the citation is not between ``` tags', () => {
+    const matchRegex = /```[\s\S]*?```/g;
+    const text = '``` This is a citation ``` another test citaiton';
+    const citation: Citation = {
+      start: 40,
+      end: 48,
+      text: 'citation',
+      document_ids: ['12345'],
+    };
+    const result = isReferenceBetweenSpecialTags(matchRegex, text, citation.start);
+    expect(result).toBe(false);
+  });
 });