tirea_agentos/composition/
bundle.rs1use 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#[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}