synth_ai_core/data/
artifacts.rs

1//! Artifact types for storing outputs and intermediate results.
2//!
3//! Artifacts can contain text, structured data, or file references.
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use std::collections::HashMap;
8
9/// Content of an artifact (text or structured data).
10#[derive(Debug, Clone, Serialize, Deserialize)]
11#[serde(untagged)]
12pub enum ArtifactContent {
13    /// Plain text content.
14    Text(String),
15    /// Structured JSON content.
16    Structured(HashMap<String, Value>),
17}
18
19impl ArtifactContent {
20    /// Create text content.
21    pub fn text(content: impl Into<String>) -> Self {
22        Self::Text(content.into())
23    }
24
25    /// Create structured content.
26    pub fn structured(data: HashMap<String, Value>) -> Self {
27        Self::Structured(data)
28    }
29
30    /// Get as text if this is text content.
31    pub fn as_text(&self) -> Option<&str> {
32        match self {
33            Self::Text(s) => Some(s),
34            Self::Structured(_) => None,
35        }
36    }
37
38    /// Get as structured if this is structured content.
39    pub fn as_structured(&self) -> Option<&HashMap<String, Value>> {
40        match self {
41            Self::Text(_) => None,
42            Self::Structured(m) => Some(m),
43        }
44    }
45
46    /// Get the size in bytes.
47    pub fn size_bytes(&self) -> usize {
48        match self {
49            Self::Text(s) => s.len(),
50            Self::Structured(m) => serde_json::to_string(m).map(|s| s.len()).unwrap_or(0),
51        }
52    }
53}
54
55impl Default for ArtifactContent {
56    fn default() -> Self {
57        Self::Text(String::new())
58    }
59}
60
61/// An artifact produced during a rollout or evaluation.
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct Artifact {
64    /// The artifact content.
65    pub content: ArtifactContent,
66    /// MIME type (e.g., "text/plain", "application/json", "image/png").
67    #[serde(default)]
68    pub content_type: Option<String>,
69    /// Additional metadata.
70    #[serde(default)]
71    pub metadata: HashMap<String, Value>,
72    /// Unique artifact ID.
73    #[serde(default)]
74    pub artifact_id: Option<String>,
75    /// Correlation ID linking to a trace.
76    #[serde(default)]
77    pub trace_correlation_id: Option<String>,
78    /// Size in bytes.
79    #[serde(default)]
80    pub size_bytes: Option<i64>,
81    /// SHA-256 hash of content.
82    #[serde(default)]
83    pub sha256: Option<String>,
84    /// Storage location information.
85    #[serde(default)]
86    pub storage: Option<HashMap<String, Value>>,
87    /// When the artifact was created.
88    #[serde(default)]
89    pub created_at: Option<String>,
90    /// Name/label for the artifact.
91    #[serde(default)]
92    pub name: Option<String>,
93    /// Description.
94    #[serde(default)]
95    pub description: Option<String>,
96}
97
98impl Artifact {
99    /// Create a new text artifact.
100    pub fn text(content: impl Into<String>) -> Self {
101        Self {
102            content: ArtifactContent::text(content),
103            content_type: Some("text/plain".to_string()),
104            metadata: HashMap::new(),
105            artifact_id: None,
106            trace_correlation_id: None,
107            size_bytes: None,
108            sha256: None,
109            storage: None,
110            created_at: None,
111            name: None,
112            description: None,
113        }
114    }
115
116    /// Create a new JSON artifact.
117    pub fn json(data: HashMap<String, Value>) -> Self {
118        Self {
119            content: ArtifactContent::structured(data),
120            content_type: Some("application/json".to_string()),
121            metadata: HashMap::new(),
122            artifact_id: None,
123            trace_correlation_id: None,
124            size_bytes: None,
125            sha256: None,
126            storage: None,
127            created_at: None,
128            name: None,
129            description: None,
130        }
131    }
132
133    /// Set the artifact ID.
134    pub fn with_id(mut self, id: impl Into<String>) -> Self {
135        self.artifact_id = Some(id.into());
136        self
137    }
138
139    /// Set the name.
140    pub fn with_name(mut self, name: impl Into<String>) -> Self {
141        self.name = Some(name.into());
142        self
143    }
144
145    /// Set the trace correlation ID.
146    pub fn with_trace_id(mut self, trace_id: impl Into<String>) -> Self {
147        self.trace_correlation_id = Some(trace_id.into());
148        self
149    }
150
151    /// Add metadata.
152    pub fn with_metadata(mut self, key: impl Into<String>, value: Value) -> Self {
153        self.metadata.insert(key.into(), value);
154        self
155    }
156
157    /// Validate artifact size against a maximum.
158    pub fn validate_size(&self, max_size_bytes: i64) -> Result<(), String> {
159        let size = self
160            .size_bytes
161            .unwrap_or_else(|| self.content.size_bytes() as i64);
162        if size > max_size_bytes {
163            return Err(format!(
164                "Artifact size {} bytes exceeds maximum {} bytes",
165                size, max_size_bytes
166            ));
167        }
168        Ok(())
169    }
170
171    /// Calculate and set size_bytes from content.
172    pub fn compute_size(&mut self) {
173        self.size_bytes = Some(self.content.size_bytes() as i64);
174    }
175}
176
177impl Default for Artifact {
178    fn default() -> Self {
179        Self::text("")
180    }
181}
182
183/// Collection of artifacts from a rollout.
184#[derive(Debug, Clone, Default, Serialize, Deserialize)]
185pub struct ArtifactBundle {
186    /// List of artifacts.
187    #[serde(default)]
188    pub artifacts: Vec<Artifact>,
189    /// Total size in bytes.
190    #[serde(default)]
191    pub total_size_bytes: Option<i64>,
192    /// Bundle metadata.
193    #[serde(default)]
194    pub metadata: HashMap<String, Value>,
195}
196
197impl ArtifactBundle {
198    /// Create a new empty bundle.
199    pub fn new() -> Self {
200        Self::default()
201    }
202
203    /// Add an artifact.
204    pub fn add(&mut self, artifact: Artifact) {
205        self.artifacts.push(artifact);
206    }
207
208    /// Get total size of all artifacts.
209    pub fn compute_total_size(&mut self) -> i64 {
210        let total: i64 = self
211            .artifacts
212            .iter()
213            .map(|a| {
214                a.size_bytes
215                    .unwrap_or_else(|| a.content.size_bytes() as i64)
216            })
217            .sum();
218        self.total_size_bytes = Some(total);
219        total
220    }
221
222    /// Get artifact by ID.
223    pub fn get_by_id(&self, id: &str) -> Option<&Artifact> {
224        self.artifacts
225            .iter()
226            .find(|a| a.artifact_id.as_deref() == Some(id))
227    }
228
229    /// Get artifact by name.
230    pub fn get_by_name(&self, name: &str) -> Option<&Artifact> {
231        self.artifacts
232            .iter()
233            .find(|a| a.name.as_deref() == Some(name))
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    #[test]
242    fn test_text_artifact() {
243        let artifact = Artifact::text("Hello, world!")
244            .with_name("greeting")
245            .with_id("art-001");
246
247        assert_eq!(artifact.content.as_text(), Some("Hello, world!"));
248        assert_eq!(artifact.name, Some("greeting".to_string()));
249        assert_eq!(artifact.content_type, Some("text/plain".to_string()));
250    }
251
252    #[test]
253    fn test_json_artifact() {
254        let mut data = HashMap::new();
255        data.insert("key".to_string(), serde_json::json!("value"));
256
257        let artifact = Artifact::json(data);
258
259        assert!(artifact.content.as_structured().is_some());
260        assert_eq!(artifact.content_type, Some("application/json".to_string()));
261    }
262
263    #[test]
264    fn test_size_validation() {
265        let artifact = Artifact::text("x".repeat(1000));
266
267        assert!(artifact.validate_size(2000).is_ok());
268        assert!(artifact.validate_size(500).is_err());
269    }
270
271    #[test]
272    fn test_artifact_bundle() {
273        let mut bundle = ArtifactBundle::new();
274        bundle.add(Artifact::text("First").with_name("first"));
275        bundle.add(Artifact::text("Second").with_name("second"));
276
277        assert_eq!(bundle.artifacts.len(), 2);
278        assert!(bundle.get_by_name("first").is_some());
279    }
280
281    #[test]
282    fn test_serde() {
283        let artifact = Artifact::text("test content").with_id("test-id");
284
285        let json = serde_json::to_string(&artifact).unwrap();
286        let parsed: Artifact = serde_json::from_str(&json).unwrap();
287
288        assert_eq!(parsed.content.as_text(), Some("test content"));
289        assert_eq!(parsed.artifact_id, Some("test-id".to_string()));
290    }
291}
synth_ai_core/data/artifacts.rs

synth_ai_core/data/
artifacts.rs
