mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(mavvrik_focus): also derive total_tokens in FOCUS Tags
total_tokens has no stored column in LiteLLM_DailyUserSpend at all, so it can't be selected like prompt_tokens/completion_tokens. Derive it as their sum in _with_token_tags, only when both source counts are present for a row, and add tests covering the sum and the partial-data case.
This commit is contained in:
parent
d4f059c9b5
commit
24dc17f1fa
2 changed files with 24 additions and 5 deletions
|
|
@ -41,21 +41,23 @@ else:
|
|||
# drops prompt_tokens/completion_tokens even though the source query selects
|
||||
# them. Mavvrik carries them through as extra keys in the existing Tags JSON
|
||||
# column (the spec's own escape hatch for non-standard fields), rather than
|
||||
# changing the shared transformer used by every FOCUS destination.
|
||||
# changing the shared transformer used by every FOCUS destination. total_tokens
|
||||
# isn't a stored column at all -- it's derived here as their sum.
|
||||
_TOKEN_TAG_KEYS = ("prompt_tokens", "completion_tokens")
|
||||
|
||||
|
||||
def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFrame:
|
||||
"""Merge prompt/completion token counts from the pre-transform frame into
|
||||
``normalized``'s Tags column. Rows correspond 1:1 and in the same order
|
||||
across both frames -- transform() only adds/renames columns, it never
|
||||
filters or reorders rows.
|
||||
"""Merge prompt/completion token counts (and their sum, total_tokens) from
|
||||
the pre-transform frame into ``normalized``'s Tags column. Rows correspond
|
||||
1:1 and in the same order across both frames -- transform() only
|
||||
adds/renames columns, it never filters or reorders rows.
|
||||
"""
|
||||
available = [k for k in _TOKEN_TAG_KEYS if k in data.columns]
|
||||
if not available or len(data) != len(normalized):
|
||||
return normalized
|
||||
|
||||
token_rows = data.select(available).to_dicts()
|
||||
has_both = "prompt_tokens" in available and "completion_tokens" in available
|
||||
|
||||
def _merge(tags_json: str, row: dict) -> str:
|
||||
tags = json.loads(tags_json) if tags_json else {}
|
||||
|
|
@ -63,6 +65,11 @@ def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFra
|
|||
value = row.get(key)
|
||||
if value is not None:
|
||||
tags[key] = str(value)
|
||||
if has_both:
|
||||
prompt = row.get("prompt_tokens")
|
||||
completion = row.get("completion_tokens")
|
||||
if prompt is not None and completion is not None:
|
||||
tags["total_tokens"] = str(prompt + completion)
|
||||
return json.dumps(tags)
|
||||
|
||||
merged_tags = pl.Series(
|
||||
|
|
|
|||
|
|
@ -84,9 +84,21 @@ def test_with_token_tags_merges_prompt_and_completion_tokens() -> None:
|
|||
"model": "azure/gpt-4o-mini",
|
||||
"prompt_tokens": "57",
|
||||
"completion_tokens": "753",
|
||||
"total_tokens": "810",
|
||||
}
|
||||
|
||||
|
||||
def test_with_token_tags_omits_total_when_only_one_token_column_present() -> None:
|
||||
data = pl.DataFrame({"prompt_tokens": [57]})
|
||||
normalized = pl.DataFrame({"Tags": [json.dumps({"model": "azure/gpt-4o-mini"})]})
|
||||
|
||||
result = _with_token_tags(data, normalized)
|
||||
|
||||
tags = json.loads(result["Tags"][0])
|
||||
assert tags == {"model": "azure/gpt-4o-mini", "prompt_tokens": "57"}
|
||||
assert "total_tokens" not in tags
|
||||
|
||||
|
||||
def test_with_token_tags_noop_when_token_columns_absent() -> None:
|
||||
data = pl.DataFrame({"model": ["azure/gpt-4o-mini"]})
|
||||
normalized = pl.DataFrame({"Tags": [json.dumps({"model": "azure/gpt-4o-mini"})]})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue