diff --git a/posthog_llma/event_builder.py b/posthog_llma/event_builder.py index 358b9ca..f4faa35 100644 --- a/posthog_llma/event_builder.py +++ b/posthog_llma/event_builder.py @@ -167,7 +167,13 @@ def build_events(parsed: dict, config: dict) -> list[dict]: if not privacy_mode: max_attr = config.get("max_attribute_length", 12000) content_blocks = [] - if gen["output_text"]: + # Prefer the structured blocks so extended-thinking content stays + # typed as "thinking" rather than being flattened into assistant + # text. Older parsed payloads have no output_blocks, so fall back + # to the flattened output_text. + if gen.get("output_blocks"): + content_blocks.extend(gen["output_blocks"]) + elif gen["output_text"]: content_blocks.append({"type": "text", "text": gen["output_text"]}) content_blocks.extend(_truncate_tool_blocks(gen.get("tool_use_blocks", []), max_attr)) if content_blocks: diff --git a/posthog_llma/events.py b/posthog_llma/events.py index c5a594d..77bbc17 100644 --- a/posthog_llma/events.py +++ b/posthog_llma/events.py @@ -69,8 +69,8 @@ def build_ai_generation( "$ai_framework": "claude-code", "$ai_project_name": project_name, "$ai_agent_name": agent_name, - "cache_read_input_tokens": cache_read_tokens, - "cache_creation_input_tokens": cache_creation_tokens, + "$ai_cache_read_input_tokens": cache_read_tokens, + "$ai_cache_creation_input_tokens": cache_creation_tokens, } if user_prompt and not privacy_mode: diff --git a/posthog_llma/parser.py b/posthog_llma/parser.py index 72e4bc7..72022d4 100644 --- a/posthog_llma/parser.py +++ b/posthog_llma/parser.py @@ -332,6 +332,10 @@ def _finalize_generation(state: dict) -> tuple[dict, list]: type_order = state.get("type_order") or ["thinking", "text", "tool_use"] text_parts = [] + # Structured counterpart to text_parts: keeps thinking and text distinct so + # downstream consumers can render/filter them separately. `output_text` + # remains the flattened join for backward compatibility. + output_blocks = [] entry_tool_uses = [] for block_type in type_order: if block_type == "thinking": @@ -339,11 +343,13 @@ def _finalize_generation(state: dict) -> tuple[dict, list]: t = item.get("thinking", "") if t: text_parts.append(t) + output_blocks.append({"type": "thinking", "thinking": t}) elif block_type == "text": for item in blocks.get("text", []): t = item.get("text", "") if t: text_parts.append(t) + output_blocks.append({"type": "text", "text": t}) elif block_type == "tool_use": for item in blocks.get("tool_use", []): entry_tool_uses.append({ @@ -374,6 +380,7 @@ def _finalize_generation(state: dict) -> tuple[dict, list]: "span_id": span_id, "msg_id": state.get("msg_id", ""), "output_text": output_text, + "output_blocks": output_blocks, "tool_use_blocks": tool_use_blocks, "is_error": stop_reason == "error", "error_message": state.get("error_message"), diff --git a/tests/test_posthog_llma.py b/tests/test_posthog_llma.py index 037a365..665288f 100644 --- a/tests/test_posthog_llma.py +++ b/tests/test_posthog_llma.py @@ -83,6 +83,43 @@ def test_timestamp_passthrough(self): ) assert ev["timestamp"] == "2026-04-12T21:00:00Z" + def test_cache_tokens_use_ai_namespace(self): + """Cache token properties must carry the $ai_ prefix. + + PostHog's LLM Analytics pipeline reads $ai_cache_read_input_tokens and + $ai_cache_creation_input_tokens when computing cost. Unprefixed keys are + ingested as ordinary custom properties and ignored by cost calculation, + which silently understates spend on prompt-cached workloads. + """ + ev = build_ai_generation( + model="claude-opus-4-6", + input_tokens=4, + output_tokens=926, + cache_read_tokens=150109, + cache_creation_tokens=75729, + trace_id="t", session_id="s", + ) + props = ev["properties"] + assert props["$ai_cache_read_input_tokens"] == 150109 + assert props["$ai_cache_creation_input_tokens"] == 75729 + # The unprefixed keys must not linger, or the same numbers get ingested + # twice under two different names. + assert "cache_read_input_tokens" not in props + assert "cache_creation_input_tokens" not in props + + def test_cache_tokens_emitted_even_when_zero(self): + """Anthropic generations always carry cache fields, matching posthog-python. + + See posthog/ai/utils.py: for the anthropic provider the official SDK + always includes both cache fields even at 0, rather than omitting them. + """ + ev = build_ai_generation( + model="m", trace_id="t", session_id="s", + ) + props = ev["properties"] + assert props["$ai_cache_read_input_tokens"] == 0 + assert props["$ai_cache_creation_input_tokens"] == 0 + def test_no_timestamp_means_no_key(self): ev = build_ai_generation( model="m", trace_id="t", session_id="s", @@ -96,8 +133,8 @@ def test_cache_tokens(self): cache_read_tokens=5, cache_creation_tokens=3, ) props = ev["properties"] - assert props["cache_read_input_tokens"] == 5 - assert props["cache_creation_input_tokens"] == 3 + assert props["$ai_cache_read_input_tokens"] == 5 + assert props["$ai_cache_creation_input_tokens"] == 3 def test_no_cost_properties(self): """Cost is calculated by PostHog ingestion, we should not send it.""" diff --git a/tests/test_session_parser.py b/tests/test_session_parser.py index e27ebc5..ad138e2 100644 --- a/tests/test_session_parser.py +++ b/tests/test_session_parser.py @@ -314,12 +314,14 @@ def test_thinking_preserved_when_split_across_chunks(self): assert len(parsed["tool_uses"]) == 1 assert parsed["tool_uses"][0]["name"] == "Bash" - # Downstream events include thinking in output_choices + # Downstream events include thinking in output_choices, typed as a + # thinking block rather than flattened into assistant text. events = build_events(parsed, DEFAULT_CONFIG) gen_ev = next(e for e in events if e["event"] == "$ai_generation") content_blocks = gen_ev["properties"]["$ai_output_choices"][0]["content"] - text_blocks = [b for b in content_blocks if b.get("type") == "text"] - assert text_blocks and "deep thoughts" in text_blocks[0]["text"] + thinking_blocks = [b for b in content_blocks if b.get("type") == "thinking"] + assert thinking_blocks and thinking_blocks[0]["thinking"] == "deep thoughts" + assert not [b for b in content_blocks if b.get("type") == "text"] finally: os.unlink(path) @@ -514,6 +516,65 @@ def test_merge_preserves_first_seen_block_order(self): # Text came first in the source, so output_text must lead # with "preface" rather than reshuffling thinking ahead of it. assert parsed["generations"][0]["output_text"] == "preface\nafterthought" + # The same order is preserved in the structured blocks, with each + # block keeping its own type rather than collapsing into text. + blocks = parsed["generations"][0]["output_blocks"] + assert blocks == [ + {"type": "text", "text": "preface"}, + {"type": "thinking", "thinking": "afterthought"}, + ] + finally: + os.unlink(path) + + def test_thinking_and_text_stay_separate_blocks(self): + """Extended thinking must not be ingested as assistant text. + + _finalize_generation joins thinking and text into a single output_text + for backward compatibility, but $ai_output_choices needs them typed + separately or thinking content renders as ordinary assistant output and + cannot be filtered on. + """ + entries = [ + {"type": "permission-mode", "permissionMode": "default", "sessionId": "s1"}, + { + "type": "user", "uuid": "u-1", "parentUuid": None, + "promptId": "p-1", "isMeta": False, + "message": {"role": "user", "content": "go"}, + "timestamp": "2026-04-12T10:00:00.000Z", + "sessionId": "s1", "version": "2.1.0", "cwd": "/tmp", + }, + { + "type": "assistant", "uuid": "a-1", "parentUuid": "u-1", + "message": { + "role": "assistant", "id": "msg-1", + "model": "claude-opus-4-6", "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 10}, + "content": [ + {"type": "thinking", "thinking": "weighing the options"}, + {"type": "text", "text": "Here is the answer."}, + ], + }, + "timestamp": "2026-04-12T10:00:01.000Z", + "sessionId": "s1", "version": "2.1.0", "cwd": "/tmp", + }, + ] + path = _write_jsonl(entries) + try: + parsed = parse_session(path, DEFAULT_CONFIG) + gen = parsed["generations"][0] + # Flattened form is unchanged. + assert gen["output_text"] == "weighing the options\nHere is the answer." + + events = build_events(parsed, DEFAULT_CONFIG) + gen_ev = next(e for e in events if e["event"] == "$ai_generation") + content = gen_ev["properties"]["$ai_output_choices"][0]["content"] + + thinking = [b for b in content if b.get("type") == "thinking"] + text = [b for b in content if b.get("type") == "text"] + assert len(thinking) == 1 and thinking[0]["thinking"] == "weighing the options" + assert len(text) == 1 and text[0]["text"] == "Here is the answer." + # And the reasoning must not leak into the visible answer text. + assert "weighing the options" not in text[0]["text"] finally: os.unlink(path)