Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_helpers/resource_display_names.py: 100%

1"""Display names for ids stored on management objects. DB rows win; config-declared servers and agents fill the gaps.""" 

2 

3from collections.abc import Mapping, Sequence 

4from types import MappingProxyType 

5from typing import Final 

6 

7from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry 

8from litellm.proxy.utils import PrismaClient 

9from litellm.repositories.table_repositories import AgentsRepository, MCPServerRepository 

10from litellm.repositories.verification_token_repository import VerificationTokenRepository 

11from litellm.types.mcp_server.mcp_server_manager import MCPServer 

12 

13 

14async def mcp_server_display_names( 

15 prisma_client: PrismaClient, 

16 server_ids: Sequence[str], 

17 config_servers: Mapping[str, MCPServer], 

18) -> Mapping[str, str]: 

19 """server_id -> alias, falling back to server_name; config-only servers also fall back to their registry name.""" 

20 if not server_ids: 

21 return MappingProxyType({}) 

22 wanted: Final = frozenset(server_ids) 

23 where: Final = {"server_id": {"in": tuple(wanted)}} # mutable-ok: prisma where is a dict 

24 rows: Final = await MCPServerRepository(prisma_client).table.find_many(where=where) 

25 from_config: Final = { 

26 server_id: server.alias or server.server_name or server.name 

27 for server_id, server in config_servers.items() 

28 if server_id in wanted 

29 } 

30 from_db: Final = {row.server_id: name for row in rows if (name := row.alias or row.server_name)} 

31 return MappingProxyType({**from_config, **from_db}) 

32 

33 

34async def agent_display_names( 

35 prisma_client: PrismaClient, 

36 agent_ids: Sequence[str], 

37 registry: AgentRegistry, 

38) -> Mapping[str, str]: 

39 """agent_id -> agent_name. The registry covers config-declared agents and their legacy ids.""" 

40 if not agent_ids: 

41 return MappingProxyType({}) 

42 wanted: Final = frozenset(agent_ids) 

43 where: Final = {"agent_id": {"in": tuple(wanted)}} # mutable-ok: prisma where is a dict 

44 rows: Final = await AgentsRepository(prisma_client).table.find_many(where=where) 

45 from_registry: Final = { 

46 alias_id: agent.agent_name 

47 for agent in registry.get_agent_list() 

48 for alias_id in registry.ids_for_agent(agent.agent_id) 

49 if alias_id in wanted 

50 } 

51 from_db: Final = {row.agent_id: row.agent_name for row in rows} 

52 return MappingProxyType({**from_registry, **from_db}) 

53 

54 

55async def key_display_names(prisma_client: PrismaClient, tokens: Sequence[str]) -> Mapping[str, str]: 

56 """token hash -> key_alias for the keys that have one.""" 

57 if not tokens: 

58 return MappingProxyType({}) 

59 where: Final = {"token": {"in": tuple(frozenset(tokens))}} # mutable-ok: prisma where is a dict 

60 rows: Final = await VerificationTokenRepository(prisma_client).table.find_many(where=where) 

61 return MappingProxyType({row.token: row.key_alias for row in rows if row.key_alias})