From 4cbf3a2837f965f285812600aae50dfb2f478145 Mon Sep 17 00:00:00 2001
From: jrmccluskey <thejackmccluskey@gmail.com>
Date: Tue, 21 Jul 2026 19:40:59 +0000
Subject: [PATCH 1/2] Fix pyrefly check unexpected-keyword

---
 sdks/python/apache_beam/ml/inference/base.py | 2 +-
 sdks/python/apache_beam/ml/rag/utils.py      | 3 +++
 sdks/python/pyproject.toml                   | 3 +--
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 62c4047a961f..84d68ef6c061 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -573,7 +573,7 @@ def __init__(self, mh_map: dict[str, ModelHandler]):
     # Map key for a model to a unique tag that will persist for the life of
     # that model in memory. A new tag will be generated if a model is swapped
     # out of memory and reloaded.
-    self._tag_map: dict[str, str] = OrderedDict()
+    self._tag_map: OrderedDict[str, str] = OrderedDict()
     # Map a tag to a multiprocessshared model object for that tag. Each entry
     # of this map should last as long as the corresponding entry in _tag_map.
     self._proxy_map: dict[str, multi_process_shared.MultiProcessShared] = {}
diff --git a/sdks/python/apache_beam/ml/rag/utils.py b/sdks/python/apache_beam/ml/rag/utils.py
index f56bfe518485..d76f9f3a854d 100644
--- a/sdks/python/apache_beam/ml/rag/utils.py
+++ b/sdks/python/apache_beam/ml/rag/utils.py
@@ -62,10 +62,13 @@ class MilvusConnectionParameters:
   token: str = field(default_factory=str)
   timeout: Optional[float] = None
   kwargs: dict[str, Any] = field(default_factory=dict)
+  db_id: Optional[str] = None
 
   def __post_init__(self):
     if not self.uri:
       raise ValueError("URI must be provided for Milvus connection")
+    if self.db_id is not None and self.db_name == "default":
+      self.db_name = self.db_id
 
     # Generate unique alias if not provided. One-to-one mapping between alias
     # and connection - each alias represents exactly one Milvus connection.
diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml
index c5e642a8b37c..a410d83a8598 100644
--- a/sdks/python/pyproject.toml
+++ b/sdks/python/pyproject.toml
@@ -211,5 +211,4 @@ deprecated = "ignore"
 not-callable = "ignore"
 invalid-argument = "ignore"
 invalid-inheritance = "ignore"
-not-iterable = "ignore"
-unexpected-keyword = "ignore"
+not-iterable = "ignore"
\ No newline at end of file

From 0d9317592861bca5830378ed22d21668fdb21f2b Mon Sep 17 00:00:00 2001
From: jrmccluskey <thejackmccluskey@gmail.com>
Date: Tue, 28 Jul 2026 14:27:31 +0000
Subject: [PATCH 2/2] fix field usage instead of adding a new field

---
 .../snippets/transforms/elementwise/enrichment_test.py         | 2 +-
 sdks/python/apache_beam/ml/rag/utils.py                        | 3 ---
 2 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
index 0b356e039930..f37f297530ce 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
@@ -380,7 +380,7 @@ def pre_milvus_enrichment() -> VectorDBContainerInfo:
           uri=db.uri,
           user=db.user,
           password=db.password,
-          db_id=db.id,
+          db_name=db.id,
           token=db.token)
       collection_name = MilvusTestHelpers.initialize_db_with_data(
           connection_params)
diff --git a/sdks/python/apache_beam/ml/rag/utils.py b/sdks/python/apache_beam/ml/rag/utils.py
index d76f9f3a854d..f56bfe518485 100644
--- a/sdks/python/apache_beam/ml/rag/utils.py
+++ b/sdks/python/apache_beam/ml/rag/utils.py
@@ -62,13 +62,10 @@ class MilvusConnectionParameters:
   token: str = field(default_factory=str)
   timeout: Optional[float] = None
   kwargs: dict[str, Any] = field(default_factory=dict)
-  db_id: Optional[str] = None
 
   def __post_init__(self):
     if not self.uri:
       raise ValueError("URI must be provided for Milvus connection")
-    if self.db_id is not None and self.db_name == "default":
-      self.db_name = self.db_id
 
     # Generate unique alias if not provided. One-to-one mapping between alias
     # and connection - each alias represents exactly one Milvus connection.