diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index 8869d2cf2db..b8d3595163e 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -47,6 +47,16 @@ class _EmbeddingData(BaseModel): data: tuple[_EmbeddingItem, ...] +class _EmbeddingRequest(BaseModel): + """The /embeddings-shaped request as the pre-call hooks (rate limits, budgets, guardrails) hand it back.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + model: str + input: tuple[str, ...] + metadata: dict[str, object] # mutable-ok: the router mutates the metadata dict it is handed + + def cosine_similarity(left: Vector, right: Vector) -> float: dot: Final = sum(a * b for a, b in zip(left, right, strict=True)) norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right)) @@ -68,15 +78,21 @@ def router_embedder( """Embeds through the router after the same key rate-limit, budget and guardrail pre-call hooks /embeddings runs.""" async def embed(texts: Sequence[str]) -> Sequence[Vector]: - batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input - metadata: Final = embedding_spend_metadata(user_api_key_dict) request: Final = { # mutable-ok: pre_call_hook mutates the request dict in place "model": embedding_model, - "input": batch, - "metadata": metadata, + "input": list(texts), # mutable-ok: Router.aembedding accepts only str | list input + "metadata": embedding_spend_metadata(user_api_key_dict), } - await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=request, call_type="aembedding") - response: Final = await router.aembedding(model=embedding_model, input=batch, metadata=metadata) + processed: Final = _EmbeddingRequest.model_validate( + await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=request, call_type="aembedding" + ) + ) + response: Final = await router.aembedding( + model=processed.model, + input=list(processed.input), # mutable-ok: Router.aembedding accepts only str | list input + metadata=processed.metadata, + ) return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data) return embed diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py index 4bec0baa1d9..3765c4a0813 100644 --- a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py +++ b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py @@ -329,6 +329,38 @@ class TestSearchSkills: assert checked["data"]["input"] == router.aembedding.await_args.kwargs["input"] assert checked["data"]["metadata"]["user_api_key"] == "hashed-caller-key" + @pytest.mark.asyncio + async def test_the_embedding_model_sees_the_request_as_the_guardrails_rewrote_it(self) -> None: + router = MagicMock() + router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": [1.0, 0.0, 0.0]} for i in range(len(input))], + ) + ) + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type: { + **data, + "input": ["[MASKED]" for _ in data["input"]], + "metadata": {**data["metadata"], "guardrail": "masked"}, + } + ) + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + sent = tuple(call.kwargs for call in router.aembedding.await_args_list) + assert sent + assert all(set(call["input"]) == {"[MASKED]"} for call in sent) + assert all(call["metadata"]["guardrail"] == "masked" for call in sent) + @pytest.mark.asyncio async def test_a_key_over_its_limit_never_reaches_the_embedding_model(self) -> None: router = _embedding_router()