-
Notifications
You must be signed in to change notification settings - Fork 73
fix(ai): preserve Anthropic cache-write TTL breakdowns #746
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| --- | ||
| pypi/posthog: patch | ||
| --- | ||
|
|
||
| Preserve Anthropic cache-write TTL breakdowns across Python SDK AI integrations. | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -40,6 +40,7 @@ class _GenerationData: | |
| output_tokens: int = 0 | ||
| cache_read_input_tokens: int = 0 | ||
| cache_creation_input_tokens: int = 0 | ||
| raw_usage: Optional[Dict[str, Any]] = None | ||
| start_time: float = 0.0 | ||
| end_time: float = 0.0 | ||
| span_id: str = field(default_factory=lambda: str(uuid.uuid4())) | ||
|
|
@@ -77,9 +78,11 @@ def process_stream_event(self, event: "StreamEvent") -> None: | |
| self._current.cache_creation_input_tokens = usage.get( | ||
| "cache_creation_input_tokens", 0 | ||
| ) | ||
| self._current.raw_usage = dict(usage) | ||
|
|
||
| elif event_type == "message_delta" and self._current is not None: | ||
| usage = raw.get("usage", {}) | ||
| self._current.raw_usage = {**(self._current.raw_usage or {}), **usage} | ||
| # message_delta usage reports cumulative output tokens | ||
| if usage.get("output_tokens"): | ||
| self._current.output_tokens = usage["output_tokens"] | ||
|
|
@@ -452,6 +455,8 @@ def _emit_generation( | |
| properties["$ai_cache_creation_input_tokens"] = ( | ||
| gen.cache_creation_input_tokens | ||
| ) | ||
| if gen.raw_usage: | ||
| properties["$ai_usage"] = gen.raw_usage | ||
|
|
||
| if gen.stop_reason is not None: | ||
| properties["$ai_stop_reason"] = gen.stop_reason | ||
|
|
@@ -511,6 +516,8 @@ def _emit_generation_from_result( | |
| properties["$ai_cache_read_input_tokens"] = cache_read | ||
| if cache_creation: | ||
| properties["$ai_cache_creation_input_tokens"] = cache_creation | ||
| if usage: | ||
| properties["$ai_usage"] = usage | ||
|
Comment on lines
+519
to
+520
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This stores |
||
|
|
||
| if result.total_cost_usd is not None: | ||
| properties["$ai_total_cost_usd"] = result.total_cost_usd | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,6 +1,6 @@ | ||
| import time | ||
| import uuid | ||
| from typing import Any, Callable, Dict, List, Optional, cast | ||
| from typing import Any, Callable, Dict, List, Optional, Tuple, cast | ||
|
|
||
| from posthog import get_tags, identify_context, new_context, tag, contexts | ||
| from posthog.ai.gateway import warn_if_posthog_ai_gateway | ||
|
|
@@ -21,12 +21,36 @@ | |
| "$ai_output_tokens", | ||
| "$ai_cache_read_input_tokens", | ||
| "$ai_cache_creation_input_tokens", | ||
| "$ai_cache_creation_5m_input_tokens", | ||
| "$ai_cache_creation_1h_input_tokens", | ||
| "$ai_total_tokens", | ||
| "$ai_reasoning_tokens", | ||
| } | ||
| ) | ||
|
|
||
|
|
||
| def _extract_cache_creation_ttl_breakdown( | ||
| cache_creation: Any, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The parameter is named Maybe rename to |
||
| ) -> Optional[Tuple[int, int]]: | ||
| """Return a complete 5-minute/1-hour cache-write pair when usable.""" | ||
| if not isinstance(cache_creation, dict): | ||
| return None | ||
|
|
||
| values = ( | ||
| cache_creation.get("ephemeral_5m_input_tokens"), | ||
| cache_creation.get("ephemeral_1h_input_tokens"), | ||
| ) | ||
| if not any(value is not None for value in values) or not all( | ||
| isinstance(value, int) and not isinstance(value, bool) and value >= 0 | ||
| for value in values | ||
| if value is not None | ||
| ): | ||
| return None | ||
|
|
||
| breakdown = (values[0] or 0, values[1] or 0) | ||
| return breakdown if sum(breakdown) > 0 else None | ||
|
|
||
|
|
||
| def _get_tokens_source( | ||
| sdk_tags: Dict[str, Any], posthog_properties: Optional[Dict[str, Any]] | ||
| ) -> str: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -7,9 +7,9 @@ | |
| from posthog import identify_context, new_context | ||
|
|
||
| try: | ||
| from anthropic.types import Message, Usage | ||
| from anthropic.types import CacheCreation, Message, Usage | ||
|
|
||
| from posthog.ai.anthropic import Anthropic, AsyncAnthropic | ||
| from posthog.ai.anthropic import Anthropic, AnthropicBedrock, AsyncAnthropic | ||
| from posthog.test.ai.utils import RecordingAsyncStream | ||
|
|
||
| ANTHROPIC_AVAILABLE = True | ||
|
|
@@ -78,6 +78,14 @@ def create_mock_response(**kwargs): | |
| return MockResponse(**kwargs) | ||
|
|
||
|
|
||
| def assert_cache_creation_ttl_breakdown_preserved(props): | ||
| assert props["$ai_cache_creation_input_tokens"] == 800 | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Worth also asserting |
||
| assert props["$ai_usage"]["cache_creation"] == { | ||
| "ephemeral_5m_input_tokens": 300, | ||
| "ephemeral_1h_input_tokens": 500, | ||
| } | ||
|
|
||
|
|
||
| # Streaming mock helpers | ||
| class MockStreamEvent: | ||
| """Reusable mock event class for streaming responses.""" | ||
|
|
@@ -227,6 +235,46 @@ def mock_anthropic_response_with_cached_tokens(): | |
| ) | ||
|
|
||
|
|
||
| @pytest.fixture | ||
| def mock_anthropic_response_with_cache_ttl(): | ||
| return Message( | ||
| id="msg_cache_ttl", | ||
| type="message", | ||
| role="assistant", | ||
| content=[{"type": "text", "text": "Test response"}], | ||
| model="claude-sonnet-4-20250514", | ||
| usage=Usage( | ||
| input_tokens=20, | ||
| output_tokens=10, | ||
| cache_creation_input_tokens=800, | ||
| cache_creation=CacheCreation( | ||
| ephemeral_5m_input_tokens=300, | ||
| ephemeral_1h_input_tokens=500, | ||
| ), | ||
| ), | ||
| stop_reason="end_turn", | ||
| stop_sequence=None, | ||
| ) | ||
|
|
||
|
|
||
| def anthropic_stream_with_cache_ttl(): | ||
| message_start = MockStreamEvent("message_start") | ||
| message_start.message = MockStreamEvent( | ||
| usage=Usage( | ||
| input_tokens=20, | ||
| output_tokens=0, | ||
| cache_creation_input_tokens=800, | ||
| cache_creation=CacheCreation( | ||
| ephemeral_5m_input_tokens=300, | ||
| ephemeral_1h_input_tokens=500, | ||
| ), | ||
| ) | ||
| ) | ||
| message_delta = MockStreamEvent("message_delta") | ||
| message_delta.usage = MockUsage(output_tokens=10) | ||
| return [message_start, message_delta, MockStreamEvent("message_stop")] | ||
|
|
||
|
|
||
| @pytest.fixture | ||
| def mock_anthropic_response_with_tool_calls(): | ||
| return Message( | ||
|
|
@@ -593,11 +641,120 @@ def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens): | |
| assert props["$ai_output_tokens"] == 10 | ||
| assert props["$ai_cache_read_input_tokens"] == 15 | ||
| assert props["$ai_cache_creation_input_tokens"] == 2 | ||
| assert "$ai_cache_creation_5m_input_tokens" not in props | ||
| assert "$ai_cache_creation_1h_input_tokens" not in props | ||
| assert props["$ai_http_status"] == 200 | ||
| assert props["foo"] == "bar" | ||
| assert isinstance(props["$ai_latency"], float) | ||
|
|
||
|
|
||
| def test_preserves_cache_creation_ttl_breakdown_non_streaming( | ||
| mock_client, mock_anthropic_response_with_cache_ttl | ||
| ): | ||
| with patch( | ||
| "anthropic.resources.Messages.create", | ||
| return_value=mock_anthropic_response_with_cache_ttl, | ||
| ): | ||
| client = Anthropic(api_key="test-key", posthog_client=mock_client) | ||
| client.messages.create( | ||
| model="claude-sonnet-4-20250514", | ||
| messages=[{"role": "user", "content": "Hello"}], | ||
| ) | ||
|
|
||
| assert_cache_creation_ttl_breakdown_preserved( | ||
| mock_client.capture.call_args.kwargs["properties"] | ||
| ) | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_preserves_cache_creation_ttl_breakdown_non_streaming_async( | ||
| mock_client, mock_anthropic_response_with_cache_ttl | ||
| ): | ||
| async def mock_async_create(**kwargs): | ||
| return mock_anthropic_response_with_cache_ttl | ||
|
|
||
| with patch( | ||
| "anthropic.resources.messages.AsyncMessages.create", | ||
| side_effect=mock_async_create, | ||
| ): | ||
| client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) | ||
| await client.messages.create( | ||
| model="claude-sonnet-4-20250514", | ||
| messages=[{"role": "user", "content": "Hello"}], | ||
| ) | ||
|
|
||
| assert_cache_creation_ttl_breakdown_preserved( | ||
| mock_client.capture.call_args.kwargs["properties"] | ||
| ) | ||
|
|
||
|
|
||
| def test_preserves_cache_creation_ttl_breakdown_streaming(mock_client): | ||
| with patch( | ||
| "anthropic.resources.Messages.create", | ||
| return_value=iter(anthropic_stream_with_cache_ttl()), | ||
| ): | ||
| client = Anthropic(api_key="test-key", posthog_client=mock_client) | ||
| response = client.messages.create( | ||
| model="claude-sonnet-4-20250514", | ||
| messages=[{"role": "user", "content": "Hello"}], | ||
| stream=True, | ||
| ) | ||
| list(response) | ||
|
|
||
| assert_cache_creation_ttl_breakdown_preserved( | ||
| mock_client.capture.call_args.kwargs["properties"] | ||
| ) | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_preserves_cache_creation_ttl_breakdown_streaming_async(mock_client): | ||
| async def mock_async_stream(): | ||
| for event in anthropic_stream_with_cache_ttl(): | ||
| yield event | ||
|
|
||
| async def mock_async_create(**kwargs): | ||
| return mock_async_stream() | ||
|
|
||
| with patch( | ||
| "anthropic.resources.messages.AsyncMessages.create", | ||
| side_effect=mock_async_create, | ||
| ): | ||
| client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) | ||
| response = await client.messages.create( | ||
| model="claude-sonnet-4-20250514", | ||
| messages=[{"role": "user", "content": "Hello"}], | ||
| stream=True, | ||
| ) | ||
| [event async for event in response] | ||
|
|
||
| assert_cache_creation_ttl_breakdown_preserved( | ||
| mock_client.capture.call_args.kwargs["properties"] | ||
| ) | ||
|
|
||
|
|
||
| def test_anthropic_bedrock_preserves_cache_creation_ttl_breakdown( | ||
| mock_client, mock_anthropic_response_with_cache_ttl | ||
| ): | ||
| with patch( | ||
| "anthropic.resources.Messages.create", | ||
| return_value=mock_anthropic_response_with_cache_ttl, | ||
| ): | ||
| client = AnthropicBedrock( | ||
| posthog_client=mock_client, | ||
| aws_access_key="test-key", | ||
| aws_secret_key="test-secret", | ||
| aws_region="us-east-1", | ||
| ) | ||
| client.messages.create( | ||
| model="anthropic.claude-sonnet-4-20250514-v1:0", | ||
| messages=[{"role": "user", "content": "Hello"}], | ||
| ) | ||
|
|
||
| assert_cache_creation_ttl_breakdown_preserved( | ||
| mock_client.capture.call_args.kwargs["properties"] | ||
| ) | ||
|
|
||
|
|
||
| def test_tool_definition(mock_client, mock_anthropic_response): | ||
| with patch( | ||
| "anthropic.resources.Messages.create", | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This adds two public attributes to
ModelUsageplus new event properties, which reads as minor rather than patch to me. Deferring to Client Libraries on the convention here.