diff --git a/pyproject.toml b/pyproject.toml index 86e60307..96168132 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,8 +26,8 @@ classifiers = [ "Topic :: Software Development :: Libraries :: Python Modules", ] dependencies = [ - "boto3>=1.43.0", - "botocore>=1.43.0", + "boto3>=1.43.31", + "botocore>=1.43.31", "pydantic>=2.0.0,<2.41.3", "urllib3>=1.26.0", "starlette>=0.46.2", diff --git a/src/bedrock_agentcore/evaluation/__init__.py b/src/bedrock_agentcore/evaluation/__init__.py index ffa88401..fece1a1a 100644 --- a/src/bedrock_agentcore/evaluation/__init__.py +++ b/src/bedrock_agentcore/evaluation/__init__.py @@ -18,6 +18,7 @@ EvaluatorStatistics, EvaluatorSummary, FailedScenario, + OnlineEvaluationDataSourceConfig, ) from bedrock_agentcore.evaluation.runner.batch.batch_evaluation_runner import ( BatchEvaluationRunner, @@ -71,6 +72,7 @@ "BatchEvaluationRunConfig", "CloudWatchOutputDataConfig", "CloudWatchDataSourceConfig", + "OnlineEvaluationDataSourceConfig", "BatchEvaluatorConfig", "BatchEvaluationSummary", "EvaluatorStatistics", diff --git a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_models.py b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_models.py index 3e5b93b1..d8138ed9 100644 --- a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_models.py +++ b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_models.py @@ -111,6 +111,46 @@ def to_data_source_config( } +class OnlineEvaluationDataSourceConfig(BaseModel, DataSourceConfig): + """Online-evaluation data source — pulls spans from an existing OnlineEvaluationConfig. + + .. warning:: + This feature is in preview and may change in future releases. + + Unlike :class:`CloudWatchDataSourceConfig`, this source does not filter by the + session IDs generated during agent invocation: the service reads sessions already + captured by the referenced OnlineEvaluationConfig, optionally narrowed to a time + window. As a result it is typically used to (re-)evaluate previously recorded + sessions rather than sessions produced by the current ``agent_invoker`` run. + + Attributes: + online_evaluation_config_arn: ARN of the OnlineEvaluationConfig whose + captured sessions are evaluated. + use_invocation_time_range: When ``True`` (default), the runner supplies the + ``sessionFilterConfig`` time window from the earliest/latest session times + observed during agent invocation. Set to ``False`` to omit the window and + let the service use the OnlineEvaluationConfig's own session selection. + """ + + online_evaluation_config_arn: str = Field(min_length=1) + use_invocation_time_range: bool = True + + def to_data_source_config( + self, + session_ids: List[str], + start_time: datetime, + end_time: datetime, + ) -> Dict[str, Any]: + """Return an onlineEvaluationConfigSource dataSourceConfig dict for the evaluation API.""" + source: Dict[str, Any] = {"onlineEvaluationConfigArn": self.online_evaluation_config_arn} + if self.use_invocation_time_range: + source["sessionFilterConfig"] = { + "startTime": start_time, + "endTime": end_time, + } + return {"onlineEvaluationConfigSource": source} + + # --------------------------------------------------------------------------- # Batch eval result models # --------------------------------------------------------------------------- @@ -219,6 +259,8 @@ class BatchEvaluationResult(BaseModel): description: Optional human-readable description of the batch evaluation job. status: Terminal status of the job (e.g. ``"COMPLETED"``). created_at: Timestamp when the batch evaluation job was created. + updated_at: Timestamp when the batch evaluation job was last updated by + the service. ``None`` if the API did not return it. evaluation_results: Aggregated per-evaluator statistics. Present when the job completed successfully; ``None`` otherwise. error_details: Service-reported error messages when the job failed. @@ -230,6 +272,9 @@ class BatchEvaluationResult(BaseModel): per-session evaluation result events. Pass to :py:meth:`BatchEvaluationRunner.fetch_evaluation_events` to read the raw OTel evaluation records. + kms_key_arn: ARN of the KMS key the service used to encrypt this batch + evaluation's data, echoed back by the API. ``None`` when an + AWS-owned key was used. """ batch_evaluation_id: str @@ -238,10 +283,12 @@ class BatchEvaluationResult(BaseModel): description: Optional[str] = None status: str created_at: datetime + updated_at: Optional[datetime] = None evaluation_results: Optional[BatchEvaluationSummary] = None error_details: Optional[List[str]] = None agent_invocation_failures: List[FailedScenario] = Field(default_factory=list) output_data_config: Optional[CloudWatchOutputDataConfig] = None + kms_key_arn: Optional[str] = None # --------------------------------------------------------------------------- @@ -281,6 +328,13 @@ class BatchEvaluationRunConfig(BaseModel): Defaults to 30 seconds. Must be less than ``polling_timeout_seconds``. simulation_config: Actor simulation settings. Required when the dataset contains SimulatedScenario entries. + kms_key_arn: ARN of the customer-managed KMS key used to encrypt the + batch evaluation's data at rest. When omitted, the service uses an + AWS-owned key. The key must be in the same region as the evaluation, + and the calling principal must have ``kms:Encrypt``/``kms:Decrypt`` + permissions on it. + tags: Optional resource tags applied to the batch evaluation job + (key/value pairs), useful for cost allocation and access control. """ model_config = ConfigDict(arbitrary_types_allowed=True) @@ -293,6 +347,8 @@ class BatchEvaluationRunConfig(BaseModel): polling_timeout_seconds: int = 1800 polling_interval_seconds: int = 30 simulation_config: Optional[SimulationConfig] = None + kms_key_arn: Optional[str] = None + tags: Optional[Dict[str, str]] = None @model_validator(mode="after") def validate_polling(self): diff --git a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py index 640676d5..0f5d8e1f 100644 --- a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py +++ b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py @@ -398,6 +398,10 @@ def run_dataset_evaluation( ) if config.description is not None: start_kwargs["description"] = config.description + if config.kms_key_arn is not None: + start_kwargs["kmsKeyArn"] = config.kms_key_arn + if config.tags: + start_kwargs["tags"] = config.tags start_response = self.data_plane_client.start_batch_evaluation(**start_kwargs) except Exception as e: error_code = self._get_boto3_error_code(e) @@ -438,11 +442,13 @@ def run_dataset_evaluation( batch_evaluation_name=response["batchEvaluationName"], status=response["status"], created_at=response["createdAt"], + updated_at=response.get("updatedAt"), description=response.get("description"), agent_invocation_failures=failed_scenarios, evaluation_results=evaluation_results, error_details=response.get("errorDetails"), output_data_config=output_data_config, + kms_key_arn=response.get("kmsKeyArn"), ) logger.info( diff --git a/tests/bedrock_agentcore/evaluation/runner/batch/test_batch_evaluation_models.py b/tests/bedrock_agentcore/evaluation/runner/batch/test_batch_evaluation_models.py index 24fa7860..c855a4fd 100644 --- a/tests/bedrock_agentcore/evaluation/runner/batch/test_batch_evaluation_models.py +++ b/tests/bedrock_agentcore/evaluation/runner/batch/test_batch_evaluation_models.py @@ -2,8 +2,12 @@ from datetime import datetime, timezone +import pytest +from pydantic import ValidationError + from bedrock_agentcore.evaluation.runner.batch.batch_evaluation_models import ( CloudWatchDataSourceConfig, + OnlineEvaluationDataSourceConfig, ) _T0 = datetime(2024, 1, 1, tzinfo=timezone.utc) @@ -15,6 +19,8 @@ ingestion_delay_seconds=0, ) +_ONLINE_CONFIG_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789012:online-evaluation-config/oec-1" + def test_cloudwatch_to_data_source_config_returns_session_ids(): result = _CW_SOURCE.to_data_source_config(["s1", "s2"], _T0, _T1) @@ -24,3 +30,29 @@ def test_cloudwatch_to_data_source_config_returns_session_ids(): assert cw["filterConfig"]["sessionIds"] == ["s1", "s2"] assert cw["filterConfig"]["timeRange"]["startTime"] == _T0 assert cw["filterConfig"]["timeRange"]["endTime"] == _T1 + + +def test_online_to_data_source_config_includes_time_range_by_default(): + source = OnlineEvaluationDataSourceConfig(online_evaluation_config_arn=_ONLINE_CONFIG_ARN) + result = source.to_data_source_config(["s1", "s2"], _T0, _T1) + online = result["onlineEvaluationConfigSource"] + assert online["onlineEvaluationConfigArn"] == _ONLINE_CONFIG_ARN + assert online["sessionFilterConfig"]["startTime"] == _T0 + assert online["sessionFilterConfig"]["endTime"] == _T1 + # Online source does not filter by session IDs. + assert "sessionIds" not in online + + +def test_online_to_data_source_config_omits_time_range_when_disabled(): + source = OnlineEvaluationDataSourceConfig( + online_evaluation_config_arn=_ONLINE_CONFIG_ARN, + use_invocation_time_range=False, + ) + result = source.to_data_source_config(["s1"], _T0, _T1) + online = result["onlineEvaluationConfigSource"] + assert online == {"onlineEvaluationConfigArn": _ONLINE_CONFIG_ARN} + + +def test_online_config_arn_must_be_non_empty(): + with pytest.raises(ValidationError): + OnlineEvaluationDataSourceConfig(online_evaluation_config_arn="") diff --git a/tests/bedrock_agentcore/evaluation/runner/batch/test_batch_evaluation_runner.py b/tests/bedrock_agentcore/evaluation/runner/batch/test_batch_evaluation_runner.py index db2f7138..63f16c49 100644 --- a/tests/bedrock_agentcore/evaluation/runner/batch/test_batch_evaluation_runner.py +++ b/tests/bedrock_agentcore/evaluation/runner/batch/test_batch_evaluation_runner.py @@ -11,6 +11,7 @@ BatchEvaluatorConfig, CloudWatchDataSourceConfig, CloudWatchOutputDataConfig, + OnlineEvaluationDataSourceConfig, ) from bedrock_agentcore.evaluation.runner.batch.batch_evaluation_runner import ( BatchEvaluationRunner, @@ -451,6 +452,180 @@ def test_run_cloudwatch_source_passes_session_ids(): assert filter_config["sessionIds"] == ["s1-session-abc"] +# --------------------------------------------------------------------------- +# kms_key_arn / tags wiring +# --------------------------------------------------------------------------- + +_KMS_KEY_ARN = "arn:aws:kms:us-west-2:123456789012:key/abcd-1234" + + +def _make_config_with_kms(kms_key_arn=_KMS_KEY_ARN, tags=None): + return BatchEvaluationRunConfig( + batch_evaluation_name="test-eval", + evaluator_config=BatchEvaluatorConfig(evaluator_ids=["Builtin.Helpfulness"]), + data_source=_make_cw_source(), + max_concurrent_scenarios=2, + polling_timeout_seconds=60, + polling_interval_seconds=5, + kms_key_arn=kms_key_arn, + tags=tags, + ) + + +def _run_with_single_session(runner, config): + with patch.object( + runner, + "_execute_scenarios_parallel", + return_value=( + [MagicMock(scenario_id="s1", session_id="s1-session-abc", start_time=_T0, end_time=_T1, ground_truth=None)], + [], + ), + ): + return runner.run_dataset_evaluation(config=config, dataset=_DATASET, agent_invoker=_make_invoker()) + + +def test_run_passes_kms_key_arn_to_start_batch_evaluation(): + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = _make_completed_response() + + _run_with_single_session(runner, _make_config_with_kms()) + + call_kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs + assert call_kwargs["kmsKeyArn"] == _KMS_KEY_ARN + + +def test_run_omits_kms_key_arn_when_not_set(): + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = _make_completed_response() + + _run_with_single_session(runner, _make_config()) + + call_kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs + assert "kmsKeyArn" not in call_kwargs + + +def test_run_surfaces_kms_key_arn_on_result(): + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = { + **_make_completed_response(), + "kmsKeyArn": _KMS_KEY_ARN, + } + + result = _run_with_single_session(runner, _make_config_with_kms()) + + assert result.kms_key_arn == _KMS_KEY_ARN + + +def test_run_kms_key_arn_none_on_result_when_absent(): + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = _make_completed_response() + + result = _run_with_single_session(runner, _make_config()) + + assert result.kms_key_arn is None + + +def test_run_passes_tags_to_start_batch_evaluation(): + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = _make_completed_response() + + _run_with_single_session(runner, _make_config_with_kms(tags={"team": "agentcore"})) + + call_kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs + assert call_kwargs["tags"] == {"team": "agentcore"} + + +def test_run_omits_tags_when_not_set(): + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = _make_completed_response() + + _run_with_single_session(runner, _make_config()) + + call_kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs + assert "tags" not in call_kwargs + + +# --------------------------------------------------------------------------- +# OnlineEvaluationDataSourceConfig wiring + updated_at +# --------------------------------------------------------------------------- + +_ONLINE_CONFIG_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789012:online-evaluation-config/oec-1" + + +def test_run_online_source_passes_online_config_arn(): + """Runner emits onlineEvaluationConfigSource (not cloudWatchLogs) for online sources.""" + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = _make_completed_response() + + online_source = OnlineEvaluationDataSourceConfig(online_evaluation_config_arn=_ONLINE_CONFIG_ARN) + + with patch.object( + runner, + "_execute_scenarios_parallel", + return_value=( + [MagicMock(scenario_id="s1", session_id="s1-session-abc", start_time=_T0, end_time=_T1, ground_truth=None)], + [], + ), + ): + runner.run_dataset_evaluation( + config=_make_config(source=online_source), dataset=_DATASET, agent_invoker=_make_invoker() + ) + + call_kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs + ds = call_kwargs["dataSourceConfig"] + assert "cloudWatchLogs" not in ds + online = ds["onlineEvaluationConfigSource"] + assert online["onlineEvaluationConfigArn"] == _ONLINE_CONFIG_ARN + assert online["sessionFilterConfig"]["startTime"] == _T0 + assert online["sessionFilterConfig"]["endTime"] == _T1 + + +def test_run_surfaces_updated_at_on_result(): + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = { + **_make_completed_response(), + "updatedAt": _T1, + } + + with patch.object( + runner, + "_execute_scenarios_parallel", + return_value=( + [MagicMock(scenario_id="s1", session_id="s1-session-abc", start_time=_T0, end_time=_T1, ground_truth=None)], + [], + ), + ): + result = runner.run_dataset_evaluation(config=_make_config(), dataset=_DATASET, agent_invoker=_make_invoker()) + + assert result.updated_at == _T1 + + +def test_run_updated_at_none_when_absent(): + runner = _make_runner() + runner.data_plane_client.start_batch_evaluation.return_value = _make_start_response() + runner.data_plane_client.get_batch_evaluation.return_value = _make_completed_response() + + with patch.object( + runner, + "_execute_scenarios_parallel", + return_value=( + [MagicMock(scenario_id="s1", session_id="s1-session-abc", start_time=_T0, end_time=_T1, ground_truth=None)], + [], + ), + ): + result = runner.run_dataset_evaluation(config=_make_config(), dataset=_DATASET, agent_invoker=_make_invoker()) + + assert result.updated_at is None + + # --------------------------------------------------------------------------- # fetch_evaluation_events (#11) # ---------------------------------------------------------------------------