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 fn new(title: impl Into<String>, chart_type: ChartType, ctx: &RequestContext) -> Self {
61        Self {
62            artifact_type: "chart".to_string(),
63            labels: Vec::new(),
64            datasets: Vec::new(),
65            chart_type,
66            title: title.into(),
67            x_axis_label: "X".to_string(),
68            y_axis_label: "Y".to_string(),
69            x_axis_type: AxisType::Category,
70            y_axis_type: AxisType::Linear,
71            metadata: ExecutionMetadata::with_request(ctx),
72        }
73    }
74
75    pub fn with_x_axis_labels(mut self, labels: Vec<String>) -> Self {
76        self.labels = labels;
77        self
78    }
79
80    pub fn with_labels(self, labels: Vec<String>) -> Self {
81        self.with_x_axis_labels(labels)
82    }
83
84    pub fn with_datasets(mut self, datasets: Vec<ChartDataset>) -> Self {
85        self.datasets = datasets;
86        self
87    }
88
89    pub fn add_dataset(mut self, dataset: ChartDataset) -> Self {
90        self.datasets.push(dataset);
91        self
92    }
93
94    pub const fn with_x_axis_type(mut self, axis_type: AxisType) -> Self {
95        self.x_axis_type = axis_type;
96        self
97    }
98
99    pub const fn with_y_axis_type(mut self, axis_type: AxisType) -> Self {
100        self.y_axis_type = axis_type;
101        self
102    }
103
104    pub fn with_axes(mut self, x_label: impl Into<String>, y_label: impl Into<String>) -> Self {
105        self.x_axis_label = x_label.into();
106        self.y_axis_label = y_label.into();
107        self
108    }
109
110    pub fn with_execution_id(mut self, id: impl Into<String>) -> Self {
111        self.metadata.execution_id = Some(id.into());
112        self
113    }
114
115    pub fn with_skill(
116        mut self,
117        skill_id: impl Into<SkillId>,
118        skill_name: impl Into<String>,
119    ) -> Self {
120        self.metadata.skill_id = Some(skill_id.into());
121        self.metadata.skill_name = Some(skill_name.into());
122        self
123    }
124}
125
126impl Artifact for ChartArtifact {
127    fn artifact_type(&self) -> ArtifactType {
128        ArtifactType::Chart
129    }
130
131    fn to_schema(&self) -> JsonValue {
132        json!({
133            "type": "object",
134            "properties": {
135                "labels": {
136                    "type": "array",
137                    "items": {"type": "string"}
138                },
139                "datasets": {
140                    "type": "array",
141                    "items": {
142                        "type": "object",
143                        "properties": {
144                            "label": {"type": "string"},
145                            "data": {"type": "array", "items": {"type": "number"}}
146                        }
147                    }
148                },
149                "_execution_id": {"type": "string"}
150            },
151            "required": ["labels", "datasets"],
152            "x-artifact-type": "chart",
153            "x-chart-hints": {
154                "chart_type": self.chart_type,
155                "title": self.title,
156                "x_axis": {
157                    "label": self.x_axis_label,
158                    "type": self.x_axis_type
159                },
160                "y_axis": {
161                    "label": self.y_axis_label,
162                    "type": self.y_axis_type
163                }
164            }
165        })
166    }
167}