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")