Skip to content

feat: GenAI SDK client - Enabling Few-shot Prompt Optimization by passing either "OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS" or "OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE" to the optimize_prompt method #6164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+422 −30
Open

feat: GenAI SDK client - Enabling Few-shot Prompt Optimization by passing either "OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS" or "OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE" to the optimize_prompt method #6164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 60 additions & 0 deletions tests/unit/vertexai/genai/replays/test_prompt_optimizer_async_optimize_prompt_return_type.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

from tests.unit.vertexai.genai.replays import pytest_helper
from vertexai._genai import types
import pandas as pd
import pytest


Expand All @@ -32,6 +33,65 @@ async def test_optimize_prompt(client):
assert response.raw_text_response


@pytest.mark.asyncio
async def test_optimize_prompt_w_optimization_target(client):
"""Tests the optimize request parameters method with optimization target."""
test_prompt = "Generate system instructions for analyzing medical articles"
response = await client.aio.prompt_optimizer.optimize_prompt(
prompt=test_prompt,
config=types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO,
),
)
assert isinstance(response, types.OptimizeResponse)
assert response.raw_text_response


@pytest.mark.asyncio
async def test_optimize_prompt_w_few_shot_optimization_target(client):
"""Tests the optimize request parameters method with few shot optimization target."""
test_prompt = "Generate system instructions for analyzing medical articles"
df = pd.DataFrame(
{
"prompt": ["prompt1", "prompt2"],
"model_response": ["response1", "response2"],
"target_response": ["target1", "target2"],
}
)
response = await client.aio.prompt_optimizer.optimize_prompt(
prompt=test_prompt,
config=types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE,
examples_dataframe=df,
),
)
assert isinstance(response, types.OptimizeResponse)
assert response.raw_text_response


@pytest.mark.asyncio
async def test_optimize_prompt_w_few_shot_optimization_rubrics(client):
"""Tests the optimize request parameters method with few shot optimization target."""
test_prompt = "Generate system instructions for analyzing medical articles"
df = pd.DataFrame(
{
"prompt": ["prompt1", "prompt2"],
"model_response": ["response1", "response2"],
"rubrics": ["rubric1", "rubric2"],
"rubrics_evaluations": ["[True, True]", "[True, False]"],
}
)
response = await client.aio.prompt_optimizer.optimize_prompt(
prompt=test_prompt,
config=types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS,
examples_dataframe=df,
),
)
assert isinstance(response, types.OptimizeResponse)
assert response.raw_text_response


pytestmark = pytest_helper.setup(
file=__file__,
globals_for_file=globals(),
Expand Down
67 changes: 55 additions & 12 deletions tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_prompt_return_type.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

from tests.unit.vertexai.genai.replays import pytest_helper
from vertexai._genai import types
import pandas as pd


def test_optimize_prompt(client):
Expand All @@ -27,18 +28,60 @@ def test_optimize_prompt(client):
assert response.raw_text_response


# def test_optimize_prompt_w_optimization_target(client):
# """Tests the optimize request parameters method with optimization target."""
# from google.genai import types as genai_types
# test_prompt = "Generate system instructions for analyzing medical articles"
# response = client.prompt_optimizer.optimize_prompt(
# prompt=test_prompt,
# config=types.OptimizeConfig(
# optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO,
# ),
# )
# assert isinstance(response, types.OptimizeResponse)
# assert response.raw_text_response
def test_optimize_prompt_w_optimization_target(client):
"""Tests the optimize request parameters method with optimization target."""
test_prompt = "Generate system instructions for analyzing medical articles"
response = client.prompt_optimizer.optimize_prompt(
prompt=test_prompt,
config=types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_GEMINI_NANO,
),
)
assert isinstance(response, types.OptimizeResponse)
assert response.raw_text_response


def test_optimize_prompt_w_few_shot_optimization_target(client):
"""Tests the optimize request parameters method with few shot optimization target."""
test_prompt = "Generate system instructions for analyzing medical articles"
df = pd.DataFrame(
{
"prompt": ["prompt1", "prompt2"],
"model_response": ["response1", "response2"],
"target_response": ["target1", "target2"],
}
)
response = client.prompt_optimizer.optimize_prompt(
prompt=test_prompt,
config=types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE,
examples_dataframe=df,
),
)
assert isinstance(response, types.OptimizeResponse)
assert response.raw_text_response


def test_optimize_prompt_w_few_shot_optimization_rubrics(client):
"""Tests the optimize request parameters method with few shot optimization target."""
test_prompt = "Generate system instructions for analyzing medical articles"
df = pd.DataFrame(
{
"prompt": ["prompt1", "prompt2"],
"model_response": ["response1", "response2"],
"rubrics": ["rubric1", "rubric2"],
"rubrics_evaluations": ["[True, True]", "[True, False]"],
}
)
response = client.prompt_optimizer.optimize_prompt(
prompt=test_prompt,
config=types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS,
examples_dataframe=df,
),
)
assert isinstance(response, types.OptimizeResponse)
assert response.raw_text_response


pytestmark = pytest_helper.setup(
Expand Down
84 changes: 84 additions & 0 deletions tests/unit/vertexai/genai/test_prompt_optimizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
from vertexai._genai import prompt_optimizer
from vertexai._genai import types
from google.genai import client
import pandas as pd
import pytest


Expand Down Expand Up @@ -91,6 +92,34 @@ def test_prompt_optimizer_optimize_prompt(
mock_client.assert_called_once()
mock_custom_optimize_prompt.assert_called_once()

@mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt")
def test_prompt_optimizer_optimize_few_shot(self, mock_custom_optimize_prompt):
"""Test that prompt_optimizer.optimize method for few shot optimizer."""
df = pd.DataFrame(
{
"prompt": ["prompt1", "prompt2"],
"model_response": ["response1", "response2"],
"target_response": ["target1", "target2"],
}
)
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
test_config = types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE,
examples_dataframe=df,
)
test_client.prompt_optimizer.optimize_prompt(
prompt="test_prompt",
config=test_config,
)
mock_custom_optimize_prompt.assert_called_once()
mock_kwargs = mock_custom_optimize_prompt.call_args.kwargs
assert (
mock_kwargs["config"].optimization_target == test_config.optimization_target
)
pd.testing.assert_frame_equal(
mock_kwargs["config"].examples_dataframe, test_config.examples_dataframe
)

@mock.patch.object(prompt_optimizer.PromptOptimizer, "_custom_optimize_prompt")
def test_prompt_optimizer_optimize_prompt_with_optimization_target(
self, mock_custom_optimize_prompt
Expand Down Expand Up @@ -138,4 +167,59 @@ async def test_async_prompt_optimizer_optimize_prompt_with_optimization_target(
config=config,
)

@pytest.mark.asyncio
@mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt")
async def test_async_prompt_optimizer_optimize_prompt_few_shot_target_response(
self, mock_custom_optimize_prompt
):
"""Test that async prompt_optimizer.optimize_prompt calls optimize_prompt with few shot target response."""
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
df = pd.DataFrame(
{
"prompt": ["prompt1", "prompt2"],
"model_response": ["response1", "response2"],
"target_response": ["target1", "target2"],
}
)
config = types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE,
examples_dataframe=df,
)
await test_client.aio.prompt_optimizer.optimize_prompt(
prompt="test_prompt",
config=config,
)
mock_custom_optimize_prompt.assert_called_once_with(
content=mock.ANY,
config=config,
)

@pytest.mark.asyncio
@mock.patch.object(prompt_optimizer.AsyncPromptOptimizer, "_custom_optimize_prompt")
async def test_async_prompt_optimizer_optimize_prompt_few_shot_rubrics(
self, mock_custom_optimize_prompt
):
"""Test that async prompt_optimizer.optimize_prompt calls optimize_prompt with few shot rubrics."""
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
df = pd.DataFrame(
{
"prompt": ["prompt1", "prompt2"],
"model_response": ["response1", "response2"],
"rubrics": ["rubric1", "rubric2"],
"rubrics_evaluations": ["[True, True]", "[True, False]"],
}
)
config = types.OptimizeConfig(
optimization_target=types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS,
examples_dataframe=df,
)
await test_client.aio.prompt_optimizer.optimize_prompt(
prompt="test_prompt",
config=config,
)
mock_custom_optimize_prompt.assert_called_once_with(
content=mock.ANY,
config=config,
)

# # TODO(b/415060797): add more tests for prompt_optimizer.optimize
Loading
Loading