mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Add litellm_metadata support
This commit is contained in:
parent
ebdc34e226
commit
b19951a0d4
2 changed files with 48 additions and 14 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue