feat: GenAI - Added support for GenerationConfig.response_modalities

Ark-kun · copybara-github · commit 78898fc70fc4 · 2024-12-13T16:49:09.000-08:00
PiperOrigin-RevId: 706042494
diff --git a/tests/unit/vertexai/test_generative_models.py b/tests/unit/vertexai/test_generative_models.py
@@ -748,6 +748,7 @@ def test_generate_content(
                 frequency_penalty=0.0,
                 logprobs=5,
                 response_logprobs=True,
+                response_modalities=["TEXT"],
             ),
             safety_settings=[
                 generative_models.SafetySetting(
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -1638,6 +1638,8 @@ class ResponseValidationError(ResponseBlockedError):
 class GenerationConfig:
     """Parameters for the generation."""
 
+    Modality = gapic_content_types.GenerationConfig.Modality
+
     def __init__(
         self,
         *,
@@ -1656,6 +1658,7 @@ def __init__(
         routing_config: Optional["RoutingConfig"] = None,
         logprobs: Optional[int] = None,
         response_logprobs: Optional[bool] = None,
+        response_modalities: Optional[List["GenerationConfig.Modality"]] = None,
     ):
         r"""Constructs a GenerationConfig object.
 
@@ -1725,6 +1728,7 @@ def __init__(
             audio_timestamp=audio_timestamp,
             logprobs=logprobs,
             response_logprobs=response_logprobs,
+            response_modalities=response_modalities,
         )
         if routing_config is not None:
             self._raw_generation_config.routing_config = (
