scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(a2a): attribute agent search embedding spend to the calling key

Browse source
This commit is contained in:
mateo-berri 2026-08-27 19:11:00 -07:00
parent 6de53732ee
commit e9cc9c9bc3
5 changed files with 56 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -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):

19
litellm/proxy/agent_endpoints/agent_search.py
View file

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

7
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -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:

1
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
View file

@ -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",

37
tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
View file

@ -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