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    /// Runtime for `model` under `model_name` with `system_prompt` as the product prompt section and default limits.
15    pub fn new(
16        model: Arc<dyn ChatModel>,
17        model_name: impl Into<String>,
18        system_prompt: impl Into<String>,
19    ) -> Self {
20        let mut prompts = af_agent::PromptRegistry::default();
21        prompts
22            .insert(PromptSection {
23                id: "product.base".into(),
24                order: 100,
25                authority: PromptAuthority::Product,
26                source: "product".into(),
27                version: env!("CARGO_PKG_VERSION").into(),
28                content: system_prompt.into(),
29            })
30            .expect("base prompt is validated by ProductRuntime");
31        Self {
32            compactor: Arc::new(ModelCompactor::new(Arc::clone(&model))),
33            model,
34            model_name: model_name.into(),
35            prompts,
36            contexts: Vec::new(),
37            tools: ToolRegistry::new(),
38            hooks: Vec::new(),
39            limits: RuntimeLimits::default(),
40            meter: Arc::new(ApproximateTokenMeter),
41            plugins: None,
42            reasoning_effort: None,
43            output_policy: None,
44        }
45    }
46
47    /// Replace the tool registry.
48    pub fn with_tools(mut self, tools: ToolRegistry) -> Self {
49        self.tools = tools;
50        self
51    }
52
53    /// Install hooks; they run in this order.
54    pub fn with_hooks(mut self, hooks: Vec<Arc<dyn Hook>>) -> Self {
55        self.hooks = hooks;
56        self
57    }
58
59    /// Replace the runtime limits.
60    pub fn with_limits(mut self, limits: RuntimeLimits) -> Self {
61        self.limits = limits;
62        self
63    }
64
65    /// Replace the token meter.
66    pub fn with_token_meter(mut self, meter: Arc<dyn TokenMeter>) -> Self {
67        self.meter = meter;
68        self
69    }
70
71    /// Replace the compactor.
72    pub fn with_compactor(mut self, compactor: Arc<dyn Compactor>) -> Self {
73        self.compactor = compactor;
74        self
75    }
76
77    /// Pin the reasoning effort for every model request.
78    pub fn with_reasoning_effort(mut self, effort: af_llm::ReasoningEffort) -> Self {
79        self.reasoning_effort = Some(effort);
80        self
81    }
82
83    /// Validate assistant output against this JSON Schema before it is persisted.
84    pub fn with_output_policy(mut self, schema: serde_json::Value) -> Result<Self, PluginError> {
85        af_agent::validate_json_schema_definition(&schema).map_err(PluginError::Conflict)?;
86        self.output_policy = Some(schema);
87        Ok(self)
88    }
89
90    /// Add a prompt section; duplicate ids are rejected.
91    pub fn with_prompt_section(mut self, section: PromptSection) -> Result<Self, PluginError> {
92        self.prompts
93            .insert(section)
94            .map_err(PluginError::Conflict)?;
95        Ok(self)
96    }
97
98    /// Add a per-step context contributor.
99    pub fn with_context_contributor(mut self, context: Arc<dyn ContextContributor>) -> Self {
100        self.contexts.push(context);
101        self
102    }
103
104    /// Adopt the prompts, tools, hooks and contributors of mounted plugins.
105    pub fn with_plugins(mut self, mounted: MountedPlugins) -> Result<Self, PluginError> {
106        self.tools
107            .extend(mounted.tools())
108            .map_err(PluginError::Conflict)?;
109        self.hooks.extend(mounted.hooks().iter().cloned());
110        self.contexts.extend(mounted.contexts().iter().cloned());
111        for prompt in mounted.prompts() {
112            self.prompts
113                .insert(prompt.clone())
114                .map_err(PluginError::Conflict)?;
115        }
116        self.plugins = Some(mounted);
117        Ok(self)
118    }
119
120    /// Narrow the tool registry to `allowed`; `None` keeps every tool.
121    pub fn with_allowed_tools(mut self, allowed: Option<&BTreeSet<String>>) -> Self {
122        if let Some(allowed) = allowed {
123            self.tools = self.tools.filtered(allowed.iter().map(String::as_str));
124        }
125        self
126    }
127}