Skip to content

Commit 0e5037d

Browse files
cleop-googlecopybara-github
cleop-google
authored and
copybara-github
committed
feat: GenAI SDK client(multimodal) - Add single_turn_template helper to GeminiRequestReadConfig.
PiperOrigin-RevId: 896573028
1 parent 8016cb6 commit 0e5037d

2 files changed

Lines changed: 147 additions & 0 deletions

File tree

‎tests/unit/vertexai/genai/test_multimodal_datasets_genai.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
from vertexai._genai import _datasets_utils
1919
from vertexai._genai import types
20+
from google.genai import types as genai_types
2021
import pytest
2122

2223

@@ -155,3 +156,63 @@ def test_to_bigframes(self, mock_import_bigframes):
155156
mock_import_bigframes.return_value.pandas.read_gbq_table.assert_called_once_with(
156157
"project.dataset.table"
157158
)
159+
160+
161+
class TestGeminiRequestReadConfig:
162+
def test_single_turn_template(self):
163+
read_config = types.GeminiRequestReadConfig.single_turn_template(
164+
model="gemini-1.5-flash",
165+
prompt="test_prompt",
166+
response="test_response",
167+
system_instruction="test_system_instruction",
168+
cached_content="test_cached_content",
169+
tools=[{"function_declarations": [{"name": "test_tool"}]}],
170+
tool_config={"function_calling_config": {"mode": "ANY"}},
171+
safety_settings=[{"category": "HARM_CATEGORY_DANGEROUS_CONTENT"}],
172+
generation_config={"temperature": 0.5},
173+
field_mapping={"test_placeholder": "test_column"},
174+
)
175+
176+
expected_read_config = types.GeminiRequestReadConfig(
177+
template_config=types.GeminiTemplateConfig(
178+
gemini_example=types.GeminiExample(
179+
model="gemini-1.5-flash",
180+
contents=[
181+
genai_types.Content(
182+
role="user",
183+
parts=[genai_types.Part.from_text(text="test_prompt")],
184+
),
185+
genai_types.Content(
186+
role="model",
187+
parts=[genai_types.Part.from_text(text="test_response")],
188+
),
189+
],
190+
system_instruction=genai_types.Content(
191+
parts=[
192+
genai_types.Part.from_text(text="test_system_instruction")
193+
],
194+
),
195+
cached_content="test_cached_content",
196+
tools=[
197+
genai_types.Tool(
198+
function_declarations=[
199+
genai_types.FunctionDeclaration(name="test_tool")
200+
]
201+
)
202+
],
203+
tool_config=genai_types.ToolConfig(
204+
function_calling_config=genai_types.FunctionCallingConfig(
205+
mode="ANY"
206+
)
207+
),
208+
safety_settings=[
209+
genai_types.SafetySetting(
210+
category="HARM_CATEGORY_DANGEROUS_CONTENT"
211+
)
212+
],
213+
generation_config=genai_types.GenerationConfig(temperature=0.5),
214+
),
215+
field_mapping={"test_placeholder": "test_column"},
216+
),
217+
)
218+
assert read_config == expected_read_config

‎vertexai/_genai/types/common.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12583,6 +12583,92 @@ class GeminiRequestReadConfig(_common.BaseModel):
1258312583
description="""Column name in the underlying BigQuery table that contains already fully assembled Gemini requests.""",
1258412584
)
1258512585

12586+
@classmethod
12587+
def single_turn_template(
12588+
cls,
12589+
*,
12590+
prompt: str,
12591+
response: Optional[str] = None,
12592+
system_instruction: Optional[str] = None,
12593+
model: Optional[str] = None,
12594+
cached_content: Optional[str] = None,
12595+
tools: Optional[list[Union[genai_types.Tool, dict[str, Any]]]] = None,
12596+
tool_config: Optional[Union[genai_types.ToolConfig, dict[str, Any]]] = None,
12597+
safety_settings: Optional[
12598+
list[Union[genai_types.SafetySetting, dict[str, Any]]]
12599+
] = None,
12600+
generation_config: Optional[
12601+
Union[genai_types.GenerationConfig, dict[str, Any]]
12602+
] = None,
12603+
field_mapping: Optional[dict[str, str]] = None,
12604+
) -> "GeminiRequestReadConfig":
12605+
"""Constructs a GeminiRequestReadConfig object for single-turn cases.
12606+
12607+
Example:
12608+
read_config = GeminiRequestReadConfig.single_turn_template(
12609+
prompt="Which flower is this {flower_image}?",
12610+
response="This is a {label}.",
12611+
system_instruction="You are a botanical classifier."
12612+
)
12613+
12614+
Args:
12615+
prompt: Required. User input.
12616+
response: Optional. Model response to user input.
12617+
system_instruction: Optional. System instructions for the model.
12618+
model: Optional. The model to use for the GeminiExample.
12619+
cached_content: Optional. The cached content to use for the GeminiExample.
12620+
tools: Optional. The tools to use for the GeminiExample.
12621+
tool_config: Optional. The tool config to use for the GeminiExample.
12622+
safety_settings: Optional. The safety settings to use for the GeminiExample.
12623+
generation_config: Optional. The generation config to use for the GeminiExample.
12624+
field_mapping: Optional. Mapping of placeholders to dataset columns.
12625+
12626+
Returns:
12627+
A GeminiRequestReadConfig object.
12628+
"""
12629+
contents = []
12630+
contents.append(
12631+
genai_types.Content(
12632+
role="user",
12633+
parts=[
12634+
genai_types.Part.from_text(text=prompt),
12635+
],
12636+
)
12637+
)
12638+
if response:
12639+
contents.append(
12640+
genai_types.Content(
12641+
role="model",
12642+
parts=[
12643+
genai_types.Part.from_text(text=response),
12644+
],
12645+
)
12646+
)
12647+
12648+
system_instruction_content = None
12649+
if system_instruction:
12650+
system_instruction_content = genai_types.Content(
12651+
parts=[
12652+
genai_types.Part.from_text(text=system_instruction),
12653+
],
12654+
)
12655+
12656+
return cls(
12657+
template_config=GeminiTemplateConfig(
12658+
gemini_example=GeminiExample(
12659+
model=model,
12660+
contents=contents,
12661+
system_instruction=system_instruction_content,
12662+
cached_content=cached_content,
12663+
tools=tools,
12664+
tool_config=tool_config,
12665+
safety_settings=safety_settings,
12666+
generation_config=generation_config,
12667+
),
12668+
field_mapping=field_mapping,
12669+
),
12670+
)
12671+
1258612672

1258712673
class GeminiRequestReadConfigDict(TypedDict, total=False):
1258812674
"""Represents the config for reading Gemini requests."""

0 commit comments

Comments
 (0)