GenAI metadata generation - resolve conflicts+add guardrails

Author: dlpzx

backend/dataall/modules/s3_datasets/aws/bedrock_prompts/metadata_generation_folder_template.txt

@@ -34,7 +34,7 @@ Input:
 
 Output:
 {{
-    "label": "NotEnoughData",
+    "label": "my-folder-1",
     "description": "NotEnoughData"
 }}
 

backend/dataall/modules/s3_datasets/aws/s3_dataset_client.py

@@ -81,7 +81,7 @@ def list_bucket_files(self, bucket_name, prefix):
                 Bucket=bucket_name,
                 Prefix=prefix,
                 ExpectedBucketOwner=dataset.AwsAccountId,
-                MaxKeys=1000,
+                MaxKeys=100,
             )
             return response.get('Contents', [])
         except ClientError as e:

backend/dataall/modules/s3_datasets/db/dataset_column_repositories.py

@@ -48,6 +48,8 @@ def paginate_active_columns_for_table(session, table_uri: str, filter: dict):
         return paginate(query=q, page=filter.get('page', 1), page_size=filter.get('pageSize', 10)).to_dict()
 
     @staticmethod
-    def list_active_columns_for_table(session, table_uri: str):
+    def list_active_columns_for_table(session, table_uri: str, limit=None):
         q = DatasetColumnRepository.query_active_columns_for_table(session, table_uri)
+        if limit:
+            q = q.limit(limit)
         return q.all()

backend/dataall/modules/s3_datasets/db/dataset_location_repositories.py

@@ -97,9 +97,12 @@ def delete_dataset_locations(session, dataset_uri) -> bool:
         return True
 
     @staticmethod
-    def get_dataset_folders(session, dataset_uri):
+    def get_dataset_folders(session, dataset_uri, limit=None):
         """return the dataset folders"""
-        return session.query(DatasetStorageLocation).filter(DatasetStorageLocation.datasetUri == dataset_uri).all()
+        query = session.query(DatasetStorageLocation).filter(DatasetStorageLocation.datasetUri == dataset_uri)
+        if limit:
+            query = query.limit(limit)
+        return query.all()
 
     @staticmethod
     def paginated_dataset_locations(session, uri, data=None) -> dict:

backend/dataall/modules/s3_datasets/db/dataset_repositories.py

@@ -186,9 +186,12 @@ def update_glue_database_status(session, dataset_uri):
         dataset.glueDatabaseCreated = True
 
     @staticmethod
-    def get_dataset_tables(session, dataset_uri):
+    def get_dataset_tables(session, dataset_uri, limit=None):
         """return the dataset tables"""
-        return session.query(DatasetTable).filter(DatasetTable.datasetUri == dataset_uri).all()
+        query = session.query(DatasetTable).filter(DatasetTable.datasetUri == dataset_uri)
+        if limit:
+            query = query.limit(limit)
+        return query.all()
 
     @staticmethod
     def delete_dataset(session, dataset) -> bool:

backend/dataall/modules/s3_datasets/services/dataset_service.py

@@ -601,8 +601,8 @@ def generate_metadata_for_dataset(uri, metadata_types):
         context = get_context()
         with context.db_engine.scoped_session() as session:
             dataset = DatasetBaseRepository.get_dataset_by_uri(session, uri)
-            tables = DatasetRepository.get_dataset_tables(session, dataset.datasetUri)
-            folders = DatasetLocationRepository.get_dataset_folders(session, dataset.datasetUri)
+            tables = DatasetRepository.get_dataset_tables(session, dataset.datasetUri, limit=50)
+            folders = DatasetLocationRepository.get_dataset_folders(session, dataset.datasetUri, limit=50)
             metadata = BedrockClient().invoke_model_dataset_metadata(
                 metadata_types=metadata_types,
                 dataset=dataset,

backend/dataall/modules/s3_datasets/services/dataset_table_service.py

@@ -33,6 +33,7 @@
 from dataall.base.db import exceptions
 from dataall.modules.s3_datasets.aws.bedrock_metadata_client import BedrockClient
 from dataall.modules.s3_datasets.db.dataset_column_repositories import DatasetColumnRepository
+from dataall.modules.s3_datasets.services.dataset_enums import MetadataGenerationTypes
 
 
 log = logging.getLogger(__name__)
@@ -207,12 +208,14 @@ def _delete_dataset_table_read_permission(session, table_uri):
     #     'metadata', 'modules.s3_datasets.features.generate_metadata_ai.max_count_per_day'
     # )
     def generate_metadata_for_table(uri, metadata_types, sample_data):
+        metadataTypesForTable = [MetadataGenerationTypes.Description, MetadataGenerationTypes.Tag]
+        table_metadata_types = [item for item in metadata_types if item in metadataTypesForTable]
         context = get_context()
         with context.db_engine.scoped_session() as session:
             table = DatasetTableRepository.get_dataset_table_by_uri(session, uri)
-            table_columns = DatasetColumnRepository.list_active_columns_for_table(session, table.tableUri)
+            table_columns = DatasetColumnRepository.list_active_columns_for_table(session, table.tableUri, limit=50)
             metadata = BedrockClient().invoke_model_table_metadata(
-                table=table, columns=table_columns, metadata_types=metadata_types, sample_data=sample_data
+                table=table, columns=table_columns, metadata_types=table_metadata_types, sample_data=sample_data
             )
 
             result = [{'targetUri': uri, 'targetType': 'Table', **metadata}]

frontend/src/modules/S3_Datasets/components/ReviewMetadataComponent.js

@@ -95,10 +95,10 @@ export const ReviewMetadataComponent = (props) => {
         if (targetIndex !== -1) {
           const updatedTarget = {
             ...targets[targetIndex],
-            description: response.data.generateMetadata[0].description,
-            label: response.data.generateMetadata[0].label,
-            tags: response.data.generateMetadata[0].tags,
-            topics: response.data.generateMetadata[0].topics
+            description: response.data.generateMetadata[0].description ?? '',
+            label: response.data.generateMetadata[0].label ?? '',
+            tags: response.data.generateMetadata[0].tags ?? '',
+            topics: response.data.generateMetadata[0].topics ?? ''
           };
 
           const updatedTargets = [...targets];