engram_storage/
q_table_store.rs1use rusqlite::params;
2
3use crate::database::Database;
4use crate::error::StorageError;
5
6impl Database {
7 pub fn upsert_q_value(
8 &self,
9 level: &str,
10 state: &str,
11 action: &str,
12 value: f32,
13 timestamp: &str,
14 ) -> Result<(), StorageError> {
15 self.connection().execute(
16 "INSERT INTO q_table (router_level, state, action, value, update_count, updated_at)
17 VALUES (?1, ?2, ?3, ?4, 1, ?5)
18 ON CONFLICT(router_level, state, action)
19 DO UPDATE SET value = ?4, update_count = update_count + 1, updated_at = ?5",
20 params![level, state, action, value, timestamp],
21 )?;
22 Ok(())
23 }
24
25 pub fn get_q_value(&self, level: &str, state: &str, action: &str) -> Result<f32, StorageError> {
26 let result = self.connection().query_row(
27 "SELECT value FROM q_table WHERE router_level = ?1 AND state = ?2 AND action = ?3",
28 params![level, state, action],
29 |row| row.get(0),
30 );
31 match result {
32 Ok(value) => Ok(value),
33 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(0.0),
34 Err(other) => Err(StorageError::Sqlite(other)),
35 }
36 }
37
38 pub fn load_q_table(
39 &self,
40 level: &str,
41 ) -> Result<Vec<(String, String, f32, u32)>, StorageError> {
42 let mut statement = self.connection().prepare(
43 "SELECT state, action, value, update_count
44 FROM q_table
45 WHERE router_level = ?1",
46 )?;
47 let rows = statement.query_map(params![level], |row| {
48 Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?))
49 })?;
50 let mut results = Vec::new();
51 for row in rows {
52 results.push(row?);
53 }
54 Ok(results)
55 }
56}