mentra 0.6.0

An agent runtime for tool-using LLM applications
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use std::{borrow::Cow, collections::HashSet, sync::Arc};

use crate::{
    Role,
    error::RuntimeError,
    provider::Provider,
    runtime::{RuntimeIntrinsicTool, handle::RuntimeHandle},
};

use super::{
    Agent, AgentConfig, AgentSpawnOptions, SpawnedAgentStatus, SpawnedAgentSummary,
    TeammateIdentity,
};

const SUBAGENT_MAX_ROUNDS: usize = 30;
const SUBAGENT_SYSTEM_PROMPT: &str = "You are a subagent working for another agent. Solve the delegated task, use tools when helpful, and finish with a concise final answer for the parent agent.";

#[derive(Clone)]
pub(crate) struct DisposableSubagentTemplate {
    runtime: RuntimeHandle,
    model: String,
    parent_name: String,
    config: AgentConfig,
    provider: Arc<dyn Provider>,
    hidden_tools: HashSet<String>,
    teammate_identity: Option<TeammateIdentity>,
}

impl DisposableSubagentTemplate {
    pub(crate) fn from_agent(agent: &Agent) -> Self {
        Self {
            runtime: agent.runtime.clone(),
            model: agent.model.clone(),
            parent_name: agent.name.clone(),
            config: agent.config.clone(),
            provider: Arc::clone(&agent.provider),
            hidden_tools: agent.hidden_tools.clone(),
            teammate_identity: agent.teammate_identity.clone(),
        }
    }

    pub(crate) fn spawn(&self) -> Result<Agent, RuntimeError> {
        let mut hidden_tools = self.hidden_tools.clone();
        hidden_tools.insert(RuntimeIntrinsicTool::Task.to_string());

        let mut config = self.config.clone();
        config.system = Some(build_subagent_system_prompt(
            self.config.system.as_deref().map(Cow::Borrowed),
        ));

        Agent::new(
            self.runtime.clone(),
            self.model.clone(),
            format!("{}::task", self.parent_name),
            config,
            Arc::clone(&self.provider),
            AgentSpawnOptions {
                hidden_tools,
                max_rounds: Some(SUBAGENT_MAX_ROUNDS),
                teammate_identity: self.teammate_identity.clone(),
            },
        )
    }
}

impl Agent {
    pub(crate) fn spawn_subagent(&self) -> Result<Self, RuntimeError> {
        self.disposable_subagent_template().spawn()
    }

    pub(crate) fn disposable_subagent_template(&self) -> DisposableSubagentTemplate {
        DisposableSubagentTemplate::from_agent(self)
    }

    pub(crate) fn register_subagent(&mut self, agent: &Agent) -> SpawnedAgentSummary {
        let summary = SpawnedAgentSummary {
            id: agent.id.clone(),
            name: agent.name.clone(),
            model: agent.model.clone(),
            status: SpawnedAgentStatus::Running,
        };
        let summary_for_snapshot = summary.clone();
        self.mutate_snapshot(|snapshot| {
            snapshot.subagents.push(summary_for_snapshot);
        });
        summary
    }

    pub(crate) fn finish_subagent(
        &mut self,
        id: &str,
        status: SpawnedAgentStatus,
    ) -> Option<SpawnedAgentSummary> {
        let mut finished = None;
        self.mutate_snapshot(|snapshot| {
            if let Some(summary) = snapshot.subagents.iter_mut().find(|agent| agent.id == id) {
                summary.status = status;
                finished = Some(summary.clone());
            }
        });
        finished
    }

    pub(crate) fn final_text_summary(&self) -> String {
        let Some(message) = self.last_message() else {
            return "(no summary)".to_string();
        };

        if message.role != Role::Assistant {
            return "(no summary)".to_string();
        }

        let text = message.text();

        if text.is_empty() {
            "(no summary)".to_string()
        } else {
            text
        }
    }
}

pub(super) fn build_subagent_system_prompt(base: Option<Cow<'_, str>>) -> String {
    match base {
        Some(system) => format!("{system}\n\n{SUBAGENT_SYSTEM_PROMPT}"),
        None => SUBAGENT_SYSTEM_PROMPT.to_string(),
    }
}