Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/policy_engine/response_retrieval.py: 39%

1from collections.abc import Mapping 

2from dataclasses import dataclass 

3from types import MappingProxyType 

4from typing import TYPE_CHECKING, Final, Literal, TypeAlias 

5 

6from pydantic import TypeAdapter 

7 

8from litellm._logging import verbose_proxy_logger 

9from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket 

10from litellm.proxy.common_utils.callback_utils import ( 

11 add_guardrail_to_applied_guardrails_header, 

12 add_policy_sources_to_metadata, 

13 add_policy_to_applied_policies_header, 

14) 

15from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body 

16from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry 

17from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher 

18from litellm.proxy.policy_engine.policy_registry import get_policy_registry 

19from litellm.proxy.policy_engine.policy_resolver import PolicyResolver 

20from litellm.responses.utils import ResponsesAPIRequestUtils 

21from litellm.router_utils.common_utils import resolve_model_group_alias 

22from litellm.types.proxy.policy_engine import PolicyMatchContext 

23from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline 

24 

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

26 from litellm.proxy._types import UserAPIKeyAuth 

27 from litellm.router import Router 

28 

29PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...] 

30 

31_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines) 

32 

33 

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

35class UngovernedRetrieval: 

36 reason: Literal["no router", "response id names no deployment", "deployment no longer in the router"] 

37 

38 

39def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | UngovernedRetrieval: 

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

41 return UngovernedRetrieval("no router") 

42 model_id: Final = ( 

43 ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) if isinstance(response_id, str) else None 

44 ) 

45 if model_id is None: 45 ↛ 47line 45 didn't jump to line 47 because the condition on line 45 was always true

46 return UngovernedRetrieval("response id names no deployment") 

47 deployment: Final = llm_router.get_deployment(model_id) 

48 if deployment is None: 

49 return UngovernedRetrieval("deployment no longer in the router") 

50 hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias) 

51 if hidden_by is not None: 

52 verbose_proxy_logger.warning( 

53 "Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), " 

54 "so a policy attached to the model name it was submitted as does not run on it", 

55 response_id, 

56 deployment.model_name, 

57 hidden_by, 

58 ) 

59 return deployment.model_name 

60 

61 

62def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None: 

63 if "*" in model_group: 

64 return "a wildcard deployment" 

65 aliases: Final = tuple( 

66 alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group 

67 ) 

68 if not aliases: 

69 return None 

70 return f"the target of model_group_alias {', '.join(aliases)}" 

71 

72 

73def _retrieval_context( 

74 data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str 

75) -> PolicyMatchContext: 

76 team_alias: Final = user_api_key_dict.team_alias 

77 key_alias: Final = user_api_key_dict.key_alias 

78 return PolicyMatchContext( 

79 team_alias=team_alias if isinstance(team_alias, str) else None, 

80 key_alias=key_alias if isinstance(key_alias, str) else None, 

81 model=model_group, 

82 tags=get_tags_from_request_body(data) or None, 

83 ) 

84 

85 

86def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]: 

87 matches: Final = get_attachment_registry().get_attached_policies_with_reasons( 

88 context, PolicyMatcher.policy_applies(context) 

89 ) 

90 if not matches: 

91 return (), MappingProxyType({}) 

92 applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( 

93 policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list 

94 context=context, 

95 ) 

96 post_call_pipelines: Final = tuple( 

97 (policy_name, pipeline) 

98 for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context( 

99 context=context, policy_names=applied_policy_names 

100 ) 

101 if pipeline.mode == "post_call" 

102 ) 

103 return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches}) 

104 

105 

106def attach_post_call_pipelines_to_retrieval( 

107 data: dict[str, object], # mutable-ok: request-state dict the policy engine hooks all write in place 

108 user_api_key_dict: "UserAPIKeyAuth", 

109 llm_router: "Router | None", 

110) -> None: 

111 if not get_policy_registry().is_initialized(): 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true

112 return 

113 model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router) 

114 if isinstance(model_group, UngovernedRetrieval): 114 ↛ 121line 114 didn't jump to line 121 because the condition on line 114 was always true

115 verbose_proxy_logger.warning( 

116 "Policy engine: background response %s is retrieved without its post_call policy pipelines (%s)", 

117 data.get("response_id"), 

118 model_group.reason, 

119 ) 

120 return 

121 context: Final = _retrieval_context(data, user_api_key_dict, model_group) 

122 post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context) 

123 _, bucket = get_or_create_metadata_bucket(data) 

124 already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ()) 

125 attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached) 

126 added: Final = tuple( 

127 (policy_name, pipeline) 

128 for policy_name, pipeline in post_call_pipelines 

129 if policy_name not in attached_policy_names 

130 ) 

131 if not added: 

132 return 

133 pipelines: Final = (*already_attached, *added) 

134 bucket["_guardrail_pipelines"] = pipelines 

135 bucket["_pipeline_managed_guardrails"] = frozenset( 

136 step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps 

137 ) 

138 for policy_name, _pipeline in added: 

139 add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name) 

140 for _policy_name, pipeline in added: 

141 for step in pipeline.steps: 

142 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail) 

143 add_policy_sources_to_metadata( 

144 request_data=data, 

145 policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict 

146 policy_name: policy_sources[policy_name] for policy_name, _pipeline in added 

147 }, 

148 ) 

149 verbose_proxy_logger.debug( 

150 "Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s", 

151 data.get("response_id"), 

152 model_group, 

153 ", ".join(policy_name for policy_name, _pipeline in added), 

154 )