from unittest.mock import Mock, patch import pytest import mlflow from mlflow.entities import TraceData, TraceInfo, TraceLocation, TraceState from mlflow.entities.assessment import Feedback from mlflow.entities.assessment_source import AssessmentSource, AssessmentSourceType from mlflow.entities.trace import Trace from mlflow.exceptions import MlflowException from mlflow.genai import scorer from mlflow.genai.evaluation.entities import EvalItem from mlflow.genai.evaluation.session_utils import ( classify_scorers, evaluate_session_level_scorers, get_first_trace_in_session, group_traces_by_session, validate_session_level_evaluation_inputs, ) from mlflow.tracing.constant import TraceMetadataKey class _MultiTurnTestScorer: """Helper class for testing multi-turn scorers.""" def __init__(self, name="test_multi_turn_scorer"): self.name = name self.is_session_level_scorer = True self.aggregations = [] def run(self, session=None, **kwargs): return True def __call__(self, traces=None, **kwargs): return 1.0 # ==================== Tests for classify_scorers ==================== def test_classify_scorers_all_single_turn(): @scorer def custom_scorer1(outputs): return 1.0 @scorer def custom_scorer2(outputs): return 2.0 scorers_list = [custom_scorer1, custom_scorer2] single_turn, multi_turn = classify_scorers(scorers_list) assert len(single_turn) == 2 assert len(multi_turn) == 0 assert single_turn == scorers_list def test_classify_scorers_all_multi_turn(): multi_turn_scorer1 = _MultiTurnTestScorer(name="multi_turn_scorer1") multi_turn_scorer2 = _MultiTurnTestScorer(name="multi_turn_scorer2") scorers_list = [multi_turn_scorer1, multi_turn_scorer2] single_turn, multi_turn = classify_scorers(scorers_list) assert len(single_turn) == 0 assert len(multi_turn) == 2 assert multi_turn == scorers_list # Verify they are actually multi-turn assert multi_turn_scorer1.is_session_level_scorer is True assert multi_turn_scorer2.is_session_level_scorer is True def test_classify_scorers_mixed(): @scorer def single_turn_scorer(outputs): return 1.0 multi_turn_scorer = _MultiTurnTestScorer(name="multi_turn_scorer") scorers_list = [single_turn_scorer, multi_turn_scorer] single_turn, multi_turn = classify_scorers(scorers_list) assert len(single_turn) == 1 assert len(multi_turn) == 1 assert single_turn[0] == single_turn_scorer assert multi_turn[0] == multi_turn_scorer # Verify properties assert single_turn_scorer.is_session_level_scorer is False assert multi_turn_scorer.is_session_level_scorer is True def test_classify_scorers_empty_list(): single_turn, multi_turn = classify_scorers([]) assert len(single_turn) == 0 assert len(multi_turn) == 0 # ==================== Tests for group_traces_by_session ==================== def _create_mock_trace(trace_id: str, session_id: str | None, request_time: int): """Helper to create a mock trace with session_id and request_time.""" trace_metadata = {} if session_id is not None: trace_metadata[TraceMetadataKey.TRACE_SESSION] = session_id trace_info = TraceInfo( trace_id=trace_id, trace_location=TraceLocation.from_experiment_id("0"), request_time=request_time, execution_duration=1000, state=TraceState.OK, trace_metadata=trace_metadata, tags={}, ) trace = Mock(spec=Trace) trace.info = trace_info trace.data = TraceData(spans=[]) return trace def _create_mock_eval_item(trace): """Helper to create a mock EvalItem with a trace.""" eval_item = Mock(spec=EvalItem) eval_item.trace = trace eval_item.source = None # Explicitly set to None so it doesn't return a Mock return eval_item def test_group_traces_by_session_single_session(): trace1 = _create_mock_trace("trace-1", "session-1", 1000) trace2 = _create_mock_trace("trace-2", "session-1", 2000) trace3 = _create_mock_trace("trace-3", "session-1", 3000) eval_item1 = _create_mock_eval_item(trace1) eval_item2 = _create_mock_eval_item(trace2) eval_item3 = _create_mock_eval_item(trace3) eval_items = [eval_item1, eval_item2, eval_item3] session_groups = group_traces_by_session(eval_items) assert len(session_groups) == 1 assert "session-1" in session_groups assert len(session_groups["session-1"]) == 3 # Check that all traces are included session_traces = [item.trace for item in session_groups["session-1"]] assert trace1 in session_traces assert trace2 in session_traces assert trace3 in session_traces def test_group_traces_by_session_multiple_sessions(): trace1 = _create_mock_trace("trace-1", "session-1", 1000) trace2 = _create_mock_trace("trace-2", "session-1", 2000) trace3 = _create_mock_trace("trace-3", "session-2", 1500) trace4 = _create_mock_trace("trace-4", "session-2", 2500) eval_items = [ _create_mock_eval_item(trace1), _create_mock_eval_item(trace2), _create_mock_eval_item(trace3), _create_mock_eval_item(trace4), ] session_groups = group_traces_by_session(eval_items) assert len(session_groups) == 2 assert "session-1" in session_groups assert "session-2" in session_groups assert len(session_groups["session-1"]) == 2 assert len(session_groups["session-2"]) == 2 def test_group_traces_by_session_excludes_no_session_id(): trace1 = _create_mock_trace("trace-1", "session-1", 1000) trace2 = _create_mock_trace("trace-2", None, 2000) # No session_id trace3 = _create_mock_trace("trace-3", "session-1", 3000) eval_items = [ _create_mock_eval_item(trace1), _create_mock_eval_item(trace2), _create_mock_eval_item(trace3), ] session_groups = group_traces_by_session(eval_items) assert len(session_groups) == 1 assert "session-1" in session_groups assert len(session_groups["session-1"]) == 2 # trace2 should not be included session_traces = [item.trace for item in session_groups["session-1"]] assert trace1 in session_traces assert trace2 not in session_traces assert trace3 in session_traces def test_group_traces_by_session_excludes_none_traces(): trace1 = _create_mock_trace("trace-1", "session-1", 1000) eval_item1 = _create_mock_eval_item(trace1) eval_item2 = Mock() eval_item2.trace = None # No trace eval_item2.source = None # No source eval_items = [eval_item1, eval_item2] session_groups = group_traces_by_session(eval_items) assert len(session_groups) == 1 assert "session-1" in session_groups assert len(session_groups["session-1"]) == 1 def test_group_traces_by_session_empty_list(): session_groups = group_traces_by_session([]) assert len(session_groups) == 0 assert session_groups == {} # ==================== Tests for get_first_trace_in_session ==================== def test_get_first_trace_in_session_chronological_order(): trace1 = _create_mock_trace("trace-1", "session-1", 3000) trace2 = _create_mock_trace("trace-2", "session-1", 1000) # Earliest trace3 = _create_mock_trace("trace-3", "session-1", 2000) eval_item1 = _create_mock_eval_item(trace1) eval_item2 = _create_mock_eval_item(trace2) eval_item3 = _create_mock_eval_item(trace3) session_items = [eval_item1, eval_item2, eval_item3] first_item = get_first_trace_in_session(session_items) assert first_item.trace == trace2 assert first_item == eval_item2 def test_get_first_trace_in_session_single_trace(): trace1 = _create_mock_trace("trace-1", "session-1", 1000) eval_item1 = _create_mock_eval_item(trace1) session_items = [eval_item1] first_item = get_first_trace_in_session(session_items) assert first_item.trace == trace1 assert first_item == eval_item1 def test_get_first_trace_in_session_same_timestamp(): # When timestamps are equal, min() will return the first one in the list trace1 = _create_mock_trace("trace-1", "session-1", 1000) trace2 = _create_mock_trace("trace-2", "session-1", 1000) trace3 = _create_mock_trace("trace-3", "session-1", 1000) eval_item1 = _create_mock_eval_item(trace1) eval_item2 = _create_mock_eval_item(trace2) eval_item3 = _create_mock_eval_item(trace3) session_items = [eval_item1, eval_item2, eval_item3] first_item = get_first_trace_in_session(session_items) # Should return one of the traces with timestamp 1000 (likely the first one) assert first_item.trace.info.request_time == 1000 # ==================== Tests for validate_session_level_evaluation_inputs ==================== def test_validate_session_level_evaluation_inputs_no_session_level_scorers(): @scorer def single_turn_scorer(outputs): return 1.0 scorers_list = [single_turn_scorer] # Should not raise any exceptions validate_session_level_evaluation_inputs( scorers=scorers_list, predict_fn=None, ) def test_validate_session_level_evaluation_inputs_with_predict_fn(): multi_turn_scorer = _MultiTurnTestScorer() scorers_list = [multi_turn_scorer] def dummy_predict_fn(): return "output" with pytest.raises( MlflowException, match=r"Session-level scorers require traces with session IDs.*" r"Either pass a ConversationSimulator to `data` with `predict_fn`", ): validate_session_level_evaluation_inputs( scorers=scorers_list, predict_fn=dummy_predict_fn, ) def test_validate_session_level_evaluation_inputs_mixed_scorers(): @scorer def single_turn_scorer(outputs): return 1.0 multi_turn_scorer = _MultiTurnTestScorer() scorers_list = [single_turn_scorer, multi_turn_scorer] # Should not raise any exceptions validate_session_level_evaluation_inputs( scorers=scorers_list, predict_fn=None, ) # ==================== Tests for evaluate_session_level_scorers ==================== def _create_test_trace(trace_id: str, request_time: int = 0) -> Trace: """Helper to create a minimal test trace""" return Trace( info=TraceInfo( trace_id=trace_id, trace_location=TraceLocation.from_experiment_id("0"), request_time=request_time, execution_duration=100, state=TraceState.OK, trace_metadata={}, tags={}, ), data=TraceData(spans=[]), ) def _create_eval_item(trace_id: str, request_time: int = 0) -> EvalItem: """Helper to create a minimal EvalItem with a trace""" trace = _create_test_trace(trace_id, request_time) return EvalItem( request_id=trace_id, trace=trace, inputs={}, outputs={}, expectations={}, ) def test_evaluate_session_level_scorers_success(): mock_scorer = Mock(spec=mlflow.genai.Scorer) mock_scorer.name = "test_scorer" mock_scorer.run.return_value = 0.8 # Test with a single session containing multiple traces session_items = [ _create_eval_item("trace1", request_time=100), _create_eval_item("trace2", request_time=200), ] with patch( "mlflow.genai.evaluation.session_utils.standardize_scorer_value" ) as mock_standardize: # Return a new Feedback object each time to avoid metadata overwriting def create_feedback(*args, **kwargs): return [ Feedback( name="test_scorer", source=AssessmentSource( source_type=AssessmentSourceType.CODE, source_id="test" ), value=0.8, ) ] mock_standardize.side_effect = create_feedback result = evaluate_session_level_scorers("session1", session_items, [mock_scorer]) # Verify scorer was called once (for the single session) assert mock_scorer.run.call_count == 1 # Verify scorer received session traces call_args = mock_scorer.run.call_args assert "session" in call_args.kwargs assert len(call_args.kwargs["session"]) == 2 # session has 2 traces # Verify result is for first item assert result.eval_item.trace.info.trace_id == "trace1" assert len(result.assessments) == 1 assert result.assessments[0].name == "test_scorer" assert result.assessments[0].value == 0.8 # Verify session_id was added to metadata assert result.assessments[0].metadata is not None assert result.assessments[0].metadata[TraceMetadataKey.TRACE_SESSION] == "session1" def test_evaluate_session_level_scorers_handles_scorer_error(): mock_scorer = Mock(spec=mlflow.genai.Scorer) mock_scorer.name = "failing_scorer" mock_scorer.run.side_effect = ValueError("Scorer failed!") session_items = [_create_eval_item("trace1", 100)] result = evaluate_session_level_scorers("session1", session_items, [mock_scorer]) # Verify error feedback was created assert result.eval_item.trace.info.trace_id == "trace1" assert len(result.assessments) == 1 feedback = result.assessments[0] assert feedback.name == "failing_scorer" assert feedback.error is not None assert feedback.error.error_code == "SCORER_ERROR" assert feedback.error.stack_trace is not None assert feedback.error.to_proto().error_message == "Scorer failed!" assert isinstance(feedback.error.error_message, str) assert feedback.error.error_message == "Scorer failed!" # Verify session_id metadata is present even on error feedbacks assert feedback.metadata is not None assert feedback.metadata[TraceMetadataKey.TRACE_SESSION] == "session1" def test_evaluate_session_level_scorers_multiple_feedbacks_per_scorer(): mock_scorer = Mock(spec=mlflow.genai.Scorer) mock_scorer.name = "multi_feedback_scorer" mock_scorer.run.return_value = {"metric1": 0.7, "metric2": 0.9} session_items = [_create_eval_item("trace1", 100)] with patch( "mlflow.genai.evaluation.session_utils.standardize_scorer_value" ) as mock_standardize: feedbacks = [ Feedback( name="multi_feedback_scorer/metric1", source=AssessmentSource(source_type=AssessmentSourceType.CODE, source_id="test"), value=0.7, ), Feedback( name="multi_feedback_scorer/metric2", source=AssessmentSource(source_type=AssessmentSourceType.CODE, source_id="test"), value=0.9, ), ] mock_standardize.return_value = feedbacks result = evaluate_session_level_scorers("session1", session_items, [mock_scorer]) # Verify both feedbacks are stored assert result.eval_item.trace.info.trace_id == "trace1" assert len(result.assessments) == 2 # Find feedbacks by name feedback_by_name = {f.name: f for f in result.assessments} assert "multi_feedback_scorer/metric1" in feedback_by_name assert "multi_feedback_scorer/metric2" in feedback_by_name assert feedback_by_name["multi_feedback_scorer/metric1"].value == 0.7 assert feedback_by_name["multi_feedback_scorer/metric2"].value == 0.9 def test_evaluate_session_level_scorers_first_trace_selection(): mock_scorer = Mock(spec=mlflow.genai.Scorer) mock_scorer.name = "first_trace_scorer" mock_scorer.run.return_value = 1.0 # Create session with traces in non-chronological order session_items = [ _create_eval_item("trace2", request_time=200), # Second chronologically _create_eval_item("trace1", request_time=100), # First chronologically _create_eval_item("trace3", request_time=300), # Third chronologically ] with patch( "mlflow.genai.evaluation.session_utils.standardize_scorer_value" ) as mock_standardize: feedback = Feedback( name="first_trace_scorer", source=AssessmentSource(source_type=AssessmentSourceType.CODE, source_id="test"), value=1.0, ) mock_standardize.return_value = [feedback] result = evaluate_session_level_scorers("session1", session_items, [mock_scorer]) # Verify assessment is for trace1 (earliest request_time) assert result.eval_item.trace.info.trace_id == "trace1" assert len(result.assessments) == 1 assert result.assessments[0].name == "first_trace_scorer" assert result.assessments[0].value == 1.0 def test_evaluate_session_level_scorers_multiple_scorers(): mock_scorer1 = Mock(spec=mlflow.genai.Scorer) mock_scorer1.name = "scorer1" mock_scorer1.run.return_value = 0.6 mock_scorer2 = Mock(spec=mlflow.genai.Scorer) mock_scorer2.name = "scorer2" mock_scorer2.run.return_value = 0.8 session_items = [_create_eval_item("trace1", 100)] with patch( "mlflow.genai.evaluation.session_utils.standardize_scorer_value" ) as mock_standardize: def create_feedback(name, value): return [ Feedback( name=name, source=AssessmentSource( source_type=AssessmentSourceType.CODE, source_id="test" ), value=value, ) ] mock_standardize.side_effect = [ create_feedback("scorer1", 0.6), create_feedback("scorer2", 0.8), ] result = evaluate_session_level_scorers( "session1", session_items, [mock_scorer1, mock_scorer2] ) # Verify both scorers were evaluated (runs in parallel) assert mock_scorer1.run.call_count == 1 assert mock_scorer2.run.call_count == 1 # Verify result contains assessments from both scorers assert result.eval_item.trace.info.trace_id == "trace1" assert len(result.assessments) == 2 # Find feedbacks by name feedback_by_name = {f.name: f for f in result.assessments} assert "scorer1" in feedback_by_name assert "scorer2" in feedback_by_name assert feedback_by_name["scorer1"].value == 0.6 assert feedback_by_name["scorer2"].value == 0.8 def test_evaluate_session_level_scorers_error_multiple_traces(): mock_scorer = Mock(spec=mlflow.genai.Scorer) mock_scorer.name = "failing_scorer" mock_scorer.run.side_effect = RuntimeError("boom") session_items = [ _create_eval_item("trace1", request_time=100), _create_eval_item("trace2", request_time=200), ] result = evaluate_session_level_scorers("session-abc", session_items, [mock_scorer]) assert result.eval_item.trace.info.trace_id == "trace1" feedback = result.assessments[0] assert feedback.error is not None assert feedback.metadata[TraceMetadataKey.TRACE_SESSION] == "session-abc"