diff --git a/products/ai_observability/backend/api/personal_spend.py b/products/ai_observability/backend/api/personal_spend.py index 62d5f66c6ae1..5018fa35c79b 100644 --- a/products/ai_observability/backend/api/personal_spend.py +++ b/products/ai_observability/backend/api/personal_spend.py @@ -37,7 +37,7 @@ from posthog.schema import HogQLQueryModifiers from posthog.hogql import ast -from posthog.hogql.parser import parse_select +from posthog.hogql.parser import parse_expr, parse_select from posthog.hogql.query import execute_hogql_query from posthog.auth import ( @@ -72,12 +72,30 @@ MAX_TIME_BUCKETS = 600 _RELATIVE_DATE_RE = re.compile(r"^-?\d+[hdwmqyHDWMQY](Start|End)?$") +# `by_input_size` buckets on `$ai_input_tokens`, ascending by lower bound. This is the +# primary cost driver: input tokens are billed per token, so large-context turns dominate +# spend regardless of which tool they called. Single source of truth for both the query's +# `multiIf` and the row labels. +INPUT_SIZE_BUCKETS: list[tuple[str, int]] = [ + ("<50k", 0), + ("50k-100k", 50_000), + ("100k-150k", 100_000), + ("150k-200k", 150_000), + ("200k-300k", 200_000), + (">300k", 300_000), +] + MIN_LIMIT = 1 MAX_LIMIT = 200 DEFAULT_LIMIT = 50 CACHE_TIMEOUT_SECONDS = 300 -CACHE_VERSION = "v2" + +# Namespaces the cache by response schema. Bump on any payload shape change so entries +# written by an older process during a rolling deploy are never served against the new +# contract (this revision added `by_input_size`, `cost_attributed_usd`, `share_attributed` +# and removed `top_traces`). +CACHE_SCHEMA_VERSION = "v2" UTC = ZoneInfo("UTC") @@ -90,9 +108,10 @@ def _cache_key( email: str, date_from: str, date_to: str | None, product: str, limit: int, bucket_minutes: int | None ) -> str: to_slot = date_to or "_now" - # Suffix only when set, so bucketless requests keep their pre-bucket_minutes cache keys. + # Only append the bucket segment when set, so bucketless and bucketed requests get + # distinct keys without perturbing the bucketless key shape. bucket_slot = f":{bucket_minutes}" if bucket_minutes else "" - return f"personal_spend:{CACHE_VERSION}:{email}:{date_from}:{to_slot}:{product}:{limit}{bucket_slot}" + return f"personal_spend:{CACHE_SCHEMA_VERSION}:{email}:{date_from}:{to_slot}:{product}:{limit}{bucket_slot}" def _parse_date_param(value: str, field: str, now: datetime.datetime) -> datetime.datetime: @@ -184,8 +203,9 @@ class _SpendQueryParamsSerializer(serializers.Serializer): allow_blank=False, max_length=MAX_PRODUCT_KEY_LENGTH, help_text=( - "Required `ai_product` key to scope the tool / model / trace breakdowns to a single product. " - f"Only the following products are currently supported: {', '.join(sorted(SUPPORTED_PRODUCTS))}." + "Required `ai_product` key to scope the tool / model / input-size / day breakdowns to a " + f"single product. Only the following products are currently supported: " + f"{', '.join(sorted(SUPPORTED_PRODUCTS))}." ), ) limit = serializers.IntegerField( @@ -245,7 +265,10 @@ class _ToolBreakdownRowSerializer(serializers.Serializer): allow_null=True, help_text=( "Individual tool name from `$ai_tools_called` (split on `,` since multi-tool generations " - "store a comma-separated list). Null = pure text response with no tool call." + "store a comma-separated list). Null = pure text response with no tool call. Skill " + "invocations all collapse to the single `Skill` token, because the specific skill name " + "isn't recorded at ingestion — treat that row as an upper bound across every skill " + "combined, not any one skill." ), ) generation_count = serializers.IntegerField( @@ -255,20 +278,39 @@ class _ToolBreakdownRowSerializer(serializers.Serializer): help_text=( "Sum of `$ai_total_cost_usd` for generations whose tool list includes this tool. Multi-tool " "generations contribute their full cost to every tool they invoked, so this sum can exceed " - "`summary.scoped_cost_usd`. Prefer `share_of_scoped` for headline percentages — it's computed " - "per row and doesn't require the totals to reconcile." + "`summary.scoped_cost_usd` and rows can't be added together. It measures co-occurrence, not " + "cost caused by the tool. Use `cost_attributed_usd` for a number that reconciles." ), ) share_of_scoped = serializers.FloatField( help_text=( - "This tool's share of `summary.scoped_cost_usd`, expressed as a float in `[0, 1]`. Independent " - "per row, so co-occurring tools can each show a substantial share — the headline number to " - "present (e.g. `'Bash drove 47% of your spend'`)." + "This tool's share of `summary.scoped_cost_usd`, expressed as a float in `[0, 1]`. Computed " + "independently per row from `cost_usd`, so it inherits the same co-occurrence overstatement: " + "a tool that's simply present in most generations (e.g. Bash) shows a large share even when " + "it isn't what drove the cost, and shares across rows don't reconcile to 1. Use " + "`share_attributed` for an honest per-tool percentage." ), ) avg_input_tokens = serializers.FloatField( help_text="Average `$ai_input_tokens` across these generations — high values signal context bloat per call.", ) + cost_attributed_usd = serializers.FloatField( + help_text=( + "This tool's fractional share of generation cost: each generation's `$ai_total_cost_usd` is " + "divided evenly across the distinct tools it called, and a generation that called no tools " + "attributes its full cost to this same null `tool` row. Summed across all rows (when not " + "truncated), this reconciles to `summary.scoped_cost_usd` for the GENERATION cost only -- " + "embeddings carry no tools, so they're excluded here and the sum can land slightly under the " + "full scoped cost. Prefer this over `cost_usd` for an honest per-tool number." + ), + ) + share_attributed = serializers.FloatField( + help_text=( + "`cost_attributed_usd` normalized to `summary.scoped_cost_usd`, as a float in `[0, 1]`. Shares " + "across rows reconcile (summing to at most 1, short by the embedding-cost fraction which " + "carries no tools) — this is the honest headline percentage, unlike `share_of_scoped`." + ), + ) class _ModelBreakdownRowSerializer(serializers.Serializer): @@ -279,6 +321,28 @@ class _ModelBreakdownRowSerializer(serializers.Serializer): output_tokens = serializers.IntegerField(help_text="Sum of `$ai_output_tokens` for this model.") +class _InputSizeBreakdownRowSerializer(serializers.Serializer): + bucket = serializers.CharField( + help_text=( + "Label for this `$ai_input_tokens` bucket (e.g. `100k-150k`). Buckets: " + f"{', '.join(label for label, _ in INPUT_SIZE_BUCKETS)}." + ) + ) + min_input_tokens = serializers.IntegerField( + help_text="Inclusive lower bound of `$ai_input_tokens` for this bucket. Use to order or re-derive the bucket." + ) + generation_count = serializers.IntegerField( + help_text="Number of $ai_generation + $ai_embedding events with `$ai_input_tokens` in this bucket." + ) + cost_usd = serializers.FloatField(help_text="Sum of `$ai_total_cost_usd` for events in this bucket.") + share_of_scoped = serializers.FloatField( + help_text="This bucket's share of `summary.scoped_cost_usd`, as a float in `[0, 1]`. Reconciles across buckets." + ) + avg_cost_per_generation = serializers.FloatField( + help_text="Average `$ai_total_cost_usd` per event in this bucket -- `cost_usd / generation_count`." + ) + + class _DayBreakdownRowSerializer(serializers.Serializer): day = serializers.DateField(help_text="UTC calendar day the events fall on (`toDate(timestamp)`).") event_count = serializers.IntegerField( @@ -355,28 +419,16 @@ class _BucketBreakdownRowSerializer(serializers.Serializer): ) -class _TopTraceRowSerializer(serializers.Serializer): - trace_id = serializers.CharField( - allow_null=True, - help_text=( - "`$ai_trace_id` of the session — opaque string scoped to the originating product. Format is " - "not stable: most are UUIDs but some SDK wrappers emit JSON-shaped strings like " - '`{"device_id":"...","session_id":"..."}`. Callers should treat this as an opaque ' - "identifier (URL-encode before linking to a trace view)." - ), - ) - generation_count = serializers.IntegerField(help_text="Number of $ai_generation events in this trace.") - cost_usd = serializers.FloatField(help_text="Total cost in USD for this trace.") - started_at = serializers.DateTimeField(allow_null=True, help_text="Timestamp of the earliest event in this trace.") - - class _SummarySerializer(serializers.Serializer): date_from = serializers.DateTimeField( help_text="Inclusive UTC start of the spend window resolved from the request." ) date_to = serializers.DateTimeField(help_text="Exclusive UTC end of the spend window resolved from the request.") product = serializers.CharField( - help_text="The `ai_product` filter applied to tool / model / trace breakdowns — echoes the request `product`.", + help_text=( + "The `ai_product` filter applied to tool / model / input-size / day breakdowns — echoes the " + "request `product`." + ), ) total_cost_usd = serializers.FloatField( help_text="Total LLM cost in USD across every `ai_product` for the user — independent of the `product` filter." @@ -422,6 +474,24 @@ class _ModelBreakdownSerializer(serializers.Serializer): ) +class _InputSizeBreakdownSerializer(serializers.Serializer): + items = _InputSizeBreakdownRowSerializer( + many=True, + help_text=( + "One row per `$ai_input_tokens` bucket that has events, ordered by `min_input_tokens` " + "ascending. This is the primary cost driver: large-context turns dominate spend regardless " + "of which tool they called, so check this before attributing spend to a tool. Buckets with " + "no events are omitted." + ), + ) + truncated = serializers.BooleanField( + help_text=( + "Always false: there are only six buckets and `by_input_size` ignores `limit` because " + "truncating by cost would drop whole buckets rather than individual rows." + ) + ) + + class _DayBreakdownSerializer(serializers.Serializer): items = _DayBreakdownRowSerializer( many=True, @@ -459,13 +529,6 @@ class _BucketBreakdownSerializer(serializers.Serializer): ) -class _TopTracesSerializer(serializers.Serializer): - items = _TopTraceRowSerializer(many=True, help_text="Rows of top traces by cost, ordered by cost descending.") - truncated = serializers.BooleanField( - help_text="True when more rows exist beyond the requested `limit`. Re-request with a larger `limit` to retrieve them." - ) - - class PersonalSpendAnalysisResponseSerializer(serializers.Serializer): """Structured personal LLM spend analysis for the requesting user.""" @@ -475,6 +538,13 @@ class PersonalSpendAnalysisResponseSerializer(serializers.Serializer): ) by_tool = _ToolBreakdownSerializer(help_text="Spend grouped by tool. Scoped to `product` when set.") by_model = _ModelBreakdownSerializer(help_text="Spend grouped by `$ai_model`. Scoped to `product` when set.") + by_input_size = _InputSizeBreakdownSerializer( + help_text=( + "Spend grouped by `$ai_input_tokens` bucket. Scoped to `product`. Not subject to `limit`. This " + "is the primary cost driver: large-context turns dominate spend regardless of which tool they " + "called, so check this before attributing spend to a tool." + ) + ) by_day = _DayBreakdownSerializer( help_text="Spend grouped by UTC day, ordered ascending. Scoped to `product`. Not subject to `limit`." ) @@ -492,14 +562,6 @@ class PersonalSpendAnalysisResponseSerializer(serializers.Serializer): "Scoped to `product`. Only present when the request set `bucket_minutes`." ), ) - top_traces = _TopTracesSerializer( - help_text=( - "Deprecated — always returns `{items: [], truncated: false}`. Trace IDs are opaque strings " - "that aren't actionable in the UI. Kept in the response shape so existing consumers don't " - "crash; remove your rendering of this field and we'll drop it from the response entirely " - "in a follow-up." - ) - ) def _email_filter(email: str) -> ast.Expr: @@ -557,6 +619,11 @@ def _truncate(rows: list[dict[str, Any]], limit: int) -> dict[str, Any]: return {"items": rows[:limit], "truncated": len(rows) > limit} +def _share_of(value: float, scoped: float) -> float: + """Fraction of scoped cost in `[0, 1]`, guarding the empty-window (`scoped == 0`) case.""" + return (value / scoped) if scoped > 0 else 0.0 + + def _fetch_summary( team: Team, email: str, from_dt: datetime.datetime, to_dt: datetime.datetime, product: str ) -> dict[str, Any]: @@ -632,6 +699,60 @@ def _fetch_by_product( return _truncate(rows, limit) +def _fetch_tool_cost_attribution( + team: Team, + email: str, + from_dt: datetime.datetime, + to_dt: datetime.datetime, + product: str, +) -> dict[str | None, float]: + """Splits each generation's cost evenly across the DISTINCT tools it called, + so the sum across tools reconciles to the scoped generation cost instead of + overcounting co-occurring tools. `$ai_tools_called` can list the same tool more + than once (tools are recorded in call order, e.g. "Bash,Read,Bash"); `arrayDistinct` + collapses those repeats before dividing, since a tool called twice in one + generation still only "caused" that generation's cost once. Generations with no + tools attribute their full cost to the same NULL bucket the unattributed + `_fetch_by_tool` rows use. Not truncated by `limit` -- callers merge by `tool` + into an already-truncated row set, so untruncated tools here simply have no + matching row to attach to. + """ + query = parse_select( + """ + SELECT + tool, + round(sum(cost_share), 6) AS cost_attributed_usd + FROM ( + SELECT + nullIf(arrayJoin(if(empty(tools), [''], tools)), '') AS tool, + cost / length(if(empty(tools), [''], tools)) AS cost_share + FROM ( + SELECT + arrayDistinct(arrayFilter(x -> x != '', splitByChar(',', coalesce(properties.$ai_tools_called, '')))) AS tools, + toFloat(properties.$ai_total_cost_usd) AS cost + FROM events + WHERE equals(event, '$ai_generation') + AND {product_filter} + AND {email_filter} + AND {timestamp_filter} + ) + ) + GROUP BY tool + """ + ) + result = execute_hogql_query( + query=query, + placeholders={ + "product_filter": _product_filter(product), + "email_filter": _email_filter(email), + "timestamp_filter": _timestamp_filter(from_dt, to_dt), + }, + team=team, + query_type="PersonalSpendToolCostAttribution", + ) + return {(row[0] if row[0] is not None else None): float(row[1] or 0.0) for row in (result.results or [])} + + def _fetch_by_tool( team: Team, email: str, @@ -681,7 +802,16 @@ def _fetch_by_tool( } for row in (result.results or []) ] - return _truncate(rows, limit) + truncated = _truncate(rows, limit) + + # Merged by `tool` rather than requeried per-row: the attribution query isn't + # itself truncated, so a tool absent from the (possibly truncated) row set above + # just has no row to attach `cost_attributed_usd` to. + attribution = _fetch_tool_cost_attribution(team, email, from_dt, to_dt, product) + for row in truncated["items"]: + row["cost_attributed_usd"] = attribution.get(row["tool"], 0.0) + + return truncated def _fetch_by_model( @@ -735,6 +865,87 @@ def _fetch_by_model( return _truncate(rows, limit) +def _input_size_bucket_case_sql() -> str: + """Builds the `multiIf(...)` branches from `INPUT_SIZE_BUCKETS`: each bucket (other + than the last) becomes `$ai_input_tokens < , '