feat(evaluation): evaluate existing sessions and add RecommendationClient (#708)

* 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.

* 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.

---------

Co-authored-by: Sundar Raghavan <sdraghav@amazon.com>

Author: sundargthb

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

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

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,

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

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,
             }
         }
 

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