potato_agent/agents/
task.rs

1use crate::{AgentResponse, PyAgentResponse};
2use potato_prompt::Prompt;
3use potato_util::PyHelperFuncs;
4use pyo3::prelude::*;
5use pyo3::IntoPyObjectExt;
6use serde::{Deserialize, Serialize};
7
8#[pyclass(eq)]
9#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
10pub enum TaskStatus {
11    Pending,
12    Running,
13    Completed,
14    Failed,
15}
16
17#[pyclass]
18#[derive(Debug, Serialize)]
19pub struct PyTask {
20    #[pyo3(get)]
21    pub id: String,
22    #[pyo3(get, set)]
23    pub prompt: Prompt,
24    #[pyo3(get, set)]
25    pub dependencies: Vec<String>,
26    #[pyo3(get)]
27    pub status: TaskStatus,
28    #[pyo3(get)]
29    pub agent_id: String,
30    #[pyo3(get)]
31    pub max_retries: u32,
32    pub result: Option<PyAgentResponse>,
33    pub retry_count: u32,
34}
35
36#[pymethods]
37impl PyTask {
38    #[getter]
39    pub fn result<'py>(&mut self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
40        match &mut self.result {
41            Some(resp) => {
42                //Ok(chat_resp)
43                Ok(Some(resp.result(py)?))
44            }
45            None => Ok(None),
46        }
47    }
48}
49
50#[pyclass]
51#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
52pub struct Task {
53    #[pyo3(get)]
54    pub id: String,
55    #[pyo3(get, set)]
56    pub prompt: Prompt,
57    #[pyo3(get, set)]
58    pub dependencies: Vec<String>,
59    #[pyo3(get)]
60    pub status: TaskStatus,
61    #[pyo3(get)]
62    pub agent_id: String,
63    pub result: Option<AgentResponse>,
64    #[pyo3(get)]
65    pub max_retries: u32,
66    pub retry_count: u32,
67}
68
69#[pymethods]
70impl Task {
71    #[new]
72    #[pyo3(signature = (agent_id, prompt, id, dependencies = None, max_retries=None))]
73    pub fn new(
74        agent_id: &str,
75        prompt: Prompt,
76        id: &str,
77        dependencies: Option<Vec<String>>,
78        max_retries: Option<u32>,
79    ) -> Self {
80        Self {
81            prompt,
82            dependencies: dependencies.unwrap_or_default(),
83            status: TaskStatus::Pending,
84            result: None,
85            id: id.to_string(),
86            agent_id: agent_id.to_string(),
87            max_retries: max_retries.unwrap_or(3),
88            retry_count: 0,
89        }
90    }
91
92    pub fn add_dependency(&mut self, dependency: String) {
93        self.dependencies.push(dependency);
94    }
95
96    pub fn set_status(&mut self, status: TaskStatus) {
97        self.status = status;
98    }
99
100    pub fn set_result(&mut self, result: AgentResponse) {
101        self.result = Some(result);
102    }
103
104    #[getter]
105    pub fn result<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
106        match &self.result {
107            Some(resp) => Ok(resp.clone().into_bound_py_any(py).map(Some)?),
108            None => Ok(None),
109        }
110    }
111
112    pub fn __str__(&self) -> String {
113        PyHelperFuncs::__str__(self)
114    }
115}
116
117impl Task {
118    pub fn increment_retry(&mut self) {
119        self.retry_count += 1;
120    }
121}
potato_agent/agents/task.rs

potato_agent/agents/
task.rs
