From f371dc69974db3945d3b4abaf41544401c93c9a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Gouv=C3=AAa?= <64869705+gouveags@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:50:18 -0300 Subject: [PATCH] fix(ai): preserve Anthropic cache-write TTL breakdowns --- .../anthropic-cache-ttl-breakdown.md | 5 + posthog/ai/claude_agent_sdk/processor.py | 7 + posthog/ai/langchain/callbacks.py | 31 +++- posthog/ai/utils.py | 26 ++- posthog/test/ai/anthropic/test_anthropic.py | 161 +++++++++++++++- .../ai/claude_agent_sdk/test_processor.py | 92 ++++++++- posthog/test/ai/langchain/test_callbacks.py | 174 ++++++++++++++++++ posthog/test/ai/test_tokens_source.py | 12 ++ references/public_api_snapshot.txt | 4 +- 9 files changed, 497 insertions(+), 15 deletions(-) create mode 100644 .sampo/changesets/anthropic-cache-ttl-breakdown.md diff --git a/.sampo/changesets/anthropic-cache-ttl-breakdown.md b/.sampo/changesets/anthropic-cache-ttl-breakdown.md new file mode 100644 index 00000000..764248e3 --- /dev/null +++ b/.sampo/changesets/anthropic-cache-ttl-breakdown.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Preserve Anthropic cache-write TTL breakdowns across Python SDK AI integrations. diff --git a/posthog/ai/claude_agent_sdk/processor.py b/posthog/ai/claude_agent_sdk/processor.py index 09d8187e..3ee26071 100644 --- a/posthog/ai/claude_agent_sdk/processor.py +++ b/posthog/ai/claude_agent_sdk/processor.py @@ -39,6 +39,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())) @@ -76,9 +77,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"] @@ -450,6 +453,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 @@ -509,6 +514,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 if result.total_cost_usd is not None: properties["$ai_total_cost_usd"] = result.total_cost_usd diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 7045aa70..822d0490 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -44,7 +44,11 @@ from posthog import setup from posthog.ai.gateway import warn_if_posthog_ai_gateway from posthog.ai.sanitization import sanitize_langchain -from posthog.ai.utils import get_model_params, with_privacy_mode +from posthog.ai.utils import ( + _extract_cache_creation_ttl_breakdown, + get_model_params, + with_privacy_mode, +) from posthog.client import Client log = logging.getLogger("posthog") @@ -657,6 +661,16 @@ def _capture_generation( event_properties["$ai_cache_creation_input_tokens"] = ( usage.cache_write_tokens ) + if ( + usage.cache_write_5m_tokens is not None + and usage.cache_write_1h_tokens is not None + ): + event_properties["$ai_cache_creation_5m_input_tokens"] = ( + usage.cache_write_5m_tokens + ) + event_properties["$ai_cache_creation_1h_input_tokens"] = ( + usage.cache_write_1h_tokens + ) event_properties["$ai_cache_read_input_tokens"] = usage.cache_read_tokens event_properties["$ai_reasoning_tokens"] = usage.reasoning_tokens @@ -771,6 +785,8 @@ class ModelUsage: cache_write_tokens: Optional[int] cache_read_tokens: Optional[int] reasoning_tokens: Optional[int] + cache_write_5m_tokens: Optional[int] = None + cache_write_1h_tokens: Optional[int] = None def _parse_usage_model( @@ -823,8 +839,15 @@ def _parse_usage_model( if "input_token_details" in usage and isinstance( usage["input_token_details"], dict ): - parsed_usage["cache_write"] = usage["input_token_details"].get("cache_creation") - parsed_usage["cache_read"] = usage["input_token_details"].get("cache_read") + input_token_details = usage["input_token_details"] + parsed_usage["cache_write"] = input_token_details.get("cache_creation") + cache_write_ttl = _extract_cache_creation_ttl_breakdown(input_token_details) + if cache_write_ttl is not None: + cache_write_5m, cache_write_1h = cache_write_ttl + parsed_usage["cache_write_5m"] = cache_write_5m + parsed_usage["cache_write_1h"] = cache_write_1h + parsed_usage["cache_write"] = cache_write_5m + cache_write_1h + parsed_usage["cache_read"] = input_token_details.get("cache_read") # Reasoning (OpenAI & langchain 0.3.9+) if "output_token_details" in usage and isinstance( @@ -844,6 +867,8 @@ def _parse_usage_model( dataclass_key: parsed_usage.get(mapped_key) or 0 for mapped_key, dataclass_key in field_mapping.items() }, + cache_write_5m_tokens=parsed_usage.get("cache_write_5m"), + cache_write_1h_tokens=parsed_usage.get("cache_write_1h"), ) # For Anthropic providers, LangChain reports input_tokens as the sum of all input tokens. # Our cost calculation expects them to be separate for Anthropic, so we subtract cache tokens. diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index 127f5d11..fde8f890 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -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 @@ -20,12 +20,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, +) -> 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: diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index a2dd451b..7ebcd5db 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -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 + 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", diff --git a/posthog/test/ai/claude_agent_sdk/test_processor.py b/posthog/test/ai/claude_agent_sdk/test_processor.py index 74bfc1e0..acc003c6 100644 --- a/posthog/test/ai/claude_agent_sdk/test_processor.py +++ b/posthog/test/ai/claude_agent_sdk/test_processor.py @@ -101,7 +101,17 @@ def _make_result_message( is_error: bool = False, cache_read: int = 0, cache_creation: int = 0, + cache_creation_ttl: Optional[Dict[str, int]] = None, ) -> ResultMessage: + usage: Dict[str, Any] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cache_read_input_tokens": cache_read, + "cache_creation_input_tokens": cache_creation, + } + if cache_creation_ttl is not None: + usage["cache_creation"] = cache_creation_ttl + return ResultMessage( subtype="success" if not is_error else "error", duration_ms=duration_ms, @@ -110,12 +120,7 @@ def _make_result_message( num_turns=num_turns, session_id="sess_123", total_cost_usd=total_cost_usd, - usage={ - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "cache_read_input_tokens": cache_read, - "cache_creation_input_tokens": cache_creation, - }, + usage=usage, ) @@ -212,6 +217,44 @@ async def test_emits_generation_from_stream_events(self, processor, mock_client) assert props["$ai_output_tokens"] == 50 assert props["$ai_cache_read_input_tokens"] == 20 + @pytest.mark.asyncio + async def test_emits_cache_creation_ttl_breakdown_from_stream_events( + self, processor, mock_client + ): + message_start = _make_message_start(input_tokens=100, cache_creation=800) + message_start.event["message"]["usage"]["cache_creation"] = { + "ephemeral_5m_input_tokens": 300, + "ephemeral_1h_input_tokens": 500, + } + messages = [ + message_start, + _make_message_delta(output_tokens=50), + _make_message_stop(), + _make_assistant_message(), + _make_result_message(), + ] + + with patch( + "posthog.ai.claude_agent_sdk.processor.original_query", + side_effect=lambda **kw: _fake_query(messages), + ): + async for _ in processor.query(prompt="Hi", options=ClaudeAgentOptions()): + pass + + generation_call = next( + call + for call in mock_client.capture.call_args_list + if (call.kwargs.get("event") or call[1].get("event")) == "$ai_generation" + ) + props = generation_call.kwargs.get("properties") or generation_call[1].get( + "properties" + ) + assert props["$ai_cache_creation_input_tokens"] == 800 + assert props["$ai_usage"]["cache_creation"] == { + "ephemeral_5m_input_tokens": 300, + "ephemeral_1h_input_tokens": 500, + } + @pytest.mark.asyncio async def test_emits_multiple_generations_for_multi_turn( self, processor, mock_client @@ -251,7 +294,15 @@ async def test_fallback_generation_from_result_when_no_stream_events( """When StreamEvents are not available, fall back to ResultMessage data.""" messages = [ _make_assistant_message(), - _make_result_message(input_tokens=200, output_tokens=100), + _make_result_message( + input_tokens=200, + output_tokens=100, + cache_creation=800, + cache_creation_ttl={ + "ephemeral_5m_input_tokens": 300, + "ephemeral_1h_input_tokens": 500, + }, + ), ] with patch( @@ -272,6 +323,11 @@ async def test_fallback_generation_from_result_when_no_stream_events( ) assert props["$ai_input_tokens"] == 200 assert props["$ai_output_tokens"] == 100 + assert props["$ai_cache_creation_input_tokens"] == 800 + assert props["$ai_usage"]["cache_creation"] == { + "ephemeral_5m_input_tokens": 300, + "ephemeral_1h_input_tokens": 500, + } class TestToolSpanEmission: @@ -379,8 +435,13 @@ async def test_privacy_mode_redacts_tool_input(self, mock_client): distinct_id="user", privacy_mode=True, ) + message_start = _make_message_start(cache_creation=800) + message_start.event["message"]["usage"]["cache_creation"] = { + "ephemeral_5m_input_tokens": 300, + "ephemeral_1h_input_tokens": 500, + } messages = [ - _make_message_start(), + message_start, _make_message_stop(), _make_assistant_message( tool_uses=[ @@ -414,6 +475,21 @@ async def test_privacy_mode_redacts_tool_input(self, mock_client): ) assert "$ai_input_state" not in props + generation_call = next( + call + for call in mock_client.capture.call_args_list + if (call.kwargs.get("event") or call[1].get("event")) == "$ai_generation" + ) + generation_props = generation_call.kwargs.get("properties") or generation_call[ + 1 + ].get("properties") + assert generation_props["$ai_input"] is None + assert generation_props.get("$ai_output_choices") is None + assert generation_props["$ai_usage"]["cache_creation"] == { + "ephemeral_5m_input_tokens": 300, + "ephemeral_1h_input_tokens": 500, + } + class TestPersonlessMode: @pytest.mark.asyncio diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 8cda547c..71a723a3 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -307,6 +307,51 @@ async def test_async_basic_chat_chain(mock_client, stream): assert "$ai_parent_id" not in trace_props +def _capture_anthropic_usage( + mock_client, + input_token_details, + *, + provider="anthropic", + model="claude-sonnet", + input_tokens=1000, + cache_read_tokens=100, +): + from langchain_core.messages import AIMessage + from langchain_core.outputs import ChatGeneration, LLMResult + + cb = CallbackHandler(mock_client) + run_id = uuid.uuid4() + cb._set_llm_metadata( + serialized={}, + run_id=run_id, + messages=[{"role": "user", "content": "test"}], + metadata={"ls_provider": provider, "ls_model_name": model}, + ) + response = LLMResult( + generations=[ + [ + ChatGeneration( + message=AIMessage(content="Response"), + generation_info={ + "usage_metadata": { + "input_tokens": input_tokens, + "output_tokens": 50, + "input_token_details": { + **input_token_details, + "cache_read": cache_read_tokens, + }, + } + }, + ) + ] + ], + llm_output={}, + ) + + cb._pop_run_and_capture_generation(run_id, None, response) + return mock_client.capture.call_args.kwargs["properties"] + + @pytest.mark.parametrize( "Model,stream", [ @@ -1683,6 +1728,135 @@ def test_anthropic_provider_subtracts_cache_write_tokens(mock_client): assert generation_args["properties"]["$ai_cache_creation_input_tokens"] == 800 +@pytest.mark.parametrize( + "cache_creation_details,provider,model,expected_5m_tokens,expected_1h_tokens", + [ + ( + { + "cache_creation": 0, + "ephemeral_5m_input_tokens": 300, + }, + "bedrock_converse", + "us.anthropic.claude-sonnet-4-20250514-v1:0", + 300, + 0, + ), + ( + { + "cache_creation": 0, + "ephemeral_1h_input_tokens": 500, + }, + "anthropic", + "claude-sonnet", + 0, + 500, + ), + ( + { + "cache_creation": 0, + "ephemeral_5m_input_tokens": 300, + "ephemeral_1h_input_tokens": 500, + }, + "anthropic", + "claude-sonnet", + 300, + 500, + ), + ], + ids=["5m-only", "1h-only", "mixed"], +) +def test_anthropic_cache_creation_ttl_breakdown( + mock_client, + cache_creation_details, + provider, + model, + expected_5m_tokens, + expected_1h_tokens, +): + """TTL counts are emitted separately without double counting cache writes.""" + cache_read_tokens = 200 + input_tokens = 2000 + expected_cache_write_tokens = expected_5m_tokens + expected_1h_tokens + generation_props = _capture_anthropic_usage( + mock_client, + cache_creation_details, + provider=provider, + model=model, + input_tokens=input_tokens, + cache_read_tokens=cache_read_tokens, + ) + assert ( + generation_props["$ai_cache_creation_input_tokens"] + == expected_cache_write_tokens + ) + assert generation_props["$ai_cache_creation_5m_input_tokens"] == expected_5m_tokens + assert generation_props["$ai_cache_creation_1h_input_tokens"] == expected_1h_tokens + assert generation_props["$ai_input_tokens"] == ( + input_tokens - cache_read_tokens - expected_cache_write_tokens + ) + + +def test_anthropic_aggregate_cache_creation_remains_supported(mock_client): + """Legacy aggregate-only usage keeps its existing capture semantics.""" + generation_props = _capture_anthropic_usage( + mock_client, + {"cache_creation": 500}, + ) + assert generation_props["$ai_cache_creation_input_tokens"] == 500 + assert generation_props["$ai_cache_read_input_tokens"] == 100 + assert generation_props["$ai_input_tokens"] == 400 + assert "$ai_cache_creation_5m_input_tokens" not in generation_props + assert "$ai_cache_creation_1h_input_tokens" not in generation_props + + +@pytest.mark.parametrize( + "ttl_details", + [ + { + "ephemeral_5m_input_tokens": None, + "ephemeral_1h_input_tokens": {"unexpected": 1}, + }, + { + "ephemeral_5m_input_tokens": True, + "ephemeral_1h_input_tokens": 0, + }, + { + "ephemeral_5m_input_tokens": -1, + "ephemeral_1h_input_tokens": 0, + }, + ], + ids=["none-and-malformed", "boolean", "negative"], +) +def test_anthropic_malformed_cache_creation_ttl_breakdown_uses_aggregate( + mock_client, ttl_details +): + """Unusable TTL details do not replace valid legacy aggregate usage.""" + generation_props = _capture_anthropic_usage( + mock_client, + {"cache_creation": 500, **ttl_details}, + ) + assert generation_props["$ai_cache_creation_input_tokens"] == 500 + assert generation_props["$ai_input_tokens"] == 400 + assert "$ai_cache_creation_5m_input_tokens" not in generation_props + assert "$ai_cache_creation_1h_input_tokens" not in generation_props + + +def test_anthropic_zero_cache_creation_ttl_breakdown_uses_aggregate(mock_client): + """An empty TTL pair does not replace a usable legacy aggregate count.""" + generation_props = _capture_anthropic_usage( + mock_client, + { + "cache_creation": 500, + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 0, + }, + ) + assert generation_props["$ai_cache_creation_input_tokens"] == 500 + assert generation_props["$ai_input_tokens"] == 400 + assert "$ai_cache_creation_5m_input_tokens" not in generation_props + assert "$ai_cache_creation_1h_input_tokens" not in generation_props + + def test_anthropic_provider_subtracts_both_cache_read_and_write_tokens(mock_client): """Test that Anthropic provider correctly subtracts both cache read and write tokens.""" from langchain_core.outputs import LLMResult, ChatGeneration diff --git a/posthog/test/ai/test_tokens_source.py b/posthog/test/ai/test_tokens_source.py index c9adcf95..1f4fe446 100644 --- a/posthog/test/ai/test_tokens_source.py +++ b/posthog/test/ai/test_tokens_source.py @@ -43,6 +43,18 @@ {"$ai_cache_creation_input_tokens": 200}, "passthrough", ), + ( + "override_cache_creation_5m", + {"$ai_input_tokens": 100}, + {"$ai_cache_creation_5m_input_tokens": 200}, + "passthrough", + ), + ( + "override_cache_creation_1h", + {"$ai_input_tokens": 100}, + {"$ai_cache_creation_1h_input_tokens": 200}, + "passthrough", + ), ( "override_reasoning_tokens", {"$ai_input_tokens": 100}, diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index f7468453..2de15101 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -364,6 +364,8 @@ attribute posthog.ai.langchain.callbacks.GenerationMetadata.posthog_properties: attribute posthog.ai.langchain.callbacks.GenerationMetadata.provider: Optional[str] = None attribute posthog.ai.langchain.callbacks.GenerationMetadata.tools: Optional[List[Dict[str, Any]]] = None attribute posthog.ai.langchain.callbacks.ModelUsage.cache_read_tokens: Optional[int] +attribute posthog.ai.langchain.callbacks.ModelUsage.cache_write_1h_tokens: Optional[int] = None +attribute posthog.ai.langchain.callbacks.ModelUsage.cache_write_5m_tokens: Optional[int] = None attribute posthog.ai.langchain.callbacks.ModelUsage.cache_write_tokens: Optional[int] attribute posthog.ai.langchain.callbacks.ModelUsage.input_tokens: Optional[int] attribute posthog.ai.langchain.callbacks.ModelUsage.output_tokens: Optional[int] @@ -821,7 +823,7 @@ class posthog.ai.gemini.gemini_converter.GeminiMessage class posthog.ai.gemini.gemini_converter.GeminiPart class posthog.ai.langchain.callbacks.CallbackHandler(client: Optional[Client] = None, *, distinct_id: Optional[Union[str, int, UUID]] = None, trace_id: Optional[Union[str, int, float, UUID]] = None, properties: Optional[Dict[str, Any]] = None, privacy_mode: bool = False, groups: Optional[Dict[str, Any]] = None) class posthog.ai.langchain.callbacks.GenerationMetadata(name: str, start_time: float, end_time: Optional[float], input: Optional[Any], provider: Optional[str] = None, model: Optional[str] = None, model_params: Optional[Dict[str, Any]] = None, base_url: Optional[str] = None, tools: Optional[List[Dict[str, Any]]] = None, posthog_properties: Optional[Dict[str, Any]] = None) -class posthog.ai.langchain.callbacks.ModelUsage(input_tokens: Optional[int], output_tokens: Optional[int], cache_write_tokens: Optional[int], cache_read_tokens: Optional[int], reasoning_tokens: Optional[int]) +class posthog.ai.langchain.callbacks.ModelUsage(input_tokens: Optional[int], output_tokens: Optional[int], cache_write_tokens: Optional[int], cache_read_tokens: Optional[int], reasoning_tokens: Optional[int], cache_write_5m_tokens: Optional[int] = None, cache_write_1h_tokens: Optional[int] = None) class posthog.ai.langchain.callbacks.SpanMetadata(name: str, start_time: float, end_time: Optional[float], input: Optional[Any]) class posthog.ai.openai.openai.OpenAI(posthog_client: Optional[PostHogClient] = None, **kwargs) class posthog.ai.openai.openai.WrappedBeta