diff --git a/src/core/config.rs b/src/core/config.rs index 9185e78..684ca10 100644 --- a/src/core/config.rs +++ b/src/core/config.rs @@ -5,10 +5,13 @@ use std::path::PathBuf; use chrono::Utc; +pub const DEFAULT_MEMORY: u64 = 100; + pub struct Config { pub path: Option, pub did: Option, pub handle: Option, + pub memory: u64, } #[derive(Debug, Deserialize)] @@ -21,6 +24,7 @@ struct BotConfig { did: Option, handle: Option, path: Option, + memory: Option, } pub fn config_file() -> PathBuf { @@ -49,6 +53,7 @@ pub fn load() -> Config { path: bot.path, did: bot.did, handle: bot.handle, + memory: bot.memory.unwrap_or(DEFAULT_MEMORY), }; } } @@ -58,6 +63,7 @@ pub fn load() -> Config { path: None, did: None, handle: None, + memory: DEFAULT_MEMORY, } } diff --git a/src/core/reader.rs b/src/core/reader.rs index c179df5..2f2e7ce 100644 --- a/src/core/reader.rs +++ b/src/core/reader.rs @@ -14,24 +14,39 @@ pub fn read_core() -> Result { Ok(record) } -pub fn read_memory() -> Result> { +pub fn read_memory_all() -> Result> { let cfg = config::load(); let dir = config::collection_dir(&cfg, "ai.syui.gpt.memory"); if !dir.exists() { - return Ok(None); + return Ok(Vec::new()); } let mut files: Vec<_> = fs::read_dir(&dir)? .filter_map(|e| e.ok()) .filter(|e| e.path().extension().is_some_and(|ext| ext == "json")) .collect(); - if files.is_empty() { - return Ok(None); - } files.sort_by_key(|e| e.file_name()); - let latest = files.last().unwrap().path(); - let content = fs::read_to_string(&latest) - .with_context(|| format!("Failed to read {}", latest.display()))?; - let record: Value = serde_json::from_str(&content) - .with_context(|| format!("Failed to parse {}", latest.display()))?; - Ok(Some(record)) + + let mut records = Vec::new(); + for entry in &files { + let path = entry.path(); + let content = fs::read_to_string(&path) + .with_context(|| format!("Failed to read {}", path.display()))?; + let record: Value = serde_json::from_str(&content) + .with_context(|| format!("Failed to parse {}", path.display()))?; + records.push(record); + } + Ok(records) +} + +pub fn memory_count() -> usize { + let cfg = config::load(); + let dir = config::collection_dir(&cfg, "ai.syui.gpt.memory"); + fs::read_dir(&dir) + .map(|entries| { + entries + .filter_map(|e| e.ok()) + .filter(|e| e.path().extension().is_some_and(|ext| ext == "json")) + .count() + }) + .unwrap_or(0) } diff --git a/src/core/writer.rs b/src/core/writer.rs index 51eb3c4..95bf906 100644 --- a/src/core/writer.rs +++ b/src/core/writer.rs @@ -4,7 +4,7 @@ use serde_json::json; use std::fs; use std::time::{SystemTime, UNIX_EPOCH}; -use crate::core::{config, reader}; +use crate::core::config; fn generate_tid() -> String { // ATProto TID: 64-bit integer as 13 base32-sortable chars @@ -23,18 +23,11 @@ fn generate_tid() -> String { tid } -fn next_version() -> u64 { - match reader::read_memory() { - Ok(Some(record)) => record["value"]["version"].as_u64().unwrap_or(0) + 1, - _ => 1, - } -} - +/// Save a single memory element as a new TID file pub fn save_memory(content: &str) -> Result<()> { let cfg = config::load(); let did = cfg.did.clone().unwrap_or_else(|| "self".to_string()); let tid = generate_tid(); - let version = next_version(); let now = Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string(); let record = json!({ @@ -46,7 +39,6 @@ pub fn save_memory(content: &str) -> Result<()> { "$type": "ai.syui.gpt.memory#markdown", "text": content }, - "version": version, "createdAt": now } }); @@ -59,3 +51,62 @@ pub fn save_memory(content: &str) -> Result<()> { fs::write(&path, json_str) .with_context(|| format!("Failed to write {}", path.display())) } + +/// Delete all memory files, then write new ones from the given items +pub fn compress_memory(items: &[String]) -> Result<()> { + let cfg = config::load(); + let did = cfg.did.clone().unwrap_or_else(|| "self".to_string()); + let dir = config::collection_dir(&cfg, "ai.syui.gpt.memory"); + + // delete all existing memory files + if dir.exists() { + for entry in fs::read_dir(&dir)? { + let entry = entry?; + if entry.path().extension().is_some_and(|ext| ext == "json") { + fs::remove_file(entry.path())?; + } + } + } + + fs::create_dir_all(&dir) + .with_context(|| format!("Failed to create {}", dir.display()))?; + + // write each item as a new TID file + for item in items { + let micros = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_micros() as u64; + // small delay to ensure unique TIDs + std::thread::sleep(std::time::Duration::from_micros(1)); + + let v = (micros << 10) & 0x7FFFFFFFFFFFFFFF; + const CHARSET: &[u8] = b"234567abcdefghijklmnopqrstuvwxyz"; + let mut tid = String::with_capacity(13); + for i in (0..13).rev() { + let idx = ((v >> (i * 5)) & 0x1f) as usize; + tid.push(CHARSET[idx] as char); + } + + let now = Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string(); + let record = json!({ + "uri": format!("at://{}/ai.syui.gpt.memory/{}", did, tid), + "value": { + "$type": "ai.syui.gpt.memory", + "did": did, + "content": { + "$type": "ai.syui.gpt.memory#markdown", + "text": item + }, + "createdAt": now + } + }); + + let path = dir.join(format!("{}.json", tid)); + let json_str = serde_json::to_string_pretty(&record)?; + fs::write(&path, json_str) + .with_context(|| format!("Failed to write {}", path.display()))?; + } + + Ok(()) +} diff --git a/src/main.rs b/src/main.rs index 30624dc..002c4cb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -18,13 +18,13 @@ enum Commands { /// Start MCP server (JSON-RPC over stdio) Server, - /// Read core.json + /// Read core record ReadCore, - /// Read latest memory record + /// Read all memory records ReadMemory, - /// Create new memory record + /// Add a single memory element SaveMemory { /// Content to write content: String, @@ -51,15 +51,19 @@ fn main() -> Result<()> { } Some(Commands::ReadMemory) => { - match reader::read_memory()? { - Some(record) => println!("{}", serde_json::to_string_pretty(&record)?), - None => println!("No memory records found"), + let records = reader::read_memory_all()?; + if records.is_empty() { + println!("No memory records found"); + } else { + for record in &records { + println!("{}", serde_json::to_string_pretty(record)?); + } } } Some(Commands::SaveMemory { content }) => { writer::save_memory(&content)?; - println!("Saved."); + println!("Saved. ({} records)", reader::memory_count()); } } @@ -72,30 +76,17 @@ fn print_status() { let handle = cfg.handle.clone().unwrap_or_else(|| "self".to_string()); let base = config::base_dir(&cfg); let id = config::identity(&cfg); - - let memory_dir = config::collection_dir(&cfg, "ai.syui.gpt.memory"); - let memory_count = std::fs::read_dir(&memory_dir) - .map(|entries| { - entries - .filter_map(|e| e.ok()) - .filter(|e| e.path().extension().is_some_and(|ext| ext == "json")) - .count() - }) - .unwrap_or(0); - - let latest_version = match reader::read_memory() { - Ok(Some(record)) => record["value"]["version"].as_u64().unwrap_or(0), - _ => 0, - }; + let count = reader::memory_count(); println!("aigpt - AI memory MCP server\n"); println!("config: {}", config::config_file().display()); println!("did: {}", did); println!("handle: {}", handle); + println!("memory: {}", cfg.memory); println!(); println!("path: {}/", base.display()); println!(" {}/{}", id, "ai.syui.gpt.core/self.json"); println!(" {}/{}", id, "ai.syui.gpt.memory/*.json"); println!(); - println!("memory: {} records (version: {})", memory_count, latest_version); + println!("records: {}/{}", count, cfg.memory); } diff --git a/src/mcp/server.rs b/src/mcp/server.rs index 2bb1583..845ef30 100644 --- a/src/mcp/server.rs +++ b/src/mcp/server.rs @@ -90,8 +90,9 @@ impl MCPServer { } } - if let Ok(Some(memory)) = reader::read_memory() { - if let Some(text) = memory["value"]["content"]["text"].as_str() { + let records = reader::read_memory_all().unwrap_or_default(); + for record in &records { + if let Some(text) = record["value"]["content"]["text"].as_str() { if !text.is_empty() { parts.push(text.to_string()); } @@ -105,7 +106,7 @@ impl MCPServer { let tools = vec![ json!({ "name": "read_core", - "description": "Read core.md - the AI's identity and instructions", + "description": "Read the AI's identity and instructions (core record)", "inputSchema": { "type": "object", "properties": {} @@ -113,7 +114,7 @@ impl MCPServer { }), json!({ "name": "read_memory", - "description": "Read memory.md - the AI's accumulated memories", + "description": "Read all memory records. Each record is a single memory element.", "inputSchema": { "type": "object", "properties": {} @@ -121,13 +122,13 @@ impl MCPServer { }), json!({ "name": "save_memory", - "description": "Overwrite memory.md with new content", + "description": "Add a single memory element as a new record", "inputSchema": { "type": "object", "properties": { "content": { "type": "string", - "description": "Content to write to memory.md" + "description": "A single memory element to save" } }, "required": ["content"] @@ -135,16 +136,17 @@ impl MCPServer { }), json!({ "name": "compress", - "description": "Compress conversation into memory. AI decides what to keep, tool writes the result to memory.md", + "description": "Replace all memory records with a compressed set. Deletes all existing records and creates new ones from the provided items.", "inputSchema": { "type": "object", "properties": { - "conversation": { - "type": "string", - "description": "Compressed memory content to save" + "items": { + "type": "array", + "items": { "type": "string" }, + "description": "Array of memory elements to keep after compression" } }, - "required": ["conversation"] + "required": ["items"] } }), ]; @@ -192,9 +194,8 @@ impl MCPServer { } fn tool_read_memory(&self) -> Value { - match reader::read_memory() { - Ok(Some(record)) => record, - Ok(None) => json!({ "content": "" }), + match reader::read_memory_all() { + Ok(records) => json!({ "records": records, "count": records.len() }), Err(e) => json!({ "error": e.to_string() }), } } @@ -202,15 +203,23 @@ impl MCPServer { fn tool_save_memory(&self, arguments: &Value) -> Value { let content = arguments["content"].as_str().unwrap_or(""); match writer::save_memory(content) { - Ok(()) => json!({ "success": true }), + Ok(()) => json!({ "success": true, "count": reader::memory_count() }), Err(e) => json!({ "error": e.to_string() }), } } fn tool_compress(&self, arguments: &Value) -> Value { - let conversation = arguments["conversation"].as_str().unwrap_or(""); - match writer::save_memory(conversation) { - Ok(()) => json!({ "success": true }), + let items: Vec = arguments["items"] + .as_array() + .map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(|s| s.to_string())) + .collect() + }) + .unwrap_or_default(); + + match writer::compress_memory(&items) { + Ok(()) => json!({ "success": true, "count": items.len() }), Err(e) => json!({ "error": e.to_string() }), } }