tirea_agentos/composition/
bundle.rs

1use super::registry::{
2    AgentRegistry, BehaviorRegistry, ModelDefinition, ModelRegistry, ProviderRegistry,
3    RegistryBundle, ToolRegistry,
4};
5use crate::composition::AgentDefinition;
6use crate::contracts::runtime::tool_call::Tool;
7use crate::contracts::runtime::AgentBehavior;
8use genai::Client;
9use std::collections::HashMap;
10use std::sync::Arc;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum BundleRegistryKind {
14    Agent,
15    Tool,
16    Behavior,
17    Provider,
18    Model,
19}
20
21impl std::fmt::Display for BundleRegistryKind {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        match self {
24            Self::Agent => write!(f, "agent"),
25            Self::Tool => write!(f, "tool"),
26            Self::Behavior => write!(f, "behavior"),
27            Self::Provider => write!(f, "provider"),
28            Self::Model => write!(f, "model"),
29        }
30    }
31}
32
33#[derive(Debug, thiserror::Error)]
34pub enum BundleComposeError {
35    #[error("bundle '{bundle_id}' contributed duplicate {kind} id: {id}")]
36    DuplicateId {
37        bundle_id: String,
38        kind: BundleRegistryKind,
39        id: String,
40    },
41
42    #[error("bundle '{bundle_id}' contributed empty {kind} id")]
43    EmptyId {
44        bundle_id: String,
45        kind: BundleRegistryKind,
46    },
47}
48
49pub struct BundleRegistryAccumulator<'a> {
50    pub agent_definitions: &'a mut HashMap<String, AgentDefinition>,
51    pub agent_registries: &'a mut Vec<Arc<dyn AgentRegistry>>,
52    pub tool_definitions: &'a mut HashMap<String, Arc<dyn Tool>>,
53    pub tool_registries: &'a mut Vec<Arc<dyn ToolRegistry>>,
54    pub behavior_definitions: &'a mut HashMap<String, Arc<dyn AgentBehavior>>,
55    pub behavior_registries: &'a mut Vec<Arc<dyn BehaviorRegistry>>,
56    pub provider_definitions: &'a mut HashMap<String, Client>,
57    pub provider_registries: &'a mut Vec<Arc<dyn ProviderRegistry>>,
58    pub model_definitions: &'a mut HashMap<String, ModelDefinition>,
59    pub model_registries: &'a mut Vec<Arc<dyn ModelRegistry>>,
60}
61
62pub struct BundleComposer;
63
64impl BundleComposer {
65    pub fn apply(
66        bundles: &[Arc<dyn RegistryBundle>],
67        acc: BundleRegistryAccumulator<'_>,
68    ) -> Result<(), BundleComposeError> {
69        let BundleRegistryAccumulator {
70            agent_definitions,
71            agent_registries,
72            tool_definitions,
73            tool_registries,
74            behavior_definitions,
75            behavior_registries,
76            provider_definitions,
77            provider_registries,
78            model_definitions,
79            model_registries,
80        } = acc;
81
82        for bundle in bundles {
83            let bundle_id = bundle.id().to_string();
84            Self::merge_named(
85                agent_definitions,
86                bundle.agent_definitions(),
87                &bundle_id,
88                BundleRegistryKind::Agent,
89            )?;
90            agent_registries.extend(bundle.agent_registries());
91
92            Self::merge_named(
93                tool_definitions,
94                bundle.tool_definitions(),
95                &bundle_id,
96                BundleRegistryKind::Tool,
97            )?;
98            tool_registries.extend(bundle.tool_registries());
99
100            Self::merge_named(
101                behavior_definitions,
102                bundle.behavior_definitions(),
103                &bundle_id,
104                BundleRegistryKind::Behavior,
105            )?;
106            behavior_registries.extend(bundle.behavior_registries());
107
108            Self::merge_named(
109                provider_definitions,
110                bundle.provider_definitions(),
111                &bundle_id,
112                BundleRegistryKind::Provider,
113            )?;
114            provider_registries.extend(bundle.provider_registries());
115
116            Self::merge_named(
117                model_definitions,
118                bundle.model_definitions(),
119                &bundle_id,
120                BundleRegistryKind::Model,
121            )?;
122            model_registries.extend(bundle.model_registries());
123        }
124
125        Ok(())
126    }
127
128    fn merge_named<V>(
129        target: &mut HashMap<String, V>,
130        incoming: HashMap<String, V>,
131        bundle_id: &str,
132        kind: BundleRegistryKind,
133    ) -> Result<(), BundleComposeError> {
134        for (id, value) in incoming {
135            if id.trim().is_empty() {
136                return Err(BundleComposeError::EmptyId {
137                    bundle_id: bundle_id.to_string(),
138                    kind,
139                });
140            }
141            if target.contains_key(&id) {
142                return Err(BundleComposeError::DuplicateId {
143                    bundle_id: bundle_id.to_string(),
144                    kind,
145                    id,
146                });
147            }
148            target.insert(id, value);
149        }
150        Ok(())
151    }
152}
153
154/// Lightweight bundle carrying only tools/behaviors contributions.
155///
156/// This is useful for runtime/system wiring where agent/model/provider registries
157/// should not be mutated, but tools/behaviors still need a uniform bundle path.
158#[derive(Clone, Default)]
159pub struct ToolBehaviorBundle {
160    id: String,
161    tools: HashMap<String, Arc<dyn Tool>>,
162    behaviors: HashMap<String, Arc<dyn AgentBehavior>>,
163}
164
165impl ToolBehaviorBundle {
166    pub fn new(id: impl Into<String>) -> Self {
167        Self {
168            id: id.into(),
169            ..Self::default()
170        }
171    }
172
173    pub fn with_tool(mut self, tool: Arc<dyn Tool>) -> Self {
174        self.tools.insert(tool.descriptor().id, tool);
175        self
176    }
177
178    pub fn with_tools(mut self, tools: HashMap<String, Arc<dyn Tool>>) -> Self {
179        self.tools.extend(tools);
180        self
181    }
182
183    pub fn with_behavior(mut self, behavior: Arc<dyn AgentBehavior>) -> Self {
184        self.behaviors.insert(behavior.id().to_string(), behavior);
185        self
186    }
187}
188
189impl RegistryBundle for ToolBehaviorBundle {
190    fn id(&self) -> &str {
191        &self.id
192    }
193
194    fn tool_definitions(&self) -> HashMap<String, Arc<dyn Tool>> {
195        self.tools.clone()
196    }
197
198    fn behavior_definitions(&self) -> HashMap<String, Arc<dyn AgentBehavior>> {
199        self.behaviors.clone()
200    }
201}