synth_ai_core/tracing/
storage.rs

1//! Trace storage trait definition.
2//!
3//! This module defines the abstract `TraceStorage` trait that storage
4//! backends must implement.
5
6use async_trait::async_trait;
7use chrono::{DateTime, Utc};
8use serde_json::Value;
9
10use super::error::TracingError;
11use super::models::{
12    EventReward, MarkovBlanketMessage, OutcomeReward, SessionTimeStep, SessionTrace, TracingEvent,
13};
14
15/// Abstract storage backend for traces.
16///
17/// Implementations can store traces in SQLite, Turso, or other backends.
18#[async_trait]
19pub trait TraceStorage: Send + Sync {
20    // ========================================================================
21    // INITIALIZATION
22    // ========================================================================
23
24    /// Initialize the storage backend (create tables, etc.).
25    async fn initialize(&self) -> Result<(), TracingError>;
26
27    /// Close the storage connection.
28    async fn close(&self) -> Result<(), TracingError>;
29
30    // ========================================================================
31    // SESSION OPERATIONS
32    // ========================================================================
33
34    /// Ensure a session exists (create if not present).
35    ///
36    /// This is idempotent - calling multiple times with the same session_id
37    /// will not create duplicates.
38    async fn ensure_session(
39        &self,
40        session_id: &str,
41        created_at: DateTime<Utc>,
42        metadata: &Value,
43    ) -> Result<(), TracingError>;
44
45    /// Get a session trace by ID.
46    async fn get_session(&self, session_id: &str) -> Result<Option<SessionTrace>, TracingError>;
47
48    /// Delete a session and all related data.
49    async fn delete_session(&self, session_id: &str) -> Result<bool, TracingError>;
50
51    // ========================================================================
52    // TIMESTEP OPERATIONS
53    // ========================================================================
54
55    /// Ensure a timestep exists and return its database ID.
56    ///
57    /// This is idempotent - calling multiple times with the same session_id
58    /// and step_id will return the same database ID.
59    async fn ensure_timestep(
60        &self,
61        session_id: &str,
62        step: &SessionTimeStep,
63    ) -> Result<i64, TracingError>;
64
65    /// Update a timestep (e.g., mark as completed).
66    async fn update_timestep(
67        &self,
68        session_id: &str,
69        step_id: &str,
70        completed_at: Option<DateTime<Utc>>,
71    ) -> Result<(), TracingError>;
72
73    // ========================================================================
74    // EVENT OPERATIONS
75    // ========================================================================
76
77    /// Insert an event and return its database ID.
78    async fn insert_event(
79        &self,
80        session_id: &str,
81        timestep_db_id: Option<i64>,
82        event: &TracingEvent,
83    ) -> Result<i64, TracingError>;
84
85    // ========================================================================
86    // MESSAGE OPERATIONS
87    // ========================================================================
88
89    /// Insert a message and return its database ID.
90    async fn insert_message(
91        &self,
92        session_id: &str,
93        timestep_db_id: Option<i64>,
94        msg: &MarkovBlanketMessage,
95    ) -> Result<i64, TracingError>;
96
97    // ========================================================================
98    // REWARD OPERATIONS
99    // ========================================================================
100
101    /// Insert an outcome (session-level) reward.
102    async fn insert_outcome_reward(
103        &self,
104        session_id: &str,
105        reward: &OutcomeReward,
106    ) -> Result<i64, TracingError>;
107
108    /// Insert an event-level reward.
109    async fn insert_event_reward(
110        &self,
111        session_id: &str,
112        event_id: i64,
113        message_id: Option<i64>,
114        turn_number: Option<i32>,
115        reward: &EventReward,
116    ) -> Result<i64, TracingError>;
117
118    // ========================================================================
119    // QUERY OPERATIONS
120    // ========================================================================
121
122    /// Execute a raw SQL query and return results as JSON values.
123    async fn query(&self, sql: &str, params: Vec<Value>) -> Result<Vec<Value>, TracingError>;
124
125    /// Update session counts (num_timesteps, num_events, num_messages).
126    async fn update_session_counts(&self, session_id: &str) -> Result<(), TracingError>;
127}
128
129/// Storage configuration.
130#[derive(Debug, Clone)]
131pub struct StorageConfig {
132    /// Database URL (file path, :memory:, or libsql:// URL)
133    pub db_url: String,
134    /// Auth token for remote databases
135    pub auth_token: Option<String>,
136    /// Whether to auto-initialize schema
137    pub auto_init: bool,
138}
139
140impl Default for StorageConfig {
141    fn default() -> Self {
142        Self {
143            db_url: ":memory:".to_string(),
144            auth_token: None,
145            auto_init: true,
146        }
147    }
148}
149
150impl StorageConfig {
151    /// Create config for in-memory database.
152    pub fn memory() -> Self {
153        Self::default()
154    }
155
156    /// Create config for a file-based database.
157    pub fn file(path: impl Into<String>) -> Self {
158        Self {
159            db_url: path.into(),
160            auth_token: None,
161            auto_init: true,
162        }
163    }
164
165    /// Create config for a remote Turso database.
166    pub fn turso(url: impl Into<String>, token: impl Into<String>) -> Self {
167        Self {
168            db_url: url.into(),
169            auth_token: Some(token.into()),
170            auto_init: true,
171        }
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    #[test]
180    fn test_storage_config_default() {
181        let config = StorageConfig::default();
182        assert_eq!(config.db_url, ":memory:");
183        assert!(config.auth_token.is_none());
184        assert!(config.auto_init);
185    }
186
187    #[test]
188    fn test_storage_config_file() {
189        let config = StorageConfig::file("/tmp/test.db");
190        assert_eq!(config.db_url, "/tmp/test.db");
191    }
192
193    #[test]
194    fn test_storage_config_turso() {
195        let config = StorageConfig::turso("libsql://test.turso.io", "token123");
196        assert_eq!(config.db_url, "libsql://test.turso.io");
197        assert_eq!(config.auth_token, Some("token123".to_string()));
198    }
199}
synth_ai_core/tracing/storage.rs

synth_ai_core/tracing/
storage.rs
