use std::sync::Arc;
use async_trait::async_trait;
use serde_json::{Value, json};
use entelix_core::ir::{ContentPart, Message, Role};
use entelix_core::tools::{Tool, ToolEffect, ToolMetadata};
use entelix_core::{AgentContext, Error, ExecutionContext, Result, SkillRegistry, ToolRegistry};
use entelix_runnable::Runnable;
use crate::agent::{AgentEventSink, Approver};
use crate::react_agent::react_agent_builder;
use crate::state::ReActState;
pub struct Subagent<M>
where
M: Runnable<Vec<Message>, Message> + 'static,
{
name: String,
description: String,
model: Arc<M>,
tool_registry: ToolRegistry,
skills: SkillRegistry,
sinks: Vec<Arc<dyn AgentEventSink<ReActState>>>,
approver: Option<Arc<dyn Approver>>,
}
#[derive(Clone, Debug)]
pub struct SubagentMetadata {
pub name: String,
pub description: String,
pub tool_count: usize,
pub tool_names: Vec<String>,
}
impl<M> std::fmt::Debug for Subagent<M>
where
M: Runnable<Vec<Message>, Message> + 'static,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Subagent")
.field("name", &self.name)
.field("description", &self.description)
.field("tool_count", &self.tool_registry.len())
.field("skill_count", &self.skills.len())
.field("sinks", &self.sinks.len())
.field("has_approver", &self.approver.is_some())
.finish_non_exhaustive()
}
}
impl<M> Subagent<M>
where
M: Runnable<Vec<Message>, Message> + 'static,
{
pub fn builder(
model: M,
parent_registry: &ToolRegistry,
name: impl Into<String>,
description: impl Into<String>,
) -> SubagentBuilder<'_, M> {
SubagentBuilder::new(model, parent_registry, name.into(), description.into())
}
pub fn name(&self) -> &str {
&self.name
}
pub fn description(&self) -> &str {
&self.description
}
#[must_use]
pub fn metadata(&self) -> SubagentMetadata {
SubagentMetadata {
name: self.name.clone(),
description: self.description.clone(),
tool_count: self.tool_registry.len(),
tool_names: self.tool_registry.names().map(str::to_owned).collect(),
}
}
#[must_use]
pub fn tool_count(&self) -> usize {
self.tool_registry.len()
}
#[must_use]
pub fn tool_names(&self) -> Vec<&str> {
self.tool_registry.names().collect()
}
#[must_use]
pub const fn tool_registry(&self) -> &ToolRegistry {
&self.tool_registry
}
#[must_use]
pub const fn skills(&self) -> &SkillRegistry {
&self.skills
}
pub fn into_react_agent(self) -> Result<crate::agent::Agent<ReActState>> {
let Self {
name: _,
description: _,
model,
tool_registry,
skills,
sinks,
approver,
} = self;
let model = ArcRunnable::new(model);
let registry_with_skills = if skills.is_empty() {
tool_registry
} else {
entelix_tools::skills::install(tool_registry, skills)?
};
let registry = match &approver {
Some(approver) => {
registry_with_skills.layer(crate::agent::ApprovalLayer::new(Arc::clone(approver)))
}
None => registry_with_skills,
};
let tool_specs = registry.tool_specs();
let defaults = if tool_specs.is_empty() {
None
} else {
Some(entelix_core::RunOverrides::new().with_tool_specs(tool_specs))
};
let mut builder = react_agent_builder(model, registry, defaults)?;
for sink in sinks {
builder = builder.add_sink_arc(sink);
}
if let Some(approver) = approver {
builder = builder
.with_execution_mode(crate::agent::ExecutionMode::Supervised)
.with_approver_arc(approver);
}
builder.build()
}
pub fn into_tool(self) -> Result<SubagentTool> {
let Self {
name, description, ..
} = &self;
let name = name.clone();
let description = description.clone();
let agent = self.into_react_agent()?;
Ok(SubagentTool::new(agent, name, description))
}
}
type ToolPredicate = Box<dyn Fn(&dyn Tool) -> bool + Send + Sync>;
enum SubagentSelection {
All,
Restrict(Vec<String>),
Filter(ToolPredicate),
}
impl SubagentSelection {
fn apply(self, parent: &ToolRegistry) -> Result<ToolRegistry> {
match self {
Self::All => Ok(parent.clone()),
Self::Restrict(allowed) => {
let refs: Vec<&str> = allowed.iter().map(String::as_str).collect();
parent.restricted_to(&refs)
}
Self::Filter(predicate) => Ok(parent.filter(|tool| predicate(tool))),
}
}
}
pub struct SubagentBuilder<'a, M>
where
M: Runnable<Vec<Message>, Message> + 'static,
{
model: M,
parent_registry: &'a ToolRegistry,
name: String,
description: String,
selection: SubagentSelection,
skills_request: Option<(&'a SkillRegistry, Vec<String>)>,
sinks: Vec<Arc<dyn AgentEventSink<ReActState>>>,
approver: Option<Arc<dyn Approver>>,
}
impl<'a, M> SubagentBuilder<'a, M>
where
M: Runnable<Vec<Message>, Message> + 'static,
{
fn new(model: M, parent_registry: &'a ToolRegistry, name: String, description: String) -> Self {
Self {
model,
parent_registry,
name,
description,
selection: SubagentSelection::All,
skills_request: None,
sinks: Vec::new(),
approver: None,
}
}
#[must_use]
pub fn restrict_to(mut self, allowed: &[&str]) -> Self {
let owned: Vec<String> = allowed.iter().map(|s| (*s).to_owned()).collect();
self.selection = SubagentSelection::Restrict(owned);
self
}
#[must_use]
pub fn filter<F>(mut self, predicate: F) -> Self
where
F: Fn(&dyn Tool) -> bool + Send + Sync + 'static,
{
self.selection = SubagentSelection::Filter(Box::new(predicate));
self
}
#[must_use]
pub fn add_sink(mut self, sink: Arc<dyn AgentEventSink<ReActState>>) -> Self {
self.sinks.push(sink);
self
}
#[must_use]
pub fn with_approver(mut self, approver: Arc<dyn Approver>) -> Self {
self.approver = Some(approver);
self
}
#[must_use]
pub fn with_skills(mut self, parent_skills: &'a SkillRegistry, allowed: &[&str]) -> Self {
let owned: Vec<String> = allowed.iter().map(|s| (*s).to_owned()).collect();
self.skills_request = Some((parent_skills, owned));
self
}
pub fn build(self) -> Result<Subagent<M>> {
let Self {
model,
parent_registry,
name,
description,
selection,
skills_request,
sinks,
approver,
} = self;
if name.is_empty() {
return Err(entelix_core::Error::config(
"SubagentBuilder: name cannot be empty",
));
}
if description.is_empty() {
return Err(entelix_core::Error::config(
"SubagentBuilder: description cannot be empty",
));
}
let tool_registry = selection.apply(parent_registry)?;
let skills = match skills_request {
None => SkillRegistry::new(),
Some((parent_skills, allowed)) => {
for name in &allowed {
entelix_core::identity::validate_config_identifier(
"SubagentBuilder::with_skills",
"skill name",
name,
)?;
}
let missing: Vec<&str> = allowed
.iter()
.map(String::as_str)
.filter(|name| !parent_skills.has(name))
.collect();
if !missing.is_empty() {
return Err(entelix_core::Error::config(format!(
"SubagentBuilder::with_skills: skill name(s) not in parent registry: {}",
missing.join(", ")
)));
}
let allowed_refs: Vec<&str> = allowed.iter().map(String::as_str).collect();
parent_skills.filter(&allowed_refs)
}
};
Ok(Subagent {
name,
description,
model: Arc::new(model),
tool_registry,
skills,
sinks,
approver,
})
}
}
pub struct SubagentTool {
inner: crate::agent::Agent<ReActState>,
metadata: ToolMetadata,
}
impl SubagentTool {
fn new(inner: crate::agent::Agent<ReActState>, name: String, description: String) -> Self {
let metadata = ToolMetadata::function(
name,
description,
json!({
"type": "object",
"required": ["task"],
"properties": {
"task": {
"type": "string",
"description": "Concrete task for the sub-agent. \
Phrased as you would phrase a user \
message to a fresh assistant."
}
},
"additionalProperties": false
}),
)
.with_effect(ToolEffect::Mutating);
Self { inner, metadata }
}
#[must_use]
pub fn with_effect(mut self, effect: ToolEffect) -> Self {
self.metadata = self.metadata.with_effect(effect);
self
}
}
impl std::fmt::Debug for SubagentTool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SubagentTool")
.field("name", &self.metadata.name)
.field("effect", &self.metadata.effect)
.finish_non_exhaustive()
}
}
#[async_trait]
impl Tool for SubagentTool {
fn metadata(&self) -> &ToolMetadata {
&self.metadata
}
async fn execute(&self, input: Value, ctx: &AgentContext<()>) -> Result<Value> {
let task = input.get("task").and_then(Value::as_str).ok_or_else(|| {
Error::invalid_request("SubagentTool: input must include a string 'task' field")
})?;
let sub_thread_id = uuid::Uuid::now_v7().to_string();
if let Some(handle) = ctx.audit_sink() {
handle
.as_sink()
.record_sub_agent_invoked(self.metadata.name.as_str(), &sub_thread_id);
}
let child_ctx = ctx.core().clone().with_thread_id(sub_thread_id);
let initial = ReActState::from_user(task);
let final_state = self.inner.invoke(initial, &child_ctx).await?;
let output_text = final_state
.messages
.iter()
.rev()
.find(|m| matches!(m.role, Role::Assistant))
.and_then(|m| {
m.content.iter().find_map(|p| match p {
ContentPart::Text { text, .. } => Some(text.clone()),
_ => None,
})
})
.unwrap_or_default();
Ok(json!({ "output": output_text }))
}
}
struct ArcRunnable<R> {
inner: Arc<R>,
}
impl<R> ArcRunnable<R> {
const fn new(inner: Arc<R>) -> Self {
Self { inner }
}
}
#[async_trait::async_trait]
impl<R> Runnable<Vec<Message>, Message> for ArcRunnable<R>
where
R: Runnable<Vec<Message>, Message>,
{
async fn invoke(&self, input: Vec<Message>, ctx: &ExecutionContext) -> Result<Message> {
self.inner.invoke(input, ctx).await
}
}