simplified example

Author: AWSChris

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

@@ -68,7 +68,7 @@ def invoke_agent(self, agent_id, agent_alias_id, session_id, prompt):
 
     # snippet-end:[python.example_code.bedrock-agent-runtime.InvokeAgent]
 
-    # snippet-start:[python.example_code.bedrock-agent-runtime.InvokeFlow]
+        # 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.
@@ -81,10 +81,10 @@ def invoke_flow(self, flow_id, flow_alias_id, input_data, execution_id):
             execution_id: Execution ID for continuing a flow. Use the value None on first run.
 
         Returns:
-            Dict containing flow_complete status, input_required info, and execution_id
+            The response
         """
         try:
-            response = None
+      
             request_params = None
 
             if execution_id is None:
@@ -109,37 +109,19 @@ def invoke_flow(self, flow_id, flow_alias_id, input_data, execution_id):
             if "executionId" not in request_params:
                 execution_id = response['executionId']
 
-            input_required = None
-            flow_status = ""
+            result = ""
 
-            # Process the streaming response
+            # Get the streaming response
             for event in response['responseStream']:
+                result = result + str(event) + '\n'
+            print(result)
 
-                # Check if flow is complete.
-                if 'flowCompletionEvent' in event:
-                    flow_status = event['flowCompletionEvent']['completionReason']
-
-                # Check if more input us needed from user.
-                elif 'flowMultiTurnInputRequestEvent' in event:
-                    input_required = event
-
-                # Print the model output.
-                elif 'flowOutputEvent' in event:
-                    print(event['flowOutputEvent']['content']['document'])
-
-                # Log trace events.
-                elif 'flowTraceEvent' in event:
-                    logger.info("Flow trace:  %s", event['flowTraceEvent'])
-
-            return {
-                "flow_status": flow_status,
-                "input_required": input_required,
-                "execution_id": execution_id
-            }
         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]