fix(skills): embed the request as the pre-call hooks returned it, not the original text

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

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()