use std::sync::{Arc, RwLock};
use khive_db::StorageBackend;
use khive_gate::{ActorRef, AllowAllGate, GateRef};
use khive_storage::{EntityStore, EventStore, GraphStore, NoteStore, SqlAccess};
use khive_types::{EdgeEndpointRule, Namespace};
use lattice_embed::{
CachedEmbeddingService, EmbeddingModel, EmbeddingService, NativeEmbeddingService,
};
use tokio::sync::OnceCell;
use crate::error::RuntimeResult;
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct BackendId(pub String);
impl BackendId {
pub const MAIN: &'static str = "main";
pub fn new(name: impl Into<String>) -> Self {
Self(name.into())
}
pub fn main() -> Self {
Self(Self::MAIN.to_string())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for BackendId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
mod private {
#[derive(Clone, Debug)]
pub(crate) struct Sealed;
}
#[derive(Clone, Debug)]
pub struct NamespaceToken {
namespace: Namespace,
actor: ActorRef,
_sealed: private::Sealed,
}
impl NamespaceToken {
pub(crate) fn mint_authorized(namespace: Namespace, actor: ActorRef) -> Self {
Self {
namespace,
actor,
_sealed: private::Sealed,
}
}
#[allow(dead_code)]
pub(crate) fn local() -> Self {
Self::mint_authorized(Namespace::local(), ActorRef::anonymous())
}
#[allow(dead_code)]
pub(crate) fn for_namespace(ns: Namespace) -> Self {
Self::mint_authorized(ns, ActorRef::anonymous())
}
pub fn namespace(&self) -> &Namespace {
&self.namespace
}
pub fn actor(&self) -> &ActorRef {
&self.actor
}
}
#[derive(Clone, Debug)]
pub struct RuntimeConfig {
pub db_path: Option<std::path::PathBuf>,
pub default_namespace: Namespace,
pub embedding_model: Option<EmbeddingModel>,
pub gate: GateRef,
pub packs: Vec<String>,
pub backend_id: BackendId,
}
pub fn parse_pack_list(s: &str) -> Vec<String> {
s.split(|c: char| c == ',' || c.is_whitespace())
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_owned)
.collect()
}
impl Default for RuntimeConfig {
fn default() -> Self {
let db_path = std::env::var("HOME")
.ok()
.map(|h| std::path::PathBuf::from(h).join(".khive/khive-graph.db"));
let embedding_model = std::env::var("KHIVE_EMBEDDING_MODEL")
.ok()
.and_then(|s| s.parse().ok())
.or(Some(EmbeddingModel::AllMiniLmL6V2));
let packs = std::env::var("KHIVE_PACKS")
.ok()
.map(|s| parse_pack_list(&s))
.filter(|v| !v.is_empty())
.unwrap_or_else(|| vec!["kg".to_string()]);
Self {
db_path,
default_namespace: Namespace::local(),
embedding_model,
gate: Arc::new(AllowAllGate),
packs,
backend_id: BackendId::main(),
}
}
}
#[derive(Clone)]
pub struct KhiveRuntime {
backend: Arc<StorageBackend>,
config: RuntimeConfig,
embedder: Arc<OnceCell<Arc<dyn EmbeddingService>>>,
edge_rules: Arc<RwLock<Vec<EdgeEndpointRule>>>,
}
impl KhiveRuntime {
pub fn new(config: RuntimeConfig) -> RuntimeResult<Self> {
let backend = match &config.db_path {
Some(path) => {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).ok();
}
StorageBackend::sqlite(path)?
}
None => StorageBackend::memory()?,
};
Ok(Self {
backend: Arc::new(backend),
config,
embedder: Arc::new(OnceCell::new()),
edge_rules: Arc::new(RwLock::new(Vec::new())),
})
}
pub fn from_backend(backend: Arc<StorageBackend>, config: RuntimeConfig) -> Self {
Self {
backend,
config,
embedder: Arc::new(OnceCell::new()),
edge_rules: Arc::new(RwLock::new(Vec::new())),
}
}
pub fn memory() -> RuntimeResult<Self> {
Self::new(RuntimeConfig {
db_path: None,
default_namespace: Namespace::local(),
embedding_model: None,
gate: Arc::new(AllowAllGate),
packs: vec!["kg".to_string()],
backend_id: BackendId::main(),
})
}
pub fn backend_id(&self) -> &BackendId {
&self.config.backend_id
}
pub fn config(&self) -> &RuntimeConfig {
&self.config
}
pub fn backend(&self) -> &StorageBackend {
&self.backend
}
pub fn entities(&self, token: &NamespaceToken) -> RuntimeResult<Arc<dyn EntityStore>> {
Ok(self
.backend
.entities_for_namespace(token.namespace().as_str())?)
}
pub fn graph(&self, token: &NamespaceToken) -> RuntimeResult<Arc<dyn GraphStore>> {
Ok(self
.backend
.graph_for_namespace(token.namespace().as_str())?)
}
pub fn notes(&self, token: &NamespaceToken) -> RuntimeResult<Arc<dyn NoteStore>> {
Ok(self
.backend
.notes_for_namespace(token.namespace().as_str())?)
}
pub fn events(&self, token: &NamespaceToken) -> RuntimeResult<Arc<dyn EventStore>> {
Ok(self
.backend
.events_for_namespace(token.namespace().as_str())?)
}
pub fn sql(&self) -> Arc<dyn SqlAccess> {
self.backend.sql()
}
pub fn vectors(
&self,
token: &NamespaceToken,
) -> RuntimeResult<Arc<dyn khive_storage::VectorStore>> {
let model = self
.config
.embedding_model
.ok_or_else(|| crate::RuntimeError::Unconfigured("embedding_model".into()))?;
Ok(self.backend.vectors_for_namespace(
&vec_model_key(model),
model.dimensions(),
token.namespace().as_str(),
)?)
}
pub fn text(
&self,
token: &NamespaceToken,
) -> RuntimeResult<Arc<dyn khive_storage::TextSearch>> {
let key = format!("entities_{}", sanitize_key(token.namespace().as_str()));
Ok(self.backend.text(&key)?)
}
pub fn text_for_notes(
&self,
token: &NamespaceToken,
) -> RuntimeResult<Arc<dyn khive_storage::TextSearch>> {
let key = format!("notes_{}", sanitize_key(token.namespace().as_str()));
Ok(self.backend.text(&key)?)
}
pub fn authorize(&self, ns: Namespace) -> NamespaceToken {
NamespaceToken::mint_authorized(ns, ActorRef::anonymous())
}
pub fn install_edge_rules(&self, rules: Vec<EdgeEndpointRule>) {
if let Ok(mut guard) = self.edge_rules.write() {
*guard = rules;
}
}
pub(crate) fn pack_edge_rules(&self) -> Vec<EdgeEndpointRule> {
self.edge_rules
.read()
.map(|g| g.clone())
.unwrap_or_default()
}
pub async fn embedder(&self) -> RuntimeResult<Arc<dyn EmbeddingService>> {
let model = self
.config
.embedding_model
.ok_or_else(|| crate::RuntimeError::Unconfigured("embedding_model".into()))?;
let service = self
.embedder
.get_or_init(|| async move {
let native = Arc::new(NativeEmbeddingService::with_model(model));
let cached = CachedEmbeddingService::with_default_cache(native);
Arc::new(cached) as Arc<dyn EmbeddingService>
})
.await
.clone();
Ok(service)
}
}
pub(crate) fn vec_model_key(model: EmbeddingModel) -> String {
sanitize_key(&model.to_string())
}
fn sanitize_key(s: &str) -> String {
s.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn memory_runtime_creates_successfully() {
let rt = KhiveRuntime::memory().expect("memory runtime should create");
assert!(rt.config().db_path.is_none());
}
#[test]
fn file_runtime_creates_successfully() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test.db");
let config = RuntimeConfig {
db_path: Some(path.clone()),
default_namespace: Namespace::parse("test").unwrap(),
embedding_model: None,
gate: Arc::new(AllowAllGate),
packs: vec!["kg".to_string()],
backend_id: BackendId::main(),
};
let rt = KhiveRuntime::new(config).expect("file runtime should create");
assert!(path.exists());
assert_eq!(rt.config().default_namespace.as_str(), "test");
}
#[test]
fn from_backend_uses_provided_backend() {
let backend = Arc::new(StorageBackend::memory().expect("memory backend"));
let config = RuntimeConfig {
db_path: None,
default_namespace: Namespace::local(),
embedding_model: None,
gate: Arc::new(AllowAllGate),
packs: vec!["kg".to_string()],
backend_id: BackendId::new("lore"),
};
let rt = KhiveRuntime::from_backend(backend, config);
assert_eq!(rt.backend_id().as_str(), "lore");
assert!(rt.config().db_path.is_none());
}
#[test]
fn backend_id_defaults_to_main() {
let rt = KhiveRuntime::memory().unwrap();
assert_eq!(rt.backend_id().as_str(), BackendId::MAIN);
}
#[test]
fn store_accessors_return_ok() {
let rt = KhiveRuntime::memory().unwrap();
let tok = NamespaceToken::local();
assert!(rt.entities(&tok).is_ok());
assert!(rt.graph(&tok).is_ok());
assert!(rt.notes(&tok).is_ok());
assert!(rt.events(&tok).is_ok());
}
#[test]
fn vectors_returns_unconfigured_without_model() {
let rt = KhiveRuntime::memory().unwrap();
let tok = NamespaceToken::local();
match rt.vectors(&tok) {
Err(crate::RuntimeError::Unconfigured(s)) => assert_eq!(s, "embedding_model"),
Err(other) => panic!("expected Unconfigured, got {:?}", other),
Ok(_) => panic!("expected Err, got Ok"),
}
}
#[test]
fn vec_model_key_sanitizes_dots_and_dashes() {
assert_eq!(
vec_model_key(EmbeddingModel::BgeSmallEnV15),
"bge_small_en_v1_5"
);
assert_eq!(
vec_model_key(EmbeddingModel::BgeBaseEnV15),
"bge_base_en_v1_5"
);
assert_eq!(
vec_model_key(EmbeddingModel::AllMiniLmL6V2),
"all_minilm_l6_v2"
);
}
#[test]
fn default_config_uses_allow_all_gate() {
let cfg = RuntimeConfig::default();
assert_eq!(cfg.default_namespace.as_str(), "local");
let _: GateRef = cfg.gate.clone();
}
#[test]
fn parse_pack_list_handles_comma_and_whitespace() {
assert_eq!(parse_pack_list("kg"), vec!["kg".to_string()]);
assert_eq!(
parse_pack_list("kg,gtd"),
vec!["kg".to_string(), "gtd".to_string()]
);
assert_eq!(
parse_pack_list(" kg , gtd "),
vec!["kg".to_string(), "gtd".to_string()]
);
assert_eq!(
parse_pack_list("kg gtd"),
vec!["kg".to_string(), "gtd".to_string()]
);
assert_eq!(parse_pack_list(",,"), Vec::<String>::new());
assert_eq!(parse_pack_list(""), Vec::<String>::new());
}
#[test]
fn default_config_packs_falls_back_to_kg() {
let prior = std::env::var("KHIVE_PACKS").ok();
unsafe {
std::env::remove_var("KHIVE_PACKS");
}
let cfg = RuntimeConfig::default();
assert_eq!(cfg.packs, vec!["kg".to_string()]);
if let Some(v) = prior {
unsafe {
std::env::set_var("KHIVE_PACKS", v);
}
}
}
#[test]
fn default_config_uses_minilm_when_env_unset() {
let prior = std::env::var("KHIVE_EMBEDDING_MODEL").ok();
unsafe {
std::env::remove_var("KHIVE_EMBEDDING_MODEL");
}
let cfg = RuntimeConfig::default();
assert_eq!(cfg.embedding_model, Some(EmbeddingModel::AllMiniLmL6V2));
if let Some(v) = prior {
unsafe {
std::env::set_var("KHIVE_EMBEDDING_MODEL", v);
}
}
}
}