From e9cc9c9bc3b06f652f35bf14e04305316a6035bd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:11:00 -0700 Subject: [PATCH] fix(a2a): attribute agent search embedding spend to the calling key --- .../_experimental/mcp_server/tool_search.py | 1 + litellm/proxy/agent_endpoints/agent_search.py | 19 ++++++++-- litellm/proxy/agent_endpoints/endpoints.py | 7 +++- .../mcp_server/test_mcp_tool_search.py | 1 + .../agent_endpoints/test_agent_search.py | 37 +++++++++++++++++-- 5 files changed, 56 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index a5259e127fa..f79765f6d01 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -142,6 +142,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI router=llm_router, embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, + user_api_key_dict=user_api_key_dict, ) match outcome: case AgentSearchHits(hits): diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py index a8fc57eb229..896a22fbf9d 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -16,6 +16,7 @@ from litellm.exceptions import BudgetExceededError from litellm.types.agents import AgentResponse if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth from litellm.router import Router DEFAULT_AGENT_SEARCH_TOP_K: Final = 5 @@ -122,10 +123,21 @@ def cosine_similarity(left: Vector, right: Vector) -> float: return dot / norms if norms else 0.0 -def router_embedder(router: Router, embedding_model: str) -> Embedder: +def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]: + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + 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, + } + + +def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder: async def embed(texts: Sequence[str]) -> Sequence[Vector]: batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input - response: Final = await router.aembedding(model=embedding_model, input=batch) + response: Final = await router.aembedding( + model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict) + ) return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data) return embed @@ -174,6 +186,7 @@ async def search_agents( router: Router | None, embedding_model: str | None, index: AgentSearchIndex, + user_api_key_dict: UserAPIKeyAuth, ) -> AgentSearchOutcome: if embedding_model is None: return AgentSearchNotConfigured( @@ -181,4 +194,4 @@ async def search_agents( ) if router is None: return AgentSearchNotConfigured(reason="agent search needs a model_list so the embedding model can be called") - return await index.search(query, agents, top_k, router_embedder(router, embedding_model)) + return await index.search(query, agents, top_k, router_embedder(router, embedding_model, user_api_key_dict)) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index cfb0597cc2a..b6c41a17503 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -231,7 +231,9 @@ def _agent_search_error(status_code: int, error: str, message: str) -> HTTPExcep return HTTPException(status_code=status_code, detail=detail) -async def _rank_agents_by_query(query: str, agents: Sequence[AgentResponse], top_k: int) -> tuple[AgentResponse, ...]: +async def _rank_agents_by_query( + query: str, agents: Sequence[AgentResponse], top_k: int, user_api_key_dict: UserAPIKeyAuth +) -> tuple[AgentResponse, ...]: from litellm.proxy.proxy_server import llm_router outcome: Final = await search_agents( @@ -241,6 +243,7 @@ async def _rank_agents_by_query(query: str, agents: Sequence[AgentResponse], top router=llm_router, embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, + user_api_key_dict=user_api_key_dict, ) match outcome: case AgentSearchHits(hits): @@ -376,7 +379,7 @@ async def get_agents( if query is None: return returned_agents - return await _rank_agents_by_query(query, returned_agents, top_k) + return await _rank_agents_by_query(query, returned_agents, top_k, user_api_key_dict) except HTTPException: raise except Exception as e: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index bf0bd3c795e..16221f44efe 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -566,6 +566,7 @@ class TestCallToolRestApiVirtualTools: result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict) assert result.isError is False + assert mock_search.await_args.kwargs["user_api_key_dict"] is user_api_key_dict assert json.loads(result.content[0].text) == [ { "agent_id": "translator", diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py index f6313eade9e..4b674eb142b 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py @@ -24,6 +24,8 @@ from litellm.proxy.agent_endpoints.auth.agent_permission_handler import Restrict from litellm.proxy.agent_endpoints.endpoints import router, user_api_key_auth from litellm.types.agents import AgentResponse +CALLER: Final = UserAPIKeyAuth(api_key="hashed-caller-key", team_id="team-1", user_id="user-1") + TRANSLATOR: Final = AgentResponse( agent_id="translator", agent_name="document-translator", @@ -150,21 +152,23 @@ class TestSearchAgents: @pytest.mark.asyncio async def test_no_embedding_model_is_not_configured(self) -> None: outcome = await search_agents( - "q", AGENTS, 5, router=MagicMock(), embedding_model=None, index=AgentSearchIndex() + "q", AGENTS, 5, router=MagicMock(), embedding_model=None, index=AgentSearchIndex(), user_api_key_dict=CALLER ) assert isinstance(outcome, AgentSearchNotConfigured) assert "agent_search_embedding_model" in outcome.reason @pytest.mark.asyncio async def test_no_router_is_not_configured(self) -> None: - outcome = await search_agents("q", AGENTS, 5, router=None, embedding_model="m", index=AgentSearchIndex()) + outcome = await search_agents( + "q", AGENTS, 5, router=None, embedding_model="m", index=AgentSearchIndex(), user_api_key_dict=CALLER + ) assert isinstance(outcome, AgentSearchNotConfigured) @pytest.mark.asyncio async def test_router_embeddings_are_read_from_the_response(self) -> None: router = MagicMock() router.aembedding = AsyncMock( - side_effect=lambda model, input: litellm.EmbeddingResponse( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( model=model, data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], ) @@ -176,11 +180,35 @@ class TestSearchAgents: router=router, embedding_model="text-embedding-3-small", index=AgentSearchIndex(), + user_api_key_dict=CALLER, ) assert isinstance(outcome, AgentSearchHits) assert [hit.agent.agent_id for hit in outcome.hits] == ["translator"] assert router.aembedding.await_args.kwargs["model"] == "text-embedding-3-small" + @pytest.mark.asyncio + async def test_embedding_spend_is_attributed_to_the_calling_key(self) -> None: + router = MagicMock() + router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], + ) + ) + await search_agents( + "language translation", + AGENTS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + ) + metadata = router.aembedding.await_args.kwargs["metadata"] + assert metadata["user_api_key"] == "hashed-caller-key" + assert metadata["user_api_key_team_id"] == "team-1" + assert metadata["user_api_key_user_id"] == "user-1" + def _client(role: LitellmUserRoles) -> TestClient: app = FastAPI() @@ -205,7 +233,7 @@ def registry(monkeypatch: pytest.MonkeyPatch) -> MagicMock: def embedding_router(monkeypatch: pytest.MonkeyPatch) -> MagicMock: router = MagicMock() router.aembedding = AsyncMock( - side_effect=lambda model, input: litellm.EmbeddingResponse( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( model=model, data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], ) @@ -231,6 +259,7 @@ class TestGetAgentsQuery: body = response.json() assert [agent["agent_id"] for agent in body] == ["translator", "trip"] assert body[0]["search_score"] > body[1]["search_score"] + assert embedding_router.aembedding.await_args.kwargs["metadata"]["user_api_key_user_id"] == "u" def test_without_query_the_list_is_unchanged_and_unscored( self, registry: MagicMock, embedding_router: MagicMock, no_db: None