Add message attributes to otel traces

python/packages/autogen-core/src/autogen_core/_routed_agent.py

@@ -18,6 +18,7 @@ from typing import (
     runtime_checkable,
 )
 
+from opentelemetry.trace import get_tracer
 from ._base_agent import BaseAgent
 from ._message_context import MessageContext
 from ._serialization import MessageSerializer, try_get_known_serializers_for_type
@@ -38,6 +39,7 @@ ProducesT = TypeVar("ProducesT", covariant=True)
 # Revisit this later to see if we can remove the ignore.
 @runtime_checkable
 class MessageHandler(Protocol[AgentT, ReceivesT, ProducesT]):  # type: ignore
+    name: str
     target_types: Sequence[type]
     produces_types: Sequence[type]
     is_message_handler: Literal[True]
@@ -157,6 +159,7 @@ def message_handler(
             return return_value
 
         wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, ProducesT], wrapper)
+        wrapper_handler.name = func.__name__
         wrapper_handler.target_types = list(target_types)
         wrapper_handler.produces_types = list(return_types)
         wrapper_handler.is_message_handler = True
@@ -276,6 +279,7 @@ def event(
             return None
 
         wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, None], wrapper)
+        wrapper_handler.name = func.__name__
         wrapper_handler.target_types = list(target_types)
         wrapper_handler.produces_types = list(return_types)
         wrapper_handler.is_message_handler = True
@@ -397,6 +401,7 @@ def rpc(
             return return_value
 
         wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, ProducesT], wrapper)
+        wrapper_handler.name = func.__name__
         wrapper_handler.target_types = list(target_types)
         wrapper_handler.produces_types = list(return_types)
         wrapper_handler.is_message_handler = True
@@ -482,7 +487,14 @@ class RoutedAgent(BaseAgent):
             # Call the first handler whose router returns True and then return the result.
             for h in handlers:
                 if h.router(message, ctx):
-                    return await h(self, message, ctx)
+                    with get_tracer("routed_agent").start_as_current_span(
+                        f"message_handler_{h.name}",
+                        attributes={
+                            "handler_name": h.name,
+                        },
+                    ):
+                        # Call the handler
+                        return await h(self, message, ctx)
         return await self.on_unhandled_message(message, ctx)  # type: ignore
 
     async def on_unhandled_message(self, message: Any, ctx: MessageContext) -> None:

python/packages/autogen-core/src/autogen_core/_single_threaded_agent_runtime.py

@@ -440,8 +440,19 @@ class SingleThreadedAgentRuntime(AgentRuntime):
                     cancellation_token=message_envelope.cancellation_token,
                     message_id=message_envelope.message_id,
                 )
-                with self._tracer_helper.trace_block("process", recipient_agent.id, parent=message_envelope.metadata):
-                    with MessageHandlerContext.populate_context(recipient_agent.id):
+                with MessageHandlerContext.populate_context(recipient_agent.id):
+                    with self._tracer_helper.trace_block(
+                        "process",
+                        recipient_agent.id,
+                        parent=message_envelope.metadata,
+                        attributes={
+                            "messaging.type": type(message_envelope.message).__name__,
+                            "messaging.content": self._try_serialize(message_envelope.message),
+                            "messaging.recipient_id": str(recipient_agent.id),
+                            "messaging.recipient_class": recipient_agent.__class__.__name__,
+                            "messaging.sender_id": str(message_envelope.sender),
+                        },
+                    ):
                         response = await recipient_agent.on_message(
                             message_envelope.message,
                             ctx=message_context,
@@ -527,8 +538,20 @@ class SingleThreadedAgentRuntime(AgentRuntime):
                     agent = await self._get_agent(agent_id)
 
                     async def _on_message(agent: Agent, message_context: MessageContext) -> Any:
-                        with self._tracer_helper.trace_block("process", agent.id, parent=message_envelope.metadata):
-                            with MessageHandlerContext.populate_context(agent.id):
+                        with MessageHandlerContext.populate_context(agent.id):
+                            with self._tracer_helper.trace_block(
+                                "process",
+                                agent.id,
+                                parent=message_envelope.metadata,
+                                attributes={
+                                    "messaging.type": type(message_envelope.message).__name__,
+                                    "messaging.content": self._try_serialize(message_envelope.message),
+                                    "messaging.recipient_id": str(agent.id),
+                                    "messaging.recipient_class": agent.__class__.__name__,
+                                    "messaging.topic_id": str(message_envelope.topic_id),
+                                    "messaging.sender_id": str(message_envelope.sender),
+                                },
+                            ):
                                 try:
                                     return await agent.on_message(
                                         message_envelope.message,