Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions src/bedrock_agentcore/_utils/polling.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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",
Expand All @@ -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.
Expand All @@ -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")
Expand Down
2 changes: 2 additions & 0 deletions src/bedrock_agentcore/evaluation/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -94,6 +95,7 @@
"CloudWatchAgentSpanCollector",
"Dataset",
"DatasetClient",
"RecommendationClient",
"DatasetProvider",
"EvaluationClient",
"EvaluationResult",
Expand Down
43 changes: 43 additions & 0 deletions src/bedrock_agentcore/evaluation/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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 = {
Expand Down Expand Up @@ -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,
Expand Down
134 changes: 134 additions & 0 deletions src/bedrock_agentcore/evaluation/recommendation_client.py
Original file line number Diff line number Diff line change
@@ -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",
)
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}

Expand Down
63 changes: 62 additions & 1 deletion tests/bedrock_agentcore/evaluation/test_client_passthrough.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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")
Loading
Loading