diff --git a/api/src/host/events.rs b/api/src/host/events.rs index 30942eb..a843db4 100644 --- a/api/src/host/events.rs +++ b/api/src/host/events.rs @@ -22,7 +22,7 @@ //! rather than growing a subscribe method here. /// Why an embedding model was reported unhealthy, and what took over. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct EmbeddingHealthReason { /// The provider that failed. pub provider: String, @@ -46,7 +46,7 @@ pub type SyncTrigger = String; /// Deliberately **not** `#[non_exhaustive]`: the host's mapping impl matches /// exhaustively on purpose, so adding a variant here is a compile error at the /// mapping site rather than an event that silently never reaches the bus. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub enum MemoryEvent { /// A sync run moved to a new stage. SyncStageChanged { diff --git a/core/src/store/client.rs b/core/src/store/client.rs index 569eee1..e93ba9c 100644 --- a/core/src/store/client.rs +++ b/core/src/store/client.rs @@ -93,14 +93,30 @@ impl MemoryClient { /// tool-scoped memory layer) without depending on the concrete /// `MemoryClient` type or holding a reference to it. /// - /// Intentionally `pub(crate)` — handing a raw `Arc` to an - /// external consumer bypasses any policy decorator wrapped around the - /// `MemoryClient` API, so the escape hatch stays in-crate. Mirrors - /// [`Self::profile_conn`]. + /// This is public for the `tinymemory-module` provider, which implements + /// the TinyMemory contract over this exact client. Product hosts must use + /// the guarded provider and must not retain this raw engine handle. pub fn memory_handle(&self) -> Arc { Arc::clone(&self.inner) as Arc } + /// Wrap an already-configured unified store and start its ingestion worker. + /// + /// This is the constructor used by the compiled module: the module first + /// resolves the host-supplied embedding route and storage configuration, + /// then gives the resulting store to the high-level client without opening + /// a second database or creating a second ingestion queue. + #[must_use] + pub fn from_unified_memory(inner: UnifiedMemory) -> Self { + let inner = Arc::new(inner); + let ingestion_queue = + ingestion_queue::start_worker_with_state(Arc::clone(&inner), IngestionState::new()); + Self { + inner, + ingestion_queue, + } + } + /// Create a new local memory client using the default `.openhuman` directory. /// /// # Errors @@ -141,18 +157,7 @@ impl MemoryClient { // Create the underlying UnifiedMemory instance. let memory = UnifiedMemory::new(&workspace_dir, embedder, None).map_err(|e| format!("{e}"))?; - let inner = Arc::new(memory); - - // Start the background worker for document ingestion and graph extraction. - // The worker shares its IngestionState with the synchronous ingest path - // below so all ingestion is singleton-serialised. - let ingestion_queue = - ingestion_queue::start_worker_with_state(Arc::clone(&inner), IngestionState::new()); - - Ok(Self { - inner, - ingestion_queue, - }) + Ok(Self::from_unified_memory(memory)) } /// Store a document in a specific namespace. diff --git a/core/src/store/entities.rs b/core/src/store/entities.rs index a571a15..47024f4 100644 --- a/core/src/store/entities.rs +++ b/core/src/store/entities.rs @@ -11,6 +11,92 @@ use crate::sync::composio::providers::profile::{is_self_identity_any_toolkit, Id use crate::tinycortex::memory_config_from; use crate::Config; +/// Aggregate entity-index row for capability providers. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TopEntity { + /// Canonical entity id. + pub id: String, + /// Stable entity kind string. + pub kind: String, + /// Representative observed surface form. + pub name: String, + /// Number of indexed observations. + pub mentions: u32, +} + +/// Entity row scoped to one memory-tree namespace. +pub fn namespace_entities( + config: &Config, + namespace: &str, + query: Option<&str>, + limit: usize, +) -> Result> { + let memory = memory_config_from(config, config.workspace_dir().clone()); + let connection = tinycortex::memory::chunks::shared_connection(&memory)?; + let guard = connection.lock(); + let pattern = query.map(|value| format!("%{}%", value.to_ascii_lowercase())); + let mut statement = guard.prepare( + "SELECT entity_id, entity_kind, MAX(surface), COUNT(*) + FROM mem_tree_entity_index + WHERE tree_id = ?1 + AND (?2 IS NULL OR LOWER(entity_id) LIKE ?2 OR LOWER(surface) LIKE ?2) + GROUP BY entity_id, entity_kind + ORDER BY COUNT(*) DESC, MAX(timestamp_ms) DESC + LIMIT ?3", + )?; + let rows = statement + .query_map( + rusqlite::params![namespace, pattern, i64::try_from(limit).unwrap_or(i64::MAX)], + |row| { + let mentions: i64 = row.get(3)?; + Ok(TopEntity { + id: row.get(0)?, + kind: row.get(1)?, + name: row.get(2)?, + mentions: u32::try_from(mentions.max(0)).unwrap_or(u32::MAX), + }) + }, + )? + .collect::>>()?; + Ok(rows) +} + +/// Co-occurrence edges scoped to one memory-tree namespace. +pub fn namespace_entity_edges( + config: &Config, + namespace: &str, + entity_id: &str, + limit: usize, +) -> Result> { + let memory = memory_config_from(config, config.workspace_dir().clone()); + let connection = tinycortex::memory::chunks::shared_connection(&memory)?; + let guard = connection.lock(); + let mut statement = guard.prepare( + "SELECT b.entity_id, COUNT(*) + FROM mem_tree_entity_index a + JOIN mem_tree_entity_index b + ON a.node_id = b.node_id AND a.tree_id = b.tree_id + WHERE a.tree_id = ?1 AND a.entity_id = ?2 AND b.entity_id <> a.entity_id + GROUP BY b.entity_id + ORDER BY COUNT(*) DESC + LIMIT ?3", + )?; + let rows = statement + .query_map( + rusqlite::params![ + namespace, + entity_id, + i64::try_from(limit).unwrap_or(i64::MAX) + ], + |row| { + let count: i64 = row.get(1)?; + Ok((row.get(0)?, u32::try_from(count.max(0)).unwrap_or(u32::MAX))) + }, + )? + .collect::>>()?; + Ok(rows) +} + pub use tinycortex::memory::store::entity_index::EntityHit; #[derive(Debug)] @@ -80,10 +166,61 @@ pub fn count_entity_index(config: &Config) -> Result { index(config)?.count_entity_index() } +/// Most frequently observed entities, with recency as the tie-breaker. +pub fn top_entities(config: &Config, limit: usize) -> Result> { + let memory = memory_config_from(config, config.workspace_dir().clone()); + let connection = tinycortex::memory::chunks::shared_connection(&memory)?; + let guard = connection.lock(); + let mut statement = guard.prepare( + "SELECT entity_id, entity_kind, MAX(surface), COUNT(*) + FROM mem_tree_entity_index + GROUP BY entity_id, entity_kind + ORDER BY COUNT(*) DESC, MAX(timestamp_ms) DESC + LIMIT ?1", + )?; + let rows = statement + .query_map([i64::try_from(limit).unwrap_or(i64::MAX)], |row| { + let mentions: i64 = row.get(3)?; + Ok(TopEntity { + id: row.get(0)?, + kind: row.get(1)?, + name: row.get(2)?, + mentions: u32::try_from(mentions.max(0)).unwrap_or(u32::MAX), + }) + })? + .collect::>>()?; + Ok(rows) +} + #[cfg(test)] mod tests { use super::*; + fn scoped_config() -> ( + tempfile::TempDir, + tinymemory_api::host::test_support::TestHostConfig, + ) { + crate::test_seams::init(); + let temp = tempfile::tempdir().expect("tempdir"); + let mut config = tinymemory_api::host::test_support::TestHostConfig::default(); + config.workspace_dir = temp.path().to_path_buf(); + (temp, config) + } + + fn insert_entity(config: &Config, tree: &str, entity: &str, node: &str, surface: &str) { + let memory = memory_config_from(config, config.workspace_dir().clone()); + let connection = tinycortex::memory::chunks::shared_connection(&memory).expect("db"); + connection + .lock() + .execute( + "INSERT INTO mem_tree_entity_index + (entity_id,node_id,node_kind,entity_kind,surface,score,timestamp_ms,tree_id) + VALUES (?1,?2,'chunk','person',?3,1.0,1,?4)", + rusqlite::params![entity, node, surface, tree], + ) + .expect("insert"); + } + #[test] fn crate_entity_hit_is_the_host_facade_type() { let hit = EntityHit { @@ -99,4 +236,29 @@ mod tests { }; assert_eq!(hit.entity_id, "person:alice"); } + + #[test] + fn namespace_entity_reads_do_not_cross_tree_ids() { + let (_temp, config) = scoped_config(); + insert_entity(&config, "team-a", "person:alice", "a1", "Alice"); + insert_entity(&config, "team-b", "person:bob", "b1", "Bob"); + + let rows = namespace_entities(&config, "team-a", None, 10).expect("entities"); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].id, "person:alice"); + assert!(namespace_entities(&config, "team-a", Some("bob"), 10) + .expect("search") + .is_empty()); + } + + #[test] + fn namespace_edges_join_only_rows_in_the_same_tree() { + let (_temp, config) = scoped_config(); + insert_entity(&config, "team-a", "person:alice", "shared", "Alice"); + insert_entity(&config, "team-a", "person:bob", "shared", "Bob"); + insert_entity(&config, "team-b", "person:mallory", "shared", "Mallory"); + + let rows = namespace_entity_edges(&config, "team-a", "person:alice", 10).expect("edges"); + assert_eq!(rows, vec![("person:bob".to_string(), 1)]); + } } diff --git a/core/src/store/factories.rs b/core/src/store/factories.rs index 4a006f8..1938b3a 100644 --- a/core/src/store/factories.rs +++ b/core/src/store/factories.rs @@ -577,6 +577,32 @@ fn create_unified_memory_full( ) } +/// Create the high-level memory client used by the compiled TinyMemory module. +/// +/// Unlike [`crate::store::MemoryClient::from_workspace_dir`], this preserves +/// the caller's resolved embedding and storage configuration. The returned +/// client and its raw [`Memory`] handle share one `UnifiedMemory` instance and +/// one ingestion worker. +pub fn create_memory_client_with_local_ai( + memory: &MemoryConfig, + local_embedding_model: Option<&str>, + embedding_api_key: &str, + embedding_routes: &[EmbeddingRouteConfig], + storage_provider: Option<&StorageProviderConfig>, + workspace_dir: &Path, +) -> anyhow::Result { + let store = create_unified_memory_full( + memory, + embedding_routes, + storage_provider, + local_embedding_model, + embedding_api_key, + workspace_dir, + "memory", + )?; + Ok(crate::store::MemoryClient::from_unified_memory(store)) +} + /// Create a memory instance specifically for migration purposes. /// /// The unified namespace memory core has a single workspace-scoped diff --git a/crates/tinymemory-module/Cargo.lock b/crates/tinymemory-module/Cargo.lock index bfb2cfc..dad224c 100644 --- a/crates/tinymemory-module/Cargo.lock +++ b/crates/tinymemory-module/Cargo.lock @@ -184,6 +184,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d262e149917187838d5b42777c8253bcb64500067342904e7d429499a6f277e" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -576,6 +578,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "git2" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddddbf932745a6be37109b6112d3ee09696106f848449069d3a57bba937ab82e" +dependencies = [ + "bitflags", + "libc", + "libgit2-sys", + "log", +] + [[package]] name = "hashbrown" version = "0.16.1" @@ -600,6 +614,12 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "http" version = "1.5.0" @@ -863,6 +883,16 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + [[package]] name = "js-sys" version = "0.3.104" @@ -880,6 +910,18 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "libgit2-sys" +version = "0.18.7+1.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23c7391e4b9f4ffab1a624223cc1d7385ff9a678f490768add717de7ea2f4d89" +dependencies = [ + "cc", + "libc", + "libz-sys", + "pkg-config", +] + [[package]] name = "libredox" version = "0.1.19" @@ -900,6 +942,18 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "libz-sys" +version = "1.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85bc9657773828b90eeb625adff10eeac83cc21bbfd8e23a03eaa8a33c9e28d9" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -1839,6 +1893,8 @@ dependencies = [ "chrono", "dirs", "futures", + "git2", + "hex", "log", "parking_lot", "rand 0.10.2", @@ -1911,6 +1967,7 @@ dependencies = [ "chrono", "dirs", "futures", + "git2", "log", "parking_lot", "rand 0.8.7", @@ -1939,10 +1996,12 @@ version = "0.1.0" dependencies = [ "anyhow", "async-trait", + "chrono", "log", "serde", "serde_json", "tempfile", + "tinyagents", "tinybus", "tinybus-module", "tinycortex", diff --git a/crates/tinymemory-module/Cargo.toml b/crates/tinymemory-module/Cargo.toml index bbda847..6c55c4c 100644 --- a/crates/tinymemory-module/Cargo.toml +++ b/crates/tinymemory-module/Cargo.toml @@ -31,9 +31,10 @@ tinymemory = { path = "../.." } # The engine and the seam that adapts it. Carrying these is the entire point of # the module: they are 14.7s of the host's critical build path, and a host that # loads this binary compiles neither. -tinymemory-core = { path = "../../core" } +tinymemory-core = { path = "../../core", features = ["memory-git"] } tinymemory-tinycortex = { path = "../../adapters/tinycortex" } tinycortex = { version = "0.1" } +tinyagents = { version = "2.1" } # TinyBus provides the typed service interface and the dynamic module host ABI. # Reached by path now that this crate is its own workspace root: the nested # checkout's `[workspace.package]` resolves correctly from here. @@ -47,6 +48,7 @@ tinybus-module = { version = "0.1.0", path = "../../vendor/tinybus/crates/tinybu async-trait = "0.1" # `EmbeddingProvider::embed` is anyhow-typed. anyhow = "1" +chrono = "0.4" # Diagnostics. Never carries a namespace key or entry content — see `service`. log = "0.4" # Module configuration is JSON supplied by the host at load time. diff --git a/crates/tinymemory-module/src/chat.rs b/crates/tinymemory-module/src/chat.rs new file mode 100644 index 0000000..3eb894f --- /dev/null +++ b/crates/tinymemory-module/src/chat.rs @@ -0,0 +1,114 @@ +//! Chat-model calls stay host-side and cross `TinyBus` as typed requests. + +use std::sync::Arc; + +use async_trait::async_trait; +use tinyagents::harness::model::{ChatModel, ModelRequest, ModelResponse}; +use tinybus::Connection; + +use crate::ModuleConfig; + +/// Well-known host service used for memory summarisation and extraction. +pub const CHAT_HOST_BUS_NAME: &str = "ai.tinyhumans.tinymemory.ChatHost"; +/// Host object path for [`CHAT_HOST_BUS_NAME`]. +pub const CHAT_HOST_OBJECT_PATH: &str = "/ai/tinyhumans/tinymemory/ChatHost"; +/// Interface name exported by the host. +pub const CHAT_HOST_INTERFACE: &str = "ai.tinyhumans.tinymemory.ChatHost"; + +/// Builds bus-backed chat models without receiving a provider credential. +pub struct BusChatHost { + connection: Connection, + provider: String, + model_id: String, +} + +impl std::fmt::Debug for BusChatHost { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("BusChatHost") + .field("provider", &self.provider) + .field("model_id", &self.model_id) + .finish_non_exhaustive() + } +} + +impl BusChatHost { + /// Build the host bridge over the module connection. + #[must_use] + pub fn new(connection: Connection, config: &ModuleConfig) -> Self { + Self { + connection, + provider: config + .memory_provider + .clone() + .unwrap_or_else(|| "host".to_string()), + model_id: config + .default_model + .clone() + .unwrap_or_else(|| "host-default".to_string()), + } + } +} + +impl tinymemory_core::chat_host::ChatHost for BusChatHost { + fn provider_for_role(&self, _role: &str, _config: &tinymemory_core::Config) -> String { + self.provider.clone() + } + + fn create_chat_model_with_model_id( + &self, + role: &str, + _config: &tinymemory_core::Config, + _temperature: f64, + ) -> Result<(Arc>, String), String> { + Ok(( + Arc::new(BusChatModel { + connection: self.connection.clone(), + role: role.to_string(), + }), + self.model_id.clone(), + )) + } + + fn usage_from_response( + &self, + _response: &ModelResponse, + ) -> Option { + None + } + + fn summarizer_available(&self, _config: &tinymemory_core::Config) -> (bool, &'static str) { + (true, "served by the TinyMemory host callback") + } +} + +struct BusChatModel { + connection: Connection, + role: String, +} + +#[async_trait] +impl ChatModel<()> for BusChatModel { + fn cache_identity(&self) -> Option { + Some(format!("tinymemory-module-host:{}", self.role)) + } + + async fn invoke( + &self, + _state: &(), + request: ModelRequest, + ) -> tinyagents::Result { + let proxy = self + .connection + .proxy( + CHAT_HOST_BUS_NAME, + CHAT_HOST_OBJECT_PATH, + CHAT_HOST_INTERFACE, + ) + .map_err(|error| tinyagents::TinyAgentsError::Model(error.to_string()))?; + proxy + .call("Complete", (self.role.clone(), request)) + .await + .map_err(|error| tinyagents::TinyAgentsError::Model(error.to_string())) + } +} diff --git a/crates/tinymemory-module/src/config.rs b/crates/tinymemory-module/src/config.rs index 65193b0..46b7066 100644 --- a/crates/tinymemory-module/src/config.rs +++ b/crates/tinymemory-module/src/config.rs @@ -35,7 +35,10 @@ use std::path::PathBuf; use serde::{Deserialize, Serialize}; -use tinymemory_api::host::{EmbeddingRouteConfig, MemoryConfig, StorageProviderConfig}; +use tinymemory_api::host::{ + EmbeddingRouteConfig, LocalAiConfig, MemoryConfig, MemoryTreeConfig, SchedulerGateConfig, + StorageProviderConfig, +}; /// Everything this module needs to bring up a memory engine. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -53,6 +56,35 @@ pub struct ModuleConfig { /// The engine's own configuration, passed through unchanged. pub memory: MemoryConfig, + /// Summary-tree and language-model settings used by tree operations. + pub memory_tree: MemoryTreeConfig, + + /// Background-work admission policy used by bounded maintenance steps. + pub scheduler_gate: SchedulerGateConfig, + + /// Local model selection. Credentials remain host-side; this contains only + /// routing and model identifiers. + pub local_ai: LocalAiConfig, + + /// Resolved route for the embedding workload. + pub embeddings_provider: Option, + + /// Resolved route for memory summarisation/extraction. + pub memory_provider: Option, + + /// Default chat model identifier used for module-side summarisation. + pub default_model: Option, + + /// Sampling temperature for module-side background language tasks. + pub default_temperature: f64, + + /// Optional output language for summaries and extracted artifacts. + pub output_language: Option, + + /// Serialized memory-source registry used by snapshot operations. + #[serde(default = "empty_array")] + pub memory_sources: serde_json::Value, + /// Per-workload embedding routes, as the host resolved them. pub embedding_routes: Vec, @@ -97,6 +129,15 @@ impl Default for ModuleConfig { Self { workspace_dir: PathBuf::new(), memory: MemoryConfig::default(), + memory_tree: MemoryTreeConfig::default(), + scheduler_gate: SchedulerGateConfig::default(), + local_ai: LocalAiConfig::default(), + embeddings_provider: None, + memory_provider: None, + default_model: None, + default_temperature: 0.0, + output_language: None, + memory_sources: empty_array(), embedding_routes: Vec::new(), storage_provider: None, ollama_base_url: String::new(), @@ -108,6 +149,10 @@ impl Default for ModuleConfig { } } +fn empty_array() -> serde_json::Value { + serde_json::Value::Array(Vec::new()) +} + impl ModuleConfig { /// Reject a configuration that cannot bring up a store. /// diff --git a/crates/tinymemory-module/src/host.rs b/crates/tinymemory-module/src/host.rs new file mode 100644 index 0000000..8147ce7 --- /dev/null +++ b/crates/tinymemory-module/src/host.rs @@ -0,0 +1,125 @@ +//! Host-owned runtime services used by the compiled memory engine. + +use std::sync::Arc; + +use async_trait::async_trait; +use tinybus::Connection; +use tinymemory_api::host::{ErrorReporter, MemoryEvent, MemoryEventSink, SpacyResponse}; + +/// Host callback routing constants. +pub const RUNTIME_HOST_BUS_NAME: &str = "ai.tinyhumans.tinymemory.RuntimeHost"; +/// Object path for the host callbacks. +pub const RUNTIME_HOST_OBJECT_PATH: &str = "/ai/tinyhumans/tinymemory/RuntimeHost"; +/// Interface exported by the host. +pub const RUNTIME_HOST_INTERFACE: &str = "ai.tinyhumans.tinymemory.RuntimeHost"; + +#[derive(Clone)] +pub(crate) struct BusRuntimeHost { + connection: Connection, +} + +impl std::fmt::Debug for BusRuntimeHost { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("BusRuntimeHost") + .finish_non_exhaustive() + } +} + +impl BusRuntimeHost { + pub(crate) fn new(connection: Connection) -> Self { + Self { connection } + } + + fn proxy(&self) -> Result { + self.connection.proxy( + RUNTIME_HOST_BUS_NAME, + RUNTIME_HOST_OBJECT_PATH, + RUNTIME_HOST_INTERFACE, + ) + } + + fn notify(&self, method: &'static str, arguments: T) + where + T: serde::Serialize + Send + 'static, + { + let host = self.clone(); + tokio::spawn(async move { + let result = match host.proxy() { + Ok(proxy) => proxy.call::<()>(method, arguments).await, + Err(error) => Err(error), + }; + if let Err(error) = result { + log::debug!("[tinymemory:module] host callback {method} failed: {error}"); + } + }); + } +} + +impl MemoryEventSink for BusRuntimeHost { + fn publish(&self, event: MemoryEvent) { + self.notify("PublishEvent", (event,)); + } +} + +impl ErrorReporter for BusRuntimeHost { + fn report_error(&self, rendered: &str, domain: &str, operation: &str, tags: &[(&str, &str)]) { + self.notify( + "ReportError", + ( + false, + rendered.to_string(), + domain.to_string(), + operation.to_string(), + owned_tags(tags), + ), + ); + } + + fn report_error_or_expected( + &self, + rendered: &str, + domain: &str, + operation: &str, + tags: &[(&str, &str)], + ) { + self.notify( + "ReportError", + ( + true, + rendered.to_string(), + domain.to_string(), + operation.to_string(), + owned_tags(tags), + ), + ); + } +} + +fn owned_tags(tags: &[(&str, &str)]) -> Vec<(String, String)> { + tags.iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect() +} + +#[async_trait] +impl tinymemory_core::nlp_host::NlpHost for BusRuntimeHost { + async fn extract_spacy( + &self, + _config: &tinymemory_core::Config, + text: &str, + ) -> Result { + let proxy = self.proxy().map_err(|error| error.to_string())?; + proxy + .call("ExtractSpacy", (text.to_string(),)) + .await + .map_err(|error| error.to_string()) + } +} + +pub(crate) fn install(connection: Connection) { + let host = Arc::new(BusRuntimeHost::new(connection)); + tinymemory_core::events::set_event_sink(Arc::clone(&host) as Arc); + tinymemory_core::observability::set_error_reporter(Arc::clone(&host) as Arc); + tinymemory_core::nlp_host::set_nlp_host(host); +} diff --git a/crates/tinymemory-module/src/lib.rs b/crates/tinymemory-module/src/lib.rs index 64f885a..9a9e2cd 100644 --- a/crates/tinymemory-module/src/lib.rs +++ b/crates/tinymemory-module/src/lib.rs @@ -43,13 +43,12 @@ //! [`config::ModuleConfig::strip_host_credentials`], not merely asserted about a //! field list. "Carried verbatim" carries credentials verbatim too. //! -//! # Scope: the mandatory three +//! # Scope: the complete TinyMemory API //! -//! The served surface is `tinymemory_api`'s mandatory capability families — -//! Core, Recall, Portability — which is exactly what `tinymemory-tinycortex` -//! can provide. The ten optional families need a host's configuration, embedding -//! compute and job queue, and a host that has those implements them itself. -//! See [`service`] for the method list. +//! The module boundary mirrors every capability family in `tinymemory_api`. +//! Host applications keep policy, scheduling, credentials, and bus/event types; +//! memory storage, retrieval, ingestion, trees, graph operations, goals, source +//! persistence, and maintenance execute inside this compiled module. // Test code may panic; library code may not. The `[lints]` table cannot be // scoped to non-test builds, so the exemption is expressed here instead. @@ -63,15 +62,20 @@ ) )] +pub mod chat; pub mod config; pub mod embedding; +mod host; +mod provider; mod service; +pub use chat::{CHAT_HOST_BUS_NAME, CHAT_HOST_INTERFACE, CHAT_HOST_OBJECT_PATH}; pub use config::ModuleConfig; pub use embedding::{ BusEmbeddingHost, BusEmbeddingProvider, EMBEDDING_HOST_BUS_NAME, EMBEDDING_HOST_INTERFACE, EMBEDDING_HOST_OBJECT_PATH, }; +pub use host::{RUNTIME_HOST_BUS_NAME, RUNTIME_HOST_INTERFACE, RUNTIME_HOST_OBJECT_PATH}; pub use service::{BUS_NAME, OBJECT_PATH}; use std::sync::atomic::{AtomicBool, Ordering}; @@ -100,8 +104,8 @@ const SETUP_FAILED_ERROR: &str = "ai.tinyhumans.tinymemory.Error.SetupFailed"; /// the host, which holds the real credential, so there is nothing to pass and /// nothing here that could leak one. async fn setup(connection: Connection, mut config: ModuleConfig) -> BusResult<()> { - claim_process_setup()?; config.validate().map_err(setup_error)?; + claim_process_setup()?; // `MemoryConfig` travels verbatim, and it contains a bearer token field for a // remote memory backend. Carried credentials are exactly what this module @@ -126,8 +130,13 @@ async fn setup(connection: Connection, mut config: ModuleConfig) -> BusResult<() connection.clone(), &config, ))); + tinymemory_core::chat_host::set_chat_host(Arc::new(chat::BusChatHost::new( + connection.clone(), + &config, + ))); + host::install(connection.clone()); - let memory = tinymemory_core::store::factories::create_memory_with_local_ai( + let client = tinymemory_core::store::factories::create_memory_client_with_local_ai( &config.memory, None, "", @@ -144,14 +153,13 @@ async fn setup(connection: Connection, mut config: ModuleConfig) -> BusResult<() setup_error("create memory store") })?; - let provider = tinymemory_tinycortex::provider(Arc::from(memory)); + let provider = provider::ModuleMemoryProvider::new(&config, Arc::new(client)); service::serve(&connection, Arc::new(provider)).await } /// Claim this process's single setup slot. /// -/// `setup` installs a **process-global** embedding host -/// (`tinymemory_core::embedding_host::set_embedding_host`), so it is not +/// `setup` installs **process-global** host callbacks, so it is not /// re-entrant the way a per-host resource would be. `ModuleHost` rejects a /// duplicate module name only within one host, and nothing stops a process from /// building a second host — a test harness is the obvious way it happens. The @@ -172,7 +180,7 @@ fn claim_process_setup() -> BusResult<()> { if CLAIMED.swap(true, Ordering::SeqCst) { return Err(setup_error( "this module is already set up in this process; it installs a \ - process-global embedding host and cannot be served twice", + process-global host callbacks and cannot be served twice", )); } Ok(()) @@ -216,6 +224,38 @@ mod exports { "Recall", "ExportPage", "ImportRecords", + "IngestDocument", + "IngestChat", + "PutDocument", + "GetDocument", + "QueryDocuments", + "Append", + "QuerySource", + "DrillDown", + "Seal", + "Cascade", + "Entities", + "EntityEdges", + "TouchEntities", + "KvGet", + "KvPut", + "KvList", + "Relations", + "PutRelation", + "CaptureSnapshot", + "Snapshots", + "Diff", + "Goals", + "SetGoals", + "ToolRules", + "PutToolRule", + "DeleteToolRule", + "AcceptSourceItems", + "ForgetSource", + "Reembed", + "Compact", + "Consolidate", + "Doctor", ], signals = [], // The host's embedder is deliberately NOT declared as `requires`. That diff --git a/crates/tinymemory-module/src/provider.rs b/crates/tinymemory-module/src/provider.rs new file mode 100644 index 0000000..357ee61 --- /dev/null +++ b/crates/tinymemory-module/src/provider.rs @@ -0,0 +1,1141 @@ +//! Complete TinyMemory provider backed by the module-owned engine. + +use std::collections::HashSet; +use std::path::PathBuf; +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::Utc; +use tinymemory::mandatory::MemoryTraitProvider; +use tinymemory_api::capabilities::Capabilities; +use tinymemory_api::chunks::Chunk; +use tinymemory_api::error::MemoryError; +use tinymemory_api::goals::GoalsDoc; +use tinymemory_api::health::MemoryHealth; +use tinymemory_api::host::{ + CloudProviderCreds, ComposioMode, LocalAiConfig, MemoryConfig, MemoryHostConfig, + MemoryTreeConfig, SchedulerGateConfig, +}; +use tinymemory_api::provider::types::{ + ChangeKind, DiffReport, EntityHit, EntityRef, ExportPage, ExportRecord, ImportOutcome, + IngestItem, IngestOutcome, MaintenanceReport, SnapshotRef, SourceChange, SourceItem, + SourceScope, +}; +use tinymemory_api::provider::{ + MemoryCore, MemoryDiff, MemoryDocuments, MemoryEntities, MemoryGoals, MemoryGraph, + MemoryIngest, MemoryMaintenance, MemoryPortability, MemoryProvider, MemoryRecall, + MemorySourceSink, MemoryToolMemory, MemoryTree, +}; +use tinymemory_api::recall::OwnedRecallOpts; +use tinymemory_api::tool_memory::ToolMemoryRule; +use tinymemory_api::tree::{IngestRequest, QueryResult, TreeStatus}; +use tinymemory_api::types::{ + GraphRelationRecord, MemoryCategory, MemoryEntry, MemoryKvRecord, MemoryTaint, + NamespaceDocumentInput, NamespaceRetrievalContext, NamespaceSummary, StoredMemoryDocument, +}; +use tinymemory_core::store::{MemoryClient, MemoryClientRef}; +use tinymemory_tinycortex::TinycortexMemory; + +use crate::ModuleConfig; + +/// The concrete, credential-free host configuration available inside a module. +#[derive(Debug, Clone)] +struct ModuleRuntimeConfig { + workspace_dir: PathBuf, + config_path: PathBuf, + memory: MemoryConfig, + memory_tree: MemoryTreeConfig, + scheduler_gate: SchedulerGateConfig, + local_ai: LocalAiConfig, + embeddings_provider: Option, + memory_provider: Option, + default_model: Option, + default_temperature: f64, + output_language: Option, + memory_sources: serde_json::Value, +} + +impl From<&ModuleConfig> for ModuleRuntimeConfig { + fn from(config: &ModuleConfig) -> Self { + Self { + workspace_dir: config.workspace_dir.clone(), + config_path: config.workspace_dir.join("config.toml"), + memory: config.memory.clone(), + memory_tree: config.memory_tree.clone(), + scheduler_gate: config.scheduler_gate.clone(), + local_ai: config.local_ai.clone(), + embeddings_provider: config.embeddings_provider.clone(), + memory_provider: config.memory_provider.clone(), + default_model: config.default_model.clone(), + default_temperature: config.default_temperature, + output_language: config.output_language.clone(), + memory_sources: config.memory_sources.clone(), + } + } +} + +#[async_trait] +impl MemoryHostConfig for ModuleRuntimeConfig { + fn workspace_dir(&self) -> &PathBuf { + &self.workspace_dir + } + fn config_path(&self) -> &PathBuf { + &self.config_path + } + fn memory_tree_content_root(&self) -> PathBuf { + self.memory_tree + .content_dir + .clone() + .unwrap_or_else(|| self.workspace_dir.join("memory_tree/content")) + } + fn memory(&self) -> &MemoryConfig { + &self.memory + } + fn memory_tree(&self) -> &MemoryTreeConfig { + &self.memory_tree + } + fn scheduler_gate(&self) -> &SchedulerGateConfig { + &self.scheduler_gate + } + fn local_ai(&self) -> &LocalAiConfig { + &self.local_ai + } + fn cloud_providers(&self) -> &Vec { + static NONE: Vec = Vec::new(); + &NONE + } + fn embeddings_provider(&self) -> Option<&str> { + self.embeddings_provider.as_deref() + } + fn memory_provider(&self) -> Option<&str> { + self.memory_provider.as_deref() + } + fn workload_local_model(&self, workload: &str) -> Option { + let route = match workload { + "memory" => self.memory_provider.as_deref(), + "embeddings" => self.embeddings_provider.as_deref(), + _ => None, + }?; + route + .strip_prefix("ollama:") + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn to_arc(&self) -> Arc { + Arc::new(self.clone()) + } + fn api_url(&self) -> Option<&str> { + None + } + fn effective_backend_api_url(&self) -> String { + String::new() + } + fn session_token(&self) -> Result, String> { + Ok(None) + } + fn default_model(&self) -> Option<&str> { + self.default_model.as_deref() + } + fn default_temperature(&self) -> f64 { + self.default_temperature + } + fn output_language(&self) -> Option<&str> { + self.output_language.as_deref() + } + fn memory_sync_interval_secs(&self) -> Option { + Some(0) + } + fn onboarding_completed(&self) -> bool { + true + } + fn secrets_encrypt(&self) -> bool { + false + } + fn composio(&self) -> ComposioMode { + ComposioMode::default() + } + fn memory_sources_json(&self) -> anyhow::Result { + Ok(self.memory_sources.clone()) + } + fn set_memory_sources_json(&mut self, value: serde_json::Value) -> anyhow::Result<()> { + self.memory_sources = value; + Ok(()) + } + fn composio_source_caps_migration_version(&self) -> u32 { + 0 + } + fn set_composio_source_caps_migration_version(&mut self, _version: u32) {} + fn apply_env_overrides(&mut self) {} + async fn save(&self) -> anyhow::Result<()> { + Ok(()) + } +} + +/// The module-owned implementation of every TinyMemory capability family. +pub(crate) struct ModuleMemoryProvider { + driver_id: String, + mandatory: MemoryTraitProvider, + client: MemoryClientRef, + config: ModuleRuntimeConfig, +} + +impl ModuleMemoryProvider { + pub(crate) fn new(config: &ModuleConfig, client: Arc) -> Self { + let memory = client.memory_handle(); + let mandatory = MemoryTraitProvider::new( + Arc::new(TinycortexMemory::new(memory)), + config.driver_id.clone(), + ); + Self { + driver_id: config.driver_id.clone(), + mandatory, + client, + config: ModuleRuntimeConfig::from(config), + } + } + + fn other(context: &'static str, error: impl std::fmt::Display) -> MemoryError { + MemoryError::Other(anyhow::anyhow!("{context}: {error}")) + } + + fn cross( + value: &A, + context: &'static str, + ) -> Result { + let value = serde_json::to_value(value).map_err(|error| Self::other(context, error))?; + serde_json::from_value(value).map_err(|error| Self::other(context, error)) + } +} + +fn validate_ingest_item(item: &IngestItem) -> Result<(), MemoryError> { + if item.taint != MemoryTaint::default() { + return Err(MemoryError::Invalid( + "ingest cannot preserve a non-default taint in the chunk tier".to_string(), + )); + } + if item.content.trim().is_empty() { + return Err(MemoryError::Invalid( + "ingest content must not be empty".to_string(), + )); + } + if let Some(mime) = item.mime.as_deref() { + let mime = mime.trim().to_ascii_lowercase(); + let base = mime.split(';').next().unwrap_or("").trim(); + if !(base.starts_with("text/") + || base.ends_with("+json") + || base.ends_with("+xml") + || matches!( + base, + "application/json" | "application/xml" | "application/x-ndjson" + )) + { + return Err(MemoryError::Invalid(format!( + "unsupported MIME '{mime}': ingest accepts decoded text only" + ))); + } + } + Ok(()) +} + +async fn blocking( + config: ModuleRuntimeConfig, + context: &'static str, + run: F, +) -> Result +where + T: Send + 'static, + F: FnOnce(&ModuleRuntimeConfig) -> anyhow::Result + Send + 'static, +{ + tokio::task::spawn_blocking(move || run(&config)) + .await + .map_err(|error| ModuleMemoryProvider::other(context, error))? + .map_err(|error| ModuleMemoryProvider::other(context, error)) +} + +#[async_trait] +impl MemoryCore for ModuleMemoryProvider { + async fn store( + &self, + namespace: &str, + key: &str, + content: &str, + category: MemoryCategory, + session_id: Option<&str>, + taint: MemoryTaint, + ) -> Result<(), MemoryError> { + self.mandatory + .store(namespace, key, content, category, session_id, taint) + .await + } + async fn get(&self, namespace: &str, key: &str) -> Result, MemoryError> { + self.mandatory.get(namespace, key).await + } + async fn forget(&self, namespace: &str, key: &str) -> Result { + self.mandatory.forget(namespace, key).await + } + async fn list( + &self, + namespace: Option<&str>, + category: Option<&MemoryCategory>, + session_id: Option<&str>, + ) -> Result, MemoryError> { + self.mandatory.list(namespace, category, session_id).await + } + async fn namespaces(&self) -> Result, MemoryError> { + self.mandatory.namespaces().await + } +} + +#[async_trait] +impl MemoryRecall for ModuleMemoryProvider { + async fn recall( + &self, + query: &str, + limit: usize, + opts: &OwnedRecallOpts, + scope: Option<&SourceScope>, + ) -> Result, MemoryError> { + self.mandatory.recall(query, limit, opts, scope).await + } +} + +#[async_trait] +impl MemoryPortability for ModuleMemoryProvider { + async fn export_page( + &self, + cursor: Option<&str>, + limit: usize, + ) -> Result { + self.mandatory.export_page(cursor, limit).await + } + async fn import_records( + &self, + records: Vec, + ) -> Result { + self.mandatory.import_records(records).await + } +} + +#[async_trait] +impl MemoryDocuments for ModuleMemoryProvider { + async fn put_document(&self, input: NamespaceDocumentInput) -> Result { + let input = Self::cross(&input, "convert document input")?; + self.client + .put_doc(input) + .await + .map_err(|error| Self::other("put_document", error)) + } + async fn get_document( + &self, + namespace: &str, + key: &str, + ) -> Result, MemoryError> { + let document = self + .client + .get_document(namespace, key) + .await + .map_err(|error| Self::other("get_document", error))?; + document + .map(|document| Self::cross(&document, "convert stored document")) + .transpose() + } + async fn query_documents( + &self, + namespace: &str, + query: &str, + limit: usize, + ) -> Result { + let limit = u32::try_from(limit).unwrap_or(u32::MAX); + let context = self + .client + .query_namespace_context_data(namespace, query, limit) + .await + .map_err(|error| Self::other("query_documents", error))?; + Self::cross(&context, "convert document query result") + } +} + +#[async_trait] +impl MemoryIngest for ModuleMemoryProvider { + async fn ingest_document(&self, item: IngestItem) -> Result { + validate_ingest_item(&item)?; + let document = tinycortex::memory::ingest::canonicalize::document::DocumentInput { + provider: item.source.as_str().to_string(), + title: String::new(), + body: item.content, + modified_at: item.timestamp.unwrap_or_else(Utc::now), + source_ref: item.source_ref.map(|source_ref| source_ref.value), + }; + let result = tinymemory_core::ingest_pipeline::ingest_document_with_scope( + &self.config, + &item.source_id, + &item.owner, + item.tags, + document, + item.path_scope, + ) + .await + .map_err(|error| Self::other("ingest document", error))?; + Ok(IngestOutcome { + written: u32::try_from(result.chunks_written).unwrap_or(u32::MAX), + skipped: if result.already_ingested { + 1 + } else { + u32::try_from(result.chunks_dropped).unwrap_or(u32::MAX) + }, + ids: result.chunk_ids, + }) + } + + async fn ingest_chat(&self, messages: Vec) -> Result { + let Some(first) = messages.first() else { + return Ok(IngestOutcome::default()); + }; + let source_id = first.source_id.clone(); + let owner = first.owner.clone(); + let tags = first.tags.clone(); + let platform = first.source.as_str().to_string(); + for item in &messages { + validate_ingest_item(item)?; + if item.source_id != source_id { + return Err(MemoryError::Invalid( + "ingest_chat batches must contain one conversation".to_string(), + )); + } + } + let batch = tinycortex::memory::ingest::canonicalize::chat::ChatBatch { + platform, + channel_label: source_id.clone(), + messages: messages + .into_iter() + .map( + |item| tinycortex::memory::ingest::canonicalize::chat::ChatMessage { + author: item.owner, + timestamp: item.timestamp.unwrap_or_else(Utc::now), + text: item.content, + source_ref: item.source_ref.map(|source_ref| source_ref.value), + }, + ) + .collect(), + }; + let result = tinymemory_core::ingest_pipeline::ingest_chat( + &self.config, + &source_id, + &owner, + tags, + batch, + ) + .await + .map_err(|error| Self::other("ingest chat", error))?; + Ok(IngestOutcome { + written: u32::try_from(result.chunks_written).unwrap_or(u32::MAX), + skipped: if result.already_ingested { + 1 + } else { + u32::try_from(result.chunks_dropped).unwrap_or(u32::MAX) + }, + ids: result.chunk_ids, + }) + } +} + +#[async_trait] +impl MemoryGraph for ModuleMemoryProvider { + async fn kv_get( + &self, + namespace: Option<&str>, + key: &str, + ) -> Result, MemoryError> { + let record = self + .client + .kv_records(namespace) + .await + .map_err(|error| Self::other("kv_get", error))? + .into_iter() + .find(|record| record.key == key); + record + .map(|record| Self::cross(&record, "convert key/value record")) + .transpose() + } + async fn kv_put( + &self, + namespace: Option<&str>, + key: &str, + value: serde_json::Value, + ) -> Result<(), MemoryError> { + self.client + .kv_set(namespace, key, &value) + .await + .map_err(|error| Self::other("kv_put", error)) + } + async fn kv_list( + &self, + namespace: Option<&str>, + prefix: Option<&str>, + limit: usize, + ) -> Result, MemoryError> { + let mut records = self + .client + .kv_records(namespace) + .await + .map_err(|error| Self::other("kv_list", error))?; + if let Some(prefix) = prefix { + records.retain(|record| record.key.starts_with(prefix)); + } + records.truncate(limit); + Self::cross(&records, "convert key/value records") + } + async fn relations( + &self, + namespace: Option<&str>, + subject: Option<&str>, + predicate: Option<&str>, + limit: usize, + ) -> Result, MemoryError> { + let mut records = self + .client + .graph_relations(namespace, subject, predicate) + .await + .map_err(|error| Self::other("relations", error))?; + records.truncate(limit); + Self::cross(&records, "convert graph relations") + } + async fn put_relation(&self, relation: GraphRelationRecord) -> Result<(), MemoryError> { + self.client + .graph_upsert( + relation.namespace.as_deref(), + &relation.subject, + &relation.predicate, + &relation.object, + &relation.attrs, + ) + .await + .map_err(|error| Self::other("put_relation", error)) + } +} + +#[async_trait] +impl MemoryGoals for ModuleMemoryProvider { + async fn goals(&self) -> Result { + let workspace = self.config.workspace_dir.clone(); + let document = + tokio::task::spawn_blocking(move || tinycortex::memory::goals::store::load(&workspace)) + .await + .map_err(|error| Self::other("join goals read", error))? + .map_err(|error| Self::other("read goals", error))?; + Self::cross(&document, "convert goals") + } + + async fn set_goals(&self, goals: GoalsDoc) -> Result<(), MemoryError> { + let workspace = self.config.workspace_dir.clone(); + let mut goals = Self::cross(&goals, "convert goals")?; + tokio::task::spawn_blocking(move || { + tinycortex::memory::goals::store::save(&workspace, &mut goals) + }) + .await + .map_err(|error| Self::other("join goals write", error))? + .map_err(|error| Self::other("write goals", error)) + } +} + +#[async_trait] +impl MemoryToolMemory for ModuleMemoryProvider { + async fn tool_rules(&self, tool_name: &str) -> Result, MemoryError> { + let rules = tinymemory_core::tool_memory::tool_memory_store(self.client.memory_handle()) + .list_rules(tool_name) + .await + .map_err(|error| Self::other("list tool rules", error))?; + Self::cross(&rules, "convert tool rules") + } + + async fn put_tool_rule(&self, rule: ToolMemoryRule) -> Result<(), MemoryError> { + let rule = Self::cross(&rule, "convert tool rule")?; + tinymemory_core::tool_memory::tool_memory_store(self.client.memory_handle()) + .put_rule(rule) + .await + .map(|_| ()) + .map_err(|error| Self::other("put tool rule", error)) + } + + async fn delete_tool_rule(&self, tool_name: &str, rule_id: &str) -> Result { + tinymemory_core::tool_memory::tool_memory_store(self.client.memory_handle()) + .delete_rule(tool_name, rule_id) + .await + .map_err(|error| Self::other("delete tool rule", error)) + } +} + +#[async_trait] +impl MemoryTree for ModuleMemoryProvider { + async fn append(&self, request: IngestRequest) -> Result<(), MemoryError> { + tinycortex::memory::tree::runtime::store::validate_namespace(&request.namespace) + .map_err(MemoryError::Invalid)?; + if request.content.trim().is_empty() { + return Err(MemoryError::Invalid( + "content must not be empty".to_string(), + )); + } + let namespace = request.namespace.trim().to_string(); + let content = request.content; + let timestamp = request.timestamp.unwrap_or_else(Utc::now); + let metadata = request.metadata; + blocking(self.config.clone(), "append tree content", move |config| { + tinymemory_core::tree::tree_runtime::store::buffer_write( + config, + &namespace, + &content, + ×tamp, + metadata.as_ref(), + ) + .map(|_| ()) + }) + .await + } + + async fn query_source( + &self, + namespace: &str, + source_id: &str, + limit: usize, + scope: Option<&SourceScope>, + ) -> Result, MemoryError> { + tinycortex::memory::tree::runtime::store::validate_namespace(namespace) + .map_err(MemoryError::Invalid)?; + let query = tinymemory_core::store::chunks::ListChunksQuery { + source_id: Some(source_id.to_string()), + source_scope: scope.map(|scope| scope.allow.iter().cloned().collect::>()), + limit: Some(limit), + exclude_dropped: true, + ..Default::default() + }; + let chunks = blocking(self.config.clone(), "query source", move |config| { + tinymemory_core::store::chunks::list_chunks(config, &query) + }) + .await?; + Self::cross(&chunks, "convert source chunks") + } + + async fn drill_down(&self, namespace: &str, node_id: &str) -> Result { + tinycortex::memory::tree::runtime::store::validate_namespace(namespace) + .map_err(MemoryError::Invalid)?; + tinycortex::memory::tree::runtime::store::validate_node_id(node_id) + .map_err(MemoryError::Invalid)?; + let namespace = namespace.trim().to_string(); + let node_id = node_id.to_string(); + let lookup_namespace = namespace.clone(); + let lookup_node = node_id.clone(); + let result = blocking(self.config.clone(), "drill down", move |config| { + let Some(node) = tinymemory_core::tree::tree_runtime::store::read_node( + config, + &lookup_namespace, + &lookup_node, + )? + else { + return Ok(None); + }; + let children = tinymemory_core::tree::tree_runtime::store::read_children( + config, + &lookup_namespace, + &lookup_node, + )?; + Ok(Some((node, children))) + }) + .await? + .ok_or_else(|| { + MemoryError::NotFound(format!("tree node '{node_id}' not found in '{namespace}'")) + })?; + Self::cross(&result, "convert tree drill-down") + .map(|(node, children)| QueryResult { node, children }) + } + + async fn seal(&self, namespace: &str) -> Result { + tinycortex::memory::tree::runtime::store::validate_namespace(namespace) + .map_err(MemoryError::Invalid)?; + let namespace = namespace.trim().to_string(); + let read_namespace = namespace.clone(); + let buffered = blocking(self.config.clone(), "read tree buffer", move |config| { + tinymemory_core::tree::tree_runtime::store::buffer_read(config, &read_namespace) + }) + .await?; + if !buffered.is_empty() { + let (model, _) = tinymemory_core::chat_host::create_chat_model_with_model_id( + "summarization", + &self.config, + self.config.default_temperature, + ) + .map_err(|error| Self::other("create summarizer", error))?; + tinymemory_core::tree::tree_runtime::engine::run_summarization( + &self.config, + model.as_ref(), + &namespace, + Utc::now(), + ) + .await + .map_err(|error| Self::other("seal tree", error))?; + } + let status = blocking(self.config.clone(), "read tree status", move |config| { + tinymemory_core::tree::tree_runtime::store::get_tree_status(config, &namespace) + }) + .await?; + Self::cross(&status, "convert tree status") + } + + async fn cascade(&self, namespace: &str) -> Result { + tinycortex::memory::tree::runtime::store::validate_namespace(namespace) + .map_err(MemoryError::Invalid)?; + let namespace = namespace.trim().to_string(); + let read_namespace = namespace.clone(); + let status = blocking(self.config.clone(), "read tree status", move |config| { + tinymemory_core::tree::tree_runtime::store::get_tree_status(config, &read_namespace) + }) + .await?; + if status.total_nodes == 0 { + return Self::cross(&status, "convert tree status"); + } + let (model, _) = tinymemory_core::chat_host::create_chat_model_with_model_id( + "summarization", + &self.config, + self.config.default_temperature, + ) + .map_err(|error| Self::other("create summarizer", error))?; + let status = tinymemory_core::tree::tree_runtime::engine::rebuild_tree( + &self.config, + model.as_ref(), + &namespace, + ) + .await + .map_err(|error| Self::other("cascade tree", error))?; + Self::cross(&status, "convert tree status") + } +} + +#[async_trait] +impl MemoryEntities for ModuleMemoryProvider { + async fn entities( + &self, + namespace: &str, + query: Option<&str>, + limit: usize, + ) -> Result, MemoryError> { + let namespace = namespace.to_string(); + let query_namespace = namespace.clone(); + let query = query.map(str::to_string); + let rows = blocking( + self.config.clone(), + "list namespace entities", + move |config| { + tinymemory_core::store::entities::namespace_entities( + config, + &query_namespace, + query.as_deref(), + limit, + ) + }, + ) + .await? + .into_iter() + .map(|hit| (hit.id, hit.kind, hit.name, hit.mentions)) + .collect::>(); + + let config = self.config.clone(); + blocking(config, "attach entity hotness", move |config| { + Ok(rows + .into_iter() + .map(|(id, kind, name, mentions)| { + let hotness_key = format!("{namespace}:{id}"); + let hotness = tinymemory_core::store::trees::hotness::get(config, &hotness_key) + .ok() + .flatten() + .map_or(0.0, |counters| { + f64::from( + tinymemory_core::tree_policy::TreePolicy::topic().topic_hotness( + &id, + &counters.stats(), + Utc::now().timestamp_millis(), + ), + ) + }); + EntityHit { + entity: EntityRef { id, kind, name }, + hotness, + mentions, + } + }) + .collect()) + }) + .await + } + + async fn entity_edges( + &self, + namespace: &str, + entity_id: &str, + limit: usize, + ) -> Result, MemoryError> { + let subject = entity_id.to_string(); + let lookup = subject.clone(); + let namespace = namespace.to_string(); + let query_namespace = namespace.clone(); + let neighbours = blocking(self.config.clone(), "read entity edges", move |config| { + tinymemory_core::store::entities::namespace_entity_edges( + config, + &query_namespace, + &lookup, + limit, + ) + }) + .await?; + Ok(neighbours + .into_iter() + .map(|(object, weight)| GraphRelationRecord { + namespace: Some(namespace.clone()), + subject: subject.clone(), + predicate: "co_occurs_with".to_string(), + object, + attrs: serde_json::Value::Null, + updated_at: 0.0, + evidence_count: weight, + order_index: None, + document_ids: Vec::new(), + chunk_ids: Vec::new(), + }) + .collect()) + } + + async fn touch_entities( + &self, + namespace: &str, + entity_ids: &[String], + ) -> Result<(), MemoryError> { + let entity_ids = entity_ids.to_vec(); + let namespace = namespace.to_string(); + blocking(self.config.clone(), "touch entities", move |config| { + let now = Utc::now().timestamp_millis(); + for entity_id in entity_ids { + let entity_id = format!("{namespace}:{entity_id}"); + let mut counters = + tinymemory_core::store::trees::hotness::get_or_fresh(config, &entity_id)?; + counters.mention_count_30d = counters.mention_count_30d.saturating_add(1); + counters.last_seen_ms = Some(now); + counters.last_updated_ms = now; + tinymemory_core::store::trees::hotness::upsert(config, &counters)?; + } + Ok(()) + }) + .await + } +} + +#[async_trait] +impl MemoryDiff for ModuleMemoryProvider { + async fn capture_snapshot(&self, source_id: &str) -> Result { + let source = tinymemory_core::sources::registry::decode_memory_sources(&self.config) + .into_iter() + .find(|source| source.id == source_id) + .ok_or_else(|| MemoryError::NotFound(source_id.to_string()))?; + let snapshot = tinymemory_core::diff::ops::take_snapshot( + &source, + &self.config, + tinymemory_core::diff::SnapshotTrigger::Manual, + ) + .await + .map_err(|error| Self::other("capture snapshot", error))?; + Ok(SnapshotRef { + id: snapshot.id, + source_id: snapshot.source_id, + label: snapshot.label, + item_count: snapshot.item_count, + taken_at_ms: snapshot.taken_at_ms, + }) + } + + async fn snapshots( + &self, + source_id: &str, + limit: usize, + ) -> Result, MemoryError> { + let snapshots = tinymemory_core::diff::ops::list_snapshots( + &self.config, + Some(source_id), + u32::try_from(limit).unwrap_or(u32::MAX), + ) + .await + .map_err(|error| Self::other("list snapshots", error))?; + Ok(snapshots + .into_iter() + .map(|snapshot| SnapshotRef { + id: snapshot.id, + source_id: snapshot.source_id, + label: snapshot.label, + item_count: snapshot.item_count, + taken_at_ms: snapshot.taken_at_ms, + }) + .collect()) + } + + async fn diff( + &self, + source_id: &str, + from: Option<&str>, + to: &str, + ) -> Result { + let result = tinymemory_core::diff::ops::compute_diff(&self.config, from, to, false) + .await + .map_err(|error| Self::other("compute diff", error))?; + if result.source_id != source_id { + return Err(MemoryError::Invalid(format!( + "snapshot '{to}' belongs to a different source" + ))); + } + let changes = result + .changes + .into_iter() + .map(|change| SourceChange { + item_id: change.item_id, + title: change.title, + kind: match change.kind { + tinymemory_core::diff::ChangeKind::Added => ChangeKind::Added, + tinymemory_core::diff::ChangeKind::Removed => ChangeKind::Removed, + tinymemory_core::diff::ChangeKind::Modified => ChangeKind::Modified, + }, + old_content_hash: change.old_content_hash, + new_content_hash: change.new_content_hash, + }) + .collect(); + Ok(DiffReport { + source_id: result.source_id, + from_snapshot_id: result.from_snapshot_id, + to_snapshot_id: result.to_snapshot_id, + added: result.summary.added, + removed: result.summary.removed, + modified: result.summary.modified, + unchanged: result.summary.unchanged, + changes, + }) + } +} + +#[async_trait] +impl MemorySourceSink for ModuleMemoryProvider { + async fn accept_source_items( + &self, + source_id: &str, + source_kind: &str, + items: Vec, + taint: MemoryTaint, + ) -> Result { + let namespace = format!("source:{source_id}"); + let mut outcome = IngestOutcome::default(); + for item in items { + if item.item_id.trim().is_empty() { + return Err(MemoryError::Invalid( + "source item_id must not be empty".to_string(), + )); + } + let title = if item.title.trim().is_empty() { + item.item_id.clone() + } else { + item.title.clone() + }; + let input = NamespaceDocumentInput { + namespace: namespace.clone(), + key: item.item_id, + title, + content: item.content, + source_type: source_kind.to_string(), + priority: "medium".to_string(), + tags: item.tags, + metadata: serde_json::json!({ + "sourceId": source_id, + "sourceKind": source_kind, + "url": item.url, + "mime": item.mime, + "updatedAtMs": item.updated_at_ms, + }), + category: "core".to_string(), + session_id: None, + document_id: None, + taint, + }; + let input = Self::cross(&input, "convert source document")?; + match self.client.put_doc(input).await { + Ok(id) => { + outcome.written = outcome.written.saturating_add(1); + outcome.ids.push(id); + } + Err(_) => { + outcome.skipped = outcome.skipped.saturating_add(1); + } + } + } + Ok(outcome) + } + + async fn forget_source(&self, source_id: &str) -> Result { + let namespace = format!("source:{source_id}"); + let listed = self + .client + .list_documents(Some(&namespace)) + .await + .map_err(|error| Self::other("list source documents", error))?; + let documents = listed + .get("documents") + .and_then(serde_json::Value::as_array) + .map_or(0, Vec::len); + if documents > 0 { + self.client + .clear_namespace(&namespace) + .await + .map_err(|error| Self::other("clear source documents", error))?; + } + let source_id = source_id.to_string(); + let chunks = blocking(self.config.clone(), "clear source chunks", move |config| { + use tinymemory_core::store::chunks::{ + delete_chunks_by_source, delete_orphaned_source_tree, SourceKind, + }; + let removed = delete_chunks_by_source(config, SourceKind::Document, &source_id)?; + delete_orphaned_source_tree(config, SourceKind::Document, &source_id)?; + Ok(removed) + }) + .await?; + Ok(u64::try_from(documents.saturating_add(chunks)).unwrap_or(u64::MAX)) + } +} + +#[async_trait] +impl MemoryMaintenance for ModuleMemoryProvider { + async fn reembed(&self) -> Result { + let (examined, changed) = + blocking(self.config.clone(), "enqueue re-embedding", move |config| { + let total = tinymemory_core::queue::count_total(config).unwrap_or(0); + let before = tinymemory_core::queue::count_by_status( + config, + tinymemory_core::queue::JobStatus::Ready, + ) + .unwrap_or(0); + tinymemory_core::queue::ensure_reembed_backfill(config); + let after = tinymemory_core::queue::count_by_status( + config, + tinymemory_core::queue::JobStatus::Ready, + ) + .unwrap_or(0); + Ok((total, after.saturating_sub(before))) + }) + .await?; + Ok(MaintenanceReport { + operation: "reembed".to_string(), + examined, + changed, + findings: vec![format!("enqueued {changed} re-embedding job(s)")], + }) + } + + async fn compact(&self) -> Result { + let (examined, changed) = + blocking(self.config.clone(), "compact memory queue", move |config| { + Ok(( + tinymemory_core::queue::count_total(config).unwrap_or(0), + u64::try_from(tinymemory_core::queue::recover_stale_locks(config).unwrap_or(0)) + .unwrap_or(u64::MAX), + )) + }) + .await?; + Ok(MaintenanceReport { + operation: "compact".to_string(), + examined, + changed, + findings: vec![format!("released {changed} stale queue lock(s)")], + }) + } + + async fn consolidate(&self) -> Result { + let (examined, enqueued) = blocking( + self.config.clone(), + "enqueue consolidation", + move |config| { + Ok(( + tinymemory_core::queue::count_total(config).unwrap_or(0), + tinymemory_core::queue::scheduler::enqueue_flush_stale_job(config) + .map_err(anyhow::Error::msg)?, + )) + }, + ) + .await?; + Ok(MaintenanceReport { + operation: "consolidate".to_string(), + examined, + changed: u64::from(enqueued), + findings: vec![if enqueued { + "enqueued a stale-buffer flush".to_string() + } else { + "a stale-buffer flush is already queued".to_string() + }], + }) + } + + async fn doctor(&self) -> Result { + let report = tinymemory_core::tree::health::async_run_doctor(&self.config).await; + Ok(MaintenanceReport { + operation: "doctor".to_string(), + examined: report.counters.total_chunks, + changed: 0, + findings: report + .stages + .into_iter() + .filter(|stage| !stage.ok) + .map(|stage| format!("{}: {}", stage.stage, stage.note)) + .collect(), + }) + } +} + +#[async_trait] +impl MemoryProvider for ModuleMemoryProvider { + fn driver_id(&self) -> &str { + &self.driver_id + } + fn capabilities(&self) -> Capabilities { + Capabilities::all() + } + async fn health(&self) -> MemoryHealth { + if self.client.memory_handle().health_check().await { + MemoryHealth::Ready + } else { + MemoryHealth::down("memory store is unavailable") + } + } + fn as_documents(&self) -> Option<&dyn MemoryDocuments> { + Some(self) + } + fn as_ingest(&self) -> Option<&dyn MemoryIngest> { + Some(self) + } + fn as_graph(&self) -> Option<&dyn MemoryGraph> { + Some(self) + } + fn as_goals(&self) -> Option<&dyn MemoryGoals> { + Some(self) + } + fn as_tool_memory(&self) -> Option<&dyn MemoryToolMemory> { + Some(self) + } + fn as_tree(&self) -> Option<&dyn MemoryTree> { + Some(self) + } + fn as_entities(&self) -> Option<&dyn MemoryEntities> { + Some(self) + } + fn as_diff(&self) -> Option<&dyn MemoryDiff> { + Some(self) + } + fn as_sources(&self) -> Option<&dyn MemorySourceSink> { + Some(self) + } + fn as_maintenance(&self) -> Option<&dyn MemoryMaintenance> { + Some(self) + } +} diff --git a/crates/tinymemory-module/src/service/mod.rs b/crates/tinymemory-module/src/service/mod.rs index d8b2bcf..d7877de 100644 --- a/crates/tinymemory-module/src/service/mod.rs +++ b/crates/tinymemory-module/src/service/mod.rs @@ -1,7 +1,7 @@ //! `TinyBus` service boundary for the memory surface. //! -//! One object, `/ai/tinyhumans/tinymemory/Memory`, exporting the mandatory -//! capability families plus the four driver-level methods: +//! One object, `/ai/tinyhumans/tinymemory/Memory`, exporting every capability +//! family plus the four driver-level methods. //! //! ```text //! DriverId() -> String @@ -21,23 +21,14 @@ //! //! # Why the method list mirrors a trait exactly //! -//! These twelve are `tinymemory_api`'s [`MemoryProvider`] plus its three -//! mandatory supertraits, with the borrows replaced by owned equivalents. That +//! These are `tinymemory_api`'s [`MemoryProvider`] and all of its capability +//! traits, with the borrows replaced by owned equivalents. That //! is deliberate: the host binds an `Arc`, so a host-side //! client that forwards each method one-for-one is a *complete* provider with no //! translation layer in between. Anything cleverer — batching, a combined //! "recall and store" call — would put engine semantics on the wire, where two //! sides could disagree about them. //! -//! # Why only the mandatory families -//! -//! `tinymemory-tinycortex` advertises Core, Recall and Portability and nothing -//! else, because the ten optional families are reached through engine entry -//! points that need a host's configuration, embedding compute and job queue. -//! This module serves exactly what that adapter can provide. Serving more would -//! mean advertising capabilities whose accessors return nothing, which -//! `audit_provider` is specifically written to catch. -//! //! # Everything travels inline //! //! A `TinyBus` frame is JSON capped at 16 MiB. That is a real constraint for a @@ -80,16 +71,26 @@ use std::sync::Arc; use tinybus::{Connection, Error as BusError, Result as BusResult}; -use tinymemory_api::capabilities::Capabilities; +use tinymemory_api::capabilities::{Capabilities, Capability}; +use tinymemory_api::chunks::Chunk; use tinymemory_api::error::MemoryError; +use tinymemory_api::goals::GoalsDoc; use tinymemory_api::health::MemoryHealth; -use tinymemory_api::provider::types::{ExportPage, ExportRecord, ImportOutcome, SourceScope}; +use tinymemory_api::provider::types::{ + DiffReport, EntityHit, ExportPage, ExportRecord, ImportOutcome, IngestItem, IngestOutcome, + MaintenanceReport, SnapshotRef, SourceItem, SourceScope, +}; // `MemoryCore`, `MemoryRecall` and `MemoryPortability` are deliberately not // imported: they are supertraits of `MemoryProvider`, so their methods are // already callable on the trait object. use tinymemory_api::provider::MemoryProvider; use tinymemory_api::recall::OwnedRecallOpts; -use tinymemory_api::types::{MemoryCategory, MemoryEntry, MemoryTaint, NamespaceSummary}; +use tinymemory_api::tool_memory::ToolMemoryRule; +use tinymemory_api::tree::{IngestRequest, QueryResult, TreeStatus}; +use tinymemory_api::types::{ + GraphRelationRecord, MemoryCategory, MemoryEntry, MemoryKvRecord, MemoryTaint, + NamespaceDocumentInput, NamespaceRetrievalContext, NamespaceSummary, StoredMemoryDocument, +}; use tinymemory_api::wire; /// Well-known name exported by the `TinyMemory` module. @@ -110,6 +111,15 @@ impl MemoryService { } } +macro_rules! require_family { + ($service:expr, $accessor:ident, $capability:expr) => { + $service + .provider + .$accessor() + .ok_or_else(|| into_bus_error(&MemoryError::unsupported($capability)))? + }; +} + #[tinybus::interface(name = "ai.tinyhumans.tinymemory.Memory")] impl MemoryService { /// The bound driver's stable identifier. @@ -270,6 +280,301 @@ impl MemoryService { .await .map_err(|error| into_bus_error(&error)) } + + async fn ingest_document(&self, item: IngestItem) -> BusResult { + require_family!(self, as_ingest, Capability::Ingest) + .ingest_document(item) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn ingest_chat(&self, messages: Vec) -> BusResult { + require_family!(self, as_ingest, Capability::Ingest) + .ingest_chat(messages) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn put_document(&self, input: NamespaceDocumentInput) -> BusResult { + require_family!(self, as_documents, Capability::Documents) + .put_document(input) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn get_document( + &self, + namespace: String, + key: String, + ) -> BusResult> { + require_family!(self, as_documents, Capability::Documents) + .get_document(&namespace, &key) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn query_documents( + &self, + namespace: String, + query: String, + limit: usize, + ) -> BusResult { + let response = require_family!(self, as_documents, Capability::Documents) + .query_documents(&namespace, &query, limit) + .await + .map_err(|error| into_bus_error(&error))?; + ensure_response_fits(&response, "QueryDocuments")?; + Ok(response) + } + + async fn append(&self, request: IngestRequest) -> BusResult<()> { + require_family!(self, as_tree, Capability::Tree) + .append(request) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn query_source( + &self, + namespace: String, + source_id: String, + limit: usize, + scope: Option, + ) -> BusResult> { + let response = require_family!(self, as_tree, Capability::Tree) + .query_source(&namespace, &source_id, limit, scope.as_ref()) + .await + .map_err(|error| into_bus_error(&error))?; + ensure_response_fits(&response, "QuerySource")?; + Ok(response) + } + + async fn drill_down(&self, namespace: String, node_id: String) -> BusResult { + require_family!(self, as_tree, Capability::Tree) + .drill_down(&namespace, &node_id) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn seal(&self, namespace: String) -> BusResult { + require_family!(self, as_tree, Capability::Tree) + .seal(&namespace) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn cascade(&self, namespace: String) -> BusResult { + require_family!(self, as_tree, Capability::Tree) + .cascade(&namespace) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn entities( + &self, + namespace: String, + query: Option, + limit: usize, + ) -> BusResult> { + let response = require_family!(self, as_entities, Capability::Entities) + .entities(&namespace, query.as_deref(), limit) + .await + .map_err(|error| into_bus_error(&error))?; + ensure_response_fits(&response, "Entities")?; + Ok(response) + } + + async fn entity_edges( + &self, + namespace: String, + entity_id: String, + limit: usize, + ) -> BusResult> { + require_family!(self, as_entities, Capability::Entities) + .entity_edges(&namespace, &entity_id, limit) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn touch_entities(&self, namespace: String, entity_ids: Vec) -> BusResult<()> { + require_family!(self, as_entities, Capability::Entities) + .touch_entities(&namespace, &entity_ids) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn kv_get( + &self, + namespace: Option, + key: String, + ) -> BusResult> { + require_family!(self, as_graph, Capability::Graph) + .kv_get(namespace.as_deref(), &key) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn kv_put( + &self, + namespace: Option, + key: String, + value: serde_json::Value, + ) -> BusResult<()> { + require_family!(self, as_graph, Capability::Graph) + .kv_put(namespace.as_deref(), &key, value) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn kv_list( + &self, + namespace: Option, + prefix: Option, + limit: usize, + ) -> BusResult> { + let response = require_family!(self, as_graph, Capability::Graph) + .kv_list(namespace.as_deref(), prefix.as_deref(), limit) + .await + .map_err(|error| into_bus_error(&error))?; + ensure_response_fits(&response, "KvList")?; + Ok(response) + } + + async fn relations( + &self, + namespace: Option, + subject: Option, + predicate: Option, + limit: usize, + ) -> BusResult> { + let response = require_family!(self, as_graph, Capability::Graph) + .relations( + namespace.as_deref(), + subject.as_deref(), + predicate.as_deref(), + limit, + ) + .await + .map_err(|error| into_bus_error(&error))?; + ensure_response_fits(&response, "Relations")?; + Ok(response) + } + + async fn put_relation(&self, relation: GraphRelationRecord) -> BusResult<()> { + require_family!(self, as_graph, Capability::Graph) + .put_relation(relation) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn capture_snapshot(&self, source_id: String) -> BusResult { + require_family!(self, as_diff, Capability::Diff) + .capture_snapshot(&source_id) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn snapshots(&self, source_id: String, limit: usize) -> BusResult> { + require_family!(self, as_diff, Capability::Diff) + .snapshots(&source_id, limit) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn diff( + &self, + source_id: String, + from: Option, + to: String, + ) -> BusResult { + require_family!(self, as_diff, Capability::Diff) + .diff(&source_id, from.as_deref(), &to) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn goals(&self) -> BusResult { + require_family!(self, as_goals, Capability::Goals) + .goals() + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn set_goals(&self, goals: GoalsDoc) -> BusResult<()> { + require_family!(self, as_goals, Capability::Goals) + .set_goals(goals) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn tool_rules(&self, tool_name: String) -> BusResult> { + require_family!(self, as_tool_memory, Capability::ToolMemory) + .tool_rules(&tool_name) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn put_tool_rule(&self, rule: ToolMemoryRule) -> BusResult<()> { + require_family!(self, as_tool_memory, Capability::ToolMemory) + .put_tool_rule(rule) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn delete_tool_rule(&self, tool_name: String, rule_id: String) -> BusResult { + require_family!(self, as_tool_memory, Capability::ToolMemory) + .delete_tool_rule(&tool_name, &rule_id) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn accept_source_items( + &self, + source_id: String, + source_kind: String, + items: Vec, + taint: MemoryTaint, + ) -> BusResult { + require_family!(self, as_sources, Capability::Sources) + .accept_source_items(&source_id, &source_kind, items, taint) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn forget_source(&self, source_id: String) -> BusResult { + require_family!(self, as_sources, Capability::Sources) + .forget_source(&source_id) + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn reembed(&self) -> BusResult { + require_family!(self, as_maintenance, Capability::Maintenance) + .reembed() + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn compact(&self) -> BusResult { + require_family!(self, as_maintenance, Capability::Maintenance) + .compact() + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn consolidate(&self) -> BusResult { + require_family!(self, as_maintenance, Capability::Maintenance) + .consolidate() + .await + .map_err(|error| into_bus_error(&error)) + } + + async fn doctor(&self) -> BusResult { + require_family!(self, as_maintenance, Capability::Maintenance) + .doctor() + .await + .map_err(|error| into_bus_error(&error)) + } } /// The response-size ceiling for a method that returns a list of entries. @@ -279,14 +584,6 @@ impl MemoryService { /// pathological string, so a response that passes this check fits with margin. pub(crate) const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024; -/// Per-entry allowance for the fields that are not `content`. -/// -/// Keys, namespaces, timestamps, category and taint. Deliberately generous: this -/// check exists to stop a response overflowing a frame, and over-estimating -/// refuses slightly early while under-estimating fails at the transport with an -/// error the caller cannot act on. -const PER_ENTRY_OVERHEAD_BYTES: usize = 512; - /// Refuse a response that would not fit in a frame. /// /// # Why a refusal and not a truncation @@ -310,25 +607,22 @@ const PER_ENTRY_OVERHEAD_BYTES: usize = 512; /// /// [`wire::BUDGET_EXCEEDED`], when the estimate exceeds [`MAX_RESPONSE_BYTES`]. /// The message names the method and the sizes, never entry content. -fn ensure_response_fits(entries: &[MemoryEntry], method: &str) -> BusResult<()> { - let estimate: usize = entries - .iter() - .map(|entry| entry.content.len().saturating_add(PER_ENTRY_OVERHEAD_BYTES)) - .sum(); +fn ensure_response_fits(response: &T, method: &str) -> BusResult<()> { + let estimate = serde_json::to_vec(response) + .map_err(|error| BusError::Protocol(error.to_string()))? + .len(); if estimate > MAX_RESPONSE_BYTES { log::warn!( - "[tinymemory:module] {method} refused: {} entries estimated at {estimate} bytes \ - exceeds the {MAX_RESPONSE_BYTES} byte response ceiling", - entries.len() + "[tinymemory:module] {method} refused: response estimated at {estimate} bytes \ + exceeds the {MAX_RESPONSE_BYTES} byte response ceiling" ); return Err(BusError::MethodFailed { name: wire::BUDGET_EXCEEDED.to_string(), message: format!( - "{method} would return {} entries (~{estimate} bytes), over the \ + "{method} would return ~{estimate} bytes, over the \ {MAX_RESPONSE_BYTES} byte response ceiling; narrow the query by \ - namespace, category or session", - entries.len() + namespace, category or session" ), }); } diff --git a/crates/tinymemory-module/src/service/test.rs b/crates/tinymemory-module/src/service/test.rs index 573892e..71d14aa 100644 --- a/crates/tinymemory-module/src/service/test.rs +++ b/crates/tinymemory-module/src/service/test.rs @@ -133,7 +133,10 @@ fn an_ordinary_list_response_is_not_refused() { #[test] fn an_empty_list_response_is_not_refused() { - assert!(super::ensure_response_fits(&[], "List").is_ok()); + assert!( + super::ensure_response_fits(&Vec::::new(), "List") + .is_ok() + ); } #[test] @@ -210,7 +213,10 @@ fn the_per_entry_overhead_is_counted_so_many_tiny_entries_still_trip_it() { // A million empty entries carry no content at all but still cannot cross a // frame — the JSON structure around each one is the payload. Counting only // `content.len()` would let this through. - let count = super::MAX_RESPONSE_BYTES / super::PER_ENTRY_OVERHEAD_BYTES + 1; + let encoded_entry = serde_json::to_vec(&entry_of(0)) + .expect("serializable") + .len(); + let count = super::MAX_RESPONSE_BYTES / encoded_entry + 1; let entries: Vec<_> = (0..count).map(|_| entry_of(0)).collect(); assert!( super::ensure_response_fits(&entries, "List").is_err(), diff --git a/crates/tinymemory-module/tests/module_e2e.rs b/crates/tinymemory-module/tests/module_e2e.rs index c0d13f4..7e2ea9c 100644 --- a/crates/tinymemory-module/tests/module_e2e.rs +++ b/crates/tinymemory-module/tests/module_e2e.rs @@ -197,7 +197,7 @@ fn proxy(connection: &Connection) -> tinybus::Proxy { #[tokio::test] #[ignore = "drives a real dlopen'ed module; must be the only such test in the process — see the module docs"] -async fn the_module_advertises_exactly_the_mandatory_families() { +async fn the_module_advertises_the_complete_tinymemory_api() { let workspace = tempfile::tempdir().expect("tempdir"); let (client, _host, _task) = admit_module(workspace.path()).await; @@ -206,18 +206,10 @@ async fn the_module_advertises_exactly_the_mandatory_families() { .await .expect("Capabilities"); - // The adapter deliberately advertises only what it can reach. Advertising - // more would make `audit_provider` fail host-side, and would register RPC - // methods that answer errors. - // - // Asserted as an exact set rather than as "the mandatory three are present - // and `Tree` is absent": that weaker pair passes while any *other* optional - // family is advertised, which is the same overstatement with a different - // name on it. assert_eq!( capabilities, - Capabilities::mandatory(), - "the module must advertise exactly the mandatory families" + Capabilities::all(), + "the compiled module must own every TinyMemory capability family" ); for mandatory in Capability::MANDATORY { assert!( @@ -500,6 +492,38 @@ const EXPECTED_METHODS: &[&str] = &[ "Recall", "ExportPage", "ImportRecords", + "IngestDocument", + "IngestChat", + "PutDocument", + "GetDocument", + "QueryDocuments", + "Append", + "QuerySource", + "DrillDown", + "Seal", + "Cascade", + "Entities", + "EntityEdges", + "TouchEntities", + "KvGet", + "KvPut", + "KvList", + "Relations", + "PutRelation", + "CaptureSnapshot", + "Snapshots", + "Diff", + "Goals", + "SetGoals", + "ToolRules", + "PutToolRule", + "DeleteToolRule", + "AcceptSourceItems", + "ForgetSource", + "Reembed", + "Compact", + "Consolidate", + "Doctor", ]; #[tokio::test]