diff --git a/docs/targets/bedrock_flows.md b/docs/targets/bedrock_flows.md new file mode 100644 index 0000000..e67a990 --- /dev/null +++ b/docs/targets/bedrock_flows.md @@ -0,0 +1,30 @@ +# Amazon Bedrock Flows + +Amazon Bedrock Flows offer the ability to link prompts, foundation models, and other AWS services into end-to-end workflows through a graphical UI. For more information, visit the AWS documentation [here](https://docs.aws.amazon.com/bedrock/latest/userguide/flows.html). + + +## Prerequisites + +The principal must have the following permissions: + +- [InvokeFlow](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_InvokeFlow.html) + + +## Configurations + +```yaml title="agenteval.yml" +target: + type: bedrock-flow + bedrock_flow_id: my-flow-id + bedrock_flow_alias_id: my-alias-id +``` + +`bedrock_flow_id` *(string)* + +The unique identifier of the Bedrock flow. Typically 10 characters uppercase alphanumeric. + +--- + +`bedrock_flow_alias_id` *(string)* + +The alias of the Bedrock flow. Typically 10 characters uppercase alphanumeric. diff --git a/docs/targets/index.md b/docs/targets/index.md index 0b9a20e..200b0a1 100644 --- a/docs/targets/index.md +++ b/docs/targets/index.md @@ -46,6 +46,7 @@ Configures the Boto3 client with the maximum number of retry attempts allowed. T ## Built-in targets - [Agents for Amazon Bedrock](bedrock_agents.md) +- [Amazon Bedrock Flows](bedrock_flows.md) - [Knowledge bases for Amazon Bedrock](bedrock_knowledge_bases.md) - [Amazon Q for Business](q_business.md) - [Amazon SageMaker endpoints](sagemaker_endpoints.md) diff --git a/src/agenteval/targets/bedrock_flow/__init__.py b/src/agenteval/targets/bedrock_flow/__init__.py new file mode 100644 index 0000000..ffc441d --- /dev/null +++ b/src/agenteval/targets/bedrock_flow/__init__.py @@ -0,0 +1,6 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .target import BedrockFlowTarget + +__all__ = ["BedrockFlowTarget"] diff --git a/src/agenteval/targets/bedrock_flow/target.py b/src/agenteval/targets/bedrock_flow/target.py new file mode 100644 index 0000000..7192ae2 --- /dev/null +++ b/src/agenteval/targets/bedrock_flow/target.py @@ -0,0 +1,72 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from agenteval.targets import Boto3Target, TargetResponse + +_SERVICE_NAME = "bedrock-agent-runtime" + + +class BedrockFlowTarget(Boto3Target): + """A target encapsulating an Amazon Bedrock Flow.""" + + def __init__(self, bedrock_flow_id: str, bedrock_flow_alias_id: str, **kwargs): + """Initialize the target. + + Args: + bedrock_flow_id (str): The unique identifier of the Bedrock flow. + bedrock_flow_alias_id (str): The alias of the Bedrock flow. + """ + super().__init__(boto3_service_name=_SERVICE_NAME, **kwargs) + self._bedrock_flow_id = bedrock_flow_id + self._bedrock_flow_alias_id = bedrock_flow_alias_id + + def invoke(self, prompt: str) -> TargetResponse: + """Invoke the target with a prompt. + + Args: + prompt (str): The prompt as a string. + + Returns: + TargetResponse + """ + args = { + "enableTrace": True, + "flowIdentifier": self._bedrock_flow_id, + "flowAliasIdentifier": self._bedrock_flow_alias_id, + "inputs": [ + { + # Although the API implies otherwise, Flows currently seem to be limited to + # follow this single-input node pattern anyway: + "content": {"document": prompt}, + "nodeName": "FlowInputNode", + "nodeOutputName": "document", + } + ], + } + + response = self.boto3_client.invoke_flow(**args) + + stream = response["responseStream"] + completion = "" + trace_data = [] + + for event in stream: + if "flowTraceEvent" in event: + trace_data.append(event["flowTraceEvent"].get("trace")) + if "flowOutputEvent" in event: + output_event = event["flowOutputEvent"] + # Testing 2024-12 suggests 'nodeType' actually not always present in API (despite + # the docs?): + if ( + "nodeType" not in output_event + or output_event["nodeType"] == "FlowOutputNode" + ): + completion += output_event.get("content", {}).get("document", "") + + errs = {k: v for k, v in event.items() if k.endswith("Exception")} + if errs: + raise ValueError(errs) + + return TargetResponse( + response=completion, data={"bedrock_flow_trace": trace_data} + ) diff --git a/src/agenteval/targets/target_factory.py b/src/agenteval/targets/target_factory.py index a20bb43..d70f69c 100644 --- a/src/agenteval/targets/target_factory.py +++ b/src/agenteval/targets/target_factory.py @@ -5,6 +5,7 @@ from agenteval.targets import BaseTarget from agenteval.targets.bedrock_agent import BedrockAgentTarget +from agenteval.targets.bedrock_flow import BedrockFlowTarget from agenteval.targets.bedrock_knowledge_base import BedrockKnowledgeBaseTarget from agenteval.targets.lexv2 import LexV2Target from agenteval.targets.q_business import QBusinessTarget @@ -13,6 +14,7 @@ _TARGET_MAP = { "bedrock-agent": BedrockAgentTarget, + "bedrock-flow": BedrockFlowTarget, "q-business": QBusinessTarget, "sagemaker-endpoint": SageMakerEndpointTarget, "bedrock-knowledge-base": BedrockKnowledgeBaseTarget, diff --git a/tests/src/agenteval/targets/bedrock_flow/__init__.py b/tests/src/agenteval/targets/bedrock_flow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/src/agenteval/targets/bedrock_flow/test_target.py b/tests/src/agenteval/targets/bedrock_flow/test_target.py new file mode 100644 index 0000000..0ce09f6 --- /dev/null +++ b/tests/src/agenteval/targets/bedrock_flow/test_target.py @@ -0,0 +1,62 @@ +import pytest + +from src.agenteval.targets.bedrock_flow import target +from src.agenteval.utils import aws + + +@pytest.fixture +def bedrock_flow_fixture(mocker): + mocker.patch.object(aws.boto3, "Session") + + fixture = target.BedrockFlowTarget( + bedrock_flow_id="test-agent-id", + bedrock_flow_alias_id="test-alias-id", + aws_profile="test-profile", + aws_region="us-west-2", + ) + + return fixture + + +class TestBedrockFlowTarget: + + def test_invoke(self, mocker, bedrock_flow_fixture): + mock_invoke_flow = mocker.patch.object( + bedrock_flow_fixture.boto3_client, "invoke_flow" + ) + + mock_invoke_flow.return_value = { + "responseStream": [ + { + "flowTraceEvent": { + "trace": {"fascinating": "trace content"}, + }, + }, + { + "flowOutputEvent": { + "content": {"document": "Meow!"}, + "nodeType": "FlowOutputNode", + }, + }, + { + "flowTraceEvent": { + "trace": {"some garbage": "or other"}, + }, + }, + { + "flowCompletionEvent": { + "completionReason": "SUCCESS", + }, + }, + ] + } + + response = bedrock_flow_fixture.invoke("test prompt") + + assert response.response == "Meow!" + assert response.data == { + "bedrock_flow_trace": [ + {"fascinating": "trace content"}, + {"some garbage": "or other"}, + ], + }