gemini_adk_fluent_rs/compose/
guards.rs

1//! G — Guard composition.
2//!
3//! Combine output guards with `+`: every guard must pass
4//! (`G::pii() + G::length(1, 2000)`).
5//!
6//! ## Wiring
7//!
8//! A [`GuardComposite`] attached via `AgentBuilder::guard` is installed on the
9//! compiled `LlmTextAgent` as an `after_model` middleware layer (see
10//! [`GuardComposite::into_middleware`]). Every model response is checked against
11//! all guards; if any guard rejects the output the agent run fails with an
12//! [`AgentError`] enumerating the violations, vetoing the response.
13
14use std::sync::Arc;
15
16use async_trait::async_trait;
17use gemini_adk_rs::error::AgentError;
18use gemini_adk_rs::llm::{BaseLlm, LlmRequest, LlmResponse};
19use gemini_adk_rs::middleware::Middleware;
20
21use crate::compose::judge::{LlmJudge, render_contents};
22
23/// A guard that validates agent output.
24#[derive(Clone)]
25pub struct GuardRule {
26    name: &'static str,
27    kind: GuardKind,
28}
29
30/// How a guard decides pass/fail.
31#[derive(Clone)]
32enum GuardKind {
33    /// Synchronous predicate over the output text.
34    Sync(#[allow(clippy::type_complexity)] Arc<dyn Fn(&str) -> Result<(), String> + Send + Sync>),
35    /// LLM-as-judge over the output (and, for grounding, the input context).
36    Judge(LlmJudge),
37}
38
39impl GuardRule {
40    fn new(
41        name: &'static str,
42        f: impl Fn(&str) -> Result<(), String> + Send + Sync + 'static,
43    ) -> Self {
44        Self {
45            name,
46            kind: GuardKind::Sync(Arc::new(f)),
47        }
48    }
49
50    fn judge(name: &'static str, judge: LlmJudge) -> Self {
51        Self {
52            name,
53            kind: GuardKind::Judge(judge),
54        }
55    }
56
57    /// Name of this guard.
58    pub fn name(&self) -> &str {
59        self.name
60    }
61
62    /// Synchronously check the output. LLM-judge guards cannot run on the sync
63    /// path and always return `Ok(())` here — use [`GuardRule::check_async`] (the
64    /// guard middleware uses the async path).
65    pub fn check(&self, output: &str) -> Result<(), String> {
66        match &self.kind {
67            GuardKind::Sync(f) => f(output),
68            GuardKind::Judge(_) => Ok(()),
69        }
70    }
71
72    /// Check the output, running an LLM judge if this is a judge guard.
73    /// `context` is the model's input history (for grounding/hallucination).
74    pub async fn check_async(&self, output: &str, context: Option<&str>) -> Result<(), String> {
75        match &self.kind {
76            GuardKind::Sync(f) => f(output),
77            GuardKind::Judge(judge) => {
78                let verdict = judge.judge(output, context).await;
79                if verdict.flagged {
80                    Err(verdict.reason)
81                } else {
82                    Ok(())
83                }
84            }
85        }
86    }
87}
88
89impl std::fmt::Debug for GuardRule {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        f.debug_struct("GuardRule")
92            .field("name", &self.name)
93            .finish()
94    }
95}
96
97/// Combine two guards with `+`: both must pass.
98impl std::ops::Add for GuardRule {
99    type Output = GuardComposite;
100
101    fn add(self, rhs: GuardRule) -> Self::Output {
102        GuardComposite {
103            guards: vec![self, rhs],
104        }
105    }
106}
107
108/// A composite of guards — all must pass for output to be accepted.
109#[derive(Clone)]
110#[non_exhaustive]
111pub struct GuardComposite {
112    /// The guards in this composite.
113    pub guards: Vec<GuardRule>,
114}
115
116impl GuardComposite {
117    /// Check all guards against the output (sync path; LLM-judge guards are
118    /// skipped — see [`GuardComposite::check_all_async`]). Returns all violations.
119    pub fn check_all(&self, output: &str) -> Vec<String> {
120        self.guards
121            .iter()
122            .filter_map(|g| g.check(output).err())
123            .collect()
124    }
125
126    /// Check all guards, running LLM-judge guards against `output` and the
127    /// optional input `context`. Returns all violations as `name: reason`.
128    pub async fn check_all_async(&self, output: &str, context: Option<&str>) -> Vec<String> {
129        let mut violations = Vec::new();
130        for g in &self.guards {
131            if let Err(reason) = g.check_async(output, context).await {
132                violations.push(format!("{}: {}", g.name(), reason));
133            }
134        }
135        violations
136    }
137
138    /// Number of guards.
139    pub fn len(&self) -> usize {
140        self.guards.len()
141    }
142
143    /// Whether empty.
144    pub fn is_empty(&self) -> bool {
145        self.guards.is_empty()
146    }
147}
148
149impl std::ops::Add<GuardRule> for GuardComposite {
150    type Output = GuardComposite;
151
152    fn add(mut self, rhs: GuardRule) -> Self::Output {
153        self.guards.push(rhs);
154        self
155    }
156}
157
158/// A single guard is a one-element composite, so `.guard(G::pii())` works
159/// without an explicit `| `.
160impl From<GuardRule> for GuardComposite {
161    fn from(guard: GuardRule) -> Self {
162        GuardComposite {
163            guards: vec![guard],
164        }
165    }
166}
167
168impl GuardComposite {
169    /// Adapt this guard composite into an `after_model` middleware layer that
170    /// vetoes any model response failing one or more guards.
171    pub fn into_middleware(self) -> Arc<dyn Middleware> {
172        Arc::new(GuardMiddleware { guards: self })
173    }
174}
175
176/// Middleware adapter that enforces a [`GuardComposite`] on every model response.
177struct GuardMiddleware {
178    guards: GuardComposite,
179}
180
181#[async_trait]
182impl Middleware for GuardMiddleware {
183    fn name(&self) -> &str {
184        "guard"
185    }
186
187    async fn after_model(
188        &self,
189        request: &LlmRequest,
190        response: &LlmResponse,
191    ) -> Result<Option<LlmResponse>, AgentError> {
192        // Render the input history so grounding/hallucination judges can see
193        // what the response is supposed to be consistent with.
194        let context = render_contents(&request.contents);
195        let violations = self
196            .guards
197            .check_all_async(&response.text(), Some(&context))
198            .await;
199        if violations.is_empty() {
200            Ok(None)
201        } else {
202            Err(AgentError::Other(format!(
203                "guard violation: {}",
204                violations.join("; ")
205            )))
206        }
207    }
208}
209
210/// The `G` namespace — static factory methods for guards.
211pub struct G;
212
213impl G {
214    /// Length guard — output must be within bounds.
215    pub fn length(min: usize, max: usize) -> GuardRule {
216        GuardRule::new("length", move |output| {
217            let len = output.len();
218            if len < min {
219                Err(format!("Output too short: {len} < {min}"))
220            } else if len > max {
221                Err(format!("Output too long: {len} > {max}"))
222            } else {
223                Ok(())
224            }
225        })
226    }
227
228    /// Regex guard — output must match (or not match) a pattern.
229    pub fn regex(pattern: &str) -> GuardRule {
230        let pattern = pattern.to_string();
231        GuardRule::new("regex", move |output| {
232            // Simple substring check — full regex requires the `regex` crate.
233            if output.contains(&pattern) {
234                Err(format!("Output matches forbidden pattern: {pattern}"))
235            } else {
236                Ok(())
237            }
238        })
239    }
240
241    /// Budget guard — output must not exceed a token estimate.
242    pub fn budget(max_tokens: usize) -> GuardRule {
243        GuardRule::new("budget", move |output| {
244            // Rough estimate: 4 chars per token.
245            let estimated_tokens = output.len() / 4;
246            if estimated_tokens > max_tokens {
247                Err(format!(
248                    "Output exceeds token budget: ~{estimated_tokens} > {max_tokens}"
249                ))
250            } else {
251                Ok(())
252            }
253        })
254    }
255
256    /// JSON guard — output must be valid JSON.
257    pub fn json() -> GuardRule {
258        GuardRule::new("json", |output| {
259            serde_json::from_str::<serde_json::Value>(output)
260                .map(|_| ())
261                .map_err(|e| format!("Invalid JSON: {e}"))
262        })
263    }
264
265    /// PII guard — checks for common PII patterns (email, phone).
266    pub fn pii() -> GuardRule {
267        GuardRule::new("pii", |output| {
268            // Simple heuristic checks for common PII patterns.
269            if output.contains('@') && output.contains('.') {
270                // Might be an email — flag it.
271                return Err("Output may contain email addresses".to_string());
272            }
273            Ok(())
274        })
275    }
276
277    /// Topic restriction guard — output must not mention denied topics.
278    pub fn topic(deny: &[&str]) -> GuardRule {
279        let deny: Vec<String> = deny.iter().map(|s| s.to_lowercase()).collect();
280        GuardRule::new("topic", move |output| {
281            let lower = output.to_lowercase();
282            for topic in &deny {
283                if lower.contains(topic.as_str()) {
284                    return Err(format!("Output mentions denied topic: {topic}"));
285                }
286            }
287            Ok(())
288        })
289    }
290
291    /// Custom guard from a validation function.
292    pub fn custom(f: impl Fn(&str) -> Result<(), String> + Send + Sync + 'static) -> GuardRule {
293        GuardRule::new("custom", f)
294    }
295
296    /// Toxicity guard — flags toxic/abusive output using an LLM judge.
297    ///
298    /// Vetoes the response if the judge model decides it contains toxic, hateful,
299    /// harassing, sexual, or abusive content (mirrors ADK's safety evaluation,
300    /// but runs locally against the provided judge LLM).
301    pub fn toxicity(judge: Arc<dyn BaseLlm>) -> GuardRule {
302        GuardRule::judge(
303            "toxicity",
304            LlmJudge::new(
305                judge,
306                "the RESPONSE contains toxic, hateful, harassing, sexual, or otherwise \
307                 abusive or unsafe content",
308            ),
309        )
310    }
311
312    /// Grounding guard — flags output not supported by the conversation context.
313    ///
314    /// The judge sees the model's input history as CONTEXT and vetoes the
315    /// response if it makes factual claims not supported by that context.
316    pub fn grounded(judge: Arc<dyn BaseLlm>) -> GuardRule {
317        GuardRule::judge(
318            "grounded",
319            LlmJudge::new(
320                judge,
321                "the RESPONSE asserts facts that are NOT supported by, or that \
322                 contradict, the provided CONTEXT",
323            )
324            .with_context("CONTEXT"),
325        )
326    }
327
328    /// Hallucination guard — flags fabricated/unverifiable claims via an LLM judge.
329    pub fn hallucination(judge: Arc<dyn BaseLlm>) -> GuardRule {
330        GuardRule::judge(
331            "hallucination",
332            LlmJudge::new(
333                judge,
334                "the RESPONSE contains fabricated, invented, or unverifiable facts \
335                 that are not supported by the CONTEXT",
336            )
337            .with_context("CONTEXT"),
338        )
339    }
340
341    /// Conditional guard — only applies `inner` when `predicate` returns true.
342    pub fn when(
343        predicate: impl Fn(&str) -> bool + Send + Sync + 'static,
344        inner: GuardRule,
345    ) -> GuardRule {
346        GuardRule::new("when", move |output| {
347            if predicate(output) {
348                inner.check(output)
349            } else {
350                Ok(())
351            }
352        })
353    }
354
355    /// LLM-as-judge content guard.
356    ///
357    /// `rubric` describes the condition that constitutes a *violation*; the judge
358    /// model vetoes the response when that condition holds. Example:
359    /// `G::llm_judge(llm, "the response gives medical advice without a disclaimer")`.
360    pub fn llm_judge(judge: Arc<dyn BaseLlm>, rubric: impl Into<String>) -> GuardRule {
361        GuardRule::judge("llm_judge", LlmJudge::new(judge, rubric))
362    }
363
364    /// Named custom judge function guard.
365    pub fn custom_judge(
366        name: &str,
367        f: impl Fn(&str) -> Result<(), String> + Send + Sync + 'static,
368    ) -> GuardRule {
369        // Leak the name to get a 'static str, matching the GuardRule field type.
370        let name: &'static str = Box::leak(name.to_string().into_boxed_str());
371        GuardRule::new(name, f)
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    use super::*;
378
379    #[test]
380    fn length_guard_passes() {
381        assert!(G::length(1, 100).check("hello").is_ok());
382    }
383
384    #[test]
385    fn length_guard_too_short() {
386        assert!(G::length(10, 100).check("hi").is_err());
387    }
388
389    #[test]
390    fn length_guard_too_long() {
391        assert!(G::length(1, 5).check("too long text").is_err());
392    }
393
394    #[test]
395    fn json_guard_valid() {
396        assert!(G::json().check(r#"{"key": "value"}"#).is_ok());
397    }
398
399    #[test]
400    fn json_guard_invalid() {
401        assert!(G::json().check("not json").is_err());
402    }
403
404    #[test]
405    fn regex_guard_blocks() {
406        assert!(G::regex("secret").check("this is a secret").is_err());
407    }
408
409    #[test]
410    fn regex_guard_passes() {
411        assert!(G::regex("secret").check("this is public").is_ok());
412    }
413
414    #[test]
415    fn budget_guard_passes() {
416        assert!(G::budget(100).check("short").is_ok());
417    }
418
419    #[test]
420    fn topic_guard_blocks() {
421        assert!(G::topic(&["violence"]).check("There was violence").is_err());
422    }
423
424    #[test]
425    fn topic_guard_passes() {
426        assert!(G::topic(&["violence"]).check("A peaceful day").is_ok());
427    }
428
429    #[test]
430    fn compose_with_bitor() {
431        let composite = G::length(1, 1000) + G::json();
432        assert_eq!(composite.len(), 2);
433    }
434
435    #[test]
436    fn check_all_returns_violations() {
437        let composite = G::length(1, 5) + G::json();
438        let violations = composite.check_all("not json and too long text here");
439        assert!(!violations.is_empty());
440    }
441
442    #[test]
443    fn custom_guard() {
444        let g = G::custom(|output| {
445            if output.contains("bad") {
446                Err("Contains 'bad'".into())
447            } else {
448                Ok(())
449            }
450        });
451        assert!(g.check("good output").is_ok());
452        assert!(g.check("bad output").is_err());
453    }
454
455    // A no-op judge LLM for constructing LLM-backed guards in unit tests
456    // (these tests exercise composition/naming, not the judge call itself).
457    fn judge_llm() -> Arc<dyn BaseLlm> {
458        use gemini_adk_rs::llm::{LlmError, LlmResponse};
459        use gemini_genai_rs::prelude::{Content, Part, Role};
460
461        struct NoopJudge;
462        #[async_trait]
463        impl BaseLlm for NoopJudge {
464            fn model_id(&self) -> &str {
465                "noop-judge"
466            }
467            async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
468                Ok(LlmResponse {
469                    content: Content {
470                        role: Some(Role::Model),
471                        parts: vec![Part::Text {
472                            text: r#"{"violation": false, "reason": "ok"}"#.to_string(),
473                        }],
474                    },
475                    finish_reason: Some("STOP".into()),
476                    usage: None,
477                })
478            }
479        }
480        Arc::new(NoopJudge)
481    }
482
483    #[test]
484    fn toxicity_guard() {
485        let g = G::toxicity(judge_llm());
486        // Sync path is a no-op for judge guards.
487        assert!(g.check("anything").is_ok());
488        assert_eq!(g.name(), "toxicity");
489    }
490
491    #[test]
492    fn grounded_guard() {
493        let g = G::grounded(judge_llm());
494        assert!(g.check("anything").is_ok());
495        assert_eq!(g.name(), "grounded");
496    }
497
498    #[test]
499    fn hallucination_guard() {
500        let g = G::hallucination(judge_llm());
501        assert!(g.check("anything").is_ok());
502        assert_eq!(g.name(), "hallucination");
503    }
504
505    #[tokio::test]
506    async fn judge_guard_runs_async() {
507        // A judge that flags everything should produce a violation via check_async.
508        use gemini_adk_rs::llm::{LlmError, LlmResponse};
509        use gemini_genai_rs::prelude::{Content, Part, Role};
510        struct FlagAll;
511        #[async_trait]
512        impl BaseLlm for FlagAll {
513            fn model_id(&self) -> &str {
514                "flag-all"
515            }
516            async fn generate(&self, _req: LlmRequest) -> Result<LlmResponse, LlmError> {
517                Ok(LlmResponse {
518                    content: Content {
519                        role: Some(Role::Model),
520                        parts: vec![Part::Text {
521                            text: r#"{"violation": true, "reason": "bad"}"#.to_string(),
522                        }],
523                    },
524                    finish_reason: Some("STOP".into()),
525                    usage: None,
526                })
527            }
528        }
529        let g = G::toxicity(Arc::new(FlagAll));
530        assert!(g.check_async("hello", None).await.is_err());
531    }
532
533    #[test]
534    fn when_guard_applies() {
535        let inner = G::length(1, 5);
536        let g = G::when(|output| output.starts_with("check:"), inner);
537        // Predicate true — inner guard runs and rejects long output.
538        assert!(g.check("check: this is way too long").is_err());
539        // Predicate false — inner guard skipped.
540        assert!(g.check("skip: this is way too long").is_ok());
541        assert_eq!(g.name(), "when");
542    }
543
544    #[test]
545    fn llm_judge_guard() {
546        let g = G::llm_judge(judge_llm(), "the response is unhelpful");
547        assert!(g.check("anything").is_ok());
548        assert_eq!(g.name(), "llm_judge");
549    }
550
551    #[test]
552    fn custom_judge_guard() {
553        let g = G::custom_judge("profanity_filter", |output| {
554            if output.contains("bad_word") {
555                Err("Profanity detected".into())
556            } else {
557                Ok(())
558            }
559        });
560        assert!(g.check("clean text").is_ok());
561        assert!(g.check("has bad_word here").is_err());
562        assert_eq!(g.name(), "profanity_filter");
563    }
564
565    #[test]
566    fn compose_new_guards_with_bitor() {
567        let composite =
568            G::toxicity(judge_llm()) + G::grounded(judge_llm()) + G::hallucination(judge_llm());
569        assert_eq!(composite.len(), 3);
570        // Sync path skips judge guards, so no violations surface synchronously.
571        assert!(composite.check_all("test").is_empty());
572    }
573}
gemini_adk_fluent_rs/compose/guards.rs

gemini_adk_fluent_rs/compose/
guards.rs
