use std::sync::Arc;
use starweaver_model::{ModelRequestParameters, ModelSettings};
use starweaver_tools::{DynTool, DynToolset, ToolRegistry};
use starweaver_usage::UsageLimits;
use crate::{
instructions::DynDynamicInstruction,
output::{DynOutputFunction, OutputValidator},
};
use super::{AgentCapability, CapabilitySpec};
pub trait CapabilityBundle: Send + Sync {
fn name(&self) -> &str;
fn spec(&self) -> CapabilitySpec {
CapabilitySpec::new(self.name())
}
fn hooks(&self) -> Vec<Arc<dyn AgentCapability>> {
Vec::new()
}
fn stream_observers(&self) -> Vec<Arc<dyn AgentCapability>> {
Vec::new()
}
fn get_instructions(&self) -> Vec<String> {
Vec::new()
}
fn dynamic_instructions(&self) -> Vec<DynDynamicInstruction> {
Vec::new()
}
fn get_tools(&self) -> Option<ToolRegistry> {
None
}
fn model_settings(&self) -> Option<ModelSettings> {
None
}
fn request_params(&self) -> Option<ModelRequestParameters> {
None
}
fn output_functions(&self) -> Vec<DynOutputFunction> {
Vec::new()
}
fn output_validators(&self) -> Vec<Arc<dyn OutputValidator>> {
Vec::new()
}
fn usage_limits(&self) -> Option<UsageLimits> {
None
}
}
#[derive(Clone, Default)]
pub struct StaticCapabilityBundle {
name: String,
spec: Option<CapabilitySpec>,
hooks: Vec<Arc<dyn AgentCapability>>,
stream_observers: Vec<Arc<dyn AgentCapability>>,
instructions: Vec<String>,
dynamic_instructions: Vec<DynDynamicInstruction>,
tools: ToolRegistry,
model_settings: Option<ModelSettings>,
request_params: Option<ModelRequestParameters>,
output_functions: Vec<DynOutputFunction>,
output_validators: Vec<Arc<dyn OutputValidator>>,
usage_limits: Option<UsageLimits>,
}
impl StaticCapabilityBundle {
#[must_use]
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
spec: None,
hooks: Vec::new(),
stream_observers: Vec::new(),
instructions: Vec::new(),
dynamic_instructions: Vec::new(),
tools: ToolRegistry::new(),
model_settings: None,
request_params: None,
output_functions: Vec::new(),
output_validators: Vec::new(),
usage_limits: None,
}
}
#[must_use]
pub fn with_spec(mut self, spec: CapabilitySpec) -> Self {
self.spec = Some(spec);
self
}
#[must_use]
pub fn with_hook(mut self, hook: Arc<dyn AgentCapability>) -> Self {
self.hooks.push(hook);
self
}
#[must_use]
pub fn with_stream_observer(mut self, observer: Arc<dyn AgentCapability>) -> Self {
self.stream_observers.push(observer);
self
}
#[must_use]
pub fn with_instruction(mut self, instruction: impl Into<String>) -> Self {
self.instructions.push(instruction.into());
self
}
#[must_use]
pub fn with_dynamic_instruction(mut self, instruction: DynDynamicInstruction) -> Self {
self.dynamic_instructions.push(instruction);
self
}
#[must_use]
pub fn with_tool(mut self, tool: DynTool) -> Self {
self.tools.insert(tool);
self
}
#[must_use]
pub fn with_toolset(mut self, toolset: &DynToolset) -> Self {
self.tools.insert_toolset(toolset);
self
}
#[must_use]
pub fn with_model_settings(mut self, settings: ModelSettings) -> Self {
self.model_settings = Some(settings);
self
}
#[must_use]
pub fn with_request_params(mut self, params: ModelRequestParameters) -> Self {
self.request_params = Some(params);
self
}
#[must_use]
pub fn with_output_function(mut self, function: DynOutputFunction) -> Self {
self.output_functions.push(function);
self
}
#[must_use]
pub fn with_output_validator(mut self, validator: Arc<dyn OutputValidator>) -> Self {
self.output_validators.push(validator);
self
}
#[must_use]
pub const fn with_usage_limits(mut self, limits: UsageLimits) -> Self {
self.usage_limits = Some(limits);
self
}
}
impl CapabilityBundle for StaticCapabilityBundle {
fn name(&self) -> &str {
&self.name
}
fn spec(&self) -> CapabilitySpec {
self.spec
.clone()
.unwrap_or_else(|| CapabilitySpec::new(self.name.clone()))
}
fn hooks(&self) -> Vec<Arc<dyn AgentCapability>> {
self.hooks.clone()
}
fn stream_observers(&self) -> Vec<Arc<dyn AgentCapability>> {
self.stream_observers.clone()
}
fn get_instructions(&self) -> Vec<String> {
self.instructions.clone()
}
fn dynamic_instructions(&self) -> Vec<DynDynamicInstruction> {
self.dynamic_instructions.clone()
}
fn get_tools(&self) -> Option<ToolRegistry> {
if self.tools.is_empty() && self.tools.get_instructions().is_empty() {
None
} else {
Some(self.tools.clone())
}
}
fn model_settings(&self) -> Option<ModelSettings> {
self.model_settings.clone()
}
fn request_params(&self) -> Option<ModelRequestParameters> {
self.request_params.clone()
}
fn output_functions(&self) -> Vec<DynOutputFunction> {
self.output_functions.clone()
}
fn output_validators(&self) -> Vec<Arc<dyn OutputValidator>> {
self.output_validators.clone()
}
fn usage_limits(&self) -> Option<UsageLimits> {
self.usage_limits.clone()
}
}