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 async fn save_experiment(
17 &self,
18 tenant_id: &str,
19 experiment: &GenomeExperiment,
20 ) -> Result<(), EvolutionStoreError>;
21
22 async fn load_recent(
24 &self,
25 tenant_id: &str,
26 limit: usize,
27 ) -> Result<Vec<GenomeExperiment>, EvolutionStoreError>;
28
29 async fn save_rule(
31 &self,
32 tenant_id: &str,
33 rule: &OptimizationRule,
34 ) -> Result<(), EvolutionStoreError>;
35
36 async fn load_rules(
38 &self,
39 tenant_id: &str,
40 ) -> Result<Vec<OptimizationRule>, EvolutionStoreError>;
41
42 async fn count_experiments(&self, tenant_id: &str) -> Result<u64, EvolutionStoreError>;
44}
45
46#[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 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(), 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(), };
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#[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}