Add document link downloader to Document Library (#236)

bedanley · web-flow · commit 95f1b948394a · 2025-02-03T19:55:45.000-07:00
* Add download link in document library UI
diff --git a/jest.config.js b/jest.config.js
@@ -23,4 +23,5 @@ module.exports = {
         '^.+\\.tsx?$': 'ts-jest'
     },
     collectCoverage: true,
+    silent: true,
 };
diff --git a/lambda/repository/lambda_functions.py b/lambda/repository/lambda_functions.py
@@ -505,6 +505,40 @@ def ingest_documents(event: dict, context: dict) -> dict:
     }
 
 
+@api_wrapper
+def download_document(event: dict, context: dict) -> str:
+    """Generate a pre-signed S3 URL for downloading a file from the RAG ingested files.
+    Args:
+        event (dict): The Lambda event object containing:
+            path_params:
+                repositoryId - the repository
+                documentId - the document
+
+    Returns:
+        url: The presigned URL response object with download fields and URL
+
+    Notes:
+        - URL expires in 300 seconds (5 mins)
+    """
+    path_params = event.get("pathParameters", {}) or {}
+    repository_id = path_params.get("repositoryId")
+    document_id = path_params.get("documentId")
+
+    ensure_repository_access(event, find_repository_by_id(repository_id))
+    doc = doc_repo.find_by_id(repository_id=repository_id, document_id=document_id)
+
+    source = doc.get("source")
+    bucket, key = source.replace("s3://", "").split("/", 1)
+
+    url: str = s3.generate_presigned_url(
+        ClientMethod="get_object",
+        Params={"Bucket": bucket, "Key": key},
+        ExpiresIn=300,
+    )
+
+    return url
+
+
 @api_wrapper
 def presigned_url(event: dict, context: dict) -> dict:
     """Generate a pre-signed URL for uploading files to the RAG ingest bucket.
diff --git a/lib/networking/vpc/security-group-factory.ts b/lib/networking/vpc/security-group-factory.ts
@@ -46,7 +46,7 @@ export class SecurityGroupFactory {
         description: string,
     ): ISecurityGroup {
         if (securityGroupOverride) {
-            console.log(`Security Role Override provided. Using ${securityGroupOverride} for ${securityGroupId}`);
+            console.debug(`Security Role Override provided. Using ${securityGroupOverride} for ${securityGroupId}`);
             const sg = SecurityGroup.fromSecurityGroupId(construct, securityGroupId, securityGroupOverride);
             // Validate the security group exists
             if (!sg) {
diff --git a/lib/rag/api/repository.ts b/lib/rag/api/repository.ts
@@ -134,6 +134,16 @@ export class RepositoryApi extends Construct {
                     ...baseEnvironment,
                 },
             },
+            {
+                name: 'download_document',
+                resource: 'repository',
+                description: 'Creates presigned url to download document within repository',
+                path: 'repository/{repositoryId}/{documentId}/download',
+                method: 'GET',
+                environment: {
+                    ...baseEnvironment,
+                },
+            },
         ];
         apis.forEach((f) => {
             registerAPIEndpoint(
diff --git a/lib/rag/index.ts b/lib/rag/index.ts
@@ -363,14 +363,14 @@ export class LisaRagStack extends Stack {
             }
 
             // Create ingest pipeline state machines for each pipeline config
-            console.log('[DEBUG] Checking pipelines configuration:', {
+            console.debug('Checking pipelines configuration:', {
                 hasPipelines: !!ragConfig.pipelines,
                 pipelinesLength: ragConfig.pipelines?.length || 0
             });
 
             if (ragConfig.pipelines) {
                 ragConfig.pipelines.forEach((pipelineConfig, index) => {
-                    console.log(`[DEBUG] Creating pipeline ${index}:`, {
+                    console.debug(`Creating pipeline ${index}:`, {
                         pipelineConfig: JSON.stringify(pipelineConfig, null, 2)
                     });
 
@@ -420,9 +420,9 @@ export class LisaRagStack extends Stack {
                             });
                             policy.attachToRole(lambdaRole);
                         }
-                        console.log(`[DEBUG] Successfully created pipeline ${index}`);
+                        console.debug(`Successfully created pipeline ${index}`);
                     } catch (error) {
-                        console.error(`[ERROR] Failed to create pipeline ${index}:`, error);
+                        console.error(`Failed to create pipeline ${index}:`, error);
                         throw error; // Re-throw to ensure CDK deployment fails
                     }
                 });
diff --git a/lib/user-interface/react/src/components/document-library/DocumentLibraryComponent.tsx b/lib/user-interface/react/src/components/document-library/DocumentLibraryComponent.tsx
@@ -24,7 +24,11 @@ import {
     TextFilter,
 } from '@cloudscape-design/components';
 import SpaceBetween from '@cloudscape-design/components/space-between';
-import { useDeleteRagDocumentsMutation, useListRagDocumentsQuery } from '../../shared/reducers/rag.reducer';
+import {
+    useDeleteRagDocumentsMutation,
+    useLazyDownloadRagDocumentQuery,
+    useListRagDocumentsQuery,
+} from '../../shared/reducers/rag.reducer';
 import Table from '@cloudscape-design/components/table';
 import ButtonDropdown from '@cloudscape-design/components/button-dropdown';
 import { DEFAULT_PREFERENCES, PAGE_SIZE_OPTIONS, TABLE_DEFINITION, TABLE_PREFERENCES } from './DocumentLibraryConfig';
@@ -35,6 +39,7 @@ import { selectCurrentUserIsAdmin, selectCurrentUsername } from '../../shared/re
 import { RagDocument } from '../types';
 import { setConfirmationModal } from '../../shared/reducers/modal.reducer';
 import { useLocalStorage } from '../../shared/hooks/use-local-storage';
+import { downloadFile } from '../../shared/util/downloader';
 
 type DocumentLibraryComponentProps = {
     repositoryId?: string;
@@ -54,9 +59,8 @@ function disabledDeleteReason (selectedItems: ReadonlyArray<RagDocument>) {
 
 export function DocumentLibraryComponent ({ repositoryId }: DocumentLibraryComponentProps): ReactElement {
     const { data: allDocs, isFetching } = useListRagDocumentsQuery({ repositoryId }, { refetchOnMountOrArgChange: 5 });
-    const [deleteMutation, {
-        isLoading: isDeleteLoading,
-    }] = useDeleteRagDocumentsMutation();
+    const [deleteMutation, { isLoading: isDeleteLoading }] = useDeleteRagDocumentsMutation();
+
     const currentUser = useAppSelector(selectCurrentUsername);
     const isAdmin = useAppSelector(selectCurrentUserIsAdmin);
     const [preferences, setPreferences] = useLocalStorage('DocumentRagPreferences', DEFAULT_PREFERENCES);
@@ -84,13 +88,18 @@ export function DocumentLibraryComponent ({ repositoryId }: DocumentLibraryCompo
             selection: { trackBy: 'document_id' },
         },
     );
-
+    const [getDownloadUrl, { isFetching: isDownloading }] = useLazyDownloadRagDocumentQuery();
     const actionItems: ButtonDropdownProps.Item[] = [
         {
             id: 'rm',
             text: 'Delete',
             disabled: !canDeleteAll(collectionProps.selectedItems, currentUser, isAdmin),
             disabledReason: disabledDeleteReason(collectionProps.selectedItems),
+        }, {
+            id: 'download',
+            text: 'Download',
+            disabled: collectionProps.selectedItems.length > 1,
+            disabledReason: 'Only one file can be downloaded at a time',
         },
     ];
     const handleAction = async (e: any) => {
@@ -109,6 +118,12 @@ export function DocumentLibraryComponent ({ repositoryId }: DocumentLibraryCompo
                 );
                 break;
             }
+            case 'download': {
+                const { document_id, document_name } = collectionProps.selectedItems[0];
+                const resp = await getDownloadUrl({ documentId: document_id, repositoryId });
+                downloadFile(resp.data, document_name);
+                break;
+            }
             default:
                 console.error('Action not implemented', e.detail.id);
         }
@@ -150,8 +165,9 @@ export function DocumentLibraryComponent ({ repositoryId }: DocumentLibraryCompo
                         >
                             <ButtonDropdown
                                 items={actionItems}
-                                loading={isDeleteLoading}
-                                onItemClick={(e) => handleAction(e)}
+                                loading={isDeleteLoading || isDownloading}
+                                disabled={collectionProps.selectedItems.length === 0}
+                                onItemClick={async (e) => handleAction(e)}
                             >
                                 Actions
                             </ButtonDropdown>
diff --git a/lib/user-interface/react/src/pages/Home.tsx b/lib/user-interface/react/src/pages/Home.tsx
@@ -58,8 +58,8 @@ export function Home ({ setNav }) {
             }
         >
             {visible && <Alert type='error'>You must sign in to access this page!</Alert>}
-            <Box float='center'>
-                <div align='center'>
+            <Box>
+                <div>
                     <figure>
                         <img
                             src={chatImg}
diff --git a/lib/user-interface/react/src/shared/reducers/rag.reducer.ts b/lib/user-interface/react/src/shared/reducers/rag.reducer.ts
@@ -135,6 +135,12 @@ export const ragApi = createApi({
             },
             invalidatesTags: ['Docs'],
         }),
+        downloadRagDocument: builder.query<string, { documentId: string, repositoryId: string }>({
+            query: ({ documentId, repositoryId }) => ({
+                url: `/repository/${repositoryId}/${documentId}/download`,
+                method: 'GET',
+            }),
+        }),
     }),
 });
 
@@ -145,5 +151,6 @@ export const {
     useIngestDocumentsMutation,
     useListRagDocumentsQuery,
     useDeleteRagDocumentsMutation,
-    useLazyGetRelevantDocumentsQuery
+    useLazyGetRelevantDocumentsQuery,
+    useLazyDownloadRagDocumentQuery,
 } = ragApi;
diff --git a/lib/user-interface/react/src/shared/util/downloader.ts b/lib/user-interface/react/src/shared/util/downloader.ts
@@ -0,0 +1,31 @@
+/**
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+/**
+ * Download a file using hidden link
+ * @param url
+ * @param filename
+ */
+export function downloadFile (url: string, filename: string) {
+    const link = document.createElement('a');
+    link.href = url;
+    link.download = filename || 'download';
+    link.hidden = true;
+
+    document.body.appendChild(link);
+    link.click();
+    document.body.removeChild(link);
+}
diff --git a/test/cdk/stacks/nag.test.ts b/test/cdk/stacks/nag.test.ts
@@ -39,7 +39,7 @@ const nagResults: NagResult = {
     LisaIAM: [0,13,0,0],
     LisaModels: [1,74,0,28],
     LisaNetworking: [1,2,3,5],
-    LisaRAG: [2,41,0,38],
+    LisaRAG: [2,43,0,38],
     LisaServe: [1,21,0,31],
     LisaUI: [0,16,0,8],
 };

Original file line number	Diff line number	Diff line change
`@@ -58,8 +58,8 @@ export function Home ({ setNav }) {`
`58`	`58`	`}`
`59`	`59`	`>`
`60`	`60`	`{visible && <Alert type='error'>You must sign in to access this page!</Alert>}`
`61`		`- <Box float='center'>`
`62`		`- <div align='center'>`
	`61`	`+ <Box>`
	`62`	`+ <div>`
`63`	`63`	`<figure>`
`64`	`64`	`<img`
`65`	`65`	`src={chatImg}`