use crate::context::{assemble_context, ContextOptions};
use crate::error::{BrainError, Result};
use crate::exporter::{BrainExporter, BrainImporter};
use crate::indexer::WorkspaceIndexer;
use crate::query::{QueryOptions, RankedHit};
use crate::storage::Database;
use crate::types::{ContextBundle, Node, SyncStats};
use std::path::{Path, PathBuf};
pub struct Brain {
workspace: PathBuf,
brain_dir: PathBuf,
db: Database,
}
impl Brain {
pub fn create(workspace: impl AsRef<Path>) -> Result<Self> {
let workspace = canonicalize_or_owned(workspace.as_ref())?;
let brain_dir = workspace.join(".brain");
std::fs::create_dir_all(&brain_dir)?;
let db_path = brain_dir.join("db.sqlite");
let db = Database::open(&db_path)?;
let marker = brain_dir.join("workspace.json");
if !marker.exists() {
let meta = serde_json::json!({
"version": 1,
"workspace": workspace.to_string_lossy(),
});
std::fs::write(&marker, serde_json::to_string_pretty(&meta)?)?;
}
Ok(Self {
workspace,
brain_dir,
db,
})
}
pub fn open(workspace: impl AsRef<Path>) -> Result<Self> {
let workspace = canonicalize_or_owned(workspace.as_ref())?;
let brain_dir = workspace.join(".brain");
let db_path = brain_dir.join("db.sqlite");
if !db_path.exists() {
return Err(BrainError::BrainNotFound { path: brain_dir });
}
let db = Database::open(&db_path)?;
Ok(Self {
workspace,
brain_dir,
db,
})
}
pub fn open_or_create(workspace: impl AsRef<Path>) -> Result<Self> {
let workspace = workspace.as_ref();
let db_path = workspace.join(".brain").join("db.sqlite");
if db_path.exists() {
Self::open(workspace)
} else {
Self::create(workspace)
}
}
pub fn workspace(&self) -> &Path {
&self.workspace
}
pub fn brain_dir(&self) -> &Path {
&self.brain_dir
}
pub fn database(&self) -> &Database {
&self.db
}
pub fn sync(&mut self) -> Result<SyncStats> {
let db_path = self.brain_dir.join("db.sqlite");
let db = Database::open(&db_path)?;
let indexer = WorkspaceIndexer::new(db, self.workspace.clone());
let stats = indexer.index_workspace()?;
self.db = Database::open(&db_path)?;
Ok(stats)
}
pub fn query(&self, q: &str) -> Result<Vec<Node>> {
let hits = self.query_ranked(q, &QueryOptions::default())?;
Ok(hits.into_iter().map(|h| h.node).collect())
}
pub fn query_ranked(&self, q: &str, opts: &QueryOptions) -> Result<Vec<RankedHit>> {
self.db.search_ranked(q, opts)
}
pub fn context_for_prompt(&self, prompt: &str, max_tokens: usize) -> Result<ContextBundle> {
let opts = ContextOptions {
max_tokens,
..ContextOptions::default()
};
self.context_for_prompt_with(prompt, &opts)
}
pub fn context_for_prompt_with(
&self,
prompt: &str,
opts: &ContextOptions,
) -> Result<ContextBundle> {
assemble_context(&self.db, &self.brain_dir, prompt, opts)
}
pub fn export(&self, out: impl AsRef<Path>, decouple_ast: bool) -> Result<()> {
BrainExporter::export_bundle(&self.db, out, decouple_ast)
}
pub fn import(&mut self, input: impl AsRef<Path>) -> Result<usize> {
let n = BrainImporter::import_bundle(&self.db, input)?;
#[cfg(feature = "mmap")]
{
let indexer = WorkspaceIndexer::new(
Database::open(self.brain_dir.join("db.sqlite"))?,
self.workspace.clone(),
);
let _ = indexer.compile_mmap(&self.brain_dir.join("graph.mmap"));
self.db = Database::open(self.brain_dir.join("db.sqlite"))?;
}
Ok(n)
}
pub fn watch(&self, debounce_ms: u64) -> Result<()> {
crate::watch::watch_workspace(
&self.workspace,
crate::watch::WatchConfig {
debounce: std::time::Duration::from_millis(debounce_ms),
verbose: true,
},
)
}
}
fn canonicalize_or_owned(path: &Path) -> Result<PathBuf> {
if path.exists() {
Ok(fs_canonicalize(path)?)
} else {
std::fs::create_dir_all(path)?;
Ok(fs_canonicalize(path)?)
}
}
fn fs_canonicalize(path: &Path) -> Result<PathBuf> {
std::fs::canonicalize(path).map_err(BrainError::from)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::ContextRole;
use tempfile::tempdir;
#[test]
fn create_sync_query_context_export() {
let dir = tempdir().unwrap();
let docs = dir.path().join("docs");
std::fs::create_dir_all(&docs).unwrap();
std::fs::write(
docs.join("raft.md"),
"---\ntags: [raft]\nnode_type: concept\n---\n# Raft\nSee [[logcompaction]].\n",
)
.unwrap();
std::fs::write(
docs.join("logcompaction.md"),
"---\ntags: [log]\nnode_type: concept\n---\n# Log Compaction\nSee [[raft]].\n",
)
.unwrap();
let mut brain = Brain::create(dir.path()).unwrap();
let stats = brain.sync().unwrap();
assert_eq!(stats.markdown_files, 2);
let hits = brain.query("raft").unwrap();
assert!(!hits.is_empty());
let ranked = brain
.query_ranked("raft", &QueryOptions::default())
.unwrap();
assert!(ranked[0].score > 0.0);
let ctx = brain.context_for_prompt("raft", 512).unwrap();
assert!(
!ctx.nodes.is_empty(),
"expected FTS hits for 'raft', got none"
);
assert!(ctx.tokens_used > 0);
assert!(ctx.nodes.iter().any(|n| n.role == ContextRole::Seed));
let xml = ctx.to_xml();
assert!(xml.contains("<rustbrain_context"));
assert!(xml.contains("tokens_used="));
let out = dir.path().join("out.brainbundle");
brain.export(&out, true).unwrap();
assert!(out.exists());
let _ = brain.sync().unwrap();
assert_eq!(brain.database().count_fts_rows().unwrap(), 2);
}
#[test]
fn note_anchors_to_symbol() {
let dir = tempdir().unwrap();
let docs = dir.path().join("docs");
let src = dir.path().join("src");
std::fs::create_dir_all(&docs).unwrap();
std::fs::create_dir_all(&src).unwrap();
std::fs::write(
dir.path().join("Cargo.toml"),
"[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
)
.unwrap();
std::fs::write(
src.join("lib.rs"),
"/// Storage engine\npub struct StorageEngine;\nimpl StorageEngine { pub fn open() {} }\n",
)
.unwrap();
std::fs::write(
docs.join("design.md"),
"---\nnode_type: adr\n---\n# Design\nUses symbol:StorageEngine for persistence.\n",
)
.unwrap();
let mut brain = Brain::create(dir.path()).unwrap();
let stats = brain.sync().unwrap();
assert!(stats.symbol_anchors >= 1);
assert!(stats.markdown_files >= 1);
let _ = brain.sync().unwrap();
let edges = brain.database().get_all_edges().unwrap();
let anchors: Vec<_> = edges
.iter()
.filter(|e| e.relation_type == "anchors")
.collect();
assert!(
!anchors.is_empty(),
"expected anchors edge note→symbol, edges={edges:?}"
);
}
}