From fb6898e63d6a652eb9906b53731985eef1c990a2 Mon Sep 17 00:00:00 2001 From: Bernat Torres Date: Fri, 4 Sep 2026 14:32:26 +0200 Subject: [PATCH 1/4] feat(ai): capture the served service tier into model parameters OpenAI responses carry the tier the call was actually served on (flex, priority, default). Recording it lets cost processing price tiered calls at their real rates; the JS SDK already does this. Co-Authored-By: Claude Fable 5 --- posthog/ai/langchain/callbacks.py | 10 +++- posthog/ai/openai/_streaming.py | 8 +++ posthog/ai/openai/openai.py | 2 + posthog/ai/openai/openai_async.py | 2 + posthog/ai/types.py | 1 + posthog/ai/utils.py | 29 +++++++++-- posthog/test/ai/langchain/test_callbacks.py | 25 ++++++++++ posthog/test/ai/openai/test_openai.py | 54 +++++++++++++++++++++ 8 files changed, 127 insertions(+), 4 deletions(-) diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 91fdece52..610c9ab69 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -616,6 +616,14 @@ def _capture_generation( output: Union[LLMResult, BaseException], parent_run_id: Optional[UUID] = None, ): + model_params = run.model_params + if isinstance(output, LLMResult) and isinstance(output.llm_output, dict): + # The tier the provider served, which langchain lifts out of the response; + # a requested tier can be refused. + 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, @@ -623,7 +631,7 @@ def _capture_generation( "$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, diff --git a/posthog/ai/openai/_streaming.py b/posthog/ai/openai/_streaming.py index 3f2f38e74..5c857b940 100644 --- a/posthog/ai/openai/_streaming.py +++ b/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, ) diff --git a/posthog/ai/openai/openai.py b/posthog/ai/openai/openai.py index 9a046831c..c1950939c 100644 --- a/posthog/ai/openai/openai.py +++ b/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) diff --git a/posthog/ai/openai/openai_async.py b/posthog/ai/openai/openai_async.py index e7a955cc1..dac3eeb4a 100644 --- a/posthog/ai/openai/openai_async.py +++ b/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) diff --git a/posthog/ai/types.py b/posthog/ai/types.py index 4c9a3e871..36424b133 100644 --- a/posthog/ai/types.py +++ b/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]] diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index 895b34ab0..3a54ebd56 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -257,6 +257,17 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]: return model_params +def with_served_service_tier( + model_params: Dict[str, Any], service_tier: Optional[str] +) -> Dict[str, Any]: + """ + Merges the service tier the provider served (a requested tier can be refused). + """ + if service_tier is not None: + model_params["service_tier"] = service_tier + return model_params + + def get_usage(response, provider: str) -> TokenUsage: """ Extract usage statistics from response based on provider. @@ -476,7 +487,12 @@ 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", + with_served_service_tier( + get_model_params(kwargs), getattr(response, "service_tier", None) + ), + ) tag( "$ai_input", with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages), @@ -639,7 +655,12 @@ 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", + with_served_service_tier( + get_model_params(kwargs), getattr(response, "service_tier", None) + ), + ) tag( "$ai_input", with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages), @@ -786,7 +807,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": with_served_service_tier( + get_model_params(event_data["kwargs"]), event_data.get("service_tier") + ), "$ai_input": with_privacy_mode( ph_client, event_data["privacy_mode"], diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 11f021099..fd2ed8dc6 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/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 diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 1444eeb12..a8aa9c616 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/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" From 49f548b2c8b29c079b1af8c4a4ad357078bc2e5f Mon Sep 17 00:00:00 2001 From: Bernat Torres Date: Fri, 4 Sep 2026 14:32:49 +0200 Subject: [PATCH 2/4] chore: add changeset Co-Authored-By: Claude Fable 5 --- .sampo/changesets/capture-served-service-tier.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .sampo/changesets/capture-served-service-tier.md diff --git a/.sampo/changesets/capture-served-service-tier.md b/.sampo/changesets/capture-served-service-tier.md new file mode 100644 index 000000000..2d3d66f5b --- /dev/null +++ b/.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. From 6a108e1a6d42eca07e9aed9b190d297c01a07367 Mon Sep 17 00:00:00 2001 From: Bernat Torres Date: Fri, 4 Sep 2026 16:20:14 +0200 Subject: [PATCH 3/4] refactor(ai): fold the served tier into get_model_params Co-Authored-By: Claude Fable 5 --- posthog/ai/langchain/callbacks.py | 3 +-- posthog/ai/utils.py | 32 +++++++++++-------------------- 2 files changed, 12 insertions(+), 23 deletions(-) diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 610c9ab69..bab85cecd 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -616,10 +616,9 @@ 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): - # The tier the provider served, which langchain lifts out of the response; - # a requested tier can be refused. served_tier = output.llm_output.get("service_tier") if served_tier is not None: model_params = {**(model_params or {}), "service_tier": served_tier} diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index 3a54ebd56..6c5d3ebba 100644 --- a/posthog/ai/utils.py +++ b/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,17 +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] - return model_params - - -def with_served_service_tier( - model_params: Dict[str, Any], service_tier: Optional[str] -) -> Dict[str, Any]: - """ - Merges the service tier the provider served (a requested tier can be refused). - """ - if service_tier is not None: - model_params["service_tier"] = service_tier + if served_service_tier is not None: + model_params["service_tier"] = served_service_tier return model_params @@ -489,9 +483,7 @@ def call_llm_and_track_usage( tag("$ai_model", kwargs.get("model") or getattr(response, "model", None)) tag( "$ai_model_parameters", - with_served_service_tier( - get_model_params(kwargs), getattr(response, "service_tier", None) - ), + get_model_params(kwargs, getattr(response, "service_tier", None)), ) tag( "$ai_input", @@ -657,9 +649,7 @@ async def call_llm_and_track_usage_async( tag("$ai_model", kwargs.get("model") or getattr(response, "model", None)) tag( "$ai_model_parameters", - with_served_service_tier( - get_model_params(kwargs), getattr(response, "service_tier", None) - ), + get_model_params(kwargs, getattr(response, "service_tier", None)), ) tag( "$ai_input", @@ -807,8 +797,8 @@ def capture_streaming_event( event_properties = { "$ai_provider": event_data["provider"], "$ai_model": event_data["model"], - "$ai_model_parameters": with_served_service_tier( - get_model_params(event_data["kwargs"]), event_data.get("service_tier") + "$ai_model_parameters": get_model_params( + event_data["kwargs"], event_data.get("service_tier") ), "$ai_input": with_privacy_mode( ph_client, From c062d7cfcf52329e296bffba64f70d962f8a35ec Mon Sep 17 00:00:00 2001 From: Bernat Torres Date: Fri, 4 Sep 2026 16:33:58 +0200 Subject: [PATCH 4/4] chore: refresh the public API snapshot Co-Authored-By: Claude Fable 5 --- references/public_api_snapshot.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 01e2ff669..06d7394ee 100644 --- a/references/public_api_snapshot.txt +++ b/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