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. diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 91fdece52..bab85cecd 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -616,6 +616,13 @@ 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, @@ -623,7 +630,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..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,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"], 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" 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