feat(ai): capture the served service tier into model parameters

.sampo/changesets/capture-served-service-tier.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+OpenAI generations now record the service tier the provider served (`service_tier` inside `$ai_model_parameters`), on non-streaming, streaming, and LangChain capture paths. LLM analytics uses it to price flex and priority calls at their real rates instead of standard; a requested tier can be refused, so the value always comes from the response.

posthog/ai/langchain/callbacks.py

@@ -616,14 +616,21 @@ def _capture_generation(
         output: Union[LLMResult, BaseException],
         parent_run_id: Optional[UUID] = None,
     ):
+        # The served tier comes from the response, because a requested tier can be refused.
+        model_params = run.model_params
+        if isinstance(output, LLMResult) and isinstance(output.llm_output, dict):
+            served_tier = output.llm_output.get("service_tier")
+            if served_tier is not None:
+                model_params = {**(model_params or {}), "service_tier": served_tier}
+
         event_properties = {
             "$ai_trace_id": trace_id,
             "$ai_span_id": run_id,
             "$ai_span_name": run.name,
             "$ai_parent_id": parent_run_id,
             "$ai_provider": run.provider,
             "$ai_model": run.model,
-            "$ai_model_parameters": run.model_params,
+            "$ai_model_parameters": model_params,
             "$ai_input": with_privacy_mode(
                 self._ph_client,
                 self._privacy_mode,

posthog/ai/openai/_streaming.py

@@ -21,11 +21,14 @@ class _ResponsesStreamState:
     output: List[Any] = field(default_factory=list)
     model: Optional[str] = None
     stop_reason: Optional[str] = None
+    service_tier: Optional[str] = None
 
     def process_chunk(self, chunk: Any) -> None:
         response = getattr(chunk, "response", None)
         if response and self.model is None and hasattr(response, "model"):
             self.model = response.model
+        if response and self.service_tier is None:
+            self.service_tier = getattr(response, "service_tier", None)
 
         chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
         if chunk_usage:
@@ -50,10 +53,13 @@ class _ChatCompletionsStreamState:
     _tool_calls: Dict[int, Dict[str, Any]] = field(default_factory=dict)
     model: Optional[str] = None
     stop_reason: Optional[str] = None
+    service_tier: Optional[str] = None
 
     def process_chunk(self, chunk: Any) -> None:
         if self.model is None and hasattr(chunk, "model"):
             self.model = chunk.model
+        if self.service_tier is None:
+            self.service_tier = getattr(chunk, "service_tier", None)
 
         chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
         if chunk_usage:
@@ -93,6 +99,7 @@ def _build_streaming_event_data(
     groups: Optional[Dict[str, Any]],
     model_from_response: Optional[str],
     stop_reason: Optional[str],
+    service_tier: Optional[str] = None,
 ) -> StreamingEventData:
     """Build the fields shared by both OpenAI streaming endpoint events."""
 
@@ -111,4 +118,5 @@ def _build_streaming_event_data(
         privacy_mode=privacy_mode,
         groups=groups,
         stop_reason=stop_reason,
+        service_tier=service_tier,
     )

posthog/ai/openai/openai.py

@@ -214,6 +214,7 @@ def _capture_streaming_event(
             groups=posthog_groups,
             model_from_response=state.model,
             stop_reason=state.stop_reason,
+            service_tier=state.service_tier,
         )
         capture_streaming_event(self._client._ph_client, event_data)
 
@@ -433,6 +434,7 @@ def _capture_streaming_event(
             groups=posthog_groups,
             model_from_response=state.model,
             stop_reason=state.stop_reason,
+            service_tier=state.service_tier,
         )
         capture_streaming_event(self._client._ph_client, event_data)
 

posthog/ai/openai/openai_async.py

@@ -216,6 +216,7 @@ async def _capture_streaming_event(
             groups=posthog_groups,
             model_from_response=state.model,
             stop_reason=state.stop_reason,
+            service_tier=state.service_tier,
         )
         capture_streaming_event(self._client._ph_client, event_data)
 
@@ -438,6 +439,7 @@ async def _capture_streaming_event(
             groups=posthog_groups,
             model_from_response=state.model,
             stop_reason=state.stop_reason,
+            service_tier=state.service_tier,
         )
         capture_streaming_event(self._client._ph_client, event_data)
 

posthog/ai/types.py

@@ -142,3 +142,4 @@ class StreamingEventData(TypedDict):
     privacy_mode: bool
     groups: Optional[Dict[str, Any]]
     stop_reason: Optional[str]
+    service_tier: NotRequired[Optional[str]]

posthog/ai/utils.py

@@ -235,9 +235,12 @@ def merge_usage_stats(
         raise ValueError(f"Invalid mode: {mode}. Must be 'incremental' or 'cumulative'")
 
 
-def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
+def get_model_params(
+    kwargs: Dict[str, Any], served_service_tier: Optional[str] = None
+) -> Dict[str, Any]:
     """
-    Extracts model parameters from the kwargs dictionary.
+    Extracts model parameters from the kwargs dictionary. The service tier comes
+    from the response instead, because a requested tier can be refused.
     """
     model_params = {}
     for param in [
@@ -254,6 +257,8 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
     ]:
         if param in kwargs and kwargs[param] is not None:
             model_params[param] = kwargs[param]
+    if served_service_tier is not None:
+        model_params["service_tier"] = served_service_tier
     return model_params
 
 
@@ -476,7 +481,10 @@ def call_llm_and_track_usage(
 
             tag("$ai_provider", provider)
             tag("$ai_model", kwargs.get("model") or getattr(response, "model", None))
-            tag("$ai_model_parameters", get_model_params(kwargs))
+            tag(
+                "$ai_model_parameters",
+                get_model_params(kwargs, getattr(response, "service_tier", None)),
+            )
             tag(
                 "$ai_input",
                 with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages),
@@ -639,7 +647,10 @@ async def call_llm_and_track_usage_async(
 
             tag("$ai_provider", provider)
             tag("$ai_model", kwargs.get("model") or getattr(response, "model", None))
-            tag("$ai_model_parameters", get_model_params(kwargs))
+            tag(
+                "$ai_model_parameters",
+                get_model_params(kwargs, getattr(response, "service_tier", None)),
+            )
             tag(
                 "$ai_input",
                 with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages),
@@ -786,7 +797,9 @@ def capture_streaming_event(
     event_properties = {
         "$ai_provider": event_data["provider"],
         "$ai_model": event_data["model"],
-        "$ai_model_parameters": get_model_params(event_data["kwargs"]),
+        "$ai_model_parameters": get_model_params(
+            event_data["kwargs"], event_data.get("service_tier")
+        ),
         "$ai_input": with_privacy_mode(
             ph_client,
             event_data["privacy_mode"],

posthog/test/ai/langchain/test_callbacks.py

@@ -2839,3 +2839,28 @@ def test_ai_lane_client_routes_through_capture_ai(mock_client):
     events = [c[1]["event"] for c in mock_client.capture_ai.call_args_list]
     assert "$ai_generation" in events
     assert "$ai_trace" in events
+
+
+def test_served_service_tier_merges_into_model_parameters(mock_client):
+    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": "openai", "ls_model_name": "gpt-5-mini"},
+        invocation_params={"temperature": 0.5},
+    )
+    response = LLMResult(
+        generations=[[ChatGeneration(message=AIMessage(content="Response"))]],
+        llm_output={"service_tier": "flex"},
+    )
+
+    cb._pop_run_and_capture_generation(run_id, None, response)
+
+    props = mock_client.capture.call_args.kwargs["properties"]
+    assert props["$ai_model_parameters"]["service_tier"] == "flex"
+    assert props["$ai_model_parameters"]["temperature"] == 0.5

posthog/test/ai/openai/test_openai.py

@@ -2664,3 +2664,57 @@ async def test_async_provider_override_embeddings(mock_client, mock_embedding_re
     props = mock_client.capture.call_args[1]["properties"]
     assert props["$ai_provider"] == "perplexity"
     assert props["$ai_model"] == "text-embedding-3-small"
+
+
+def test_served_service_tier_lands_in_model_parameters(
+    mock_client, mock_openai_response
+):
+    mock_openai_response.service_tier = "flex"
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+        )
+
+        props = mock_client.capture.call_args[1]["properties"]
+        assert props["$ai_model_parameters"]["service_tier"] == "flex"
+
+
+def test_response_without_service_tier_omits_it(mock_client, mock_openai_response):
+    with patch(
+        "openai.resources.chat.completions.Completions.create",
+        return_value=mock_openai_response,
+    ):
+        client = OpenAI(api_key="test-key", posthog_client=mock_client)
+        client.chat.completions.create(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "Hello"}],
+            posthog_distinct_id="test-id",
+        )
+
+        props = mock_client.capture.call_args[1]["properties"]
+        assert "service_tier" not in props["$ai_model_parameters"]
+
+
+def test_streaming_state_tracks_served_service_tier():
+    from types import SimpleNamespace
+
+    from posthog.ai.openai._streaming import (
+        _ChatCompletionsStreamState,
+        _ResponsesStreamState,
+    )
+
+    chat_state = _ChatCompletionsStreamState()
+    chat_state.process_chunk(SimpleNamespace(service_tier="flex", choices=[]))
+    assert chat_state.service_tier == "flex"
+
+    responses_state = _ResponsesStreamState()
+    responses_state.process_chunk(
+        SimpleNamespace(response=SimpleNamespace(service_tier="flex"), type="other")
+    )
+    assert responses_state.service_tier == "flex"

references/public_api_snapshot.txt

@@ -487,6 +487,7 @@ attribute posthog.ai.types.StreamingEventData.model: str
 attribute posthog.ai.types.StreamingEventData.privacy_mode: bool
 attribute posthog.ai.types.StreamingEventData.properties: Optional[Dict[str, Any]]
 attribute posthog.ai.types.StreamingEventData.provider: str
+attribute posthog.ai.types.StreamingEventData.service_tier: NotRequired[Optional[str]]
 attribute posthog.ai.types.StreamingEventData.stop_reason: Optional[str]
 attribute posthog.ai.types.StreamingEventData.trace_id: Optional[str]
 attribute posthog.ai.types.StreamingEventData.usage_stats: TokenUsage
@@ -1074,7 +1075,7 @@ function posthog.ai.utils.extract_available_tool_calls(provider: str, kwargs: Di
 function posthog.ai.utils.extract_stop_reason(response: Any, provider: str) -> Optional[str]
 function posthog.ai.utils.finalize_ai_content(value: Any, ph_client: Any = None) -> Any
 function posthog.ai.utils.format_response(response, provider: str)
-function posthog.ai.utils.get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]
+function posthog.ai.utils.get_model_params(kwargs: Dict[str, Any], served_service_tier: Optional[str] = None) -> Dict[str, Any]
 function posthog.ai.utils.get_usage(response, provider: str) -> TokenUsage
 function posthog.ai.utils.merge_system_prompt(kwargs: Dict[str, Any], provider: str) -> List[FormattedMessage]
 function posthog.ai.utils.merge_usage_stats(target: TokenUsage, source: TokenUsage, mode: str = 'incremental') -> None