From 353a08514c5fafebe225571125f9ef5081109c16 Mon Sep 17 00:00:00 2001 From: poshinchen Date: Mon, 11 Aug 2025 15:40:11 -0400 Subject: [PATCH 1/2] feat: add structured_output_span --- src/strands/agent/agent.py | 65 ++++++++++++++++++++----------- tests/strands/agent/test_agent.py | 39 +++++++++++++++++++ 2 files changed, 82 insertions(+), 22 deletions(-) diff --git a/src/strands/agent/agent.py b/src/strands/agent/agent.py index 2022142c6..77a212e04 100644 --- a/src/strands/agent/agent.py +++ b/src/strands/agent/agent.py @@ -33,7 +33,7 @@ from ..models.model import Model from ..session.session_manager import SessionManager from ..telemetry.metrics import EventLoopMetrics -from ..telemetry.tracer import get_tracer +from ..telemetry.tracer import get_tracer, serialize from ..tools.registry import ToolRegistry from ..tools.watcher import ToolWatcher from ..types.content import ContentBlock, Message, Messages @@ -445,27 +445,48 @@ async def structured_output_async( ValueError: If no conversation history or prompt is provided. """ self.hooks.invoke_callbacks(BeforeInvocationEvent(agent=self)) - - try: - if not self.messages and not prompt: - raise ValueError("No conversation history or prompt provided") - - # Create temporary messages array if prompt is provided - if prompt: - content: list[ContentBlock] = [{"text": prompt}] if isinstance(prompt, str) else prompt - temp_messages = self.messages + [{"role": "user", "content": content}] - else: - temp_messages = self.messages - - events = self.model.structured_output(output_model, temp_messages, system_prompt=self.system_prompt) - async for event in events: - if "callback" in event: - self.callback_handler(**cast(dict, event["callback"])) - - return event["output"] - - finally: - self.hooks.invoke_callbacks(AfterInvocationEvent(agent=self)) + with self.tracer.tracer.start_as_current_span( + "execute_structured_output", kind=trace_api.SpanKind.CLIENT + ) as structured_output_span: + try: + if not self.messages and not prompt: + raise ValueError("No conversation history or prompt provided") + # Create temporary messages array if prompt is provided + if prompt: + content: list[ContentBlock] = [{"text": prompt}] if isinstance(prompt, str) else prompt + temp_messages = self.messages + [{"role": "user", "content": content}] + else: + temp_messages = self.messages + + structured_output_span.set_attributes( + { + "gen_ai.system": "strands-agents", + "gen_ai.agent.name": self.name, + "gen_ai.agent.id": self.agent_id, + "gen_ai.operation.name": "execute_structured_output", + } + ) + for message in temp_messages: + structured_output_span.add_event( + f"gen_ai.{message['role']}.message", + attributes={"role": message["role"], "content": serialize(message["content"])}, + ) + if self.system_prompt: + structured_output_span.add_event( + "gen_ai.tool.message", + attributes={"role": "tool", "content": serialize([{"text": self.system_prompt}])}, + ) + events = self.model.structured_output(output_model, temp_messages, system_prompt=self.system_prompt) + async for event in events: + if "callback" in event: + self.callback_handler(**cast(dict, event["callback"])) + structured_output_span.add_event( + "gen_ai.choice", attributes={"message": serialize(event["output"].model_dump())} + ) + return event["output"] + + finally: + self.hooks.invoke_callbacks(AfterInvocationEvent(agent=self)) async def stream_async(self, prompt: Union[str, list[ContentBlock]], **kwargs: Any) -> AsyncIterator[Any]: """Process a natural language prompt and yield events as an async iterator. diff --git a/tests/strands/agent/test_agent.py b/tests/strands/agent/test_agent.py index c27243dfe..fdce7c368 100644 --- a/tests/strands/agent/test_agent.py +++ b/tests/strands/agent/test_agent.py @@ -980,6 +980,14 @@ def test_agent_callback_handler_custom_handler_used(): def test_agent_structured_output(agent, system_prompt, user, agenerator): + # Setup mock tracer and span + mock_strands_tracer = unittest.mock.MagicMock() + mock_otel_tracer = unittest.mock.MagicMock() + mock_span = unittest.mock.MagicMock() + mock_strands_tracer.tracer = mock_otel_tracer + mock_otel_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span + agent.tracer = mock_strands_tracer + agent.model.structured_output = unittest.mock.Mock(return_value=agenerator([{"output": user}])) prompt = "Jane Doe is 30 years old and her email is jane@doe.com" @@ -999,8 +1007,34 @@ def test_agent_structured_output(agent, system_prompt, user, agenerator): type(user), [{"role": "user", "content": [{"text": prompt}]}], system_prompt=system_prompt ) + mock_span.set_attributes.assert_called_once_with( + { + "gen_ai.system": "strands-agents", + "gen_ai.agent.name": "Strands Agents", + "gen_ai.agent.id": "default", + "gen_ai.operation.name": "execute_structured_output", + } + ) + + mock_span.add_event.assert_any_call( + "gen_ai.user.message", + attributes={"role": "user", "content": '[{"text": "Jane Doe is 30 years old and her email is jane@doe.com"}]'}, + ) + + mock_span.add_event.assert_called_with( + "gen_ai.choice", + attributes={"message": json.dumps(user.model_dump())}, + ) + def test_agent_structured_output_multi_modal_input(agent, system_prompt, user, agenerator): + # Setup mock tracer and span + mock_strands_tracer = unittest.mock.MagicMock() + mock_otel_tracer = unittest.mock.MagicMock() + mock_span = unittest.mock.MagicMock() + mock_strands_tracer.tracer = mock_otel_tracer + mock_otel_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span + agent.tracer = mock_strands_tracer agent.model.structured_output = unittest.mock.Mock(return_value=agenerator([{"output": user}])) prompt = [ @@ -1030,6 +1064,11 @@ def test_agent_structured_output_multi_modal_input(agent, system_prompt, user, a type(user), [{"role": "user", "content": prompt}], system_prompt=system_prompt ) + mock_span.add_event.assert_called_with( + "gen_ai.choice", + attributes={"message": json.dumps(user.model_dump())}, + ) + @pytest.mark.asyncio async def test_agent_structured_output_in_async_context(agent, user, agenerator): From 0b9811276095d6c1a8fe043607863cde865c575e Mon Sep 17 00:00:00 2001 From: poshinchen Date: Mon, 11 Aug 2025 15:40:11 -0400 Subject: [PATCH 2/2] feat: add structured_output_span --- src/strands/agent/agent.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/strands/agent/agent.py b/src/strands/agent/agent.py index 77a212e04..43b5cbf8c 100644 --- a/src/strands/agent/agent.py +++ b/src/strands/agent/agent.py @@ -473,8 +473,8 @@ async def structured_output_async( ) if self.system_prompt: structured_output_span.add_event( - "gen_ai.tool.message", - attributes={"role": "tool", "content": serialize([{"text": self.system_prompt}])}, + "gen_ai.system.message", + attributes={"role": "system", "content": serialize([{"text": self.system_prompt}])}, ) events = self.model.structured_output(output_model, temp_messages, system_prompt=self.system_prompt) async for event in events: