Skip to main content

vex_persist/
evolution_store.rs

1use async_trait::async_trait;
2use thiserror::Error;
3use vex_core::{GenomeExperiment, OptimizationRule};
4
5#[derive(Debug, Error)]
6pub enum EvolutionStoreError {
7    #[error("Database error: {0}")]
8    Database(#[from] sqlx::Error),
9    #[error("Serialization error: {0}")]
10    Serialization(#[from] serde_json::Error),
11}
12
13#[async_trait]
14pub trait EvolutionStore: Send + Sync + std::fmt::Debug {
15    /// Save an experiment to persistent storage
16    async fn save_experiment(
17        &self,
18        tenant_id: &str,
19        experiment: &GenomeExperiment,
20    ) -> Result<(), EvolutionStoreError>;
21
22    /// Load recent experiments
23    async fn load_recent(
24        &self,
25        tenant_id: &str,
26        limit: usize,
27    ) -> Result<Vec<GenomeExperiment>, EvolutionStoreError>;
28
29    /// Save an optimization rule (semantic lesson)
30    async fn save_rule(
31        &self,
32        tenant_id: &str,
33        rule: &OptimizationRule,
34    ) -> Result<(), EvolutionStoreError>;
35
36    /// Load available optimization rules
37    async fn load_rules(
38        &self,
39        tenant_id: &str,
40    ) -> Result<Vec<OptimizationRule>, EvolutionStoreError>;
41
42    /// Count the number of experiments for a tenant
43    async fn count_experiments(&self, tenant_id: &str) -> Result<u64, EvolutionStoreError>;
44}
45
46/// SQL implementation of EvolutionStore
47#[cfg(feature = "sqlite")]
48#[derive(Debug)]
49pub struct SqliteEvolutionStore {
50    pool: sqlx::SqlitePool,
51}
52
53#[cfg(feature = "sqlite")]
54impl SqliteEvolutionStore {
55    pub fn new(pool: sqlx::SqlitePool) -> Self {
56        Self { pool }
57    }
58}
59
60#[cfg(feature = "sqlite")]
61#[async_trait]
62impl EvolutionStore for SqliteEvolutionStore {
63    async fn save_experiment(
64        &self,
65        tenant_id: &str,
66        experiment: &GenomeExperiment,
67    ) -> Result<(), EvolutionStoreError> {
68        let traits_json = serde_json::to_string(&experiment.traits)?;
69        let trait_names_json = serde_json::to_string(&experiment.trait_names)?;
70        let fitness_json = serde_json::to_string(&experiment.fitness_scores)?;
71        let task_summary = &experiment.task_summary;
72        let overall_fitness = experiment.overall_fitness;
73
74        sqlx::query(
75            r#"
76            INSERT INTO evolution_experiments (
77                id, tenant_id, traits, trait_names, fitness_scores, task_summary, overall_fitness, created_at
78            ) VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'))
79            "#,
80        )
81        .bind(experiment.id.to_string())
82        .bind(tenant_id)
83        .bind(traits_json)
84        .bind(trait_names_json)
85        .bind(fitness_json)
86        .bind(task_summary)
87        .bind(overall_fitness)
88        .execute(&self.pool)
89        .await?;
90
91        Ok(())
92    }
93
94    async fn load_recent(
95        &self,
96        tenant_id: &str,
97        limit: usize,
98    ) -> Result<Vec<GenomeExperiment>, EvolutionStoreError> {
99        use sqlx::Row;
100
101        let rows = sqlx::query(
102            r#"
103            SELECT 
104                id, traits, trait_names, fitness_scores, task_summary, overall_fitness, created_at
105            FROM evolution_experiments
106            WHERE tenant_id = ?
107            ORDER BY created_at DESC
108            LIMIT ?
109            "#,
110        )
111        .bind(tenant_id)
112        .bind(limit as i64)
113        .fetch_all(&self.pool)
114        .await?;
115
116        let mut experiments = Vec::new();
117        for row in rows {
118            let traits_str: String = row.try_get("traits")?;
119            let trait_names_str: String = row.try_get("trait_names")?;
120            let fitness_scores_str: String = row.try_get("fitness_scores")?;
121            let id_str: String = row.try_get("id")?;
122
123            let traits = serde_json::from_str(&traits_str)?;
124            let trait_names = serde_json::from_str(&trait_names_str)?;
125            let fitness_scores = serde_json::from_str(&fitness_scores_str)?;
126
127            // We reconstruct the experiment
128            let exp = GenomeExperiment {
129                id: uuid::Uuid::parse_str(&id_str).unwrap_or_default(),
130                traits,
131                trait_names,
132                fitness_scores,
133                task_summary: row.try_get("task_summary")?,
134                overall_fitness: row.try_get("overall_fitness")?,
135                timestamp: chrono::Utc::now(), // Use current time as parsing SQL datetime can be tricky without types
136                successful: row.try_get::<f64, _>("overall_fitness")? > 0.6,
137            };
138            experiments.push(exp);
139        }
140
141        Ok(experiments)
142    }
143
144    async fn save_rule(
145        &self,
146        tenant_id: &str,
147        rule: &OptimizationRule,
148    ) -> Result<(), EvolutionStoreError> {
149        let traits_json = serde_json::to_string(&rule.affected_traits)?;
150
151        sqlx::query(
152            r#"
153            INSERT INTO optimization_rules (
154                id, tenant_id, rule_description, affected_traits, confidence, source_count, created_at
155            ) VALUES (?, ?, ?, ?, ?, ?, datetime('now'))
156            "#,
157        )
158        .bind(rule.id.to_string())
159        .bind(tenant_id)
160        .bind(&rule.rule_description)
161        .bind(traits_json)
162        .bind(rule.confidence)
163        .bind(rule.source_experiments_count as i64)
164        .execute(&self.pool)
165        .await?;
166
167        Ok(())
168    }
169
170    async fn load_rules(
171        &self,
172        tenant_id: &str,
173    ) -> Result<Vec<OptimizationRule>, EvolutionStoreError> {
174        use sqlx::Row;
175
176        let rows = sqlx::query(
177            r#"
178            SELECT 
179                id, rule_description, affected_traits, confidence, source_count, created_at
180            FROM optimization_rules
181            WHERE tenant_id = ?
182            ORDER BY confidence DESC, created_at DESC
183            LIMIT 50
184            "#,
185        )
186        .bind(tenant_id)
187        .fetch_all(&self.pool)
188        .await?;
189
190        let mut rules = Vec::new();
191        for row in rows {
192            let id_str: String = row.try_get("id")?;
193            let traits_str: String = row.try_get("affected_traits")?;
194
195            let rules_obj = OptimizationRule {
196                id: uuid::Uuid::parse_str(&id_str).unwrap_or_default(),
197                rule_description: row.try_get("rule_description")?,
198                affected_traits: serde_json::from_str(&traits_str)?,
199                confidence: row.try_get("confidence")?,
200                source_experiments_count: row.try_get::<i64, _>("source_count")? as usize,
201                created_at: chrono::Utc::now(), // Simplified
202            };
203            rules.push(rules_obj);
204        }
205
206        Ok(rules)
207    }
208
209    async fn count_experiments(&self, tenant_id: &str) -> Result<u64, EvolutionStoreError> {
210        use sqlx::Row;
211        let count: i64 =
212            sqlx::query("SELECT COUNT(*) as count FROM evolution_experiments WHERE tenant_id = ?")
213                .bind(tenant_id)
214                .fetch_one(&self.pool)
215                .await?
216                .try_get("count")?;
217
218        Ok(count as u64)
219    }
220}
221
222/// PostgreSQL implementation of EvolutionStore
223#[cfg(feature = "postgres")]
224#[derive(Debug)]
225pub struct PostgresEvolutionStore {
226    pool: sqlx::PgPool,
227}
228
229#[cfg(feature = "postgres")]
230impl PostgresEvolutionStore {
231    pub fn new(pool: sqlx::PgPool) -> Self {
232        Self { pool }
233    }
234}
235
236#[cfg(feature = "postgres")]
237#[async_trait]
238impl EvolutionStore for PostgresEvolutionStore {
239    async fn save_experiment(
240        &self,
241        tenant_id: &str,
242        experiment: &GenomeExperiment,
243    ) -> Result<(), EvolutionStoreError> {
244        let traits_json = serde_json::to_string(&experiment.traits)?;
245        let trait_names_json = serde_json::to_string(&experiment.trait_names)?;
246        let fitness_json = serde_json::to_string(&experiment.fitness_scores)?;
247
248        sqlx::query(
249            r#"
250            INSERT INTO evolution_experiments (
251                id, tenant_id, traits, trait_names, fitness_scores, task_summary, overall_fitness, created_at
252            ) VALUES ($1, $2, $3, $4, $5, $6, $7, NOW())
253            "#,
254        )
255        .bind(experiment.id.to_string())
256        .bind(tenant_id)
257        .bind(traits_json)
258        .bind(trait_names_json)
259        .bind(fitness_json)
260        .bind(&experiment.task_summary)
261        .bind(experiment.overall_fitness)
262        .execute(&self.pool)
263        .await?;
264
265        Ok(())
266    }
267
268    async fn load_recent(
269        &self,
270        tenant_id: &str,
271        limit: usize,
272    ) -> Result<Vec<GenomeExperiment>, EvolutionStoreError> {
273        use sqlx::Row;
274
275        let rows = sqlx::query(
276            r#"
277            SELECT id, traits, trait_names, fitness_scores, task_summary, overall_fitness
278            FROM evolution_experiments
279            WHERE tenant_id = $1
280            ORDER BY created_at DESC
281            LIMIT $2
282            "#,
283        )
284        .bind(tenant_id)
285        .bind(limit as i64)
286        .fetch_all(&self.pool)
287        .await?;
288
289        let mut experiments = Vec::new();
290        for row in rows {
291            let traits_str: String = row.try_get("traits")?;
292            let trait_names_str: String = row.try_get("trait_names")?;
293            let fitness_scores_str: String = row.try_get("fitness_scores")?;
294            let id_str: String = row.try_get("id")?;
295
296            let exp = GenomeExperiment {
297                id: uuid::Uuid::parse_str(&id_str).unwrap_or_default(),
298                traits: serde_json::from_str(&traits_str)?,
299                trait_names: serde_json::from_str(&trait_names_str)?,
300                fitness_scores: serde_json::from_str(&fitness_scores_str)?,
301                task_summary: row.try_get("task_summary")?,
302                overall_fitness: row.try_get("overall_fitness")?,
303                timestamp: chrono::Utc::now(),
304                successful: row.try_get::<f64, _>("overall_fitness")? > 0.6,
305            };
306            experiments.push(exp);
307        }
308
309        Ok(experiments)
310    }
311
312    async fn save_rule(
313        &self,
314        tenant_id: &str,
315        rule: &OptimizationRule,
316    ) -> Result<(), EvolutionStoreError> {
317        let traits_json = serde_json::to_string(&rule.affected_traits)?;
318
319        sqlx::query(
320            r#"
321            INSERT INTO optimization_rules (
322                id, tenant_id, rule_description, affected_traits, confidence, source_count, created_at
323            ) VALUES ($1, $2, $3, $4, $5, $6, NOW())
324            "#,
325        )
326        .bind(rule.id.to_string())
327        .bind(tenant_id)
328        .bind(&rule.rule_description)
329        .bind(traits_json)
330        .bind(rule.confidence)
331        .bind(rule.source_experiments_count as i64)
332        .execute(&self.pool)
333        .await?;
334
335        Ok(())
336    }
337
338    async fn load_rules(
339        &self,
340        tenant_id: &str,
341    ) -> Result<Vec<OptimizationRule>, EvolutionStoreError> {
342        use sqlx::Row;
343
344        let rows = sqlx::query(
345            r#"
346            SELECT id, rule_description, affected_traits, confidence, source_count
347            FROM optimization_rules
348            WHERE tenant_id = $1
349            ORDER BY confidence DESC, created_at DESC
350            LIMIT 50
351            "#,
352        )
353        .bind(tenant_id)
354        .fetch_all(&self.pool)
355        .await?;
356
357        let mut rules = Vec::new();
358        for row in rows {
359            let id_str: String = row.try_get("id")?;
360            let traits_str: String = row.try_get("affected_traits")?;
361
362            rules.push(OptimizationRule {
363                id: uuid::Uuid::parse_str(&id_str).unwrap_or_default(),
364                rule_description: row.try_get("rule_description")?,
365                affected_traits: serde_json::from_str(&traits_str)?,
366                confidence: row.try_get("confidence")?,
367                source_experiments_count: row.try_get::<i64, _>("source_count")? as usize,
368                created_at: chrono::Utc::now(),
369            });
370        }
371
372        Ok(rules)
373    }
374
375    async fn count_experiments(&self, tenant_id: &str) -> Result<u64, EvolutionStoreError> {
376        use sqlx::Row;
377        let count: i64 =
378            sqlx::query("SELECT COUNT(*) as count FROM evolution_experiments WHERE tenant_id = $1")
379                .bind(tenant_id)
380                .fetch_one(&self.pool)
381                .await?
382                .try_get("count")?;
383
384        Ok(count as u64)
385    }
386}