feat: GenAI SDK client(multimodal) - Add single_turn_template helper to GeminiRequestReadConfig.

PiperOrigin-RevId: 890422497

Author: cleop-google

tests/unit/vertexai/genai/test_multimodal_datasets_genai.py

@@ -15,6 +15,7 @@
 """Tests for multimodal datasets."""
 
 from vertexai._genai import types
+from google.genai import types as genai_types
 
 
 class TestMultimodalDataset:
@@ -126,3 +127,59 @@ def test_set_bigquery_uri_preserves_other_fields(self):
             dataset.metadata.gemini_request_read_config.assembled_request_column_name
             == "test_column"
         )
+
+
+class TestGeminiRequestReadConfig:
+    def test_single_turn_template(self):
+        read_config = types.GeminiRequestReadConfig.single_turn_template(
+            model="gemini-1.5-flash",
+            prompt="test_prompt",
+            response="test_response",
+            system_instruction="test_system_instruction",
+            cached_content="test_cached_content",
+            tools={"function_declarations": [{"name": "test_tool"}]},
+            tool_config={"function_calling_config": {"mode": "ANY"}},
+            safety_settings={"category": "HARM_CATEGORY_DANGEROUS_CONTENT"},
+            generation_config={"temperature": 0.5},
+            field_mapping={"test_placeholder": "test_column"},
+        )
+
+        expected_read_config = types.GeminiRequestReadConfig(
+            template_config=types.GeminiTemplateConfig(
+                gemini_example=types.GeminiExample(
+                    model="gemini-1.5-flash",
+                    contents=[
+                        genai_types.Content(
+                            role="user",
+                            parts=[genai_types.Part.from_text(text="test_prompt")],
+                        ),
+                        genai_types.Content(
+                            role="model",
+                            parts=[genai_types.Part.from_text(text="test_response")],
+                        ),
+                    ],
+                    system_instruction=genai_types.Content(
+                        parts=[
+                            genai_types.Part.from_text(text="test_system_instruction")
+                        ],
+                    ),
+                    cached_content="test_cached_content",
+                    tools=genai_types.Tool(
+                        function_declarations=[
+                            genai_types.FunctionDeclaration(name="test_tool")
+                        ]
+                    ),
+                    tool_config=genai_types.ToolConfig(
+                        function_calling_config=genai_types.FunctionCallingConfig(
+                            mode="ANY"
+                        )
+                    ),
+                    safety_settings=genai_types.SafetySetting(
+                        category="HARM_CATEGORY_DANGEROUS_CONTENT"
+                    ),
+                    generation_config=genai_types.GenerationConfig(temperature=0.5),
+                ),
+                field_mapping={"test_placeholder": "test_column"},
+            ),
+        )
+        assert read_config == expected_read_config

vertexai/_genai/types/common.py

@@ -11873,6 +11873,104 @@ class GeminiRequestReadConfig(_common.BaseModel):
         description="""Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""",
     )
 
+    @classmethod
+    def single_turn_template(
+        cls,
+        *,
+        prompt: Optional[str] = None,
+        response: Optional[str] = None,
+        system_instruction: Optional[str] = None,
+        model: Optional[str] = None,
+        cached_content: Optional[str] = None,
+        tools: Optional[list[Union[genai_types.Tool, dict[str, Any]]]] = None,
+        tool_config: Optional[Union[genai_types.ToolConfig, dict[str, Any]]] = None,
+        safety_settings: Optional[
+            list[Union[genai_types.SafetySetting, dict[str, Any]]]
+        ] = None,
+        generation_config: Optional[
+            Union[genai_types.GenerationConfig, dict[str, Any]]
+        ] = None,
+        field_mapping: Optional[dict[str, str]] = None,
+    ) -> "GeminiRequestReadConfig":
+        """Constructs a GeminiRequestReadConfig object for single-turn cases.
+
+        Example:
+            template_config = datasets.construct_single_turn_template(
+                    prompt = "Which flower is this {flower_image} ?",
+                    response="This is a {label}.",
+                    system_instruction="You are a botanical classifier."
+            )
+
+        Args:
+
+            prompt (str):
+                Required. User input.
+            response (str):
+                Optional. Model response to user input.
+            system_instruction (str):
+                Optional. System instructions for the model.
+            model (str):
+                Optional. The model to use for the GeminiExample.
+            cached_content (str):
+                Optional. The cached content to use for the GeminiExample.
+            tools (List[Tool]):
+                Optional. The tools to use for the GeminiExample.
+            tool_config (ToolConfig):
+                Optional. The tool config to use for the GeminiExample.
+            safety_settings (List[SafetySetting]):
+                Optional. The safety settings to use for the GeminiExample.
+            generation_config (GenerationConfig):
+                Optional. The generation config to use for the GeminiExample.
+            field_mapping (dict[str, str]):
+                Optional. Mapping of placeholders to dataset columns.
+
+        Returns:
+            A GeminiRequestReadConfig object.
+        """
+        contents = []
+        if prompt:
+            contents.append(
+                genai_types.Content(
+                    role="user",
+                    parts=[
+                        genai_types.Part.from_text(text=prompt),
+                    ],
+                )
+            )
+        if response:
+            contents.append(
+                genai_types.Content(
+                    role="model",
+                    parts=[
+                        genai_types.Part.from_text(text=response),
+                    ],
+                )
+            )
+
+        sys_inst = None
+        if system_instruction:
+            sys_inst = genai_types.Content(
+                parts=[
+                    genai_types.Part.from_text(text=system_instruction),
+                ],
+            )
+
+        return cls(
+            template_config=GeminiTemplateConfig(
+                gemini_example=GeminiExample(
+                    model=model,
+                    contents=contents,
+                    system_instruction=sys_inst,
+                    cached_content=cached_content,
+                    tools=tools,
+                    tool_config=tool_config,
+                    safety_settings=safety_settings,
+                    generation_config=generation_config,
+                ),
+                field_mapping=field_mapping,
+            ),
+        )
+
 
 class GeminiRequestReadConfigDict(TypedDict, total=False):
     """Represents the config for reading Gemini requests."""