diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index b8d3595163e..d3fe68e65f7 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -68,7 +68,7 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj return { # mutable-ok: the router mutates the metadata dict it is handed **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict), - "user_api_key": user_api_key_dict.api_key, + "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), } diff --git a/tests/unit/llms/litellm_proxy/skills/test_skill_search.py b/tests/unit/llms/litellm_proxy/skills/test_skill_search.py index 3f1fe0d5d68..e28e927cf61 100644 --- a/tests/unit/llms/litellm_proxy/skills/test_skill_search.py +++ b/tests/unit/llms/litellm_proxy/skills/test_skill_search.py @@ -304,6 +304,29 @@ class TestSearchSkills: assert metadata["user_api_key_team_id"] == "team-1" assert metadata["user_api_key_user_id"] == "user-1" + @pytest.mark.asyncio + async def test_cli_session_embedding_spend_is_attributed_to_the_per_user_alias_not_the_login_token(self) -> None: + router = _embedding_router() + session = UserAPIKeyAuth( + api_key="cli-session-Qm7xJ2kP9sLw4vT1nR8yAa", + user_id="alice", + key_alias="cli-session-alice", + is_session_token=True, + ) + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=session, + proxy_logging_obj=_pass_through_key_limits(), + ) + metadata = router.aembedding.await_args.kwargs["metadata"] + assert metadata["user_api_key"] == "cli-session-alice" + assert metadata["user_api_key_hash"] == "cli-session-alice" + @pytest.mark.asyncio async def test_key_limits_are_checked_against_the_real_embedding_call_before_it_runs(self) -> None: router = _embedding_router()