pw_core/
artifacts.rs

1//! Artifact Storage Types
2//!
3//! Types for storing and retrieving LLM artifacts (code, documents, summaries, etc.)
4
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize};
7use uuid::Uuid;
8
9// ============================================================================
10// Artifact Types
11// ============================================================================
12
13/// Stored artifact with embedding for semantic search
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct Artifact {
16    pub id: Uuid,
17    pub user_id: String,
18    pub artifact_type: ArtifactType,
19    pub title: Option<String>,
20    pub content: String,
21    pub metadata: serde_json::Value,
22    pub chunk_index: i32,
23    pub total_chunks: i32,
24    pub parent_id: Option<Uuid>,
25    pub token_count: i32,
26    pub created_at: DateTime<Utc>,
27}
28
29impl Artifact {
30    /// Create a new artifact
31    pub fn new(
32        user_id: impl Into<String>,
33        artifact_type: ArtifactType,
34        content: impl Into<String>,
35    ) -> Self {
36        Self {
37            id: Uuid::new_v4(),
38            user_id: user_id.into(),
39            artifact_type,
40            title: None,
41            content: content.into(),
42            metadata: serde_json::json!({}),
43            chunk_index: 0,
44            total_chunks: 1,
45            parent_id: None,
46            token_count: 0,
47            created_at: Utc::now(),
48        }
49    }
50
51    /// Set title
52    pub fn with_title(mut self, title: impl Into<String>) -> Self {
53        self.title = Some(title.into());
54        self
55    }
56
57    /// Set metadata
58    pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self {
59        self.metadata = metadata;
60        self
61    }
62}
63
64/// Type of artifact content
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
66#[serde(rename_all = "snake_case")]
67#[cfg_attr(feature = "sqlx", derive(sqlx::Type))]
68#[cfg_attr(feature = "sqlx", sqlx(type_name = "artifact_type", rename_all = "snake_case"))]
69pub enum ArtifactType {
70    /// LLM chat response
71    ChatResponse,
72    /// Document (markdown, text, etc.)
73    Document,
74    /// Code snippet
75    CodeSnippet,
76    /// Summary of other content
77    Summary,
78    /// Analysis or insight
79    Analysis,
80    /// Custom/other type
81    Custom,
82}
83
84impl ArtifactType {
85    /// Get all artifact types
86    pub fn all() -> &'static [ArtifactType] {
87        &[
88            ArtifactType::ChatResponse,
89            ArtifactType::Document,
90            ArtifactType::CodeSnippet,
91            ArtifactType::Summary,
92            ArtifactType::Analysis,
93            ArtifactType::Custom,
94        ]
95    }
96
97    /// Get type as string
98    pub fn as_str(&self) -> &'static str {
99        match self {
100            ArtifactType::ChatResponse => "chat_response",
101            ArtifactType::Document => "document",
102            ArtifactType::CodeSnippet => "code_snippet",
103            ArtifactType::Summary => "summary",
104            ArtifactType::Analysis => "analysis",
105            ArtifactType::Custom => "custom",
106        }
107    }
108}
109
110impl std::fmt::Display for ArtifactType {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        write!(f, "{}", self.as_str())
113    }
114}
115
116// ============================================================================
117// Search Types
118// ============================================================================
119
120/// Search result with similarity score
121#[derive(Debug, Clone, Serialize, Deserialize)]
122pub struct SearchResult {
123    pub artifact: Artifact,
124    pub similarity: f32,
125    pub highlights: Vec<String>,
126}
127
128/// Search query parameters
129#[derive(Debug, Clone, Serialize, Deserialize)]
130pub struct SearchQuery {
131    /// Search query text
132    pub query: String,
133    /// Filter by artifact types
134    #[serde(default, skip_serializing_if = "Option::is_none")]
135    pub artifact_types: Option<Vec<ArtifactType>>,
136    /// Minimum similarity threshold (0-1)
137    #[serde(default = "default_min_similarity")]
138    pub min_similarity: f32,
139    /// Maximum results to return
140    #[serde(default = "default_limit")]
141    pub limit: usize,
142    /// Filter by tags in metadata
143    #[serde(default, skip_serializing_if = "Option::is_none")]
144    pub tags: Option<Vec<String>>,
145}
146
147fn default_min_similarity() -> f32 {
148    0.5
149}
150
151fn default_limit() -> usize {
152    10
153}
154
155impl Default for SearchQuery {
156    fn default() -> Self {
157        Self {
158            query: String::new(),
159            artifact_types: None,
160            min_similarity: default_min_similarity(),
161            limit: default_limit(),
162            tags: None,
163        }
164    }
165}
166
167impl SearchQuery {
168    /// Create a new search query
169    pub fn new(query: impl Into<String>) -> Self {
170        Self {
171            query: query.into(),
172            ..Default::default()
173        }
174    }
175
176    /// Filter by artifact types
177    pub fn with_types(mut self, types: Vec<ArtifactType>) -> Self {
178        self.artifact_types = Some(types);
179        self
180    }
181
182    /// Set minimum similarity
183    pub fn with_min_similarity(mut self, threshold: f32) -> Self {
184        self.min_similarity = threshold;
185        self
186    }
187
188    /// Set result limit
189    pub fn with_limit(mut self, limit: usize) -> Self {
190        self.limit = limit;
191        self
192    }
193}
194
195// ============================================================================
196// API Request/Response Types
197// ============================================================================
198
199/// Request to store an artifact
200#[derive(Debug, Clone, Serialize, Deserialize)]
201pub struct StoreArtifactRequest {
202    /// Type of artifact
203    pub artifact_type: ArtifactType,
204    /// Content to store
205    pub content: String,
206    /// Optional title
207    #[serde(skip_serializing_if = "Option::is_none")]
208    pub title: Option<String>,
209    /// Optional metadata
210    #[serde(default)]
211    pub metadata: serde_json::Value,
212}
213
214/// Response after storing an artifact
215#[derive(Debug, Clone, Serialize, Deserialize)]
216pub struct StoreArtifactResponse {
217    /// IDs of created artifacts (multiple if chunked)
218    pub ids: Vec<Uuid>,
219    /// Number of chunks created
220    pub chunks: usize,
221    /// Total tokens in content
222    pub token_count: i32,
223}
224
225/// Request to search artifacts
226#[derive(Debug, Clone, Serialize, Deserialize)]
227pub struct SearchArtifactsRequest {
228    /// Search query
229    pub query: String,
230    /// Filter by artifact types
231    #[serde(skip_serializing_if = "Option::is_none")]
232    pub artifact_types: Option<Vec<ArtifactType>>,
233    /// Minimum similarity (0-1)
234    #[serde(default = "default_min_similarity")]
235    pub min_similarity: f32,
236    /// Maximum results
237    #[serde(default = "default_limit")]
238    pub limit: usize,
239}
240
241/// Response from artifact search
242#[derive(Debug, Clone, Serialize, Deserialize)]
243pub struct SearchArtifactsResponse {
244    pub results: Vec<SearchResult>,
245    pub total: usize,
246    pub query: String,
247}
248
249// ============================================================================
250// Tests
251// ============================================================================
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256
257    #[test]
258    fn test_artifact_creation() {
259        let artifact = Artifact::new("user123", ArtifactType::CodeSnippet, "fn main() {}")
260            .with_title("Main function")
261            .with_metadata(serde_json::json!({"language": "rust"}));
262
263        assert_eq!(artifact.user_id, "user123");
264        assert_eq!(artifact.artifact_type, ArtifactType::CodeSnippet);
265        assert_eq!(artifact.title, Some("Main function".to_string()));
266    }
267
268    #[test]
269    fn test_artifact_type_serialization() {
270        let types = vec![
271            (ArtifactType::ChatResponse, "\"chat_response\""),
272            (ArtifactType::Document, "\"document\""),
273            (ArtifactType::CodeSnippet, "\"code_snippet\""),
274            (ArtifactType::Summary, "\"summary\""),
275            (ArtifactType::Analysis, "\"analysis\""),
276            (ArtifactType::Custom, "\"custom\""),
277        ];
278
279        for (artifact_type, expected) in types {
280            let json = serde_json::to_string(&artifact_type).unwrap();
281            assert_eq!(json, expected);
282        }
283    }
284
285    #[test]
286    fn test_search_query_builder() {
287        let query = SearchQuery::new("rust error handling")
288            .with_types(vec![ArtifactType::CodeSnippet, ArtifactType::Document])
289            .with_min_similarity(0.7)
290            .with_limit(5);
291
292        assert_eq!(query.query, "rust error handling");
293        assert_eq!(query.artifact_types.unwrap().len(), 2);
294        assert_eq!(query.min_similarity, 0.7);
295        assert_eq!(query.limit, 5);
296    }
297
298    #[test]
299    fn test_search_result_serialization() {
300        let result = SearchResult {
301            artifact: Artifact::new("user", ArtifactType::Document, "Hello world"),
302            similarity: 0.95,
303            highlights: vec!["Hello".to_string()],
304        };
305
306        let json = serde_json::to_value(&result).unwrap();
307        assert!((json["similarity"].as_f64().unwrap() - 0.95).abs() < 0.01);
308        assert_eq!(json["highlights"][0], "Hello");
309    }
310
311    #[test]
312    fn test_store_request_serialization() {
313        let request = StoreArtifactRequest {
314            artifact_type: ArtifactType::CodeSnippet,
315            content: "let x = 42;".to_string(),
316            title: Some("Variable".to_string()),
317            metadata: serde_json::json!({"language": "rust"}),
318        };
319
320        let json = serde_json::to_value(&request).unwrap();
321        assert_eq!(json["artifact_type"], "code_snippet");
322        assert_eq!(json["content"], "let x = 42;");
323    }
324}
pw_core/artifacts.rs

pw_core/
artifacts.rs
