From 2aa484b303fe210c0a5c62c815176db52c75a1cf Mon Sep 17 00:00:00 2001 From: Sundar Raghavan Date: Thu, 22 Oct 2026 11:45:31 -0700 Subject: [PATCH 1/2] feat(evaluation): evaluate existing sessions and add RecommendationClient Add BatchEvaluationRunner.run_session_evaluation to score sessions already in CloudWatch without invoking the agent, and RecommendationClient with start_recommendation_and_wait for system prompt and tool description recommendations. --- src/bedrock_agentcore/evaluation/__init__.py | 2 + .../evaluation/recommendation_client.py | 134 ++++++++++++++++++ .../runner/batch/batch_evaluation_models.py | 17 +-- .../runner/batch/batch_evaluation_runner.py | 92 +++++++++++- .../batch/test_batch_evaluation_runner.py | 93 ++++++++++++ .../evaluation/test_recommendation_client.py | 122 ++++++++++++++++ 6 files changed, 445 insertions(+), 15 deletions(-) create mode 100644 src/bedrock_agentcore/evaluation/recommendation_client.py create mode 100644 tests/bedrock_agentcore/evaluation/test_recommendation_client.py diff --git a/src/bedrock_agentcore/evaluation/__init__.py b/src/bedrock_agentcore/evaluation/__init__.py index 40bf85a5..d79a8a4b 100644 --- a/src/bedrock_agentcore/evaluation/__init__.py +++ b/src/bedrock_agentcore/evaluation/__init__.py @@ -9,6 +9,7 @@ custom_code_based_evaluator, ) from bedrock_agentcore.evaluation.dataset_client import DatasetClient +from bedrock_agentcore.evaluation.recommendation_client import RecommendationClient from bedrock_agentcore.evaluation.runner.batch.batch_evaluation_models import ( BatchEvaluationResult, BatchEvaluationRunConfig, @@ -94,6 +95,7 @@ "CloudWatchAgentSpanCollector", "Dataset", "DatasetClient", + "RecommendationClient", "DatasetProvider", "EvaluationClient", "EvaluationResult", diff --git a/src/bedrock_agentcore/evaluation/recommendation_client.py b/src/bedrock_agentcore/evaluation/recommendation_client.py new file mode 100644 index 00000000..c0c30ec2 --- /dev/null +++ b/src/bedrock_agentcore/evaluation/recommendation_client.py @@ -0,0 +1,134 @@ +"""RecommendationClient for AgentCore system prompt and tool description recommendations.""" + +import logging +from typing import Any, Dict, Optional + +import boto3 +from botocore.config import Config + +from bedrock_agentcore._utils.config import WaitConfig +from bedrock_agentcore._utils.endpoints import DP_ENDPOINT_OVERRIDE +from bedrock_agentcore._utils.polling import wait_until +from bedrock_agentcore._utils.snake_case import accept_snake_case_kwargs, convert_kwargs +from bedrock_agentcore._utils.user_agent import build_user_agent_suffix + +logger = logging.getLogger(__name__) + + +class RecommendationClient: + """Client for AgentCore recommendations. + + Provides pass-through access to the recommendation APIs on the + bedrock-agentcore service, plus ``start_recommendation_and_wait`` for the + asynchronous start-then-poll flow every recommendation needs. + + .. warning:: + This feature is in preview and may change in future releases. + + Example:: + + client = RecommendationClient(region_name="us-west-2") + + recommendation = client.start_recommendation_and_wait( + name="my-prompt-rec", + type="SYSTEM_PROMPT_RECOMMENDATION", + recommendationConfig={"systemPromptRecommendationConfig": { + "systemPrompt": {"text": "You are a helpful assistant."}, + "agentTraces": {"batchEvaluation": {"batchEvaluationArn": arn}}, + "evaluationConfig": {"evaluators": [ + {"evaluatorArn": "arn:aws:bedrock-agentcore:::evaluator/Builtin.GoalSuccessRate"} + ]}, + }}, + ) + print(recommendation["recommendationResult"]) + + # Pass-through to any recommendation API + client.list_recommendations(maxResults=10) + """ + + _ALLOWED_DP_METHODS = { + "start_recommendation", + "get_recommendation", + "list_recommendations", + "delete_recommendation", + } + + def __init__( + self, + region_name: Optional[str] = None, + integration_source: Optional[str] = None, + boto3_session: Optional[boto3.Session] = None, + ): + """Initialize the RecommendationClient. + + Args: + region_name: AWS region. Falls back to boto3 session region or us-west-2. + integration_source: Optional integration framework identifier for telemetry. + boto3_session: Optional boto3 Session. If not provided, a default is created. + """ + session = boto3_session if boto3_session else boto3.Session() + self.region_name = region_name or session.region_name or "us-west-2" + self.integration_source = integration_source + + client_kwargs: Dict[str, Any] = { + "region_name": self.region_name, + "config": Config(user_agent_extra=build_user_agent_suffix(integration_source)), + } + if DP_ENDPOINT_OVERRIDE: + client_kwargs["endpoint_url"] = DP_ENDPOINT_OVERRIDE + self._dp_client = session.client("bedrock-agentcore", **client_kwargs) + + logger.info("Initialized RecommendationClient in region %s", self.region_name) + + def __getattr__(self, name: str): + """Dynamically forward allowlisted method calls to the boto3 client.""" + if "_dp_client" not in self.__dict__: + raise AttributeError(name) + + if name in self._ALLOWED_DP_METHODS and hasattr(self._dp_client, name): + method = getattr(self._dp_client, name) + logger.debug("Forwarding method '%s' to _dp_client", name) + return accept_snake_case_kwargs(method) + + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{name}'. " + f"Method not found on data plane client. " + f"Available methods can be found in the boto3 documentation for " + f"'bedrock-agentcore' service." + ) + + def start_recommendation_and_wait( + self, + wait_config: Optional[WaitConfig] = None, + **kwargs, + ) -> Dict[str, Any]: + """Start a recommendation and wait for it to reach COMPLETED status. + + Recommendations typically take several minutes, so the default + ``WaitConfig`` (5 minutes) may be too short. Pass a longer ``max_wait`` + for large trace sets. + + Args: + wait_config: Optional WaitConfig for polling behavior. + **kwargs: Arguments forwarded to the start_recommendation API. + + Returns: + The GetRecommendation response when COMPLETED. The generated result is + under ``recommendationResult``. Note that ``recommendationConfig`` + echoes the input you submitted, so read the result from + ``recommendationResult`` rather than searching the response. + + Raises: + RuntimeError: If the recommendation reaches FAILED status. + TimeoutError: If the recommendation doesn't complete within max_wait. + """ + response = self._dp_client.start_recommendation(**convert_kwargs(kwargs)) + recommendation_id = response["recommendationId"] + logger.info("Started recommendation %s", recommendation_id) + return wait_until( + lambda: self._dp_client.get_recommendation(recommendationId=recommendation_id), + "COMPLETED", + {"FAILED"}, + wait_config, + error_field="recommendationResult", + ) 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 18f4452a..958c8fb4 100644 --- a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_models.py +++ b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_models.py @@ -95,18 +95,19 @@ def to_data_source_config( start_time: datetime, end_time: datetime, ) -> Dict[str, Any]: - """Return a cloudWatchLogs dataSourceConfig dict for the evaluation API.""" + """Return a cloudWatchLogs dataSourceConfig dict for the evaluation API. + + ``sessionIds`` is omitted when empty, so that every session in the time + range is evaluated. + """ + filter_config: Dict[str, Any] = {"timeRange": {"startTime": start_time, "endTime": end_time}} + if session_ids: + filter_config["sessionIds"] = session_ids return { "cloudWatchLogs": { "serviceNames": self.service_names, "logGroupNames": self.log_group_names, - "filterConfig": { - "sessionIds": session_ids, - "timeRange": { - "startTime": start_time, - "endTime": end_time, - }, - }, + "filterConfig": filter_config, } } 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 7f83a407..3bf8e34f 100644 --- a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py +++ b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py @@ -387,18 +387,39 @@ def run_dataset_evaluation( for session in successful_sessions ] + return self._start_and_collect( + config, + data_source_config=config.data_source.to_data_source_config( + [s.session_id for s in successful_sessions], + min(s.start_time for s in successful_sessions), + max(s.end_time for s in successful_sessions), + ), + evaluation_metadata={"sessionMetadata": session_metadata_list}, + agent_invocation_failures=failed_scenarios, + ) + + def _start_and_collect( + self, + config: BatchEvaluationRunConfig, + data_source_config: Dict[str, Any], + evaluation_metadata: Optional[Dict[str, Any]] = None, + agent_invocation_failures: Optional[List[FailedScenario]] = None, + ) -> BatchEvaluationResult: + """Start a batch evaluation, poll it to a terminal state, and build the result. + + Shared by :py:meth:`run_dataset_evaluation` and + :py:meth:`run_session_evaluation`, which differ only in how they choose the + sessions to evaluate. + """ logger.info("Calling StartBatchEvaluation (name=%s)", config.batch_evaluation_name) try: start_kwargs: Dict[str, Any] = dict( batchEvaluationName=config.batch_evaluation_name, evaluators=[{"evaluatorId": eid} for eid in config.evaluator_config.evaluator_ids], - dataSourceConfig=config.data_source.to_data_source_config( - [s.session_id for s in successful_sessions], - min(s.start_time for s in successful_sessions), - max(s.end_time for s in successful_sessions), - ), - evaluationMetadata={"sessionMetadata": session_metadata_list}, + dataSourceConfig=data_source_config, ) + if evaluation_metadata is not None: + start_kwargs["evaluationMetadata"] = evaluation_metadata if config.description is not None: start_kwargs["description"] = config.description if config.kms_key_arn is not None: @@ -454,7 +475,7 @@ def run_dataset_evaluation( created_at=response["createdAt"], updated_at=response.get("updatedAt"), description=response.get("description"), - agent_invocation_failures=failed_scenarios, + agent_invocation_failures=agent_invocation_failures or [], evaluation_results=evaluation_results, insight_results=insight_results, error_details=response.get("errorDetails"), @@ -471,6 +492,63 @@ def run_dataset_evaluation( ) return result + def run_session_evaluation( + self, + config: BatchEvaluationRunConfig, + start_time: datetime, + end_time: datetime, + session_ids: Optional[List[str]] = None, + ) -> BatchEvaluationResult: + """Run a batch evaluation on sessions that already exist. + + Use this to evaluate production traffic or any previously recorded + sessions, for example to take a baseline before a change or to compare + before and after it. Unlike :py:meth:`run_dataset_evaluation`, nothing is + invoked: the service reads the sessions from ``config.data_source``. + + Analysis evaluators such as ``Builtin.FailureAnalysis`` can run alongside + scoring evaluators in ``config.evaluator_config``; their patterns are + returned in ``result.insight_results``. + + The data source's pre-run hook (for example the CloudWatch ingestion + delay) is not called, because there is no agent invocation to wait for. + Sessions must already be ingested. + + Args: + config: Evaluation name, evaluator IDs, session source, and polling + timeouts. ``max_concurrent_scenarios`` is not used. + start_time: Start of the time window to read sessions from. + end_time: End of the time window. + session_ids: Optional session IDs to narrow the window to. When + omitted, every session in the window is evaluated. + + Returns: + :class:`BatchEvaluationResult` with job IDs, status, + ``evaluation_results``, ``insight_results``, and + ``output_data_config``. + + Raises: + ValueError: If ``start_time`` is not earlier than ``end_time``. + RuntimeError: If API calls fail. + TimeoutError: If the job exceeds ``config.polling_timeout_seconds``. + """ + if start_time >= end_time: + raise ValueError("start_time must be earlier than end_time") + logger.info( + "Starting session batch evaluation: %s to %s, %s session IDs", + start_time, + end_time, + len(session_ids) if session_ids else "all", + ) + return self._start_and_collect( + config, + data_source_config=config.data_source.to_data_source_config( + list(session_ids or []), + start_time, + end_time, + ), + ) + def fetch_evaluation_events(self, result: BatchEvaluationResult) -> List[Dict[str, Any]]: """Fetch per-turn evaluation events from CloudWatch. 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 a13158a9..10c632f8 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 @@ -1068,3 +1068,96 @@ class ResourceNotFoundException(Exception): with pytest.raises(LookupError, match="CloudWatch log stream not found"): runner.fetch_evaluation_events(_make_fetch_result(_CW_OUTPUT_CONFIG)) + + +# --------------------------------------------------------------------------- +# run_session_evaluation: sessions that already exist +# --------------------------------------------------------------------------- + + +def _run_sessions(session_ids=None, insight_results=None, start=_T0, end=_T1): + 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_eval_results=True, insight_results=insight_results + ) + result = runner.run_session_evaluation( + config=_make_config(), start_time=start, end_time=end, session_ids=session_ids + ) + return runner, result + + +def test_run_session_evaluation_whole_window_omits_session_ids(): + """With no session IDs, every session in the window is evaluated.""" + runner, result = _run_sessions() + kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs + filter_config = kwargs["dataSourceConfig"]["cloudWatchLogs"]["filterConfig"] + assert "sessionIds" not in filter_config + assert filter_config["timeRange"] == {"startTime": _T0, "endTime": _T1} + assert result.status == "COMPLETED" + assert result.evaluation_results.total_number_of_sessions == 1 + + +def test_run_session_evaluation_passes_session_ids(): + runner, _ = _run_sessions(session_ids=["s-1", "s-2"]) + filter_config = runner.data_plane_client.start_batch_evaluation.call_args.kwargs["dataSourceConfig"][ + "cloudWatchLogs" + ]["filterConfig"] + assert filter_config["sessionIds"] == ["s-1", "s-2"] + + +def test_run_session_evaluation_does_not_invoke_or_wait(): + """Nothing is invoked, and the ingestion hook is not run for existing sessions.""" + 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() + config = _make_config(source=_make_cw_source()) + with ( + patch.object(runner, "_execute_scenarios_parallel") as executed, + patch.object(type(config.data_source), "pre_evaluation_run_hook") as hook, + ): + runner.run_session_evaluation(config=config, start_time=_T0, end_time=_T1) + executed.assert_not_called() + hook.assert_not_called() + + +def test_run_session_evaluation_sends_no_session_metadata(): + runner, _ = _run_sessions() + assert "evaluationMetadata" not in runner.data_plane_client.start_batch_evaluation.call_args.kwargs + + +def test_run_session_evaluation_returns_insight_results(): + _, result = _run_sessions(insight_results=_INSIGHT_RESULTS) + assert [i.name for i in result.insight_results] == [ + "IAM Permission Denial Errors", + "Query Timeout Errors", + ] + assert result.agent_invocation_failures == [] + + +def test_run_session_evaluation_rejects_inverted_window(): + runner = _make_runner() + with pytest.raises(ValueError, match="start_time must be earlier"): + runner.run_session_evaluation(config=_make_config(), start_time=_T1, end_time=_T0) + runner.data_plane_client.start_batch_evaluation.assert_not_called() + + +def test_dataset_evaluation_still_sends_session_ids_and_metadata(): + """The shared start-and-collect path keeps the dataset run's request unchanged.""" + 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)], + [], + ), + ): + runner.run_dataset_evaluation(config=_make_config(), dataset=_DATASET, agent_invoker=_make_invoker()) + kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs + assert kwargs["dataSourceConfig"]["cloudWatchLogs"]["filterConfig"]["sessionIds"] == ["s1-session-abc"] + assert kwargs["evaluationMetadata"] == { + "sessionMetadata": [{"sessionId": "s1-session-abc", "testScenarioId": "s1"}] + } diff --git a/tests/bedrock_agentcore/evaluation/test_recommendation_client.py b/tests/bedrock_agentcore/evaluation/test_recommendation_client.py new file mode 100644 index 00000000..86b60ba7 --- /dev/null +++ b/tests/bedrock_agentcore/evaluation/test_recommendation_client.py @@ -0,0 +1,122 @@ +"""Tests for RecommendationClient.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from bedrock_agentcore._utils.config import WaitConfig +from bedrock_agentcore.evaluation.recommendation_client import RecommendationClient + + +@pytest.fixture +def mock_boto3(): + with patch("bedrock_agentcore.evaluation.recommendation_client.boto3") as m: + yield m + + +@pytest.fixture +def client_and_dp(mock_boto3): + """Return a (RecommendationClient, mock_dp_client) tuple.""" + mock_session = MagicMock() + mock_boto3.Session.return_value = mock_session + mock_dp = mock_session.client.return_value + client = RecommendationClient() + return client, mock_dp + + +_FAST = WaitConfig(max_wait=5, poll_interval=1) + + +class TestRecommendationClientInit: + def test_uses_data_plane_service(self, mock_boto3): + mock_session = MagicMock() + mock_boto3.Session.return_value = mock_session + RecommendationClient(region_name="eu-west-1") + args, kwargs = mock_session.client.call_args + assert args == ("bedrock-agentcore",) + assert kwargs["region_name"] == "eu-west-1" + + def test_default_region_from_session(self, mock_boto3): + mock_session = MagicMock() + mock_session.region_name = "us-east-1" + mock_boto3.Session.return_value = mock_session + assert RecommendationClient().region_name == "us-east-1" + + def test_custom_session(self, mock_boto3): + custom_session = MagicMock() + custom_session.region_name = "ap-southeast-1" + client = RecommendationClient(boto3_session=custom_session) + assert client.region_name == "ap-southeast-1" + custom_session.client.assert_called_once() + + def test_endpoint_override(self, mock_boto3): + mock_session = MagicMock() + mock_boto3.Session.return_value = mock_session + with patch( + "bedrock_agentcore.evaluation.recommendation_client.DP_ENDPOINT_OVERRIDE", + "https://example.test", + ): + RecommendationClient(region_name="us-west-2") + assert mock_session.client.call_args.kwargs["endpoint_url"] == "https://example.test" + + +class TestRecommendationClientPassthrough: + def test_allowed_method_forwarded(self, client_and_dp): + client, mock_dp = client_and_dp + mock_dp.list_recommendations.return_value = {"recommendations": []} + client.list_recommendations(maxResults=10) + mock_dp.list_recommendations.assert_called_once_with(maxResults=10) + + def test_snake_case_kwargs_converted(self, client_and_dp): + client, mock_dp = client_and_dp + client.get_recommendation(recommendation_id="rec-1") + mock_dp.get_recommendation.assert_called_once_with(recommendationId="rec-1") + + def test_disallowed_method_raises(self, client_and_dp): + client, _ = client_and_dp + with pytest.raises(AttributeError, match="no attribute 'start_batch_evaluation'"): + client.start_batch_evaluation() + + +class TestStartRecommendationAndWait: + def test_returns_completed_response(self, client_and_dp): + client, mock_dp = client_and_dp + mock_dp.start_recommendation.return_value = {"recommendationId": "rec-1", "status": "PENDING"} + mock_dp.get_recommendation.side_effect = [ + {"recommendationId": "rec-1", "status": "IN_PROGRESS"}, + { + "recommendationId": "rec-1", + "status": "COMPLETED", + "recommendationResult": {"systemPromptRecommendationResult": {"recommendedSystemPrompt": "new"}}, + }, + ] + with patch("bedrock_agentcore._utils.polling.time.sleep"): + result = client.start_recommendation_and_wait( + wait_config=_FAST, name="r", type="SYSTEM_PROMPT_RECOMMENDATION", recommendation_config={} + ) + assert result["recommendationResult"]["systemPromptRecommendationResult"]["recommendedSystemPrompt"] == "new" + mock_dp.start_recommendation.assert_called_once_with( + name="r", type="SYSTEM_PROMPT_RECOMMENDATION", recommendationConfig={} + ) + mock_dp.get_recommendation.assert_called_with(recommendationId="rec-1") + + def test_failed_raises_with_result(self, client_and_dp): + client, mock_dp = client_and_dp + mock_dp.start_recommendation.return_value = {"recommendationId": "rec-1"} + mock_dp.get_recommendation.return_value = { + "status": "FAILED", + "recommendationResult": {"systemPromptRecommendationResult": {"errorMessage": "no traces found"}}, + } + with pytest.raises(RuntimeError, match="no traces found"): + client.start_recommendation_and_wait(wait_config=_FAST, name="r") + + def test_timeout(self, client_and_dp): + client, mock_dp = client_and_dp + mock_dp.start_recommendation.return_value = {"recommendationId": "rec-1"} + mock_dp.get_recommendation.return_value = {"status": "IN_PROGRESS"} + with ( + patch("bedrock_agentcore._utils.polling.time.sleep"), + patch("bedrock_agentcore._utils.polling.time.time", side_effect=[0, 0, 10]), + pytest.raises(TimeoutError), + ): + client.start_recommendation_and_wait(wait_config=_FAST, name="r") From 72d2bb78d0eac20359ff1cbfb6495b9bf1bd860f Mon Sep 17 00:00:00 2001 From: Sundar Raghavan Date: Thu, 22 Oct 2026 15:56:08 -0700 Subject: [PATCH 2/2] refactor(evaluation): move existing-session runs to EvaluationClient Allowlist the batch evaluation APIs on EvaluationClient and add start_batch_evaluation_and_wait. Revert run_session_evaluation from BatchEvaluationRunner, which is for dataset runs. wait_until accepts a set of target statuses. --- src/bedrock_agentcore/_utils/polling.py | 9 +- src/bedrock_agentcore/evaluation/client.py | 43 +++++++++ .../runner/batch/batch_evaluation_runner.py | 92 ++---------------- .../batch/test_batch_evaluation_runner.py | 93 ------------------- .../evaluation/test_client_passthrough.py | 63 ++++++++++++- tests/unit/test_utils.py | 7 ++ 6 files changed, 124 insertions(+), 183 deletions(-) diff --git a/src/bedrock_agentcore/_utils/polling.py b/src/bedrock_agentcore/_utils/polling.py index 7ff10435..e1d65392 100644 --- a/src/bedrock_agentcore/_utils/polling.py +++ b/src/bedrock_agentcore/_utils/polling.py @@ -2,7 +2,7 @@ import logging import time -from typing import Any, Callable, Dict, Optional, Set +from typing import Any, Callable, Dict, Optional, Set, Union from .config import WaitConfig @@ -11,7 +11,7 @@ def wait_until( poll_fn: Callable[[], Dict[str, Any]], - target: str, + target: Union[str, Set[str]], failed: Set[str], wait_config: Optional[WaitConfig] = None, error_field: str = "statusReasons", @@ -20,7 +20,7 @@ def wait_until( Args: poll_fn: Zero-arg callable that returns the resource's current state. - target: The status to wait for (e.g. "ACTIVE", "READY"). + target: The status to wait for (e.g. "ACTIVE", "READY"), or a set of them. failed: Statuses that indicate terminal failure. wait_config: Optional WaitConfig for polling behavior. error_field: Response field containing error details. @@ -33,13 +33,14 @@ def wait_until( TimeoutError: If target status is not reached within max_wait. """ wait = wait_config or WaitConfig() + targets = {target} if isinstance(target, str) else target start_time = time.time() while True: resp = poll_fn() status = resp.get("status") if status is None: logger.warning("Response missing 'status' field: %s", resp) - if status == target: + if status in targets: return resp if status in failed: reason = resp.get(error_field, "Unknown") diff --git a/src/bedrock_agentcore/evaluation/client.py b/src/bedrock_agentcore/evaluation/client.py index 0b9c3c10..16b7d2ad 100644 --- a/src/bedrock_agentcore/evaluation/client.py +++ b/src/bedrock_agentcore/evaluation/client.py @@ -22,6 +22,8 @@ POLL_INTERVAL_SECONDS = 2 _EVALUATOR_FAILED_STATUSES = {"CREATE_FAILED", "UPDATE_FAILED"} +_BATCH_EVALUATION_DONE_STATUSES = {"COMPLETED", "COMPLETED_WITH_ERRORS"} +_BATCH_EVALUATION_FAILED_STATUSES = {"FAILED", "STOPPED"} class ReferenceInputs(BaseModel): @@ -102,6 +104,12 @@ def __init__( # ------------------------------------------------------------------------- _ALLOWED_DP_METHODS = { "evaluate", + # Batch evaluation + "start_batch_evaluation", + "get_batch_evaluation", + "list_batch_evaluations", + "stop_batch_evaluation", + "delete_batch_evaluation", } _ALLOWED_CP_METHODS = { @@ -372,6 +380,41 @@ def _build_reference_inputs( # *_and_wait methods # ------------------------------------------------------------------------- + def start_batch_evaluation_and_wait( + self, + wait_config: Optional[WaitConfig] = None, + **kwargs, + ) -> Dict[str, Any]: + """Start a batch evaluation and wait for it to finish. + + Scores sessions that already exist in CloudWatch, without invoking the + agent. To run a dataset against an agent, use ``BatchEvaluationRunner``. + Batch evaluations typically take several minutes, so pass a ``max_wait`` + longer than the default 5 minutes. + + Args: + wait_config: Optional WaitConfig for polling behavior. + **kwargs: Arguments forwarded to the start_batch_evaluation API. + + Returns: + The GetBatchEvaluation response when COMPLETED or COMPLETED_WITH_ERRORS. + Scores are under ``evaluationResults``, and insights from analysis + evaluators such as ``Builtin.FailureAnalysis`` under ``insightResults``. + + Raises: + RuntimeError: If the batch evaluation reaches FAILED or STOPPED status. + TimeoutError: If the batch evaluation doesn't finish within max_wait. + """ + response = self._dp_client.start_batch_evaluation(**convert_kwargs(kwargs)) + batch_id = response["batchEvaluationId"] + return wait_until( + lambda: self._dp_client.get_batch_evaluation(batchEvaluationId=batch_id), + _BATCH_EVALUATION_DONE_STATUSES, + _BATCH_EVALUATION_FAILED_STATUSES, + wait_config, + error_field="errorDetails", + ) + def create_evaluator_and_wait( self, wait_config: Optional[WaitConfig] = None, 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 3bf8e34f..7f83a407 100644 --- a/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py +++ b/src/bedrock_agentcore/evaluation/runner/batch/batch_evaluation_runner.py @@ -387,39 +387,18 @@ def run_dataset_evaluation( for session in successful_sessions ] - return self._start_and_collect( - config, - data_source_config=config.data_source.to_data_source_config( - [s.session_id for s in successful_sessions], - min(s.start_time for s in successful_sessions), - max(s.end_time for s in successful_sessions), - ), - evaluation_metadata={"sessionMetadata": session_metadata_list}, - agent_invocation_failures=failed_scenarios, - ) - - def _start_and_collect( - self, - config: BatchEvaluationRunConfig, - data_source_config: Dict[str, Any], - evaluation_metadata: Optional[Dict[str, Any]] = None, - agent_invocation_failures: Optional[List[FailedScenario]] = None, - ) -> BatchEvaluationResult: - """Start a batch evaluation, poll it to a terminal state, and build the result. - - Shared by :py:meth:`run_dataset_evaluation` and - :py:meth:`run_session_evaluation`, which differ only in how they choose the - sessions to evaluate. - """ logger.info("Calling StartBatchEvaluation (name=%s)", config.batch_evaluation_name) try: start_kwargs: Dict[str, Any] = dict( batchEvaluationName=config.batch_evaluation_name, evaluators=[{"evaluatorId": eid} for eid in config.evaluator_config.evaluator_ids], - dataSourceConfig=data_source_config, + dataSourceConfig=config.data_source.to_data_source_config( + [s.session_id for s in successful_sessions], + min(s.start_time for s in successful_sessions), + max(s.end_time for s in successful_sessions), + ), + evaluationMetadata={"sessionMetadata": session_metadata_list}, ) - if evaluation_metadata is not None: - start_kwargs["evaluationMetadata"] = evaluation_metadata if config.description is not None: start_kwargs["description"] = config.description if config.kms_key_arn is not None: @@ -475,7 +454,7 @@ def _start_and_collect( created_at=response["createdAt"], updated_at=response.get("updatedAt"), description=response.get("description"), - agent_invocation_failures=agent_invocation_failures or [], + agent_invocation_failures=failed_scenarios, evaluation_results=evaluation_results, insight_results=insight_results, error_details=response.get("errorDetails"), @@ -492,63 +471,6 @@ def _start_and_collect( ) return result - def run_session_evaluation( - self, - config: BatchEvaluationRunConfig, - start_time: datetime, - end_time: datetime, - session_ids: Optional[List[str]] = None, - ) -> BatchEvaluationResult: - """Run a batch evaluation on sessions that already exist. - - Use this to evaluate production traffic or any previously recorded - sessions, for example to take a baseline before a change or to compare - before and after it. Unlike :py:meth:`run_dataset_evaluation`, nothing is - invoked: the service reads the sessions from ``config.data_source``. - - Analysis evaluators such as ``Builtin.FailureAnalysis`` can run alongside - scoring evaluators in ``config.evaluator_config``; their patterns are - returned in ``result.insight_results``. - - The data source's pre-run hook (for example the CloudWatch ingestion - delay) is not called, because there is no agent invocation to wait for. - Sessions must already be ingested. - - Args: - config: Evaluation name, evaluator IDs, session source, and polling - timeouts. ``max_concurrent_scenarios`` is not used. - start_time: Start of the time window to read sessions from. - end_time: End of the time window. - session_ids: Optional session IDs to narrow the window to. When - omitted, every session in the window is evaluated. - - Returns: - :class:`BatchEvaluationResult` with job IDs, status, - ``evaluation_results``, ``insight_results``, and - ``output_data_config``. - - Raises: - ValueError: If ``start_time`` is not earlier than ``end_time``. - RuntimeError: If API calls fail. - TimeoutError: If the job exceeds ``config.polling_timeout_seconds``. - """ - if start_time >= end_time: - raise ValueError("start_time must be earlier than end_time") - logger.info( - "Starting session batch evaluation: %s to %s, %s session IDs", - start_time, - end_time, - len(session_ids) if session_ids else "all", - ) - return self._start_and_collect( - config, - data_source_config=config.data_source.to_data_source_config( - list(session_ids or []), - start_time, - end_time, - ), - ) - def fetch_evaluation_events(self, result: BatchEvaluationResult) -> List[Dict[str, Any]]: """Fetch per-turn evaluation events from CloudWatch. 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 10c632f8..a13158a9 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 @@ -1068,96 +1068,3 @@ class ResourceNotFoundException(Exception): with pytest.raises(LookupError, match="CloudWatch log stream not found"): runner.fetch_evaluation_events(_make_fetch_result(_CW_OUTPUT_CONFIG)) - - -# --------------------------------------------------------------------------- -# run_session_evaluation: sessions that already exist -# --------------------------------------------------------------------------- - - -def _run_sessions(session_ids=None, insight_results=None, start=_T0, end=_T1): - 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_eval_results=True, insight_results=insight_results - ) - result = runner.run_session_evaluation( - config=_make_config(), start_time=start, end_time=end, session_ids=session_ids - ) - return runner, result - - -def test_run_session_evaluation_whole_window_omits_session_ids(): - """With no session IDs, every session in the window is evaluated.""" - runner, result = _run_sessions() - kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs - filter_config = kwargs["dataSourceConfig"]["cloudWatchLogs"]["filterConfig"] - assert "sessionIds" not in filter_config - assert filter_config["timeRange"] == {"startTime": _T0, "endTime": _T1} - assert result.status == "COMPLETED" - assert result.evaluation_results.total_number_of_sessions == 1 - - -def test_run_session_evaluation_passes_session_ids(): - runner, _ = _run_sessions(session_ids=["s-1", "s-2"]) - filter_config = runner.data_plane_client.start_batch_evaluation.call_args.kwargs["dataSourceConfig"][ - "cloudWatchLogs" - ]["filterConfig"] - assert filter_config["sessionIds"] == ["s-1", "s-2"] - - -def test_run_session_evaluation_does_not_invoke_or_wait(): - """Nothing is invoked, and the ingestion hook is not run for existing sessions.""" - 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() - config = _make_config(source=_make_cw_source()) - with ( - patch.object(runner, "_execute_scenarios_parallel") as executed, - patch.object(type(config.data_source), "pre_evaluation_run_hook") as hook, - ): - runner.run_session_evaluation(config=config, start_time=_T0, end_time=_T1) - executed.assert_not_called() - hook.assert_not_called() - - -def test_run_session_evaluation_sends_no_session_metadata(): - runner, _ = _run_sessions() - assert "evaluationMetadata" not in runner.data_plane_client.start_batch_evaluation.call_args.kwargs - - -def test_run_session_evaluation_returns_insight_results(): - _, result = _run_sessions(insight_results=_INSIGHT_RESULTS) - assert [i.name for i in result.insight_results] == [ - "IAM Permission Denial Errors", - "Query Timeout Errors", - ] - assert result.agent_invocation_failures == [] - - -def test_run_session_evaluation_rejects_inverted_window(): - runner = _make_runner() - with pytest.raises(ValueError, match="start_time must be earlier"): - runner.run_session_evaluation(config=_make_config(), start_time=_T1, end_time=_T0) - runner.data_plane_client.start_batch_evaluation.assert_not_called() - - -def test_dataset_evaluation_still_sends_session_ids_and_metadata(): - """The shared start-and-collect path keeps the dataset run's request unchanged.""" - 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)], - [], - ), - ): - runner.run_dataset_evaluation(config=_make_config(), dataset=_DATASET, agent_invoker=_make_invoker()) - kwargs = runner.data_plane_client.start_batch_evaluation.call_args.kwargs - assert kwargs["dataSourceConfig"]["cloudWatchLogs"]["filterConfig"]["sessionIds"] == ["s1-session-abc"] - assert kwargs["evaluationMetadata"] == { - "sessionMetadata": [{"sessionId": "s1-session-abc", "testScenarioId": "s1"}] - } diff --git a/tests/bedrock_agentcore/evaluation/test_client_passthrough.py b/tests/bedrock_agentcore/evaluation/test_client_passthrough.py index c752c125..d09240e5 100644 --- a/tests/bedrock_agentcore/evaluation/test_client_passthrough.py +++ b/tests/bedrock_agentcore/evaluation/test_client_passthrough.py @@ -59,7 +59,14 @@ def test_all_cp_methods_in_allowlist(self): assert expected == EvaluationClient._ALLOWED_CP_METHODS def test_all_dp_methods_in_allowlist(self): - expected = {"evaluate"} + expected = { + "evaluate", + "start_batch_evaluation", + "get_batch_evaluation", + "list_batch_evaluations", + "stop_batch_evaluation", + "delete_batch_evaluation", + } assert expected == EvaluationClient._ALLOWED_DP_METHODS @@ -153,3 +160,57 @@ def test_delete_online_eval_config_and_wait(self, _mock_sleep): onlineEvaluationConfigId="c-123", ) client._cp_client.delete_online_evaluation_config.assert_called_once() + + +class TestBatchEvaluationAndWait: + """Tests for batch evaluation passthrough and start_batch_evaluation_and_wait.""" + + def _make_client(self): + with patch("boto3.client"): + client = EvaluationClient(region_name="us-west-2") + client._cp_client = Mock() + client._dp_client = Mock() + return client + + @pytest.mark.parametrize( + "method", + [ + "start_batch_evaluation", + "get_batch_evaluation", + "list_batch_evaluations", + "stop_batch_evaluation", + "delete_batch_evaluation", + ], + ) + def test_batch_methods_forwarded(self, method): + client = self._make_client() + getattr(client, method)(batch_evaluation_id="b-1") + getattr(client._dp_client, method).assert_called_once_with(batchEvaluationId="b-1") + + @pytest.mark.parametrize("done", ["COMPLETED", "COMPLETED_WITH_ERRORS"]) + @patch("bedrock_agentcore._utils.polling.time.sleep") + def test_returns_when_done(self, _mock_sleep, done): + client = self._make_client() + client._dp_client.start_batch_evaluation.return_value = {"batchEvaluationId": "b-1"} + client._dp_client.get_batch_evaluation.side_effect = [ + {"status": "IN_PROGRESS"}, + {"status": done, "insightResults": [{"name": "Timeouts"}]}, + ] + result = client.start_batch_evaluation_and_wait( + batch_evaluation_name="run", + evaluators=[{"evaluatorId": "Builtin.FailureAnalysis"}], + ) + assert result["status"] == done + assert result["insightResults"] == [{"name": "Timeouts"}] + client._dp_client.start_batch_evaluation.assert_called_once_with( + batchEvaluationName="run", evaluators=[{"evaluatorId": "Builtin.FailureAnalysis"}] + ) + client._dp_client.get_batch_evaluation.assert_called_with(batchEvaluationId="b-1") + + @pytest.mark.parametrize("failed", ["FAILED", "STOPPED"]) + def test_raises_on_failure(self, failed): + client = self._make_client() + client._dp_client.start_batch_evaluation.return_value = {"batchEvaluationId": "b-1"} + client._dp_client.get_batch_evaluation.return_value = {"status": failed, "errorDetails": ["no spans"]} + with pytest.raises(RuntimeError, match="no spans"): + client.start_batch_evaluation_and_wait(batch_evaluation_name="run") diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 9b136480..fba4d2bb 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -193,3 +193,10 @@ def test_neither_returns_none(self): # Caller is responsible for combining this with namespace_path; helper just # reports "no legacy/exact value was provided." assert resolve_namespace_prefix_deprecation() is None + + +class TestWaitUntilTargetSet: + def test_any_status_in_target_set(self): + poll_fn = Mock(return_value={"status": "COMPLETED_WITH_ERRORS"}) + result = wait_until(poll_fn, {"COMPLETED", "COMPLETED_WITH_ERRORS"}, {"FAILED"}) + assert result["status"] == "COMPLETED_WITH_ERRORS"