1use std::sync::Arc;
4
5use starweaver_context::{AgentContext, AgentInfo, ModelConfig, ToolConfig};
6use starweaver_core::{AgentId, CancellationToken};
7use starweaver_model::{ModelAdapter, ModelRequestParameters, ModelSettings, OutputMode};
8use starweaver_tools::{DynToolset, ToolRegistry};
9
10use starweaver_usage::UsageLimits;
11
12use crate::{
13 agent::helpers::merge_request_params,
14 capability::{resolve_capability_order, AgentCapability, CapabilityBundle},
15 executor::{DirectAgentExecutor, DynAgentExecutor},
16 graph::{inspect_graph, AgentGraphTrace, AgentNode, GraphError},
17 instructions::DynDynamicInstruction,
18 output::{
19 DynOutputFunction, OutputPolicy, OutputSchema, OutputValidator, SchemaOutputFunction,
20 },
21 trace::{DynTraceRecorder, NoopTraceRecorder},
22};
23
24mod helpers;
25mod overrides;
26mod run_loop;
27mod run_loop_helpers;
28mod runtime_helpers;
29mod types;
30
31pub use overrides::AgentOverride;
32pub use types::{AgentEndStrategy, AgentError, AgentInput, AgentResult, AgentRuntimePolicy};
33
34#[derive(Clone)]
36pub struct Agent {
37 default_id: AgentId,
38 default_name: String,
39 model: Arc<dyn ModelAdapter>,
40 instructions: Vec<String>,
41 dynamic_instructions: Vec<DynDynamicInstruction>,
42 model_settings: Option<ModelSettings>,
43 request_params: ModelRequestParameters,
44 output_schema: Option<OutputSchema>,
45 output_validators: Vec<Arc<dyn OutputValidator>>,
46 output_functions: Vec<DynOutputFunction>,
47 usage_limits: Option<UsageLimits>,
48 tools: ToolRegistry,
49 toolsets: Vec<DynToolset>,
50 capabilities: Vec<Arc<dyn AgentCapability>>,
51 stream_observers: Vec<Arc<dyn AgentCapability>>,
52 cancellation_token: Option<CancellationToken>,
53 executor: DynAgentExecutor,
54 trace_recorder: DynTraceRecorder,
55 policy: AgentRuntimePolicy,
56 model_config: Option<ModelConfig>,
57 tool_config: Option<ToolConfig>,
58}
59
60impl Agent {
61 #[must_use]
63 pub fn new(model: Arc<dyn ModelAdapter>) -> Self {
64 Self {
65 default_id: AgentId::default(),
66 default_name: AgentId::default().as_str().to_string(),
67 model,
68 instructions: Vec::new(),
69 dynamic_instructions: Vec::new(),
70 model_settings: None,
71 request_params: ModelRequestParameters::default(),
72 output_schema: None,
73 output_validators: Vec::new(),
74 output_functions: Vec::new(),
75 usage_limits: None,
76 tools: ToolRegistry::new(),
77 toolsets: Vec::new(),
78 capabilities: Vec::new(),
79 stream_observers: Vec::new(),
80 cancellation_token: None,
81 executor: Arc::new(DirectAgentExecutor),
82 trace_recorder: Arc::new(NoopTraceRecorder),
83 policy: AgentRuntimePolicy::default(),
84 model_config: None,
85 tool_config: None,
86 }
87 }
88
89 #[must_use]
91 pub const fn agent_id(&self) -> &AgentId {
92 &self.default_id
93 }
94
95 #[must_use]
97 pub fn agent_name(&self) -> &str {
98 &self.default_name
99 }
100
101 #[must_use]
103 pub fn with_agent_id(mut self, agent_id: AgentId) -> Self {
104 self.default_id = agent_id;
105 self
106 }
107
108 #[must_use]
110 pub fn with_agent_name(mut self, agent_name: impl Into<String>) -> Self {
111 self.default_name = agent_name.into();
112 self
113 }
114
115 #[must_use]
117 pub fn with_agent_identity(mut self, agent_id: AgentId, agent_name: impl Into<String>) -> Self {
118 self.default_id = agent_id;
119 self.default_name = agent_name.into();
120 self
121 }
122
123 #[must_use]
125 pub fn new_context(&self) -> AgentContext {
126 let mut context = AgentContext::new(self.default_id.clone());
127 self.apply_context_identity(&mut context);
128 context
129 }
130
131 fn apply_context_identity(&self, context: &mut AgentContext) {
132 context.agent_registry.insert(
133 self.default_id.as_str().to_string(),
134 AgentInfo::new(self.default_id.as_str(), self.default_name.clone()),
135 );
136 if self.default_name != self.default_id.as_str() {
137 context.metadata.insert(
138 "agent_name".to_string(),
139 serde_json::json!(self.default_name.as_str()),
140 );
141 }
142 }
143
144 #[must_use]
146 pub fn with_instruction(mut self, instruction: impl Into<String>) -> Self {
147 self.instructions.push(instruction.into());
148 self
149 }
150
151 #[must_use]
153 pub fn with_dynamic_instruction(mut self, instruction: DynDynamicInstruction) -> Self {
154 self.dynamic_instructions.push(instruction);
155 self
156 }
157
158 #[must_use]
160 pub fn with_model_settings(mut self, settings: ModelSettings) -> Self {
161 self.model_settings = Some(settings);
162 self
163 }
164
165 #[must_use]
167 pub fn with_request_params(mut self, params: ModelRequestParameters) -> Self {
168 self.request_params = params;
169 self
170 }
171
172 #[must_use]
174 pub fn with_tools(mut self, tools: ToolRegistry) -> Self {
175 self.tools = tools;
176 self
177 }
178
179 #[must_use]
181 pub fn with_toolset(mut self, toolset: DynToolset) -> Self {
182 self.toolsets.push(toolset);
183 self
184 }
185
186 #[must_use]
188 pub fn with_toolsets(mut self, toolsets: impl IntoIterator<Item = DynToolset>) -> Self {
189 self.toolsets.extend(toolsets);
190 self
191 }
192
193 #[must_use]
195 pub fn with_appended_tools(mut self, tools: &ToolRegistry) -> Self {
196 self.tools.insert_registry(tools);
197 self
198 }
199
200 #[must_use]
202 pub fn tools(&self) -> ToolRegistry {
203 let mut tools = self.tools.clone();
204 for toolset in &self.toolsets {
205 tools.insert_toolset(toolset);
206 }
207 tools
208 }
209
210 #[must_use]
212 pub fn with_tool_retries(mut self, max_retries: usize) -> Self {
213 self.tools.set_max_retries(max_retries);
214 self
215 }
216
217 #[must_use]
219 pub fn with_output_schema(mut self, schema: OutputSchema) -> Self {
220 self.output_schema = Some(schema);
221 self
222 }
223
224 #[must_use]
226 pub fn with_output_policy(mut self, policy: OutputPolicy) -> Self {
227 let (schema, validators, functions, retries, mode, allow_text_output, allow_image_output) =
228 policy.into_parts();
229 let schema_output_function = match (schema.as_ref(), mode) {
230 (Some(schema), Some(OutputMode::Tool | OutputMode::ToolOrText)) => {
231 Some(Arc::new(SchemaOutputFunction::new(schema.clone())) as DynOutputFunction)
232 }
233 _ => None,
234 };
235 if let Some(schema) = schema {
236 self.output_schema = Some(schema);
237 }
238 self.output_validators.extend(validators);
239 if let Some(function) = schema_output_function {
240 self.output_functions.push(function);
241 }
242 self.output_functions.extend(functions);
243 if let Some(retries) = retries {
244 self.policy.output_retries = retries;
245 }
246 if let Some(mode) = mode {
247 self.request_params.output_mode = Some(mode);
248 }
249 if let Some(allow_text_output) = allow_text_output {
250 self.request_params.allow_text_output = Some(allow_text_output);
251 }
252 if let Some(allow_image_output) = allow_image_output {
253 self.request_params.allow_image_output = Some(allow_image_output);
254 }
255 self
256 }
257
258 #[must_use]
260 pub fn with_output_validator(mut self, validator: Arc<dyn OutputValidator>) -> Self {
261 self.output_validators.push(validator);
262 self
263 }
264
265 #[must_use]
267 pub fn with_output_function(mut self, function: DynOutputFunction) -> Self {
268 self.output_functions.push(function);
269 self
270 }
271
272 #[must_use]
274 pub const fn with_usage_limits(mut self, limits: UsageLimits) -> Self {
275 self.usage_limits = Some(limits);
276 self
277 }
278
279 #[must_use]
281 pub fn with_model_config(mut self, model_config: ModelConfig) -> Self {
282 self.model_config = Some(model_config);
283 self
284 }
285
286 #[must_use]
288 pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
289 self.tool_config = Some(tool_config);
290 self
291 }
292
293 #[must_use]
295 pub fn with_context_window(mut self, context_window: u64) -> Self {
296 let mut model_config = self.model_config.unwrap_or_default();
297 model_config.context_window = Some(context_window);
298 self.model_config = Some(model_config);
299 self
300 }
301
302 #[must_use]
304 pub fn with_capability(mut self, capability: Arc<dyn AgentCapability>) -> Self {
305 self.capabilities.push(capability);
306 self
307 }
308
309 pub(super) fn ordered_capabilities(&self) -> Result<Vec<Arc<dyn AgentCapability>>, AgentError> {
310 resolve_capability_order(&self.capabilities).map_err(AgentError::from)
311 }
312
313 pub(super) fn ordered_stream_observers(
314 &self,
315 ) -> Result<Vec<Arc<dyn AgentCapability>>, AgentError> {
316 resolve_capability_order(&self.stream_observers).map_err(AgentError::from)
317 }
318
319 pub(super) fn ordered_capabilities_for_validation(
320 &self,
321 ) -> Result<Vec<Arc<dyn AgentCapability>>, crate::capability::CapabilityError> {
322 resolve_capability_order(&self.capabilities)
323 .map_err(|error| crate::capability::CapabilityError::Failed(error.to_string()))
324 }
325
326 #[must_use]
328 pub fn with_stream_observer(mut self, observer: Arc<dyn AgentCapability>) -> Self {
329 self.stream_observers.push(observer);
330 self
331 }
332
333 #[must_use]
335 pub fn with_cancellation_token(mut self, token: CancellationToken) -> Self {
336 self.cancellation_token = Some(token);
337 self
338 }
339
340 #[must_use]
342 pub fn with_capability_bundle(mut self, bundle: &dyn CapabilityBundle) -> Self {
343 self.apply_capability_bundle(bundle);
344 self
345 }
346
347 #[must_use]
349 pub fn with_executor(mut self, executor: DynAgentExecutor) -> Self {
350 self.executor = executor;
351 self
352 }
353
354 #[must_use]
356 pub fn with_trace_recorder(mut self, recorder: DynTraceRecorder) -> Self {
357 self.trace_recorder = recorder;
358 self
359 }
360
361 #[must_use]
363 pub const fn with_policy(mut self, policy: AgentRuntimePolicy) -> Self {
364 self.policy = policy;
365 self
366 }
367
368 pub fn inspect_graph(
374 &self,
375 start: AgentNode,
376 state: &crate::run::AgentRunState,
377 ) -> Result<AgentGraphTrace, GraphError> {
378 inspect_graph(
379 start,
380 state,
381 self.policy.max_steps,
382 self.policy.max_steps.saturating_mul(4),
383 )
384 }
385
386 #[must_use]
388 pub fn override_config(&self) -> AgentOverride {
389 AgentOverride::new(self.clone())
390 }
391
392 fn apply_capability_bundle(&mut self, bundle: &dyn CapabilityBundle) {
393 self.capabilities.extend(bundle.hooks());
394 self.stream_observers.extend(bundle.stream_observers());
395 self.instructions.extend(bundle.get_instructions());
396 self.dynamic_instructions
397 .extend(bundle.dynamic_instructions());
398 if let Some(tools) = bundle.get_tools() {
399 self.tools.insert_registry(&tools);
400 }
401 if let Some(settings) = bundle.model_settings() {
402 self.model_settings = Some(match &self.model_settings {
403 Some(current) => current.merge(&settings),
404 None => settings,
405 });
406 }
407 if let Some(params) = bundle.request_params() {
408 self.request_params = merge_request_params(&self.request_params, ¶ms);
409 }
410 self.output_functions.extend(bundle.output_functions());
411 self.output_validators.extend(bundle.output_validators());
412 if let Some(limits) = bundle.usage_limits() {
413 self.usage_limits = Some(limits);
414 }
415 }
416}