diff --git a/Cargo.lock b/Cargo.lock index a10d4c99..6ba2e32a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -21,6 +21,7 @@ dependencies = [ name = "ai-mate" version = "0.3.4" dependencies = [ + "anndists", "anyhow", "bytes", "clap", @@ -33,6 +34,7 @@ dependencies = [ "flate2", "futures-util", "hex", + "hnsw_rs", "hound", "kokoro-micro", "log", @@ -42,6 +44,7 @@ dependencies = [ "serde_ini", "serde_json", "sha2", + "sled", "tar", "tokio", "url", @@ -49,6 +52,12 @@ dependencies = [ "whisper-rs", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "alsa" version = "0.9.1" @@ -71,6 +80,23 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "anndists" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8238f99889a837cd6641360f9f3ead18f70b07bf6ce1f04a319bc6bd8a2f48f1" +dependencies = [ + "anyhow", + "cfg-if", + "cpu-time", + "env_logger", + "lazy_static", + "log", + "num-traits", + "num_cpus", + "rayon", +] + [[package]] name = "anstream" version = "0.6.21" @@ -160,6 +186,15 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + [[package]] name = "bindgen" version = "0.69.5" @@ -441,7 +476,17 @@ dependencies = [ "wasm-bindgen", "wasm-bindgen-futures", "web-sys", - "windows", + "windows 0.54.0", +] + +[[package]] +name = "cpu-time" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9e393a7668fe1fad3075085b86c781883000b4ede868f43627b34a87c8b7ded" +dependencies = [ + "libc", + "winapi", ] [[package]] @@ -471,6 +516,25 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -487,7 +551,7 @@ dependencies = [ "crossterm_winapi", "libc", "mio 0.8.11", - "parking_lot", + "parking_lot 0.12.5", "signal-hook", "signal-hook-mio", "winapi", @@ -596,6 +660,18 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "enum-as-inner" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1e6a265c649f3f5979b601d26f1d05ada116434c87741c9493cb56218f76cbc" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "env_filter" version = "1.0.0" @@ -738,6 +814,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -806,6 +892,15 @@ dependencies = [ "slab", ] +[[package]] +name = "fxhash" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c31b6d751ae2c7f11320402d34e41349dd1016f8d5d45e48c4312bc8625af50c" +dependencies = [ + "byteorder", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -887,6 +982,8 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ + "allocator-api2", + "equivalent", "foldhash", ] @@ -902,6 +999,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hex" version = "0.4.3" @@ -914,6 +1017,31 @@ version = "1.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f" +[[package]] +name = "hnsw_rs" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a5258f079b97bf2e8311ff9579e903c899dcbac0d9a138d62e9a066778bd07" +dependencies = [ + "anndists", + "anyhow", + "bincode", + "cfg-if", + "cpu-time", + "env_logger", + "hashbrown 0.15.5", + "indexmap", + "lazy_static", + "log", + "mmap-rs", + "num-traits", + "num_cpus", + "parking_lot 0.12.5", + "rand", + "rayon", + "serde", +] + [[package]] name = "home" version = "0.5.12" @@ -1158,9 +1286,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.13.0" +version = "2.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" +checksum = "0ad4bb2b565bca0645f4d68c5c9af97fba094e9791da685bf83cb5f3ce74acf2" dependencies = [ "equivalent", "hashbrown 0.16.1", @@ -1168,6 +1296,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "instant" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" +dependencies = [ + "cfg-if", +] + [[package]] name = "ipnet" version = "2.11.0" @@ -1465,6 +1602,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "mmap-rs" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ecce9d566cb9234ae3db9e249c8b55665feaaf32b0859ff1e27e310d2beb3d8" +dependencies = [ + "bitflags 2.11.0", + "combine", + "libc", + "mach2", + "nix", + "sysctl", + "thiserror 2.0.18", + "widestring", + "windows 0.48.0", +] + [[package]] name = "native-tls" version = "0.2.16" @@ -1540,6 +1694,18 @@ dependencies = [ "jni-sys", ] +[[package]] +name = "nix" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +dependencies = [ + "bitflags 2.11.0", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "nom" version = "7.1.3" @@ -1598,6 +1764,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "num_enum" version = "0.7.5" @@ -1739,6 +1915,17 @@ dependencies = [ "ureq", ] +[[package]] +name = "parking_lot" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d17b78036a60663b797adeaee46f5c9dfebb86948d1255007a1d6be0271ff99" +dependencies = [ + "instant", + "lock_api", + "parking_lot_core 0.8.6", +] + [[package]] name = "parking_lot" version = "0.12.5" @@ -1746,7 +1933,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" dependencies = [ "lock_api", - "parking_lot_core", + "parking_lot_core 0.9.12", +] + +[[package]] +name = "parking_lot_core" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a2cfe6f0ad2bfc16aefa463b497d5c7a5ecd44a23efa72aa342d90177356dc" +dependencies = [ + "cfg-if", + "instant", + "libc", + "redox_syscall 0.2.16", + "smallvec", + "winapi", ] [[package]] @@ -2017,6 +2218,35 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" +[[package]] +name = "rayon" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" +dependencies = [ + "bitflags 1.3.2", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2408,6 +2638,22 @@ version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" +[[package]] +name = "sled" +version = "0.34.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f96b4737c2ce5987354855aed3797279def4ebf734436c6aa4552cf8e169935" +dependencies = [ + "crc32fast", + "crossbeam-epoch", + "crossbeam-utils", + "fs2", + "fxhash", + "libc", + "log", + "parking_lot 0.11.2", +] + [[package]] name = "smallvec" version = "1.15.1" @@ -2484,6 +2730,20 @@ dependencies = [ "syn", ] +[[package]] +name = "sysctl" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01198a2debb237c62b6826ec7081082d951f46dbb64b0e8c7649a452230d1dfc" +dependencies = [ + "bitflags 2.11.0", + "byteorder", + "enum-as-inner", + "libc", + "thiserror 1.0.69", + "walkdir", +] + [[package]] name = "system-configuration" version = "0.7.0" @@ -3108,6 +3368,12 @@ dependencies = [ "fs_extra", ] +[[package]] +name = "widestring" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471" + [[package]] name = "winapi" version = "0.3.9" @@ -3139,6 +3405,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f" +dependencies = [ + "windows-targets 0.48.5", +] + [[package]] name = "windows" version = "0.54.0" diff --git a/Cargo.toml b/Cargo.toml index e58bad86..dfb737ec 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,9 @@ hex = "0.4" directories = "5" log = "0.4" env_logger = "0.11" +hnsw_rs = "0.3.4" +anndists = "0.1.4" +sled = "0.34" [build-dependencies] sha2 = "0.10" diff --git a/README.md b/README.md index 476e0d74..54253c56 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,8 @@ You can also initialize ai-mate at specific agent: ai-mate --agent planner ``` +If you want to use OpenTTS, start the docker service first: `docker run --rm --platform=linux/amd64 -p 5500:5500 synesthesiam/opentts:all` (it will pull the image the first time). Adjust the platform as needed depending on your hardware. + If you need help: ``` diff --git a/demo/ai-mate 0.3.0 - English speech demo.mov b/demo/ai-mate 0.3.0 - English speech demo.mov new file mode 100644 index 00000000..0a8cc94d Binary files /dev/null and b/demo/ai-mate 0.3.0 - English speech demo.mov differ diff --git a/demo/ai-mate 0.3.0 - Spanish speech demo.mov b/demo/ai-mate 0.3.0 - Spanish speech demo.mov new file mode 100644 index 00000000..1024c2d2 Binary files /dev/null and b/demo/ai-mate 0.3.0 - Spanish speech demo.mov differ diff --git a/demo/e-thumb.png b/demo/e-thumb.png new file mode 100644 index 00000000..788b953c Binary files /dev/null and b/demo/e-thumb.png differ diff --git a/demo/s-thumb.png b/demo/s-thumb.png new file mode 100644 index 00000000..467e9ca3 Binary files /dev/null and b/demo/s-thumb.png differ diff --git a/spec/acceleration.txt b/spec/acceleration.txt index 51b1d86a..a7d688cf 100644 --- a/spec/acceleration.txt +++ b/spec/acceleration.txt @@ -21,4 +21,4 @@ ----------------------------------------------------------------- ARM64: Metal support + Accelerate + eigen ------------------------------------------------------------------ \ No newline at end of file +----------------------------------------------------------------- diff --git a/spec/memory.md b/spec/memory.md new file mode 100644 index 00000000..c4b4df18 --- /dev/null +++ b/spec/memory.md @@ -0,0 +1,95 @@ +# Memory System Overview + +The **Memory** module is a lightweight knowledge‑base that stores *knowledge units* together with their embeddings. It is used by the AI‑assistant to recall past events and facts. + +## 1. Data Model + +- **KnowledgeUnit** – a semantic triple consisting of `subject`, `predicate`, `object`, optional `location` and a timestamp. +- **VecKnowledgeUnit** – the unit plus a 128‑dimensional embedding vector. +- **Memory** – holds the HNSW graph (`hnsw`), a map of IDs to `VecKnowledgeUnit` (`index_map`), and three *inverted indexes* (subject, predicate, object) that map a value to the set of IDs that contain it. + +## 2. Persistence + +The entire state is stored in a single **sled** database file named `memory.json` (the name is a bit of a misnomer – it is a sled DB, not a plain JSON file). The DB contains one key: + +``` +"memory" => +``` + +The bytes are a JSON object with the following fields: + +|-------------------------------------------------------------------------------------| +| Field | Type | Meaning | +|---------------------|----------------------------|----------------------------------| +| `next_id` | u64 | Next available ID for a new unit | +| `units` | map | All stored units, keyed by ID | +| `subject_index` | map> | Inverted index for subjects | +| `predicate_index` | map> | Inverted index for predicates | +| `object_index` | map> | Inverted index for objects | +--------------------------------------------------------------------------------------- + +When `Memory::save_to_file` is called it serialises this structure with `serde_json::to_vec` and writes it to sled. Loading reverses the process: the bytes are read, deserialised, and the HNSW graph is rebuilt from the embeddings. + +## 3. Storing a Unit + +``` +let id = memory.next_id; +memory.index_map.insert(id, VecKnowledgeUnit{embedding, knowledge: unit.clone()}); +memory.hnsw.insert((&embedding, id)); +memory.next_id += 1; +// update inverted indexes +subject_index.entry(unit.subject).or_default().push(id); +predicate_index.entry(unit.predicate.name).or_default().push(id); +object_index.entry(unit.object).or_default().push(id); +``` + +After the in‑memory state is updated, `save_to_file` persists the data. + +## 4. Querying + +### a. Full‑text / Embedding query + +`Memory::query(query, k, ef_search)` performs a vector similarity search: +1. The query string is turned into a 128‑dimensional vector via `embed_text` (hash‑based word frequency). +2. HNSW returns the `k` nearest neighbour IDs. +3. The corresponding `KnowledgeUnit`s are returned. +4. A final filter keeps only results that contain any word from the query (case‑insensitive, alphanumeric only). + +### b. Inverted‑index lookup + +Functions `get_by_subject`, `get_by_predicate`, `get_by_object`, and `get_by_location` simply look up the relevant ID list in the inverted index and fetch the units. This is O(1) to get the list plus O(1) per result. + +## 5. CLI Commands + +The binary exposes four new flags: + +``` +--get-memories-by-subject +--get-memories-by-predicate +--get-memories-by-object +--get-memories-by-location +``` + +Each flag loads `memory.json`, calls the appropriate `get_by_*` method, and prints each unit’s context via `Memory::build_context_from_units`. + +## 6. Summary + +* Units are stored with an embedding and an ID. +* All data is persisted as a single sled database entry containing JSON. +* Retrieval can be done either by embedding similarity or by fast inverted‑index lookups. +* The CLI provides convenient commands for inspecting and querying the memory. + +This design keeps the runtime lightweight while enabling efficient queries over thousands of units. + +## 7. LLM integration + +Store memories: + +1. when sending llm prompt, provide "story_memory" tool with available predicates +2. execute tool calls (new memories to store) from LLM response + +Retrieve memories: + +1. Before submitting prompt, query memories +2. Collect results (knowledge units) and produce text readable sentences using build_context_from_units(&retrieved_units); +3. Provide context to the llm diff --git a/spec/messaging.txt b/spec/messaging.txt new file mode 100644 index 00000000..a353d6bf --- /dev/null +++ b/spec/messaging.txt @@ -0,0 +1,55 @@ +Threads and the channels they use + +------------------------------------------------------------------------------------ + Main (src/main.rs) +------------------------------------------------------------------------------------ +- Sends the global stop signal through stop_all_tx. +- Sends raw microphone audio to the recording thread via tx_rec. +- Sends recognized utterances to the conversation thread via tx_utt. +- Sends raw microphone audio to the conversation thread via conv_tx (for synchronous processing). +- Sends playback audio to the playback thread through tx_play. +- Sends playback‑stop requests to the recording and playback threads via stop_play_tx. +- Sends the playback‑stop receiver to the playback thread via stop_play_rx. + +------------------------------------------------------------------------------------ + UI thread (src/ui.rs) +------------------------------------------------------------------------------------ +- Receives the global stop signal on stop_all_rx. +- Reads the shared status_line, print_lock, and peak to update the terminal. + +------------------------------------------------------------------------------------ + Playback thread (src/playback.rs) +------------------------------------------------------------------------------------ +- Receives the global stop signal on stop_all_rx. +- Receives playback audio chunks on rx_play. +- Receives a playback‑stop signal on stop_play_rx. +- Reads control flags: playback_active, gate_until_ms, paused. +- Reads UI state (ui) and updates it. +- Reads and updates the shared volume. + +------------------------------------------------------------------------------------ + Record thread (src/record.rs) +------------------------------------------------------------------------------------ +- Receives the global stop signal on stop_all_rx. +- Reads the interrupt counter via interrupt_counter. +- Reads and updates shared state: volume, recording_paused, and peak. +- Sends microphone audio to the conversation thread via conv_tx. +- Sends recognized utterance audio to the conversation thread via tx_utt. +- Sends a playback‑stop request on stop_play_tx. + +------------------------------------------------------------------------------------ + Conversation thread (src/conversation.rs) +------------------------------------------------------------------------------------ +- Receives the global stop signal on stop_all_rx. +- Receives utterance chunks from the record thread on rx_utt. +- Sends processed speech audio to the playback thread via tx_play. +- Sends a global stop signal on stop_all_tx when needed. + +------------------------------------------------------------------------------------ + Keyboard thread (src/keyword.rs) +------------------------------------------------------------------------------------ +- Receives the global stop signal on stop_all_rx. +- Sends the global stop signal via stop_all_tx. +- Reads control flags: paused, playback_active. +- Reads the recording‑paused flag via recording_paused. +- Reads configuration: voice, args.tts, args.language. diff --git a/src/config.rs b/src/config.rs index db2cb6ff..a912c677 100644 --- a/src/config.rs +++ b/src/config.rs @@ -100,6 +100,29 @@ pub struct Args { #[arg(long)] pub ptt: Option, + + #[arg(long, action=clap::ArgAction::SetTrue)] + pub get_memories: bool, + + /// Get memories by subject + #[arg(long, value_name="SUBJECT", action=clap::ArgAction::Set)] + pub get_memories_by_subject: Option, + + /// Get memories by predicate + #[arg(long, value_name="PREDICATE", action=clap::ArgAction::Set)] + pub get_memories_by_predicate: Option, + + /// Get memories by object + #[arg(long, value_name="OBJECT", action=clap::ArgAction::Set)] + pub get_memories_by_object: Option, + + /// Get memories by location + #[arg(long, value_name="LOCATION", action=clap::ArgAction::Set)] + pub get_memories_by_location: Option, + + /// Query memory using embeddings and print results + #[arg(long, value_name="QUERY", action=clap::ArgAction::Set)] + pub query_memory: Option, } // internal static values diff --git a/src/conversation.rs b/src/conversation.rs index ae626a34..85f80a5b 100644 --- a/src/conversation.rs +++ b/src/conversation.rs @@ -4,6 +4,7 @@ use crate::state::GLOBAL_STATE; use crate::START_INSTANT; +use crate::llm::ChatMessage; use crossbeam_channel::{select, Receiver, Sender}; use std::cell::Cell; use std::sync::OnceLock; @@ -17,12 +18,6 @@ static WHISPER_CTX: OnceLock = OnceLock::new(); // API // ------------------------------------------------------------------ -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ChatMessage { - pub role: String, - pub content: String, -} - pub type ConversationHistory = std::sync::Arc>>; /// Initialise the Whisper context once, performing a warm‑up. @@ -44,13 +39,15 @@ pub fn conversation_thread( model_path: String, settings: crate::config::AgentSettings, ui: crate::state::UiState, - conversation_history: ConversationHistory, + conversation_history: std::sync::Arc>>, tx_ui: Sender, tts_tx: Sender<(String, u64)>, ) -> Result<(), Box> { let ctx = init_whisper_context(&model_path); crate::log::log("info", &format!("LLM model: {}", settings.model)); + let tts_tx_cloned = tts_tx.clone(); + let interrupt_counter_cloned = interrupt_counter.clone(); loop { select! { recv(stop_all_rx) -> _ => break, @@ -116,7 +113,7 @@ pub fn conversation_thread( let _ = tx_ui.send(format!("line|{}", user_text)); let _ = tx_ui.send("line|\n".to_string()); - conversation_history.lock().unwrap().push(ChatMessage{role:"user".to_string(), content:user_text.clone()}); + conversation_history.lock().unwrap().push(ChatMessage { role: "user".into(), content: user_text.clone() }); ui.thinking.store(true, Ordering::Relaxed); // Snapshot interruption counter for this assistant turn. @@ -173,46 +170,109 @@ pub fn conversation_thread( let _ = tx_ui_cloned_for_closure.send(format!("stream|{}", piece)); }; - let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); + let user_text2 = user_text.clone(); + let rt_llm_call = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); let stop_all_rx_cloned = stop_all_rx.clone(); - let ollama_url = state.baseurl.lock().unwrap().clone(); + let ollama_url = settings.baseurl.clone(); let interrupt_counter_cloned = interrupt_counter.clone(); - let llama_url = state.baseurl.lock().unwrap().clone(); - let model = state.model.lock().unwrap().clone(); - let engine_type = state.provider.lock().unwrap().clone(); + // llama_url not needed in this context, use settings.baseurl when required + let model = settings.model.clone(); + let engine_type = settings.provider.clone(); if *state.provider.lock().unwrap() == "llama-server" { let on_piece_cloned = std::sync::Arc::new(std::sync::Mutex::new(on_piece)); + let on_piece_cloned1 = on_piece_cloned.clone(); + let conv_hist_clone = conversation_history.clone(); + let value = user_text.clone(); + + let rt_llm_tools_call1 = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); + let conv_hist_clone2a = conversation_history.clone(); + let llama_url2a = settings.baseurl.clone(); + let model2a = settings.model.clone(); + let engine_type2a = settings.provider.clone(); + let my_interrupt2a = my_interrupt.clone(); + let stop_all_rx_cloned2a = stop_all_rx.clone(); + let interrupt_counter_cloned2a = interrupt_counter.clone(); + let on_piece_cloned12a = on_piece_cloned1.clone(); + let value2a = user_text.clone(); + let handle = std::thread::spawn(move || { - rt.block_on(async { - crate::log::log("info", "eoo"); - match crate::llm::llama_server_stream_response_into ( - &messages, - llama_url.as_str(), - model.as_str(), - engine_type.as_str(), - &stop_all_rx_cloned, - interrupt_counter_cloned.clone(), - my_interrupt, - &mut *on_piece_cloned.lock().unwrap() - ).await { - Ok(_) => Ok(()), - Err(e) => { - crate::log::log("error", &format!("llama server error: {e}. Make sure llama-server / llamafile is running")); - Err(e) + if crate::llm::TOOLS_SUPPORTED.get().copied().unwrap_or(false) { + rt_llm_tools_call1.block_on(async { + match crate::llm::llama_server_stream_response_into ( + &conv_hist_clone2a.lock().unwrap(), + false, // send only this user message + true, // include tools + &value2a, + llama_url2a.as_str(), + model2a.as_str(), + engine_type2a.as_str(), + &stop_all_rx_cloned2a, + interrupt_counter_cloned2a.clone(), + my_interrupt2a, + &mut *on_piece_cloned12a.lock().unwrap() + ).await { + Ok(_) => (), + Err(e) => { + crate::log::log("error", &format!("llama server error: {e}. Make sure llama-server / llamafile is running")); + } } - } - }) + }) + } + }); + + let rt_llm_tools_call2 = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); + let conv_hist_clone2b = conversation_history.clone(); + let llama_url2b = settings.baseurl.clone(); + let model2b = settings.model.clone(); + let engine_type2b = settings.provider.clone(); + let my_interrupt2b = my_interrupt.clone(); + let stop_all_rx_cloned2b = stop_all_rx.clone(); + let interrupt_counter_cloned2b = interrupt_counter.clone(); + let on_piece_cloned12b = on_piece_cloned1.clone(); + let value2b = user_text.clone(); + + let handle2 = std::thread::spawn(move || { + if crate::llm::TOOLS_SUPPORTED.get().copied().unwrap_or(false) { + rt_llm_tools_call2.block_on(async { + match crate::llm::llama_server_stream_response_into ( + &conv_hist_clone2b.lock().unwrap(), + true, // send only this user message + true, // include tools + &value2b, + llama_url2b.as_str(), + model2b.as_str(), + engine_type2b.as_str(), + &stop_all_rx_cloned2b, + interrupt_counter_cloned2b.clone(), + my_interrupt2b, + &mut *on_piece_cloned12b.lock().unwrap() + ).await { + Ok(_) => (), + Err(e) => { + crate::log::log("error", &format!("llama server error: {e}. Make sure llama-server / llamafile is running")); + } + } + }) + } }); - // ignore join result to prevent panic on llama server error - let _join_result = handle.join(); + + handle.join().unwrap(); + handle2.join().unwrap(); + } else { let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap(); let on_piece_cloned = std::sync::Arc::new(std::sync::Mutex::new(on_piece)); + let conv_hist_clone = conversation_history.clone(); + let handle = std::thread::spawn(move || { + // llm stream request rt.block_on(async { match crate::llm::llama_server_stream_response_into ( - &messages, + &conv_hist_clone.lock().unwrap(), + true, // include full history + false, // do not include tools + &user_text, ollama_url.as_str(), model.as_str(), engine_type.as_str(), @@ -223,14 +283,41 @@ pub fn conversation_thread( ).await { Ok(_) => Ok(()), Err(e) => { - crate::log::log("error", &format!("ollama error. {}. Make sure ollama is running and model '{}' is available", e, model.as_str())); + crate::log::log("error", &format!("ollama error. {e}. Make sure ollama is running")); Err(e) } } - }) + }); + + // ask for store memory tool calls using this last message + if crate::llm::TOOLS_SUPPORTED.get().copied().unwrap_or(false) { + rt.block_on(async { + match crate::llm::llama_server_stream_response_into ( + &conv_hist_clone.lock().unwrap(), + false, // send only this user message + true, // include tools + &user_text, + ollama_url.as_str(), + model.as_str(), + engine_type.as_str(), + &stop_all_rx_cloned, + interrupt_counter_cloned.clone(), + my_interrupt, + &mut *on_piece_cloned.lock().unwrap() + ).await { + Ok(_) => Ok(()), + Err(e) => { + crate::log::log("error", &format!("ollama error. {e}. Make sure ollama is running")); + Err(e) + } + } + }) + } else { + Ok(()) + } }); - // ignore join result to prevent panic on llama server error - let _join_result = handle.join(); + + handle.join().unwrap(); } ui_thinking_cloned_for_closure.store(false, Ordering::Relaxed); if let Some(phrase) = speaker_arc.lock().unwrap().flush() { @@ -298,7 +385,8 @@ fn strip_special_chars(s: &str) -> String { if !inside { result.extend(part.chars().filter(|c| { ![ - '+', '.', '~', '*', '&', '-', ',', ';', ':', '(', ')', '[', ']', '{', '}', '"', '\'', '#', '`', '|' + '+', '.', '~', '*', '&', '-', ',', ';', ':', '(', ')', '[', ']', '{', '}', '"', '\'', + '#', '`', '|', ] .contains(c) })); diff --git a/src/llm.rs b/src/llm.rs index c727ebd7..94df0357 100644 --- a/src/llm.rs +++ b/src/llm.rs @@ -2,29 +2,57 @@ // LLM handling // ------------------------------------------------------------------ -use std::sync::{Arc, atomic::AtomicU64}; +use std::sync::{Arc, atomic::AtomicU64, OnceLock}; use crossbeam_channel::Receiver; use reqwest::StatusCode; use futures_util::StreamExt; use bytes::Bytes; -use serde_json::json; +use serde_json::{json, Value}; +use serde::{Deserialize, Serialize}; +use crate::tools::store_memory::StoreMemoryTool; +use crate::tools::Tool; +use crate::tools::get_available_tools; +use crate::tools::handle_tool_call; +use std::collections::HashMap; +use crate::tools::remember::RememberTool; + +pub static TOOLS_SUPPORTED: OnceLock = OnceLock::new(); + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ChatMessage { + pub role: String, + pub content: String, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct ChatMessageRef<'a> { + role: &'a str, + content: &'a str, +} /// Stream response from Llama/Ollama endpoints, fallback if one fails, and mid-stream cancellation support pub async fn llama_server_stream_response_into( - messages: &Vec, + conversation_history: &[ChatMessage], + full_history: bool, + include_tools: bool, + user_prompt: &str, llama_host: &str, llama_model: &str, server_type: &str, stop_all_rx: &Receiver<()>, interrupt_counter: Arc, expected_interrupt: u64, - on_piece: &mut dyn FnMut(&str), + on_piece: &mut dyn FnMut(&str) ) -> Result<(), Box> { + let mut full_payload = String::new(); + let mut partial_tool_calls_buf: HashMap = HashMap::new(); + let mut name_map: HashMap = HashMap::new(); + let mut last_key: Option = None; + #[derive(Clone, Copy, Debug)] enum ApiKind { OaiChat, OllamaGenerate, OllamaChat } - fn should_fallback_status(code: StatusCode) -> bool { matches!( code, @@ -59,6 +87,7 @@ pub async fn llama_server_stream_response_into( let client = reqwest::Client::new(); let tries = candidates(llama_host, server_type); let mut last_err: Option = None; + let tool_schemas = get_available_tools()?; for (url, kind) in tries { if stop_all_rx.try_recv().is_ok() || @@ -68,30 +97,109 @@ pub async fn llama_server_stream_response_into( crate::log::log("info", &format!("Trying endpoint: {}", url)); + let new_user_prompt = ChatMessageRef { + role: "user", + content: user_prompt, + }; + + // contains the full prompt (with previous history) + let mut full_messages_vec: Vec = Vec::new(); + + // prepare messages + full_messages_vec.push(ChatMessageRef { + role: "system", + content: "You are a helpful assistant.", + }); + + // include the full chat history in the prompt + if full_history { + for m in conversation_history.iter() { + full_messages_vec.push(ChatMessageRef { + role: m.role.as_str(), + content: &m.content, + }); + } + } + + // add latest message at the end of history + full_messages_vec.push(ChatMessageRef { + role: "user", + content: user_prompt, + }); + + let messages = full_messages_vec; + let tools_supported = crate::llm::TOOLS_SUPPORTED.get().copied().unwrap_or(false); + let req = match kind { ApiKind::OaiChat => { let payload = json!({ "model": llama_model, - "messages": messages.iter().map(|m| json!({ "role": m.role, "content": m.content })).collect::>(), - "stream": true + "messages": messages, + "stream": true, + "tools": if include_tools && tools_supported { + Some(vec![StoreMemoryTool::json_schema()?]) + } else { + None + }, + "tool_choice": if include_tools && tools_supported { + Some("auto") + } else { + None + }, + "parallel_tool_calls": if include_tools && tools_supported { + Some(true) + } else { + None + }, }); client.post(&url).json(&payload) } + ApiKind::OllamaGenerate => { - let prompt_str = messages.iter().map(|m| m.content.as_str()).collect::>().join("\n"); let payload = json!({ "model": llama_model, - "prompt": prompt_str, + "prompt": user_prompt, "stream": true, - "max_tokens": 1024 + "max_tokens": 1024, + "tools": if include_tools && tools_supported { + Some(vec![StoreMemoryTool::json_schema()?]) + } else { + None + }, + "tool_choice": if include_tools && tools_supported { + Some("auto") + } else { + None + }, + "parallel_tool_calls": if include_tools && tools_supported { + Some(true) + } else { + None + } }); client.post(&url).json(&payload) } + ApiKind::OllamaChat => { let payload = json!({ "model": llama_model, - "messages": messages.iter().map(|m| json!({ "role": m.role, "content": m.content })).collect::>(), - "stream": true + "messages": messages, + "stream": true, + "tools": if include_tools && tools_supported { + Some(vec![StoreMemoryTool::json_schema()?]) + } else { + None + }, + "tool_choice": if include_tools && tools_supported { + Some("auto") + } else { + None + }, + "parallel_tool_calls": if include_tools && tools_supported { + Some(true) + } else { + None + } }); client.post(&url).json(&payload) } @@ -115,7 +223,11 @@ pub async fn llama_server_stream_response_into( let status = resp.status(); last_err = Some(format!("Endpoint {} returned HTTP {}", url, status)); log::warn!("{}", last_err.as_ref().unwrap()); - if should_fallback_status(status) { continue; } else { return Err(last_err.clone().unwrap().into()); } + if should_fallback_status(status) { + continue; + } else { + return Err(last_err.clone().unwrap().into()); + } } crate::log::log("info", &format!("Streaming response from: {}", url)); @@ -142,14 +254,36 @@ pub async fn llama_server_stream_response_into( // crate::log::log("debug", &format!("chunk: {}", text)); for line in text.lines() { let payload = line.trim().strip_prefix("data:").unwrap_or(line).trim(); - if payload == "[DONE]" { return Ok(()); } + if payload == "[DONE]" { + // crate::log::log("info", &format!("Full payload: {}", full_payload.to_string())); + return Ok(()); + } // parse JSON safely if let Ok(v) = serde_json::from_str::(payload) { // Handle new Llama3.2 style: {"message":{"content":...}} if let Some(message) = v.get("message") { if let Some(content) = message.get("content").and_then(|c| c.as_str()) { - if !content.is_empty() { on_piece(content); } + if !content.is_empty() { + on_piece(content); + full_payload.push_str(&content); + } + } + + if let Some(tool_calls_value) = message.get("tool_calls") { + if tools_supported { + match tool_calls_value { + serde_json::Value::Array(arr) => { + process_tool_calls_array( + arr, + &mut partial_tool_calls_buf, + &mut name_map, + &mut last_key, + ); + } + _ => {} + } + } } } else { match kind { @@ -158,10 +292,30 @@ pub async fn llama_server_stream_response_into( for choice in choices { if let Some(delta) = choice.get("delta") { if let Some(content) = delta.get("content").and_then(|c| c.as_str()) { - if !content.is_empty() { on_piece(content); } + if !content.is_empty() { + on_piece(content); + full_payload.push_str(&content); + } + if let Some(tool_calls_value) = delta.get("tool_calls") { + if tools_supported { + // tool_calls can be a string or an array of objects. Handle both. + match tool_calls_value { + serde_json::Value::Array(arr) => { + process_tool_calls_array( + arr, + &mut partial_tool_calls_buf, + &mut name_map, + &mut last_key, + ); + } + _ => {} + } + } + } } } if choice.get("finish_reason").and_then(|r| r.as_str()) == Some("stop") { + // crate::log::log("info", &format!("Full payload: {}", full_payload.to_string())); return Ok(()); } } @@ -169,6 +323,7 @@ pub async fn llama_server_stream_response_into( if v.get("done").and_then(|x| x.as_bool()) == Some(true) || v.get("status").and_then(|x| x.as_str()) == Some("completed") { + // crate::log::log("info", &format!("Full payload: {}", full_payload.to_string())); return Ok(()); } } @@ -185,4 +340,142 @@ pub async fn llama_server_stream_response_into( // all endpoints failed Err(last_err.unwrap_or_else(|| "No endpoint candidates succeeded".to_string()).into()) +} + + +pub async fn supports_tool_calls( + model: &str, + llm_engine_type: &str, + base_url: &str, +) -> Result> { + let client = reqwest::Client::new(); + + match llm_engine_type { + "ollama" => { + // Use /api/show with POST + let endpoint = format!("{}/api/show", base_url); + let payload = serde_json::json!({ "name": model }); + + let resp = client.post(&endpoint).json(&payload).send().await?; + + if !resp.status().is_success() { + return Ok(false); + } + + let data: serde_json::Value = resp.json().await?; + if let Some(capabilities) = data.get("capabilities").and_then(|v| v.as_array()) { + // Check if "tools" is in capabilities + return Ok(capabilities.iter().any(|c| c.as_str() == Some("tools"))); + } + + Ok(false) + } + + "llama-server" => { + // Minimal test: send a prompt asking for JSON tool call + let payload = serde_json::json!({ + "stream": false, + "messages": [ + { "role": "system", "content": "You are a helpful assistant." }, + { "role": "user", "content": "Calculate 1 + 1 using the available tools" } + ], + "max_tokens": 200, + "tools": [ + { + "type": "function", + "function": { + "name": "calculator", + "description": "Perform a calculation", + "parameters": { + "type": "object", + "properties":{ + "operation": { + "type": "string", + "description": "The operation to be calculated" + } + }, + "required":["operation"] + } + } + } + ] + }); + + let endpoint = format!("{}/v1/chat/completions", base_url); // wrapper endpoint + let resp = client.post(&endpoint).json(&payload).send().await?; + let data: Value = resp.json().await?; + + // crate::log::log("info", &format!("data: {:?}", data)); + + if let Some(choice) = data.get("choices").and_then(|c: &serde_json::Value| c.get(0)) { + if let Some(message) = choice.get("message") { + if message + .get("tool_calls") + .and_then(|arr: &serde_json::Value| arr.as_array()) + .map_or(false, |a: &Vec| !a.is_empty()) + { + return Ok(true); + } + } + } + + Ok(false) + } + _ => Ok(false), + } +} + + +// processes an array of tool calls, accumulating arguments across chunks +fn process_tool_calls_array( + arr: &Vec, + args_map: &mut HashMap, + name_map: &mut HashMap, + last_key: &mut Option, +) { + for tc in arr { + let key = if let Some(id) = tc.get("id").and_then(|v| v.as_str()) { + id.to_string() + } else { + last_key.clone().unwrap_or("__no_id__".to_string()) + }; + *last_key = Some(key.clone()); + if let Some(func_obj) = tc.get("function") { + // extract the tool name from the function object + let name = func_obj + .get("name") + .and_then(|n| n.as_str()) + .unwrap_or("unknown"); + // store name if not already + name_map + .entry(key.clone()) + .or_insert_with(|| name.to_string()); + if let Some(args_val) = func_obj.get("arguments") { + // buffers partial string arguments + // validates them as JSON and executes the tool call once a complete arguments payload is available + match args_val { + serde_json::Value::String(s) => { + let buf = args_map.entry(key.clone()).or_insert_with(String::new); + buf.push_str(s); + if serde_json::from_str::(buf).is_ok() { + let stored_name = name_map.get(&key).map(|s: &String| s.as_str()).unwrap_or("unknown"); + *buf = serde_json::to_string( + &serde_json::from_str::(buf).unwrap_or(serde_json::Value::Null), + ) + .unwrap_or("{}".to_string()); + let payload = format!(r#"{{"name":"{}","arguments":{}}}"#, stored_name, buf); + let _ = handle_tool_call(&payload); + buf.clear(); + } + } + _ => { + let args_str = serde_json::to_string(args_val).unwrap_or("{}".to_string()); + let stored_name = name_map.get(&key).map(|s: &String| s.as_str()).unwrap_or("unknown"); + let payload = format!(r#"{{"name":"{}","arguments":{}}}"#, stored_name, args_str); + let _ = handle_tool_call(&payload); + } + } + } + } + } } \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index c444a539..78385d06 100644 --- a/src/main.rs +++ b/src/main.rs @@ -23,14 +23,19 @@ mod stt; mod tts; mod ui; mod util; +mod tools; +mod memory; static START_INSTANT: OnceLock = OnceLock::new(); -fn main() -> Result<(), Box> { +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<(), Box> { let args = crate::config::Args::parse(); crate::log::set_verbose(args.verbose || false); let _ = START_INSTANT.get_or_init(Instant::now); + let args = crate::config::Args::parse(); + let memory_path = memory::ensure_memory_path(); // make sure piper phonemes are unpacked assets::ensure_piper_espeak_env(); @@ -71,6 +76,133 @@ fn main() -> Result<(), Box> { env_logger::init(); whisper_rs::install_logging_hooks(); + if args.get_memories { + use crate::memory::Memory; + use std::path::Path; + let path = memory_path.as_str(); + if Path::new(path).exists() { + let memory = Memory::load_from_file(path).expect("failed to load memory"); + for v in memory.index_map.values() { + let unit = v.knowledge.clone(); + let text = Memory::build_context_from_units(&[unit]); + println!("{}", text); + } + } else { + eprintln!("No memory file found"); + } + process::exit(0); + } + + // Get by subject + if let Some(subj) = args.get_memories_by_subject { + use crate::memory::Memory; + use std::path::Path; + let path = memory_path.as_str(); + if !Path::new(path).exists() { + eprintln!("No memory file found"); + process::exit(1); + } + let memory = Memory::load_from_file(path).expect("failed to load memory"); + let results = memory.get_by_subject(&subj, None, None, None); + for unit in results { + let text = Memory::build_context_from_units(&[unit]); + println!("{}", text); + } + process::exit(0); + } + + // Get by predicate + if let Some(pred) = args.get_memories_by_predicate { + use crate::memory::Memory; + use std::path::Path; + let path = memory_path.as_str(); + if !Path::new(path).exists() { + eprintln!("No memory file found"); + process::exit(1); + } + let memory = Memory::load_from_file(path).expect("failed to load memory"); + let results = memory.get_by_predicate(&pred, None, None, None); + for unit in results { + let text = Memory::build_context_from_units(&[unit]); + println!("{}", text); + } + process::exit(0); + } + + // Get by object + if let Some(obj) = args.get_memories_by_object { + use crate::memory::Memory; + use std::path::Path; + let path = memory_path.as_str(); + if !Path::new(path).exists() { + eprintln!("No memory file found"); + process::exit(1); + } + let memory = Memory::load_from_file(path).expect("failed to load memory"); + let results = memory.get_by_object(&obj, None, None, None); + for unit in results { + let text = Memory::build_context_from_units(&[unit]); + println!("{}", text); + } + process::exit(0); + } + + // Get by location + if let Some(loc) = args.get_memories_by_location { + use crate::memory::Memory; + use std::path::Path; + let path = memory_path.as_str(); + if !Path::new(path).exists() { + eprintln!("No memory file found"); + process::exit(1); + } + let memory = Memory::load_from_file(path).expect("failed to load memory"); + let results = memory.get_by_location(&loc, None, None); + for unit in results { + let text = Memory::build_context_from_units(&[unit]); + println!("{}", text); + } + process::exit(0); + } + + // Query memory + if let Some(query) = args.query_memory { + use crate::memory::Memory; + use std::path::Path; + let path = memory_path.as_str(); + if !Path::new(path).exists() { + eprintln!("No memory file found"); + process::exit(1); + } + let memory = Memory::load_from_file(path).expect("failed to load memory"); + let top_k = memory.index_map.len(); + let ef_search = 200; + let results = memory.query(&query, top_k, ef_search); + let cleaned_query = |s: &str| { + s.to_lowercase() + .replace(|c: char| !c.is_ascii_alphanumeric(), " ") + }; + let cleaned_query_str = cleaned_query(&query); + let query_words: Vec<&str> = cleaned_query_str.split_whitespace().collect(); + for unit in results { + let combined = format!( + "{} {} {} {}", + unit.subject, + unit.predicate.name, + unit.object, + unit.location.clone().unwrap_or_default() + ); + let cleaned_combined = cleaned_query(&combined); + if query_words.iter().any(|w| cleaned_combined.contains(w)) { + let text = Memory::build_context_from_units(&[unit]); + println!("{}", text); + } + } + process::exit(0); + } + + + // --------------------------------------------------- // Load Settings // --------------------------------------------------- @@ -105,6 +237,14 @@ fn main() -> Result<(), Box> { let ui = state.ui.clone(); let status_line = state.status_line.clone(); + // Check if the LLM supports tool calls and set the global flag + let tools_supported = crate::llm::supports_tool_calls(&settings.model, &settings.provider, &settings.baseurl).await?; + crate::llm::TOOLS_SUPPORTED.set(tools_supported).ok(); + log::log( + "info", + &format!("Model supports tools: {}", tools_supported), + ); + // --------------------------------------------------- // Thread: UI // --------------------------------------------------- @@ -179,6 +319,8 @@ fn main() -> Result<(), Box> { "info", &format!("Playback stream SR (truth): {}", out_sample_rate), ); + // Set global UI sender for memory notifications + crate::memory::TX_UI.set(tx_ui.clone()).unwrap(); log::log("info", &format!("TTS: {}", settings.tts)); log::log("info", &format!("Language: {}", settings.language)); @@ -186,7 +328,7 @@ fn main() -> Result<(), Box> { log::log("info", &format!("LLM provider: {}", settings.provider)); if settings.tts == "kokoro" { - tts::start_kokoro_engine()?; + tts::start_kokoro_engine().await?; } if settings.provider == "ollama" { diff --git a/src/memory.rs b/src/memory.rs new file mode 100644 index 00000000..3c481f56 --- /dev/null +++ b/src/memory.rs @@ -0,0 +1,463 @@ +// ------------------------------------------------------------------ +// Memory +// ------------------------------------------------------------------ + +use anndists::dist::DistL2; // L2 distance implementation +// bincode removed +use crossbeam_channel::Sender; +use hnsw_rs::prelude::*; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use serde_json::json; +use sled; +use std::collections::{HashMap, HashSet}; +use std::hash::{Hash, Hasher}; +// File and BufReader/BufWriter removed +use crate::util; +use std::sync::OnceLock; +use std::thread::{self, JoinHandle}; +use std::time::Duration; +use std::time::{SystemTime, UNIX_EPOCH}; +pub static TX_UI: OnceLock> = OnceLock::new(); + +// API +// ------------------------------------------------------------------ + +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct Predicate { + pub name: String, + pub inverse: String, +} + +impl Predicate { + pub fn to_string(&self) -> String { + self.name.clone() + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct KnowledgeUnit { + pub subject: String, + pub predicate: Predicate, + pub object: String, + pub location: Option, + pub timestamp: SystemTime, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct VecKnowledgeUnit { + pub embedding: Vec, + pub knowledge: KnowledgeUnit, +} + +pub struct Memory { + hnsw: Hnsw<'static, f32, DistL2>, + pub index_map: HashMap, + next_id: usize, + // Inverted indexes for disk‑side search + subject_index: HashMap>, + predicate_index: HashMap>, + object_index: HashMap>, +} + +impl Memory { + pub fn new(expected_elements: usize) -> Self { + // HNSW parameters + let max_nb_connection = 16; + let max_layer = std::cmp::max(1, 16.min((expected_elements as f32).ln().trunc() as usize)); + let ef_construction = 200; + + let hnsw: Hnsw<'static, f32, DistL2> = Hnsw::new( + max_nb_connection, + expected_elements, + max_layer, + ef_construction, + DistL2 {}, + ); + + Memory { + hnsw, + index_map: HashMap::new(), + next_id: 0, + subject_index: HashMap::new(), + predicate_index: HashMap::new(), + object_index: HashMap::new(), + } + } + + fn embed_text(text: &str) -> Vec { + let mut vec = vec![0.0f32; 128]; + for word in text.split_whitespace() { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + word.hash(&mut hasher); + let idx = (hasher.finish() as usize) % 128; + vec[idx] += 1.0; + } + vec + } + + // returns context as sentences (for llm usage) + pub fn build_context_from_units(units: &[KnowledgeUnit]) -> String { + let mut context_phrases = Vec::new(); + for unit in units { + let location = unit.location.clone().unwrap_or("unknown location".into()); + + // Convert SystemTime -> seconds since UNIX_EPOCH + let duration_since_epoch = unit + .timestamp + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + + // Optional: simple date formatting (UTC) + let secs = duration_since_epoch; + let days = secs / 86400; + let hours = (secs % 86400) / 3600; + let minutes = (secs % 3600) / 60; + + let time_str = format!("{}d {}h {}m since epoch", days, hours, minutes); + + // Build sentence + let phrase = format!( + "Subject: '\x1b[32m{}\x1b[0m' Predicate: '\x1b[34m{}\x1b[0m' Object: '\x1b[35m{}\x1b[0m' in '\x1b[36m{}\x1b[0m' at '\x1b[33m{}\x1b[0m'.", + unit.subject, unit.predicate.name, unit.object, location, time_str + ); + context_phrases.push(phrase); + } + context_phrases.join("\n") + } + + pub fn store(&mut self, unit: KnowledgeUnit) { + let text = format!( + "{} {} {}", + unit.subject, + unit.predicate.to_string(), + unit.object + ); + + let embedding = Memory::embed_text(&text); + let id = self.next_id; + + let unit_clone = unit.clone(); + self.index_map.insert( + id, + VecKnowledgeUnit { + embedding: embedding.clone(), + knowledge: unit_clone, + }, + ); + + self.hnsw.insert((&embedding, id)); + self.next_id += 1; + // Update inverted indexes + self + .subject_index + .entry(unit.subject.clone()) + .or_default() + .push(id); + self + .predicate_index + .entry(unit.predicate.name.clone()) + .or_default() + .push(id); + self + .object_index + .entry(unit.object.clone()) + .or_default() + .push(id); + // Send UI notification + if let Some(sender) = TX_UI.get() { + // Send UI notification using a clone to avoid moving the original unit + let _ = sender.send(format!( + "line|🧠 Memory saved: {} {} {}", + unit.subject.clone(), + unit.predicate.name.clone(), + unit.object.clone() + )); + } + } + + fn filter_units<'a>( + &'a self, + units: impl Iterator, + location: Option<&str>, + start: Option, + end: Option, + ) -> Vec { + units + .filter(|v| { + let k = &v.knowledge; + let loc_ok = match (location, &k.location) { + (Some(l), Some(kl)) => kl == l, + (Some(_), None) => false, + (None, _) => true, + }; + let ts_ok = match (start, end) { + (Some(s), Some(e)) => k.timestamp >= s && k.timestamp <= e, + (Some(s), None) => k.timestamp >= s, + (None, Some(e)) => k.timestamp <= e, + (None, None) => true, + }; + loc_ok && ts_ok + }) + .map(|v| v.knowledge.clone()) + .collect() + } + + pub fn get_by_subject( + &self, + subject: &str, + location: Option<&str>, + start: Option, + end: Option, + ) -> Vec { + let units = self + .index_map + .values() + .filter(move |v| v.knowledge.subject == subject); + self.filter_units(units, location, start, end) + } + + pub fn get_by_predicate( + &self, + predicate_name: &str, + location: Option<&str>, + start: Option, + end: Option, + ) -> Vec { + let units = self + .index_map + .values() + .filter(move |v| v.knowledge.predicate.name == predicate_name); + self.filter_units(units, location, start, end) + } + + pub fn get_by_object( + &self, + object: &str, + location: Option<&str>, + start: Option, + end: Option, + ) -> Vec { + let units = self + .index_map + .values() + .filter(move |v| v.knowledge.object == object); + self.filter_units(units, location, start, end) + } + + pub fn get_by_location( + &self, + location: &str, + start: Option, + end: Option, + ) -> Vec { + let units = self.index_map.values(); + self.filter_units(units, Some(location), start, end) + } + + pub fn query(&self, query: &str, k: usize, ef_search: usize) -> Vec { + let q_embed = Memory::embed_text(query); + let neighbors = self.hnsw.search(&q_embed, k, ef_search); + + neighbors + .into_iter() + .map(|neigh| self.index_map[&neigh.d_id].knowledge.clone()) + .collect() + } + + /// Disk‑side search using the inverted indexes + pub fn query_disk(&self, field: &str, value: &str) -> Vec { + let id_list = match field { + "subject" => self.subject_index.get(value), + "predicate" => self.predicate_index.get(value), + "object" => self.object_index.get(value), + _ => None, + }; + + match id_list { + Some(ids) => ids + .iter() + .map(|&id| self.index_map[&id].knowledge.clone()) + .collect(), + None => vec![], + } + } + + pub fn to_json_graph(&self) -> serde_json::Value { + let mut nodes_set = HashSet::new(); + let mut nodes = vec![]; + let mut edges = vec![]; + + for v in self.index_map.values() { + let k = &v.knowledge; + + if nodes_set.insert(k.subject.clone()) { + nodes.push(json!({ "id": k.subject, "label": k.subject })); + } + if nodes_set.insert(k.object.clone()) { + nodes.push(json!({ "id": k.object, "label": k.object })); + } + + let mut edge = json!({ + "from": k.subject, + "to": k.object, + "label": k.predicate.name, + }); + + let timestamp = k.timestamp.duration_since(UNIX_EPOCH).unwrap().as_secs(); + edge["timestamp"] = json!(timestamp); + + if let Some(loc) = &k.location { + edge["location"] = json!(loc); + } + + edges.push(edge); + } + + json!({ "nodes": nodes, "edges": edges }) + } + + /// Save memory to disk (both embeddings & knowledge units) + pub fn save_to_file(&self, path: &str) -> anyhow::Result<()> { + // Store in sled under a single key + let db = sled::open(path)?; + let persist = serde_json::json!({ + "next_id": self.next_id, + "units": self.index_map, + "subject_index": self.subject_index, + "predicate_index": self.predicate_index, + "object_index": self.object_index + }); + let bytes = serde_json::to_vec(&persist)?; + db.insert(b"memory", bytes)?; + db.flush()?; + Ok(()) + } + + /// Load memory from disk, reconstructing HNSW graph + pub fn load_from_file(path: &str) -> anyhow::Result { + let db = sled::open(path)?; + let data = db + .get(b"memory")? + .ok_or_else(|| anyhow::anyhow!("memory data not found"))?; + let persist: serde_json::Value = serde_json::from_slice(&data)?; + + let units_map: HashMap = + serde_json::from_value(persist["units"].clone())?; + let next_id = persist["next_id"].as_u64().unwrap_or(0) as usize; + let expected_elements = units_map.len().max(1); + + let max_nb_connection = 16; + let max_layer = std::cmp::max(1, 16.min((expected_elements as f32).ln().trunc() as usize)); + let ef_construction = 200; + + let hnsw: Hnsw<'static, f32, DistL2> = Hnsw::new( + max_nb_connection, + expected_elements, + max_layer, + ef_construction, + DistL2 {}, + ); + + for (id, v) in &units_map { + hnsw.insert((&v.embedding, *id)); + } + + let mut memory = Memory { + hnsw, + index_map: units_map, + next_id, + subject_index: serde_json::from_value(persist["subject_index"].clone())?, + predicate_index: serde_json::from_value(persist["predicate_index"].clone())?, + object_index: serde_json::from_value(persist["object_index"].clone())?, + }; + + Ok(memory) + } + + pub fn autosave(self, path: String, interval_sec: u64) -> JoinHandle<()> + where + Self: Send + 'static, + { + thread::spawn(move || { + loop { + if let Err(e) = self.save_to_file(&path) { + eprintln!("Failed to autosave memory: {:?}", e); + } + thread::sleep(Duration::from_secs(interval_sec)); + } + }) + } +} + +// PRIVATE +// ------------------------------------------------------------------ + +pub fn ensure_memory_path() -> String { + let home = util::get_user_home_path().expect("Failed to get home dir"); + let dir = home.join(".ai-mate/agents/default"); + std::fs::create_dir_all(&dir).expect("Failed to create memory directory"); + dir.join("memory").to_string_lossy().into_owned() +} + +use std::path::Path; + +pub fn ensure_memory_file(path: &str) -> anyhow::Result { + if Path::new(path).exists() { + Memory::load_from_file(path) + } else { + Ok(Memory::new(1000)) + } +} + +static AVAILABLE_PREDICATES: &[(&str, &str)] = &[ + // predicate // inverse + ("believed", "was believed by"), + ("assumed", "was assumed by"), + ("made", "was made by"), + ("saw", "was seen by"), + ("said", "was said by"), + ("said to", "was told by"), + ("failed at", "was a failure of"), + ("wanted", "was wanted by"), + ("thought", "was thought of by"), + ("asked about", "was asked about by"), + ("planned", "was planned"), + ("requested", "was requested by"), + ("ordered", "was ordered by"), + ("complained about", "received a complaint from"), + ("ocurred at", "was created by"), + ("created", "was created by"), + ("met with", "was met by"), + ("destroyed", "was destroyed by"), + ("modified", "was modified by"), + ("examined", "was examined by"), + ("inspected", "was inspected by"), + ("evaluated", "was evaluated by"), + ("tested", "was tested by"), + ("analyzed", "was analyzed by"), + ("calculated", "was calculated by"), + ("estimated", "was estimated by"), + ("predicted", "was predicted by"), + ("performed", "was performed by"), + ("executed", "was executed by"), + ("completed", "was completed by"), + ("succeeded", "was succeeded by"), + ("confirmed", "was confirmed by"), + ("approved", "was approved by"), + ("denied", "was denied by"), + ("received", "was received by"), + ("sent", "was sent by"), + ("delivered", "was delivered by"), + ("communicated", "was communicated to"), + ("informed", "was informed by"), + ("informed about", "was informed about by"), + ("questioned", "was questioned by"), + ("inquired", "was inquired by"), + ("participated in", "was participated in by"), + ("attended", "was attended by"), + ("presented", "was presented by"), + ("displayed", "was displayed by"), + ("demonstrated", "was demonstrated by"), +]; diff --git a/src/state.rs b/src/state.rs index 98611691..3b57220a 100644 --- a/src/state.rs +++ b/src/state.rs @@ -33,7 +33,7 @@ pub struct AppState { pub voice: Arc>, pub ui: UiState, pub speed: AtomicU32, - pub conversation_history: crate::conversation::ConversationHistory, + pub conversation_history: std::sync::Arc>>, pub agent_name: Arc>, pub agents: Arc>, pub tts: Arc>, diff --git a/src/tools.rs b/src/tools.rs new file mode 100644 index 00000000..40b943af --- /dev/null +++ b/src/tools.rs @@ -0,0 +1,103 @@ +// ------------------------------------------------------------------ +// Tool handling +// ------------------------------------------------------------------ + +use remember::RememberTool; +use serde_json::{Value, json}; +use store_memory::StoreMemoryTool; + +// API +// ------------------------------------------------------------------ + +pub mod remember; +pub mod store_memory; + +pub trait Tool { + fn name(&self) -> &str; + fn handle( + &self, + tool_call_args: &Value, + ) -> Result>; + fn json_schema() -> Result>; +} + +pub fn get_available_tools() -> Result, Box> { + Ok(vec![StoreMemoryTool::json_schema()?]) +} + +pub fn validate_tool_call( + tool_call_args: &Value, + schema: &Value, +) -> Result<(), Box> { + // Locate parameters section: function.parameters + let schema_params = schema.get("parameters"); + if let Some(schema_params) = schema_params { + // required fields + if let Some(required) = schema_params.get("required").and_then(|v| v.as_array()) { + for key in required { + if let Some(k) = key.as_str() { + if !tool_call_args.get(k).is_some() { + return Err(format!("Missing required field: {}", k).into()); + } + } + } + } + // properties validation + if let Some(schema_props) = schema_params.get("properties").and_then(|v| v.as_object()) { + for (name, schema_prop_def) in schema_props { + if let Some(arg_val) = tool_call_args.get(name) { + if let Some(type_str) = schema_prop_def.get("type").and_then(|v| v.as_str()) { + let type_ok = match type_str { + "string" => arg_val.is_string(), + "number" => arg_val.is_number(), + "integer" => arg_val.is_i64() || arg_val.is_u64() || arg_val.is_f64(), + "boolean" => arg_val.is_boolean(), + "array" => arg_val.is_array(), + "object" => arg_val.is_object(), + _ => true, + }; + if !type_ok { + return Err( + format!( + "Field '{}' has incorrect type. Expected {}, got {}", + name, type_str, arg_val + ) + .into(), + ); + } + } + } + } + } + } + Ok(()) +} + +pub fn handle_tool_call( + call_json: &str, +) -> Result> { + crate::log::log("debug", &format!("handling tool: {}", call_json)); + + let v: Value = serde_json::from_str(call_json)?; + let name = v + .get("name") + .and_then(|x| x.as_str()) + .ok_or("Missing tool name")?; + let args = v.get("arguments").ok_or("Missing arguments")?; + + let schema = match name { + "store_memory" => StoreMemoryTool::json_schema()?, + "remember" => RememberTool::json_schema()?, + _ => return Err(format!("Unknown tool: {}", name).into()), + }; + + crate::log::log("debug", &format!("validating tool args")); + // validate the tool call args against the schema + validate_tool_call(args, &schema)?; + + match name { + "store_memory" => StoreMemoryTool::new().handle(args), + "remember" => RememberTool::new().handle(args), + _ => unreachable!(), + } +} diff --git a/src/tools/remember.rs b/src/tools/remember.rs new file mode 100644 index 00000000..c530b6e4 --- /dev/null +++ b/src/tools/remember.rs @@ -0,0 +1,78 @@ +// ------------------------------------------------------------------ +// Tool: Remember +// ------------------------------------------------------------------ + +use super::Tool; +use crate::memory::Memory; +use crate::util; +use serde_json::Value; +use serde_json::json; +// use std::path::Path; + +// API +// ------------------------------------------------------------------ + +pub struct RememberTool; + +// PRIVATE +// ------------------------------------------------------------------ + +impl RememberTool { + pub fn new() -> Self { + RememberTool + } +} + +impl Tool for RememberTool { + fn name(&self) -> &str { + "remember" + } + + fn handle( + &self, + tool_call_args: &Value, + ) -> Result> { + // Validation handled in tools.rs + // Extract query string + let query = tool_call_args + .get("query") + .and_then(|v| v.as_str()) + .ok_or("Missing or invalid 'query'")?; + + // Load or create memory + let memory_path = crate::memory::ensure_memory_path(); + let path = memory_path.as_str(); + let memory = crate::memory::ensure_memory_file(path)?; + + // Perform query + let top_k = 5; + let ef_search = 50; + let retrieved_units = memory.query(query, top_k, ef_search); + + // Build context string + let context_text = Memory::build_context_from_units(&retrieved_units); + + Ok(context_text) + } + + fn json_schema() -> Result> { + Ok(json!({ + "type": "function", + "function": { + "name": "remember", + "description": "Retrieves relevant memories based on a query", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string" + } + }, + "required": [ + "query" + ] + } + } + })) + } +} diff --git a/src/tools/store_memory.rs b/src/tools/store_memory.rs new file mode 100644 index 00000000..dac12cf4 --- /dev/null +++ b/src/tools/store_memory.rs @@ -0,0 +1,177 @@ +// ------------------------------------------------------------------ +// Tool: Store memory +// ------------------------------------------------------------------ + +use super::Tool; +use crate::memory::{KnowledgeUnit, Memory, Predicate}; +use crate::util; +use serde_json::Value; +use serde_json::json; +// use std::path::Path; +use std::time::{SystemTime, UNIX_EPOCH}; + +// API +// ------------------------------------------------------------------ + +pub struct StoreMemoryTool; + +// PRIVATE +// ------------------------------------------------------------------ + +impl StoreMemoryTool { + pub fn new() -> Self { + StoreMemoryTool + } +} + +impl Tool for StoreMemoryTool { + fn name(&self) -> &str { + "store_memory" + } + + fn handle( + &self, + tool_call_args: &Value, + ) -> Result> { + // Extract required fields + let subject = tool_call_args + .get("subject") + .and_then(|v| v.as_str()) + .ok_or("Missing or invalid 'subject'")?; + let predicate = tool_call_args + .get("predicate") + .and_then(|v| v.as_str()) + .ok_or("Missing or invalid 'predicate'")?; + let object = tool_call_args + .get("object") + .and_then(|v| v.as_str()) + .ok_or("Missing or invalid 'object'")?; + + // Optional fields + let location = tool_call_args.get("location").and_then(|v| v.as_str()); + let timestamp = match tool_call_args.get("timestamp") { + Some(v) => { + let secs = v.as_i64().ok_or("Invalid 'timestamp'")?; + Some(SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(secs as u64)) + } + None => None, + }; + + let ts = timestamp.unwrap_or_else(|| SystemTime::now()); + + // Build knowledge unit + let unit = KnowledgeUnit { + subject: subject.to_string(), + predicate: Predicate { + name: predicate.to_string(), + inverse: "".to_string(), + }, + object: object.to_string(), + location: location.map(|s| s.to_string()), + timestamp: ts, + }; + + // Load or create memory + let memory_path = crate::memory::ensure_memory_path(); + let path = memory_path.as_str(); + let mut memory = crate::memory::ensure_memory_file(path)?; + + // Store unit and persist + memory.store(unit); + memory.save_to_file(path)?; + + Ok(format!( + "Stored memory: {} {} {}", + subject, predicate, object + )) + } + + fn json_schema() -> Result> { + Ok(json!({ + "type": "function", + "function": { + "name": "store_memory", + "description": "Use this tool to store relevant memories from the latest user message that can be useful in the future. Use accurate Subject and Object and pick the right Predicate from the allowed values.", + "parameters": { + "type": "object", + "properties": { + "subject": { + "type": "string", + "description": "the person or thing performing the action." + }, + "predicate": { + "type": "string", + "description": "the verb", + "enum": [ + "believed", + "assumed", + "made", + "saw", + "said", + "said to", + "failed at", + "wanted", + "thought", + "asked about", + "planned", + "requested", + "ordered", + "complained about", + "ocurred at", + "created", + "met with", + "destroyed", + "modified", + "examined", + "inspected", + "evaluated", + "tested", + "analyzed", + "calculated", + "estimated", + "predicted", + "performed", + "executed", + "completed", + "succeeded", + "confirmed", + "approved", + "denied", + "received", + "sent", + "delivered", + "communicated", + "informed", + "informed about", + "questioned", + "inquired", + "participated in", + "attended", + "presented", + "displayed", + "demonstrated" + ] + }, + "object": { + "type": "string", + "description": "the person or thing that receives the action of the verb" + }, + "location": { + "type": "string", + "description": "optional location if its specified or inferred from the message" + }, + "timestamp": { + "type": "integer", + "description": "optional timestamp if its specified or inferred from the message" + } + }, + "required": [ + "subject", + "predicate", + "object" + ] + } + } + })) + } +} diff --git a/src/tts.rs b/src/tts.rs index e70268c1..18b2cc39 100644 --- a/src/tts.rs +++ b/src/tts.rs @@ -422,11 +422,8 @@ pub fn speak_via_kokoro_stream( } } -pub fn start_kokoro_engine() -> Result<(), Box> { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build()?; - let engine = rt.block_on(TtsEngine::new())?; +pub async fn start_kokoro_engine() -> Result<(), Box> { + let engine: TtsEngine = TtsEngine::new().await?; KOKORO_ENGINE.set(Arc::new(Mutex::new(engine))).ok(); Ok(()) } diff --git a/src/ui.rs b/src/ui.rs index c5987c56..265de9a4 100644 --- a/src/ui.rs +++ b/src/ui.rs @@ -19,7 +19,6 @@ use std::thread; use std::time::Duration; // API -// ------------------------------------------------------------------ pub static STOP_STREAM: AtomicBool = AtomicBool::new(false); @@ -34,6 +33,7 @@ pub fn spawn_ui_thread( status_line: Arc>, rx_ui: Receiver, ) -> thread::JoinHandle<()> { + // separate thread for bottom bar update + render thread::spawn(move || { let mut ui_state = ui_state; let mut out = io::stdout(); @@ -432,7 +432,6 @@ fn render_bottom_bar( "\x1b[37m" }; let bar = format!("{}{}\x1b[0m", bar_color, "█".repeat(bar_len)); - let spaces = cols.saturating_sub( get_visible_len_for(&status) + 2