wehub-resource-sync
662 lines
23 KiB
Python
662 lines
23 KiB
Python
import json
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest.mock import MagicMock, Mock, patch
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.genai.optimize.optimizers.gepa_optimizer import GepaPromptOptimizer
|
|
from mlflow.genai.optimize.types import EvaluationResultRecord, PromptOptimizerOutput
|
|
|
|
|
|
@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 {
|
|
"system_prompt": "You are a helpful assistant.",
|
|
"instruction": "Answer the following question: {{question}}",
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_eval_fn():
|
|
def eval_fn(candidate_prompts: dict[str, str], dataset: list[dict[str, Any]]):
|
|
return [
|
|
EvaluationResultRecord(
|
|
inputs=record["inputs"],
|
|
outputs="outputs",
|
|
expectations=record["outputs"],
|
|
score=0.8,
|
|
trace={"info": "mock trace"},
|
|
rationales={"score": "mock rationale"},
|
|
individual_scores={"accuracy": 0.9, "relevance": 0.7},
|
|
)
|
|
for record in dataset
|
|
]
|
|
|
|
return eval_fn
|
|
|
|
|
|
def test_gepa_optimizer_initialization():
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
assert optimizer.reflection_model == "openai:/gpt-4o"
|
|
assert optimizer.max_metric_calls == 100
|
|
assert optimizer.display_progress_bar is False
|
|
assert optimizer.gepa_kwargs == {}
|
|
|
|
|
|
def test_gepa_optimizer_initialization_with_custom_params():
|
|
optimizer = GepaPromptOptimizer(
|
|
reflection_model="openai:/gpt-4o",
|
|
max_metric_calls=100,
|
|
display_progress_bar=True,
|
|
)
|
|
assert optimizer.reflection_model == "openai:/gpt-4o"
|
|
assert optimizer.max_metric_calls == 100
|
|
assert optimizer.display_progress_bar is True
|
|
assert optimizer.gepa_kwargs == {}
|
|
|
|
|
|
def test_gepa_optimizer_initialization_with_gepa_kwargs():
|
|
gepa_kwargs_example = {"foo": "bar"}
|
|
optimizer = GepaPromptOptimizer(
|
|
reflection_model="openai:/gpt-4o",
|
|
gepa_kwargs=gepa_kwargs_example,
|
|
)
|
|
assert optimizer.reflection_model == "openai:/gpt-4o"
|
|
assert optimizer.max_metric_calls == 100
|
|
assert optimizer.display_progress_bar is False
|
|
assert optimizer.gepa_kwargs == gepa_kwargs_example
|
|
|
|
|
|
def test_gepa_optimizer_optimize(
|
|
sample_train_data: list[dict[str, Any]],
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
):
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = {
|
|
"system_prompt": "You are a highly skilled assistant.",
|
|
"instruction": "Please answer this question carefully: {{question}}",
|
|
}
|
|
mock_result.val_aggregate_scores = [0.5, 0.6, 0.8, 0.9] # Mock scores for testing
|
|
mock_result.val_aggregate_subscores = [
|
|
{"accuracy": 0.4, "relevance": 0.6}, # Initial (index 0)
|
|
{"accuracy": 0.5, "relevance": 0.7}, # Index 1
|
|
{"accuracy": 0.7, "relevance": 0.9}, # Index 2
|
|
{"accuracy": 0.85, "relevance": 0.95}, # Final best (index 3, max score)
|
|
]
|
|
mock_gepa_module.optimize.return_value = mock_result
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
optimizer = GepaPromptOptimizer(
|
|
reflection_model="openai:/gpt-4o-mini", max_metric_calls=50, display_progress_bar=True
|
|
)
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
result = optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
# Verify result
|
|
assert isinstance(result, PromptOptimizerOutput)
|
|
assert result.optimized_prompts == mock_result.best_candidate
|
|
assert "system_prompt" in result.optimized_prompts
|
|
assert "instruction" in result.optimized_prompts
|
|
# Verify aggregated scores are extracted
|
|
assert result.initial_eval_score == 0.5 # First score
|
|
assert result.final_eval_score == 0.9 # Max score
|
|
# Verify per-scorer scores are extracted
|
|
assert result.initial_eval_score_per_scorer == {"accuracy": 0.4, "relevance": 0.6}
|
|
assert result.final_eval_score_per_scorer == {"accuracy": 0.85, "relevance": 0.95}
|
|
|
|
# Verify GEPA was called with correct parameters
|
|
mock_gepa_module.optimize.assert_called_once()
|
|
call_kwargs = mock_gepa_module.optimize.call_args.kwargs
|
|
|
|
assert call_kwargs["seed_candidate"] == sample_target_prompts
|
|
assert call_kwargs["adapter"] is not None
|
|
assert call_kwargs["max_metric_calls"] == 50
|
|
assert call_kwargs["reflection_lm"] == "openai/gpt-4o-mini"
|
|
assert call_kwargs["display_progress_bar"] is True
|
|
assert len(call_kwargs["trainset"]) == 4
|
|
|
|
|
|
def test_gepa_optimizer_optimize_with_custom_reflection_model(
|
|
sample_train_data: list[dict[str, Any]],
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
):
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = sample_target_prompts
|
|
mock_result.val_aggregate_scores = []
|
|
mock_result.val_aggregate_subscores = None
|
|
mock_gepa_module.optimize.return_value = mock_result
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
|
|
optimizer = GepaPromptOptimizer(
|
|
reflection_model="anthropic:/claude-3-5-sonnet-20241022",
|
|
)
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
call_kwargs = mock_gepa_module.optimize.call_args.kwargs
|
|
assert call_kwargs["reflection_lm"] == "anthropic/claude-3-5-sonnet-20241022"
|
|
|
|
|
|
def test_gepa_optimizer_optimize_with_custom_gepa_params(
|
|
sample_train_data: list[dict[str, Any]],
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
):
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = sample_target_prompts
|
|
mock_result.val_aggregate_scores = []
|
|
mock_result.val_aggregate_subscores = None
|
|
mock_gepa_module.optimize.return_value = mock_result
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
|
|
optimizer = GepaPromptOptimizer(
|
|
reflection_model="openai:/gpt-4o-mini", gepa_kwargs={"foo": "bar"}
|
|
)
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
call_kwargs = mock_gepa_module.optimize.call_args.kwargs
|
|
assert call_kwargs["foo"] == "bar"
|
|
|
|
|
|
def test_gepa_optimizer_optimize_model_name_parsing(
|
|
sample_train_data: list[dict[str, Any]],
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
):
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = sample_target_prompts
|
|
mock_result.val_aggregate_scores = []
|
|
mock_result.val_aggregate_subscores = None
|
|
mock_gepa_module.optimize.return_value = mock_result
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
call_kwargs = mock_gepa_module.optimize.call_args.kwargs
|
|
assert call_kwargs["reflection_lm"] == "openai/gpt-4o"
|
|
|
|
|
|
def test_gepa_optimizer_import_error(
|
|
sample_train_data: list[dict[str, Any]],
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
):
|
|
with patch.dict("sys.modules", {"gepa": None}):
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
with pytest.raises(ImportError, match="GEPA >= 0.0.26 is required"):
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
|
|
def test_gepa_optimizer_requires_train_data(
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
):
|
|
from mlflow.exceptions import MlflowException
|
|
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="GEPA optimizer requires `train_data` to be provided",
|
|
):
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=[],
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
|
|
def test_gepa_optimizer_single_record_dataset(
|
|
sample_target_prompts: dict[str, str], mock_eval_fn: Any
|
|
):
|
|
single_record_data = [
|
|
{
|
|
"inputs": {"question": "What is 2+2?"},
|
|
"outputs": "4",
|
|
}
|
|
]
|
|
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = sample_target_prompts
|
|
mock_result.val_aggregate_scores = []
|
|
mock_result.val_aggregate_subscores = None
|
|
mock_gepa_module.optimize.return_value = mock_result
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=single_record_data,
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
call_kwargs = mock_gepa_module.optimize.call_args.kwargs
|
|
assert len(call_kwargs["trainset"]) == 1
|
|
|
|
|
|
def test_gepa_optimizer_custom_adapter_evaluate(
|
|
sample_train_data: list[dict[str, Any]],
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
):
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = sample_target_prompts
|
|
mock_result.val_aggregate_scores = []
|
|
mock_result.val_aggregate_subscores = None
|
|
mock_gepa_module.optimize.return_value = mock_result
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
result = optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
call_kwargs = mock_gepa_module.optimize.call_args.kwargs
|
|
assert "adapter" in call_kwargs
|
|
assert call_kwargs["adapter"] is not None
|
|
assert result.optimized_prompts == sample_target_prompts
|
|
|
|
|
|
def test_make_reflective_dataset_with_traces(
|
|
sample_target_prompts: dict[str, str], mock_eval_fn: Any
|
|
):
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
mock_gepa_module.GEPAAdapter = object
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
captured_adapter = None
|
|
|
|
def mock_optimize_fn(**kwargs):
|
|
nonlocal captured_adapter
|
|
captured_adapter = kwargs["adapter"]
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = sample_target_prompts
|
|
mock_result.val_aggregate_scores = []
|
|
mock_result.val_aggregate_subscores = None
|
|
return mock_result
|
|
|
|
mock_gepa_module.optimize = mock_optimize_fn
|
|
|
|
# Call optimize to create the inner adapter
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=[{"inputs": {"question": "test"}, "outputs": "test"}],
|
|
target_prompts=sample_target_prompts,
|
|
)
|
|
|
|
# Now test make_reflective_dataset with the captured adapter
|
|
mock_trace = Mock()
|
|
mock_span1 = Mock()
|
|
mock_span1.name = "llm_call"
|
|
mock_span1.inputs = {"prompt": "What is 2+2?"}
|
|
mock_span1.outputs = {"response": "4"}
|
|
|
|
mock_span2 = Mock()
|
|
mock_span2.name = "retrieval"
|
|
mock_span2.inputs = {"query": "math"}
|
|
mock_span2.outputs = {"documents": ["doc1", "doc2"]}
|
|
|
|
mock_trace.data.spans = [mock_span1, mock_span2]
|
|
|
|
# Create mock trajectories
|
|
mock_trajectory1 = Mock()
|
|
mock_trajectory1.trace = mock_trace
|
|
mock_trajectory1.inputs = {"question": "What is 2+2?"}
|
|
mock_trajectory1.outputs = "4"
|
|
mock_trajectory1.expectations = {"expected_response": "4"}
|
|
|
|
mock_trajectory2 = Mock()
|
|
mock_trajectory2.trace = None
|
|
mock_trajectory2.inputs = {"question": "What is the capital of France?"}
|
|
mock_trajectory2.outputs = "Paris"
|
|
mock_trajectory2.expectations = {"expected_response": "Paris"}
|
|
|
|
# Create mock evaluation batch
|
|
mock_eval_batch = Mock()
|
|
mock_eval_batch.trajectories = [mock_trajectory1, mock_trajectory2]
|
|
mock_eval_batch.scores = [0.9, 0.7]
|
|
|
|
# Test make_reflective_dataset
|
|
candidate = {"system_prompt": "You are helpful"}
|
|
components_to_update = ["system_prompt", "instruction"]
|
|
|
|
result = captured_adapter.make_reflective_dataset(
|
|
candidate, mock_eval_batch, components_to_update
|
|
)
|
|
|
|
# Verify result structure
|
|
assert isinstance(result, dict)
|
|
assert "system_prompt" in result
|
|
assert "instruction" in result
|
|
|
|
system_data = result["system_prompt"]
|
|
assert len(system_data) == 2
|
|
assert system_data[0]["component_name"] == "system_prompt"
|
|
assert system_data[0]["current_text"] == "You are helpful"
|
|
assert system_data[0]["score"] == 0.9
|
|
assert system_data[0]["inputs"] == {"question": "What is 2+2?"}
|
|
assert system_data[0]["outputs"] == "4"
|
|
assert system_data[0]["expectations"] == {"expected_response": "4"}
|
|
assert system_data[0]["index"] == 0
|
|
|
|
# Verify trace spans
|
|
assert len(system_data[0]["trace"]) == 2
|
|
assert system_data[0]["trace"][0]["name"] == "llm_call"
|
|
assert system_data[0]["trace"][0]["inputs"] == {"prompt": "What is 2+2?"}
|
|
assert system_data[0]["trace"][0]["outputs"] == {"response": "4"}
|
|
assert system_data[0]["trace"][1]["name"] == "retrieval"
|
|
|
|
# Verify second record (no trace)
|
|
assert system_data[1]["trace"] == []
|
|
assert system_data[1]["score"] == 0.7
|
|
assert system_data[1]["inputs"] == {"question": "What is the capital of France?"}
|
|
assert system_data[1]["outputs"] == "Paris"
|
|
assert system_data[1]["expectations"] == {"expected_response": "Paris"}
|
|
|
|
|
|
@pytest.mark.parametrize("enable_tracking", [True, False])
|
|
def test_gepa_optimizer_passes_use_mlflow(
|
|
sample_train_data: list[dict[str, Any]],
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
enable_tracking: bool,
|
|
):
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = sample_target_prompts
|
|
mock_result.val_aggregate_scores = []
|
|
mock_result.val_aggregate_subscores = None
|
|
mock_gepa_module.optimize.return_value = mock_result
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
enable_tracking=enable_tracking,
|
|
)
|
|
|
|
call_kwargs = mock_gepa_module.optimize.call_args.kwargs
|
|
assert "use_mlflow" in call_kwargs
|
|
assert call_kwargs["use_mlflow"] == enable_tracking
|
|
|
|
|
|
def test_gepa_optimizer_logs_prompt_candidates(
|
|
sample_train_data: list[dict[str, Any]],
|
|
sample_target_prompts: dict[str, str],
|
|
mock_eval_fn: Any,
|
|
):
|
|
mock_gepa_module = MagicMock()
|
|
mock_modules = {
|
|
"gepa": mock_gepa_module,
|
|
"gepa.core": MagicMock(),
|
|
"gepa.core.adapter": MagicMock(),
|
|
}
|
|
mock_gepa_module.EvaluationBatch = MagicMock()
|
|
mock_gepa_module.GEPAAdapter = object
|
|
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
|
|
logged_artifacts = []
|
|
logged_tables = []
|
|
logged_metrics = []
|
|
|
|
with patch.dict(sys.modules, mock_modules):
|
|
captured_adapter = None
|
|
|
|
def mock_optimize_fn(**kwargs):
|
|
nonlocal captured_adapter
|
|
captured_adapter = kwargs["adapter"]
|
|
mock_result = Mock()
|
|
mock_result.best_candidate = sample_target_prompts
|
|
mock_result.val_aggregate_scores = [0.8]
|
|
mock_result.val_aggregate_subscores = None
|
|
return mock_result
|
|
|
|
mock_gepa_module.optimize = mock_optimize_fn
|
|
|
|
with mlflow.start_run():
|
|
with (
|
|
patch(
|
|
"mlflow.genai.optimize.optimizers.gepa_optimizer.mlflow.log_artifact"
|
|
) as mock_log_artifact,
|
|
patch(
|
|
"mlflow.genai.optimize.optimizers.gepa_optimizer.mlflow.log_table"
|
|
) as mock_log_table,
|
|
patch(
|
|
"mlflow.genai.optimize.optimizers.gepa_optimizer.mlflow.log_metrics"
|
|
) as mock_log_metrics,
|
|
):
|
|
|
|
def capture_artifact(path, artifact_path=None):
|
|
with open(path) as f:
|
|
logged_artifacts.append({
|
|
"path": str(path),
|
|
"artifact_path": artifact_path,
|
|
"content": f.read(),
|
|
})
|
|
|
|
def capture_table(data, artifact_file):
|
|
logged_tables.append({"data": data, "artifact_file": artifact_file})
|
|
|
|
def capture_metrics(metrics, step=None):
|
|
logged_metrics.append({"metrics": metrics, "step": step})
|
|
|
|
mock_log_artifact.side_effect = capture_artifact
|
|
mock_log_table.side_effect = capture_table
|
|
mock_log_metrics.side_effect = capture_metrics
|
|
|
|
optimizer.optimize(
|
|
eval_fn=mock_eval_fn,
|
|
train_data=sample_train_data,
|
|
target_prompts=sample_target_prompts,
|
|
enable_tracking=True,
|
|
)
|
|
|
|
# First: minibatch evaluation (should NOT log any artifacts)
|
|
minibatch = sample_train_data[:2]
|
|
captured_adapter.evaluate(
|
|
minibatch, {"system_prompt": "Test"}, capture_traces=False
|
|
)
|
|
|
|
# Second: full dataset validation (should log artifacts)
|
|
candidate = {"system_prompt": "Optimized prompt", "instruction": "New instruction"}
|
|
captured_adapter.evaluate(sample_train_data, candidate, capture_traces=False)
|
|
|
|
# Verify scores.json was logged
|
|
scores_artifact = next((a for a in logged_artifacts if "scores.json" in a["path"]), None)
|
|
assert scores_artifact is not None
|
|
assert scores_artifact["artifact_path"] == "prompt_candidates/iteration_0"
|
|
scores_content = json.loads(scores_artifact["content"])
|
|
assert scores_content["aggregate"] == 0.8
|
|
assert scores_content["per_scorer"] == {"accuracy": 0.9, "relevance": 0.7}
|
|
|
|
# Verify prompt text files were logged
|
|
prompt_artifacts = [a for a in logged_artifacts if a["path"].endswith(".txt")]
|
|
assert len(prompt_artifacts) == 2 # system_prompt.txt and instruction.txt
|
|
for a in prompt_artifacts:
|
|
assert a["artifact_path"] == "prompt_candidates/iteration_0"
|
|
prompt_contents = {Path(a["path"]).stem: a["content"] for a in prompt_artifacts}
|
|
assert prompt_contents["system_prompt"] == "Optimized prompt"
|
|
assert prompt_contents["instruction"] == "New instruction"
|
|
|
|
# Verify eval results table was logged
|
|
assert len(logged_tables) == 1
|
|
table = logged_tables[0]
|
|
assert table["artifact_file"] == "prompt_candidates/iteration_0/eval_results.json"
|
|
data = table["data"]
|
|
assert "inputs" in data
|
|
assert "output" in data
|
|
assert "expectation" in data
|
|
assert "aggregate_score" in data
|
|
assert "accuracy" in data
|
|
assert "relevance" in data
|
|
assert len(data["inputs"]) == len(sample_train_data)
|
|
assert all(score == 0.9 for score in data["accuracy"])
|
|
assert all(score == 0.7 for score in data["relevance"])
|
|
|
|
# Verify metrics were logged with step for time progression
|
|
assert len(logged_metrics) == 1
|
|
metrics = logged_metrics[0]
|
|
assert metrics["step"] == 0
|
|
assert metrics["metrics"]["eval_score"] == 0.8
|
|
assert metrics["metrics"]["eval_score.accuracy"] == 0.9
|
|
assert metrics["metrics"]["eval_score.relevance"] == 0.7
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("val_aggregate_scores", "val_aggregate_subscores", "expected"),
|
|
[
|
|
# No scores at all
|
|
([], None, (None, None, {}, {})),
|
|
(None, None, (None, None, {}, {})),
|
|
# Only aggregate scores, no subscores
|
|
([0.5, 0.7, 0.9], None, (0.5, 0.9, {}, {})),
|
|
# Both aggregate and per-scorer scores
|
|
(
|
|
[0.5, 0.7, 0.9],
|
|
[
|
|
{"Correctness": 0.4, "Safety": 0.6},
|
|
{"Correctness": 0.6, "Safety": 0.8},
|
|
{"Correctness": 0.85, "Safety": 0.95},
|
|
],
|
|
(0.5, 0.9, {"Correctness": 0.4, "Safety": 0.6}, {"Correctness": 0.85, "Safety": 0.95}),
|
|
),
|
|
# Empty subscores dict at index 0
|
|
(
|
|
[0.5, 0.9],
|
|
[{}, {"Correctness": 0.9}],
|
|
(0.5, 0.9, {}, {"Correctness": 0.9}),
|
|
),
|
|
# Best score not at last index
|
|
(
|
|
[0.5, 0.95, 0.8],
|
|
[
|
|
{"A": 0.4},
|
|
{"A": 0.95}, # Best score at index 1
|
|
{"A": 0.7},
|
|
],
|
|
(0.5, 0.95, {"A": 0.4}, {"A": 0.95}),
|
|
),
|
|
],
|
|
)
|
|
def test_extract_eval_scores_per_scorer(val_aggregate_scores, val_aggregate_subscores, expected):
|
|
optimizer = GepaPromptOptimizer(reflection_model="openai:/gpt-4o")
|
|
mock_result = Mock()
|
|
mock_result.val_aggregate_scores = val_aggregate_scores
|
|
mock_result.val_aggregate_subscores = val_aggregate_subscores
|
|
|
|
result = optimizer._extract_eval_scores(mock_result)
|
|
assert result == expected
|