Skip to main content

systemprompt_models/artifacts/chart/
mod.rs

1//! Chart artifact.
2//!
3//! A [`ChartArtifact`] carries labelled axes and one or more [`ChartDataset`]s
4//! for a given [`ChartType`]. Axis configuration and presentation hints are
5//! serialized into the emitted JSON schema; the artifact implements
6//! [`Artifact`].
7
8use 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}