feat: GenAI - Added response_style to GenerationConfig

PiperOrigin-RevId: 630462525
googleapis · May 3, 2024 · 7279dab · 7279dab
1 parent 3c3727b
commit 7279dab
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 0 deletions.
diff --git a/google/cloud/aiplatform_v1beta1/types/content.py b/google/cloud/aiplatform_v1beta1/types/content.py
@@ -309,8 +309,30 @@ class GenerationConfig(proto.Message):
                The model needs to be prompted to output the appropriate
                response type, otherwise the behavior is undefined. This
                is a preview feature.
+
+        response_style (google.cloud.aiplatform_v1beta1.types.GenerationConfig.ResponseStyle):
+            Control Three levels of creativity in the model output.
+            Default: RESPONSE_STYLE_BALANCED
     """
 
+    class ResponseStyle(proto.Enum):
+        r"""Choices of the response style.
+
+        Values:
+            RESPONSE_STYLE_UNSPECIFIED (0):
+                response style unspecified.
+            RESPONSE_STYLE_PRECISE (1):
+                Precise response.
+            RESPONSE_STYLE_BALANCED (2):
+                Default response style.
+            RESPONSE_STYLE_CREATIVE (3):
+                Creative response style.
+        """
+        RESPONSE_STYLE_UNSPECIFIED = 0
+        RESPONSE_STYLE_PRECISE = 1
+        RESPONSE_STYLE_BALANCED = 2
+        RESPONSE_STYLE_CREATIVE = 3
+
     temperature: float = proto.Field(
         proto.FLOAT,
         number=1,
@@ -354,6 +376,11 @@ class GenerationConfig(proto.Message):
         proto.STRING,
         number=13,
     )
+    response_style: ResponseStyle = proto.Field(
+        proto.ENUM,
+        number=14,
+        enum=ResponseStyle,
+    )
 
 
 class SafetySetting(proto.Message):

diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -1181,6 +1181,7 @@ class ResponseValidationError(ResponseBlockedError):
 
 class GenerationConfig:
     """Parameters for the generation."""
+    ResponseStyle = gapic_content_types.GenerationConfig.ResponseStyle
 
     def __init__(
         self,
@@ -1194,6 +1195,7 @@ def __init__(
         presence_penalty: Optional[float] = None,
         frequency_penalty: Optional[float] = None,
         response_mime_type: Optional[str] = None,
+        response_style: Optional["GenerationConfig.ResponseStyle"] = None,
     ):
         r"""Constructs a GenerationConfig object.
 
@@ -1216,6 +1218,7 @@ def __init__(
 
                 The model needs to be prompted to output the appropriate
                 response type, otherwise the behavior is undefined.
+            response_style: Control three levels of creativity in the model output.
 
         Usage:
             ```
@@ -1228,6 +1231,7 @@ def __init__(
                     candidate_count=1,
                     max_output_tokens=100,
                     stop_sequences=["\n\n\n"],
+                    response_style=ResponseStyle.RESPONSE_STYLE_PRECISE,
                 )
             )
             ```
@@ -1242,6 +1246,7 @@ def __init__(
             presence_penalty=presence_penalty,
             frequency_penalty=frequency_penalty,
             response_mime_type=response_mime_type,
+            response_style=response_style,
         )
 
     @classmethod