From 4aa2bd990a1749135ccb675afa1d39dc5b49c327 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Tue, 10 Nov 2026 00:29:02 +0000 Subject: [PATCH] feat(policy-engine): resolve proxy policies in rust and delete the python engine --- litellm-rust/Cargo.lock | 9 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/policy-engine/Cargo.toml | 13 + .../crates/policy-engine/src/attachment.rs | 193 ++++++ .../crates/policy-engine/src/condition.rs | 292 ++++++++ litellm-rust/crates/policy-engine/src/lib.rs | 646 ++++++++++++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 5 + .../crates/python-bridge/src/policy_engine.rs | 238 +++++++ litellm/proxy/common_utils/callback_utils.py | 4 +- litellm/proxy/litellm_pre_call_utils.py | 286 +++----- .../policy_endpoints/endpoints.py | 26 +- litellm/proxy/policy_engine/__init__.py | 6 - litellm/proxy/policy_engine/architecture.md | 21 +- .../policy_engine/attachment_registry.py | 131 +--- .../policy_engine/condition_evaluator.py | 107 --- litellm/proxy/policy_engine/init_policies.py | 5 +- litellm/proxy/policy_engine/policy_matcher.py | 218 ------ .../proxy/policy_engine/policy_registry.py | 105 +-- .../policy_engine/policy_resolve_endpoints.py | 54 +- .../proxy/policy_engine/policy_resolver.py | 296 -------- .../proxy/policy_engine/response_retrieval.py | 23 +- litellm/proxy/utils.py | 8 +- litellm/rust_bridge/_native.pyi | 57 ++ litellm/rust_bridge/policy_engine.py | 330 +++++++++ .../test_policy_endpoints.py | 18 +- .../policy_engine/test_attachment_registry.py | 43 +- .../policy_engine/test_condition_evaluator.py | 145 ---- .../proxy/policy_engine/test_policy_engine.py | 559 +++++++++++++++ .../policy_engine/test_policy_matcher.py | 434 ------------ .../policy_engine/test_policy_resolver.py | 255 ------- .../policy_engine/test_policy_versioning.py | 26 +- .../unit/proxy/test_litellm_pre_call_utils.py | 164 +++-- tests/unit/realtime_api/test_main.py | 20 +- tests/unit/rust_bridge/test_policy_engine.py | 61 ++ 35 files changed, 2823 insertions(+), 1977 deletions(-) create mode 100644 litellm-rust/crates/policy-engine/Cargo.toml create mode 100644 litellm-rust/crates/policy-engine/src/attachment.rs create mode 100644 litellm-rust/crates/policy-engine/src/condition.rs create mode 100644 litellm-rust/crates/policy-engine/src/lib.rs create mode 100644 litellm-rust/crates/python-bridge/src/policy_engine.rs delete mode 100644 litellm/proxy/policy_engine/condition_evaluator.py delete mode 100644 litellm/proxy/policy_engine/policy_matcher.py delete mode 100644 litellm/proxy/policy_engine/policy_resolver.py create mode 100644 litellm/rust_bridge/policy_engine.py delete mode 100644 tests/unit/proxy/policy_engine/test_condition_evaluator.py create mode 100644 tests/unit/proxy/policy_engine/test_policy_engine.py delete mode 100644 tests/unit/proxy/policy_engine/test_policy_matcher.py delete mode 100644 tests/unit/proxy/policy_engine/test_policy_resolver.py create mode 100644 tests/unit/rust_bridge/test_policy_engine.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c20454c7058..3e09a970617 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4371,6 +4371,14 @@ dependencies = [ "url", ] +[[package]] +name = "litellm-policy-engine" +version = "0.1.0" +dependencies = [ + "fancy-regex 0.19.2", + "rstest", +] + [[package]] name = "litellm-python-bridge" version = "0.1.0" @@ -4402,6 +4410,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-policy-engine", "litellm-python-compat", "litellm-redis-client", "litellm-redis-connect", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 72d6694ba79..ff6cef5f537 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -14,6 +14,7 @@ litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } litellm-router-types = { path = "crates/router-types" } litellm-router-core = { path = "crates/router-core" } +litellm-policy-engine = { path = "crates/policy-engine" } litellm-tracing = { path = "crates/tracing" } litellm-spend-clickhouse = { path = "crates/spend-clickhouse" } litellm-storage-clickhouse = { path = "crates/storage-clickhouse" } diff --git a/litellm-rust/crates/policy-engine/Cargo.toml b/litellm-rust/crates/policy-engine/Cargo.toml new file mode 100644 index 00000000000..c80a36af4b2 --- /dev/null +++ b/litellm-rust/crates/policy-engine/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "litellm-policy-engine" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "The proxy policy engine of litellm/proxy/policy_engine: which policies apply to a request and the guardrails and pipelines they resolve to" + +[dependencies] +fancy-regex.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/policy-engine/src/attachment.rs b/litellm-rust/crates/policy-engine/src/attachment.rs new file mode 100644 index 00000000000..6a6f96e864d --- /dev/null +++ b/litellm-rust/crates/policy-engine/src/attachment.rs @@ -0,0 +1,193 @@ +use crate::MatchContext; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct Attachment { + pub policy: String, + pub global: bool, + pub teams: Vec, + pub keys: Vec, + pub models: Vec, + pub tags: Vec, + pub priority: Option, + pub default: bool, +} + +pub(crate) fn matches_pattern(value: Option<&str>, patterns: &[String]) -> bool { + let Some(value) = value else { + return patterns.iter().any(|pattern| pattern == "*"); + }; + patterns + .iter() + .any(|pattern| match pattern.strip_suffix('*') { + Some(prefix) => value.starts_with(prefix), + None => value == pattern, + }) +} + +fn matches_or_wildcard(value: Option<&str>, patterns: &[String]) -> bool { + if patterns.is_empty() { + return true; + } + matches_pattern(value, patterns) +} + +impl Attachment { + pub(crate) fn in_scope(&self, context: &MatchContext) -> bool { + if self.global { + return true; + } + matches_or_wildcard(context.team_alias.as_deref(), &self.teams) + && matches_or_wildcard(context.key_alias.as_deref(), &self.keys) + && matches_or_wildcard(context.model.as_deref(), &self.models) + && (self.tags.is_empty() + || context + .tags + .iter() + .any(|tag| matches_pattern(Some(tag), &self.tags))) + } + + fn specificity(&self) -> (u8, u8) { + if self.global { + return (0, 0); + } + let dimensions = [ + (&self.teams, 1), + (&self.keys, 2), + (&self.tags, 3), + (&self.models, 4), + ] + .into_iter() + .filter(|(values, _)| !values.is_empty()) + .map(|(_, specificity)| specificity); + dimensions.fold((0, 0), |(max, count), specificity| { + (max.max(specificity), count + 1) + }) + } + + pub(crate) fn order(&self) -> (bool, i32, u8, u8) { + let (specificity, dimensions) = self.specificity(); + ( + self.priority.is_none(), + self.priority.unwrap_or(0), + specificity, + dimensions, + ) + } + + pub(crate) fn match_reason(&self, context: &MatchContext) -> String { + let reason = self.scope_reason(context); + if self.default { + format!("default:{reason}") + } else { + reason + } + } + + fn scope_reason(&self, context: &MatchContext) -> String { + if self.global { + return "scope:*".to_owned(); + } + let tag = (!self.tags.is_empty()) + .then(|| { + context + .tags + .iter() + .find(|tag| matches_pattern(Some(tag), &self.tags)) + }) + .flatten() + .map(|tag| format!("tag:{tag}")); + let named = |values: &[String], label: &str, value: &Option| { + value + .as_deref() + .filter(|value| !values.is_empty() && !value.is_empty()) + .map(|value| format!("{label}:{value}")) + }; + let reasons: Vec = [ + tag, + named(&self.teams, "team", &context.team_alias), + named(&self.keys, "key", &context.key_alias), + named(&self.models, "model", &context.model), + ] + .into_iter() + .flatten() + .collect(); + if reasons.is_empty() { + "scope:default".to_owned() + } else { + reasons.join("+") + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn strings(values: &[&str]) -> Vec { + values.iter().map(|value| (*value).to_owned()).collect() + } + + #[rstest] + #[case::none_matches_only_star(None, &["*"], true)] + #[case::none_misses_prefix(None, &["team-*"], false)] + #[case::prefix(Some("team-a"), &["team-*"], true)] + #[case::star_is_empty_prefix(Some(""), &["*"], true)] + #[case::exact(Some("team-a"), &["team-a"], true)] + #[case::star_only_at_end(Some("team-a"), &["*-a"], false)] + #[case::any_of(Some("b"), &["a", "b"], true)] + fn wildcard_patterns( + #[case] value: Option<&str>, + #[case] patterns: &[&str], + #[case] expected: bool, + ) { + assert_eq!(matches_pattern(value, &strings(patterns)), expected); + } + + fn context() -> MatchContext { + MatchContext { + team_alias: Some("team-a".to_owned()), + key_alias: Some(String::new()), + model: Some("gpt-4".to_owned()), + tags: strings(&["env:dev", "health-x"]), + } + } + + #[rstest] + #[case::global(Attachment { global: true, teams: strings(&["nope"]), ..Default::default() }, true, "scope:*")] + #[case::unscoped_matches_everything(Attachment::default(), true, "scope:default")] + #[case::team(Attachment { teams: strings(&["team-*"]), ..Default::default() }, true, "team:team-a")] + #[case::empty_key_alias_matches_but_gives_no_reason(Attachment { keys: strings(&["*"]), ..Default::default() }, true, "scope:default")] + #[case::first_matching_tag_then_dimensions(Attachment { tags: strings(&["health-*", "env:*"]), models: strings(&["gpt-4"]), ..Default::default() }, true, "tag:env:dev+model:gpt-4")] + #[case::tag_miss(Attachment { tags: strings(&["prod"]), ..Default::default() }, false, "scope:default")] + #[case::default_prefix(Attachment { teams: strings(&["team-a"]), default: true, ..Default::default() }, true, "default:team:team-a")] + fn scope_and_reason( + #[case] attachment: Attachment, + #[case] in_scope: bool, + #[case] reason: &str, + ) { + assert_eq!(attachment.in_scope(&context()), in_scope); + assert_eq!(attachment.match_reason(&context()), reason); + } + + #[test] + fn tags_never_match_an_untagged_request() { + let attachment = Attachment { + tags: strings(&["*"]), + ..Default::default() + }; + assert!(!attachment.in_scope(&MatchContext::default())); + } + + #[rstest] + #[case::priority_before_none(Attachment { priority: Some(9), models: strings(&["m"]), ..Default::default() }, Attachment { global: true, ..Default::default() })] + #[case::lower_priority_first(Attachment { priority: Some(-1), ..Default::default() }, Attachment { priority: Some(0), ..Default::default() })] + #[case::global_before_scoped(Attachment { global: true, teams: strings(&["t"]), ..Default::default() }, Attachment { teams: strings(&["t"]), ..Default::default() })] + #[case::broader_dimension_first(Attachment { teams: strings(&["t"]), ..Default::default() }, Attachment { keys: strings(&["k"]), ..Default::default() })] + #[case::fewer_dimensions_first(Attachment { models: strings(&["m"]), ..Default::default() }, Attachment { models: strings(&["m"]), teams: strings(&["t"]), ..Default::default() })] + #[case::most_specific_dimension_decides(Attachment { tags: strings(&["x"]), keys: strings(&["k"]), teams: strings(&["t"]), ..Default::default() }, Attachment { models: strings(&["m"]), ..Default::default() })] + fn ordering(#[case] earlier: Attachment, #[case] later: Attachment) { + assert!(earlier.order() < later.order()); + } +} diff --git a/litellm-rust/crates/policy-engine/src/condition.rs b/litellm-rust/crates/policy-engine/src/condition.rs new file mode 100644 index 00000000000..fc5c97fb493 --- /dev/null +++ b/litellm-rust/crates/policy-engine/src/condition.rs @@ -0,0 +1,292 @@ +use std::fmt; +use std::sync::Arc; + +use fancy_regex::Regex; + +const ESCAPES: &[u8] = br".^$*+?{}[]|()\-/dDwWsSbBA"; +const CLASS_ESCAPES: &[u8] = br".^$*+?{}[]|()\-/dDwWsS"; +const SCOPED_FLAGS: &[u8] = b"ims"; + +/// Whether fancy-regex reads `pattern` exactly as Python's `re` does on [`plain_ascii`] values. +/// Only a vetted subset qualifies: printable ASCII, the metacharacters, the escapes above, counted +/// repeats with digits only, non-possessive quantifiers, `(?:`, the `i`/`m`/`s` flags and classes +/// without nesting, set operations or a leading `]`. Everything else is matched by Python. +fn read_alike(pattern: &str) -> bool { + let bytes = pattern.as_bytes(); + if !bytes.iter().all(u8::is_ascii_graphic_or_space) { + return false; + } + let mut at = 0; + let mut after_quantifier = false; + while let Some(&byte) = bytes.get(at) { + let (width, quantifier) = match byte { + b'\\' => match bytes.get(at + 1) { + Some(escaped) if ESCAPES.contains(escaped) => (2, false), + _ => return false, + }, + b'[' => match class_width(&bytes[at..]) { + Some(width) => (width, false), + None => return false, + }, + b'(' if bytes.get(at + 1) == Some(&b'?') => match flags_width(&bytes[at + 2..]) { + Some(width) => (width + 2, false), + None => return false, + }, + b'{' => match repeat_width(&bytes[at..]) { + Some(width) => (width, true), + None => return false, + }, + b'+' if after_quantifier => return false, + b'*' | b'+' | b'?' => (1, true), + _ => (1, false), + }; + at += width; + after_quantifier = quantifier; + } + true +} + +trait GraphicOrSpace { + fn is_ascii_graphic_or_space(&self) -> bool; +} + +impl GraphicOrSpace for u8 { + fn is_ascii_graphic_or_space(&self) -> bool { + self.is_ascii_graphic() || *self == b' ' + } +} + +fn class_width(class: &[u8]) -> Option { + let mut at = if class.get(1) == Some(&b'^') { 2 } else { 1 }; + if class.get(at) == Some(&b']') { + return None; + } + loop { + match *class.get(at)? { + b']' => return Some(at + 1), + b'[' => return None, + b'\\' if CLASS_ESCAPES.contains(class.get(at + 1)?) => at += 2, + b'\\' => return None, + operator @ (b'&' | b'-' | b'~') if class.get(at + 1) == Some(&operator) => return None, + _ => at += 1, + } + } +} + +fn flags_width(group: &[u8]) -> Option { + let flags = group + .iter() + .take_while(|byte| SCOPED_FLAGS.contains(byte)) + .count(); + match group.get(flags)? { + b':' | b')' => Some(flags + 1), + _ => None, + } +} + +fn repeat_width(repeat: &[u8]) -> Option { + let body = &repeat[1..]; + let close = body.iter().position(|&byte| byte == b'}')?; + let (low, high) = match body[..close].iter().position(|&byte| byte == b',') { + Some(comma) => (&body[..comma], &body[comma + 1..close]), + None => (&body[..close], &b""[..]), + }; + let digits = |part: &[u8]| part.iter().all(u8::is_ascii_digit); + (!low.is_empty() && digits(low) && digits(high)).then_some(close + 2) +} + +/// Values the read-alike subset matches the same way in both engines. Python's `$` also matches +/// before a final newline, its `\s` covers `\x1c`-`\x1f`, its `\B` never matches an empty string +/// and its Unicode tables differ from Rust's, so anything else is matched by Python. +fn plain_ascii(value: &str) -> bool { + !value.is_empty() && value.bytes().all(|byte| byte.is_ascii_graphic_or_space()) +} + +#[derive(Clone)] +pub struct HostRegex(Arc bool + Send + Sync>); + +impl HostRegex { + pub fn new(full_match: impl Fn(&str) -> bool + Send + Sync + 'static) -> Self { + Self(Arc::new(full_match)) + } +} + +impl fmt::Debug for HostRegex { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("HostRegex") + } +} + +#[derive(Debug, Clone)] +pub enum PatternSpec { + Literal(String), + Regex { text: String, host: HostRegex }, +} + +#[derive(Debug)] +enum Pattern { + Literal(String), + Regex { + text: String, + regex: Regex, + host: HostRegex, + }, + Host { + text: String, + host: HostRegex, + }, +} + +impl Pattern { + fn new(spec: PatternSpec) -> Self { + let (text, host) = match spec { + PatternSpec::Literal(text) => return Self::Literal(text), + PatternSpec::Regex { text, host } => (text, host), + }; + if !read_alike(&text) { + return Self::Host { text, host }; + } + match Regex::new(&format!(r"\A(?:{text})\z")) { + Ok(regex) => Self::Regex { text, regex, host }, + Err(_) => Self::Host { text, host }, + } + } + + fn matches(&self, value: &str) -> bool { + match self { + Self::Literal(text) => text == value, + Self::Regex { text, regex, host } => { + text == value + || if plain_ascii(value) { + regex.is_match(value).unwrap_or_else(|_| (host.0)(value)) + } else { + (host.0)(value) + } + } + Self::Host { text, host } => text == value || (host.0)(value), + } + } +} + +#[derive(Debug)] +pub struct ModelCondition(Vec); + +impl ModelCondition { + pub fn new(patterns: impl IntoIterator) -> Self { + Self(patterns.into_iter().map(Pattern::new).collect()) + } + + pub fn matches(&self, model: Option<&str>) -> bool { + model.is_some_and(|model| self.0.iter().any(|pattern| pattern.matches(model))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use rstest::rstest; + + use super::*; + + fn regex(text: &str) -> PatternSpec { + PatternSpec::Regex { + text: text.to_owned(), + host: HostRegex::new(|_| panic!("Rust reads this pattern itself")), + } + } + + #[rstest] + #[case::exact(vec![regex("gpt-4")], Some("gpt-4"), true)] + #[case::regex_is_anchored_at_both_ends(vec![regex("gpt-4")], Some("gpt-4o"), false)] + #[case::regex_full_match(vec![regex("gpt-4.*")], Some("gpt-4o-mini"), true)] + #[case::regex_needs_whole_value(vec![regex("gpt")], Some("my-gpt"), false)] + #[case::alternation_backtracks_to_full_match(vec![regex("a|ab")], Some("ab"), true)] + #[case::scoped_inline_flag(vec![regex("(?i)GPT-4")], Some("gpt-4"), true)] + #[case::literal_is_exact(vec![PatternSpec::Literal("gpt-4(".to_owned())], Some("gpt-4("), true)] + #[case::literal_is_never_regex(vec![PatternSpec::Literal("abc(?i)".to_owned())], Some("ABC"), false)] + #[case::any_of_list(vec![regex("claude-.*"), regex("gpt-4")], Some("gpt-4"), true)] + #[case::empty_list_never_matches(vec![], Some("gpt-4"), false)] + #[case::missing_model_never_matches(vec![regex(".*")], None, false)] + fn matches_like_python_fullmatch( + #[case] patterns: Vec, + #[case] model: Option<&str>, + #[case] expected: bool, + ) { + assert_eq!(ModelCondition::new(patterns).matches(model), expected); + } + + #[rstest] + #[case::nested_set("[[]")] + #[case::ascii_flag(r"(?a)\w")] + #[case::verbose_comment_swallows_anchor("(?x) a b # c")] + #[case::set_intersection("[a&&b]")] + #[case::set_difference("[a--b]")] + #[case::symmetric_difference("[a~~b]")] + #[case::named_character(r"\N{LATIN CAPITAL LETTER A}")] + #[case::posix_class("[[:alpha:]]")] + #[case::word_start_escape(r"\")] + #[case::word_boundary_form(r"\b{start}a")] + #[case::class_opening_with_bracket_range("[]-z]")] + #[case::negated_class_opening_with_bracket("[^]a]")] + #[case::octal_escape_after_a_group(r"(a)\01")] + #[case::backreference(r"(a)\1")] + #[case::verbose_counted_repeat("(?x).{ 1 }")] + #[case::possessive_counted_repeat("(?:a+){2}+")] + #[case::possessive_repeat("a++")] + #[case::counted_repeat_without_minimum("a{,2}")] + #[case::lookahead("a(?=b)")] + #[case::named_group("(?Pa)")] + #[case::absolute_end(r"a\Z")] + #[case::non_ascii_pattern("é.*")] + #[case::newline_in_pattern("a$\n")] + fn patterns_rust_reads_differently_go_to_the_host(#[case] text: &str) { + let calls = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&calls); + let condition = ModelCondition::new([PatternSpec::Regex { + text: text.to_owned(), + host: HostRegex::new(move |value| { + seen.fetch_add(1, Ordering::SeqCst); + value == "host says yes" + }), + }]); + assert!(condition.matches(Some("host says yes"))); + assert!(!condition.matches(Some("anything else"))); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[rstest] + #[case::information_separator(r"a\sb", "a\x1cb")] + #[case::non_ascii_value(r"a\wb", "a\u{b2}b")] + #[case::case_folding(r"(?i)i", "\u{131}")] + #[case::trailing_newline("a$", "a\n")] + #[case::empty_value(r"\B", "")] + fn values_the_engines_read_differently_go_to_the_host(#[case] text: &str, #[case] value: &str) { + let condition = ModelCondition::new([PatternSpec::Regex { + text: text.to_owned(), + host: HostRegex::new(|value| value != "a"), + }]); + assert!(condition.matches(Some(value))); + } + + #[rstest] + #[case::scoped_flags("(?i:gpt)-4", "GPT-4")] + #[case::global_flags("(?ims)gpt.*", "GPT-4o")] + #[case::counted_repeats(r"gpt-\d{1,2}", "gpt-42")] + #[case::class_with_escapes(r"[\w.\-]+", "gpt-4.1")] + #[case::lazy_quantifier("gpt.*?o", "gpt-4o")] + #[case::word_boundary(r"gpt\b.*", "gpt 4")] + fn read_alike_patterns_stay_in_rust(#[case] text: &str, #[case] value: &str) { + assert!(ModelCondition::new([regex(text)]).matches(Some(value))); + } + + #[test] + fn exact_text_matches_without_asking_the_host() { + let condition = ModelCondition::new([PatternSpec::Regex { + text: "[[]".to_owned(), + host: HostRegex::new(|_| panic!("an exact match never needs the host")), + }]); + assert!(condition.matches(Some("[[]"))); + } +} diff --git a/litellm-rust/crates/policy-engine/src/lib.rs b/litellm-rust/crates/policy-engine/src/lib.rs new file mode 100644 index 00000000000..5029525bd8c --- /dev/null +++ b/litellm-rust/crates/policy-engine/src/lib.rs @@ -0,0 +1,646 @@ +//! The proxy policy engine of `litellm/proxy/policy_engine`: which policies apply to a request, +//! and the guardrails and pipelines they resolve to. Pure logic over an immutable snapshot of the +//! policy and attachment registries, so the Python proxy and the gateway can share it. + +mod attachment; +mod condition; + +use std::collections::{HashMap, HashSet}; + +pub use attachment::Attachment; +pub use condition::{HostRegex, ModelCondition, PatternSpec}; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct MatchContext { + pub team_alias: Option, + pub key_alias: Option, + pub model: Option, + pub tags: Vec, +} + +#[derive(Debug, Default)] +pub struct Policy { + pub inherit: Option, + pub add: Vec, + pub remove: Vec, + pub condition: Option, + pub pipeline_guardrails: Option>, +} + +impl Policy { + fn applies(&self, context: &MatchContext) -> bool { + self.condition + .as_ref() + .is_none_or(|condition| condition.matches(context.model.as_deref())) + } +} + +#[derive(Debug)] +pub struct PolicyVersion { + pub id: String, + pub name: String, + pub policy: Policy, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PipelineSource { + pub policy: String, + pub version_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AncestorAdmission { + pub policy: String, + pub ancestor: String, +} + +#[derive(Debug, Default, PartialEq, Eq)] +pub struct Resolution { + pub attached_policies: Vec, + pub unapplied_policies: Vec, + pub condition_misses: Vec, + pub policy_guardrails: Vec<(String, Vec)>, + pub applied_policies: Vec, + pub policy_sources: Vec<(String, String)>, + pub guardrails: Vec, + pub pipelines: Vec, + pub pipeline_managed_guardrails: Vec, + pub ancestor_admissions: Vec, +} + +#[derive(Debug, Default, PartialEq, Eq)] +pub struct ResolvedPolicy { + pub guardrails: Vec, + pub inheritance_chain: Vec, +} + +#[derive(Debug, Default, PartialEq, Eq)] +pub struct AttachmentMatches { + pub matches: Vec<(String, String)>, + pub ancestor_admissions: Vec, +} + +#[derive(Debug, Default)] +pub struct PolicyEngine { + policies: HashMap, + versions: HashMap, + attachments: Vec, +} + +struct Policies<'a> { + engine: &'a PolicyEngine, + overrides: HashMap<&'a str, (&'a str, &'a Policy)>, +} + +impl<'a> Policies<'a> { + fn get(&self, name: &str) -> Option<&'a Policy> { + self.overrides + .get(name) + .map(|(_, policy)| *policy) + .or_else(|| self.engine.policies.get(name)) + } + + fn version_id(&self, name: &str) -> Option<&'a str> { + self.overrides.get(name).map(|(id, _)| *id) + } + + fn inheritance_chain(&self, name: &'a str) -> Vec<&'a str> { + let mut seen = HashSet::new(); + let mut chain = Vec::new(); + let mut current = Some(name); + while let Some(name) = current { + if !seen.insert(name) { + break; + } + let Some(policy) = self.get(name) else { + break; + }; + chain.push(name); + current = policy.inherit.as_deref(); + } + chain.reverse(); + chain + } + + fn applying_chain_members(&self, name: &'a str, context: &MatchContext) -> Vec<&'a str> { + self.inheritance_chain(name) + .into_iter() + .filter(|member| { + self.get(member) + .is_some_and(|policy| policy.applies(context)) + }) + .collect() + } + + fn guardrails(&self, name: &'a str, context: Option<&MatchContext>) -> Vec<&'a str> { + let mut guardrails = Vec::new(); + for policy in self + .inheritance_chain(name) + .into_iter() + .filter_map(|member| self.get(member)) + .filter(|policy| context.is_none_or(|context| policy.applies(context))) + { + for guardrail in &policy.add { + push_unique(&mut guardrails, guardrail.as_str()); + } + guardrails.retain(|kept| !policy.remove.iter().any(|removed| removed == kept)); + } + guardrails + } +} + +fn push_unique<'a>(values: &mut Vec<&'a str>, value: &'a str) { + if !values.contains(&value) { + values.push(value); + } +} + +fn owned(values: Vec<&str>) -> Vec { + values.into_iter().map(str::to_owned).collect() +} + +impl PolicyEngine { + pub fn new( + policies: impl IntoIterator, + versions: impl IntoIterator, + attachments: Vec, + ) -> Self { + Self { + policies: policies.into_iter().collect(), + versions: versions + .into_iter() + .map(|version| (version.id, (version.name, version.policy))) + .collect(), + attachments, + } + } + + fn attached<'a>( + &'a self, + context: &MatchContext, + policies: Option<&Policies<'a>>, + admissions: &mut Vec, + ) -> Vec<&'a Attachment> { + let mut in_scope = Vec::new(); + for attachment in &self.attachments { + if !attachment.in_scope(context) { + continue; + } + let Some(policies) = policies else { + in_scope.push(attachment); + continue; + }; + let applying = policies.applying_chain_members(&attachment.policy, context); + let Some(first) = applying.first() else { + continue; + }; + if !applying.contains(&attachment.policy.as_str()) { + admissions.push(AncestorAdmission { + policy: attachment.policy.clone(), + ancestor: (*first).to_owned(), + }); + } + in_scope.push(attachment); + } + let any_non_default = in_scope.iter().any(|attachment| !attachment.default); + in_scope.retain(|attachment| attachment.default != any_non_default); + in_scope.sort_by_key(|attachment| attachment.order()); + let mut seen = HashSet::new(); + in_scope.retain(|attachment| seen.insert(attachment.policy.as_str())); + in_scope + } + + fn base(&self) -> Policies<'_> { + Policies { + engine: self, + overrides: HashMap::new(), + } + } + + /// Registry policies whose inheritance chain loops back on itself, sorted by name + pub fn circular_policies(&self) -> Vec { + let mut circular: Vec = self + .policies + .keys() + .filter(|name| self.inherits_in_a_loop(name)) + .cloned() + .collect(); + circular.sort_unstable(); + circular + } + + fn inherits_in_a_loop(&self, name: &str) -> bool { + let mut seen = HashSet::new(); + let mut current = Some(name); + while let Some(name) = current { + if !seen.insert(name) { + return true; + } + current = self + .policies + .get(name) + .and_then(|policy| policy.inherit.as_deref()); + } + false + } + + pub fn policy_guardrails(&self, name: &str, context: Option<&MatchContext>) -> ResolvedPolicy { + let base = self.base(); + ResolvedPolicy { + guardrails: owned(base.guardrails(name, context)), + inheritance_chain: owned(base.inheritance_chain(name)), + } + } + + pub fn attached_policies( + &self, + context: &MatchContext, + only_applying: bool, + ) -> AttachmentMatches { + let base = self.base(); + let mut ancestor_admissions = Vec::new(); + let matches = self + .attached( + context, + only_applying.then_some(&base), + &mut ancestor_admissions, + ) + .into_iter() + .map(|attachment| (attachment.policy.clone(), attachment.match_reason(context))) + .collect(); + AttachmentMatches { + matches, + ancestor_admissions, + } + } + + pub fn resolve( + &self, + context: &MatchContext, + names: &[String], + version_ids: &[String], + ) -> Resolution { + let mut overrides = HashMap::new(); + let mut fetched = Vec::new(); + for id in version_ids { + if let Some((name, policy)) = self.versions.get(id) { + overrides.insert(name.as_str(), (id.as_str(), policy)); + fetched.push(name.as_str()); + } + } + let merged = Policies { + engine: self, + overrides, + }; + let mut ancestor_admissions = Vec::new(); + let attached = self.attached(context, Some(&merged), &mut ancestor_admissions); + + let mut candidates = Vec::new(); + for name in attached + .iter() + .map(|attachment| attachment.policy.as_str()) + .chain(names.iter().map(String::as_str)) + .chain(fetched) + { + push_unique(&mut candidates, name); + } + let (applied, unapplied): (Vec<&str>, Vec<&str>) = candidates + .into_iter() + .partition(|name| !merged.applying_chain_members(name, context).is_empty()); + let mut condition_misses = Vec::new(); + for member in applied + .iter() + .flat_map(|name| merged.inheritance_chain(name)) + .filter(|member| { + merged + .get(member) + .is_some_and(|policy| !policy.applies(context)) + }) + { + push_unique(&mut condition_misses, member); + } + let policy_sources = applied + .iter() + .filter_map(|name| { + attached + .iter() + .find(|attachment| attachment.policy == *name) + .map(|attachment| ((*name).to_owned(), attachment.match_reason(context))) + }) + .collect(); + + let contributions: Vec<(&str, Vec<&str>)> = applied + .iter() + .map(|name| (*name, merged.guardrails(name, Some(context)))) + .collect(); + let mut guardrails = Vec::new(); + for guardrail in contributions.iter().flat_map(|(_, added)| added.iter()) { + push_unique(&mut guardrails, guardrail); + } + + let pipelines: Vec<(&str, &Vec)> = applied + .iter() + .filter_map(|name| { + let policy = merged.get(name).filter(|policy| policy.applies(context))?; + policy + .pipeline_guardrails + .as_ref() + .map(|steps| (*name, steps)) + }) + .collect(); + let mut pipeline_managed_guardrails = Vec::new(); + for guardrail in pipelines.iter().flat_map(|(_, steps)| steps.iter()) { + push_unique(&mut pipeline_managed_guardrails, guardrail.as_str()); + } + + Resolution { + attached_policies: attached + .iter() + .map(|attachment| attachment.policy.clone()) + .collect(), + unapplied_policies: owned(unapplied), + condition_misses: owned(condition_misses), + policy_guardrails: contributions + .into_iter() + .map(|(name, added)| (name.to_owned(), owned(added))) + .collect(), + applied_policies: owned(applied), + policy_sources, + guardrails: owned(guardrails), + pipelines: pipelines + .into_iter() + .map(|(name, _)| PipelineSource { + policy: name.to_owned(), + version_id: merged.version_id(name).map(str::to_owned), + }) + .collect(), + pipeline_managed_guardrails: owned(pipeline_managed_guardrails), + ancestor_admissions, + } + } +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + + use super::*; + + fn strings(values: &[&str]) -> Vec { + values.iter().map(|value| (*value).to_owned()).collect() + } + + fn policy(inherit: Option<&str>, add: &[&str], remove: &[&str]) -> Policy { + Policy { + inherit: inherit.map(str::to_owned), + add: strings(add), + remove: strings(remove), + ..Policy::default() + } + } + + fn only_for(model: &str, policy: Policy) -> Policy { + Policy { + condition: Some(ModelCondition::new([PatternSpec::Regex { + text: model.to_owned(), + host: HostRegex::new(|_| unreachable!()), + }])), + ..policy + } + } + + fn global(policy: &str) -> Attachment { + Attachment { + policy: policy.to_owned(), + global: true, + ..Attachment::default() + } + } + + fn context(model: &str) -> MatchContext { + MatchContext { + model: Some(model.to_owned()), + ..MatchContext::default() + } + } + + #[fixture] + fn engine() -> PolicyEngine { + PolicyEngine::new( + [ + ("base".to_owned(), policy(None, &["pii", "toxicity"], &[])), + ( + "child".to_owned(), + only_for("gpt-.*", policy(Some("base"), &["audit"], &["toxicity"])), + ), + ("loop-a".to_owned(), policy(Some("loop-b"), &["a"], &[])), + ("loop-b".to_owned(), policy(Some("loop-a"), &["b"], &[])), + ("orphan".to_owned(), policy(Some("missing"), &["o"], &[])), + ( + "piped".to_owned(), + Policy { + pipeline_guardrails: Some(strings(&["pii", "judge", "pii"])), + ..policy(None, &["judge"], &[]) + }, + ), + ], + [PolicyVersion { + id: "v1".to_owned(), + name: "base".to_owned(), + policy: only_for("claude-.*", policy(None, &["draft"], &[])), + }], + vec![global("child")], + ) + } + + #[rstest] + fn resolution_explains_what_was_attached_skipped_and_contributed(engine: PolicyEngine) { + let resolution = + engine.resolve(&context("claude-3"), &strings(&["missing", "orphan"]), &[]); + assert_eq!(resolution.attached_policies, strings(&["child"])); + assert_eq!(resolution.applied_policies, strings(&["child", "orphan"])); + assert_eq!(resolution.unapplied_policies, strings(&["missing"])); + assert_eq!(resolution.condition_misses, strings(&["child"])); + assert_eq!( + resolution.policy_guardrails, + vec![ + ("child".to_owned(), strings(&["pii", "toxicity"])), + ("orphan".to_owned(), strings(&["o"])), + ] + ); + } + + #[rstest] + fn child_condition_hit_adds_and_removes_through_the_chain(engine: PolicyEngine) { + let resolution = engine.resolve(&context("gpt-4"), &[], &[]); + assert_eq!(resolution.applied_policies, strings(&["child"])); + assert_eq!( + resolution.policy_sources, + vec![("child".to_owned(), "scope:*".to_owned())] + ); + assert_eq!(resolution.guardrails, strings(&["pii", "audit"])); + assert!(resolution.ancestor_admissions.is_empty()); + } + + #[rstest] + fn child_condition_miss_still_applies_through_its_ancestor(engine: PolicyEngine) { + let resolution = engine.resolve(&context("claude-3"), &[], &[]); + assert_eq!(resolution.applied_policies, strings(&["child"])); + assert_eq!(resolution.guardrails, strings(&["pii", "toxicity"])); + assert_eq!( + resolution.ancestor_admissions, + vec![AncestorAdmission { + policy: "child".to_owned(), + ancestor: "base".to_owned(), + }] + ); + } + + #[rstest] + #[case::cycle_keeps_the_walk_up_to_the_repeat(&["loop-a"], &["b", "a"])] + #[case::missing_parent_ends_the_chain(&["orphan"], &["o"])] + #[case::unknown_request_policy_is_dropped(&["nope"], &[])] + fn inheritance_edges(#[case] names: &[&str], #[case] expected_guardrails: &[&str]) { + let engine = PolicyEngine::new(engine().policies.into_iter(), [], Vec::new()); + let resolution = engine.resolve(&context("m"), &strings(names), &[]); + assert_eq!(resolution.guardrails, strings(expected_guardrails)); + assert!(resolution.policy_sources.is_empty()); + } + + #[rstest] + fn version_override_replaces_the_named_policy_for_this_request(engine: PolicyEngine) { + let resolution = engine.resolve(&context("claude-3"), &[], &strings(&["v1", "unknown"])); + assert_eq!(resolution.applied_policies, strings(&["child", "base"])); + assert_eq!(resolution.guardrails, strings(&["draft"])); + } + + #[rstest] + fn a_version_whose_condition_misses_switches_its_policy_off() { + let engine = PolicyEngine::new( + [("base".to_owned(), policy(None, &["base-guard"], &[]))], + [PolicyVersion { + id: "v1".to_owned(), + name: "base".to_owned(), + policy: only_for("claude-.*", policy(None, &["draft"], &[])), + }], + vec![global("base")], + ); + let with_version = engine.resolve(&context("gpt-4"), &[], &strings(&["v1"])); + assert_eq!(with_version.applied_policies, Vec::::new()); + assert_eq!(with_version.guardrails, Vec::::new()); + assert_eq!(with_version.unapplied_policies, strings(&["base"])); + let without_version = engine.resolve(&context("gpt-4"), &[], &[]); + assert_eq!(without_version.guardrails, strings(&["base-guard"])); + } + + #[rstest] + fn pipelines_record_their_policy_and_managed_guardrails(engine: PolicyEngine) { + let resolution = engine.resolve(&context("m"), &strings(&["piped"]), &[]); + assert_eq!( + resolution.pipelines, + vec![PipelineSource { + policy: "piped".to_owned(), + version_id: None, + }] + ); + assert_eq!( + resolution.pipeline_managed_guardrails, + strings(&["pii", "judge"]) + ); + } + + #[rstest] + fn defaults_apply_only_when_no_other_attachment_matches() { + let engine = PolicyEngine::new( + [ + ("fallback".to_owned(), policy(None, &["f"], &[])), + ("team".to_owned(), policy(None, &["t"], &[])), + ], + [], + vec![ + Attachment { + default: true, + ..global("fallback") + }, + Attachment { + policy: "team".to_owned(), + teams: strings(&["a"]), + ..Attachment::default() + }, + ], + ); + let in_team = MatchContext { + team_alias: Some("a".to_owned()), + ..MatchContext::default() + }; + assert_eq!( + engine.resolve(&in_team, &[], &[]).applied_policies, + strings(&["team"]) + ); + let elsewhere = engine.resolve(&MatchContext::default(), &[], &[]); + assert_eq!(elsewhere.applied_policies, strings(&["fallback"])); + assert_eq!( + elsewhere.policy_sources, + vec![("fallback".to_owned(), "default:scope:*".to_owned())] + ); + } + + #[rstest] + #[case::without_context_every_chain_member_counts(None, &["pii", "audit"])] + #[case::with_context_conditions_filter_the_chain(Some("claude-3"), &["pii", "toxicity"])] + fn one_policy_resolves_through_its_chain( + engine: PolicyEngine, + #[case] model: Option<&str>, + #[case] guardrails: &[&str], + ) { + let context = model.map(context); + let resolved = engine.policy_guardrails("child", context.as_ref()); + assert_eq!(resolved.guardrails, strings(guardrails)); + assert_eq!(resolved.inheritance_chain, strings(&["base", "child"])); + } + + #[rstest] + fn unknown_policy_resolves_to_nothing(engine: PolicyEngine) { + assert_eq!( + engine.policy_guardrails("nope", None), + ResolvedPolicy::default() + ); + } + + #[rstest] + #[case::every_in_scope_attachment(false, &["ghost", "base"])] + #[case::only_policies_that_apply(true, &["base"])] + fn attachment_matching(#[case] only_applying: bool, #[case] expected: &[&str]) { + let engine = PolicyEngine::new( + [("base".to_owned(), policy(None, &["pii"], &[]))], + [], + vec![global("ghost"), global("base")], + ); + let matched = engine.attached_policies(&context("m"), only_applying); + let names: Vec<&str> = matched + .matches + .iter() + .map(|(name, _)| name.as_str()) + .collect(); + assert_eq!(names, expected); + } + + #[test] + fn circular_policies_names_every_policy_whose_chain_loops() { + let engine = PolicyEngine::new( + [ + ("a".to_owned(), policy(Some("b"), &[], &[])), + ("b".to_owned(), policy(Some("a"), &[], &[])), + ("into-loop".to_owned(), policy(Some("a"), &[], &[])), + ("self".to_owned(), policy(Some("self"), &[], &[])), + ("plain".to_owned(), policy(Some("missing"), &[], &[])), + ], + [], + Vec::new(), + ); + assert_eq!( + engine.circular_policies(), + strings(&["a", "b", "into-loop", "self"]) + ); + } +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 79773ca5559..ec801f9eaea 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -58,6 +58,7 @@ litellm-secrets-types.workspace = true litellm-llms-types.workspace = true litellm-router-types.workspace = true litellm-router-core.workspace = true +litellm-policy-engine.workspace = true litellm-host-python.workspace = true litellm-model-catalog.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 7513ff6b87b..3e4206276ec 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -10,6 +10,7 @@ mod http; mod logger; mod marshal; mod model_registry; +mod policy_engine; mod preflight; mod python_settings; mod redis; @@ -38,6 +39,8 @@ mod _native { #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::policy_engine::{NativePolicyEngine, PolicyResolution}; + #[pymodule_export] use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states}; #[pymodule_export] use crate::redis::client::NativeRedisClient; @@ -170,6 +173,8 @@ mod tests { "Tokenizer", "RouterCoreHandle", "RouterCall", + "PolicyEngine", + "PolicyResolution", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm-rust/crates/python-bridge/src/policy_engine.rs b/litellm-rust/crates/python-bridge/src/policy_engine.rs new file mode 100644 index 00000000000..2cfe21eb037 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/policy_engine.rs @@ -0,0 +1,238 @@ +//! The Python face of `litellm-policy-engine`. `litellm/rust_bridge/policy_engine.py` builds one +//! `PolicyEngine` per registry snapshot from plain tuples, reading the pydantic models with the +//! same getters the registries use. A condition pattern arrives with its compiled `re.Pattern`, +//! which Rust calls only for syntax it cannot read the way Python does. An exception raised by +//! that call is raised from the engine method that made it, as `re.fullmatch` would have. +use litellm_policy_engine::{ + Attachment, AttachmentMatches, HostRegex, MatchContext, ModelCondition, PatternSpec, Policy, + PolicyEngine, PolicyVersion, Resolution, ResolvedPolicy, +}; +use std::cell::RefCell; + +use pyo3::prelude::*; + +type NamedPairs = Vec<(String, String)>; + +thread_local! { + static HOST_ERROR: RefCell> = const { RefCell::new(None) }; +} + +fn host_regex(compiled: Py) -> HostRegex { + HostRegex::new(move |value| { + Python::attach( + |py| match compiled.bind(py).call_method1("fullmatch", (value,)) { + Ok(matched) => !matched.is_none(), + Err(error) => { + HOST_ERROR.with(|slot| { + slot.borrow_mut().get_or_insert(error); + }); + false + } + }, + ) + }) +} + +fn raising_host_errors(run: impl FnOnce() -> T) -> PyResult { + HOST_ERROR.with(|slot| slot.borrow_mut().take()); + let value = run(); + match HOST_ERROR.with(|slot| slot.borrow_mut().take()) { + Some(error) => Err(error), + None => Ok(value), + } +} + +#[derive(FromPyObject)] +struct PolicyRow( + Option, + Vec, + Vec, + Option>)>>, + Option>, +); + +impl From for Policy { + fn from(row: PolicyRow) -> Self { + let PolicyRow(inherit, add, remove, condition, pipeline_guardrails) = row; + Self { + inherit, + add, + remove, + condition: condition.map(|patterns| { + ModelCondition::new(patterns.into_iter().map(|(text, compiled)| match compiled { + Some(compiled) => PatternSpec::Regex { + text, + host: host_regex(compiled), + }, + None => PatternSpec::Literal(text), + })) + }), + pipeline_guardrails, + } + } +} + +#[derive(FromPyObject)] +struct AttachmentRow( + String, + bool, + Vec, + Vec, + Vec, + Vec, + Option, + bool, +); + +impl From for Attachment { + fn from(row: AttachmentRow) -> Self { + let AttachmentRow(policy, global, teams, keys, models, tags, priority, default) = row; + Self { + policy, + global, + teams, + keys, + models, + tags, + priority, + default, + } + } +} + +#[derive(FromPyObject)] +struct ContextRow(Option, Option, Option, Vec); + +impl From for MatchContext { + fn from(row: ContextRow) -> Self { + let ContextRow(team_alias, key_alias, model, tags) = row; + Self { + team_alias, + key_alias, + model, + tags, + } + } +} + +#[pyclass(frozen, name = "PolicyEngine")] +pub(crate) struct NativePolicyEngine { + engine: PolicyEngine, +} + +#[pymethods] +impl NativePolicyEngine { + #[new] + fn new( + policies: Vec<(String, PolicyRow)>, + versions: Vec<(String, String, PolicyRow)>, + attachments: Vec, + ) -> Self { + let engine = PolicyEngine::new( + policies + .into_iter() + .map(|(name, row)| (name, Policy::from(row))), + versions.into_iter().map(|(id, name, row)| PolicyVersion { + id, + name, + policy: row.into(), + }), + attachments.into_iter().map(Attachment::from).collect(), + ); + Self { engine } + } + + fn resolve( + &self, + context: ContextRow, + names: Vec, + version_ids: Vec, + ) -> PyResult { + raising_host_errors(|| { + self.engine + .resolve(&context.into(), &names, &version_ids) + .into() + }) + } + + #[pyo3(signature = (name, context=None))] + fn policy_guardrails( + &self, + name: &str, + context: Option, + ) -> PyResult<(Vec, Vec)> { + let context = context.map(MatchContext::from); + raising_host_errors(|| { + let ResolvedPolicy { + guardrails, + inheritance_chain, + } = self.engine.policy_guardrails(name, context.as_ref()); + (guardrails, inheritance_chain) + }) + } + + fn circular_policies(&self) -> Vec { + self.engine.circular_policies() + } + + fn attached_policies( + &self, + context: ContextRow, + only_applying: bool, + ) -> PyResult<(NamedPairs, NamedPairs)> { + raising_host_errors(|| { + let AttachmentMatches { + matches, + ancestor_admissions, + } = self + .engine + .attached_policies(&context.into(), only_applying); + ( + matches, + ancestor_admissions + .into_iter() + .map(|admission| (admission.policy, admission.ancestor)) + .collect(), + ) + }) + } +} + +#[pyclass(frozen, get_all, name = "PolicyResolution")] +pub(crate) struct PolicyResolution { + attached_policies: Vec, + unapplied_policies: Vec, + condition_misses: Vec, + policy_guardrails: Vec<(String, Vec)>, + applied_policies: Vec, + policy_sources: Vec<(String, String)>, + guardrails: Vec, + pipelines: Vec<(String, Option)>, + pipeline_managed_guardrails: Vec, + ancestor_admissions: Vec<(String, String)>, +} + +impl From for PolicyResolution { + fn from(resolution: Resolution) -> Self { + Self { + attached_policies: resolution.attached_policies, + unapplied_policies: resolution.unapplied_policies, + condition_misses: resolution.condition_misses, + policy_guardrails: resolution.policy_guardrails, + applied_policies: resolution.applied_policies, + policy_sources: resolution.policy_sources, + guardrails: resolution.guardrails, + pipelines: resolution + .pipelines + .into_iter() + .map(|source| (source.policy, source.version_id)) + .collect(), + pipeline_managed_guardrails: resolution.pipeline_managed_guardrails, + ancestor_admissions: resolution + .ancestor_admissions + .into_iter() + .map(|admission| (admission.policy, admission.ancestor)) + .collect(), + } + } +} diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c195779aa0a..2cfa57dcb99 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -639,7 +639,7 @@ def add_guardrail_scan_id( _metadata[GUARDRAIL_SCAN_METADATA_METADATA_KEY] = (*entries, entry) -def add_policy_to_applied_policies_header(request_data: dict, policy_name: str | None): +def add_policy_to_applied_policies_header(request_data: dict[str, object], policy_name: str | None) -> None: """ Add a policy name to the applied_policies list in request metadata. @@ -656,7 +656,7 @@ def add_policy_to_applied_policies_header(request_data: dict, policy_name: str | _metadata["applied_policies"] = [policy_name] -def add_policy_sources_to_metadata(request_data: dict, policy_sources: dict[str, str]): +def add_policy_sources_to_metadata(request_data: dict[str, object], policy_sources: Mapping[str, str]) -> None: """ Store policy match reasons in metadata for x-litellm-policy-sources header. diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index e2367ab8179..430a13af3be 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2,6 +2,7 @@ import asyncio import copy import itertools import json +import logging import re import time from collections import OrderedDict @@ -354,10 +355,10 @@ if TYPE_CHECKING: from litellm.integrations.otel.model.destination import OtelDestination from litellm.models.tag import LiteLLM_TagTable from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.proxy.utils import PrismaClient - from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext + from litellm.rust_bridge.policy_engine import PolicyEngineResolution + from litellm.types.proxy.policy_engine import PolicyMatchContext ProxyConfig = _ProxyConfig else: @@ -4043,8 +4044,8 @@ def _add_guardrails_from_policies_in_metadata( """ from litellm._logging import verbose_proxy_logger from litellm.proxy.policy_engine.policy_registry import get_policy_registry - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.proxy.utils import premium_user_check + from litellm.rust_bridge.policy_engine import current_policy_engine from litellm.types.proxy.policy_engine import PolicyMatchContext # Collect policy names from key and team metadata @@ -4082,18 +4083,13 @@ def _add_guardrails_from_policies_in_metadata( # Build context for policy resolution (model from request data) context: Final = PolicyMatchContext(model=data.get("model")) - # Get all policies from registry - all_policies: Final = registry.get_all_policies() + engine: Final = current_policy_engine() # Resolve guardrails from the specified policies resolved_guardrails: Final[set] = set() for policy_name in policy_names: if registry.has_policy(policy_name): - resolved_policy = PolicyResolver.resolve_policy_guardrails( - policy_name=policy_name, - policies=all_policies, - context=context, - ) + resolved_policy = engine.policy_guardrails(policy_name, context) resolved_guardrails.update(resolved_policy.guardrails) verbose_proxy_logger.debug( "Policy engine: resolved guardrails from policy '%s': %s", policy_name, resolved_policy.guardrails @@ -4236,126 +4232,6 @@ def _extract_policy_id(s: str) -> str | None: return s[len(POLICY_VERSION_ID_PREFIX) :].strip() or None -def _match_and_track_policies( - data: dict, - context: "PolicyMatchContext", - request_body_policies: Sequence[str], - policies_override: dict[str, "Policy"] | None = None, - attachment_registry_override: "AttachmentRegistry | None" = None, -) -> tuple[list[str], dict[str, str]]: - """ - Match policies via attachments and request body, track them in metadata. - - Returns: - Tuple of (applied_policy_names, policy_reasons) - """ - from litellm._logging import verbose_proxy_logger - from litellm.proxy.common_utils.callback_utils import ( - add_policy_sources_to_metadata, - add_policy_to_applied_policies_header, - ) - from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry - from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - - # Get matching policies via attachments (with match reasons for attribution) - attachment_registry: Final = ( - attachment_registry_override if attachment_registry_override is not None else get_attachment_registry() - ) - matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons( - context, PolicyMatcher.policy_applies(context, policies_override) - ) - matching_policy_names: Final = [m["policy_name"] for m in matches_with_reasons] - policy_reasons: Final = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons} - - verbose_proxy_logger.debug("Policy engine: matched policies via attachments: %s", matching_policy_names) - - # Combine attachment-based policies with dynamic request body policies - request_body_policies_list: Final = ( - tuple(request_body_policies) if request_body_policies and isinstance(request_body_policies, list) else () - ) - all_policy_names: Final = tuple(dict.fromkeys((*matching_policy_names, *request_body_policies_list))) - if request_body_policies_list: - verbose_proxy_logger.debug("Policy engine: added dynamic policies from request body: %s", request_body_policies) - - if not all_policy_names: - return [], {} - - # Filter to only policies whose conditions match the context - applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( - policy_names=list(all_policy_names), - context=context, - policies=policies_override, - ) - - verbose_proxy_logger.debug("Policy engine: applied policies (conditions matched): %s", applied_policy_names) - - # Track applied policies in metadata for response headers - for policy_name in applied_policy_names: - add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name) - - # Track policy attribution sources for x-litellm-policy-sources header - applied_reasons: Final = {name: policy_reasons[name] for name in applied_policy_names if name in policy_reasons} - add_policy_sources_to_metadata(request_data=data, policy_sources=applied_reasons) - - return applied_policy_names, policy_reasons - - -def _apply_resolved_guardrails_to_metadata( - data: dict, - metadata_variable_name: str, - context: "PolicyMatchContext", - policy_names: list[str] | None = None, - policies: dict[str, "Policy"] | None = None, -) -> None: - """Apply resolved guardrails and pipelines to request metadata.""" - from litellm._logging import verbose_proxy_logger - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver - - # Resolve guardrails from matching policies - resolved_guardrails: Final = PolicyResolver.resolve_guardrails_for_context( - context=context, - policies=policies, - policy_names=policy_names, - ) - - verbose_proxy_logger.debug("Policy engine: resolved guardrails: %s", resolved_guardrails) - - # Resolve pipelines from matching policies - pipelines: Final = PolicyResolver.resolve_pipelines_for_context( - context=context, - policies=policies, - policy_names=policy_names, - ) - - # Add resolved guardrails to request metadata - if metadata_variable_name not in data: - data[metadata_variable_name] = {} - - # Record the pipelines and the guardrails they step; the hook loops skip those per pipeline mode - if pipelines: - pipeline_managed_guardrails: Final = PolicyResolver.get_pipeline_managed_guardrails(pipelines) - data[metadata_variable_name]["_guardrail_pipelines"] = pipelines - data[metadata_variable_name]["_pipeline_managed_guardrails"] = pipeline_managed_guardrails - verbose_proxy_logger.debug( - "Policy engine: resolved %s pipeline(s), managed guardrails: %s", - len(pipelines), - pipeline_managed_guardrails, - ) - - if not resolved_guardrails and not pipelines: - return - - existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) - existing_guardrails_list: Final = ( - _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () - ) - - combined: Final = merged_guardrail_entries(existing_guardrails_list, resolved_guardrails) - data[metadata_variable_name]["guardrails"] = combined - - verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) - - def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_name: str) -> list[str] | None: from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body from litellm.router_strategy.tag_based_routing import without_baggage_mapped_tags @@ -4366,6 +4242,97 @@ def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_nam return list(tags) if tags else None +def _record_policy_engine_resolution( + data: dict[str, object], # mutable-ok: records into the live request dict, the same contract as the Python path + metadata_variable_name: str, + resolution: "PolicyEngineResolution", +) -> None: + from litellm.proxy.common_utils.callback_utils import ( + add_policy_sources_to_metadata, + add_policy_to_applied_policies_header, + ) + + for policy_name in resolution.applied_policies: + add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name) + add_policy_sources_to_metadata(request_data=data, policy_sources=resolution.policy_sources) + + bucket: Final = cast( # cast-ok: the proxy keeps a dict bucket under the request's metadata name + "dict[str, object]", data.setdefault(metadata_variable_name, {}) + ) + pipeline_fields: Final = ( + { + "_guardrail_pipelines": list(resolution.pipelines), + "_pipeline_managed_guardrails": set(resolution.pipeline_managed_guardrails), + } + if resolution.pipelines + else {} + ) + if not resolution.guardrails and not resolution.pipelines: + return + existing_guardrails: Final = bucket.get("guardrails", []) + existing_guardrails_list: Final = ( + _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () + ) + guardrails: Final = merged_guardrail_entries(existing_guardrails_list, resolution.guardrails) + bucket.update({**pipeline_fields, "guardrails": guardrails}) + verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", guardrails) + + +def _log_policy_engine_resolution( + context: "PolicyMatchContext", + names: Sequence[str], + version_ids: Sequence[str], + versions: Mapping[str, tuple[str, object]], + resolution: "PolicyEngineResolution", +) -> None: + if not verbose_proxy_logger.isEnabledFor(logging.DEBUG): + return + for policy_id in version_ids: + version = versions.get(policy_id) + if version is None: + verbose_proxy_logger.debug("Policy engine: policy version %s not found in cache, skipping", policy_id) + else: + verbose_proxy_logger.debug("Policy engine: loaded version by ID policy_%s -> %s", policy_id, version[0]) + verbose_proxy_logger.debug( + "Policy engine: matched policies via attachments: %s", list(resolution.attached_policies) + ) + if names: + verbose_proxy_logger.debug("Policy engine: added dynamic policies from request body: %s", list(names)) + for policy_name in resolution.condition_misses: + verbose_proxy_logger.debug( + "Policy '%s' condition did not match model=%s, skipping its guardrails and pipeline", + policy_name, + context.model, + ) + if resolution.unapplied_policies: + verbose_proxy_logger.debug( + "Policy engine: skipped policies with no applicable chain member: %s", + list(resolution.unapplied_policies), + ) + if not resolution.applied_policies: + verbose_proxy_logger.debug( + "No policies match context: team_alias=%s, key_alias=%s, model=%s", + context.team_alias, + context.key_alias, + context.model, + ) + return + verbose_proxy_logger.debug( + "Policy engine: applied policies (conditions matched): %s", list(resolution.applied_policies) + ) + for policy_name, guardrails in resolution.policy_guardrails.items(): + verbose_proxy_logger.debug("Policy '%s' contributes guardrails: %s", policy_name, list(guardrails)) + verbose_proxy_logger.debug("Policy engine: resolved guardrails: %s", list(resolution.guardrails)) + for policy_name, pipeline in resolution.pipelines: + verbose_proxy_logger.debug("Policy '%s' has pipeline with %s steps", policy_name, len(pipeline.steps)) + if resolution.pipelines: + verbose_proxy_logger.debug( + "Policy engine: resolved %s pipeline(s), managed guardrails: %s", + len(resolution.pipelines), + set(resolution.pipeline_managed_guardrails), + ) + + async def add_guardrails_from_policy_engine( data: dict[str, object], # mutable-ok: pops the request's policies and records the resolved guardrails in place metadata_variable_name: str, @@ -4427,53 +4394,22 @@ async def add_guardrails_from_policy_engine( context.tags, ) - # Separate policy names from policy version IDs (policy_) - request_body_names: Final[list[str]] = [] - request_body_version_ids: Final[list[str]] = [] - if request_body_policies_raw and isinstance(request_body_policies_raw, list): - for item in request_body_policies_raw: - if not isinstance(item, str): - continue - if _is_policy_version_id(item): - policy_id = _extract_policy_id(item) - if policy_id: - request_body_version_ids.append(policy_id) - else: - request_body_names.append(item) - - # Resolve policy versions by ID from in-memory cache (populated by sync job; no DB in hot path) - merged_policies: Final[dict[str, Policy]] = dict(registry.get_all_policies()) - fetched_policy_names: Final[list[str]] = [] - for policy_id in request_body_version_ids: - result = registry.get_policy_by_id_for_request(policy_id=policy_id) - if result is not None: - pname, policy = result - merged_policies[pname] = policy - fetched_policy_names.append(pname) - verbose_proxy_logger.debug("Policy engine: loaded version by ID policy_%s -> %s", policy_id, pname) - else: - verbose_proxy_logger.debug("Policy engine: policy version %s not found in cache, skipping", policy_id) - - # Build request body list: names + policy names from fetched versions - request_body_policies: Final = request_body_names + fetched_policy_names - - # Match and track policies (with merged_policies when we have version overrides) - applied_policy_names, _ = _match_and_track_policies( - data, - context, - request_body_policies, - policies_override=merged_policies if request_body_version_ids else None, + request_body_items: Final = ( + tuple(item for item in _REQUEST_ENTRIES.validate_python(request_body_policies_raw) if isinstance(item, str)) + if isinstance(request_body_policies_raw, list) + else () + ) + request_body_names: Final = tuple(item for item in request_body_items if not _is_policy_version_id(item)) + request_body_version_ids: Final = tuple( + policy_id for item in request_body_items if (policy_id := _extract_policy_id(item)) ) - # Resolve and apply guardrails. Use applied_policy_names so request-body policies - # (names + version IDs) are included. Use merged_policies when we have version overrides. - _apply_resolved_guardrails_to_metadata( - data, - metadata_variable_name, - context, - policy_names=applied_policy_names if applied_policy_names else None, - policies=merged_policies if request_body_version_ids else None, - ) + from litellm.rust_bridge.policy_engine import current_policy_engine + + engine: Final = current_policy_engine() + resolution: Final = engine.resolve(context, request_body_names, request_body_version_ids) + _log_policy_engine_resolution(context, request_body_names, request_body_version_ids, engine.versions, resolution) + _record_policy_engine_resolution(data, metadata_variable_name, resolution) _ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join( diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 7cd36d51511..e0109cef221 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -45,7 +45,7 @@ from litellm.proxy.guardrails.guardrail_hooks.custom_code import ( from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.policy_engine.policy_registry import get_policy_registry -from litellm.proxy.policy_engine.policy_resolver import PolicyResolver +from litellm.rust_bridge.policy_engine import current_policy_engine from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.policy_engine import ( PolicyGuardrailsResponse, @@ -156,14 +156,9 @@ async def apply_policies( "apply_policies: policy engine not initialized, skipping policy-resolved guardrails" ) else: - policies: Final = registry.get_all_policies() + engine: Final = current_policy_engine() for policy_name in policy_names: - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name=policy_name, - policies=policies, - context=None, - ) - guardrail_name_set.update(resolved.guardrails) + guardrail_name_set.update(engine.policy_guardrails(policy_name).guardrails) if not guardrail_name_set: return {"inputs": inputs, "guardrail_errors": guardrail_errors} @@ -528,7 +523,6 @@ async def get_policy_info( - Inheritance chain """ from litellm.proxy.policy_engine.policy_registry import get_policy_registry - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver registry: Final = get_policy_registry() @@ -545,7 +539,7 @@ async def get_policy_info( detail=f"Policy '{policy_name}' not found", ) - resolved = PolicyResolver.resolve_policy_guardrails(policy_name=policy_name, policies=registry.get_all_policies()) + resolved: Final = current_policy_engine().policy_guardrails(policy_name) return PolicyInfoResponse( policy_name=policy_name, @@ -594,9 +588,8 @@ async def test_policy_matching( - matching_policies: List of policy names that match - resolved_guardrails: Final list of guardrails that would be applied """ - from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver registry: Final = get_policy_registry() @@ -608,13 +601,8 @@ async def test_policy_matching( message="Policy engine not initialized. No policies loaded.", ) - policies: Final = registry.get_all_policies() - - # Get matching policies - matching_policy_names: Final = PolicyMatcher.get_matching_policies(context=context) - - # Resolve guardrails - resolved_guardrails: Final = PolicyResolver.resolve_guardrails_for_context(context=context, policies=policies) + matching_policy_names: Final = get_attachment_registry().get_attached_policies(context, only_applying=True) + resolved_guardrails: Final = list(current_policy_engine().resolve(context).guardrails) return PolicyTestResponse( context=context, diff --git a/litellm/proxy/policy_engine/__init__.py b/litellm/proxy/policy_engine/__init__.py index 18b37dc4852..1a59268b819 100644 --- a/litellm/proxy/policy_engine/__init__.py +++ b/litellm/proxy/policy_engine/__init__.py @@ -37,21 +37,15 @@ from litellm.proxy.policy_engine.attachment_registry import ( AttachmentRegistry, get_attachment_registry, ) -from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator -from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher from litellm.proxy.policy_engine.policy_registry import ( PolicyRegistry, get_policy_registry, ) -from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.proxy.policy_engine.policy_validator import PolicyValidator __all__ = [ "AttachmentRegistry", - "ConditionEvaluator", - "PolicyMatcher", "PolicyRegistry", - "PolicyResolver", "PolicyValidator", "get_attachment_registry", "get_policy_registry", diff --git a/litellm/proxy/policy_engine/architecture.md b/litellm/proxy/policy_engine/architecture.md index fa9cbeecf5d..dc8c2cdf4f4 100644 --- a/litellm/proxy/policy_engine/architecture.md +++ b/litellm/proxy/policy_engine/architecture.md @@ -15,8 +15,8 @@ flowchart TD subgraph PolicyEngine["Policy Engine"] PR[PolicyRegistry] PV[PolicyValidator] - PM[PolicyMatcher] - PRe[PolicyResolver] + AR[AttachmentRegistry] + PE[PolicyEngine, Rust] end subgraph Request["Incoming Request"] @@ -31,11 +31,10 @@ flowchart TD PC -->|validate| PV PV -->|errors/warnings| PR - CTX -->|match| PM - PM -->|matching policies| PRe - PR -->|policies| PM - PR -->|policies| PRe - PRe -->|resolve inheritance + add/remove| GR + PR -->|policies + versions| PE + AR -->|attachments| PE + CTX -->|resolve| PE + PE -->|match attachments, resolve inheritance + add/remove| GR ``` ## Components @@ -44,11 +43,11 @@ flowchart TD |-----------|------|-------------| | **PolicyRegistry** | `policy_registry.py` | In-memory singleton store for parsed policies | | **PolicyValidator** | `policy_validator.py` | Validates configs (guardrails, inheritance, teams/keys/models) | -| **PolicyMatcher** | `policy_matcher.py` | Matches request context against policy scopes | -| **PolicyResolver** | `policy_resolver.py` | Resolves final guardrails via inheritance chain | +| **AttachmentRegistry** | `attachment_registry.py` | In-memory store for policy attachments (where policies apply) | +| **PolicyEngine** | `litellm/rust_bridge/policy_engine.py`, crate `litellm-policy-engine` | Matches attachments, evaluates conditions and resolves inheritance in Rust over a snapshot of both registries | ## Flow 1. **Startup**: `init_policies()` loads policies from config, validates, and populates `PolicyRegistry` -2. **Request**: `PolicyMatcher` finds policies matching the request's team/key/model -3. **Resolution**: `PolicyResolver` traverses inheritance and applies add/remove to get final guardrails +2. **Request**: `current_policy_engine().resolve(context)` finds the attachments matching the request's team/key/model/tags +3. **Resolution**: the same call traverses inheritance, applies add/remove and returns the final guardrails and pipelines. The engine is rebuilt only when a registry replaces its policies or attachments diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 15e7b7e8482..74bb3308967 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -5,9 +5,8 @@ Attachments define WHERE policies apply, separate from the policy definitions. This allows the same policy to be attached to multiple scopes. """ -from collections.abc import Callable, Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypedDict from litellm._logging import verbose_proxy_logger @@ -20,8 +19,6 @@ from litellm.types.proxy.policy_engine import ( ) if TYPE_CHECKING: - from collections.abc import Sequence - from prisma.models import LiteLLM_PolicyAttachmentTable from litellm.proxy.utils import PrismaClient @@ -32,30 +29,6 @@ class PolicyAttachmentMatch(TypedDict): matched_via: str -def _attachment_specificity(attachment: PolicyAttachment) -> tuple[int, int]: - if attachment.is_global(): - return (0, 0) - - dims: Final = tuple( - specificity - for values, specificity in ( - (attachment.teams, 1), - (attachment.keys, 2), - (attachment.tags, 3), - (attachment.models, 4), - ) - if values - ) - return (max(dims, default=0), len(dims)) - - -def _attachment_sort_key(attachment: PolicyAttachment) -> tuple[int, int, int, int]: - specificity: Final = _attachment_specificity(attachment) - if attachment.priority is not None: - return (0, attachment.priority, *specificity) - return (1, 0, *specificity) - - class AttachmentRegistry: """ In-memory registry for storing and managing policy attachments. @@ -88,20 +61,21 @@ class AttachmentRegistry: attachments_config: List of attachment dictionaries from YAML. """ self._attachments = [] - - for attachment_data in attachments_config: - try: - attachment = self._parse_attachment(attachment_data) - self._attachments.append(attachment) - verbose_proxy_logger.debug("Loaded attachment for policy: %s", attachment.policy) - except Exception as e: - verbose_proxy_logger.error("Error loading attachment: %s", e) - raise ValueError(f"Invalid attachment: {e}") from e + self._attachments = [self._parse_config_attachment(attachment_data) for attachment_data in attachments_config] self._config_attachments = tuple(self._attachments) self._initialized = True verbose_proxy_logger.info("Loaded %s policy attachments", len(self._attachments)) + def _parse_config_attachment(self, attachment_data: Mapping[str, object]) -> PolicyAttachment: + try: + attachment: Final = self._parse_attachment(attachment_data) + except Exception as e: + verbose_proxy_logger.error("Error loading attachment: %s", e) + raise ValueError(f"Invalid attachment: {e}") from e + verbose_proxy_logger.debug("Loaded attachment for policy: %s", attachment.policy) + return attachment + def _parse_attachment(self, attachment_data: Mapping[str, object]) -> PolicyAttachment: """ Parse an attachment from raw configuration data. @@ -125,90 +99,32 @@ class AttachmentRegistry: } ) - def get_attached_policies( - self, - context: PolicyMatchContext, - policy_applies: Callable[[str], bool] | None = None, - ) -> list[str]: + def get_attached_policies(self, context: PolicyMatchContext, *, only_applying: bool = False) -> list[str]: """ Get list of policy names attached to the given context. Args: context: The request context to match against - policy_applies: Optional predicate; attachments whose policy does not apply are ignored + only_applying: Ignore attachments whose policy does not apply to the context Returns: List of policy names that are attached to matching scopes """ - return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context, policy_applies)] + return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context, only_applying=only_applying)] def get_attached_policies_with_reasons( - self, - context: PolicyMatchContext, - policy_applies: Callable[[str], bool] | None = None, + self, context: PolicyMatchContext, *, only_applying: bool = False ) -> list[PolicyAttachmentMatch]: """ Get list of policy names and match reasons for the given context. Returns a list of dicts with 'policy_name' and 'matched_via' keys. The 'matched_via' describes which dimension caused the match. - Attachments whose policy fails `policy_applies` are dropped before defaults are considered. + With `only_applying`, attachments whose policy does not apply are dropped before defaults are considered. """ - from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher + from litellm.rust_bridge.policy_engine import current_policy_engine - in_scope: Final = tuple( - attachment - for attachment in self._attachments - if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context) - and (policy_applies is None or policy_applies(attachment.policy)) - ) - non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default) - matching_attachments: Final = sorted( - non_default or tuple(attachment for attachment in in_scope if attachment.default), - key=_attachment_sort_key, - ) - broadest_attachment_by_policy: Final = MappingProxyType( - {attachment.policy: attachment for attachment in reversed(matching_attachments)} - ) - unique_attachments: Final = tuple( - broadest_attachment_by_policy[policy_name] - for policy_name in dict.fromkeys(attachment.policy for attachment in matching_attachments) - ) - - return [ - { - "policy_name": attachment.policy, - "matched_via": self._describe_match_reason(attachment, context), - } - for attachment in unique_attachments - ] - - @staticmethod - def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: - """Describe why an attachment matched the context.""" - reason: Final = AttachmentRegistry._describe_scope_match(attachment, context) - return f"default:{reason}" if attachment.default else reason - - @staticmethod - def _describe_scope_match(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: - from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - - if attachment.is_global(): - return "scope:*" - - reasons: Final = [] - if attachment.tags and context.tags: - matching_tags: Final = [t for t in context.tags if PolicyMatcher.matches_pattern(t, attachment.tags)] - if matching_tags: - reasons.append(f"tag:{matching_tags[0]}") - if attachment.teams and context.team_alias: - reasons.append(f"team:{context.team_alias}") - if attachment.keys and context.key_alias: - reasons.append(f"key:{context.key_alias}") - if attachment.models and context.model: - reasons.append(f"model:{context.model}") - - return "+".join(reasons) if reasons else "scope:default" + return list(current_policy_engine(self._attachments).attached_policies(context, only_applying=only_applying)) def is_policy_attached(self, policy_name: str, context: PolicyMatchContext) -> bool: """ @@ -224,6 +140,11 @@ class AttachmentRegistry: attached: Final = self.get_attached_policies(context) return policy_name in attached + def current_attachments(self) -> Sequence[PolicyAttachment]: + """The live attachments container. Every change replaces it instead of mutating it, so callers may cache + anything derived from it for as long as it stays the same object""" + return self._attachments + def get_all_attachments(self) -> list[PolicyAttachment]: """ Get all loaded attachments. @@ -278,7 +199,7 @@ class AttachmentRegistry: Args: attachment: PolicyAttachment object to add """ - self._attachments.append(attachment) + self._attachments = [*self._attachments, attachment] self._initialized = True verbose_proxy_logger.debug("Added attachment for policy: %s", attachment.policy) @@ -533,7 +454,9 @@ class AttachmentRegistry: ) for attachment_response in attachments ] - self._attachments = [*self._config_attachments, *db_attachments] + synced: Final = [*self._config_attachments, *db_attachments] + if synced != self._attachments: + self._attachments = synced self._initialized = True verbose_proxy_logger.info( diff --git a/litellm/proxy/policy_engine/condition_evaluator.py b/litellm/proxy/policy_engine/condition_evaluator.py deleted file mode 100644 index 2f3a4fa8015..00000000000 --- a/litellm/proxy/policy_engine/condition_evaluator.py +++ /dev/null @@ -1,107 +0,0 @@ -""" -Condition Evaluator - Evaluates policy conditions. - -Supports model-based conditions with exact match or regex patterns. -""" - -import re - -from litellm._logging import verbose_proxy_logger -from litellm.types.proxy.policy_engine import ( - PolicyCondition, - PolicyMatchContext, -) - - -class ConditionEvaluator: - """ - Evaluates policy conditions against request context. - - Supports model conditions with: - - Exact string match: "gpt-4" - - Regex pattern: "gpt-4.*" - - List of values: ["gpt-4", "gpt-4-turbo"] - """ - - @staticmethod - def evaluate( - condition: PolicyCondition | None, - context: PolicyMatchContext, - ) -> bool: - """ - Evaluate a policy condition against a request context. - - Args: - condition: The condition to evaluate (None = always matches) - context: The request context with team, key, model - - Returns: - True if condition matches, False otherwise - """ - # No condition means always matches - if condition is None: - return True - - # Check model condition - if condition.model is not None: - if not ConditionEvaluator._evaluate_model_condition( - condition=condition.model, - model=context.model, - ): - verbose_proxy_logger.debug( - "Condition failed: model=%s did not match %s", context.model, condition.model - ) - return False - - return True - - @staticmethod - def _evaluate_model_condition( - condition: str | list[str], - model: str | None, - ) -> bool: - """ - Evaluate a model condition. - - Args: - condition: String (exact or regex) or list of strings - model: The model name to check - - Returns: - True if model matches condition, False otherwise - """ - if model is None: - return False - - # Handle list of values - if isinstance(condition, list): - return any(ConditionEvaluator._matches_pattern(pattern, model) for pattern in condition) - - # Single value - check as pattern - return ConditionEvaluator._matches_pattern(condition, model) - - @staticmethod - def _matches_pattern(pattern: str, value: str) -> bool: - """ - Check if value matches pattern (exact match or regex). - - Args: - pattern: Pattern to match (exact string or regex) - value: Value to check - - Returns: - True if matches, False otherwise - """ - # First try exact match - if pattern == value: - return True - - # Try as regex pattern - try: - if re.fullmatch(pattern, value): - return True - except re.error: - # Invalid regex, treat as literal string (already checked above) - pass - - return False diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index 2d0e9d25368..6c37941a477 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -228,7 +228,7 @@ def get_policies_summary() -> dict[str, object]: Returns: Dictionary with policy information """ - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver + from litellm.rust_bridge.policy_engine import current_policy_engine policy_registry: Final = get_policy_registry() attachment_registry: Final = get_attachment_registry() @@ -236,7 +236,8 @@ def get_policies_summary() -> dict[str, object]: if not policy_registry.is_initialized(): return {"initialized": False, "policies": {}, "attachments": []} - resolved: Final = PolicyResolver.get_all_resolved_policies() + engine: Final = current_policy_engine() + resolved: Final = {name: engine.policy_guardrails(name) for name in policy_registry.current_policies()} return { "initialized": True, diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py deleted file mode 100644 index 2ea2def8331..00000000000 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ /dev/null @@ -1,218 +0,0 @@ -""" -Policy Matcher - Matches requests against policy attachments. - -Uses existing wildcard pattern matching helpers to determine which policies -apply to a given request based on team alias, key alias, and model. - -Policies are matched via policy_attachments which define WHERE each policy applies. -""" - -from collections.abc import Callable, Sequence -from typing import Final - -from litellm._logging import verbose_proxy_logger -from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.policy_engine.policy_resolver import PolicyResolver -from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext, PolicyScope - - -class PolicyMatcher: - """ - Matches incoming requests against policy attachments. - - Supports wildcard patterns: - - "*" matches everything - - "prefix-*" matches anything starting with "prefix-" - - Uses policy_attachments to determine which policies apply to a request. - """ - - @staticmethod - def matches_pattern(value: str | None, patterns: list[str]) -> bool: - """ - Check if a value matches any of the given patterns. - - Uses the existing RouteChecks.route_matches_wildcard_pattern helper. - - Args: - value: The value to check (e.g., team alias, key alias, model) - patterns: List of patterns to match against - - Returns: - True if value matches any pattern, False otherwise - """ - # If no value provided, only match if patterns include "*" - if value is None: - return "*" in patterns - - for pattern in patterns: - # Use existing wildcard pattern matching helper - if RouteChecks.route_matches_wildcard_pattern(route=value, pattern=pattern): - return True - - return False - - @staticmethod - def scope_matches(scope: PolicyScope, context: PolicyMatchContext) -> bool: - """ - Check if a policy scope matches the given context. - - A scope matches if ALL of its fields match: - - teams matches context.team_alias - - keys matches context.key_alias - - models matches context.model - - Args: - scope: The policy scope to check - context: The request context - - Returns: - True if scope matches context, False otherwise - """ - # Check teams - if not PolicyMatcher.matches_pattern(context.team_alias, scope.get_teams()): - return False - - # Check keys - if not PolicyMatcher.matches_pattern(context.key_alias, scope.get_keys()): - return False - - # Check models - if not PolicyMatcher.matches_pattern(context.model, scope.get_models()): - return False - - # Check tags (only if scope specifies tags) - # Unlike teams/keys/models, empty tags means "do not check" rather than "match all" - scope_tags: Final = scope.get_tags() - if scope_tags: - if not context.tags: - return False - # Match if ANY context tag matches ANY scope tag pattern - if not any(PolicyMatcher.matches_pattern(tag, scope_tags) for tag in context.tags): - return False - - return True - - @staticmethod - def get_matching_policies( - context: PolicyMatchContext, - ) -> list[str]: - """ - Get list of policy names that match the given context via attachments. - - Args: - context: The request context to match against - - Returns: - List of policy names that match the context - """ - from litellm.proxy.policy_engine.attachment_registry import ( - get_attachment_registry, - ) - - registry: Final = get_attachment_registry() - if not registry.is_initialized(): - verbose_proxy_logger.debug("AttachmentRegistry not initialized, returning empty list") - return [] - - return registry.get_attached_policies(context, PolicyMatcher.policy_applies(context)) - - @staticmethod - def get_matching_policies_from_registry( - context: PolicyMatchContext, - ) -> list[str]: - """ - Get list of policy names that match the given context from the global registry. - - Args: - context: The request context to match against - - Returns: - List of policy names that match the context - """ - return PolicyMatcher.get_matching_policies(context=context) - - @staticmethod - def policy_applies( - context: PolicyMatchContext, - policies: dict[str, Policy] | None = None, - ) -> Callable[[str], bool]: - """ - Predicate telling whether a policy exists and any policy in its - inheritance chain applies to the context. Admissions where the - policy's own condition missed but an ancestor applies are logged at - INFO, once per attachment scan. - """ - resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies() - - def applies(policy_name: str) -> bool: - applying: Final = PolicyMatcher._applying_chain_members( - policy_name=policy_name, context=context, policies=resolved - ) - if applying and policy_name not in applying: - verbose_proxy_logger.info( - "Policy '%s' applied through ancestor '%s' although its own condition did not match " - "(team_alias=%s, key_alias=%s, model=%s)", - policy_name, - applying[0], - context.team_alias, - context.key_alias, - context.model, - ) - return bool(applying) - - return applies - - @staticmethod - def _applying_chain_members( - policy_name: str, - context: PolicyMatchContext, - policies: dict[str, Policy], - ) -> tuple[str, ...]: - from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator - - chain: Final = PolicyResolver.resolve_inheritance_chain(policy_name=policy_name, policies=policies) - return tuple( - name - for name in chain - if (policy := policies.get(name)) is not None - and (policy.condition is None or ConditionEvaluator.evaluate(policy.condition, context)) - ) - - @staticmethod - def _registry_policies() -> dict[str, Policy]: - from litellm.proxy.policy_engine.policy_registry import get_policy_registry - - registry: Final = get_policy_registry() - return registry.get_all_policies() if registry.is_initialized() else {} - - @staticmethod - def get_policies_with_matching_conditions( - policy_names: Sequence[str], - context: PolicyMatchContext, - policies: dict[str, Policy] | None = None, - ) -> list[str]: - """ - Filter policies to only those that apply to the given context. - - A policy applies when any policy in its inheritance chain has no - condition or a condition that evaluates to True for the context. The - resolver then drops only the chain members whose own condition fails, - so a child whose condition misses still contributes the guardrails of - its unconditional ancestors. A missing policy resolves to an empty - chain and does not apply. - - Args: - policy_names: List of policy names to filter - context: The request context to evaluate conditions against - policies: Dictionary of all policies (if None, uses global registry) - - Returns: - List of policy names that apply to the context - """ - resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies() - return [ - policy_name - for policy_name in policy_names - if PolicyMatcher._applying_chain_members(policy_name, context, resolved) - ] diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index c5f7e14e61d..173199b41dd 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -18,6 +18,7 @@ from typing import ( Optional, Protocol, TypedDict, + TypeVar, Union, cast, # noqa: TID251 # prisma types the condition/pipeline Json columns as str, but reads return decoded values ) @@ -171,6 +172,17 @@ def _attribute_of(source: object, name: str) -> object: return attribute["value"] +_Key: Final = TypeVar("_Key") +_Value: Final = TypeVar("_Value") + + +def _same_or_new( + current: dict[_Key, _Value], synced: dict[_Key, _Value] +) -> dict[_Key, _Value]: # mutable-ok: registry containers are dicts + """The current container when a sync produced the same entries, so the cached policy engine stays valid""" + return current if list(current.items()) == list(synced.items()) else synced + + class PolicyRegistry: """ In-memory registry for storing and managing policies. @@ -204,20 +216,25 @@ class PolicyRegistry: self._sources = {} self._policies_by_id = {} - for policy_name, policy_data in policies_config.items(): - try: - policy = self.parse_policy(policy_name, policy_data) - self._policies[policy_name] = policy - verbose_proxy_logger.debug("Loaded policy: %s", policy_name) - except Exception as e: - verbose_proxy_logger.error("Error loading policy '%s': %s", policy_name, e) - raise ValueError(f"Invalid policy '{policy_name}': {e}") from e + self._policies = { + policy_name: self._parse_config_policy(policy_name, policy_data) + for policy_name, policy_data in policies_config.items() + } self._config_policies = dict(self._policies) self._sources = {policy_name: "config" for policy_name in self._policies} self._initialized = True verbose_proxy_logger.info("Loaded %s policies", len(self._policies)) + def _parse_config_policy(self, policy_name: str, policy_data: Mapping[str, object]) -> Policy: + try: + policy: Final = self.parse_policy(policy_name, dict(policy_data)) + except Exception as e: + verbose_proxy_logger.error("Error loading policy '%s': %s", policy_name, e) + raise ValueError(f"Invalid policy '{policy_name}': {e}") from e + verbose_proxy_logger.debug("Loaded policy: %s", policy_name) + return policy + def parse_policy(self, policy_name: str, policy_data: dict[str, Any]) -> Policy: """ Parse a policy from raw configuration data. @@ -296,6 +313,15 @@ class PolicyRegistry: """ return self._policies.copy() + def current_policies(self) -> Mapping[str, Policy]: + """The live policies container. Every change replaces it instead of mutating it, so callers may cache + anything derived from it for as long as it stays the same object""" + return self._policies + + def current_versions(self) -> Mapping[str, tuple[str, Policy]]: + """The live policy-version cache, replaced on every change like ``current_policies``""" + return self._policies_by_id + def get_policy_names(self) -> list[str]: """ Get list of all policy names. @@ -356,7 +382,7 @@ class PolicyRegistry: policy: Policy object to add source: Provenance of the policy ("db" or "config") """ - self._policies[policy_name] = policy + self._policies = {**self._policies, policy_name: policy} self._sources = {**self._sources, policy_name: source} if source == "config": self._config_policies = {**self._config_policies, policy_name: policy} @@ -378,11 +404,11 @@ class PolicyRegistry: return False config_fallback: Final = self._config_policies.get(policy_name) if config_fallback is not None: - self._policies[policy_name] = config_fallback + self._policies = {**self._policies, policy_name: config_fallback} self._sources = {**self._sources, policy_name: "config"} verbose_proxy_logger.debug("Removed policy: %s; restored config-defined version", policy_name) return True - del self._policies[policy_name] + self._policies = {name: policy for name, policy in self._policies.items() if name != policy_name} self._sources = {name: source for name, source in self._sources.items() if name != policy_name} verbose_proxy_logger.debug("Removed policy: %s", policy_name) return True @@ -600,6 +626,21 @@ class PolicyRegistry: verbose_proxy_logger.exception("Error getting policy from DB: %s", e) raise Exception(f"Error getting policy from DB: {e}") + def _parse_version_row(self, row: _PolicyRow) -> Policy: + return self.parse_policy( + row.policy_name, + { + "inherit": row.inherit, + "description": row.description, + "guardrails": { + "add": row.guardrails_add or [], + "remove": row.guardrails_remove or [], + }, + "condition": row.condition, + "pipeline": row.pipeline, + }, + ) + def get_policy_by_id_for_request(self, policy_id: str) -> tuple[str, Policy] | None: """ Return a policy version by ID from in-memory cache (no DB access). @@ -683,29 +724,17 @@ class PolicyRegistry: ) config_sources: Mapping[str, Literal["db", "config"]] = {name: "config" for name in self._config_policies} db_sources: Final[Mapping[str, Literal["db", "config"]]] = {name: "db" for name in db_policies} - self._policies = {**self._config_policies, **db_policies} + self._policies = _same_or_new(self._policies, {**self._config_policies, **db_policies}) self._sources = {**config_sources, **db_sources} - self._policies_by_id = {} non_production: Final = await _policy_table(prisma_client).find_many( where={"version_status": {"in": ["draft", "published"]}}, order={"created_at": "desc"}, ) - for row in non_production: - policy = self.parse_policy( - row.policy_name, - { - "inherit": row.inherit, - "description": row.description, - "guardrails": { - "add": row.guardrails_add or [], - "remove": row.guardrails_remove or [], - }, - "condition": row.condition, - "pipeline": row.pipeline, - }, - ) - self._policies_by_id[row.policy_id] = (row.policy_name, policy) + self._policies_by_id = _same_or_new( + self._policies_by_id, + {row.policy_id: (row.policy_name, self._parse_version_row(row)) for row in non_production}, + ) self._initialized = True verbose_proxy_logger.info( @@ -726,7 +755,7 @@ class PolicyRegistry: """ Resolve all guardrails for a policy from the database. - Uses the existing PolicyResolver to handle inheritance chain resolution. + Resolves the inheritance chain with the policy engine over the production versions. Args: policy_name: Name of the policy to resolve @@ -735,16 +764,14 @@ class PolicyRegistry: Returns: List of resolved guardrail names """ - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver + from litellm.rust_bridge.policy_engine import build_policy_engine try: # Load only production versions so inheritance resolves against production policies: Final = await self.get_all_policies_from_db(prisma_client, version_status="production") - # Build a temporary in-memory map for resolution - temp_policies: Final = {} - for policy_response in policies: - policy = self.parse_policy( + temp_policies: Final = { + policy_response.policy_name: self.parse_policy( policy_response.policy_name, { "inherit": policy_response.inherit, @@ -757,14 +784,10 @@ class PolicyRegistry: "pipeline": policy_response.pipeline, }, ) - temp_policies[policy_response.policy_name] = policy + for policy_response in policies + } - # Use the existing PolicyResolver to resolve guardrails - resolved_policy: Final = PolicyResolver.resolve_policy_guardrails( - policy_name=policy_name, - policies=temp_policies, - context=None, # No context needed for simple resolution - ) + resolved_policy: Final = build_policy_engine(temp_policies).policy_guardrails(policy_name) return sorted(resolved_policy.guardrails) except Exception as e: diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 898e42635c5..eac70a128da 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -243,9 +243,8 @@ async def resolve_policies_for_context( }' ``` """ - from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.proxy.proxy_server import prisma_client + from litellm.rust_bridge.policy_engine import current_policy_engine if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") @@ -264,49 +263,20 @@ async def resolve_policies_for_context( tags=request.tags, ) - # Get matching policies with reasons - match_results: Final = get_attachment_registry().get_attached_policies_with_reasons( - context=context, policy_applies=PolicyMatcher.policy_applies(context) + engine: Final = current_policy_engine() + resolution: Final = engine.resolve(context) + matched_policies: Final = tuple( + PolicyMatchDetail( + policy_name=policy_name, + matched_via=resolution.policy_sources[policy_name], + guardrails_added=engine.policy_guardrails(policy_name, context).guardrails, + ) + for policy_name in resolution.applied_policies ) - if not match_results: - return PolicyResolveResponse( - effective_guardrails=[], - matched_policies=[], - ) - - # Filter by conditions - policy_names: Final = [r["policy_name"] for r in match_results] - applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( - policy_names=policy_names, - context=context, - ) - - # Resolve guardrails for each applied policy - matched_policies: Final = [] - all_guardrails: Final[set] = set() - for result in match_results: - pname = result["policy_name"] - if pname not in applied_policy_names: - continue - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name=pname, - policies=get_policy_registry().get_all_policies(), - context=context, - ) - guardrails = resolved.guardrails if resolved else [] - all_guardrails.update(guardrails) - matched_policies.append( - PolicyMatchDetail( - policy_name=pname, - matched_via=result["matched_via"], - guardrails_added=guardrails, - ) - ) - return PolicyResolveResponse( - effective_guardrails=sorted(all_guardrails), - matched_policies=matched_policies, + effective_guardrails=sorted(resolution.guardrails), + matched_policies=list(matched_policies), ) except HTTPException: raise diff --git a/litellm/proxy/policy_engine/policy_resolver.py b/litellm/proxy/policy_engine/policy_resolver.py deleted file mode 100644 index e1422d79e15..00000000000 --- a/litellm/proxy/policy_engine/policy_resolver.py +++ /dev/null @@ -1,296 +0,0 @@ -""" -Policy Resolver - Resolves final guardrail list from policies. - -Handles: -- Inheritance chain resolution (inherit with add/remove) -- Applying add/remove guardrails -- Evaluating model conditions -- Combining guardrails from multiple matching policies -""" - -from typing import Final - -from litellm._logging import verbose_proxy_logger -from litellm.types.proxy.policy_engine import ( - GuardrailPipeline, - Policy, - PolicyMatchContext, - ResolvedPolicy, -) - - -class PolicyResolver: - """ - Resolves the final list of guardrails from policies. - - Handles: - - Inheritance chains with add/remove operations - - Model-based conditions - """ - - @staticmethod - def resolve_inheritance_chain( - policy_name: str, - policies: dict[str, Policy], - visited: set[str] | None = None, - ) -> list[str]: - """ - Get the inheritance chain for a policy (from root to policy). - - Args: - policy_name: Name of the policy - policies: Dictionary of all policies - visited: Set of visited policies (for cycle detection) - - Returns: - List of policy names from root ancestor to the given policy - """ - if visited is None: - visited = set() - - if policy_name in visited: - verbose_proxy_logger.warning("Circular inheritance detected for policy '%s'", policy_name) - return [] - - policy: Final = policies.get(policy_name) - if policy is None: - return [] - - visited.add(policy_name) - - if policy.inherit: - parent_chain: Final = PolicyResolver.resolve_inheritance_chain( - policy_name=policy.inherit, policies=policies, visited=visited - ) - return parent_chain + [policy_name] - - return [policy_name] - - @staticmethod - def resolve_policy_guardrails( - policy_name: str, - policies: dict[str, Policy], - context: PolicyMatchContext | None = None, - ) -> ResolvedPolicy: - """ - Resolve the final guardrails for a single policy, including inheritance. - - This method: - 1. Resolves the inheritance chain - 2. Applies add/remove from each policy in the chain - 3. Evaluates model conditions (if context provided) - - Args: - policy_name: Name of the policy to resolve - policies: Dictionary of all policies - context: Optional request context for evaluating conditions - - Returns: - ResolvedPolicy with final guardrails list - """ - from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator - - inheritance_chain: Final = PolicyResolver.resolve_inheritance_chain(policy_name=policy_name, policies=policies) - - # Start with empty set of guardrails - guardrails: Final[set[str]] = set() - - # Apply each policy in the chain (from root to leaf) - for chain_policy_name in inheritance_chain: - policy = policies.get(chain_policy_name) - if policy is None: - continue - - # Check if policy condition matches (if context provided) - if context is not None and policy.condition is not None: - if not ConditionEvaluator.evaluate( - condition=policy.condition, - context=context, - ): - verbose_proxy_logger.debug( - "Policy '%s' condition did not match, skipping guardrails", chain_policy_name - ) - continue - - # Add guardrails from guardrails.add - for guardrail in policy.guardrails.get_add(): - guardrails.add(guardrail) - - # Remove guardrails from guardrails.remove - for guardrail in policy.guardrails.get_remove(): - guardrails.discard(guardrail) - - return ResolvedPolicy( - policy_name=policy_name, - guardrails=list(guardrails), - inheritance_chain=inheritance_chain, - ) - - @staticmethod - def resolve_guardrails_for_context( - context: PolicyMatchContext, - policies: dict[str, Policy] | None = None, - policy_names: list[str] | None = None, - ) -> list[str]: - """ - Resolve the final list of guardrails for a request context. - - This: - 1. Finds all policies that match the context via policy_attachments (or policy_names if provided) - 2. Resolves each policy's guardrails (including inheritance) - 3. Evaluates model conditions - 4. Combines all guardrails (union) - - Args: - context: The request context - policies: Dictionary of all policies (if None, uses global registry) - policy_names: If provided, use this list instead of attachment matching - - Returns: - List of guardrail names to apply - """ - from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - from litellm.proxy.policy_engine.policy_registry import get_policy_registry - - if policies is None: - registry: Final = get_policy_registry() - if not registry.is_initialized(): - return [] - policies = registry.get_all_policies() - - # Use provided policy names or get matching policies via attachments - matching_policy_names: Final = ( - policy_names if policy_names is not None else PolicyMatcher.get_matching_policies(context=context) - ) - - if not matching_policy_names: - verbose_proxy_logger.debug( - "No policies match context: team_alias=%s, key_alias=%s, model=%s", - context.team_alias, - context.key_alias, - context.model, - ) - return [] - - # Resolve each matching policy and combine guardrails - all_guardrails: Final[set[str]] = set() - - for policy_name in matching_policy_names: - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name=policy_name, - policies=policies, - context=context, - ) - all_guardrails.update(resolved.guardrails) - verbose_proxy_logger.debug("Policy '%s' contributes guardrails: %s", policy_name, resolved.guardrails) - - result: Final = list(all_guardrails) - verbose_proxy_logger.debug("Final guardrails for context: %s", result) - - return result - - @staticmethod - def resolve_pipelines_for_context( - context: PolicyMatchContext, - policies: dict[str, Policy] | None = None, - policy_names: list[str] | None = None, - ) -> list[tuple[str, GuardrailPipeline]]: - """ - Resolve pipelines from matching policies for a request context. - - Returns (policy_name, pipeline) tuples for policies that have pipelines. - Guardrails managed by pipelines should be excluded from the flat - guardrails list to avoid double execution. - - Args: - context: The request context - policies: Dictionary of all policies (if None, uses global registry) - policy_names: If provided, use this list instead of attachment matching - - Returns: - List of (policy_name, GuardrailPipeline) tuples - """ - from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator - from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - from litellm.proxy.policy_engine.policy_registry import get_policy_registry - - if policies is None: - registry: Final = get_policy_registry() - if not registry.is_initialized(): - return [] - policies = registry.get_all_policies() - - matching_policy_names: Final = ( - policy_names if policy_names is not None else PolicyMatcher.get_matching_policies(context=context) - ) - if not matching_policy_names: - return [] - - pipelines: Final[list[tuple[str, GuardrailPipeline]]] = [] - for policy_name in matching_policy_names: - policy = policies.get(policy_name) - if policy is None: - continue - if policy.condition is not None and not ConditionEvaluator.evaluate( - condition=policy.condition, context=context - ): - verbose_proxy_logger.debug("Policy '%s' condition did not match, skipping pipeline", policy_name) - continue - if policy.pipeline is not None: - pipelines.append((policy_name, policy.pipeline)) - verbose_proxy_logger.debug( - "Policy '%s' has pipeline with %s steps", policy_name, len(policy.pipeline.steps) - ) - - return pipelines - - @staticmethod - def get_pipeline_managed_guardrails( - pipelines: list[tuple[str, GuardrailPipeline]], - ) -> set[str]: - """ - Get the set of guardrail names managed by pipelines. - - These guardrails should be excluded from normal independent execution. - """ - managed: Final[set[str]] = set() - for _policy_name, pipeline in pipelines: - for step in pipeline.steps: - managed.add(step.guardrail) - return managed - - @staticmethod - def get_all_resolved_policies( - policies: dict[str, Policy] | None = None, - context: PolicyMatchContext | None = None, - ) -> dict[str, ResolvedPolicy]: - """ - Resolve all policies and return their final guardrails. - - Useful for debugging and displaying policy configurations. - - Args: - policies: Dictionary of all policies (if None, uses global registry) - context: Optional context for evaluating conditions - - Returns: - Dictionary mapping policy names to ResolvedPolicy objects - """ - from litellm.proxy.policy_engine.policy_registry import get_policy_registry - - if policies is None: - registry: Final = get_policy_registry() - if not registry.is_initialized(): - return {} - policies = registry.get_all_policies() - - resolved: Final[dict[str, ResolvedPolicy]] = {} - - for policy_name in policies: - resolved[policy_name] = PolicyResolver.resolve_policy_guardrails( - policy_name=policy_name, - policies=policies, - context=context, - ) - - return resolved diff --git a/litellm/proxy/policy_engine/response_retrieval.py b/litellm/proxy/policy_engine/response_retrieval.py index 514f05d0bcd..88323562b98 100644 --- a/litellm/proxy/policy_engine/response_retrieval.py +++ b/litellm/proxy/policy_engine/response_retrieval.py @@ -1,6 +1,5 @@ from collections.abc import Mapping from dataclasses import dataclass -from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import TypeAdapter @@ -14,11 +13,9 @@ from litellm.proxy.common_utils.callback_utils import ( add_policy_to_applied_policies_header, ) from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body -from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry -from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher from litellm.proxy.policy_engine.policy_registry import get_policy_registry -from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.rust_bridge.policy_engine import current_policy_engine from litellm.types.proxy.policy_engine import PolicyMatchContext from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline @@ -85,23 +82,11 @@ def _retrieval_context( def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]: - matches: Final = get_attachment_registry().get_attached_policies_with_reasons( - context, PolicyMatcher.policy_applies(context) - ) - if not matches: - return (), MappingProxyType({}) - applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions( - policy_names=[match["policy_name"] for match in matches], - context=context, - ) + resolution: Final = current_policy_engine().resolve(context) post_call_pipelines: Final = tuple( - (policy_name, pipeline) - for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context( - context=context, policy_names=applied_policy_names - ) - if pipeline.mode == "post_call" + (policy_name, pipeline) for policy_name, pipeline in resolution.pipelines if pipeline.mode == "post_call" ) - return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches}) + return post_call_pipelines, resolution.policy_sources def attach_post_call_pipelines_to_retrieval( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 97c327cbe4e..0c80d01924b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -283,8 +283,6 @@ from litellm.proxy.policy_engine.pipeline_executor import ( pipeline_step_is_detect_only, recorded_guardrail_information, ) -from litellm.proxy.policy_engine.policy_registry import get_policy_registry -from litellm.proxy.policy_engine.policy_resolver import PolicyResolver from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.spend_log_insert import RejectedSpendLogRow @@ -993,9 +991,9 @@ def _is_pending_background_response(response: LLMResponseTypes) -> bool: def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]: - resolved: Final = PolicyResolver.resolve_policy_guardrails( - policy_name=policy_name, policies=get_policy_registry().get_all_policies() - ) + from litellm.rust_bridge.policy_engine import current_policy_engine + + resolved: Final = current_policy_engine().policy_guardrails(policy_name) return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 114d2df026e..0a59a682b0c 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -204,6 +204,63 @@ class _RateLimitAtomicCounter(Protocol): _RateLimitLocal = int | float | str | dict[str, float] +class _FullMatcher(Protocol): + def fullmatch(self, string: str, /) -> object: ... + +_PolicyRow = tuple[ + str | None, + Sequence[str], + Sequence[str], + Sequence[tuple[str, _FullMatcher | None]] | None, + Sequence[str] | None, +] +_AttachmentRow = tuple[str, bool, Sequence[str], Sequence[str], Sequence[str], Sequence[str], int | None, bool] +_PolicyContext = tuple[str | None, str | None, str | None, Sequence[str]] + +@final +class PolicyEngine: + def __new__( + cls, + policies: Sequence[tuple[str, _PolicyRow]], + versions: Sequence[tuple[str, str, _PolicyRow]], + attachments: Sequence[_AttachmentRow], + ) -> PolicyEngine: ... + def resolve( + self, context: _PolicyContext, names: Sequence[str], version_ids: Sequence[str] + ) -> PolicyResolution: ... + # (guardrails, inheritance chain); without a context every chain member counts. + def policy_guardrails( + self, name: str, context: _PolicyContext | None = None + ) -> tuple[list[str], list[str]]: ... + # ([(policy, matched_via)], [(policy, ancestor)] admitted through an ancestor). + def attached_policies( + self, context: _PolicyContext, only_applying: bool + ) -> tuple[list[tuple[str, str]], list[tuple[str, str]]]: ... + def circular_policies(self) -> list[str]: ... + +@final +class PolicyResolution: + @property + def attached_policies(self) -> list[str]: ... + @property + def unapplied_policies(self) -> list[str]: ... + @property + def condition_misses(self) -> list[str]: ... + @property + def policy_guardrails(self) -> list[tuple[str, list[str]]]: ... + @property + def applied_policies(self) -> list[str]: ... + @property + def policy_sources(self) -> list[tuple[str, str]]: ... + @property + def guardrails(self) -> list[str]: ... + @property + def pipelines(self) -> list[tuple[str, str | None]]: ... + @property + def pipeline_managed_guardrails(self) -> list[str]: ... + @property + def ancestor_admissions(self) -> list[tuple[str, str]]: ... + @final class WindowCount: def __new__(cls, window_start: str, count: int) -> WindowCount: ... diff --git a/litellm/rust_bridge/policy_engine.py b/litellm/rust_bridge/policy_engine.py new file mode 100644 index 00000000000..db53b8f0739 --- /dev/null +++ b/litellm/rust_bridge/policy_engine.py @@ -0,0 +1,330 @@ +"""The proxy policy engine, resolved in Rust over a snapshot of the policy and attachment registries. + +The engine for the live registries is rebuilt whenever a registry replaces one of its containers. +""" + +import re +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, TypeVar + +from litellm._logging import verbose_proxy_logger +from litellm.rust_bridge import _native +from litellm.types.proxy.policy_engine import ( + GuardrailPipeline, + Policy, + PolicyAttachment, + PolicyMatchContext, + ResolvedPolicy, +) + +if TYPE_CHECKING: + from litellm.proxy.policy_engine.attachment_registry import PolicyAttachmentMatch + + +class _FullMatcher(Protocol): + def fullmatch(self, string: str, /) -> object: ... + + +_PatternRow: TypeAlias = tuple[str, _FullMatcher | None] +_Names: TypeAlias = Sequence[str] +_PolicyRow: TypeAlias = tuple[str | None, _Names, _Names, Sequence[_PatternRow] | None, _Names | None] +_AttachmentRow: TypeAlias = tuple[str, bool, _Names, _Names, _Names, _Names, int | None, bool] +_ContextRow: TypeAlias = tuple[str | None, str | None, str | None, _Names] + +_NO_POLICIES: Final[Mapping[str, Policy]] = MappingProxyType({}) +_NO_VERSIONS: Final[Mapping[str, tuple[str, Policy]]] = MappingProxyType({}) + +_FIRST_SURROGATE: Final = 0xD800 +_FIRST_STAND_IN: Final = 0x10F800 +_TO_STAND_IN: Final = MappingProxyType( + {code: _FIRST_STAND_IN + code - _FIRST_SURROGATE for code in range(_FIRST_SURROGATE, 0xE000)} +) +_FROM_STAND_IN: Final = MappingProxyType({stand_in: code for code, stand_in in _TO_STAND_IN.items()}) +_NOT_UTF8: Final = (UnicodeEncodeError, TypeError) + + +@dataclass(frozen=True, slots=True) +class PolicyEngineResolution: + attached_policies: tuple[str, ...] + unapplied_policies: tuple[str, ...] + condition_misses: tuple[str, ...] + policy_guardrails: Mapping[str, tuple[str, ...]] + applied_policies: tuple[str, ...] + policy_sources: Mapping[str, str] + guardrails: tuple[str, ...] + pipelines: tuple[tuple[str, GuardrailPipeline], ...] + pipeline_managed_guardrails: frozenset[str] + + +def _same(text: str) -> str: + return text + + +def _stand_in(text: str) -> str: + return text.translate(_TO_STAND_IN) + + +def _restore(text: str) -> str: + return text.translate(_FROM_STAND_IN) + + +@dataclass(frozen=True, slots=True) +class _RestoringMatcher: + compiled: re.Pattern[str] + + def fullmatch(self, string: str, /) -> re.Match[str] | None: + return self.compiled.fullmatch(_restore(string)) + + +@dataclass(frozen=True, slots=True) +class _Text: + """How strings cross into Rust. Rust strings are UTF-8, so lone surrogates, which Python strings and YAML + allow, cross as the last 2048 code points of plane 16's private use area and are restored on the way back. + Text that already holds those private use code points reads as the surrogates they stand in for""" + + into_rust: Callable[[str], str] + from_rust: Callable[[str], str] + + def optional(self, text: str | None) -> str | None: + return self.into_rust(text) if text is not None else None + + def names(self, names: Iterable[str]) -> tuple[str, ...]: + return tuple(map(self.into_rust, names)) + + def restored(self, names: Iterable[str]) -> tuple[str, ...]: + return tuple(map(self.from_rust, names)) + + def pairs(self, pairs: Iterable[tuple[str, str]]) -> tuple[tuple[str, str], ...]: + return tuple((self.from_rust(first), self.from_rust(second)) for first, second in pairs) + + def condition_row(self, policy: Policy) -> tuple[_PatternRow, ...] | None: + model: Final = policy.condition.model if policy.condition is not None else None + if model is None: + return None + patterns: Final = (model,) if isinstance(model, str) else model + return tuple((self.into_rust(pattern), _matcher(pattern)) for pattern in patterns) + + def policy_row(self, policy: Policy) -> _PolicyRow: + return ( + self.optional(policy.inherit or None), + self.names(policy.guardrails.get_add()), + self.names(policy.guardrails.get_remove()), + self.condition_row(policy), + self.names(step.guardrail for step in policy.pipeline.steps) if policy.pipeline is not None else None, + ) + + def attachment_row(self, attachment: PolicyAttachment) -> _AttachmentRow: + return ( + self.into_rust(attachment.policy), + attachment.is_global(), + self.names(attachment.teams or ()), + self.names(attachment.keys or ()), + self.names(attachment.models or ()), + self.names(attachment.tags or ()), + attachment.priority, + attachment.default, + ) + + def context_row(self, context: PolicyMatchContext) -> _ContextRow: + return ( + self.optional(context.team_alias), + self.optional(context.key_alias), + self.optional(context.model), + self.names(context.tags or ()), + ) + + def native_engine( + self, + policies: Mapping[str, Policy], + versions: Mapping[str, tuple[str, Policy]], + attachments: Sequence[PolicyAttachment], + ) -> _native.PolicyEngine: + return _native.PolicyEngine( + tuple((self.into_rust(name), self.policy_row(policy)) for name, policy in policies.items()), + tuple( + (self.into_rust(policy_id), self.into_rust(name), self.policy_row(policy)) + for policy_id, (name, policy) in versions.items() + ), + tuple(self.attachment_row(attachment) for attachment in attachments), + ) + + +_AS_IS: Final = _Text(_same, _same) +_STOOD_IN: Final = _Text(_stand_in, _restore) + + +def _compiled(pattern: str) -> re.Pattern[str] | None: + try: + return re.compile(pattern) + except (re.error, OverflowError, RecursionError, ValueError): + return None + + +def _matcher(pattern: str) -> _FullMatcher | None: + compiled: Final = _compiled(pattern) + return _RestoringMatcher(compiled) if compiled is not None else None + + +def _log_ancestor_admissions(admissions: Sequence[tuple[str, str]], context: PolicyMatchContext) -> None: + for policy_name, ancestor in admissions: + verbose_proxy_logger.info( + "Policy '%s' applied through ancestor '%s' although its own condition did not match " + "(team_alias=%s, key_alias=%s, model=%s)", + policy_name, + ancestor, + context.team_alias, + context.key_alias, + context.model, + ) + + +def _pipeline(policy: Policy) -> GuardrailPipeline: + pipeline: Final = policy.pipeline + assert pipeline is not None, "Rust only reports pipelines of policies that have one" + return pipeline + + +_Result: Final = TypeVar("_Result") + + +@dataclass(frozen=True, slots=True) +class PolicyEngine: + policies: Mapping[str, Policy] + versions: Mapping[str, tuple[str, Policy]] + attachments: Sequence[PolicyAttachment] + native: _native.PolicyEngine + text: _Text = _AS_IS + + def _call(self, call: Callable[[_Text], _Result]) -> tuple[_Text, _Result]: + if self.text is _STOOD_IN: + return _STOOD_IN, call(_STOOD_IN) + try: + return _AS_IS, call(_AS_IS) + except _NOT_UTF8: + return _STOOD_IN, call(_STOOD_IN) + + def resolve( + self, + context: PolicyMatchContext, + names: Sequence[str] = (), + version_ids: Sequence[str] = (), + ) -> PolicyEngineResolution: + """The policies a request applies, from its attachments plus the names and version ids it sent""" + text, resolution = self._call( + lambda text: self.native.resolve(text.context_row(context), text.names(names), text.names(version_ids)) + ) + _log_ancestor_admissions(text.pairs(resolution.ancestor_admissions), context) + return PolicyEngineResolution( + attached_policies=text.restored(resolution.attached_policies), + unapplied_policies=text.restored(resolution.unapplied_policies), + condition_misses=text.restored(resolution.condition_misses), + policy_guardrails=MappingProxyType( + {text.from_rust(name): text.restored(added) for name, added in resolution.policy_guardrails} + ), + applied_policies=text.restored(resolution.applied_policies), + policy_sources=MappingProxyType(dict(text.pairs(resolution.policy_sources))), + guardrails=text.restored(resolution.guardrails), + pipelines=tuple( + self._pipeline_of(text.from_rust(name), version_id, text) for name, version_id in resolution.pipelines + ), + pipeline_managed_guardrails=frozenset(text.restored(resolution.pipeline_managed_guardrails)), + ) + + def _pipeline_of(self, name: str, version_id: str | None, text: _Text) -> tuple[str, GuardrailPipeline]: + policy: Final = self.versions[text.from_rust(version_id)][1] if version_id is not None else self.policies[name] + return name, _pipeline(policy) + + def policy_guardrails(self, policy_name: str, context: PolicyMatchContext | None = None) -> ResolvedPolicy: + """One policy's guardrails through its inheritance chain, skipping members whose condition misses the + context. Without a context every member counts""" + text, (guardrails, chain) = self._call( + lambda text: self.native.policy_guardrails( + text.into_rust(policy_name), text.context_row(context) if context is not None else None + ) + ) + return ResolvedPolicy( + policy_name=policy_name, + guardrails=list(text.restored(guardrails)), + inheritance_chain=list(text.restored(chain)), + ) + + def attached_policies( + self, context: PolicyMatchContext, *, only_applying: bool + ) -> tuple["PolicyAttachmentMatch", ...]: + """The policies attached to the context, one per policy with the reason it matched. With + ``only_applying``, attachments whose policy does not apply are dropped before defaults are considered""" + text, (matches, admissions) = self._call( + lambda text: self.native.attached_policies(text.context_row(context), only_applying) + ) + _log_ancestor_admissions(text.pairs(admissions), context) + return tuple( + {"policy_name": policy_name, "matched_via": matched_via} for policy_name, matched_via in text.pairs(matches) + ) + + +def build_policy_engine( + policies: Mapping[str, Policy], + versions: Mapping[str, tuple[str, Policy]] = _NO_VERSIONS, + attachments: Sequence[PolicyAttachment] = (), +) -> PolicyEngine: + engine: Final = _built_policy_engine(policies, versions, attachments) + for policy_name in engine.text.restored(engine.native.circular_policies()): + verbose_proxy_logger.warning("Circular inheritance detected for policy '%s'", policy_name) + return engine + + +def _built_policy_engine( + policies: Mapping[str, Policy], + versions: Mapping[str, tuple[str, Policy]], + attachments: Sequence[PolicyAttachment], +) -> PolicyEngine: + try: + return PolicyEngine(policies, versions, attachments, _AS_IS.native_engine(policies, versions, attachments)) + except _NOT_UTF8: + return PolicyEngine( + policies, versions, attachments, _STOOD_IN.native_engine(policies, versions, attachments), _STOOD_IN + ) + + +class _CurrentPolicyEngine: + """Holds the engine for the registries' current containers""" + + def __init__(self) -> None: + self._engine: PolicyEngine | None = None + + def engine_for( + self, + policies: Mapping[str, Policy], + versions: Mapping[str, tuple[str, Policy]], + attachments: Sequence[PolicyAttachment], + ) -> PolicyEngine: + cached: Final = self._engine + if ( + cached is not None + and cached.policies is policies + and cached.versions is versions + and cached.attachments is attachments + ): + return cached + engine: Final = build_policy_engine(policies, versions, attachments) + self._engine = engine + return engine + + +_CURRENT: Final = _CurrentPolicyEngine() + + +def current_policy_engine(attachments: Sequence[PolicyAttachment] | None = None) -> PolicyEngine: + """The engine over the live policy registry and the given attachments, the live attachment registry's + by default""" + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + policy_registry: Final = get_policy_registry() + return _CURRENT.engine_for( + policy_registry.current_policies() if policy_registry.is_initialized() else _NO_POLICIES, + policy_registry.current_versions(), + attachments if attachments is not None else get_attachment_registry().current_attachments(), + ) diff --git a/tests/unit/proxy/management_endpoints/test_policy_endpoints.py b/tests/unit/proxy/management_endpoints/test_policy_endpoints.py index e334305dd98..8d19b0ae447 100644 --- a/tests/unit/proxy/management_endpoints/test_policy_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_policy_endpoints.py @@ -129,7 +129,7 @@ class TestApplyPoliciesEarlyReturn: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", return_value=ResolvedPolicy( policy_name="p", guardrails=[], inheritance_chain=[] ), @@ -177,7 +177,7 @@ class TestApplyPoliciesWithGuardrails: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", return_value=ResolvedPolicy( policy_name="p", guardrails=["my_guardrail"], @@ -236,7 +236,7 @@ class TestApplyPoliciesWithGuardrails: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", return_value=ResolvedPolicy( policy_name="p", guardrails=["guardrail_a", "guardrail_b"], @@ -278,7 +278,7 @@ class TestApplyPoliciesWithGuardrails: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", return_value=ResolvedPolicy( policy_name="p", guardrails=["missing_guardrail"], @@ -330,7 +330,7 @@ class TestApplyPoliciesWithGuardrails: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", return_value=ResolvedPolicy( policy_name="p", guardrails=["failing_guardrail"], @@ -388,7 +388,7 @@ class TestApplyPoliciesWithGuardrails: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", return_value=ResolvedPolicy( policy_name="p", guardrails=["no_apply_guardrail"], @@ -454,7 +454,7 @@ class TestApplyPoliciesWithGuardrails: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", return_value=ResolvedPolicy( policy_name="p", guardrails=["guardrail_a", "guardrail_b"], @@ -524,7 +524,7 @@ class TestApplyPoliciesMultiplePolicies: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", side_effect=resolve_returns, ), patch( @@ -618,7 +618,7 @@ class TestApplyPoliciesDirectGuardrailNames: return_value=mock_registry, ), patch( - "litellm.proxy.management_endpoints.policy_endpoints.endpoints.PolicyResolver.resolve_policy_guardrails", + "litellm.rust_bridge.policy_engine.PolicyEngine.policy_guardrails", return_value=ResolvedPolicy( policy_name="p", guardrails=["from_policy"], diff --git a/tests/unit/proxy/policy_engine/test_attachment_registry.py b/tests/unit/proxy/policy_engine/test_attachment_registry.py index aefb73eea70..4a47a020490 100644 --- a/tests/unit/proxy/policy_engine/test_attachment_registry.py +++ b/tests/unit/proxy/policy_engine/test_attachment_registry.py @@ -5,7 +5,7 @@ Tests the main entry point: get_attached_policies() """ import sys -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from types import FrameType from typing import Final, Protocol @@ -13,11 +13,12 @@ from unittest.mock import AsyncMock, MagicMock import pytest +import litellm.proxy.policy_engine.policy_registry as policy_registry_module from litellm.proxy.policy_engine.attachment_registry import ( AttachmentRegistry, get_attachment_registry, ) -from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.proxy.policy_engine.policy_registry import PolicyRegistry from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails, PolicyMatchContext @@ -532,6 +533,13 @@ class TestMatchAttribution: assert "catch-all" in attached +def _install_policies(monkeypatch: pytest.MonkeyPatch, policies: Mapping[str, Policy]) -> None: + registry: Final = PolicyRegistry() + for name, policy in policies.items(): + registry.add_policy(name, policy) + monkeypatch.setattr(policy_registry_module, "get_policy_registry", lambda: registry) + + class TestDefaultAttachments: """`default: true` attachments apply only when no non-default attachment matches.""" @@ -598,37 +606,34 @@ class TestDefaultAttachments: assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}] - def test_inapplicable_opt_in_policy_does_not_suppress_default(self): + def test_inapplicable_opt_in_policy_does_not_suppress_default(self, monkeypatch: pytest.MonkeyPatch): context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) policies = { "guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])), "guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="claude.*")), } + _install_policies(monkeypatch, policies) - results = self._registry().get_attached_policies_with_reasons( - context, PolicyMatcher.policy_applies(context, policies) - ) + results = self._registry().get_attached_policies_with_reasons(context, only_applying=True) assert results == [{"policy_name": "guardrail-y", "matched_via": "default:scope:*"}] - def test_attachment_to_missing_policy_does_not_suppress_default(self): + def test_attachment_to_missing_policy_does_not_suppress_default(self, monkeypatch: pytest.MonkeyPatch): context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) policies = {"guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"]))} + _install_policies(monkeypatch, policies) - assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [ - "guardrail-y" - ] + assert self._registry().get_attached_policies(context, only_applying=True) == ["guardrail-y"] - def test_applicable_opt_in_policy_still_wins_with_predicate(self): + def test_applicable_opt_in_policy_still_wins_with_predicate(self, monkeypatch: pytest.MonkeyPatch): context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.2", tags=["opt-in"]) policies = { "guardrail-y": Policy(guardrails=PolicyGuardrails(add=["y"])), "guardrail-x": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="gpt.*")), } + _install_policies(monkeypatch, policies) - assert self._registry().get_attached_policies(context, PolicyMatcher.policy_applies(context, policies)) == [ - "guardrail-x" - ] + assert self._registry().get_attached_policies(context, only_applying=True) == ["guardrail-x"] def test_default_defaults_to_false_when_omitted(self): registry = AttachmentRegistry() @@ -748,6 +753,16 @@ class TestConfigAttachmentsPreservedAcrossDbSync: assert registry.get_all_attachments()[0].default is True + @pytest.mark.asyncio + async def test_an_unchanged_sync_keeps_the_attachment_container(self): + registry = AttachmentRegistry() + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([_make_db_attachment_row()])) + before = registry.current_attachments() + + await registry.sync_attachments_from_db(_prisma_with_attachment_rows([_make_db_attachment_row()])) + + assert registry.current_attachments() is before + @pytest.mark.asyncio async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self): registry = AttachmentRegistry() diff --git a/tests/unit/proxy/policy_engine/test_condition_evaluator.py b/tests/unit/proxy/policy_engine/test_condition_evaluator.py deleted file mode 100644 index 3c418c7e50b..00000000000 --- a/tests/unit/proxy/policy_engine/test_condition_evaluator.py +++ /dev/null @@ -1,145 +0,0 @@ -""" -Unit tests for ConditionEvaluator - tests model condition evaluation. - -Tests: -- Exact model match -- Regex pattern match -- List of models -""" - -import pytest - -from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator -from litellm.types.proxy.policy_engine import ( - PolicyCondition, - PolicyMatchContext, -) - - -class TestConditionEvaluator: - """Test condition evaluation.""" - - def test_no_condition_always_matches(self): - """Test that None condition always matches.""" - context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") - assert ConditionEvaluator.evaluate(None, context) is True - - def test_exact_model_match(self): - """Test exact model string match.""" - condition = PolicyCondition(model="gpt-4") - - # Match - context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") - assert ConditionEvaluator.evaluate(condition, context) is True - - # No match - context_other = PolicyMatchContext( - team_alias="team", key_alias="key", model="gpt-3.5" - ) - assert ConditionEvaluator.evaluate(condition, context_other) is False - - def test_regex_pattern_match(self): - """Test regex pattern matching.""" - condition = PolicyCondition(model="gpt-4.*") - - # Matches - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4"), - ) - is True - ) - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4-turbo"), - ) - is True - ) - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4o"), - ) - is True - ) - - # No match - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-3.5"), - ) - is False - ) - - def test_list_of_models_match(self): - """Test list of model values.""" - condition = PolicyCondition(model=["gpt-4", "gpt-4-turbo", "claude-3"]) - - # Matches - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4"), - ) - is True - ) - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="claude-3"), - ) - is True - ) - - # No match - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-3.5"), - ) - is False - ) - - def test_list_with_regex_patterns(self): - """Test list can contain regex patterns.""" - condition = PolicyCondition(model=["gpt-4.*", "claude-.*"]) - - # Matches - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4-turbo"), - ) - is True - ) - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="claude-3"), - ) - is True - ) - - # No match - assert ( - ConditionEvaluator.evaluate( - condition, - PolicyMatchContext(team_alias="t", key_alias="k", model="llama-2"), - ) - is False - ) - - def test_none_model_does_not_match(self): - """Test that None model value doesn't match conditions.""" - condition = PolicyCondition(model="gpt-4") - context = PolicyMatchContext(team_alias="t", key_alias="k", model=None) - assert ConditionEvaluator.evaluate(condition, context) is False - - def test_empty_condition_always_matches(self): - """Test condition with no model field always matches.""" - condition = PolicyCondition() # No model specified - context = PolicyMatchContext(team_alias="t", key_alias="k", model="any-model") - assert ConditionEvaluator.evaluate(condition, context) is True diff --git a/tests/unit/proxy/policy_engine/test_policy_engine.py b/tests/unit/proxy/policy_engine/test_policy_engine.py new file mode 100644 index 00000000000..18ea52ff3ef --- /dev/null +++ b/tests/unit/proxy/policy_engine/test_policy_engine.py @@ -0,0 +1,559 @@ +import logging +from collections.abc import Mapping, Sequence +from typing import Final + +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +import litellm.proxy.policy_engine.attachment_registry as attachment_registry_module +import litellm.proxy.policy_engine.policy_registry as policy_registry_module +from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry +from litellm.proxy.policy_engine.policy_registry import PolicyRegistry +from litellm.rust_bridge.policy_engine import PolicyEngine, build_policy_engine, current_policy_engine +from litellm.types.proxy.policy_engine import ( + GuardrailPipeline, + PipelineStep, + Policy, + PolicyAttachment, + PolicyCondition, + PolicyGuardrails, + PolicyMatchContext, +) + + +def _context( + model: str | None = "gpt-4", team: str | None = "t", key: str | None = "k", tags: Sequence[str] | None = None +) -> PolicyMatchContext: + return PolicyMatchContext( + team_alias=team, key_alias=key, model=model, tags=list(tags) if tags is not None else None + ) + + +def _condition_applies(condition: PolicyCondition | None, model: str | None) -> bool: + policy: Final = Policy(guardrails=PolicyGuardrails(add=["g"]), condition=condition) + return build_policy_engine({"p": policy}).policy_guardrails("p", _context(model)).guardrails == ["g"] + + +class TestModelCondition: + def test_no_condition_always_matches(self) -> None: + assert _condition_applies(None, "gpt-4") is True + + def test_empty_condition_always_matches(self) -> None: + assert _condition_applies(PolicyCondition(), "any-model") is True + + @pytest.mark.parametrize(("model", "expected"), [("gpt-4", True), ("gpt-3.5", False)]) + def test_exact_model(self, model: str, expected: bool) -> None: + assert _condition_applies(PolicyCondition(model="gpt-4"), model) is expected + + @pytest.mark.parametrize( + ("model", "expected"), [("gpt-4", True), ("gpt-4-turbo", True), ("gpt-4o", True), ("gpt-3.5", False)] + ) + def test_regex_pattern(self, model: str, expected: bool) -> None: + assert _condition_applies(PolicyCondition(model="gpt-4.*"), model) is expected + + @pytest.mark.parametrize(("model", "expected"), [("gpt-4", True), ("claude-3", True), ("gpt-3.5", False)]) + def test_list_of_models(self, model: str, expected: bool) -> None: + assert _condition_applies(PolicyCondition(model=["gpt-4", "gpt-4-turbo", "claude-3"]), model) is expected + + @pytest.mark.parametrize(("model", "expected"), [("gpt-4-turbo", True), ("claude-3", True), ("llama-2", False)]) + def test_list_of_regex_patterns(self, model: str, expected: bool) -> None: + assert _condition_applies(PolicyCondition(model=["gpt-4.*", "claude-.*"]), model) is expected + + def test_none_model_does_not_match(self) -> None: + assert _condition_applies(PolicyCondition(model="gpt-4"), None) is False + + @pytest.mark.parametrize( + ("pattern", "model", "expected"), + [ + pytest.param("[[]", "[", True, id="syntax-only-python-reads"), + pytest.param("[a&&b]", "&", True, id="class-python-reads-literally"), + pytest.param("[a&&b]", "c", False, id="class-python-reads-literally-miss"), + pytest.param("gpt-4(", "gpt-4(", True, id="invalid-regex-still-matches-literally"), + ], + ) + def test_patterns_follow_python_re(self, pattern: str, model: str, expected: bool) -> None: + assert _condition_applies(PolicyCondition(model=pattern), model) is expected + + +def _attached(attachment: PolicyAttachment, context: PolicyMatchContext) -> bool: + engine: Final = build_policy_engine({"p": Policy()}, attachments=(attachment,)) + return [match["policy_name"] for match in engine.attached_policies(context, only_applying=False)] == ["p"] + + +def _scoped( + teams: Sequence[str] = ("*",), + keys: Sequence[str] = ("*",), + models: Sequence[str] = ("*",), + tags: Sequence[str] | None = None, +) -> PolicyAttachment: + return PolicyAttachment( + policy="p", teams=list(teams), keys=list(keys), models=list(models), tags=list(tags) if tags else None + ) + + +class TestAttachmentScope: + @pytest.mark.parametrize( + ("value", "patterns", "expected"), + [ + ("healthcare-team", ("healthcare-team",), True), + ("finance-team", ("healthcare-team",), False), + ("any-team", ("*",), True), + ("dev-key-123", ("dev-key-*",), True), + ("prod-key-123", ("dev-key-*",), False), + (None, ("*",), True), + (None, ("specific",), False), + ], + ) + def test_team_patterns(self, value: str | None, patterns: Sequence[str], expected: bool) -> None: + assert _attached(_scoped(teams=patterns), _context(team=value)) is expected + + def test_all_fields_match(self) -> None: + context: Final = _context(team="healthcare-team", key="any-key", model="gpt-4") + assert _attached(_scoped(teams=["healthcare-team"], models=["gpt-4"]), context) is True + + def test_team_mismatch(self) -> None: + context: Final = _context(team="finance-team", key="any-key", model="gpt-4") + assert _attached(_scoped(teams=["healthcare-team"]), context) is False + + def test_wildcard_patterns(self) -> None: + context: Final = _context(team="any-team", key="dev-key-123", model="bedrock/claude-3") + assert _attached(_scoped(keys=["dev-key-*"], models=["bedrock/*"]), context) is True + + def test_global_wildcard(self) -> None: + assert _attached(_scoped(), _context(team="any-team", key="any-key", model="any-model")) is True + + @pytest.mark.parametrize( + ("scope_tags", "context_tags", "expected"), + [ + (["healthcare"], ["healthcare", "internal"], True), + (["health-*"], ["health-prod"], True), + (["healthcare"], ["finance"], False), + (["healthcare"], None, False), + (None, ["healthcare", "internal"], True), + ], + ) + def test_tags(self, scope_tags: Sequence[str] | None, context_tags: Sequence[str] | None, expected: bool) -> None: + assert _attached(_scoped(tags=scope_tags), _context(tags=context_tags)) is expected + + @pytest.mark.parametrize( + ("team", "tags", "expected"), + [("team-a", ["healthcare"], True), ("team-b", ["healthcare"], False), ("team-a", ["finance"], False)], + ) + def test_tags_and_team_both_required(self, team: str, tags: Sequence[str], expected: bool) -> None: + assert _attached(_scoped(teams=["team-a"], tags=["healthcare"]), _context(team=team, tags=tags)) is expected + + +class TestAttachmentRegistryMatching: + def test_team_and_global_attachments_both_match(self) -> None: + registry: Final = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "healthcare-policy", "teams": ["healthcare-team"]}, + {"policy": "global-policy", "scope": "*"}, + ] + ) + attached: Final = registry.get_attached_policies(_context(team="healthcare-team")) + assert "healthcare-policy" in attached + assert "global-policy" in attached + + def test_no_match(self) -> None: + registry: Final = AttachmentRegistry() + registry.load_attachments([{"policy": "healthcare-policy", "teams": ["healthcare-team"]}]) + assert "healthcare-policy" not in registry.get_attached_policies(_context(team="finance-team")) + + +class TestInheritance: + def test_simple_policy(self) -> None: + engine: Final = build_policy_engine( + {"global": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "toxicity_filter"]))} + ) + resolved: Final = engine.policy_guardrails("global") + assert set(resolved.guardrails) == {"pii_blocker", "toxicity_filter"} + assert resolved.inheritance_chain == ["global"] + + def test_child_inherits_and_adds(self) -> None: + engine: Final = build_policy_engine( + { + "base": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"])), + "healthcare": Policy(inherit="base", guardrails=PolicyGuardrails(add=["hipaa_audit"])), + } + ) + resolved: Final = engine.policy_guardrails("healthcare") + assert set(resolved.guardrails) == {"pii_blocker", "hipaa_audit"} + assert resolved.inheritance_chain == ["base", "healthcare"] + + def test_child_removes(self) -> None: + engine: Final = build_policy_engine( + { + "base": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "phi_blocker"])), + "dev": Policy( + inherit="base", guardrails=PolicyGuardrails(add=["toxicity_filter"], remove=["phi_blocker"]) + ), + } + ) + assert set(engine.policy_guardrails("dev").guardrails) == {"pii_blocker", "toxicity_filter"} + + def test_deep_chain(self) -> None: + engine: Final = build_policy_engine( + { + "root": Policy(guardrails=PolicyGuardrails(add=["root_guardrail"])), + "middle": Policy(inherit="root", guardrails=PolicyGuardrails(add=["middle_guardrail"])), + "leaf": Policy(inherit="middle", guardrails=PolicyGuardrails(add=["leaf_guardrail"])), + } + ) + resolved: Final = engine.policy_guardrails("leaf") + assert set(resolved.guardrails) == {"root_guardrail", "middle_guardrail", "leaf_guardrail"} + assert resolved.inheritance_chain == ["root", "middle", "leaf"] + + +class TestConditions: + @pytest.mark.parametrize(("model", "expected"), [("gpt-4", ["toxicity_filter"]), ("gpt-3.5", [])]) + def test_own_condition(self, model: str, expected: Sequence[str]) -> None: + engine: Final = build_policy_engine( + { + "gpt4-policy": Policy( + guardrails=PolicyGuardrails(add=["toxicity_filter"]), condition=PolicyCondition(model="gpt-4.*") + ) + } + ) + assert engine.policy_guardrails("gpt4-policy", _context(model)).guardrails == expected + + @pytest.mark.parametrize( + ("model", "expected"), [("gpt-4", {"pii_blocker", "child_guardrail"}), ("gpt-3.5", {"pii_blocker"})] + ) + def test_inherited_condition(self, model: str, expected: set[str]) -> None: + engine: Final = build_policy_engine( + { + "base": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"])), + "child": Policy( + inherit="base", + guardrails=PolicyGuardrails(add=["child_guardrail"]), + condition=PolicyCondition(model="gpt-4"), + ), + } + ) + assert set(engine.policy_guardrails("child", _context(model)).guardrails) == expected + + @pytest.mark.parametrize(("model", "expected"), [("gpt-5.5", {"y"}), ("claude-haiku", {"x", "y"})]) + def test_condition_missing_child_keeps_inherited_parent(self, model: str, expected: set[str]) -> None: + engine: Final = build_policy_engine( + { + "parent": Policy(guardrails=PolicyGuardrails(add=["y"])), + "child": Policy( + inherit="parent", + guardrails=PolicyGuardrails(add=["x"]), + condition=PolicyCondition(model="claude.*"), + ), + } + ) + assert set(engine.resolve(_context(model), names=["child"]).guardrails) == expected + + def test_pipeline_skipped_when_own_condition_misses(self) -> None: + pipeline: Final = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="child-guard")]) + engine: Final = build_policy_engine( + { + "parent": Policy(guardrails=PolicyGuardrails(add=["y"])), + "child": Policy(inherit="parent", pipeline=pipeline, condition=PolicyCondition(model="gpt-5.5")), + } + ) + assert engine.resolve(_context("gpt-4o"), names=["child"]).pipelines == () + assert engine.resolve(_context("gpt-5.5"), names=["child"]).pipelines == (("child", pipeline),) + + def test_unknown_policy_name_applies_nothing(self) -> None: + engine: Final = build_policy_engine( + {"real": Policy(guardrails=PolicyGuardrails(add=["g"]), condition=PolicyCondition(model="claude.*"))} + ) + assert engine.resolve(_context("gpt-5.5"), names=["nope"]).applied_policies == () + + +def _install( + monkeypatch: pytest.MonkeyPatch, + policies_config: Mapping[str, Mapping[str, object]], + attachments_config: Sequence[Mapping[str, object]], +) -> None: + policies: Final = PolicyRegistry() + policies.load_policies({name: dict(config) for name, config in policies_config.items()}) + attachments: Final = AttachmentRegistry() + attachments.load_attachments([dict(attachment) for attachment in attachments_config]) + monkeypatch.setattr(policy_registry_module, "get_policy_registry", lambda: policies) + monkeypatch.setattr(attachment_registry_module, "get_attachment_registry", lambda: attachments) + + +def _matching(context: PolicyMatchContext) -> list[str]: + return [match["policy_name"] for match in current_policy_engine().attached_policies(context, only_applying=True)] + + +def _install_opt_in(monkeypatch: pytest.MonkeyPatch) -> None: + _install( + monkeypatch, + { + "guardrail-y": {"guardrails": {"add": ["y"]}}, + "guardrail-x": {"guardrails": {"add": ["x"]}, "condition": {"model": "claude.*"}}, + }, + [{"policy": "guardrail-x", "tags": ["opt-in"]}, {"policy": "guardrail-y", "scope": "*", "default": True}], + ) + + +def _install_inherited(monkeypatch: pytest.MonkeyPatch, parent_condition: Mapping[str, object] | None = None) -> None: + _install( + monkeypatch, + { + "parent": {"guardrails": {"add": ["y"]}, **({"condition": parent_condition} if parent_condition else {})}, + "child": {"inherit": "parent", "guardrails": {"add": ["x"]}, "condition": {"model": "claude.*"}}, + "fallback": {"guardrails": {"add": ["z"]}}, + }, + [{"policy": "child", "scope": "*"}, {"policy": "fallback", "scope": "*", "default": True}], + ) + + +class TestDefaultFallback: + def test_condition_failing_opt_in_falls_back_to_default(self, monkeypatch: pytest.MonkeyPatch) -> None: + _install_opt_in(monkeypatch) + assert _matching(_context("gpt-5.5", tags=["opt-in"])) == ["guardrail-y"] + + def test_condition_passing_opt_in_suppresses_default(self, monkeypatch: pytest.MonkeyPatch) -> None: + _install_opt_in(monkeypatch) + assert _matching(_context("claude-haiku", tags=["opt-in"])) == ["guardrail-x"] + + def test_condition_missing_child_with_unconditional_parent_still_matches( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + _install_inherited(monkeypatch) + assert _matching(_context("gpt-5.5")) == ["child"] + + def test_child_whose_whole_chain_misses_falls_back_to_default(self, monkeypatch: pytest.MonkeyPatch) -> None: + _install_inherited(monkeypatch, parent_condition={"model": "claude.*"}) + assert _matching(_context("gpt-5.5")) == ["fallback"] + + +_MODELS: Final = ("gpt-4o", "gpt-5.5", "claude-opus-4-1") + + +@st.composite +def _forest_and_request(draw: st.DrawFn) -> tuple[Mapping[str, Policy], tuple[str, ...], PolicyMatchContext]: + names: Final = tuple(f"p{i}" for i in range(draw(st.integers(min_value=1, max_value=6)))) + policies: Final = { + name: Policy( + inherit=draw(st.sampled_from((None, *names[:i]))), + guardrails=PolicyGuardrails(add=[f"g-{name}"]), + condition=draw(st.sampled_from((None, *(PolicyCondition(model=m) for m in _MODELS)))), + ) + for i, name in enumerate(names) + } + attached: Final = tuple(draw(st.lists(st.sampled_from(sorted(policies)), unique=True))) + return policies, attached, _context(draw(st.sampled_from(_MODELS))) + + +def _own_condition_applies(policy: Policy, context: PolicyMatchContext) -> bool: + return policy.condition is None or policy.condition.model == context.model + + +def _applicable_chain( + engine: PolicyEngine, policies: Mapping[str, Policy], name: str, context: PolicyMatchContext +) -> tuple[str, ...]: + chain: Final = engine.policy_guardrails(name).inheritance_chain + return tuple(member for member in chain if _own_condition_applies(policies[member], context)) + + +class TestChainMatchingProperties: + @given(_forest_and_request()) + @settings(max_examples=400, deadline=None) + def test_chain_matching_only_widens_to_applicable_ancestor_guardrails( + self, case: tuple[Mapping[str, Policy], tuple[str, ...], PolicyMatchContext] + ) -> None: + policies, attached, context = case + engine: Final = build_policy_engine(policies) + head: Final = engine.resolve(context, names=attached).applied_policies + base: Final = tuple(name for name in attached if _own_condition_applies(policies[name], context)) + expected_head: Final = tuple(name for name in attached if _applicable_chain(engine, policies, name, context)) + + assert head == expected_head, "a policy applies exactly when some chain member's own condition applies" + assert frozenset(base) <= frozenset(head), "head must never drop a policy base applied" + + for name in head: + resolved = engine.policy_guardrails(name, context) + assert sorted(resolved.guardrails) == sorted( + f"g-{member}" for member in _applicable_chain(engine, policies, name, context) + ) + if name not in base: + assert f"g-{name}" not in resolved.guardrails, "a condition-missed child must not add its own guardrail" + + +def _ancestor_logs(caplog: pytest.LogCaptureFixture) -> list[str]: + return [r.getMessage() for r in caplog.records if "applied through ancestor" in r.getMessage()] + + +class TestAncestorAdmissionLogging: + @staticmethod + def _attached_engine(child_condition: str, parent_condition: str | None = None) -> PolicyEngine: + return build_policy_engine( + { + "parent": Policy( + guardrails=PolicyGuardrails(add=["g-parent"]), + condition=PolicyCondition(model=parent_condition) if parent_condition else None, + ), + "child": Policy( + inherit="parent", + guardrails=PolicyGuardrails(add=["g-child"]), + condition=PolicyCondition(model=child_condition), + ), + }, + attachments=(PolicyAttachment(policy="child", scope="*"),), + ) + + def test_logs_when_admitted_through_ancestor_only(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + matches: Final = self._attached_engine("gpt-5.5").attached_policies(_context("gpt-4o"), only_applying=True) + logs: Final = _ancestor_logs(caplog) + assert [match["policy_name"] for match in matches] == ["child"] + assert len(logs) == 1 + assert "applied through ancestor 'parent'" in logs[0] + assert "'child'" in logs[0] + + def test_no_log_when_own_condition_matches(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + matches: Final = self._attached_engine("gpt-5.5").attached_policies(_context("gpt-5.5"), only_applying=True) + assert [match["policy_name"] for match in matches] == ["child"] + assert not _ancestor_logs(caplog) + + def test_no_log_when_no_chain_member_applies(self, caplog: pytest.LogCaptureFixture) -> None: + engine: Final = self._attached_engine("gpt-5.5", parent_condition="claude-opus-4-1") + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + matches: Final = engine.attached_policies(_context("gpt-4o"), only_applying=True) + assert matches == () + assert not _ancestor_logs(caplog) + + def test_requested_names_log_nothing(self, caplog: pytest.LogCaptureFixture) -> None: + engine: Final = build_policy_engine( + { + "parent": Policy(guardrails=PolicyGuardrails(add=["g-parent"])), + "child": Policy( + inherit="parent", + guardrails=PolicyGuardrails(add=["g-child"]), + condition=PolicyCondition(model="gpt-5.5"), + ), + } + ) + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + resolution: Final = engine.resolve(_context("gpt-4o"), names=["child"]) + assert resolution.applied_policies == ("child",) + assert not _ancestor_logs(caplog) + + +@pytest.mark.parametrize( + "pattern", + [ + pytest.param("a{4294967296}", id="repeat-count-overflows"), + pytest.param("(" * 1000 + "a" + ")" * 1000, id="nesting-too-deep"), + pytest.param("(?a)(?u)x", id="incompatible-flags"), + ], +) +def test_a_pattern_python_cannot_compile_does_not_break_other_policies(pattern: str) -> None: + engine: Final = build_policy_engine( + { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii"])), + "odd": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model=pattern)), + }, + attachments=(PolicyAttachment(policy="baseline", scope="*"), PolicyAttachment(policy="odd", teams=["other"])), + ) + assert engine.resolve(_context("gpt-4")).guardrails == ("pii",) + assert engine.policy_guardrails("odd", _context(pattern)).guardrails == ["x"] + + +class _Interrupted(Exception): + pass + + +class _RaisingPattern: + def fullmatch(self, value: str) -> None: + raise _Interrupted(value) + + +def _raising_pattern(pattern: str) -> _RaisingPattern: + return _RaisingPattern() + + +def test_an_exception_from_a_python_matched_pattern_is_raised_like_re_fullmatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("litellm.rust_bridge.policy_engine._compiled", _raising_pattern) + engine: Final = build_policy_engine( + {"p": Policy(guardrails=PolicyGuardrails(add=["g"]), condition=PolicyCondition(model=r"\N{DIGIT ONE}"))}, + attachments=(PolicyAttachment(policy="p", scope="*"),), + ) + with pytest.raises(_Interrupted): + engine.resolve(_context("1")) + with pytest.raises(_Interrupted): + engine.policy_guardrails("p", _context("1")) + with pytest.raises(_Interrupted): + engine.attached_policies(_context("1"), only_applying=True) + assert engine.resolve(_context(None)).guardrails == () + + +def test_lone_surrogates_resolve_like_python_strings() -> None: + engine: Final = build_policy_engine( + { + "p\ud800": Policy(guardrails=PolicyGuardrails(add=["g\ud800"]), condition=PolicyCondition(model="m.")), + "q": Policy(guardrails=PolicyGuardrails(add=["h"]), condition=PolicyCondition(model="m[\ud800-\udfff]")), + "x\ufffd": Policy(guardrails=PolicyGuardrails(add=["replacement"])), + }, + attachments=( + PolicyAttachment(policy="p\ud800", teams=["t\udc00"]), + PolicyAttachment(policy="q", models=["m\ud800"]), + ), + ) + context: Final = _context("m\ud800", team="t\udc00") + + resolution: Final = engine.resolve(context, names=["x\ud800"]) + + assert resolution.applied_policies == ("p\ud800", "q") + assert dict(resolution.policy_sources) == {"p\ud800": "team:t\udc00", "q": "model:m\ud800"} + assert resolution.guardrails == ("g\ud800", "h") + assert engine.policy_guardrails("p\ud800", context).guardrails == ["g\ud800"] + assert engine.resolve(_context("m\ufffd", team="t\udc00")).applied_policies == ("p\ud800",) + assert engine.resolve(_context("plain")).applied_policies == () + + +def test_a_lone_surrogate_in_a_request_matches_only_what_python_would() -> None: + engine: Final = build_policy_engine( + { + "x�": Policy(guardrails=PolicyGuardrails(add=["replacement"])), + "m": Policy(guardrails=PolicyGuardrails(add=["g"]), condition=PolicyCondition(model="m.")), + }, + attachments=(PolicyAttachment(policy="m", tags=["env-*"]),), + ) + + resolution: Final = engine.resolve(_context("m\udc80", team="t\udc80", tags=["env-\udc80"]), names=["x\udc80"]) + + assert resolution.applied_policies == ("m",) + assert dict(resolution.policy_sources) == {"m": "tag:env-\udc80"} + + +@pytest.mark.parametrize( + ("pattern", "expected"), + [(r"m\ud800", ("p",)), (r"m[\ud800-\udfff]", ("p",)), (r"m\U0010f800", ()), (r"m[^\U0010f800]", ("p",))], +) +def test_python_matches_a_lone_surrogate_request_against_the_original_text( + pattern: str, expected: tuple[str, ...] +) -> None: + engine: Final = build_policy_engine( + {"p": Policy(guardrails=PolicyGuardrails(add=["g"]), condition=PolicyCondition(model=pattern))}, + attachments=(PolicyAttachment(policy="p", scope="*"),), + ) + assert engine.resolve(_context("m\ud800")).applied_policies == expected + + +def test_building_an_engine_warns_about_each_circular_inheritance(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + build_policy_engine( + { + "a": Policy(inherit="b"), + "b": Policy(inherit="a"), + "fine": Policy(inherit="missing"), + } + ) + assert [record.getMessage() for record in caplog.records] == [ + "Circular inheritance detected for policy 'a'", + "Circular inheritance detected for policy 'b'", + ] diff --git a/tests/unit/proxy/policy_engine/test_policy_matcher.py b/tests/unit/proxy/policy_engine/test_policy_matcher.py deleted file mode 100644 index 862b5793eba..00000000000 --- a/tests/unit/proxy/policy_engine/test_policy_matcher.py +++ /dev/null @@ -1,434 +0,0 @@ -""" -Unit tests for PolicyMatcher - tests wildcard pattern matching via attachments. - -Tests: -- Wildcard matching (*, prefix-*) -- Scope matching via attachments (teams, keys, models) -""" - -import logging -from typing import Final - -import pytest -from hypothesis import given, settings -from hypothesis import strategies as st - -import litellm.proxy.policy_engine.attachment_registry as attachment_registry_module -import litellm.proxy.policy_engine.policy_registry as policy_registry_module -from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry -from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher -from litellm.proxy.policy_engine.policy_registry import PolicyRegistry -from litellm.proxy.policy_engine.policy_resolver import PolicyResolver -from litellm.types.proxy.policy_engine import ( - Policy, - PolicyCondition, - PolicyGuardrails, - PolicyMatchContext, - PolicyScope, -) - - -class TestPolicyMatcherPatternMatching: - """Test pattern matching utilities.""" - - def test_matches_pattern_exact(self): - """Test exact pattern matching.""" - assert ( - PolicyMatcher.matches_pattern("healthcare-team", ["healthcare-team"]) - is True - ) - assert ( - PolicyMatcher.matches_pattern("finance-team", ["healthcare-team"]) is False - ) - - def test_matches_pattern_wildcard(self): - """Test wildcard pattern matching.""" - assert PolicyMatcher.matches_pattern("any-team", ["*"]) is True - assert PolicyMatcher.matches_pattern("dev-key-123", ["dev-key-*"]) is True - assert PolicyMatcher.matches_pattern("prod-key-123", ["dev-key-*"]) is False - - def test_matches_pattern_none_value(self): - """Test None value only matches '*'.""" - assert PolicyMatcher.matches_pattern(None, ["*"]) is True - assert PolicyMatcher.matches_pattern(None, ["specific"]) is False - - -class TestPolicyMatcherScopeMatching: - """Test scope matching against context.""" - - def test_scope_matches_all_fields(self): - """Test scope matches when all fields match.""" - scope = PolicyScope(teams=["healthcare-team"], keys=["*"], models=["gpt-4"]) - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="any-key", model="gpt-4" - ) - assert PolicyMatcher.scope_matches(scope, context) is True - - def test_scope_does_not_match_team(self): - """Test scope doesn't match when team doesn't match.""" - scope = PolicyScope(teams=["healthcare-team"], keys=["*"], models=["*"]) - context = PolicyMatchContext( - team_alias="finance-team", key_alias="any-key", model="gpt-4" - ) - assert PolicyMatcher.scope_matches(scope, context) is False - - def test_scope_matches_with_wildcard_patterns(self): - """Test scope matches with wildcard patterns.""" - scope = PolicyScope(teams=["*"], keys=["dev-key-*"], models=["bedrock/*"]) - context = PolicyMatchContext( - team_alias="any-team", key_alias="dev-key-123", model="bedrock/claude-3" - ) - assert PolicyMatcher.scope_matches(scope, context) is True - - def test_scope_global_wildcard(self): - """Test global scope with all wildcards.""" - scope = PolicyScope(teams=["*"], keys=["*"], models=["*"]) - context = PolicyMatchContext( - team_alias="any-team", key_alias="any-key", model="any-model" - ) - assert PolicyMatcher.scope_matches(scope, context) is True - - -class TestPolicyMatcherScopeMatchingWithTags: - """Test scope matching with tag patterns.""" - - def test_scope_tag_matching(self): - """Test scope tag matching: exact, wildcard, no-match, and empty context tags.""" - # Exact match - scope = PolicyScope(teams=["*"], keys=["*"], models=["*"], tags=["healthcare"]) - context = PolicyMatchContext( - team_alias="team", - key_alias="key", - model="gpt-4", - tags=["healthcare", "internal"], - ) - assert PolicyMatcher.scope_matches(scope, context) is True - - # Wildcard match - scope_wc = PolicyScope(teams=["*"], keys=["*"], models=["*"], tags=["health-*"]) - context_wc = PolicyMatchContext( - team_alias="team", - key_alias="key", - model="gpt-4", - tags=["health-prod"], - ) - assert PolicyMatcher.scope_matches(scope_wc, context_wc) is True - - # No match — wrong tag - context_wrong = PolicyMatchContext( - team_alias="team", - key_alias="key", - model="gpt-4", - tags=["finance"], - ) - assert PolicyMatcher.scope_matches(scope, context_wrong) is False - - # No match — context has no tags - context_none = PolicyMatchContext( - team_alias="team", - key_alias="key", - model="gpt-4", - tags=None, - ) - assert PolicyMatcher.scope_matches(scope, context_none) is False - - # Scope without tags matches any context (opt-in semantics) - scope_no_tags = PolicyScope(teams=["*"], keys=["*"], models=["*"]) - assert PolicyMatcher.scope_matches(scope_no_tags, context) is True - - def test_scope_tags_and_team_combined(self): - """Test scope with both tags and team — both must match (AND logic).""" - scope = PolicyScope( - teams=["team-a"], keys=["*"], models=["*"], tags=["healthcare"] - ) - - # Both match - context_both = PolicyMatchContext( - team_alias="team-a", - key_alias="key", - model="gpt-4", - tags=["healthcare"], - ) - assert PolicyMatcher.scope_matches(scope, context_both) is True - - # Tag matches, team doesn't - context_wrong_team = PolicyMatchContext( - team_alias="team-b", - key_alias="key", - model="gpt-4", - tags=["healthcare"], - ) - assert PolicyMatcher.scope_matches(scope, context_wrong_team) is False - - # Team matches, tag doesn't - context_wrong_tag = PolicyMatchContext( - team_alias="team-a", - key_alias="key", - model="gpt-4", - tags=["finance"], - ) - assert PolicyMatcher.scope_matches(scope, context_wrong_tag) is False - - -class TestPolicyMatcherWithAttachments: - """Test getting matching policies via attachments.""" - - def test_get_matching_policies_via_attachments(self): - """Test matching policies through attachment registry.""" - # Create and configure attachment registry - registry = AttachmentRegistry() - registry.load_attachments( - [ - {"policy": "healthcare-policy", "teams": ["healthcare-team"]}, - {"policy": "global-policy", "scope": "*"}, - ] - ) - - # Test matching via the registry directly - context = PolicyMatchContext( - team_alias="healthcare-team", key_alias="k", model="gpt-4" - ) - attached = registry.get_attached_policies(context) - - assert "healthcare-policy" in attached - assert "global-policy" in attached - - def test_get_matching_policies_no_match(self): - """Test no policies match when attachments don't match context.""" - registry = AttachmentRegistry() - registry.load_attachments( - [ - {"policy": "healthcare-policy", "teams": ["healthcare-team"]}, - ] - ) - - context = PolicyMatchContext( - team_alias="finance-team", key_alias="k", model="gpt-4" - ) - attached = registry.get_attached_policies(context) - - assert "healthcare-policy" not in attached - - -def _global_registries(monkeypatch): - policies = PolicyRegistry() - policies.load_policies( - { - "guardrail-y": {"guardrails": {"add": ["y"]}}, - "guardrail-x": {"guardrails": {"add": ["x"]}, "condition": {"model": "claude.*"}}, - } - ) - attachments = AttachmentRegistry() - attachments.load_attachments( - [ - {"policy": "guardrail-x", "tags": ["opt-in"]}, - {"policy": "guardrail-y", "scope": "*", "default": True}, - ] - ) - monkeypatch.setattr(policy_registry_module, "get_policy_registry", lambda: policies) - monkeypatch.setattr(attachment_registry_module, "get_attachment_registry", lambda: attachments) - return policies - - -def _inherited_registries(monkeypatch, parent_condition=None): - policies = PolicyRegistry() - policies.load_policies( - { - "parent": { - "guardrails": {"add": ["y"]}, - **({"condition": parent_condition} if parent_condition else {}), - }, - "child": { - "inherit": "parent", - "guardrails": {"add": ["x"]}, - "condition": {"model": "claude.*"}, - }, - "fallback": {"guardrails": {"add": ["z"]}}, - } - ) - attachments = AttachmentRegistry() - attachments.load_attachments( - [ - {"policy": "child", "scope": "*"}, - {"policy": "fallback", "scope": "*", "default": True}, - ] - ) - monkeypatch.setattr(policy_registry_module, "get_policy_registry", lambda: policies) - monkeypatch.setattr(attachment_registry_module, "get_attachment_registry", lambda: attachments) - return policies - - -class TestGetMatchingPoliciesFallback: - def test_condition_failing_opt_in_falls_back_to_default(self, monkeypatch): - _global_registries(monkeypatch) - context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"]) - - assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-y"] - - def test_condition_passing_opt_in_suppresses_default(self, monkeypatch): - _global_registries(monkeypatch) - context = PolicyMatchContext(team_alias="t", key_alias="k", model="claude-haiku", tags=["opt-in"]) - - assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-x"] - - def test_policy_applies_reads_registry_once(self, monkeypatch): - policies = _global_registries(monkeypatch) - calls = [] - original = policies.get_all_policies - monkeypatch.setattr(policies, "get_all_policies", lambda: calls.append(1) or original()) - context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"]) - - PolicyMatcher.get_matching_policies(context=context) - - assert len(calls) == 1 - - def test_condition_missing_child_with_unconditional_parent_still_matches(self, monkeypatch): - _inherited_registries(monkeypatch) - context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5") - - assert PolicyMatcher.get_matching_policies(context=context) == ["child"] - - def test_child_whose_whole_chain_misses_falls_back_to_default(self, monkeypatch): - _inherited_registries(monkeypatch, parent_condition={"model": "claude.*"}) - context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5") - - assert PolicyMatcher.get_matching_policies(context=context) == ["fallback"] - - def test_get_policies_with_matching_conditions_keeps_missing_policy_out(self): - policies = { - "real": Policy( - guardrails=PolicyGuardrails(add=["g"]), - condition=PolicyCondition(model="claude.*"), - ), - } - context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5") - - assert ( - PolicyMatcher.get_policies_with_matching_conditions( - policy_names=["nope"], context=context, policies=policies - ) - == [] - ) - - -_MODELS: Final = ("gpt-4o", "gpt-5.5", "claude-opus-4-1") - - -def _policy_forest(draw: st.DrawFn) -> dict[str, Policy]: # mutable-ok: PolicyResolver takes dict[str, Policy] - names: Final = tuple(f"p{i}" for i in range(draw(st.integers(min_value=1, max_value=6)))) - return { - name: Policy( - inherit=draw(st.sampled_from((None, *names[:i]))), - guardrails=PolicyGuardrails(add=[f"g-{name}"]), - condition=draw(st.sampled_from((None, *(PolicyCondition(model=m) for m in _MODELS)))), - ) - for i, name in enumerate(names) - } - - -@st.composite -def _forest_and_request( - draw: st.DrawFn, -) -> tuple[dict[str, Policy], tuple[str, ...], PolicyMatchContext]: # mutable-ok: PolicyResolver takes dict - policies: Final = _policy_forest(draw) - attached: Final = tuple(draw(st.lists(st.sampled_from(sorted(policies)), unique=True))) - context: Final = PolicyMatchContext(team_alias="t", key_alias="k", model=draw(st.sampled_from(_MODELS))) - return policies, attached, context - - -def _own_condition_applies(policy: Policy, context: PolicyMatchContext) -> bool: - return policy.condition is None or policy.condition.model == context.model - - -def _applicable_chain( - policies: dict[str, Policy], # mutable-ok: PolicyResolver takes dict[str, Policy] - name: str, - context: PolicyMatchContext, -) -> tuple[str, ...]: - chain: Final = PolicyResolver.resolve_inheritance_chain(policy_name=name, policies=policies) - return tuple(member for member in chain if _own_condition_applies(policies[member], context)) - - -class TestChainMatchingProperties: - @given(_forest_and_request()) - @settings(max_examples=400, deadline=None) - def test_chain_matching_only_widens_to_applicable_ancestor_guardrails( - self, - case: tuple[dict[str, Policy], tuple[str, ...], PolicyMatchContext], # mutable-ok: PolicyResolver takes dict - ): - policies, attached, context = case - head: Final = tuple( - PolicyMatcher.get_policies_with_matching_conditions( - policy_names=attached, context=context, policies=policies - ) - ) - base: Final = tuple(name for name in attached if _own_condition_applies(policies[name], context)) - expected_head: Final = tuple(name for name in attached if _applicable_chain(policies, name, context)) - - assert head == expected_head, "a policy applies exactly when some chain member's own condition applies" - assert frozenset(base) <= frozenset(head), "head must never drop a policy base applied" - - for name in head: - resolved = PolicyResolver.resolve_policy_guardrails(policy_name=name, policies=policies, context=context) - assert sorted(resolved.guardrails) == sorted( - f"g-{member}" for member in _applicable_chain(policies, name, context) - ) - if name not in base: - assert f"g-{name}" not in resolved.guardrails, "a condition-missed child must not add its own guardrail" - - -class TestAncestorAdmissionLogging: - @staticmethod - def _chain() -> dict[str, Policy]: # mutable-ok: PolicyResolver takes dict[str, Policy] - return { - "parent": Policy(guardrails=PolicyGuardrails(add=["g-parent"])), - "child": Policy( - inherit="parent", - guardrails=PolicyGuardrails(add=["g-child"]), - condition=PolicyCondition(model="gpt-5.5"), - ), - } - - def test_logs_when_admitted_through_ancestor_only(self, caplog): - context: Final = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4o") - with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): - result: Final = PolicyMatcher.policy_applies(context, self._chain())("child") - records: Final = [r for r in caplog.records if "applied through ancestor" in r.getMessage()] - assert result is True - assert len(records) == 1 - assert "applied through ancestor 'parent'" in records[0].getMessage() - assert "'child'" in records[0].getMessage() - - def test_no_log_when_own_condition_matches(self, caplog): - context: Final = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5") - with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): - result: Final = PolicyMatcher.policy_applies(context, self._chain())("child") - assert result is True - assert not [r for r in caplog.records if "applied through ancestor" in r.getMessage()] - - def test_no_log_when_no_chain_member_applies(self, caplog): - policies: Final = { - "parent": Policy( - guardrails=PolicyGuardrails(add=["g-parent"]), - condition=PolicyCondition(model="claude-opus-4-1"), - ), - "child": Policy( - inherit="parent", - guardrails=PolicyGuardrails(add=["g-child"]), - condition=PolicyCondition(model="gpt-5.5"), - ), - } - context: Final = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4o") - with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): - result: Final = PolicyMatcher.policy_applies(context, policies)("child") - assert result is False - assert not [r for r in caplog.records if "applied through ancestor" in r.getMessage()] - - def test_condition_filter_logs_nothing(self, caplog): - context: Final = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4o") - with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): - result: Final = PolicyMatcher.get_policies_with_matching_conditions( - policy_names=["child"], context=context, policies=self._chain() - ) - assert result == ["child"] - assert not [r for r in caplog.records if "applied through ancestor" in r.getMessage()] diff --git a/tests/unit/proxy/policy_engine/test_policy_resolver.py b/tests/unit/proxy/policy_engine/test_policy_resolver.py deleted file mode 100644 index 3d2f547a744..00000000000 --- a/tests/unit/proxy/policy_engine/test_policy_resolver.py +++ /dev/null @@ -1,255 +0,0 @@ -""" -Unit tests for PolicyResolver - tests guardrail resolution. - -Tests: -- Inheritance chain resolution -- Inheritance with add/remove -- Model conditions -""" - -import pytest - -from litellm.proxy.policy_engine.policy_resolver import PolicyResolver -from litellm.types.proxy.policy_engine import ( - GuardrailPipeline, - PipelineStep, - Policy, - PolicyCondition, - PolicyGuardrails, - PolicyMatchContext, -) - - -class TestPolicyResolverInheritance: - """Test resolve_policy_guardrails - inheritance and add/remove.""" - - def test_resolve_simple_policy(self): - """Test resolving guardrails for a simple policy.""" - policies = { - "global": Policy( - guardrails=PolicyGuardrails(add=["pii_blocker", "toxicity_filter"]), - ), - } - - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name="global", policies=policies - ) - - assert set(resolved.guardrails) == {"pii_blocker", "toxicity_filter"} - assert resolved.inheritance_chain == ["global"] - - def test_resolve_with_inheritance(self): - """Test child policy inherits and adds guardrails from parent.""" - policies = { - "base": Policy( - guardrails=PolicyGuardrails(add=["pii_blocker"]), - ), - "healthcare": Policy( - inherit="base", - guardrails=PolicyGuardrails(add=["hipaa_audit"]), - ), - } - - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name="healthcare", policies=policies - ) - - # Healthcare inherits pii_blocker from base and adds hipaa_audit - assert set(resolved.guardrails) == {"pii_blocker", "hipaa_audit"} - assert resolved.inheritance_chain == ["base", "healthcare"] - - def test_resolve_with_remove(self): - """Test child policy can remove guardrails from parent.""" - policies = { - "base": Policy( - guardrails=PolicyGuardrails(add=["pii_blocker", "phi_blocker"]), - ), - "dev": Policy( - inherit="base", - guardrails=PolicyGuardrails( - add=["toxicity_filter"], remove=["phi_blocker"] - ), - ), - } - - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name="dev", policies=policies - ) - - # dev inherits pii_blocker from base, adds toxicity_filter, removes phi_blocker - assert "pii_blocker" in resolved.guardrails - assert "toxicity_filter" in resolved.guardrails - assert "phi_blocker" not in resolved.guardrails - - def test_resolve_deep_inheritance_chain(self): - """Test multi-level inheritance chain.""" - policies = { - "root": Policy( - guardrails=PolicyGuardrails(add=["root_guardrail"]), - ), - "middle": Policy( - inherit="root", - guardrails=PolicyGuardrails(add=["middle_guardrail"]), - ), - "leaf": Policy( - inherit="middle", - guardrails=PolicyGuardrails(add=["leaf_guardrail"]), - ), - } - - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name="leaf", policies=policies - ) - - assert set(resolved.guardrails) == { - "root_guardrail", - "middle_guardrail", - "leaf_guardrail", - } - assert resolved.inheritance_chain == ["root", "middle", "leaf"] - - -class TestPolicyResolverWithConditions: - """Test resolve_policy_guardrails with model conditions.""" - - def test_condition_matches(self): - """Test guardrails are added when condition matches.""" - policies = { - "gpt4-policy": Policy( - guardrails=PolicyGuardrails(add=["toxicity_filter"]), - condition=PolicyCondition(model="gpt-4.*"), - ), - } - - # GPT-4 should get guardrails - context = PolicyMatchContext(team_alias="team", key_alias="k", model="gpt-4") - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name="gpt4-policy", - policies=policies, - context=context, - ) - - assert "toxicity_filter" in resolved.guardrails - - def test_condition_does_not_match(self): - """Test guardrails are NOT added when condition doesn't match.""" - policies = { - "gpt4-policy": Policy( - guardrails=PolicyGuardrails(add=["toxicity_filter"]), - condition=PolicyCondition(model="gpt-4.*"), - ), - } - - # GPT-3.5 should NOT get guardrails - context = PolicyMatchContext(team_alias="team", key_alias="k", model="gpt-3.5") - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name="gpt4-policy", - policies=policies, - context=context, - ) - - assert "toxicity_filter" not in resolved.guardrails - - def test_no_condition_always_applies(self): - """Test policy without condition always applies.""" - policies = { - "global": Policy( - guardrails=PolicyGuardrails(add=["pii_blocker"]), - ), - } - - context = PolicyMatchContext(team_alias="any", key_alias="any", model="any") - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name="global", - policies=policies, - context=context, - ) - - assert "pii_blocker" in resolved.guardrails - - def test_inheritance_with_condition(self): - """Test inheritance works with conditions.""" - policies = { - "base": Policy( - guardrails=PolicyGuardrails(add=["pii_blocker"]), - ), - "child": Policy( - inherit="base", - guardrails=PolicyGuardrails(add=["child_guardrail"]), - condition=PolicyCondition(model="gpt-4"), - ), - } - - # GPT-4 should get both base and child guardrails - context_gpt4 = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4") - resolved_gpt4 = PolicyResolver.resolve_policy_guardrails( - policy_name="child", - policies=policies, - context=context_gpt4, - ) - assert "pii_blocker" in resolved_gpt4.guardrails - assert "child_guardrail" in resolved_gpt4.guardrails - - # GPT-3.5 should only get base guardrails (child condition doesn't match) - context_gpt35 = PolicyMatchContext( - team_alias="t", key_alias="k", model="gpt-3.5" - ) - resolved_gpt35 = PolicyResolver.resolve_policy_guardrails( - policy_name="child", - policies=policies, - context=context_gpt35, - ) - assert "pii_blocker" in resolved_gpt35.guardrails - assert "child_guardrail" not in resolved_gpt35.guardrails - - def test_resolve_guardrails_for_context_with_condition_missing_child_keeps_inherited_parent(self): - """Test a matched child whose condition misses still contributes unconditional parent guardrails.""" - policies = { - "parent": Policy( - guardrails=PolicyGuardrails(add=["y"]), - ), - "child": Policy( - inherit="parent", - guardrails=PolicyGuardrails(add=["x"]), - condition=PolicyCondition(model="claude.*"), - ), - } - - context_miss = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5") - assert PolicyResolver.resolve_guardrails_for_context( - context=context_miss, policies=policies, policy_names=["child"] - ) == ["y"] - - context_hit = PolicyMatchContext(team_alias="t", key_alias="k", model="claude-haiku") - assert set( - PolicyResolver.resolve_guardrails_for_context( - context=context_hit, policies=policies, policy_names=["child"] - ) - ) == {"x", "y"} - - def test_resolve_pipelines_for_context_skips_pipeline_when_own_condition_misses(self): - """Test a matched child whose own condition misses does not run its pipeline.""" - pipeline = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="child-guard")]) - policies = { - "parent": Policy( - guardrails=PolicyGuardrails(add=["y"]), - ), - "child": Policy( - inherit="parent", - pipeline=pipeline, - condition=PolicyCondition(model="gpt-5.5"), - ), - } - - context_miss = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4o") - assert ( - PolicyResolver.resolve_pipelines_for_context( - context=context_miss, policies=policies, policy_names=["child"] - ) - == [] - ) - - context_hit = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5") - assert PolicyResolver.resolve_pipelines_for_context( - context=context_hit, policies=policies, policy_names=["child"] - ) == [("child", pipeline)] diff --git a/tests/unit/proxy/policy_engine/test_policy_versioning.py b/tests/unit/proxy/policy_engine/test_policy_versioning.py index b6633779326..07d9343857b 100644 --- a/tests/unit/proxy/policy_engine/test_policy_versioning.py +++ b/tests/unit/proxy/policy_engine/test_policy_versioning.py @@ -476,6 +476,24 @@ class TestConfigPoliciesPreservedAcrossDbSync: assert policy.guardrails.add == ["tooling"] assert registry.get_source("config-policy") == "config" + @pytest.mark.asyncio + async def test_an_unchanged_sync_keeps_the_policy_containers(self): + registry = PolicyRegistry() + registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}}) + + def rows(): + return _prisma_with_policy_rows( + [_make_row(policy_id="db-1", policy_name="db-policy", guardrails_add=["db-guard"])], + [_make_row(policy_id="draft-1", policy_name="db-policy", version_status="draft")], + ) + + await registry.sync_policies_from_db(rows()) + policies, versions = registry.current_policies(), registry.current_versions() + await registry.sync_policies_from_db(rows()) + + assert registry.current_policies() is policies + assert registry.current_versions() is versions + @pytest.mark.asyncio async def test_sync_merges_db_policies_with_config_policies(self): registry = PolicyRegistry() @@ -518,18 +536,14 @@ class TestConfigPoliciesPreservedAcrossDbSync: @pytest.mark.asyncio async def test_config_policy_resolves_guardrails_after_sync(self): - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver + from litellm.rust_bridge.policy_engine import build_policy_engine registry = PolicyRegistry() registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}}) await registry.sync_policies_from_db(_prisma_with_policy_rows([])) - resolved = PolicyResolver.resolve_policy_guardrails( - policy_name="config-policy", - policies=registry.get_all_policies(), - context=None, - ) + resolved = build_policy_engine(registry.current_policies()).policy_guardrails("config-policy") assert resolved.guardrails == ["tooling"] @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 557ac72572f..87be23de096 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -2,6 +2,7 @@ import asyncio import base64 import copy import json +import logging import os import time from datetime import datetime, timezone @@ -45,8 +46,8 @@ from litellm.proxy.litellm_pre_call_utils import ( _apply_credential_overrides_from_model_config, _extract_credential_from_entry, _get_enforced_params, - _match_and_track_policies, _promoted_trace_control_fields, + _record_policy_engine_resolution, _resolve_credential_from_model_config, _resolve_provider_from_deployment, _update_model_if_key_alias_exists, @@ -5473,67 +5474,56 @@ async def test_baggage_mapped_tags_match_tag_policies_only_when_the_caller_or_ke assert data["metadata"].get("guardrails", []) == expected_guardrails -def test_match_and_track_policies_preserves_attachment_and_request_body_order(): - from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry - from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext +def test_policy_engine_preserves_attachment_and_request_body_order(): + from litellm.rust_bridge.policy_engine import build_policy_engine + from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyMatchContext attachment_policy_names = [f"attachment-policy-{index}" for index in range(8)] request_body_policy_names = ["body-policy-1", "body-policy-2"] policy_names = [*attachment_policy_names, *request_body_policy_names] - policies = {policy_name: Policy() for policy_name in policy_names} - attachment_registry = AttachmentRegistry() - attachment_registry.load_attachments( - [{"policy": policy_name, "scope": "*"} for policy_name in attachment_policy_names] + engine = build_policy_engine( + {policy_name: Policy() for policy_name in policy_names}, + attachments=tuple(PolicyAttachment(policy=policy_name, scope="*") for policy_name in attachment_policy_names), ) - applied_policy_names, _ = _match_and_track_policies( - data={"metadata": {}}, - context=PolicyMatchContext(model="gpt-4"), - request_body_policies=request_body_policy_names, - policies_override=policies, - attachment_registry_override=attachment_registry, - ) + resolution = engine.resolve(PolicyMatchContext(model="gpt-4"), request_body_policy_names) - assert applied_policy_names == policy_names + assert resolution.applied_policies == tuple(policy_names) -def test_match_and_track_policies_keeps_condition_missing_child_alongside_unconditional_sibling(): - from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry +def test_policy_engine_keeps_condition_missing_child_alongside_unconditional_sibling(): + from litellm.rust_bridge.policy_engine import build_policy_engine from litellm.types.proxy.policy_engine import ( Policy, + PolicyAttachment, PolicyCondition, PolicyGuardrails, PolicyMatchContext, ) - policies = { - "baseline": Policy(guardrails=PolicyGuardrails(add=["baseline_guardrail"])), - "parent": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"])), - "child": Policy( - inherit="parent", - guardrails=PolicyGuardrails(add=["child_guard"]), - condition=PolicyCondition(model="claude.*"), - ), + engine = build_policy_engine( + { + "baseline": Policy(guardrails=PolicyGuardrails(add=["baseline_guardrail"])), + "parent": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"])), + "child": Policy( + inherit="parent", + guardrails=PolicyGuardrails(add=["child_guard"]), + condition=PolicyCondition(model="claude.*"), + ), + }, + attachments=(PolicyAttachment(policy="baseline", scope="*"), PolicyAttachment(policy="child", scope="*")), + ) + data: dict[str, object] = {"metadata": {}} + + resolution = engine.resolve(PolicyMatchContext(model="gpt-5.5")) + _record_policy_engine_resolution(data, "metadata", resolution) + + assert resolution.applied_policies == ("baseline", "child") + assert data["metadata"] == { + "applied_policies": ["baseline", "child"], + "policy_sources": {"baseline": "scope:*", "child": "scope:*"}, + "guardrails": ["baseline_guardrail", "pii_blocker"], } - attachment_registry = AttachmentRegistry() - attachment_registry.load_attachments( - [ - {"policy": "baseline", "scope": "*"}, - {"policy": "child", "scope": "*"}, - ] - ) - data = {"metadata": {}} - - applied_policy_names, _ = _match_and_track_policies( - data=data, - context=PolicyMatchContext(model="gpt-5.5"), - request_body_policies=[], - policies_override=policies, - attachment_registry_override=attachment_registry, - ) - - assert applied_policy_names == ["baseline", "child"] - assert data["metadata"]["applied_policies"] == ["baseline", "child"] @pytest.mark.asyncio @@ -5578,6 +5568,66 @@ async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_ assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"] +@pytest.mark.asyncio +async def test_add_guardrails_from_policy_engine_explains_its_resolution_in_debug_logs(caplog): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import ( + GuardrailPipeline, + PipelineStep, + Policy, + PolicyAttachment, + PolicyCondition, + PolicyGuardrails, + ) + + data = {"model": "gpt-4", "metadata": {}, "policies": ["claude-only", "policy_v1", "policy_gone"]} + policy_registry = get_policy_registry() + policy_registry._policies = { + "base": Policy(guardrails=PolicyGuardrails(add=["pii"])), + "child": Policy( + inherit="base", guardrails=PolicyGuardrails(add=["audit"]), condition=PolicyCondition(model="claude.*") + ), + "claude-only": Policy(guardrails=PolicyGuardrails(add=["x"]), condition=PolicyCondition(model="claude.*")), + "draft": Policy(), + } + policy_registry._policies_by_id = { + "v1": ("draft", Policy(pipeline=GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="judge")]))) + } + policy_registry._initialized = True + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [PolicyAttachment(policy="child", scope="*")] + attachment_registry._initialized = True + + try: + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await add_guardrails_from_policy_engine( + data=data, metadata_variable_name="metadata", user_api_key_dict=UserAPIKeyAuth(api_key="test-key") + ) + finally: + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + messages = [record.getMessage() for record in caplog.records] + for expected in ( + "Policy engine: loaded version by ID policy_v1 -> draft", + "Policy engine: policy version gone not found in cache, skipping", + "Policy engine: matched policies via attachments: ['child']", + "Policy engine: added dynamic policies from request body: ['claude-only']", + "Policy 'child' condition did not match model=gpt-4, skipping its guardrails and pipeline", + "Policy engine: skipped policies with no applicable chain member: ['claude-only']", + "Policy engine: applied policies (conditions matched): ['child', 'draft']", + "Policy 'child' contributes guardrails: ['pii']", + "Policy 'draft' has pipeline with 1 steps", + "Policy engine: resolved 1 pipeline(s), managed guardrails: {'judge'}", + "Policy engine: added guardrails to request metadata: ['pii']", + ): + assert expected in messages + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_keeps_configured_guardrails_and_skips_ones_already_selected(): from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry @@ -5841,32 +5891,28 @@ async def test_add_guardrails_from_policy_engine_policy_version_by_id(): policy_registry = get_policy_registry() policy_registry._policies = {} + policy_registry._policies_by_id = {policy_version_uuid: ("test-policy-from-version", published_version_policy)} policy_registry._initialized = True attachment_registry = get_attachment_registry() attachment_registry._attachments = [] attachment_registry._initialized = True - with patch.object( - policy_registry, - "get_policy_by_id_for_request", - return_value=("test-policy-from-version", published_version_policy), - ): + try: await add_guardrails_from_policy_engine( data=data, metadata_variable_name="metadata", user_api_key_dict=user_api_key_dict, ) - # Verify guardrails from the specific version were applied - assert "metadata" in data - assert "guardrails" in data["metadata"] - assert "published_version_guardrail" in data["metadata"]["guardrails"] - assert "policies" not in data - - # Clean up - policy_registry._policies = {} - policy_registry._initialized = False + assert "metadata" in data + assert "guardrails" in data["metadata"] + assert "published_version_guardrail" in data["metadata"]["guardrails"] + assert "policies" not in data + finally: + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False @pytest.mark.asyncio diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index ad3036560a3..720aa2b1e9b 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -28,9 +28,12 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.openai.workload_identity import _workload_identity_auth from litellm.models.credentials import CredentialItem from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy -from litellm.proxy.litellm_pre_call_utils import _apply_resolved_guardrails_to_metadata +from litellm.proxy.litellm_pre_call_utils import ( + _record_policy_engine_resolution, # pyright: ignore[reportPrivateUsage] # writes a resolution into request metadata +) from litellm.realtime_api import main as realtime_main from litellm.realtime_api.main import with_resolved_session_model +from litellm.rust_bridge.policy_engine import build_policy_engine from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyMatchContext from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep @@ -2402,17 +2405,18 @@ async def test_translation_rejects_only_applicable_policy_pipelines( ) -> None: monkeypatch.setattr(litellm, "callbacks", []) request_data: Final[dict[str, object]] = {"metadata": {}, "litellm_metadata": {}} - _apply_resolved_guardrails_to_metadata( - data=request_data, - metadata_variable_name=metadata_bucket, - context=PolicyMatchContext(model="translation-deployment"), - policy_names=["translation-policy"] if selected else [], - policies={ + engine: Final = build_policy_engine( + { "translation-policy": Policy( condition=PolicyCondition(model="translation-deployment" if condition_matches else "unrelated-model"), pipeline=GuardrailPipeline(mode=mode, steps=[PipelineStep(guardrail="speech-policy", on_fail="block")]), ) - }, + } + ) + _record_policy_engine_resolution( + request_data, + metadata_bucket, + engine.resolve(PolicyMatchContext(model="translation-deployment"), ["translation-policy"] if selected else []), ) metadata: Final = realtime_main._build_litellm_metadata(request_data) blocked: Final = selected and condition_matches diff --git a/tests/unit/rust_bridge/test_policy_engine.py b/tests/unit/rust_bridge/test_policy_engine.py new file mode 100644 index 00000000000..a985bb861b6 --- /dev/null +++ b/tests/unit/rust_bridge/test_policy_engine.py @@ -0,0 +1,61 @@ +from collections.abc import Iterator +from typing import Final + +import pytest + +from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry, get_attachment_registry +from litellm.proxy.policy_engine.policy_registry import PolicyRegistry, get_policy_registry +from litellm.rust_bridge.policy_engine import current_policy_engine +from litellm.types.proxy.policy_engine import ( + GuardrailPipeline, + PipelineStep, + Policy, + PolicyAttachment, + PolicyGuardrails, + PolicyMatchContext, +) + + +@pytest.fixture +def registries() -> Iterator[tuple[PolicyRegistry, AttachmentRegistry]]: + policy_registry: Final = get_policy_registry() + attachment_registry: Final = get_attachment_registry() + policy_registry.clear() + attachment_registry.clear() + yield policy_registry, attachment_registry + policy_registry.clear() + attachment_registry.clear() + + +def test_a_registry_change_is_resolved_with_the_new_policies( + registries: tuple[PolicyRegistry, AttachmentRegistry], +) -> None: + policy_registry, attachment_registry = registries + policy_registry.add_policy("base", Policy(guardrails=PolicyGuardrails(add=["pii"]))) + attachment_registry.add_attachment(PolicyAttachment(policy="base", scope="*")) + context: Final = PolicyMatchContext(model="gpt-4") + + first: Final = current_policy_engine() + assert first.resolve(context).guardrails == ("pii",) + assert current_policy_engine() is first + policy_registry.add_policy("base", Policy(guardrails=PolicyGuardrails(add=["audit"]))) + assert current_policy_engine().resolve(context).guardrails == ("audit",) + attachment_registry.remove_attachments_for_policy("base") + assert current_policy_engine().resolve(context).guardrails == () + + +def test_pipelines_come_back_as_the_policy_objects_of_the_registry_or_version( + registries: tuple[PolicyRegistry, AttachmentRegistry], +) -> None: + policy_registry: Final = registries[0] + pipeline: Final = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="judge")]) + version_pipeline: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="pii")]) + policy_registry.add_policy("base", Policy(pipeline=pipeline)) + policy_registry.add_policy("draft", Policy()) + policy_registry._policies_by_id = {"v1": ("draft", Policy(pipeline=version_pipeline))} # pyright: ignore[reportPrivateUsage] # no public setter for the version cache + + resolution: Final = current_policy_engine().resolve(PolicyMatchContext(model="gpt-4"), ("base",), ("v1",)) + + assert resolution.pipelines == (("base", pipeline), ("draft", version_pipeline)) + assert resolution.pipelines[0][1] is pipeline + assert resolution.pipeline_managed_guardrails == frozenset({"judge", "pii"})