Skip to main content

systemprompt_models/artifacts/chart/
mod.rs

1use crate::artifacts::metadata::ExecutionMetadata;
2use crate::artifacts::traits::Artifact;
3use crate::artifacts::types::{ArtifactType, AxisType, ChartType};
4use crate::execution::context::RequestContext;
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7use serde_json::{json, Value as JsonValue};
8use systemprompt_identifiers::SkillId;
9
10fn default_artifact_type() -> String {
11    "chart".to_string()
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
15pub struct ChartDataset {
16    pub label: String,
17    pub data: Vec<f64>,
18}
19
20impl ChartDataset {
21    pub fn new(label: impl Into<String>, data: Vec<f64>) -> Self {
22        Self {
23            label: label.into(),
24            data,
25        }
26    }
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
30pub struct ChartArtifact {
31    #[serde(rename = "x-artifact-type")]
32    #[serde(default = "default_artifact_type")]
33    pub artifact_type: String,
34    pub labels: Vec<String>,
35    pub datasets: Vec<ChartDataset>,
36    #[serde(skip)]
37    #[schemars(skip)]
38    chart_type: ChartType,
39    #[serde(skip)]
40    #[schemars(skip)]
41    title: String,
42    #[serde(skip)]
43    #[schemars(skip)]
44    x_axis_label: String,
45    #[serde(skip)]
46    #[schemars(skip)]
47    y_axis_label: String,
48    #[serde(skip)]
49    #[schemars(skip)]
50    x_axis_type: AxisType,
51    #[serde(skip)]
52    #[schemars(skip)]
53    y_axis_type: AxisType,
54    #[serde(skip)]
55    #[schemars(skip)]
56    metadata: ExecutionMetadata,
57}
58
59impl ChartArtifact {
60    pub const ARTIFACT_TYPE_STR: &'static str = "chart";
61
62    pub fn new(title: impl Into<String>, chart_type: ChartType, ctx: &RequestContext) -> Self {
63        Self {
64            artifact_type: "chart".to_string(),
65            labels: Vec::new(),
66            datasets: Vec::new(),
67            chart_type,
68            title: title.into(),
69            x_axis_label: "X".to_string(),
70            y_axis_label: "Y".to_string(),
71            x_axis_type: AxisType::Category,
72            y_axis_type: AxisType::Linear,
73            metadata: ExecutionMetadata::with_request(ctx),
74        }
75    }
76
77    pub fn with_x_axis_labels(mut self, labels: Vec<String>) -> Self {
78        self.labels = labels;
79        self
80    }
81
82    pub fn with_labels(self, labels: Vec<String>) -> Self {
83        self.with_x_axis_labels(labels)
84    }
85
86    pub fn with_datasets(mut self, datasets: Vec<ChartDataset>) -> Self {
87        self.datasets = datasets;
88        self
89    }
90
91    pub fn add_dataset(mut self, dataset: ChartDataset) -> Self {
92        self.datasets.push(dataset);
93        self
94    }
95
96    pub const fn with_x_axis_type(mut self, axis_type: AxisType) -> Self {
97        self.x_axis_type = axis_type;
98        self
99    }
100
101    pub const fn with_y_axis_type(mut self, axis_type: AxisType) -> Self {
102        self.y_axis_type = axis_type;
103        self
104    }
105
106    pub fn with_axes(mut self, x_label: impl Into<String>, y_label: impl Into<String>) -> Self {
107        self.x_axis_label = x_label.into();
108        self.y_axis_label = y_label.into();
109        self
110    }
111
112    pub fn with_execution_id(mut self, id: impl Into<String>) -> Self {
113        self.metadata.execution_id = Some(id.into());
114        self
115    }
116
117    pub fn with_skill(
118        mut self,
119        skill_id: impl Into<SkillId>,
120        skill_name: impl Into<String>,
121    ) -> Self {
122        self.metadata.skill_id = Some(skill_id.into());
123        self.metadata.skill_name = Some(skill_name.into());
124        self
125    }
126}
127
128impl Artifact for ChartArtifact {
129    fn artifact_type(&self) -> ArtifactType {
130        ArtifactType::Chart
131    }
132
133    fn to_schema(&self) -> JsonValue {
134        json!({
135            "type": "object",
136            "properties": {
137                "labels": {
138                    "type": "array",
139                    "items": {"type": "string"}
140                },
141                "datasets": {
142                    "type": "array",
143                    "items": {
144                        "type": "object",
145                        "properties": {
146                            "label": {"type": "string"},
147                            "data": {"type": "array", "items": {"type": "number"}}
148                        }
149                    }
150                },
151                "_execution_id": {"type": "string"}
152            },
153            "required": ["labels", "datasets"],
154            "x-artifact-type": "chart",
155            "x-chart-hints": {
156                "chart_type": self.chart_type,
157                "title": self.title,
158                "x_axis": {
159                    "label": self.x_axis_label,
160                    "type": self.x_axis_type
161                },
162                "y_axis": {
163                    "label": self.y_axis_label,
164                    "type": self.y_axis_type
165                }
166            }
167        })
168    }
169}