Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/auto_router_compression.py: 42%

1""" 

2Decouples prompt compression between an auto router's routing decision and the model 

3it routes to, via ``auto_router_routing_compression`` / ``auto_router_model_compression`` 

4on the marker deployment: a guardrail name, or ``"none"``. 

5 

6Neither key set inherits today's behaviour. Either key set makes the auto router 

7authoritative and suppresses every other compression guardrail for that request. 

8""" 

9 

10import contextvars 

11from collections.abc import Iterable, Mapping, Sequence 

12from dataclasses import dataclass 

13from typing import TYPE_CHECKING, Final 

14 

15from litellm._logging import verbose_proxy_logger 

16from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket 

17from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX 

18from litellm.types.utils import GenericGuardrailAPIInputs 

19 

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

21 from litellm.integrations.custom_guardrail import CustomGuardrail 

22 from litellm.router import Router 

23 

24COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr", "typesafe"}) 

25_NO_COMPRESSION: Final = "none" 

26 

27# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a 

28# suppression list they can read is one they can replay to disable any guardrail. 

29_suppressed_compression_guardrails: Final[contextvars.ContextVar[frozenset[str]]] = contextvars.ContextVar( 

30 "litellm_auto_router_suppressed_compression_guardrails", default=frozenset() 

31) 

32 

33 

34def suppressed_compression_guardrails() -> frozenset[str]: 

35 """Names of the compression guardrails this request's auto router suppresses.""" 

36 return _suppressed_compression_guardrails.get() 

37 

38 

39# Only the proxy calls `arm_pre_call`, so on the SDK path nothing arms and nothing 

40# compresses; the router must not assume the model hop already ran. 

41_model_hop_armed: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( 

42 "litellm_auto_router_model_hop_armed", default=False 

43) 

44 

45 

46def model_hop_compression_armed() -> bool: 

47 """True when this request's model-side compression guardrail was actually armed.""" 

48 return _model_hop_armed.get() 

49 

50 

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

52class AutoRouterCompressionPolicy: 

53 """An auto router's compression choice for each hop. ``None`` means no compression.""" 

54 

55 routing: str | None 

56 model: str | None 

57 

58 @property 

59 def is_same(self) -> bool: 

60 return self.routing == self.model 

61 

62 

63def _normalized_compression_choice(raw: object) -> str | None: 

64 if not isinstance(raw, str) or not raw: 

65 return None 

66 return None if raw.strip().lower() == _NO_COMPRESSION else raw 

67 

68 

69def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRouterCompressionPolicy | None: 

70 raw_routing: Final = litellm_params.get("auto_router_routing_compression") 

71 raw_model: Final = litellm_params.get("auto_router_model_compression") 

72 if raw_routing is None and raw_model is None: 

73 return None 

74 return AutoRouterCompressionPolicy( 

75 routing=_normalized_compression_choice(raw_routing), 

76 model=_normalized_compression_choice(raw_model), 

77 ) 

78 

79 

80def policy_for_model( 

81 llm_router: "Router | None", 

82 model_alias: str, 

83 request_kwargs: Mapping[str, object], 

84 request_tags: Sequence[str], 

85) -> AutoRouterCompressionPolicy | None: 

86 """The compression policy of the auto router marker `model_alias` resolves to for this caller. 

87 

88 Pre-call arming and the routing hook both resolve through here, and here resolves through the 

89 router's own request-scoped deployment lookup, so an alias with several tag-scoped markers 

90 cannot suppress under one and then route under another, and a team router reached by its 

91 public name carries its policy for every principal that can reach it. 

92 """ 

93 if llm_router is None: 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true

94 return None 

95 deployments: Final = llm_router.deployments_for_request(model_alias, request_kwargs) 

96 markers: Final = tuple( 

97 litellm_params 

98 for deployment in deployments 

99 if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) # pyright: ignore[reportUnnecessaryIsInstance] # filters out non-Mapping 

100 and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) 

101 ) 

102 requested: Final = frozenset(request_tags) 

103 tag_matched: Final = tuple( 

104 params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) 

105 ) 

106 # Untagged only: a marker scoped to tags this request lacks describes other traffic. 

107 untagged: Final = tuple(params for params in markers if not params.get("tags")) 

108 # Lazy, so the first marker carrying a policy wins and the rest are never read. 

109 candidates: Final = (policy_from_litellm_params(params) for params in (*tag_matched, *untagged)) 

110 return next((policy for policy in candidates if policy is not None), None) 

111 

112 

113def _compression_guardrail_classes() -> tuple[type, ...]: 

114 """The registered guardrail classes whose provider compresses prompts.""" 

115 from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry 

116 

117 return tuple(cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS) 

118 

119 

120def is_compression_guardrail(guardrail: object) -> bool: 

121 """Whether `guardrail` is an instance of a compression guardrail provider. 

122 

123 Both hops validate through here: the policy fields are operator-supplied names, and 

124 an unvalidated one would get handed the conversation and invoked. 

125 """ 

126 classes: Final = _compression_guardrail_classes() 

127 return bool(classes) and isinstance(guardrail, classes) 

128 

129 

130def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: 

131 """Every currently-active guardrail whose type is a compression guardrail.""" 

132 import litellm 

133 from litellm.integrations.custom_guardrail import CustomGuardrail 

134 

135 if not _compression_guardrail_classes(): 

136 return () 

137 active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) 

138 return tuple(cb for cb in active if is_compression_guardrail(cb) and cb.guardrail_name) 

139 

140 

141async def arm_pre_call( 

142 data: dict[str, object], # mutable-ok: arms the live request dict in place 

143 llm_router: "Router | None", 

144) -> None: 

145 """Apply an auto router's compression policy, if any, before guardrails run. 

146 

147 Suppresses every other compression guardrail and re-enables the model-side 

148 guardrail the policy names (if any) even when it isn't ``default_on``. 

149 """ 

150 _suppressed_compression_guardrails.set(frozenset()) 

151 _model_hop_armed.set(False) 

152 if llm_router is None: 

153 return 

154 

155 model_alias: Final = data.get("model") 

156 if not isinstance(model_alias, str) or not model_alias: 

157 return 

158 

159 from litellm.router_strategy.tag_based_routing import ( 

160 _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # used in router.py and budget_limiter.py too 

161 ) 

162 

163 policy: Final = policy_for_model( 

164 llm_router=llm_router, 

165 model_alias=model_alias, 

166 request_kwargs=data, 

167 request_tags=_get_tags_from_request_kwargs(data), 

168 ) 

169 if policy is None: 169 ↛ 172line 169 didn't jump to line 172 because the condition on line 169 was always true

170 return 

171 

172 _suppressed_compression_guardrails.set( 

173 frozenset( 

174 name 

175 for guardrail in _active_compression_guardrails() 

176 if (name := guardrail.guardrail_name) and name != policy.model 

177 ) 

178 ) 

179 

180 # Arming adds the name to `metadata["guardrails"]`, which runs it even if not default_on. 

181 armed_model_hop: Final = policy.model is not None and any( 

182 guardrail.guardrail_name == policy.model for guardrail in _active_compression_guardrails() 

183 ) 

184 if policy.model is not None and not armed_model_hop: 

185 verbose_proxy_logger.warning( 

186 "AutoRouter compression: '%s' is not an active compression guardrail; the model hop is uncompressed", 

187 policy.model, 

188 ) 

189 

190 if armed_model_hop: 

191 _model_hop_armed.set(True) 

192 _, metadata = get_or_create_metadata_bucket(data) 

193 requested: Final = metadata.get("guardrails") 

194 existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () 

195 if policy.model not in existing: 

196 # A list: litellm_pre_call_utils isinstance-checks this key and drops a tuple. 

197 metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list 

198 

199 

200def _as_routing_messages( 

201 messages: Iterable[Mapping[str, object]], 

202) -> list[dict[str, object]]: # mutable-ok: shape fixed by the pre-routing hook protocol 

203 """A fresh, independently mutable copy, the shape the pre-routing hook takes.""" 

204 return [dict(message) for message in messages] # mutable-ok: shape fixed by the pre-routing hook protocol 

205 

206 

207async def messages_for_routing( 

208 policy: AutoRouterCompressionPolicy | None, 

209 # list[dict], not Sequence[Mapping]: fixed by the async_pre_routing_hook protocol. 

210 messages: list[dict[str, object]] | None, # mutable-ok: shape fixed by the pre-routing hook protocol 

211 request_kwargs: Mapping[str, object], 

212) -> list[dict[str, object]] | None: # mutable-ok: shape fixed by the pre-routing hook protocol 

213 """Messages to use for a routing decision, per `policy.routing`. None means the 

214 caller should route on whatever it already has. 

215 

216 Reads the live messages, never a pre-guardrail copy: this compresses through a real 

217 guardrail that POSTs the text out, so routing on a pre-masking snapshot would leak 

218 what the masking guardrail stripped. When the model hop already compressed and the 

219 hops differ, routing therefore reads the compressed text rather than the original. 

220 """ 

221 if policy is None or policy.routing is None: 

222 return None 

223 

224 if not messages: 

225 return None 

226 

227 from litellm.proxy.common_utils.registry_read_through import ( 

228 get_initialized_guardrail_with_read_through, 

229 ) 

230 

231 guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing) 

232 if guardrail is None: 

233 verbose_proxy_logger.warning( 

234 "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing 

235 ) 

236 return _as_routing_messages(messages) 

237 

238 if not is_compression_guardrail(guardrail): 

239 verbose_proxy_logger.warning( 

240 "AutoRouter compression: guardrail '%s' is not a compression guardrail; routing on uncompressed messages", 

241 policy.routing, 

242 ) 

243 return _as_routing_messages(messages) 

244 

245 inputs: Final[GenericGuardrailAPIInputs] = { 

246 "structured_messages": _as_routing_messages(messages) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape 

247 } 

248 model: Final = request_kwargs.get("model") 

249 # Throwaway: apply_guardrail writes stats here, so routing never double-counts into 

250 # extract_compression_saved_tokens. 

251 stats_sink: Final = {"messages": messages, "model": model} # mutable-ok: apply_guardrail writes its stats here 

252 result: Final = await guardrail.apply_guardrail( 

253 inputs=inputs, 

254 request_data=stats_sink, 

255 input_type="request", 

256 ) 

257 compressed: Final = result.get("structured_messages") 

258 return compressed if isinstance(compressed, list) else _as_routing_messages(messages)