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