the-grid 0.2.1

An AI-powered agentic operating environment
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

use crate::Event;
use chrono::{DateTime, Utc};
use rusqlite::{params, Connection, Result};
use std::collections::HashMap;
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct AgentState {
    pub name: String,
    pub personality: String,
    pub memory: Vec<Event>,
    pub last_seen: DateTime<Utc>,
    pub mood: String,
    pub xp: u32,
    pub active_task: Option<String>,
}

pub struct Database {
    conn: Connection,
}

impl Database {
    pub fn new() -> Result<Self> {
        let conn = Connection::open("the_grid.db")?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS agents (
                name        TEXT PRIMARY KEY,
                personality TEXT NOT NULL,
                memory      TEXT,
                last_seen   TEXT NOT NULL,
                mood        TEXT NOT NULL DEFAULT 'bored'
            )",
            [],
        )?;

        // Attempt to add column in case the table already existed without it (migration)
        let _ = conn.execute("ALTER TABLE agents ADD COLUMN mood TEXT NOT NULL DEFAULT 'bored'", []);
        let _ = conn.execute("ALTER TABLE agents ADD COLUMN xp INTEGER NOT NULL DEFAULT 0", []);
        let _ = conn.execute("ALTER TABLE agents ADD COLUMN active_task TEXT", []);

        // Add relationships table
        conn.execute(
            "CREATE TABLE IF NOT EXISTS relationships (
                source_agent TEXT NOT NULL,
                target_agent TEXT NOT NULL,
                affinity     INTEGER NOT NULL,
                PRIMARY KEY (source_agent, target_agent)
            )",
            [],
        )?;

        Ok(Database { conn })
    }

    pub fn save_agent_state(&self, state: &AgentState) -> Result<()> {
        let memory_json = serde_json::to_string(&state.memory)
            .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;

        self.conn.execute(
            "INSERT OR REPLACE INTO agents (name, personality, memory, last_seen, mood, xp, active_task) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
            params![
                &state.name,
                &state.personality,
                memory_json,
                &state.last_seen.to_rfc3339(),
                &state.mood,
                &state.xp,
                &state.active_task,
            ],
        )?;
        Ok(())
    }

    pub fn get_agent_state(&self, name: &str) -> Result<Option<AgentState>> {
        let mut stmt = self.conn.prepare("SELECT personality, memory, last_seen, mood, xp, active_task FROM agents WHERE name = ?1")?;
        
        if let Some(row_result) = stmt.query_map([name], |row| {
            Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?, row.get(4)?, row.get(5)?))
        })?.next() {
            let (personality, memory_json, last_seen_str, mood, xp, active_task): (String, String, String, String, u32, Option<String>) = row_result?;
            
            let memory: Vec<Event> = serde_json::from_str(&memory_json).unwrap_or_default();
            let last_seen = DateTime::parse_from_rfc3339(&last_seen_str)
                .map(|dt| dt.with_timezone(&Utc))
                .unwrap_or_else(|_| Utc::now());

            return Ok(Some(AgentState {
                name: name.to_string(),
                personality,
                memory,
                last_seen,
                mood,
                xp,
                active_task,
            }));
        }

        Ok(None)
    }

    pub fn update_relationship(&self, source: &str, target: &str, change: i32) -> Result<()> {
        let mut stmt = self.conn.prepare("SELECT affinity FROM relationships WHERE source_agent = ?1 AND target_agent = ?2")?;
        let current_affinity: i32 = stmt.query_map(params![source, target], |row| row.get(0))?.next().unwrap_or(Ok(0))?;
        
        let new_affinity = (current_affinity + change).clamp(-100, 100);

        self.conn.execute(
            "INSERT OR REPLACE INTO relationships (source_agent, target_agent, affinity) VALUES (?1, ?2, ?3)",
            params![source, target, new_affinity],
        )?;
        Ok(())
    }

    pub fn get_relationships(&self, source: &str) -> Result<HashMap<String, i32>> {
        let mut stmt = self.conn.prepare("SELECT target_agent, affinity FROM relationships WHERE source_agent = ?1")?;
        let relationships = stmt.query_map([source], |row| {
            Ok((row.get(0)?, row.get(1)?))
        })?.filter_map(Result::ok).collect::<HashMap<String, i32>>();
        Ok(relationships)
    }
}