potato_agent/agents/
task.rs

1use crate::agents::types::ChatResponse;
2use potato_prompt::Prompt;
3use potato_util::PyHelperFuncs;
4use pyo3::prelude::*;
5use serde::{Deserialize, Serialize};
6use std::sync::Arc;
7
8#[pyclass]
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    pub result: Option<ChatResponse>,
31    #[pyo3(get)]
32    pub max_retries: u32,
33    pub retry_count: u32,
34
35    #[serde(skip)]
36    pub response_type: Option<Arc<PyObject>>,
37}
38
39#[pyclass]
40#[derive(Debug, Serialize, Deserialize, Clone)]
41pub struct Task {
42    #[pyo3(get)]
43    pub id: String,
44    #[pyo3(get, set)]
45    pub prompt: Prompt,
46    #[pyo3(get, set)]
47    pub dependencies: Vec<String>,
48    #[pyo3(get)]
49    pub status: TaskStatus,
50    #[pyo3(get)]
51    pub agent_id: String,
52    pub result: Option<ChatResponse>,
53    #[pyo3(get)]
54    pub max_retries: u32,
55    pub retry_count: u32,
56}
57
58#[pymethods]
59impl Task {
60    #[new]
61    #[pyo3(signature = (agent_id, prompt, id, dependencies = None, max_retries=None))]
62    pub fn new(
63        agent_id: &str,
64        prompt: Prompt,
65        id: &str,
66        dependencies: Option<Vec<String>>,
67        max_retries: Option<u32>,
68    ) -> Self {
69        Self {
70            prompt,
71            dependencies: dependencies.unwrap_or_default(),
72            status: TaskStatus::Pending,
73            result: None,
74            id: id.to_string(),
75            agent_id: agent_id.to_string(),
76            max_retries: max_retries.unwrap_or(3),
77            retry_count: 0,
78        }
79    }
80
81    pub fn add_dependency(&mut self, dependency: String) {
82        self.dependencies.push(dependency);
83    }
84
85    pub fn set_status(&mut self, status: TaskStatus) {
86        self.status = status;
87    }
88
89    pub fn set_result(&mut self, result: ChatResponse) {
90        self.result = Some(result);
91    }
92
93    #[getter]
94    pub fn result<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
95        match &self.result {
96            Some(resp) => Ok(resp.to_python(py).map(Some)?),
97            None => Ok(None),
98        }
99    }
100
101    pub fn __str__(&self) -> String {
102        PyHelperFuncs::__str__(self)
103    }
104}
105
106impl Task {
107    pub fn increment_retry(&mut self) {
108        self.retry_count += 1;
109    }
110}
potato_agent/agents/task.rs

potato_agent/agents/
task.rs
