Merge pull request #96 from athewsey/feat/flow

Add Target for Bedrock Flows

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.

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)

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

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}
+        )

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,

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"},
+            ],
+        }