feat: GenAI Client(evals): Update to enable agent run functionality in eval management service

PiperOrigin-RevId: 877600757

Author: vertex-sdk-bot

tests/unit/vertexai/genai/test_evals.py

@@ -5822,3 +5822,264 @@ async def test_async_generate_user_scenarios(self):
         assert len(eval_dataset.eval_dataset_df) == 2
 
         self.mock_api_client.async_request.assert_called_once()
+
+
+class TestCreateEvaluationSetFromDataFrame:
+    """Unit tests for the _create_evaluation_set_from_dataframe function."""
+
+    def setup_method(self):
+        self.mock_api_client = mock.Mock(spec=client.Client)
+        self.mock_api_client.project = "test-project"
+        self.mock_api_client.location = "us-central1"
+
+    @mock.patch.object(_evals_common, "evals")
+    @mock.patch.object(_evals_common, "_gcs_utils")
+    def test_create_evaluation_set_with_intermediate_events(
+        self, mock_gcs_utils, mock_evals_module
+    ):
+        intermediate_events = [
+            {
+                "content": {"parts": [{"text": "thought 1"}]},
+                "timestamp": "2024-01-01T00:00:00Z",
+            },
+            {
+                "content": {"parts": [{"functionCall": {"name": "foo"}}]},
+                "timestamp": "2024-01-01T00:00:01Z",
+            },
+        ]
+
+        eval_df = pd.DataFrame(
+            [
+                {
+                    "prompt": "test prompt",
+                    "response": "test response",
+                    "intermediate_events": intermediate_events,
+                }
+            ]
+        )
+
+        mock_gcs_instance = mock_gcs_utils.GcsUtils.return_value
+        mock_gcs_instance.upload_json_to_prefix.return_value = (
+            "gs://bucket/path/request.json"
+        )
+
+        mock_evals_instance = mock_evals_module.Evals.return_value
+        mock_eval_item = mock.Mock()
+        mock_eval_item.name = "eval_item_1"
+        mock_evals_instance.create_evaluation_item.return_value = mock_eval_item
+
+        mock_eval_set = mock.Mock()
+        mock_evals_instance.create_evaluation_set.return_value = mock_eval_set
+
+        result = _evals_common._create_evaluation_set_from_dataframe(
+            api_client=self.mock_api_client,
+            gcs_dest_prefix="gs://bucket/prefix",
+            eval_df=eval_df,
+            candidate_name="test-candidate",
+        )
+
+        assert result == mock_eval_set
+
+        mock_gcs_instance.upload_json_to_prefix.assert_called_once()
+        call_args = mock_gcs_instance.upload_json_to_prefix.call_args
+        uploaded_data = call_args.kwargs["data"]
+
+        candidate_responses = uploaded_data["candidate_responses"]
+        assert len(candidate_responses) == 1
+        candidate_response = candidate_responses[0]
+        assert candidate_response["candidate"] == "test-candidate"
+        assert candidate_response["text"] == "test response"
+
+        expected_events = [
+            {"parts": [{"text": "thought 1"}]},
+            {"parts": [{"function_call": {"name": "foo"}}]},
+        ]
+        assert candidate_response["events"] == expected_events
+
+    @mock.patch.object(_evals_common, "evals")
+    @mock.patch.object(_evals_common, "_gcs_utils")
+    def test_create_evaluation_set_with_user_scenario(
+        self, mock_gcs_utils, mock_evals_module
+    ):
+        eval_df = pd.DataFrame(
+            [
+                {
+                    "starting_prompt": "test starting prompt",
+                    "conversation_plan": "test conversation plan",
+                }
+            ]
+        )
+
+        mock_gcs_instance = mock_gcs_utils.GcsUtils.return_value
+        mock_gcs_instance.upload_json_to_prefix.return_value = (
+            "gs://bucket/path/request.json"
+        )
+
+        mock_evals_instance = mock_evals_module.Evals.return_value
+        mock_eval_item = mock.Mock()
+        mock_eval_item.name = "eval_item_1"
+        mock_evals_instance.create_evaluation_item.return_value = mock_eval_item
+
+        mock_eval_set = mock.Mock()
+        mock_evals_instance.create_evaluation_set.return_value = mock_eval_set
+
+        result = _evals_common._create_evaluation_set_from_dataframe(
+            api_client=self.mock_api_client,
+            gcs_dest_prefix="gs://bucket/prefix",
+            eval_df=eval_df,
+            candidate_name="test-candidate",
+        )
+
+        assert result == mock_eval_set
+
+        mock_gcs_instance.upload_json_to_prefix.assert_called_once()
+        call_args = mock_gcs_instance.upload_json_to_prefix.call_args
+        uploaded_data = call_args.kwargs["data"]
+
+        assert uploaded_data.get("candidate_responses") is None
+        assert uploaded_data["prompt"]["user_scenario"] == {
+            "starting_prompt": "test starting prompt",
+            "conversation_plan": "test conversation plan",
+        }
+
+    @mock.patch.object(_evals_common, "evals")
+    @mock.patch.object(_evals_common, "_gcs_utils")
+    def test_create_evaluation_set_with_agent_data(
+        self, mock_gcs_utils, mock_evals_module
+    ):
+        agent_data = {"turns": [{"turn_id": "turn1", "events": []}]}
+        eval_df = pd.DataFrame(
+            [
+                {
+                    "prompt": "test prompt",
+                    "agent_data": agent_data,
+                }
+            ]
+        )
+
+        mock_gcs_instance = mock_gcs_utils.GcsUtils.return_value
+        mock_gcs_instance.upload_json_to_prefix.return_value = (
+            "gs://bucket/path/request.json"
+        )
+
+        mock_evals_instance = mock_evals_module.Evals.return_value
+        mock_eval_item = mock.Mock()
+        mock_eval_item.name = "eval_item_1"
+        mock_evals_instance.create_evaluation_item.return_value = mock_eval_item
+
+        mock_eval_set = mock.Mock()
+        mock_evals_instance.create_evaluation_set.return_value = mock_eval_set
+
+        result = _evals_common._create_evaluation_set_from_dataframe(
+            api_client=self.mock_api_client,
+            gcs_dest_prefix="gs://bucket/prefix",
+            eval_df=eval_df,
+            candidate_name="test-candidate",
+        )
+
+        assert result == mock_eval_set
+
+        mock_gcs_instance.upload_json_to_prefix.assert_called_once()
+        call_args = mock_gcs_instance.upload_json_to_prefix.call_args
+        uploaded_data = call_args.kwargs["data"]
+
+        assert uploaded_data["prompt"]["text"] == "test prompt"
+        candidate_responses = uploaded_data["candidate_responses"]
+        assert len(candidate_responses) == 1
+        candidate_response = candidate_responses[0]
+        assert candidate_response["candidate"] == "test-candidate"
+        assert candidate_response["agent_data"] == agent_data
+
+    @mock.patch.object(_evals_common, "evals")
+    @mock.patch.object(_evals_common, "_gcs_utils")
+    def test_create_evaluation_set_with_user_scenario(
+        self, mock_gcs_utils, mock_evals_module
+    ):
+        eval_df = pd.DataFrame(
+            [
+                {
+                    "starting_prompt": "test starting prompt",
+                    "conversation_plan": "test conversation plan",
+                }
+            ]
+        )
+
+        mock_gcs_instance = mock_gcs_utils.GcsUtils.return_value
+        mock_gcs_instance.upload_json_to_prefix.return_value = (
+            "gs://bucket/path/request.json"
+        )
+
+        mock_evals_instance = mock_evals_module.Evals.return_value
+        mock_eval_item = mock.Mock()
+        mock_eval_item.name = "eval_item_1"
+        mock_evals_instance.create_evaluation_item.return_value = mock_eval_item
+
+        mock_eval_set = mock.Mock()
+        mock_evals_instance.create_evaluation_set.return_value = mock_eval_set
+
+        result = _evals_common._create_evaluation_set_from_dataframe(
+            api_client=self.mock_api_client,
+            gcs_dest_prefix="gs://bucket/prefix",
+            eval_df=eval_df,
+            candidate_name="test-candidate",
+        )
+
+        assert result == mock_eval_set
+
+        mock_gcs_instance.upload_json_to_prefix.assert_called_once()
+        call_args = mock_gcs_instance.upload_json_to_prefix.call_args
+        uploaded_data = call_args.kwargs["data"]
+
+        assert uploaded_data.get("candidate_responses") is None
+        assert uploaded_data["prompt"]["user_scenario"] == {
+            "starting_prompt": "test starting prompt",
+            "conversation_plan": "test conversation plan",
+        }
+
+    @mock.patch.object(_evals_common, "evals")
+    @mock.patch.object(_evals_common, "_gcs_utils")
+    def test_create_evaluation_set_with_agent_data(
+        self, mock_gcs_utils, mock_evals_module
+    ):
+        agent_data = {"turns": [{"turn_id": "turn1", "events": []}]}
+        eval_df = pd.DataFrame(
+            [
+                {
+                    "prompt": "test prompt",
+                    "agent_data": agent_data,
+                }
+            ]
+        )
+
+        mock_gcs_instance = mock_gcs_utils.GcsUtils.return_value
+        mock_gcs_instance.upload_json_to_prefix.return_value = (
+            "gs://bucket/path/request.json"
+        )
+
+        mock_evals_instance = mock_evals_module.Evals.return_value
+        mock_eval_item = mock.Mock()
+        mock_eval_item.name = "eval_item_1"
+        mock_evals_instance.create_evaluation_item.return_value = mock_eval_item
+
+        mock_eval_set = mock.Mock()
+        mock_evals_instance.create_evaluation_set.return_value = mock_eval_set
+
+        result = _evals_common._create_evaluation_set_from_dataframe(
+            api_client=self.mock_api_client,
+            gcs_dest_prefix="gs://bucket/prefix",
+            eval_df=eval_df,
+            candidate_name="test-candidate",
+        )
+
+        assert result == mock_eval_set
+
+        mock_gcs_instance.upload_json_to_prefix.assert_called_once()
+        call_args = mock_gcs_instance.upload_json_to_prefix.call_args
+        uploaded_data = call_args.kwargs["data"]
+
+        assert uploaded_data["prompt"]["text"] == "test prompt"
+        candidate_responses = uploaded_data["candidate_responses"]
+        assert len(candidate_responses) == 1
+        candidate_response = candidate_responses[0]
+        assert candidate_response["candidate"] == "test-candidate"
+        assert candidate_response["agent_data"] == agent_data

vertexai/_genai/_evals_common.py

@@ -35,6 +35,7 @@
 from google.genai.models import Models
 import pandas as pd
 from tqdm import tqdm
+from pydantic import ValidationError
 
 from . import _evals_constant
 from . import _evals_data_converters
@@ -2247,7 +2248,52 @@ def _create_evaluation_set_from_dataframe(
             for event in row[_evals_constant.INTERMEDIATE_EVENTS]:
                 if CONTENT in event:
                     intermediate_events.append(event[CONTENT])
-        if _evals_constant.CONTEXT in row or _evals_constant.HISTORY in row:
+
+        agent_data_obj = None
+        if _evals_constant.AGENT_DATA in row:
+            agent_data_val = row[AGENT_DATA]
+            if isinstance(agent_data_val, str):
+                try:
+                    agent_data_val = json.loads(agent_data_val)
+                except json.JSONDecodeError:
+                    pass
+            if isinstance(agent_data_val, dict):
+                try:
+                    agent_data_obj = types.evals.AgentData.model_validate(
+                        agent_data_val
+                    )
+                except ValidationError:
+                    pass
+            elif isinstance(agent_data_val, types.evals.AgentData):
+                agent_data_obj = agent_data_val
+
+        candidate_responses = []
+        if (
+            _evals_constant.RESPONSE in row
+            or agent_data_obj
+            or intermediate_events
+        ):
+            candidate_responses.append(
+                types.CandidateResponse(
+                    candidate=candidate_name or "Candidate 1",
+                    text=row.get(_evals_constant.RESPONSE) or None,
+                    events=intermediate_events or None,
+                    agent_data=agent_data_obj,
+                )
+            )
+
+        prompt = None
+        if (
+            _evals_constant.STARTING_PROMPT in row
+            and _evals_constant.CONVERSATION_PLAN in row
+        ):
+            prompt = types.EvaluationPrompt(
+                user_scenario=types.evals.UserScenario(
+                    starting_prompt=row[_evals_constant.STARTING_PROMPT],
+                    conversation_plan=row[_evals_constant.CONVERSATION_PLAN],
+                )
+            )
+        elif _evals_constant.CONTEXT in row or _evals_constant.HISTORY in row:
             values = {}
             if _evals_constant.CONTEXT in row:
                 values[_evals_constant.CONTEXT] = _get_content(
@@ -2266,15 +2312,7 @@ def _create_evaluation_set_from_dataframe(
             )
         elif _evals_constant.PROMPT in row:
             prompt = types.EvaluationPrompt(text=row[_evals_constant.PROMPT])
-        candidate_responses = []
-        if _evals_constant.RESPONSE in row:
-            candidate_responses.append(
-                types.CandidateResponse(
-                    candidate=candidate_name or "Candidate 1",
-                    text=row[_evals_constant.RESPONSE],
-                    events=intermediate_events or None,
-                )
-            )
+
         eval_item_requests.append(
             types.EvaluationItemRequest(
                 prompt=prompt or None,
@@ -2283,7 +2321,9 @@ def _create_evaluation_set_from_dataframe(
                     if _evals_constant.REFERENCE in row
                     else None
                 ),
-                candidate_responses=candidate_responses,
+                candidate_responses=(
+                    candidate_responses if candidate_responses else None
+                ),
             )
         )
     logger.info("Writing evaluation item requests to GCS.")

vertexai/_genai/_evals_constant.py

@@ -57,6 +57,8 @@
 PARTS = "parts"
 USER_AUTHOR = "user"
 AGENT_DATA = "agent_data"
+STARTING_PROMPT = "starting_prompt"
+CONVERSATION_PLAN = "conversation_plan"
 HISTORY = "conversation_history"
 
 COMMON_DATASET_COLUMNS = frozenset(

vertexai/_genai/types/__init__.py

@@ -130,6 +130,9 @@
 from .common import AgentEngineSessionOperation
 from .common import AgentEngineSessionOperationDict
 from .common import AgentEngineSessionOperationOrDict
+from .common import AgentRunConfig
+from .common import AgentRunConfigDict
+from .common import AgentRunConfigOrDict
 from .common import AgentServerMode
 from .common import AggregatedMetricResult
 from .common import AggregatedMetricResultDict
@@ -1391,6 +1394,9 @@
     "EvaluationRunAgentConfig",
     "EvaluationRunAgentConfigDict",
     "EvaluationRunAgentConfigOrDict",
+    "AgentRunConfig",
+    "AgentRunConfigDict",
+    "AgentRunConfigOrDict",
     "EvaluationRunInferenceConfig",
     "EvaluationRunInferenceConfigDict",
     "EvaluationRunInferenceConfigOrDict",