diff --git a/assets/batch_score/components/driver/tests/requirements.txt b/assets/batch_score/components/driver/tests/requirements.txt
index 3e37cd2986..50deb3e200 100644
--- a/assets/batch_score/components/driver/tests/requirements.txt
+++ b/assets/batch_score/components/driver/tests/requirements.txt
@@ -16,6 +16,7 @@ requests>=2.28.1
 tiktoken>=0.3.3
 StrEnum>=0.4.15
 pyDispatcher>=2.0.7
+setuptools>=72.1.0
 
 # batch_score requirements
 -r ../src/batch_score_oss/requirements.txt
diff --git a/assets/batch_score/components/driver/tests/unit/common/header_providers/test_user_agent_header_provider.py b/assets/batch_score/components/driver/tests/unit/common/header_providers/test_user_agent_header_provider.py
index a1ba74a51f..b9c7f7424a 100644
--- a/assets/batch_score/components/driver/tests/unit/common/header_providers/test_user_agent_header_provider.py
+++ b/assets/batch_score/components/driver/tests/unit/common/header_providers/test_user_agent_header_provider.py
@@ -32,7 +32,8 @@
      (f"{constants.BATCH_SCORE_USER_AGENT}:{TEST_COMPONENT_VERSION}"
       f"/Run:{TEST_RUN_ID}",
       None),
-     (f"{constants.BATCH_SCORE_USER_AGENT}:{TEST_COMPONENT_VERSION}/Run:{TEST_RUN_ID}/{TEST_UA_SEGMENT}", TEST_UA_SEGMENT),
+     (f"{constants.BATCH_SCORE_USER_AGENT}:{TEST_COMPONENT_VERSION}/Run:{TEST_RUN_ID}/{TEST_UA_SEGMENT}",
+      TEST_UA_SEGMENT),
      (f"{constants.BATCH_SCORE_USER_AGENT}:{TEST_COMPONENT_VERSION}/Run:{TEST_RUN_ID}", None),
 ])
 def test_get_headers(expected_user_agent_string, user_agent_segment):
diff --git a/assets/batch_score/components/driver/tests/unit/common/request_modification/test_input_type_modifier.py b/assets/batch_score/components/driver/tests/unit/common/request_modification/test_input_type_modifier.py
index 4f2b1c8a69..180fb84170 100644
--- a/assets/batch_score/components/driver/tests/unit/common/request_modification/test_input_type_modifier.py
+++ b/assets/batch_score/components/driver/tests/unit/common/request_modification/test_input_type_modifier.py
@@ -11,13 +11,46 @@ def test_get_input_type():
     """Test input type detection."""
     assert InputTypeModifier.get_input_type(request_obj={"messages": []}) == InputType.Unknown
     assert InputTypeModifier.get_input_type(request_obj={"no_messages": []}) == InputType.Unknown
-    assert InputTypeModifier.get_input_type(request_obj={"messages": [{"content": "just text"}, {"content": [{"type": "text"}]}]}) == InputType.TextOnly
-    assert InputTypeModifier.get_input_type(request_obj={"messages": [{"content": [{"type": "image"}, "just text"]}, {"content": [{"type": "text"}]}]}) == InputType.ImageAndText
-    assert InputTypeModifier.get_input_type(request_obj={"messages": [{"content": "just text"}, {"content": [{"type": "text"}]}, {"content": [{"type": "image"}]}]}) == InputType.ImageAndText
-    assert InputTypeModifier.get_input_type(request_obj={"messages": [{"content": [{"type": "image_url"}, "just text"]}, {"content": [{"type": "text"}]}]}) == InputType.ImageAndText
-    assert InputTypeModifier.get_input_type(request_obj={"messages": [{"content": "just text"}, {"content": [{"type": "text"}]}, {"content": [{"type": "image_url"}]}]}) == InputType.ImageAndText
-    assert InputTypeModifier.get_input_type(request_obj={"messages": [{"content": [{"type": "image"}]}, {"content": [{"type": "image_url"}]}]}) == InputType.Image
-    assert InputTypeModifier.get_input_type(request_obj={"messages": [{"content": [{"type": "image"}, {"type": "image_url"}]}]}) == InputType.Image
+    assert InputTypeModifier.get_input_type(request_obj={
+        "messages": [{"content": "just text"}, {"content": [{"type": "text"}]}]
+    }) == InputType.TextOnly
+    assert InputTypeModifier.get_input_type(request_obj={
+        "messages": [
+            {"content": [{"type": "image"}, "just text"]},
+            {"content": [{"type": "text"}]}
+        ]
+    }) == InputType.ImageAndText
+    assert InputTypeModifier.get_input_type(request_obj={
+        "messages": [
+            {"content": "just text"}, 
+            {"content": [{"type": "text"}]}, 
+            {"content": [{"type": "image"}]}
+        ]
+    }) == InputType.ImageAndText
+    assert InputTypeModifier.get_input_type(request_obj={
+        "messages": [
+            {"content": [{"type": "image_url"}, "just text"]},
+            {"content": [{"type": "text"}]}
+        ]
+    }) == InputType.ImageAndText
+    assert InputTypeModifier.get_input_type(request_obj={
+        "messages": [
+            {"content": "just text"}, 
+            {"content": [{"type": "text"}]},
+            {"content": [{"type": "image_url"}]}
+        ]
+    }) == InputType.ImageAndText
+    assert InputTypeModifier.get_input_type(request_obj={
+        "messages": [
+            {"content": [{"type": "image"}]},
+            {"content": [{"type": "image_url"}]}
+        ]
+    }) == InputType.Image
+    assert InputTypeModifier.get_input_type(request_obj={
+        "messages": [
+            {"content": [{"type": "image"}, {"type": "image_url"}]}
+        ]
+    }) == InputType.Image
 
 
 def test_modify():