models - structured output - iterative callback

Author: pgrayy

src/strands/models/anthropic.py

@@ -387,15 +387,15 @@ def structured_output(
             prompt(Messages): The prompt messages to use for the agent.
             callback_handler(Optional[Callable]): Optional callback handler for processing events. Defaults to None.
         """
+        callback_handler = callback_handler or PrintingCallbackHandler()
         tool_spec = convert_pydantic_to_tool_spec(output_model)
 
         response = self.converse(messages=prompt, tool_specs=[tool_spec])
-        # process the stream and get the tool use input
-        results = process_stream(
-            response, callback_handler=callback_handler or PrintingCallbackHandler(), messages=prompt
-        )
-
-        stop_reason, messages, _, _, _ = results
+        for event in process_stream(response, prompt):
+            if "callback" in event:
+                callback_handler(**event["callback"])
+        else:
+            stop_reason, messages, _, _ = event["stop"]
 
         if stop_reason != "tool_use":
             raise ValueError("No valid tool use or tool use input was found in the Anthropic response.")

src/strands/models/bedrock.py

@@ -504,15 +504,15 @@ def structured_output(
             prompt(Messages): The prompt messages to use for the agent.
             callback_handler(Optional[Callable]): Optional callback handler for processing events. Defaults to None.
         """
+        callback_handler = callback_handler or PrintingCallbackHandler()
         tool_spec = convert_pydantic_to_tool_spec(output_model)
 
         response = self.converse(messages=prompt, tool_specs=[tool_spec])
-        # process the stream and get the tool use input
-        results = process_stream(
-            response, callback_handler=callback_handler or PrintingCallbackHandler(), messages=prompt
-        )
-
-        stop_reason, messages, _, _, _ = results
+        for event in process_stream(response, prompt):
+            if "callback" in event:
+                callback_handler(**event["callback"])
+        else:
+            stop_reason, messages, _, _, _ = event["stop"]
 
         if stop_reason != "tool_use":
             raise ValueError("No valid tool use or tool use input was found in the Bedrock response.")

tests-integ/test_model_anthropic.py

@@ -34,7 +34,7 @@ def tool_weather() -> str:
 
 @pytest.fixture
 def system_prompt():
-    return "You are an AI assistant that uses & instead of ."
+    return "You are an AI assistant."
 
 
 @pytest.fixture
@@ -47,7 +47,7 @@ def test_agent(agent):
     result = agent("What is the time and weather in New York?")
     text = result.message["content"][0]["text"].lower()
 
-    assert all(string in text for string in ["12:00", "sunny", "&"])
+    assert all(string in text for string in ["12:00", "sunny"])
 
 
 @pytest.mark.skipif("ANTHROPIC_API_KEY" not in os.environ, reason="ANTHROPIC_API_KEY environment variable missing")