feat: GenAI - Added Anthropic models support in GenAI batch prediction

PiperOrigin-RevId: 688995187

Author: vertex-sdk-bot

tests/unit/vertexai/test_batch_prediction.py

@@ -52,6 +52,10 @@
 _TEST_PALM_MODEL_RESOURCE_NAME = f"publishers/google/models/{_TEST_PALM_MODEL_NAME}"
 _TEST_LLAMA_MODEL_NAME = "llama3-405b-instruct-maas"
 _TEST_LLAMA_MODEL_RESOURCE_NAME = f"publishers/meta/models/{_TEST_LLAMA_MODEL_NAME}"
+_TEST_CLAUDE_MODEL_NAME = "claude-3-opus"
+_TEST_CLAUDE_MODEL_RESOURCE_NAME = (
+    f"publishers/anthropic/models/{_TEST_CLAUDE_MODEL_NAME}"
+)
 
 _TEST_GCS_INPUT_URI = "gs://test-bucket/test-input.jsonl"
 _TEST_GCS_INPUT_URI_2 = "gs://test-bucket/test-input-2.jsonl"
@@ -146,6 +150,23 @@ def get_batch_prediction_job_with_llama_model_mock():
         yield get_job_mock
 
 
+@pytest.fixture
+def get_batch_prediction_job_with_claude_model_mock():
+    with mock.patch.object(
+        job_service_client.JobServiceClient, "get_batch_prediction_job"
+    ) as get_job_mock:
+        get_job_mock.return_value = gca_batch_prediction_job_compat.BatchPredictionJob(
+            name=_TEST_BATCH_PREDICTION_JOB_NAME,
+            display_name=_TEST_DISPLAY_NAME,
+            model=_TEST_CLAUDE_MODEL_RESOURCE_NAME,
+            state=_TEST_JOB_STATE_SUCCESS,
+            output_info=gca_batch_prediction_job_compat.BatchPredictionJob.OutputInfo(
+                gcs_output_directory=_TEST_GCS_OUTPUT_PREFIX
+            ),
+        )
+        yield get_job_mock
+
+
 @pytest.fixture
 def get_batch_prediction_job_with_tuned_gemini_model_mock():
     with mock.patch.object(
@@ -281,6 +302,16 @@ def test_init_batch_prediction_job_with_llama_model(
             name=_TEST_BATCH_PREDICTION_JOB_NAME, retry=aiplatform_base._DEFAULT_RETRY
         )
 
+    def test_init_batch_prediction_job_with_claude_model(
+        self,
+        get_batch_prediction_job_with_claude_model_mock,
+    ):
+        batch_prediction.BatchPredictionJob(_TEST_BATCH_PREDICTION_JOB_ID)
+
+        get_batch_prediction_job_with_claude_model_mock.assert_called_once_with(
+            name=_TEST_BATCH_PREDICTION_JOB_NAME, retry=aiplatform_base._DEFAULT_RETRY
+        )
+
     def test_init_batch_prediction_job_with_tuned_gemini_model(
         self,
         get_batch_prediction_job_with_tuned_gemini_model_mock,
@@ -509,6 +540,39 @@ def test_submit_batch_prediction_job_with_llama_model(
             timeout=None,
         )
 
+    def test_submit_batch_prediction_job_with_claude_model(
+        self,
+        create_batch_prediction_job_mock,
+    ):
+        job = batch_prediction.BatchPredictionJob.submit(
+            source_model=_TEST_CLAUDE_MODEL_RESOURCE_NAME,
+            input_dataset=_TEST_BQ_INPUT_URI,
+        )
+
+        assert job.gca_resource == _TEST_GAPIC_BATCH_PREDICTION_JOB
+
+        expected_gapic_batch_prediction_job = gca_batch_prediction_job_compat.BatchPredictionJob(
+            display_name=_TEST_DISPLAY_NAME,
+            model=_TEST_CLAUDE_MODEL_RESOURCE_NAME,
+            input_config=gca_batch_prediction_job_compat.BatchPredictionJob.InputConfig(
+                instances_format="bigquery",
+                bigquery_source=gca_io_compat.BigQuerySource(
+                    input_uri=_TEST_BQ_INPUT_URI
+                ),
+            ),
+            output_config=gca_batch_prediction_job_compat.BatchPredictionJob.OutputConfig(
+                bigquery_destination=gca_io_compat.BigQueryDestination(
+                    output_uri=_TEST_BQ_OUTPUT_PREFIX
+                ),
+                predictions_format="bigquery",
+            ),
+        )
+        create_batch_prediction_job_mock.assert_called_once_with(
+            parent=_TEST_PARENT,
+            batch_prediction_job=expected_gapic_batch_prediction_job,
+            timeout=None,
+        )
+
     @pytest.mark.usefixtures("create_batch_prediction_job_mock")
     def test_submit_batch_prediction_job_with_tuned_model(
         self,

vertexai/batch_prediction/_batch_prediction.py

@@ -34,6 +34,7 @@
 
 _GEMINI_MODEL_PATTERN = r"publishers/google/models/gemini"
 _LLAMA_MODEL_PATTERN = r"publishers/meta/models/llama"
+_CLAUDE_MODEL_PATTERN = r"publishers/anthropic/models/claude"
 _GEMINI_TUNED_MODEL_PATTERN = r"^projects/[0-9]+?/locations/[0-9a-z-]+?/models/[0-9]+?$"
 
 
@@ -287,6 +288,7 @@ def _reconcile_model_name(cls, model_name: str) -> str:
             # publisher model full name
             not model_name.startswith("publishers/google/models/")
             and not model_name.startswith("publishers/meta/models/")
+            and not model_name.startswith("publishers/anthropic/models/")
             # tuned model full resource name
             and not re.search(_GEMINI_TUNED_MODEL_PATTERN, model_name)
         ):
@@ -314,6 +316,10 @@ def _is_genai_model(cls, model_name: str) -> bool:
             # Model is a Llama3 model.
             return True
 
+        if re.search(_CLAUDE_MODEL_PATTERN, model_name):
+            # Model is a claude model.
+            return True
+
         return False
 
     @classmethod