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