Resolve comments.

sdks/python/apache_beam/ml/rag/enrichment/bigquery_vector_search.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 #
 
+from collections import defaultdict
 from dataclasses import dataclass
 from typing import Any
 from typing import Dict
@@ -187,12 +188,10 @@ def format_query(self, chunks: List[Chunk]) -> str:
     """ if self.metadata_restriction_template else ""
 
     # Group chunks by their metadata conditions
-    condition_groups = {}
+    condition_groups = defaultdict(list)
     if self.metadata_restriction_template:
       for chunk in chunks:
         condition = self.metadata_restriction_template.format(**chunk.metadata)
-        if condition not in condition_groups:
-          condition_groups[condition] = []
         condition_groups[condition].append(chunk)
     else:
       # No metadata filtering - all chunks in one group

sdks/python/apache_beam/ml/rag/enrichment/bigquery_vector_search_it_test.py

@@ -23,18 +23,21 @@
 import apache_beam as beam
 from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper
 from apache_beam.io.gcp.internal.clients import bigquery
-from apache_beam.ml.rag.enrichment.bigquery_vector_search import BigQueryVectorSearchEnrichmentHandler
-from apache_beam.ml.rag.enrichment.bigquery_vector_search import BigQueryVectorSearchParameters
 from apache_beam.ml.rag.types import Chunk
 from apache_beam.ml.rag.types import Content
 from apache_beam.ml.rag.types import Embedding
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
-from apache_beam.transforms.enrichment import Enrichment
 
+# pylint: disable=ungrouped-imports
 try:
   from google.api_core.exceptions import BadRequest
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.ml.rag.enrichment.bigquery_vector_search import \
+    BigQueryVectorSearchEnrichmentHandler
+  from apache_beam.ml.rag.enrichment.bigquery_vector_search import \
+    BigQueryVectorSearchParameters
 except ImportError:
   raise unittest.SkipTest('BigQuery dependencies not installed')