Skip to main content

af_agent_runtime/
builder.rs

1use std::collections::BTreeSet;
2use std::sync::Arc;
3
4use af_agent::{
5    ChatModel, ContextContributor, Hook, MountedPlugins, PluginError, PromptAuthority,
6    PromptSection, ToolRegistry,
7};
8
9use crate::{
10    AgentRuntime, ApproximateTokenMeter, Compactor, ModelCompactor, RuntimeLimits, TokenMeter,
11};
12
13impl AgentRuntime {
14    /// Tighten the per-request image budget without enlarging host limits.
15    pub fn with_image_limit(mut self, limit: usize) -> Self {
16        self.limits.max_images = self
17            .limits
18            .max_images
19            .min(limit)
20            .min(af_llm::images::MAX_INPUT_IMAGES);
21        self
22    }
23
24    /// Runtime for `model` under `model_name` with `system_prompt` as the product prompt section and default limits.
25    pub fn new(
26        model: Arc<dyn ChatModel>,
27        model_name: impl Into<String>,
28        system_prompt: impl Into<String>,
29    ) -> Self {
30        let mut prompts = af_agent::PromptRegistry::default();
31        prompts
32            .insert(PromptSection {
33                id: "product.base".into(),
34                order: 100,
35                authority: PromptAuthority::Product,
36                source: "product".into(),
37                version: env!("CARGO_PKG_VERSION").into(),
38                content: system_prompt.into(),
39            })
40            .expect("base prompt is validated by ProductRuntime");
41        Self {
42            compactor: Arc::new(ModelCompactor::new(Arc::clone(&model))),
43            model,
44            model_name: model_name.into(),
45            prompts,
46            contexts: Vec::new(),
47            tools: ToolRegistry::new(),
48            hooks: Vec::new(),
49            limits: RuntimeLimits::default(),
50            meter: Arc::new(ApproximateTokenMeter),
51            plugins: None,
52            reasoning_effort: None,
53            output_policy: None,
54        }
55    }
56
57    /// Replace the tool registry.
58    pub fn with_tools(mut self, tools: ToolRegistry) -> Self {
59        self.tools = tools;
60        self
61    }
62
63    /// Install hooks; they run in this order.
64    pub fn with_hooks(mut self, hooks: Vec<Arc<dyn Hook>>) -> Self {
65        self.hooks = hooks;
66        self
67    }
68
69    /// Remove denied tools after assembly so they are neither advertised nor executable.
70    pub fn with_denied_tools(mut self, denied: &BTreeSet<String>) -> Self {
71        self.tools = self.tools.filtered(
72            self.tools
73                .names()
74                .into_iter()
75                .filter(|name| !denied.contains(*name)),
76        );
77        self
78    }
79
80    /// Replace the runtime limits.
81    pub fn with_limits(mut self, limits: RuntimeLimits) -> Self {
82        self.limits = limits;
83        self
84    }
85
86    /// Replace the token meter.
87    pub fn with_token_meter(mut self, meter: Arc<dyn TokenMeter>) -> Self {
88        self.meter = meter;
89        self
90    }
91
92    /// Replace the compactor.
93    pub fn with_compactor(mut self, compactor: Arc<dyn Compactor>) -> Self {
94        self.compactor = compactor;
95        self
96    }
97
98    /// Pin the reasoning effort for every model request.
99    pub fn with_reasoning_effort(mut self, effort: af_llm::ReasoningEffort) -> Self {
100        self.reasoning_effort = Some(effort);
101        self
102    }
103
104    /// Validate assistant output against this JSON Schema before it is persisted.
105    pub fn with_output_policy(mut self, schema: serde_json::Value) -> Result<Self, PluginError> {
106        af_agent::validate_json_schema_definition(&schema).map_err(PluginError::Conflict)?;
107        self.output_policy = Some(schema);
108        Ok(self)
109    }
110
111    /// Add a prompt section; duplicate ids are rejected.
112    pub fn with_prompt_section(mut self, section: PromptSection) -> Result<Self, PluginError> {
113        self.prompts
114            .insert(section)
115            .map_err(PluginError::Conflict)?;
116        Ok(self)
117    }
118
119    /// Add a per-step context contributor.
120    pub fn with_context_contributor(mut self, context: Arc<dyn ContextContributor>) -> Self {
121        self.contexts.push(context);
122        self
123    }
124
125    /// Adopt the prompts, tools, hooks and contributors of mounted plugins.
126    pub fn with_plugins(mut self, mounted: MountedPlugins) -> Result<Self, PluginError> {
127        self.tools
128            .extend(mounted.tools())
129            .map_err(PluginError::Conflict)?;
130        self.hooks.extend(mounted.hooks().iter().cloned());
131        self.contexts.extend(mounted.contexts().iter().cloned());
132        for prompt in mounted.prompts() {
133            self.prompts
134                .insert(prompt.clone())
135                .map_err(PluginError::Conflict)?;
136        }
137        self.plugins = Some(mounted);
138        Ok(self)
139    }
140
141    /// Narrow the tool registry to `allowed`; `None` keeps every tool.
142    pub fn with_allowed_tools(mut self, allowed: Option<&BTreeSet<String>>) -> Self {
143        if let Some(allowed) = allowed {
144            self.tools = self.tools.filtered(allowed.iter().map(String::as_str));
145        }
146        self
147    }
148}