1  from mlflow.entities.dataset_record_source import DatasetRecordSource, DatasetRecordSourceType
 2  from mlflow.genai.evaluation.entities import EvalItem
 3  
 4  
 5  def test_eval_item_from_dataset_row_extracts_source():
 6      source = DatasetRecordSource(
 7          source_type=DatasetRecordSourceType.TRACE,
 8          source_data={"trace_id": "tr-123", "session_id": "session_1"},
 9      )
10  
11      row = {
12          "inputs": {"question": "test"},
13          "outputs": "answer",
14          "expectations": {},
15          "source": source,
16      }
17  
18      eval_item = EvalItem.from_dataset_row(row)
19  
20      assert eval_item.source == source
21      assert eval_item.source.source_data["session_id"] == "session_1"
22      assert eval_item.inputs == {"question": "test"}
23      assert eval_item.outputs == "answer"
24  
25  
26  def test_eval_item_from_dataset_row_handles_missing_source():
27      row = {
28          "inputs": {"question": "test"},
29          "outputs": "answer",
30          "expectations": {},
31      }
32  
33      eval_item = EvalItem.from_dataset_row(row)
34  
35      assert eval_item.source is None
36      assert eval_item.inputs == {"question": "test"}
37      assert eval_item.outputs == "answer"