From b435331393628e0bda5ca5f61383f546f8762dd6 Mon Sep 17 00:00:00 2001 From: adriantqw <31589955+adriantqw@users.noreply.github.com> Date: Fri, 20 Jun 2025 11:09:36 +1000 Subject: [PATCH 1/2] Fix `_parse_usage` function for Langchain`CallbackHandler` For models with reasoning tokens the `_parse_usage` function does not successfully parse Langchain token usage metadata to PostHog. --- posthog/ai/langchain/callbacks.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 763d6880..c58725da 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -756,12 +756,13 @@ def _parse_usage(response: LLMResult) -> ModelUsage: break for generation_chunk in generation: - if generation_chunk.generation_info and ( - "usage_metadata" in generation_chunk.generation_info - ): - llm_usage = _parse_usage_model( - generation_chunk.generation_info["usage_metadata"] - ) + usage_metadata = ( + (generation_chunk.generation_info or {}).get("usage_metadata") or + (generation_chunk.message or {}).get("usage_metadata") or + {} + ) + if usage_metadata: + llm_usage = _parse_usage_model(usage_metadata) break message_chunk = getattr(generation_chunk, "message", {}) From 615068d53c7c3d286d68d9805b695fb4aeac4a4c Mon Sep 17 00:00:00 2001 From: adriantqw <31589955+adriantqw@users.noreply.github.com> Date: Fri, 20 Jun 2025 15:08:02 +1000 Subject: [PATCH 2/2] Fix issue --- posthog/ai/langchain/callbacks.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index c58725da..b7617790 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -756,14 +756,23 @@ def _parse_usage(response: LLMResult) -> ModelUsage: break for generation_chunk in generation: - usage_metadata = ( - (generation_chunk.generation_info or {}).get("usage_metadata") or - (generation_chunk.message or {}).get("usage_metadata") or - {} - ) - if usage_metadata: - llm_usage = _parse_usage_model(usage_metadata) - break + if generation_chunk.generation_info and ( + "usage_metadata" in generation_chunk.generation_info + ): + if generation_chunk.generation_info["usage_metadata"]: + llm_usage = _parse_usage_model( + generation_chunk.generation_info["usage_metadata"] + ) + break + + elif generation_chunk.message and ( + "usage_metadata" in generation_chunk.message + ): + if generation_chunk.message["usage_metadata"]: + llm_usage = _parse_usage_model( + generation_chunk.message["usage_metadata"] + ) + break message_chunk = getattr(generation_chunk, "message", {}) response_metadata = getattr(message_chunk, "response_metadata", {})