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

1"""Semantic ranking over the in-memory A2A agent registry, shared by GET /v1/agents?query= and the agent_search MCP tool.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Sequence 

6from dataclasses import dataclass 

7from typing import TYPE_CHECKING, Final, TypeAlias 

8 

9from pydantic import BaseModel, ConfigDict, ValidationError 

10 

11from litellm.proxy.common_utils.semantic_text_index import ( 

12 Embedder, 

13 EmbeddingFailed, 

14 SemanticTextIndex, 

15 router_embedder, 

16) 

17from litellm.types.agents import AgentResponse 

18 

19if TYPE_CHECKING: 19 ↛ 20line 19 didn't jump to line 20 because the condition on line 19 was never true

20 from litellm.proxy._types import UserAPIKeyAuth 

21 from litellm.proxy.utils import ProxyLogging 

22 from litellm.router import Router 

23 

24DEFAULT_AGENT_SEARCH_TOP_K: Final = 5 

25 

26 

27@dataclass(frozen=True, slots=True) 

28class AgentSearchHit: 

29 agent: AgentResponse 

30 score: float 

31 

32 

33@dataclass(frozen=True, slots=True) 

34class AgentSearchHits: 

35 hits: tuple[AgentSearchHit, ...] 

36 

37 

38@dataclass(frozen=True, slots=True) 

39class AgentSearchNotConfigured: 

40 reason: str 

41 

42 

43@dataclass(frozen=True, slots=True) 

44class AgentSearchEmbeddingFailed: 

45 reason: str 

46 

47 

48AgentSearchOutcome: TypeAlias = AgentSearchHits | AgentSearchNotConfigured | AgentSearchEmbeddingFailed 

49 

50 

51class _SearchableSkill(BaseModel): 

52 model_config = ConfigDict(frozen=True, extra="ignore") 

53 

54 name: str = "" 

55 description: str = "" 

56 tags: tuple[str, ...] = () 

57 

58 

59class _SearchableCard(BaseModel): 

60 model_config = ConfigDict(frozen=True, extra="ignore") 

61 

62 description: str = "" 

63 skills: tuple[_SearchableSkill, ...] = () 

64 

65 

66class AgentSearchResult(BaseModel): 

67 model_config = ConfigDict(frozen=True) 

68 

69 agent_id: str 

70 agent_name: str 

71 description: str 

72 skills: tuple[_SearchableSkill, ...] 

73 score: float 

74 

75 

76def _searchable_card(agent: AgentResponse) -> _SearchableCard: 

77 try: 

78 return _SearchableCard.model_validate(agent.agent_card_params) 

79 except ValidationError: 

80 return _SearchableCard() 

81 

82 

83def _skill_text(skill: _SearchableSkill) -> str: 

84 return " ".join(part for part in (skill.name, skill.description, " ".join(skill.tags)) if part) 

85 

86 

87def agent_search_text(agent: AgentResponse) -> str: 

88 card: Final = _searchable_card(agent) 

89 skill_lines: Final = tuple(_skill_text(skill) for skill in card.skills) 

90 return "\n".join(part for part in (agent.agent_name, card.description, *skill_lines) if part) 

91 

92 

93def agent_search_result(hit: AgentSearchHit) -> AgentSearchResult: 

94 card: Final = _searchable_card(hit.agent) 

95 return AgentSearchResult( 

96 agent_id=hit.agent.agent_id, 

97 agent_name=hit.agent.agent_name, 

98 description=card.description, 

99 skills=card.skills, 

100 score=hit.score, 

101 ) 

102 

103 

104class AgentSearchIndex: 

105 """Caches one vector per distinct agent text per embedding model, so repeat searches only embed the query.""" 

106 

107 def __init__(self) -> None: 

108 self._index: Final = SemanticTextIndex() 

109 

110 async def search( 

111 self, query: str, agents: Sequence[AgentResponse], top_k: int, embed: Embedder, embedding_model: str 

112 ) -> AgentSearchHits | AgentSearchEmbeddingFailed: 

113 texts: Final = tuple(agent_search_text(agent) for agent in agents) 

114 scores: Final = await self._index.scores(query, texts, embed, embedding_model) 

115 if isinstance(scores, EmbeddingFailed): 

116 return AgentSearchEmbeddingFailed(reason=scores.reason) 

117 ranked: Final = sorted( 

118 (AgentSearchHit(agent=agent, score=score) for agent, score in zip(agents, scores, strict=True)), 

119 key=lambda hit: hit.score, 

120 reverse=True, 

121 ) 

122 return AgentSearchHits(hits=tuple(ranked[:top_k])) 

123 

124 

125global_agent_search_index: Final = AgentSearchIndex() 

126 

127 

128async def search_agents( 

129 query: str, 

130 agents: Sequence[AgentResponse], 

131 top_k: int, 

132 router: Router | None, 

133 embedding_model: str | None, 

134 index: AgentSearchIndex, 

135 user_api_key_dict: UserAPIKeyAuth, 

136 proxy_logging_obj: ProxyLogging, 

137) -> AgentSearchOutcome: 

138 if embedding_model is None: 138 ↛ 142line 138 didn't jump to line 142 because the condition on line 138 was always true

139 return AgentSearchNotConfigured( 

140 reason="agent search needs litellm_settings.agent_search_embedding_model set to an embedding model from model_list" 

141 ) 

142 if router is None: 

143 return AgentSearchNotConfigured(reason="agent search needs a model_list so the embedding model can be called") 

144 embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) 

145 return await index.search(query, agents, top_k, embed, embedding_model)