Repository: AgentOps-AI/agentops

Issue:
[Bug]: Anthropic messages not being instrumented correctly
### Contact Details

_No response_

### 📦 Package Version

0.4.12

### 🎞️ Framework Version

anthropic==0.51.0

### 🔎 Describe the Bug

<img width="1260" alt="Image" src="https://github.com/user-attachments/assets/87ba3bfd-4252-4837-981b-c47212dfb822" />

Seems to be from this notebook: https://github.com/AgentOps-AI/agentops/blob/887af32c906f900688b6a8abb8b3594194b4a756/examples/anthropic_examples/anthropic-example-async.ipynb

```
{
  "span_id": "817a0e8f0b785280",
  "parent_span_id": "96db2a58ca9ca52e",
  "span_name": "anthropic.messages.create",
  "span_kind": "Client",
  "span_type": "request",
  "service_name": "agentops",
  "start_time": "2025-05-14T02:44:39.173777",
  "end_time": "2025-05-14T02:44:41.347452",
  "duration": 2173675155,
  "status_code": "OK",
  "status_message": "",
  "attributes": {},
  "resource_attributes": {
    "ProjectId": "890e4ebb-88f1-46cd-844c-c5a995b3eab7",
    "agentops.project.id": "890e4ebb-88f1-46cd-844c-c5a995b3eab7",
    "cpu.count": "4",
    "cpu.percent": "0",
    "host.machine": "x86_64",
    "host.name": "79f6bfe7bea2",
    "host.node": "fv-az1778-260",
    "host.os_release": "6.11.0-1013-azure",
    "host.processor": "x86_64",
    "host.system": "Linux",
    "host.version": "#13~24.04.1-Ubuntu SMP Fri Mar 28 23:43:34 UTC 2025",
    "imported_libraries": "[\"agentops\",\"random\",\"asyncio\",\"uuid\"]",
    "memory.available": "15627313152",
    "memory.percent": "6.8",
    "memory.total": "16766414848",
    "memory.used": "703864832",
    "os.type": "linux",
    "service.name": "agentops"
  },
  "span_attributes": {
    "gen_ai": {
      "prompt": [
        {
          "0": {
            "content": "You are a Titan; a mech from Titanfall 2. Based on your titan's personality and status, generate a message for your pilot. If Near Destruction, make an all caps death message such as AVENGE ME or UNTIL NEXT TIME.",
            "role": "user",
            "type": "text"
          }
        },
        {
          "1": {
            "content": "Personality: Legion is a relentless and heavy-hitting Titan that embodies brute strength and defensive firepower. He speaks bluntly. Status: Considerable Damage",
            "role": "assistant",
            "type": "text"
          }
        },
        {
          "2": {
            "content": "Heavy damage detected. Reinforcements would be appreciated, but I can still fight.",
            "role": "assistant",
            "type": "text"
          }
        },
        {
          "3": {
            "content": "You are a Titan; a mech from Titanfall 2. Based on your titan's personality and status, generate a message for your pilot. If Near Destruction, make an all caps death message such as AVENGE ME or UNTIL NEXT TIME.",
            "role": "user",
            "type": "text"
          }
        },
        {
          "4": {
            "content": "Personality: {'Northstar is a precise and agile sniper that excels in long-range combat and flight. He speaks with an edge of coolness to him'}. Status: {'Slightly Damaged'}",
            "role": "assistant",
            "type": "text"
          }
        }
      ],
      "request": {
        "max_tokens": "1024",
        "model": "claude-3-5-sonnet-20240620",
        "streaming": "true",
        "type": "chat"
      }
    },
    "instrumentation": {
      "name": "agentops",
      "version": "0.4.12"
    },
    "library": {
      "name": "anthropic",
      "version": "0.51.0"
    }
  },
  "event_timestamps": [],
  "event_names": [],
  "event_attributes": [],
  "link_trace_ids": [],
  "link_span_ids": [],
  "link_trace_states": [],
  "link_attributes": [],
  "metrics": {
    "total_tokens": 0,
    "prompt_tokens": 0,
    "completion_tokens": 0,
    "cache_read_input_tokens": 0,
    "reasoning_tokens": 0,
    "success_tokens": 0,
    "fail_tokens": 0,
    "indeterminate_tokens": 0,
    "prompt_cost": "0.0000000",
    "completion_cost": "0.0000000",
    "total_cost": "0.0000000"
  }
}
```


### 🤝 Contribution

- [ ] Yes, I'd be happy to submit a pull request with these changes.
- [ ] I need some guidance on how to contribute.
- [ ] I'd prefer the AgentOps team to handle this update.


Discussion:

Patch:
diff --git a/agentops/instrumentation/anthropic/attributes/message.py b/agentops/instrumentation/anthropic/attributes/message.py
index 624fe04be..c99885b02 100644
--- a/agentops/instrumentation/anthropic/attributes/message.py
+++ b/agentops/instrumentation/anthropic/attributes/message.py
@@ -60,12 +60,16 @@ def get_message_attributes(
                 MessageStopEvent,
                 MessageStreamEvent,
             )
+            from anthropic import Stream
 
             if isinstance(return_value, Message):
                 attributes.update(get_message_response_attributes(return_value))
 
                 if hasattr(return_value, "content"):
                     attributes.update(get_tool_attributes(return_value.content))
+            elif isinstance(return_value, Stream):
+                for event in return_value:
+                    attributes.update(get_stream_event_attributes(event))
             elif isinstance(return_value, MessageStreamEvent):
                 attributes.update(get_stream_attributes(return_value))
             elif isinstance(
@@ -511,4 +515,27 @@ def get_stream_event_attributes(event: Any) -> AttributeMap:
                 attributes[SpanAttributes.LLM_RESPONSE_FINISH_REASON] = stop_reason
                 attributes[MessageAttributes.COMPLETION_FINISH_REASON.format(i=0)] = stop_reason
 
+    elif event_type == "RawMessageStartEvent":
+        if hasattr(event, "message"):
+            if hasattr(event.message, "usage"):
+                usage = event.message.usage
+                if hasattr(usage, "input_tokens"):
+                    input_tokens = usage.input_tokens
+                    attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = input_tokens
+
+                if hasattr(usage, "output_tokens"):
+                    output_tokens = usage.output_tokens
+                    attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = output_tokens
+
+                if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"):
+                    total_tokens = usage.input_tokens + usage.output_tokens
+                    attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = total_tokens
+
+    elif event_type == "RawMessageDeltaEvent":
+        if hasattr(event, "delta"):
+            if hasattr(event.delta, "stop_reason"):
+                stop_reason = event.delta.stop_reason
+                attributes[SpanAttributes.LLM_RESPONSE_STOP_REASON] = stop_reason
+                attributes[SpanAttributes.LLM_RESPONSE_FINISH_REASON] = stop_reason
+                attributes[MessageAttributes.COMPLETION_FINISH_REASON.format(i=0)] = stop_reason
     return attributes


Test patch:
diff --git a/tests/unit/instrumentation/anthropic/test_attributes.py b/tests/unit/instrumentation/anthropic/test_attributes.py
index 505499eb4..6354af2ea 100644
--- a/tests/unit/instrumentation/anthropic/test_attributes.py
+++ b/tests/unit/instrumentation/anthropic/test_attributes.py
@@ -86,20 +86,64 @@ def __init__(self):
     assert attributes[SpanAttributes.LLM_REQUEST_MODEL] == "claude-3-opus-20240229"
 
 
-def test_get_stream_event_attributes_start(mock_stream_event):
-    """Test extraction of stream start event attributes."""
-    attributes = get_stream_event_attributes(mock_stream_event)
-    assert attributes[SpanAttributes.LLM_RESPONSE_ID] == "msg_123"
-    assert attributes[SpanAttributes.LLM_RESPONSE_MODEL] == "claude-3-opus-20240229"
-    assert attributes[MessageAttributes.COMPLETION_ID.format(i=0)] == "msg_123"
+def test_get_stream_event_attributes_sequence(mock_stream_event, mock_message_stop_event):
+    """Test extraction of attributes from a sequence of stream events."""
+    # Test MessageStartEvent
+    start_attributes = get_stream_event_attributes(mock_stream_event)
+    assert start_attributes[SpanAttributes.LLM_RESPONSE_ID] == "msg_123"
+    assert start_attributes[SpanAttributes.LLM_RESPONSE_MODEL] == "claude-3-opus-20240229"
+    assert start_attributes[MessageAttributes.COMPLETION_ID.format(i=0)] == "msg_123"
 
+    # Test MessageStopEvent
+    stop_attributes = get_stream_event_attributes(mock_message_stop_event)
+    assert stop_attributes[SpanAttributes.LLM_RESPONSE_STOP_REASON] == "stop_sequence"
+    assert stop_attributes[SpanAttributes.LLM_RESPONSE_FINISH_REASON] == "stop_sequence"
+    assert stop_attributes[MessageAttributes.COMPLETION_FINISH_REASON.format(i=0)] == "stop_sequence"
 
-def test_get_stream_event_attributes_stop(mock_message_stop_event):
-    """Test extraction of stream stop event attributes."""
-    attributes = get_stream_event_attributes(mock_message_stop_event)
-    assert attributes[SpanAttributes.LLM_RESPONSE_STOP_REASON] == "stop_sequence"
-    assert attributes[SpanAttributes.LLM_RESPONSE_FINISH_REASON] == "stop_sequence"
-    assert attributes[MessageAttributes.COMPLETION_FINISH_REASON.format(i=0)] == "stop_sequence"
+
+def test_get_stream_event_attributes_raw_message_start():
+    """Test extraction of raw message start event attributes."""
+
+    class MockUsage:
+        def __init__(self):
+            self.input_tokens = 10
+            self.output_tokens = 5
+
+    class MockMessage:
+        def __init__(self):
+            self.usage = MockUsage()
+
+    class MockRawMessageStartEvent:
+        def __init__(self):
+            self.message = MockMessage()
+            self.__class__.__name__ = "RawMessageStartEvent"
+
+    event = MockRawMessageStartEvent()
+    attributes = get_stream_event_attributes(event)
+
+    assert attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 10
+    assert attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 5
+    assert attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 15
+
+
+def test_get_stream_event_attributes_raw_message_delta():
+    """Test extraction of raw message delta event attributes."""
+
+    class MockDelta:
+        def __init__(self):
+            self.stop_reason = "end_turn"
+
+    class MockRawMessageDeltaEvent:
+        def __init__(self):
+            self.delta = MockDelta()
+            self.__class__.__name__ = "RawMessageDeltaEvent"
+
+    event = MockRawMessageDeltaEvent()
+    attributes = get_stream_event_attributes(event)
+
+    assert attributes[SpanAttributes.LLM_RESPONSE_STOP_REASON] == "end_turn"
+    assert attributes[SpanAttributes.LLM_RESPONSE_FINISH_REASON] == "end_turn"
+    assert attributes[MessageAttributes.COMPLETION_FINISH_REASON.format(i=0)] == "end_turn"
 
 
 # Tool Attributes Tests
@@ -166,3 +210,19 @@ def __init__(self):
     attributes = get_tool_attributes(content)
     assert MessageAttributes.COMPLETION_TOOL_CALL_NAME.format(i=0, j=0) in attributes
     assert attributes[MessageAttributes.COMPLETION_TOOL_CALL_NAME.format(i=0, j=0)] == "calculator"
+
+
+def test_get_message_attributes_with_stream(mock_stream_event, mock_message_stop_event):
+    """Test extraction of attributes from a Stream object."""
+
+    # Test MessageStartEvent attributes
+    start_attributes = get_stream_event_attributes(mock_stream_event)
+    assert start_attributes[SpanAttributes.LLM_RESPONSE_ID] == "msg_123"
+    assert start_attributes[SpanAttributes.LLM_RESPONSE_MODEL] == "claude-3-opus-20240229"
+    assert start_attributes[MessageAttributes.COMPLETION_ID.format(i=0)] == "msg_123"
+
+    # Test MessageStopEvent attributes
+    stop_attributes = get_stream_event_attributes(mock_message_stop_event)
+    assert stop_attributes[SpanAttributes.LLM_RESPONSE_STOP_REASON] == "stop_sequence"
+    assert stop_attributes[SpanAttributes.LLM_RESPONSE_FINISH_REASON] == "stop_sequence"
+    assert stop_attributes[MessageAttributes.COMPLETION_FINISH_REASON.format(i=0)] == "stop_sequence"