neuromance_agent/task/
mod.rs

1use std::sync::Arc;
2
3use anyhow::Result;
4use neuromance_tools::{ToolImplementation, ToolRegistry};
5use uuid::Uuid;
6
7use neuromance_client::LLMClient;
8use neuromance_common::agents::AgentResponse;
9use neuromance_common::tools::Tool;
10
11use crate::Agent;
12
13mod agents;
14use agents::{ActionAgent, ContextAgent, VerifierAgent};
15
16/// /// Represents a complete agent task composed of three specialized agents:
17/// - ContextAgent: Gathers and analyzes context
18/// - ActionAgent: Takes actions based on context
19/// - VerifierAgent: Verifies actions were successful
20pub struct AgentTask<C: LLMClient> {
21    pub id: String,
22    pub conversation_id: Uuid,
23    pub task_description: String,
24    pub context_agent: ContextAgent<C>,
25    pub action_agent: ActionAgent<C>,
26    pub verifier_agent: VerifierAgent<C>,
27    pub state: TaskState,
28    pub tool_registry: ToolRegistry,
29}
30
31#[derive(Debug, Clone, Default)]
32pub struct TaskState {
33    pub context_gathered: bool,
34    pub action_taken: bool,
35    pub verified: bool,
36    pub context_response: Option<AgentResponse>,
37    pub action_response: Option<AgentResponse>,
38    pub verification_response: Option<AgentResponse>,
39}
40
41/// Response from executing the full agent task
42#[derive(Debug, Clone)]
43pub struct TaskResponse {
44    pub context_response: AgentResponse,
45    pub action_response: AgentResponse,
46    pub verification_response: AgentResponse,
47    pub success: bool,
48}
49
50impl<C: LLMClient + Send + Sync> AgentTask<C> {
51    pub fn new(id: impl Into<String>, task_description: impl Into<String>, client: C) -> Self
52    where
53        C: Clone,
54    {
55        let conversation_id = Uuid::new_v4();
56        let id = id.into();
57        let task_description = task_description.into();
58
59        Self {
60            id: id.clone(),
61            conversation_id,
62            task_description,
63            context_agent: ContextAgent::new(id.clone(), client.clone()),
64            action_agent: ActionAgent::new(id.clone(), client.clone()),
65            verifier_agent: VerifierAgent::new(id, client),
66            state: TaskState::default(),
67            tool_registry: ToolRegistry::new(),
68        }
69    }
70
71    /// Add tool to AgentTask ToolRegistry for ActionAgent
72    pub fn add_tool<T: ToolImplementation + 'static>(&self, tool: T) {
73        self.tool_registry.register(Arc::new(tool));
74    }
75
76    /// Add a pre-wrapped tool (`Arc<dyn ToolImplementation>`) to the registry
77    pub fn add_tool_arc(&self, tool: Arc<dyn ToolImplementation>) {
78        self.tool_registry.register(tool);
79    }
80
81    /// Remove Tool from AgentTask ToolRegistry
82    pub async fn remove_tool(&mut self, name: &str) -> Result<Option<Arc<dyn ToolImplementation>>> {
83        let tool = self.tool_registry.remove(name);
84        Ok(tool)
85    }
86
87    /// Get all Tools from AgentTask ToolRegistry
88    pub fn get_all_tools(&self) -> Vec<Tool> {
89        self.tool_registry.get_all_definitions()
90    }
91
92    /// Format tools for inclusion in system prompt
93    fn format_tools_for_prompt(&self, tools: &[Tool]) -> String {
94        if tools.is_empty() {
95            return String::from("\n\nNo tools are currently available.");
96        }
97
98        let mut description = String::from("\n\nAvailable tools for the ActionAgent:\n");
99
100        for tool in tools {
101            description.push_str(&format!(
102                "\n- {}: {}\n",
103                tool.function.name, tool.function.description
104            ));
105
106            // NOTE parameters seem excessive and context wasting, revisit later
107            //
108            // Add parameter information if available
109            // description.push_str(&format!("  Parameters: {}\n",
110            //     serde_json::to_string_pretty(&tool.function.parameters).unwrap_or_else(|_| "N/A".to_string())
111            // ));
112        }
113
114        description
115    }
116
117    /// Execute the context gathering phase
118    pub async fn gather_context(&mut self) -> Result<AgentResponse> {
119        // Get all available tools from the registry
120        let available_tools = self.get_all_tools();
121
122        // Format tools for the context agent's system prompt
123        let tools_description = self.format_tools_for_prompt(&available_tools);
124
125        // Pass tools information to context agent
126        let response = self
127            .context_agent
128            .analyze_with_tools(self.task_description.clone(), tools_description)
129            .await?;
130        self.state.context_gathered = true;
131        self.state.context_response = Some(response.clone());
132        Ok(response)
133    }
134
135    /// Execute the action phase using context from previous phase
136    pub async fn take_action(&mut self) -> Result<AgentResponse> {
137        if !self.state.context_gathered {
138            return Err(anyhow::anyhow!(
139                "Cannot take action before gathering context"
140            ));
141        }
142
143        // Add all tools from the registry to the action agent
144        for tool_name in self.tool_registry.tool_names() {
145            if let Some(tool) = self.tool_registry.get(&tool_name) {
146                // Add tool to the action agent's tool executor
147                // We need to clone the Arc to share the tool
148                self.action_agent
149                    .agent
150                    .core
151                    .tool_executor
152                    .add_tool_arc(tool);
153            }
154        }
155
156        let context = self
157            .state
158            .context_response
159            .as_ref()
160            .map(|r| r.content.content.clone())
161            .unwrap_or_default();
162
163        let response = self
164            .action_agent
165            .execute_task(self.task_description.clone(), context)
166            .await?;
167        self.state.action_taken = true;
168        self.state.action_response = Some(response.clone());
169        Ok(response)
170    }
171
172    /// Execute the verification phase
173    pub async fn verify(&mut self) -> Result<AgentResponse> {
174        if !self.state.action_taken {
175            return Err(anyhow::anyhow!("Cannot verify before action is taken"));
176        }
177
178        let action_result = self
179            .state
180            .action_response
181            .as_ref()
182            .map(|r| r.content.content.clone())
183            .unwrap_or_default();
184
185        let (verified, response) = self
186            .verifier_agent
187            .verify(self.task_description.clone(), action_result)
188            .await?;
189
190        self.state.verified = verified;
191        self.state.verification_response = Some(response.clone());
192        Ok(response)
193    }
194
195    /// Execute the full task pipeline: context -> action -> verify
196    pub async fn execute_full(&mut self) -> Result<TaskResponse> {
197        // Phase 1: Gather context
198        let context_response = self.gather_context().await?;
199
200        // Phase 2: Take action based on context
201        let action_response = self.take_action().await?;
202
203        // Phase 3: Verify the action results
204        let verification_response = self.verify().await?;
205
206        Ok(TaskResponse {
207            context_response,
208            action_response,
209            verification_response,
210            success: self.state.verified,
211        })
212    }
213
214    /// Reset the task to initial state
215    pub async fn reset(&mut self) -> Result<()> {
216        self.context_agent.agent.reset().await?;
217        self.action_agent.agent.reset().await?;
218        self.verifier_agent.agent.reset().await?;
219        self.state = TaskState::default();
220        self.conversation_id = Uuid::new_v4();
221        Ok(())
222    }
223}