[ML] Add AnthropicModelHandler for direct Claude inference

Add ModelHandler for Anthropic Claude models using the Messages API.
This enables direct inference without requiring Vertex AI/GCP setup.

Author: bvolpato

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

@@ -0,0 +1,234 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""A ModelHandler for Anthropic Claude models using the Messages API.
+
+This module provides an integration between Apache Beam's RunInference
+transform and Anthropic's Claude models, enabling batch inference in
+Beam pipelines.
+
+Example usage::
+
+  from apache_beam.ml.inference.anthropic_inference import (
+      AnthropicModelHandler,
+      message_from_string,
+  )
+  from apache_beam.ml.inference.base import RunInference
+
+  model_handler = AnthropicModelHandler(
+      model_name='claude-haiku-4-5',
+      api_key='your-api-key',
+      request_fn=message_from_string,
+  )
+
+  with beam.Pipeline() as p:
+    results = (
+        p
+        | beam.Create(['What is Apache Beam?', 'Explain MapReduce.'])
+        | RunInference(model_handler)
+    )
+"""
+
+import logging
+from collections.abc import Callable
+from collections.abc import Iterable
+from collections.abc import Sequence
+from typing import Any
+from typing import Optional
+
+from anthropic import Anthropic
+from anthropic import APIStatusError
+
+from apache_beam.ml.inference import utils
+from apache_beam.ml.inference.base import PredictionResult
+from apache_beam.ml.inference.base import RemoteModelHandler
+
+LOGGER = logging.getLogger("AnthropicModelHandler")
+
+
+def _retry_on_appropriate_error(exception: Exception) -> bool:
+  """Retry filter that returns True for 5xx errors or 429 (rate limiting).
+
+  Args:
+    exception: the exception encountered during the request/response loop.
+
+  Returns:
+    True if the exception is retriable (429 or 5xx), False otherwise.
+  """
+  if not isinstance(exception, APIStatusError):
+    return False
+  return exception.status_code == 429 or exception.status_code >= 500
+
+
+def message_from_string(
+    model_name: str,
+    batch: Sequence[str],
+    client: Anthropic,
+    inference_args: dict[str, Any]):
+  """Request function that sends string prompts to Claude as user messages.
+
+  Each string in the batch is sent as a separate request. The results are
+  returned as a list of response objects.
+
+  Args:
+    model_name: the Claude model to use (e.g. 'claude-haiku-4-5').
+    batch: the string inputs to send to Claude.
+    client: the Anthropic client instance.
+    inference_args: additional arguments passed to the messages.create call.
+      Common args include 'max_tokens', 'system', 'temperature', 'top_p'.
+  """
+  max_tokens = inference_args.pop('max_tokens', 1024)
+  responses = []
+  for prompt in batch:
+    response = client.messages.create(
+        model=model_name,
+        max_tokens=max_tokens,
+        messages=[{"role": "user", "content": prompt}],
+        **inference_args)
+    responses.append(response)
+  return responses
+
+
+def message_from_conversation(
+    model_name: str,
+    batch: Sequence[list[dict[str, str]]],
+    client: Anthropic,
+    inference_args: dict[str, Any]):
+  """Request function that sends multi-turn conversations to Claude.
+
+  Each element in the batch is a list of message dicts with 'role' and
+  'content' keys, representing a multi-turn conversation.
+
+  Args:
+    model_name: the Claude model to use.
+    batch: a sequence of conversations (each a list of message dicts).
+    client: the Anthropic client instance.
+    inference_args: additional arguments passed to the messages.create call.
+  """
+  max_tokens = inference_args.pop('max_tokens', 1024)
+  responses = []
+  for messages in batch:
+    response = client.messages.create(
+        model=model_name,
+        max_tokens=max_tokens,
+        messages=messages,
+        **inference_args)
+    responses.append(response)
+  return responses
+
+
+class AnthropicModelHandler(RemoteModelHandler[Any, PredictionResult,
+                                               Anthropic]):
+  def __init__(
+      self,
+      model_name: str,
+      request_fn: Callable[
+          [str, Sequence[Any], Anthropic, dict[str, Any]], Any],
+      api_key: Optional[str] = None,
+      *,
+      min_batch_size: Optional[int] = None,
+      max_batch_size: Optional[int] = None,
+      max_batch_duration_secs: Optional[int] = None,
+      max_batch_weight: Optional[int] = None,
+      element_size_fn: Optional[Callable[[Any], int]] = None,
+      **kwargs):
+    """Implementation of the ModelHandler interface for Anthropic Claude.
+
+    **NOTE:** This API and its implementation are under development and
+    do not provide backward compatibility guarantees.
+
+    This handler connects to the Anthropic Messages API to run inference
+    using Claude models. It supports text generation from string prompts
+    or multi-turn conversations.
+
+    Args:
+      model_name: the Claude model to send requests to (e.g.
+        'claude-sonnet-4-6', 'claude-haiku-4-5').
+      request_fn: the function to use to send requests. Should take the
+        model name, batch, client, and inference_args and return the
+        responses from Claude. Built-in options are message_from_string
+        and message_from_conversation.
+      api_key: the Anthropic API key. If not provided, the client will
+        look for the ANTHROPIC_API_KEY environment variable.
+      min_batch_size: optional. the minimum batch size to use when
+        batching inputs.
+      max_batch_size: optional. the maximum batch size to use when
+        batching inputs.
+      max_batch_duration_secs: optional. the maximum amount of time to
+        buffer a batch before emitting; used in streaming contexts.
+      max_batch_weight: optional. the maximum total weight of a batch.
+      element_size_fn: optional. a function that returns the size
+        (weight) of an element.
+    """
+    self._batching_kwargs = {}
+    self._env_vars = kwargs.get('env_vars', {})
+    if min_batch_size is not None:
+      self._batching_kwargs["min_batch_size"] = min_batch_size
+    if max_batch_size is not None:
+      self._batching_kwargs["max_batch_size"] = max_batch_size
+    if max_batch_duration_secs is not None:
+      self._batching_kwargs["max_batch_duration_secs"] = max_batch_duration_secs
+    if max_batch_weight is not None:
+      self._batching_kwargs["max_batch_weight"] = max_batch_weight
+    if element_size_fn is not None:
+      self._batching_kwargs['element_size_fn'] = element_size_fn
+
+    self.model_name = model_name
+    self.request_fn = request_fn
+    self.api_key = api_key
+
+    super().__init__(
+        namespace='AnthropicModelHandler',
+        retry_filter=_retry_on_appropriate_error,
+        **kwargs)
+
+  def batch_elements_kwargs(self):
+    return self._batching_kwargs
+
+  def create_client(self) -> Anthropic:
+    """Creates the Anthropic client used to send requests.
+
+    If api_key was provided at construction time, it is used directly.
+    Otherwise, the client will fall back to the ANTHROPIC_API_KEY
+    environment variable.
+    """
+    if self.api_key:
+      return Anthropic(api_key=self.api_key)
+    return Anthropic()
+
+  def request(
+      self,
+      batch: Sequence[Any],
+      model: Anthropic,
+      inference_args: Optional[dict[str, Any]] = None
+  ) -> Iterable[PredictionResult]:
+    """Sends a prediction request to the Anthropic API.
+
+    Args:
+      batch: a sequence of inputs to be passed to the request function.
+      model: an Anthropic client instance.
+      inference_args: additional arguments to send as part of the
+        prediction request (e.g. max_tokens, temperature, system).
+
+    Returns:
+      An iterable of PredictionResults.
+    """
+    if inference_args is None:
+      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

@@ -0,0 +1,107 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""End-to-End test for Anthropic Claude Remote Inference"""
+
+import logging
+import os
+import unittest
+
+import pytest
+
+try:
+  from anthropic import Anthropic
+
+  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
+except ImportError:
+  raise unittest.SkipTest("Anthropic dependencies are not installed")
+
+import apache_beam as beam
+from apache_beam.ml.inference.base import RunInference
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import is_not_empty
+
+_ANTHROPIC_API_KEY = os.environ.get('ANTHROPIC_API_KEY', None)
+_TEST_MODEL = 'claude-haiku-4-5'
+
+
+def _extract_text(prediction_result):
+  return prediction_result.inference.content[0].text
+
+
+@pytest.mark.anthropic_postcommit
+class AnthropicInferenceIT(unittest.TestCase):
+  @unittest.skipIf(
+      _ANTHROPIC_API_KEY is None,
+      'ANTHROPIC_API_KEY environment variable is not set')
+  def test_anthropic_text_generation(self):
+    handler = AnthropicModelHandler(
+        model_name=_TEST_MODEL,
+        request_fn=message_from_string,
+        api_key=_ANTHROPIC_API_KEY,
+        max_batch_size=1,
+    )
+
+    prompts = [
+        'What is Apache Beam in one sentence?',
+        'Name three distributed computing frameworks.',
+    ]
+
+    with TestPipeline() as p:
+      results = (
+          p
+          | beam.Create(prompts)
+          | RunInference(handler)
+          | beam.Map(_extract_text)
+      )
+      assert_that(results, is_not_empty())
+
+  @unittest.skipIf(
+      _ANTHROPIC_API_KEY is None,
+      'ANTHROPIC_API_KEY environment variable is not set')
+  def test_anthropic_conversation(self):
+    handler = AnthropicModelHandler(
+        model_name=_TEST_MODEL,
+        request_fn=message_from_conversation,
+        api_key=_ANTHROPIC_API_KEY,
+        max_batch_size=1,
+    )
+
+    conversations = [
+        [
+            {"role": "user", "content": "What is 2 + 2?"},
+            {"role": "assistant", "content": "4"},
+            {"role": "user", "content": "Add 3 to that."},
+        ],
+    ]
+
+    with TestPipeline() as p:
+      results = (
+          p
+          | beam.Create(conversations)
+          | RunInference(handler)
+          | beam.Map(_extract_text)
+      )
+      assert_that(results, is_not_empty())
+
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.DEBUG)
+  unittest.main()