Add litellm_metadata support

This commit is contained in:
Sameer Kankute 2026-04-17 18:32:33 +05:30
parent ebdc34e226
commit b19951a0d4
No known key found for this signature in database
2 changed files with 48 additions and 14 deletions

View file

@ -31,21 +31,28 @@ def vertex_request_labels_from_litellm_params(
litellm_params: Optional[dict],
) -> Optional[Dict[str, str]]:
"""
Build Vertex/GCP billing labels from LiteLLM ``litellm_params["metadata"]``,
Build Vertex/GCP billing labels from LiteLLM user metadata on ``litellm_params``:
``metadata`` (``completion(..., metadata=...)``) or ``litellm_metadata``,
using ``requester_metadata`` string key-value pairs (same convention as Gemini).
``metadata`` is tried first when both are present.
"""
if not litellm_params or "metadata" not in litellm_params:
if not litellm_params:
return None
metadata = litellm_params["metadata"]
if metadata is None or not isinstance(metadata, dict):
return None
if "requester_metadata" not in metadata:
return None
rm = metadata["requester_metadata"]
if not isinstance(rm, dict):
return None
labels = {k: v for k, v in rm.items() if isinstance(v, str)}
return labels if labels else None
for key in ("metadata", "litellm_metadata"):
if key not in litellm_params:
continue
metadata = litellm_params[key]
if metadata is None or not isinstance(metadata, dict):
continue
if "requester_metadata" not in metadata:
continue
rm = metadata["requester_metadata"]
if not isinstance(rm, dict):
continue
labels = {k: v for k, v in rm.items() if isinstance(v, str)}
if labels:
return labels
return None
def pop_vertex_request_labels(
@ -54,8 +61,8 @@ def pop_vertex_request_labels(
) -> Optional[Dict[str, str]]:
"""
Resolve labels from optional ``labels`` (Gemini-style) and/or
``litellm_params["metadata"]["requester_metadata"]``. Pops ``labels`` from
optional_params when present.
``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``
(``requester_metadata``). Pops ``labels`` from optional_params when present.
"""
labels: Optional[Dict[str, str]] = None
if optional_params is not None and "labels" in optional_params:

View file

@ -1415,6 +1415,23 @@ def test_vertex_request_labels_from_litellm_params_extracts_requester_metadata()
assert vertex_request_labels_from_litellm_params(lp) == {"team": "analytics"}
def test_vertex_request_labels_from_litellm_params_accepts_litellm_metadata():
lp = {
"litellm_metadata": {
"requester_metadata": {"team": "platform", "count": 3}
}
}
assert vertex_request_labels_from_litellm_params(lp) == {"team": "platform"}
def test_vertex_request_labels_prefers_metadata_over_litellm_metadata():
lp = {
"metadata": {"requester_metadata": {"source": "metadata"}},
"litellm_metadata": {"requester_metadata": {"source": "litellm_metadata"}},
}
assert vertex_request_labels_from_litellm_params(lp) == {"source": "metadata"}
def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata():
optional = {"labels": {"env": "prod"}}
litellm_params = {"metadata": {"requester_metadata": {"team": "x"}}}
@ -1425,6 +1442,16 @@ def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata():
assert pop_vertex_request_labels(optional2, litellm_params) == {"team": "x"}
def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent():
optional: dict = {}
litellm_params = {
"litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}
}
assert pop_vertex_request_labels(optional, litellm_params) == {
"team": "from_litellm_meta"
}
def test_vertex_text_embedding_request_includes_labels_from_metadata():
import litellm