gemini_adk_fluent_rs/compose/
guards.rs

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

gemini_adk_fluent_rs/compose/
guards.rs
