feat(ph-ai): ai_billable metadata per-generation

Author: tatoalo

CHANGELOG.md

@@ -1,3 +1,7 @@
+# 6.7.15 - 2025-11-04
+
+- feat(ph-ai): ai_billing metadata per-generation
+
 # 6.7.14 - 2025-11-03
 
 - fix(django): Handle request.user access in async middleware context to prevent SynchronousOnlyOperation errors in Django 5+ (fixes #355)

posthog/ai/langchain/callbacks.py

@@ -79,6 +79,8 @@ class GenerationMetadata(SpanMetadata):
     """Base URL of the provider's API used in the run."""
     tools: Optional[List[Dict[str, Any]]] = None
     """Tools provided to the model."""
+    metadata: Optional[Dict[str, Any]] = None
+    """Metadata from the invocation."""
 
 
 RunMetadata = Union[SpanMetadata, GenerationMetadata]
@@ -416,6 +418,7 @@ def _set_llm_metadata(
             if tools := invocation_params.get("tools"):
                 generation.tools = tools
         if isinstance(metadata, dict):
+            generation.metadata = metadata
             if model := metadata.get("ls_model_name"):
                 generation.model = model
             if provider := metadata.get("ls_provider"):
@@ -603,6 +606,9 @@ def _capture_generation(
         if self._properties:
             event_properties.update(self._properties)
 
+        if run.metadata and "$ai_billable" in run.metadata:
+            event_properties["$ai_billable"] = run.metadata["$ai_billable"]
+
         if self._distinct_id is None:
             event_properties["$process_person_profile"] = False
 

posthog/test/ai/langchain/test_callbacks.py

@@ -113,6 +113,7 @@ def test_metadata_capture(mock_client):
         base_url="https://us.posthog.com",
         name="test",
         end_time=None,
+        metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
     )
     assert callbacks._runs[run_id] == expected
     with patch("time.time", return_value=1234567891):
@@ -1269,6 +1270,7 @@ def test_metadata_tools(mock_client):
         name="test",
         tools=tools,
         end_time=None,
+        metadata={"ls_model_name": "hog-mini", "ls_provider": "posthog"},
     )
     assert callbacks._runs[run_id] == expected
     with patch("time.time", return_value=1234567891):
@@ -1839,6 +1841,7 @@ def test_tool_definition(mock_client):
         name="test",
         tools=tools,
         end_time=None,
+        metadata={"ls_model_name": "gpt-4o-mini", "ls_provider": "openai"},
     )
     assert callbacks._runs[run_id] == expected
 
@@ -2126,3 +2129,88 @@ def test_agent_action_and_finish_imports():
     assert mock_client.capture.call_count == 1
     call_args = mock_client.capture.call_args[1]
     assert call_args["event"] == "$ai_span"
+
+
+def test_ai_billable_metadata(mock_client):
+    """Test that $ai_billable metadata is properly captured in generation events."""
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("user", "What is the weather?"),
+        ]
+    )
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="It's sunny!",
+                usage_metadata={
+                    "input_tokens": 10,
+                    "output_tokens": 5,
+                    "total_tokens": 15,
+                },
+            )
+        ]
+    )
+
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model
+
+    # Invoke with $ai_billable metadata set to False
+    result = chain.invoke(
+        {}, config={"callbacks": callbacks, "metadata": {"$ai_billable": False}}
+    )
+
+    assert result.content == "It's sunny!"
+    assert mock_client.capture.call_count == 3
+
+    generation_args = mock_client.capture.call_args_list[1][1]
+    generation_props = generation_args["properties"]
+
+    assert generation_args["event"] == "$ai_generation"
+    assert generation_props["$ai_billable"] is False
+
+    # Test with $ai_billable set to True
+    mock_client.reset_mock()
+
+    result = chain.invoke(
+        {}, config={"callbacks": callbacks, "metadata": {"$ai_billable": True}}
+    )
+
+    assert mock_client.capture.call_count == 3
+    generation_args = mock_client.capture.call_args_list[1][1]
+    generation_props = generation_args["properties"]
+
+    assert generation_args["event"] == "$ai_generation"
+    assert generation_props["$ai_billable"] is True
+
+
+def test_ai_billable_metadata_not_present_when_not_set(mock_client):
+    """Test that $ai_billable is not in event properties when not set in metadata."""
+    prompt = ChatPromptTemplate.from_messages([("user", "Hello")])
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="Hi!",
+                usage_metadata={
+                    "input_tokens": 5,
+                    "output_tokens": 2,
+                    "total_tokens": 7,
+                },
+            )
+        ]
+    )
+
+    callbacks = [CallbackHandler(mock_client)]
+    chain = prompt | model
+
+    # Invoke without $ai_billable metadata
+    result = chain.invoke({}, config={"callbacks": callbacks})
+
+    assert result.content == "Hi!"
+    assert mock_client.capture.call_count == 3
+
+    generation_args = mock_client.capture.call_args_list[1][1]
+    generation_props = generation_args["properties"]
+
+    assert generation_args["event"] == "$ai_generation"
+    # $ai_billable should not be present at all
+    assert "$ai_billable" not in generation_props

posthog/version.py

@@ -1,4 +1,4 @@
-VERSION = "6.7.14"
+VERSION = "6.7.15"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201