wehub-resource-sync
523 lines
19 KiB
Python
523 lines
19 KiB
Python
import json
|
|
from typing import Any
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.genai.optimize.optimizers.metaprompt_optimizer import MetaPromptOptimizer
|
|
from mlflow.genai.optimize.types import EvaluationResultRecord, PromptOptimizerOutput
|
|
|
|
_CALL_LLM = "mlflow.genai.utils.llm_utils._call_llm"
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_train_data():
|
|
return [
|
|
{
|
|
"inputs": {"question": "What is 2+2?"},
|
|
"outputs": "4",
|
|
},
|
|
{
|
|
"inputs": {"question": "What is the capital of France?"},
|
|
"outputs": "Paris",
|
|
},
|
|
{
|
|
"inputs": {"question": "What is 3*3?"},
|
|
"outputs": "9",
|
|
},
|
|
{
|
|
"inputs": {"question": "What color is the sky?"},
|
|
"outputs": "Blue",
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_target_prompts():
|
|
return {
|
|
"instruction": "Answer the following question: {{question}}",
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_target_prompts_multiple():
|
|
return {
|
|
"system_prompt": "You are a helpful assistant.",
|
|
"instruction": "Answer the following question: {{question}}",
|
|
}
|
|
|
|
|
|
def mock_eval_fn(candidate_prompts: dict[str, str], dataset: list[dict[str, Any]]):
|
|
"""Mock evaluation function that returns varied scores."""
|
|
# Return varied scores for diverse sampling
|
|
scores = [0.9, 0.7, 0.4, 0.2] # High to low
|
|
return [
|
|
EvaluationResultRecord(
|
|
inputs=record["inputs"],
|
|
outputs="mock output",
|
|
expectations=record["outputs"],
|
|
score=scores[i % len(scores)],
|
|
trace=Mock(), # Use Mock for trace
|
|
rationales={"correctness": f"Score {scores[i % len(scores)]}"},
|
|
)
|
|
for i, record in enumerate(dataset)
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_llm_response():
|
|
"""Mock LLM response with improved prompts."""
|
|
mock_response = Mock()
|
|
mock_response.choices = [Mock()]
|
|
mock_response.choices[0].message = Mock()
|
|
mock_response.choices[0].message.content = json.dumps({
|
|
"instruction": "Improved: Answer this question carefully: {{question}}"
|
|
})
|
|
return mock_response
|
|
|
|
|
|
def test_metaprompt_optimizer_initialization():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
assert optimizer.reflection_model == "openai:/gpt-4o"
|
|
assert optimizer.lm_kwargs == {}
|
|
|
|
|
|
def test_metaprompt_optimizer_initialization_with_custom_params():
|
|
optimizer = MetaPromptOptimizer(
|
|
reflection_model="anthropic:/claude-3-5-sonnet-20241022",
|
|
lm_kwargs={"temperature": 0.9, "max_tokens": 4096},
|
|
)
|
|
assert optimizer.reflection_model == "anthropic:/claude-3-5-sonnet-20241022"
|
|
assert optimizer.lm_kwargs == {"temperature": 0.9, "max_tokens": 4096}
|
|
|
|
|
|
def test_metaprompt_optimizer_invalid_lm_kwargs():
|
|
with pytest.raises(MlflowException, match="`lm_kwargs` must be a dictionary"):
|
|
MetaPromptOptimizer(reflection_model="openai:/gpt-4o", lm_kwargs="invalid")
|
|
|
|
with pytest.raises(MlflowException, match="`lm_kwargs` must be a dictionary"):
|
|
MetaPromptOptimizer(reflection_model="openai:/gpt-4o", lm_kwargs=123)
|
|
|
|
|
|
def test_extract_template_variables():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
prompts = {
|
|
"instruction": "Answer {{question}} about {{topic}}",
|
|
"system": "You are a {{role}}",
|
|
}
|
|
|
|
variables = optimizer._extract_template_variables(prompts)
|
|
|
|
assert variables["instruction"] == {"question", "topic"}
|
|
assert variables["system"] == {"role"}
|
|
|
|
|
|
def test_validate_template_variables_success():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
original = {"instruction": "Answer {{question}}"}
|
|
new = {"instruction": "Please answer this {{question}} carefully"}
|
|
|
|
# Should not raise
|
|
assert optimizer._validate_template_variables(original, new) is True
|
|
|
|
|
|
def test_validate_template_variables_missing_var():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
original = {"instruction": "Answer {{question}}"}
|
|
new = {"instruction": "Answer the question"} # Missing {{question}}
|
|
|
|
with pytest.raises(MlflowException, match="Missing.*question"):
|
|
optimizer._validate_template_variables(original, new)
|
|
|
|
|
|
def test_validate_template_variables_extra_var_stripped():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
original = {"instruction": "Answer {{question}}"}
|
|
new = {"instruction": "Answer {{question}} about {{topic}}"}
|
|
|
|
assert optimizer._validate_template_variables(original, new) is True
|
|
assert new["instruction"] == "Answer {{question}} about "
|
|
|
|
|
|
def test_validate_template_variables_extra_var_no_original_vars():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
original = {"system": "You are a helpful assistant."}
|
|
new = {"system": "You are {{role}}, a helpful assistant."}
|
|
|
|
assert optimizer._validate_template_variables(original, new) is True
|
|
assert new["system"] == "You are , a helpful assistant."
|
|
|
|
|
|
def test_validate_template_variables_multiple_extra_vars_stripped():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
original = {"instruction": "Answer {{question}}"}
|
|
new = {"instruction": "As a {{role}}, answer {{question}} about {{topic}}"}
|
|
|
|
assert optimizer._validate_template_variables(original, new) is True
|
|
assert new["instruction"] == "As a , answer {{question}} about "
|
|
|
|
|
|
def test_validate_prompt_names_missing():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
original = {"instruction": "Answer {{question}}", "system": "You are helpful"}
|
|
new = {"instruction": "Answer {{question}}"}
|
|
|
|
with pytest.raises(MlflowException, match="Prompts missing.*system"):
|
|
optimizer._validate_prompt_names(original, new)
|
|
|
|
|
|
def test_validate_prompt_names_unexpected():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
original = {"instruction": "Answer {{question}}"}
|
|
new = {
|
|
"instruction": "Answer {{question}}",
|
|
"extra_prompt": "This is unexpected",
|
|
}
|
|
|
|
with pytest.raises(MlflowException, match="Unexpected prompts.*extra_prompt"):
|
|
optimizer._validate_prompt_names(original, new)
|
|
|
|
|
|
def test_validate_prompt_names_success():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
original = {"instruction": "Answer {{question}}", "system": "You are helpful"}
|
|
new = {"instruction": "Answer {{question}}", "system": "You are an expert"}
|
|
|
|
assert optimizer._validate_prompt_names(original, new) is True
|
|
|
|
|
|
def test_build_zero_shot_meta_prompt(sample_target_prompts):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
template_vars = optimizer._extract_template_variables(sample_target_prompts)
|
|
|
|
meta_prompt = optimizer._build_zero_shot_meta_prompt(sample_target_prompts, template_vars)
|
|
|
|
assert "PROMPT ENGINEERING BEST PRACTICES" in meta_prompt
|
|
assert "{{question}}" in meta_prompt or "question" in meta_prompt
|
|
assert "instruction" in meta_prompt
|
|
assert "JSON" in meta_prompt
|
|
|
|
|
|
def test_build_few_shot_meta_prompt(sample_train_data, sample_target_prompts):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
template_vars = optimizer._extract_template_variables(sample_target_prompts)
|
|
eval_results = mock_eval_fn(sample_target_prompts, sample_train_data)
|
|
|
|
meta_prompt = optimizer._build_few_shot_meta_prompt(
|
|
sample_target_prompts, template_vars, eval_results
|
|
)
|
|
|
|
assert "EVALUATION EXAMPLES" in meta_prompt
|
|
assert "Example 1:" in meta_prompt
|
|
assert "Score:" in meta_prompt
|
|
assert "JSON" in meta_prompt
|
|
|
|
|
|
def test_build_few_shot_meta_prompt_empty_eval_results(sample_target_prompts):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
template_vars = optimizer._extract_template_variables(sample_target_prompts)
|
|
|
|
with pytest.raises(MlflowException, match="Few-shot metaprompting requires evaluation results"):
|
|
optimizer._build_few_shot_meta_prompt(sample_target_prompts, template_vars, [])
|
|
|
|
|
|
def test_format_examples(sample_train_data):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
eval_results = mock_eval_fn({}, sample_train_data)
|
|
|
|
formatted = optimizer._format_examples(eval_results[:2])
|
|
|
|
assert "Example 1:" in formatted
|
|
assert "Example 2:" in formatted
|
|
assert "Input:" in formatted
|
|
assert "Output:" in formatted
|
|
assert "Score:" in formatted
|
|
|
|
|
|
def test_call_reflection_model_success(mock_llm_response):
|
|
with patch(_CALL_LLM, return_value=mock_llm_response) as mock_call:
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
result = optimizer._call_reflection_model("test prompt")
|
|
|
|
assert isinstance(result, dict)
|
|
assert "instruction" in result
|
|
assert "{{question}}" in result["instruction"]
|
|
|
|
mock_call.assert_called_once()
|
|
args, kwargs = mock_call.call_args
|
|
assert args[0] == "openai:/gpt-4o"
|
|
assert kwargs["json_mode"] is True
|
|
|
|
|
|
def test_call_reflection_model_with_markdown():
|
|
# Test response with markdown code blocks
|
|
mock_response = Mock()
|
|
mock_response.choices = [Mock()]
|
|
mock_response.choices[0].message = Mock()
|
|
mock_response.choices[0].message.content = """```json
|
|
{
|
|
"instruction": "Improved: Answer {{question}}"
|
|
}
|
|
```"""
|
|
|
|
with patch(_CALL_LLM, return_value=mock_response):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
result = optimizer._call_reflection_model("test prompt")
|
|
|
|
assert isinstance(result, dict)
|
|
assert "instruction" in result
|
|
|
|
|
|
def test_call_reflection_model_llm_failure():
|
|
with patch(_CALL_LLM, side_effect=Exception("API error")):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
with pytest.raises(MlflowException, match="Failed to call reflection model"):
|
|
optimizer._call_reflection_model("test prompt")
|
|
|
|
|
|
def test_call_reflection_model_with_lm_kwargs(mock_llm_response):
|
|
custom_lm_kwargs = {"temperature": 0.5, "max_tokens": 2048, "top_p": 0.9}
|
|
|
|
with patch(_CALL_LLM, return_value=mock_llm_response) as mock_call:
|
|
optimizer = MetaPromptOptimizer(
|
|
reflection_model="openai:/gpt-4o", lm_kwargs=custom_lm_kwargs
|
|
)
|
|
result = optimizer._call_reflection_model("test prompt")
|
|
|
|
assert isinstance(result, dict)
|
|
|
|
# Verify that custom lm_kwargs were passed as inference_params
|
|
mock_call.assert_called_once()
|
|
_, kwargs = mock_call.call_args
|
|
assert kwargs["inference_params"] == custom_lm_kwargs
|
|
|
|
|
|
def test_optimize_zero_shot_mode(sample_target_prompts, mock_llm_response):
|
|
with patch(_CALL_LLM, return_value=mock_llm_response) as mock_call:
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
result = optimizer.optimize(
|
|
eval_fn=Mock(), # Not used in zero-shot
|
|
train_data=[], # Empty triggers zero-shot
|
|
target_prompts=sample_target_prompts,
|
|
enable_tracking=False,
|
|
)
|
|
|
|
assert isinstance(result, PromptOptimizerOutput)
|
|
assert result.initial_eval_score is None # No evaluation in zero-shot
|
|
assert result.final_eval_score is None
|
|
assert "instruction" in result.optimized_prompts
|
|
assert "{{question}}" in result.optimized_prompts["instruction"]
|
|
# Zero-shot uses single pass
|
|
assert mock_call.call_count == 1
|
|
|
|
|
|
def test_optimize_few_shot_mode(sample_train_data, sample_target_prompts, mock_llm_response):
|
|
with patch(_CALL_LLM, return_value=mock_llm_response) as mock_call:
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
result = optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
enable_tracking=False,
|
|
)
|
|
|
|
assert isinstance(result, PromptOptimizerOutput)
|
|
assert result.initial_eval_score is not None
|
|
assert result.final_eval_score is not None # Sanity check evaluation on train data
|
|
assert "instruction" in result.optimized_prompts
|
|
assert mock_call.call_count == 1 # Single pass
|
|
|
|
|
|
def test_optimize_few_shot_with_baseline_eval(sample_train_data, sample_target_prompts):
|
|
# Mock LLM to return improved prompts
|
|
mock_response = Mock()
|
|
mock_response.choices = [Mock()]
|
|
mock_response.choices[0].message = Mock()
|
|
mock_response.choices[0].message.content = json.dumps({
|
|
"instruction": "Better: Answer {{question}}"
|
|
})
|
|
|
|
# Mock eval_fn that returns scores
|
|
def mock_eval_fn(candidate_prompts, dataset):
|
|
return [
|
|
EvaluationResultRecord(
|
|
inputs=record["inputs"],
|
|
outputs="mock output",
|
|
expectations=record["outputs"],
|
|
score=0.7,
|
|
trace=Mock(),
|
|
rationales={},
|
|
)
|
|
for record in dataset
|
|
]
|
|
|
|
with patch(_CALL_LLM, return_value=mock_response):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
result = optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
enable_tracking=False,
|
|
)
|
|
|
|
# Should have both baseline and final eval scores (sanity check)
|
|
assert result.initial_eval_score is not None
|
|
assert result.final_eval_score is not None
|
|
assert "Better" in result.optimized_prompts["instruction"]
|
|
|
|
|
|
def test_optimize_strips_extra_vars_from_no_variable_prompt(sample_train_data):
|
|
mock_response = Mock()
|
|
mock_response.choices = [Mock()]
|
|
mock_response.choices[0].message = Mock()
|
|
mock_response.choices[0].message.content = json.dumps({
|
|
"system": "You are an expert {{topic}} assistant.",
|
|
"instruction": "Please answer: {{question}}",
|
|
})
|
|
|
|
prompts = {
|
|
"system": "You are a helpful assistant.",
|
|
"instruction": "Answer {{question}}",
|
|
}
|
|
|
|
with patch("litellm.completion", return_value=mock_response):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
result = optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=prompts,
|
|
enable_tracking=False,
|
|
)
|
|
|
|
assert result.optimized_prompts["system"] == "You are an expert assistant."
|
|
assert result.optimized_prompts["instruction"] == "Please answer: {{question}}"
|
|
|
|
|
|
def test_optimize_preserves_template_variables(sample_train_data):
|
|
# Mock response that drops the {{question}} variable
|
|
mock_response = Mock()
|
|
mock_response.choices = [Mock()]
|
|
mock_response.choices[0].message = Mock()
|
|
mock_response.choices[0].message.content = json.dumps(
|
|
{"instruction": "Answer the question"} # Missing {{question}}
|
|
)
|
|
|
|
prompts = {"instruction": "Answer {{question}}"}
|
|
|
|
with patch(_CALL_LLM, return_value=mock_response):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
result = optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=prompts,
|
|
enable_tracking=False,
|
|
)
|
|
|
|
# Should keep original prompts due to validation failure
|
|
# (caught as exception and logged as warning)
|
|
assert "{{question}}" in result.optimized_prompts["instruction"]
|
|
|
|
|
|
def test_optimize_with_multiple_prompts(sample_train_data, sample_target_prompts_multiple):
|
|
mock_response = Mock()
|
|
mock_response.choices = [Mock()]
|
|
mock_response.choices[0].message = Mock()
|
|
mock_response.choices[0].message.content = json.dumps({
|
|
"system_prompt": "Improved: You are an expert assistant.",
|
|
"instruction": "Improved: Answer {{question}}",
|
|
})
|
|
|
|
with patch(_CALL_LLM, return_value=mock_response):
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
result = optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts_multiple,
|
|
enable_tracking=False,
|
|
)
|
|
|
|
assert "system_prompt" in result.optimized_prompts
|
|
assert "instruction" in result.optimized_prompts
|
|
assert "{{question}}" in result.optimized_prompts["instruction"]
|
|
|
|
|
|
def test_build_zero_shot_meta_prompt_with_guidelines(sample_target_prompts):
|
|
custom_guidelines = "Focus on concise, accurate answers for finance domain."
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o", guidelines=custom_guidelines)
|
|
template_vars = optimizer._extract_template_variables(sample_target_prompts)
|
|
|
|
meta_prompt = optimizer._build_zero_shot_meta_prompt(sample_target_prompts, template_vars)
|
|
|
|
# Verify structure
|
|
assert "CUSTOM GUIDELINES:" in meta_prompt
|
|
assert custom_guidelines in meta_prompt
|
|
assert "TEMPLATE VARIABLES:" in meta_prompt
|
|
assert "PROMPT ENGINEERING BEST PRACTICES:" in meta_prompt
|
|
|
|
|
|
def test_build_few_shot_meta_prompt_with_guidelines(sample_target_prompts):
|
|
custom_guidelines = "Focus on concise, accurate answers for finance domain."
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o", guidelines=custom_guidelines)
|
|
template_vars = optimizer._extract_template_variables(sample_target_prompts)
|
|
|
|
# Create sample evaluation results
|
|
eval_results = [
|
|
EvaluationResultRecord(
|
|
inputs={"question": "test"},
|
|
outputs="answer",
|
|
expectations="answer",
|
|
score=0.8,
|
|
trace=Mock(),
|
|
rationales={"correctness": "Good"},
|
|
)
|
|
]
|
|
|
|
meta_prompt = optimizer._build_few_shot_meta_prompt(
|
|
sample_target_prompts, template_vars, eval_results
|
|
)
|
|
|
|
# Verify structure
|
|
assert "CUSTOM GUIDELINES:" in meta_prompt
|
|
assert custom_guidelines in meta_prompt
|
|
assert "TEMPLATE VARIABLES:" in meta_prompt
|
|
assert "EVALUATION EXAMPLES" in meta_prompt # Now includes score in header
|
|
assert "Current Score:" in meta_prompt
|
|
|
|
|
|
def test_compute_per_scorer_scores():
|
|
optimizer = MetaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
# Test with multiple results having individual scores
|
|
eval_results = [
|
|
EvaluationResultRecord(
|
|
inputs={"q": "1"},
|
|
outputs="a",
|
|
expectations="a",
|
|
score=0.8,
|
|
trace=Mock(),
|
|
rationales={},
|
|
individual_scores={"Correctness": 0.9, "Safety": 0.7},
|
|
),
|
|
EvaluationResultRecord(
|
|
inputs={"q": "2"},
|
|
outputs="b",
|
|
expectations="b",
|
|
score=0.6,
|
|
trace=Mock(),
|
|
rationales={},
|
|
individual_scores={"Correctness": 0.7, "Safety": 0.5},
|
|
),
|
|
]
|
|
|
|
per_scorer = optimizer._compute_per_scorer_scores(eval_results)
|
|
|
|
assert per_scorer == {"Correctness": 0.8, "Safety": 0.6} # Average of each scorer
|