Python: InvokeFlow action Amazon Bedrock (awsdocs#7216)

Author: AWSChris

.doc_gen/metadata/bedrock-agent-runtime_metadata.yaml

@@ -30,6 +30,14 @@ bedrock-agent-runtime_InvokeFlow:
             - description:
               snippet_files:
                 - javascriptv3/example_code/bedrock-agent-runtime/actions/invoke-flow.js
+    Python:
+      versions:
+        - sdk_version: 3
+          github: python/example_code/bedrock-agent-runtime
+          excerpts:
+            - description: Invoke a flow.
+              snippet_tags:
+                - python.example_code.bedrock-agent-runtime.InvokeFlow
   services:
     bedrock-agent-runtime: {InvokeFlow}
 

python/example_code/bedrock-agent-runtime/README.md

@@ -46,6 +46,7 @@ Code examples that show you how to perform the essential operations within a ser
 Code excerpts that show you how to call individual service functions.
 
 - [InvokeAgent](bedrock_agent_runtime_wrapper.py#L33)
+- [InvokeFlow](bedrock_agent_runtime_wrapper.py#L71)
 
 
 <!--custom.examples.start-->

python/example_code/bedrock-agent-runtime/bedrock_agent_runtime_wrapper.py

@@ -68,4 +68,58 @@ def invoke_agent(self, agent_id, agent_alias_id, session_id, prompt):
 
     # snippet-end:[python.example_code.bedrock-agent-runtime.InvokeAgent]
 
-# snippet-end:[python.example_code.bedrock-agent-runtime.BedrockAgentsRuntimeWrapper.class]
+    # snippet-start:[python.example_code.bedrock-agent-runtime.InvokeFlow]
+    def invoke_flow(self, flow_id, flow_alias_id, input_data, execution_id):
+        """
+        Invoke an Amazon Bedrock flow and handle the response stream.
+
+        Args:
+            param flow_id: The ID of the flow to invoke.
+            param flow_alias_id: The alias ID of the flow.
+            param input_data: Input data for the flow.
+            param execution_id: Execution ID for continuing a flow. Use the value None on first run.
+
+        Return: Response from the flow.
+        """
+        try:
+      
+            request_params = None
+
+            if execution_id is None:
+                # Don't pass execution ID for first run.
+                request_params = {
+                    "flowIdentifier": flow_id,
+                    "flowAliasIdentifier": flow_alias_id,
+                    "inputs": input_data,
+                    "enableTrace": True
+                }
+            else:
+                request_params = {
+                    "flowIdentifier": flow_id,
+                    "flowAliasIdentifier": flow_alias_id,
+                    "executionId": execution_id,
+                    "inputs": input_data,
+                    "enableTrace": True
+                }
+
+            response = self.agents_runtime_client.invoke_flow(**request_params)
+
+            if "executionId" not in request_params:
+                execution_id = response['executionId']
+
+            result = ""
+
+            # Get the streaming response
+            for event in response['responseStream']:
+                result = result + str(event) + '\n'
+            print(result)
+
+        except ClientError as e:
+            logger.error("Couldn't invoke flow %s.", {e})
+            raise
+
+        return result
+
+    # snippet-end:[python.example_code.bedrock-agent-runtime.InvokeFlow]
+
+# snippet-end:[python.example_code.bedrock-agent-runtime.BedrockAgentsRuntimeWrapper.class]

python/example_code/bedrock-agent-runtime/requirements.txt

@@ -1,5 +1,5 @@
-boto3==1.33.8
-botocore==1.33.8
+boto3==1.36.13
+botocore==1.36.13
 colorama==0.4.6
 iniconfig==2.0.0
 jmespath==1.0.1

python/example_code/bedrock-agent-runtime/test/test_bedrock_agent_runtime_wrapper.py

@@ -40,8 +40,52 @@ async def test_invoke_agent(make_stubber, error_code):
     if error_code is None:
         wrapper.invoke_agent(agent_id, agent_alias_id, session_id, prompt)
     else:
-        with pytest.raises(ClientError):
+        with pytest.raises(ClientError) as exc_info:
             async for _ in wrapper.invoke_agent(
                 agent_id, agent_alias_id, session_id, prompt
             ):
                 assert exc_info.value.response["Error"]["Code"] == error_code
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("error_code", [None, "ClientError"])
+async def test_invoke_flow(make_stubber, error_code):
+    runtime_client = boto3.client(
+        service_name="bedrock-agent-runtime", region_name="us-east-1"
+    )
+    stubber = make_stubber(runtime_client)
+    wrapper = BedrockAgentRuntimeWrapper(runtime_client)
+
+    flow_id = "FAKE_AGENT_ID"
+    flow_alias_id = "FAKE_AGENT_ALIAS_ID"
+    execution_id = "FAKE_SESSION_ID"
+
+
+    inputs = [
+        {
+            "content": {
+                "document": "hello!"
+            },
+            "nodeName": "FlowInputNode",
+            "nodeOutputName": "document"
+        }
+    ]
+
+    expected_params = {
+        "enableTrace" : True,
+        "flowIdentifier": flow_id,
+        "flowAliasIdentifier": flow_alias_id,
+        "executionId": execution_id,
+        "inputs": inputs
+    }
+    response = {"responseStream": {}, "executionId": execution_id}
+
+    stubber.stub_invoke_flow(expected_params, response, error_code=error_code)
+
+    if error_code is None:
+        result = wrapper.invoke_flow(flow_id, flow_alias_id, inputs, execution_id)
+        assert result is not None
+    else:
+        with pytest.raises(ClientError) as exc_info:
+             wrapper.invoke_flow(flow_id, flow_alias_id, inputs, execution_id)
+         
+        assert exc_info.value.response["Error"]["Code"] == error_code

python/test_tools/bedrock_agent_runtime_stubber.py

@@ -31,3 +31,7 @@ def stub_invoke_agent(self, expected_params, response, error_code=None):
         self._stub_bifurcator(
             "invoke_agent", expected_params, response, error_code=error_code
         )
+    def stub_invoke_flow(self, expected_params, response, error_code=None):
+        self._stub_bifurcator(
+            "invoke_flow", expected_params, response, error_code=error_code
+        )