diff --git a/packages/brick_build/lib/src/builders/aggregate_builder.dart b/packages/brick_build/lib/src/builders/aggregate_builder.dart
index e4103fce..ecfb5459 100644
--- a/packages/brick_build/lib/src/builders/aggregate_builder.dart
+++ b/packages/brick_build/lib/src/builders/aggregate_builder.dart
@@ -67,9 +67,9 @@ class AggregateBuilder implements Builder {
         imports.addAll(findAllImports(contents));
         final newContents = contents
             .replaceAll(importRegex, '')
-            .replaceAll(RegExp("part of '.*';"), '')
-            .replaceAll(RegExp(r"^part\s'.*';", multiLine: true), '')
-            .replaceAll(RegExp(r'^export\s.*;', multiLine: true), '');
+            .replaceAll(RegExp(r'''^part of\s+["'][^"']+["'];''', multiLine: true), '')
+            .replaceAll(RegExp(r'''^part\s+["'][^"']+["'];''', multiLine: true), '')
+            .replaceAll(RegExp(r'^export\s+[^\n;]+;', multiLine: true), '');
         files.add(newContents);
       }
     }
diff --git a/packages/brick_build/lib/src/builders/model_dictionary_builder.dart b/packages/brick_build/lib/src/builders/model_dictionary_builder.dart
index 21caeff4..7054fd99 100644
--- a/packages/brick_build/lib/src/builders/model_dictionary_builder.dart
+++ b/packages/brick_build/lib/src/builders/model_dictionary_builder.dart
@@ -64,7 +64,11 @@ class ModelDictionaryBuilder<_ClassAnnotation> extends BaseBuilder<_ClassAnnotat
       {
         for (final annotation in annotations)
           '${annotation.element.name}': filesToContents.entries
-              .firstWhere((entry) => entry.value.contains(RegExp(r'class\s+${annotation.element.name}\s*')))
+              .firstWhere(
+                (entry) =>
+                    entry.key.endsWith('.model.dart') && // Only use model files
+                    entry.value.contains(RegExp('class\\s+${annotation.element.name}\\b')),
+              )
               .key
               // Make relative from the `brick/` folder
               .replaceAll(RegExp('^lib/'), '../'),