[ML] Apply yapf formatting and fix isort order for Anthropic inference

Run yapf with Beam's setup.cfg style config to fix all formatting
discrepancies. Fix import ordering (message_from_conversation before
message_from_string) in both test files to satisfy isort.

Author: bvolpato

sdks/python/apache_beam/ml/inference/anthropic_inference.py

@@ -121,7 +121,9 @@ def message_from_string(
     response = client.messages.create(
         model=model_name,
         max_tokens=max_tokens,
-        messages=[{"role": "user", "content": prompt}],
+        messages=[{
+            "role": "user", "content": prompt
+        }],
         **inference_args)
     responses.append(response)
   return responses
@@ -160,8 +162,8 @@ class AnthropicModelHandler(RemoteModelHandler[Any, PredictionResult,
   def __init__(
       self,
       model_name: str,
-      request_fn: Callable[
-          [str, Sequence[Any], Anthropic, dict[str, Any]], Any],
+      request_fn: Callable[[str, Sequence[Any], Anthropic, dict[str, Any]],
+                           Any],
       api_key: Optional[str] = None,
       *,
       system: Optional[Union[str, list[dict[str, str]]]] = None,
@@ -290,6 +292,5 @@ def request(
       inference_args['system'] = self.system
     if self.output_config is not None and 'output_config' not in inference_args:
       inference_args['output_config'] = self.output_config
-    responses = self.request_fn(
-        self.model_name, batch, model, inference_args)
+    responses = self.request_fn(self.model_name, batch, model, inference_args)
     return utils._convert_to_result(batch, responses, self.model_name)

sdks/python/apache_beam/ml/inference/anthropic_inference_it_test.py

@@ -25,8 +25,8 @@
 
 try:
   from apache_beam.ml.inference.anthropic_inference import AnthropicModelHandler
-  from apache_beam.ml.inference.anthropic_inference import message_from_string
   from apache_beam.ml.inference.anthropic_inference import message_from_conversation
+  from apache_beam.ml.inference.anthropic_inference import message_from_string
 except ImportError:
   raise unittest.SkipTest("Anthropic dependencies are not installed")
 
@@ -67,8 +67,7 @@ def test_anthropic_text_generation(self):
           p
           | beam.Create(prompts)
           | RunInference(handler)
-          | beam.Map(_extract_text)
-      )
+          | beam.Map(_extract_text))
       assert_that(results, is_not_empty())
 
   @unittest.skipIf(
@@ -84,9 +83,15 @@ def test_anthropic_conversation(self):
 
     conversations = [
         [
-            {"role": "user", "content": "What is 2 + 2?"},
-            {"role": "assistant", "content": "4"},
-            {"role": "user", "content": "Add 3 to that."},
+            {
+                "role": "user", "content": "What is 2 + 2?"
+            },
+            {
+                "role": "assistant", "content": "4"
+            },
+            {
+                "role": "user", "content": "Add 3 to that."
+            },
         ],
     ]
 
@@ -95,8 +100,7 @@ def test_anthropic_conversation(self):
           p
           | beam.Create(conversations)
           | RunInference(handler)
-          | beam.Map(_extract_text)
-      )
+          | beam.Map(_extract_text))
       assert_that(results, is_not_empty())
 
   @unittest.skipIf(
@@ -118,8 +122,7 @@ def test_anthropic_with_system_prompt(self):
           p
           | beam.Create(prompts)
           | RunInference(handler)
-          | beam.Map(_extract_text)
-      )
+          | beam.Map(_extract_text))
       assert_that(results, is_not_empty())
 
   @unittest.skipIf(
@@ -133,8 +136,7 @@ def test_anthropic_system_prompt_with_structured_output(self):
         system=(
             "You are a counting bot. When asked to count objects, convert "
             "responses such that numbers that are multiples of 3 are written "
-            "as 'Fizz' instead of the number."
-        ),
+            "as 'Fizz' instead of the number."),
         output_config={
             'format': {
                 'type': 'json_schema',
@@ -146,8 +148,12 @@ def test_anthropic_system_prompt_with_structured_output(self):
                             'items': {
                                 'type': 'object',
                                 'properties': {
-                                    'name': {'type': 'string'},
-                                    'count': {'type': 'string'},
+                                    'name': {
+                                        'type': 'string'
+                                    },
+                                    'count': {
+                                        'type': 'string'
+                                    },
                                 },
                                 'required': ['name', 'count'],
                                 'additionalProperties': False,
@@ -169,8 +175,7 @@ def test_anthropic_system_prompt_with_structured_output(self):
           p
           | beam.Create(prompts)
           | RunInference(handler)
-          | beam.Map(_extract_text)
-      )
+          | beam.Map(_extract_text))
 
       def verify_fizz(response_text):
         import json
@@ -187,24 +192,19 @@ def verify_fizz(response_text):
           if 'banana' in name:
             found_banana = True
             if count != 'Fizz':
-              raise ValueError(
-                  'Expected banana count Fizz, '
-                  'got %s' % count)
+              raise ValueError('Expected banana count Fizz, '
+                               'got %s' % count)
           elif 'orange' in name:
             found_orange = True
             if count != 'Fizz':
-              raise ValueError(
-                  'Expected orange count Fizz, '
-                  'got %s' % count)
+              raise ValueError('Expected orange count Fizz, '
+                               'got %s' % count)
           elif 'apple' in name:
             if count != '2':
-              raise ValueError(
-                  'Expected apple count 2, '
-                  'got %s' % count)
+              raise ValueError('Expected apple count 2, '
+                               'got %s' % count)
         if not found_banana or not found_orange:
-          raise ValueError(
-              'Missing expected items: %s'
-              % response_text)
+          raise ValueError('Missing expected items: %s' % response_text)
         return response_text
 
       _ = results | beam.Map(verify_fizz)

sdks/python/apache_beam/ml/inference/anthropic_inference_test.py

@@ -25,8 +25,8 @@
 
   from apache_beam.ml.inference.anthropic_inference import AnthropicModelHandler
   from apache_beam.ml.inference.anthropic_inference import _retry_on_appropriate_error
-  from apache_beam.ml.inference.anthropic_inference import message_from_string
   from apache_beam.ml.inference.anthropic_inference import message_from_conversation
+  from apache_beam.ml.inference.anthropic_inference import message_from_string
 except ImportError:
   raise unittest.SkipTest('Anthropic dependencies are not installed')
 
@@ -106,23 +106,24 @@ def test_sends_each_prompt(self):
         _make_fake_response("answer 1"),
         _make_fake_response("answer 2"),
     ]
-    results = message_from_string(
-        _TEST_MODEL, ['hello', 'world'], client, {})
+    results = message_from_string(_TEST_MODEL, ['hello', 'world'], client, {})
     self.assertEqual(len(results), 2)
     self.assertEqual(client.messages.create.call_count, 2)
 
     call_args = client.messages.create.call_args_list[0]
     self.assertEqual(call_args.kwargs['model'], _TEST_MODEL)
     self.assertEqual(
-        call_args.kwargs['messages'],
-        [{"role": "user", "content": "hello"}])
+        call_args.kwargs['messages'], [{
+            "role": "user", "content": "hello"
+        }])
 
   def test_passes_inference_args(self):
     client = mock.MagicMock()
     client.messages.create.return_value = _make_fake_response("ok")
     message_from_string(
-        _TEST_MODEL, ['test'], client,
-        {'max_tokens': 2048, 'temperature': 0.5})
+        _TEST_MODEL, ['test'], client, {
+            'max_tokens': 2048, 'temperature': 0.5
+        })
     call_args = client.messages.create.call_args
     self.assertEqual(call_args.kwargs['max_tokens'], 2048)
     self.assertEqual(call_args.kwargs['temperature'], 0.5)
@@ -140,10 +141,11 @@ def test_sends_conversation(self):
     client = mock.MagicMock()
     client.messages.create.return_value = _make_fake_response("Paris!")
     convo = [
-        {"role": "user", "content": "What is the capital of France?"},
+        {
+            "role": "user", "content": "What is the capital of France?"
+        },
     ]
-    results = message_from_conversation(
-        _TEST_MODEL, [convo], client, {})
+    results = message_from_conversation(_TEST_MODEL, [convo], client, {})
     self.assertEqual(len(results), 1)
     call_args = client.messages.create.call_args
     self.assertEqual(call_args.kwargs['messages'], convo)
@@ -162,16 +164,13 @@ def test_create_client_with_api_key(self, mock_anthropic):
   @mock.patch('apache_beam.ml.inference.anthropic_inference.Anthropic')
   def test_create_client_from_env(self, mock_anthropic):
     handler = AnthropicModelHandler(
-        model_name=_TEST_MODEL,
-        request_fn=message_from_string)
+        model_name=_TEST_MODEL, request_fn=message_from_string)
     handler.create_client()
     mock_anthropic.assert_called_once_with()
 
   def test_request_returns_prediction_results(self):
     handler = AnthropicModelHandler(
-        model_name=_TEST_MODEL,
-        request_fn=message_from_string,
-        api_key='fake')
+        model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake')
     mock_client = mock.MagicMock()
     resp1 = _make_fake_response("answer 1")
     resp2 = _make_fake_response("answer 2")
@@ -201,8 +200,10 @@ def test_batch_elements_kwargs(self):
 
 def _fake_request_fn(model_name, batch, client, inference_args):
   """A picklable request function that returns fake responses."""
-  return [FakeMessage(content=[FakeContentBlock(text=f'answer for: {p}')])
-          for p in batch]
+  return [
+      FakeMessage(content=[FakeContentBlock(text=f'answer for: {p}')])
+      for p in batch
+  ]
 
 
 class SystemPromptTest(unittest.TestCase):
@@ -229,18 +230,14 @@ def test_system_prompt_not_overridden_by_handler(self):
     mock_client = mock.MagicMock()
     mock_client.messages.create.return_value = _make_fake_response("ok")
 
-    handler.request(
-        ['test'], mock_client,
-        {'system': 'Per-request override.'})
+    handler.request(['test'], mock_client, {'system': 'Per-request override.'})
 
     call_args = mock_client.messages.create.call_args
     self.assertEqual(call_args.kwargs['system'], 'Per-request override.')
 
   def test_no_system_prompt_when_none(self):
     handler = AnthropicModelHandler(
-        model_name=_TEST_MODEL,
-        request_fn=message_from_string,
-        api_key='fake')
+        model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake')
     mock_client = mock.MagicMock()
     mock_client.messages.create.return_value = _make_fake_response("ok")
 
@@ -256,7 +253,11 @@ class OutputConfigTest(unittest.TestCase):
           'type': 'json_schema',
           'schema': {
               'type': 'object',
-              'properties': {'answer': {'type': 'string'}},
+              'properties': {
+                  'answer': {
+                      'type': 'string'
+                  }
+              },
               'required': ['answer'],
               'additionalProperties': False,
           },
@@ -295,9 +296,7 @@ def test_output_config_not_overridden_by_handler(self):
 
   def test_no_output_config_when_none(self):
     handler = AnthropicModelHandler(
-        model_name=_TEST_MODEL,
-        request_fn=message_from_string,
-        api_key='fake')
+        model_name=_TEST_MODEL, request_fn=message_from_string, api_key='fake')
     mock_client = mock.MagicMock()
     mock_client.messages.create.return_value = _make_fake_response("ok")
 
@@ -324,8 +323,7 @@ def test_pipeline_e2e(self):
           p
           | beam.Create(prompts)
           | RunInference(handler)
-          | beam.Map(lambda r: r.example)
-      )
+          | beam.Map(lambda r: r.example))
       assert_that(results, equal_to(prompts))
 
   def test_pipeline_with_system_prompt(self):
@@ -345,8 +343,7 @@ def test_pipeline_with_system_prompt(self):
           p
           | beam.Create(prompts)
           | RunInference(handler)
-          | beam.Map(lambda r: r.example)
-      )
+          | beam.Map(lambda r: r.example))
       assert_that(results, equal_to(prompts))