wehub-resource-sync
478 lines
14 KiB
Python
478 lines
14 KiB
Python
from unittest import mock
|
|
|
|
import click
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from mlflow.cli.genai_eval_utils import (
|
|
NA_VALUE,
|
|
Assessment,
|
|
EvalResult,
|
|
extract_assessments_from_results,
|
|
format_table_output,
|
|
resolve_scorers,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.tracing.constant import AssessmentMetadataKey
|
|
|
|
|
|
def test_format_single_trace_with_result_and_rationale():
|
|
output_data = [
|
|
EvalResult(
|
|
trace_id="tr-123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result="yes",
|
|
rationale="The answer is relevant",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
table_output = format_table_output(output_data)
|
|
|
|
# Headers should use assessment names from output_data
|
|
assert table_output.headers == ["trace_id", "RelevanceToQuery"]
|
|
assert len(table_output.rows) == 1
|
|
assert table_output.rows[0][0].value == "tr-123"
|
|
assert "value: yes" in table_output.rows[0][1].value
|
|
assert "rationale: The answer is relevant" in table_output.rows[0][1].value
|
|
|
|
|
|
def test_format_multiple_traces_multiple_scorers():
|
|
output_data = [
|
|
EvalResult(
|
|
trace_id="tr-123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result="yes",
|
|
rationale="Relevant",
|
|
),
|
|
Assessment(name="Safety", result="yes", rationale="Safe"),
|
|
],
|
|
),
|
|
EvalResult(
|
|
trace_id="tr-456",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result="no",
|
|
rationale="Not relevant",
|
|
),
|
|
Assessment(name="Safety", result="yes", rationale="Safe"),
|
|
],
|
|
),
|
|
]
|
|
|
|
table_output = format_table_output(output_data)
|
|
|
|
# Assessment names should be sorted
|
|
assert table_output.headers == ["trace_id", "RelevanceToQuery", "Safety"]
|
|
assert len(table_output.rows) == 2
|
|
assert table_output.rows[0][0].value == "tr-123"
|
|
assert table_output.rows[1][0].value == "tr-456"
|
|
assert "value: yes" in table_output.rows[0][1].value
|
|
assert "value: no" in table_output.rows[1][1].value
|
|
|
|
|
|
def test_format_long_rationale_not_truncated():
|
|
long_rationale = "x" * 150
|
|
output_data = [
|
|
EvalResult(
|
|
trace_id="tr-123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result="yes",
|
|
rationale=long_rationale,
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
table_output = format_table_output(output_data)
|
|
|
|
assert long_rationale in table_output.rows[0][1].value
|
|
assert len(table_output.rows[0][1].value) >= len(long_rationale)
|
|
|
|
|
|
def test_format_error_message_formatting():
|
|
output_data = [
|
|
EvalResult(
|
|
trace_id="tr-123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result=None,
|
|
rationale=None,
|
|
error="OpenAI API error",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
table_output = format_table_output(output_data)
|
|
|
|
assert table_output.rows[0][1].value == "error: OpenAI API error"
|
|
|
|
|
|
def test_format_na_for_missing_results():
|
|
output_data = [
|
|
EvalResult(
|
|
trace_id="tr-123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result=None,
|
|
rationale=None,
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
table_output = format_table_output(output_data)
|
|
|
|
assert table_output.rows[0][1].value == NA_VALUE
|
|
|
|
|
|
def test_format_result_only_without_rationale():
|
|
output_data = [
|
|
EvalResult(
|
|
trace_id="tr-123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result="yes",
|
|
rationale=None,
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
table_output = format_table_output(output_data)
|
|
|
|
assert table_output.rows[0][1].value == "value: yes"
|
|
|
|
|
|
def test_format_rationale_only_without_result():
|
|
output_data = [
|
|
EvalResult(
|
|
trace_id="tr-123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result=None,
|
|
rationale="Some reasoning",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
table_output = format_table_output(output_data)
|
|
|
|
assert table_output.rows[0][1].value == "rationale: Some reasoning"
|
|
|
|
|
|
def test_format_with_different_assessment_names():
|
|
# This test demonstrates that assessment names (e.g., "relevance_to_query")
|
|
# are used in headers, not scorer class names (e.g., "RelevanceToQuery")
|
|
output_data = [
|
|
EvalResult(
|
|
trace_id="tr-123",
|
|
assessments=[
|
|
Assessment(
|
|
name="relevance_to_query", # Different from scorer name
|
|
result="yes",
|
|
rationale="The answer is relevant",
|
|
),
|
|
Assessment(
|
|
name="safety_check", # Different from scorer name
|
|
result="safe",
|
|
rationale="Content is safe",
|
|
),
|
|
],
|
|
)
|
|
]
|
|
|
|
table_output = format_table_output(output_data)
|
|
|
|
# Headers should use actual assessment names from output_data (sorted)
|
|
assert table_output.headers == ["trace_id", "relevance_to_query", "safety_check"]
|
|
assert len(table_output.rows) == 1
|
|
assert table_output.rows[0][0].value == "tr-123"
|
|
assert "value: yes" in table_output.rows[0][1].value
|
|
assert "value: safe" in table_output.rows[0][2].value
|
|
|
|
|
|
# Tests for resolve_scorers function
|
|
|
|
|
|
def test_resolve_builtin_scorer():
|
|
# Test with real built-in scorer names
|
|
scorers = resolve_scorers(["Correctness"], "experiment_123")
|
|
|
|
assert len(scorers) == 1
|
|
assert scorers[0].__class__.__name__ == "Correctness"
|
|
|
|
|
|
def test_resolve_builtin_scorer_snake_case():
|
|
# Test with snake_case name
|
|
scorers = resolve_scorers(["correctness"], "experiment_123")
|
|
|
|
assert len(scorers) == 1
|
|
assert scorers[0].__class__.__name__ == "Correctness"
|
|
|
|
|
|
def test_resolve_registered_scorer():
|
|
mock_registered = mock.Mock()
|
|
|
|
with (
|
|
mock.patch(
|
|
"mlflow.cli.genai_eval_utils.get_all_scorers", return_value=[]
|
|
) as mock_get_all_scorers,
|
|
mock.patch(
|
|
"mlflow.cli.genai_eval_utils.get_scorer", return_value=mock_registered
|
|
) as mock_get_scorer,
|
|
):
|
|
scorers = resolve_scorers(["CustomScorer"], "experiment_123")
|
|
|
|
assert len(scorers) == 1
|
|
assert scorers[0] == mock_registered
|
|
# Verify mocks were called as expected
|
|
mock_get_all_scorers.assert_called_once()
|
|
mock_get_scorer.assert_called_once_with(name="CustomScorer", experiment_id="experiment_123")
|
|
|
|
|
|
def test_resolve_mixed_scorers():
|
|
# Setup built-in scorer
|
|
mock_builtin = mock.Mock()
|
|
mock_builtin.__class__.__name__ = "Safety"
|
|
mock_builtin.name = None
|
|
|
|
# Setup registered scorer
|
|
mock_registered = mock.Mock()
|
|
|
|
with (
|
|
mock.patch(
|
|
"mlflow.cli.genai_eval_utils.get_all_scorers", return_value=[mock_builtin]
|
|
) as mock_get_all_scorers,
|
|
mock.patch(
|
|
"mlflow.cli.genai_eval_utils.get_scorer", return_value=mock_registered
|
|
) as mock_get_scorer,
|
|
):
|
|
scorers = resolve_scorers(["Safety", "CustomScorer"], "experiment_123")
|
|
|
|
assert len(scorers) == 2
|
|
assert scorers[0] == mock_builtin
|
|
assert scorers[1] == mock_registered
|
|
# Verify mocks were called as expected
|
|
mock_get_all_scorers.assert_called_once()
|
|
mock_get_scorer.assert_called_once_with(name="CustomScorer", experiment_id="experiment_123")
|
|
|
|
|
|
def test_resolve_scorer_not_found_raises_error():
|
|
with (
|
|
mock.patch(
|
|
"mlflow.cli.genai_eval_utils.get_all_scorers", return_value=[]
|
|
) as mock_get_all_scorers,
|
|
mock.patch(
|
|
"mlflow.cli.genai_eval_utils.get_scorer",
|
|
side_effect=MlflowException("Not found"),
|
|
) as mock_get_scorer,
|
|
):
|
|
with pytest.raises(click.UsageError, match="Could not identify Scorer 'UnknownScorer'"):
|
|
resolve_scorers(["UnknownScorer"], "experiment_123")
|
|
|
|
# Verify mocks were called as expected
|
|
mock_get_all_scorers.assert_called_once()
|
|
mock_get_scorer.assert_called_once_with(
|
|
name="UnknownScorer", experiment_id="experiment_123"
|
|
)
|
|
|
|
|
|
def test_resolve_empty_scorers_raises_error():
|
|
with pytest.raises(click.UsageError, match="No valid scorers"):
|
|
resolve_scorers([], "experiment_123")
|
|
|
|
|
|
# Tests for extract_assessments_from_results function
|
|
|
|
|
|
def test_extract_with_matching_run_id():
|
|
results_df = pd.DataFrame([
|
|
{
|
|
"trace_id": "tr-abc123",
|
|
"assessments": [
|
|
{
|
|
"assessment_name": "RelevanceToQuery",
|
|
"feedback": {"value": "yes"},
|
|
"rationale": "The answer is relevant",
|
|
"metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
|
|
}
|
|
],
|
|
}
|
|
])
|
|
|
|
result = extract_assessments_from_results(results_df, "run-123")
|
|
|
|
expected = [
|
|
EvalResult(
|
|
trace_id="tr-abc123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result="yes",
|
|
rationale="The answer is relevant",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
assert result == expected
|
|
|
|
|
|
def test_extract_with_different_assessment_name():
|
|
results_df = pd.DataFrame([
|
|
{
|
|
"trace_id": "tr-abc123",
|
|
"assessments": [
|
|
{
|
|
"assessment_name": "relevance_to_query",
|
|
"feedback": {"value": "yes"},
|
|
"rationale": "Relevant answer",
|
|
"metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
|
|
}
|
|
],
|
|
}
|
|
])
|
|
|
|
result = extract_assessments_from_results(results_df, "run-123")
|
|
|
|
expected = [
|
|
EvalResult(
|
|
trace_id="tr-abc123",
|
|
assessments=[
|
|
Assessment(
|
|
name="relevance_to_query",
|
|
result="yes",
|
|
rationale="Relevant answer",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
assert result == expected
|
|
|
|
|
|
def test_extract_filter_out_assessments_with_different_run_id():
|
|
results_df = pd.DataFrame([
|
|
{
|
|
"trace_id": "tr-abc123",
|
|
"assessments": [
|
|
{
|
|
"assessment_name": "RelevanceToQuery",
|
|
"feedback": {"value": "yes"},
|
|
"rationale": "Current evaluation",
|
|
"metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
|
|
},
|
|
{
|
|
"assessment_name": "Safety",
|
|
"feedback": {"value": "yes"},
|
|
"rationale": "Old evaluation",
|
|
"metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-456"},
|
|
},
|
|
],
|
|
}
|
|
])
|
|
|
|
result = extract_assessments_from_results(results_df, "run-123")
|
|
|
|
expected = [
|
|
EvalResult(
|
|
trace_id="tr-abc123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result="yes",
|
|
rationale="Current evaluation",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
assert result == expected
|
|
|
|
|
|
def test_extract_no_assessments_for_run_id():
|
|
results_df = pd.DataFrame([
|
|
{
|
|
"trace_id": "tr-abc123",
|
|
"assessments": [
|
|
{
|
|
"assessment_name": "RelevanceToQuery",
|
|
"metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-456"},
|
|
}
|
|
],
|
|
}
|
|
])
|
|
|
|
result = extract_assessments_from_results(results_df, "run-123")
|
|
|
|
assert len(result) == 1
|
|
assert len(result[0].assessments) == 1
|
|
assert result[0].assessments[0].result is None
|
|
assert result[0].assessments[0].rationale is None
|
|
assert result[0].assessments[0].error is not None
|
|
|
|
|
|
def test_extract_multiple_assessments_from_same_run():
|
|
results_df = pd.DataFrame([
|
|
{
|
|
"trace_id": "tr-abc123",
|
|
"assessments": [
|
|
{
|
|
"assessment_name": "RelevanceToQuery",
|
|
"feedback": {"value": "yes"},
|
|
"rationale": "Relevant",
|
|
"metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
|
|
},
|
|
{
|
|
"assessment_name": "Safety",
|
|
"feedback": {"value": "yes"},
|
|
"rationale": "Safe",
|
|
"metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
|
|
},
|
|
],
|
|
}
|
|
])
|
|
|
|
result = extract_assessments_from_results(results_df, "run-123")
|
|
|
|
expected = [
|
|
EvalResult(
|
|
trace_id="tr-abc123",
|
|
assessments=[
|
|
Assessment(
|
|
name="RelevanceToQuery",
|
|
result="yes",
|
|
rationale="Relevant",
|
|
),
|
|
Assessment(
|
|
name="Safety",
|
|
result="yes",
|
|
rationale="Safe",
|
|
),
|
|
],
|
|
)
|
|
]
|
|
assert result == expected
|
|
|
|
|
|
def test_extract_no_assessments_on_trace_shows_error():
|
|
results_df = pd.DataFrame([{"trace_id": "tr-abc123", "assessments": []}])
|
|
|
|
result = extract_assessments_from_results(results_df, "run-123")
|
|
|
|
assert len(result) == 1
|
|
assert len(result[0].assessments) == 1
|
|
assert result[0].assessments[0].error == "No assessments found on trace"
|