Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/agent_endpoints/model_list_helpers.py: 51%

1""" 

2Helper functions for appending A2A agents to model lists. 

3 

4Used by proxy model endpoints to make agents appear in UI alongside models. 

5""" 

6 

7from litellm._logging import verbose_proxy_logger 

8from litellm.proxy._types import UserAPIKeyAuth 

9from litellm.types.proxy.management_endpoints.model_management_endpoints import ( 

10 ModelGroupInfoProxy, 

11) 

12 

13 

14async def append_agents_to_model_group( 

15 model_groups: list[ModelGroupInfoProxy], 

16 user_api_key_dict: UserAPIKeyAuth, 

17) -> list[ModelGroupInfoProxy]: 

18 """ 

19 Append A2A agents to model groups list for UI display. 

20 

21 Converts agents to model format with "a2a/<agent-name>" naming 

22 so they appear in playground and work with LiteLLM routing. 

23 """ 

24 try: 

25 from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry 

26 from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( 

27 AgentRequestHandler, 

28 RestrictedAgentAccess, 

29 ) 

30 

31 match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): 

32 case RestrictedAgentAccess(allowed_agent_ids): 32 ↛ 33line 32 didn't jump to line 33 because the pattern on line 32 never matched

33 for agent_id in allowed_agent_ids: 

34 agent = global_agent_registry.get_agent_by_id(agent_id) 

35 if agent is not None: 

36 model_groups.append( 

37 ModelGroupInfoProxy( 

38 model_group=f"a2a/{agent.agent_name}", 

39 mode="chat", 

40 providers=["a2a"], 

41 ) 

42 ) 

43 case _: 

44 pass 

45 except Exception as e: 

46 verbose_proxy_logger.debug("Error appending agents to model_group/info: %s", e) 

47 

48 return model_groups 

49 

50 

51async def append_agents_to_model_info( 

52 models: list[dict], 

53 user_api_key_dict: UserAPIKeyAuth, 

54) -> list[dict]: 

55 """ 

56 Append A2A agents to model info list for UI display. 

57 

58 Converts agents to model format with "a2a/<agent-name>" naming 

59 so they appear in models page and work with LiteLLM routing. 

60 """ 

61 try: 

62 from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry 

63 from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( 

64 AgentRequestHandler, 

65 RestrictedAgentAccess, 

66 ) 

67 

68 match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): 

69 case RestrictedAgentAccess(allowed_agent_ids): 69 ↛ 70line 69 didn't jump to line 70 because the pattern on line 69 never matched

70 for agent_id in allowed_agent_ids: 

71 agent = global_agent_registry.get_agent_by_id(agent_id) 

72 if agent is not None: 

73 models.append( 

74 { 

75 "model_name": f"a2a/{agent.agent_name}", 

76 "litellm_params": { 

77 "model": f"a2a/{agent.agent_name}", 

78 "custom_llm_provider": "a2a", 

79 }, 

80 "model_info": { 

81 "id": agent.agent_id, 

82 "mode": "chat", 

83 "db_model": True, 

84 "created_by": agent.created_by, 

85 "created_at": agent.created_at, 

86 "updated_at": agent.updated_at, 

87 }, 

88 } 

89 ) 

90 case _: 

91 pass 

92 except Exception as e: 

93 verbose_proxy_logger.debug("Error appending agents to v2/model/info: %s", e) 

94 

95 return models