@@ -5856,3 +5856,170 @@ async def test_async_generate_user_scenarios(self):
58565856 assert len (eval_dataset .eval_dataset_df ) == 2
58575857
58585858 self .mock_api_client .async_request .assert_called_once ()
5859+
5860+
5861+ class TestCreateEvaluationSetFromDataFrame :
5862+ """Unit tests for the _create_evaluation_set_from_dataframe function."""
5863+
5864+ def setup_method (self ):
5865+ self .mock_api_client = mock .Mock (spec = client .Client )
5866+ self .mock_api_client .project = "test-project"
5867+ self .mock_api_client .location = "us-central1"
5868+
5869+ @mock .patch .object (_evals_common , "evals" )
5870+ @mock .patch .object (_evals_common , "_gcs_utils" )
5871+ def test_create_evaluation_set_with_intermediate_events (
5872+ self , mock_gcs_utils , mock_evals_module
5873+ ):
5874+ intermediate_events = [
5875+ {
5876+ "content" : {"parts" : [{"text" : "thought 1" }]},
5877+ "timestamp" : "2024-01-01T00:00:00Z" ,
5878+ },
5879+ {
5880+ "content" : {"parts" : [{"functionCall" : {"name" : "foo" }}]},
5881+ "timestamp" : "2024-01-01T00:00:01Z" ,
5882+ },
5883+ ]
5884+
5885+ eval_df = pd .DataFrame (
5886+ [
5887+ {
5888+ "prompt" : "test prompt" ,
5889+ "response" : "test response" ,
5890+ "intermediate_events" : intermediate_events ,
5891+ }
5892+ ]
5893+ )
5894+
5895+ mock_gcs_instance = mock_gcs_utils .GcsUtils .return_value
5896+ mock_gcs_instance .upload_json_to_prefix .return_value = (
5897+ "gs://bucket/path/request.json"
5898+ )
5899+
5900+ mock_evals_instance = mock_evals_module .Evals .return_value
5901+ mock_eval_item = mock .Mock ()
5902+ mock_eval_item .name = "eval_item_1"
5903+ mock_evals_instance .create_evaluation_item .return_value = mock_eval_item
5904+
5905+ mock_eval_set = mock .Mock ()
5906+ mock_evals_instance .create_evaluation_set .return_value = mock_eval_set
5907+
5908+ result = _evals_common ._create_evaluation_set_from_dataframe (
5909+ api_client = self .mock_api_client ,
5910+ gcs_dest_prefix = "gs://bucket/prefix" ,
5911+ eval_df = eval_df ,
5912+ candidate_name = "test-candidate" ,
5913+ )
5914+
5915+ assert result == mock_eval_set
5916+
5917+ mock_gcs_instance .upload_json_to_prefix .assert_called_once ()
5918+ call_args = mock_gcs_instance .upload_json_to_prefix .call_args
5919+ uploaded_data = call_args .kwargs ["data" ]
5920+
5921+ candidate_responses = uploaded_data ["candidate_responses" ]
5922+ assert len (candidate_responses ) == 1
5923+ candidate_response = candidate_responses [0 ]
5924+ assert candidate_response ["candidate" ] == "test-candidate"
5925+ assert candidate_response ["text" ] == "test response"
5926+
5927+ expected_events = [
5928+ {"parts" : [{"text" : "thought 1" }]},
5929+ {"parts" : [{"function_call" : {"name" : "foo" }}]},
5930+ ]
5931+ assert candidate_response ["events" ] == expected_events
5932+
5933+ @mock .patch .object (_evals_common , "evals" )
5934+ @mock .patch .object (_evals_common , "_gcs_utils" )
5935+ def test_create_evaluation_set_with_user_scenario (
5936+ self , mock_gcs_utils , mock_evals_module
5937+ ):
5938+ eval_df = pd .DataFrame (
5939+ [
5940+ {
5941+ "starting_prompt" : "test starting prompt" ,
5942+ "conversation_plan" : "test conversation plan" ,
5943+ }
5944+ ]
5945+ )
5946+
5947+ mock_gcs_instance = mock_gcs_utils .GcsUtils .return_value
5948+ mock_gcs_instance .upload_json_to_prefix .return_value = (
5949+ "gs://bucket/path/request.json"
5950+ )
5951+
5952+ mock_evals_instance = mock_evals_module .Evals .return_value
5953+ mock_eval_item = mock .Mock ()
5954+ mock_eval_item .name = "eval_item_1"
5955+ mock_evals_instance .create_evaluation_item .return_value = mock_eval_item
5956+
5957+ mock_eval_set = mock .Mock ()
5958+ mock_evals_instance .create_evaluation_set .return_value = mock_eval_set
5959+
5960+ result = _evals_common ._create_evaluation_set_from_dataframe (
5961+ api_client = self .mock_api_client ,
5962+ gcs_dest_prefix = "gs://bucket/prefix" ,
5963+ eval_df = eval_df ,
5964+ candidate_name = "test-candidate" ,
5965+ )
5966+
5967+ assert result == mock_eval_set
5968+
5969+ mock_gcs_instance .upload_json_to_prefix .assert_called_once ()
5970+ call_args = mock_gcs_instance .upload_json_to_prefix .call_args
5971+ uploaded_data = call_args .kwargs ["data" ]
5972+
5973+ assert uploaded_data .get ("candidate_responses" ) is None
5974+ assert uploaded_data ["prompt" ]["user_scenario" ] == {
5975+ "starting_prompt" : "test starting prompt" ,
5976+ "conversation_plan" : "test conversation plan" ,
5977+ }
5978+
5979+ @mock .patch .object (_evals_common , "evals" )
5980+ @mock .patch .object (_evals_common , "_gcs_utils" )
5981+ def test_create_evaluation_set_with_agent_data (
5982+ self , mock_gcs_utils , mock_evals_module
5983+ ):
5984+ agent_data = {"turns" : [{"turn_id" : "turn1" , "events" : []}]}
5985+ eval_df = pd .DataFrame (
5986+ [
5987+ {
5988+ "prompt" : "test prompt" ,
5989+ "agent_data" : agent_data ,
5990+ }
5991+ ]
5992+ )
5993+
5994+ mock_gcs_instance = mock_gcs_utils .GcsUtils .return_value
5995+ mock_gcs_instance .upload_json_to_prefix .return_value = (
5996+ "gs://bucket/path/request.json"
5997+ )
5998+
5999+ mock_evals_instance = mock_evals_module .Evals .return_value
6000+ mock_eval_item = mock .Mock ()
6001+ mock_eval_item .name = "eval_item_1"
6002+ mock_evals_instance .create_evaluation_item .return_value = mock_eval_item
6003+
6004+ mock_eval_set = mock .Mock ()
6005+ mock_evals_instance .create_evaluation_set .return_value = mock_eval_set
6006+
6007+ result = _evals_common ._create_evaluation_set_from_dataframe (
6008+ api_client = self .mock_api_client ,
6009+ gcs_dest_prefix = "gs://bucket/prefix" ,
6010+ eval_df = eval_df ,
6011+ candidate_name = "test-candidate" ,
6012+ )
6013+
6014+ assert result == mock_eval_set
6015+
6016+ mock_gcs_instance .upload_json_to_prefix .assert_called_once ()
6017+ call_args = mock_gcs_instance .upload_json_to_prefix .call_args
6018+ uploaded_data = call_args .kwargs ["data" ]
6019+
6020+ assert uploaded_data ["prompt" ]["text" ] == "test prompt"
6021+ candidate_responses = uploaded_data ["candidate_responses" ]
6022+ assert len (candidate_responses ) == 1
6023+ candidate_response = candidate_responses [0 ]
6024+ assert candidate_response ["candidate" ] == "test-candidate"
6025+ assert candidate_response ["agent_data" ] == agent_data
0 commit comments