mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(mavvrik_focus): also carry cache token counts in FOCUS Tags
cache_creation_input_tokens and cache_read_input_tokens are selected by the same database.py query as prompt_tokens/completion_tokens and dropped by the same transformer. Add them to _TOKEN_TAG_KEYS.
This commit is contained in:
parent
03bc06079e
commit
4b3905cb8b
2 changed files with 46 additions and 10 deletions
|
|
@ -38,19 +38,26 @@ else:
|
|||
AsyncIOScheduler = Any
|
||||
|
||||
# FOCUS v1.2 has no standard column for token counts; core's transformer
|
||||
# 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. total_tokens
|
||||
# isn't a stored column at all -- it's derived here as their sum.
|
||||
_TOKEN_TAG_KEYS = ("prompt_tokens", "completion_tokens")
|
||||
# drops prompt_tokens/completion_tokens/cache_creation_input_tokens/
|
||||
# cache_read_input_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. total_tokens isn't a
|
||||
# stored column at all -- it's derived here as the sum of prompt and
|
||||
# completion tokens.
|
||||
_TOKEN_TAG_KEYS = (
|
||||
"prompt_tokens",
|
||||
"completion_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
"cache_read_input_tokens",
|
||||
)
|
||||
|
||||
|
||||
def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFrame:
|
||||
"""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.
|
||||
"""Merge 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) or "Tags" not in normalized.columns:
|
||||
|
|
@ -77,6 +84,11 @@ def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFra
|
|||
tags["total_tokens"] = str(prompt + completion)
|
||||
return json.dumps(tags)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Mavvrik FOCUS export: merging token tags for %d row(s) (keys=%s)",
|
||||
len(token_rows),
|
||||
available,
|
||||
)
|
||||
merged_tags = pl.Series(
|
||||
[_merge(tags_json, row) for tags_json, row in zip(normalized["Tags"].to_list(), token_rows)]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -88,6 +88,30 @@ def test_with_token_tags_merges_prompt_and_completion_tokens() -> None:
|
|||
}
|
||||
|
||||
|
||||
def test_with_token_tags_merges_cache_token_columns() -> None:
|
||||
data = pl.DataFrame(
|
||||
{
|
||||
"prompt_tokens": [57],
|
||||
"completion_tokens": [753],
|
||||
"cache_creation_input_tokens": [10],
|
||||
"cache_read_input_tokens": [5],
|
||||
}
|
||||
)
|
||||
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",
|
||||
"completion_tokens": "753",
|
||||
"cache_creation_input_tokens": "10",
|
||||
"cache_read_input_tokens": "5",
|
||||
"total_tokens": "810",
|
||||
}
|
||||
|
||||
|
||||
def test_with_token_tags_recovers_from_malformed_tags_json() -> None:
|
||||
data = pl.DataFrame({"prompt_tokens": [57], "completion_tokens": [753]})
|
||||
normalized = pl.DataFrame({"Tags": ["not-valid-json"]})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue