Skip to main content

rs_adk/plugin/
context_filter.rs

1//! Context filter plugin — filters conversation context before model calls.
2//!
3//! Mirrors ADK-Python's `context_filter_plugin`. Allows trimming,
4//! filtering, or transforming the conversation history before it
5//! reaches the model.
6
7use async_trait::async_trait;
8
9use super::{Plugin, PluginResult};
10use crate::context::InvocationContext;
11
12/// Plugin that filters or transforms conversation context before model calls.
13///
14/// Can be used to:
15/// - Limit context window size
16/// - Remove tool call/result turns
17/// - Filter out specific content types
18/// - Inject additional context
19pub struct ContextFilterPlugin {
20    /// Maximum number of turns to keep in context.
21    max_turns: Option<usize>,
22    /// Whether to exclude tool call/result turns from context.
23    exclude_tool_turns: bool,
24}
25
26impl ContextFilterPlugin {
27    /// Create a new context filter plugin.
28    pub fn new() -> Self {
29        Self {
30            max_turns: None,
31            exclude_tool_turns: false,
32        }
33    }
34
35    /// Set the maximum number of turns to keep.
36    pub fn with_max_turns(mut self, max_turns: usize) -> Self {
37        self.max_turns = Some(max_turns);
38        self
39    }
40
41    /// Set whether to exclude tool call/result turns.
42    pub fn with_exclude_tool_turns(mut self, exclude: bool) -> Self {
43        self.exclude_tool_turns = exclude;
44        self
45    }
46
47    /// Returns the configured max turns limit.
48    pub fn max_turns(&self) -> Option<usize> {
49        self.max_turns
50    }
51
52    /// Returns whether tool turns are excluded.
53    pub fn exclude_tool_turns(&self) -> bool {
54        self.exclude_tool_turns
55    }
56}
57
58impl Default for ContextFilterPlugin {
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64#[async_trait]
65impl Plugin for ContextFilterPlugin {
66    fn name(&self) -> &str {
67        "context_filter"
68    }
69
70    async fn before_model(
71        &self,
72        _request: &crate::llm::LlmRequest,
73        _ctx: &InvocationContext,
74    ) -> PluginResult {
75        // The actual context filtering is applied by the runtime
76        // when constructing the LLM request, using this plugin's
77        // configuration. The plugin signals Continue to allow
78        // the (already filtered) request to proceed.
79        PluginResult::Continue
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn default_config() {
89        let plugin = ContextFilterPlugin::new();
90        assert!(plugin.max_turns().is_none());
91        assert!(!plugin.exclude_tool_turns());
92    }
93
94    #[test]
95    fn custom_config() {
96        let plugin = ContextFilterPlugin::new()
97            .with_max_turns(10)
98            .with_exclude_tool_turns(true);
99        assert_eq!(plugin.max_turns(), Some(10));
100        assert!(plugin.exclude_tool_turns());
101    }
102
103    #[test]
104    fn plugin_name() {
105        let plugin = ContextFilterPlugin::new();
106        assert_eq!(plugin.name(), "context_filter");
107    }
108}