Skip to main content

engram_storage/
q_table_store.rs

1use 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}