alith_core/
task.rs

1use crate::mcp::MCPError;
2use crate::{agent::Agent, chat::Completion};
3use std::sync::Arc;
4use std::time::{SystemTime, UNIX_EPOCH};
5use tokio::sync::RwLock;
6use uuid::Uuid;
7
8#[derive(Clone)]
9pub struct Task<M: Completion> {
10    pub id: Uuid,
11    pub prompt: String,
12    pub output: Option<String>,
13    pub agent: Arc<RwLock<Agent<M>>>, // Use RwLock for thread-safe mutable access
14    pub metadata: Option<TaskMetadata>,
15}
16
17#[derive(Debug, Clone)]
18pub struct TaskMetadata {
19    pub priority: usize,
20    pub tags: Vec<String>,
21    pub created_at: u64,
22}
23
24impl<M: Completion> Task<M> {
25    /// Creates a new task.
26    pub fn new(agent: Arc<RwLock<Agent<M>>>, prompt: String) -> Self {
27        Self {
28            id: Uuid::default(),
29            prompt,
30            output: None,
31            agent,
32            metadata: None,
33        }
34    }
35
36    /// Attaches metadata to the task.
37    pub fn with_metadata(mut self, metadata: TaskMetadata) -> Self {
38        self.metadata = Some(metadata);
39        self
40    }
41
42    /// Executes the task using the agent.
43    pub async fn execute(&mut self) -> Result<String, TaskError> {
44        // Get a write lock for mutable access to the Agent
45        let agent = self.agent.read().await;
46        // Call `execute_task` on the Agent
47        let result = agent
48            .prompt(&self.prompt.clone())
49            .await
50            .map_err(|err| TaskError::ExecutionError(err.to_string()))?;
51
52        // Set the output of the task
53        self.output = Some(result);
54        Ok(self.output.clone().unwrap())
55    }
56}
57
58impl TaskMetadata {
59    /// Creates a new `TaskMetadata` instance.
60    pub fn new(priority: usize, tags: Vec<String>) -> Self {
61        let created_at = SystemTime::now()
62            .duration_since(UNIX_EPOCH)
63            .unwrap()
64            .as_secs(); // Get current time as seconds since UNIX epoch
65
66        Self {
67            priority,
68            tags,
69            created_at,
70        }
71    }
72}
73
74#[derive(Debug, thiserror::Error)]
75pub enum TaskError {
76    #[error("Failed to execute the task: {0}")]
77    ExecutionError(String),
78    #[error("Failed to acquire lock on the agent")]
79    LockError,
80    #[error("An unknown error occurred: {0}")]
81    Unknown(String),
82    #[error("MCP error: {0}")]
83    MCPError(#[from] MCPError),
84}