semtree_rag/

session.rs

use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use semtree_core::{Chunk, ChunkKind, Language};
use semtree_embed::Embedder;
use semtree_store::VectorStore;

use crate::{
    ChunkRegistry, ContextWindow, FileManifest, HybridSearcher, Indexer, LexicalIndex, RagError,
    SearchEngine, SearchFilters, SearchMode,
};

/// Why an index is being rebuilt from scratch instead of updated in place.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RebuildReason {
    /// Nothing on disk yet.
    Missing,
    /// The caller asked for a full rebuild.
    Requested,
    /// The stored vectors came out of a different pipeline, so reusing them
    /// would mix incompatible vectors or leave stale chunk ids behind.
    Incompatible {
        /// Embedder and store the index was built with.
        was: String,
        /// Embedder and store in use now.
        now: String,
    },
}

impl std::fmt::Display for RebuildReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Missing => f.write_str("no existing index"),
            Self::Requested => f.write_str("full rebuild requested"),
            Self::Incompatible { was, now } => {
                write!(f, "index was built with {was}, now running {now}")
            }
        }
    }
}

/// What an [`IndexSession::index`] pass did.
#[derive(Debug, Clone)]
pub struct IndexReport {
    /// Chunks embedded and stored during this pass. Unchanged files contribute
    /// nothing, so an incremental pass over an untouched tree reports zero.
    pub chunks_indexed: usize,
    /// Set when the pass rebuilt from scratch; `None` for an incremental update.
    pub rebuilt: Option<RebuildReason>,
}

impl IndexReport {
    pub fn was_incremental(&self) -> bool {
        self.rebuilt.is_none()
    }
}

/// A breakdown of what an index currently holds.
#[derive(Debug, Clone)]
pub struct IndexStats {
    pub chunks: usize,
    pub files: usize,
    /// Vectors held by the store, when one is open. Should track `chunks`; a
    /// gap means the store and the chunk metadata have drifted apart. `None`
    /// after a metadata-only read, where the store was never loaded.
    pub vectors: Option<usize>,
    /// Chunk counts per language, most frequent first.
    pub by_language: Vec<(Language, usize)>,
    /// Chunk counts per kind, most frequent first.
    pub by_kind: Vec<(ChunkKind, usize)>,
    /// Fingerprint of the embedder that produced the vectors.
    pub embedder: String,
    /// Fingerprint of the store holding them, i.e. its distance metric.
    pub store: String,
}

impl IndexStats {
    /// Summarize the index at `index_dir` from its metadata alone.
    ///
    /// Reads the chunk registry and the manifest, and nothing else: no
    /// embedding model is loaded and no vectors are touched, so inspecting an
    /// index costs milliseconds. The trade-off is that `vectors` stays `None`,
    /// because only the store knows how many it holds.
    pub fn open(index_dir: &Path) -> Result<Self, RagError> {
        let registry = ChunkRegistry::open(index_dir)?;
        let manifest = FileManifest::load(index_dir);
        Ok(Self::summarize(
            &registry,
            manifest.embedder().to_string(),
            manifest.store().to_string(),
            None,
        ))
    }

    fn summarize(
        registry: &ChunkRegistry,
        embedder: String,
        store: String,
        vectors: Option<usize>,
    ) -> Self {
        let mut by_language: HashMap<Language, usize> = HashMap::new();
        let mut by_kind: HashMap<ChunkKind, usize> = HashMap::new();
        let mut files: HashSet<&Path> = HashSet::new();

        for chunk in registry.iter() {
            *by_language.entry(chunk.language).or_default() += 1;
            *by_kind.entry(chunk.kind).or_default() += 1;
            files.insert(chunk.path.as_path());
        }

        Self {
            chunks: registry.len(),
            files: files.len(),
            vectors,
            by_language: sorted_by_count(by_language),
            by_kind: sorted_by_count(by_kind),
            embedder,
            store,
        }
    }
}

/// A ranked hit with the chunk it points at already resolved.
#[derive(Debug, Clone, Copy)]
pub struct SearchResult<'a> {
    pub score: f32,
    pub chunk: &'a Chunk,
}

/// An index on disk together with the backends that query and refresh it.
///
/// This is the whole lifecycle in one place - open, check the index still
/// matches the pipeline that built it, update it incrementally, persist it,
/// search it - so callers do not each reimplement it and drift apart. The
/// [`semtree`] CLI and the MCP server are both thin shells over this type.
///
/// ```no_run
/// use std::sync::Arc;
/// use semtree_rag::{IndexSession, SearchFilters, SearchMode};
/// use semtree_embed::fastembed::FastEmbedder;
/// use semtree_store::usearch::UsearchStore;
/// use semtree_embed::Embedder;
///
/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
/// let embedder = Arc::new(FastEmbedder::new()?);
/// let store = Arc::new(UsearchStore::new(embedder.dimension())?);
///
/// let mut session = IndexSession::open(embedder, store, std::path::Path::new(".semtree"))?;
/// session.index(std::path::Path::new("."), false, |_, _| {}).await?;
/// session.save()?;
///
/// let hits = session
///     .search("where do we skip unchanged files", 5, SearchMode::Hybrid, &SearchFilters::default())
///     .await?;
/// # Ok(())
/// # }
/// ```
///
/// [`semtree`]: https://crates.io/crates/semtree
pub struct IndexSession {
    embedder: Arc<dyn Embedder>,
    store: Arc<dyn VectorStore>,
    registry: ChunkRegistry,
    manifest: FileManifest,
    searcher: HybridSearcher,
    index_dir: PathBuf,
    pending_rebuild: Option<RebuildReason>,
}

impl IndexSession {
    /// Open the index at `index_dir`, or start an empty session if there is
    /// none yet.
    ///
    /// An index built by a different embedder, store, or chunker is *not*
    /// loaded: it is left on disk and flagged for rebuild, so the next
    /// [`index`](Self::index) call replaces it instead of mixing vectors that
    /// cannot be compared. Query it before that and it looks empty, which is
    /// the honest answer.
    pub fn open(
        embedder: Arc<dyn Embedder>,
        store: Arc<dyn VectorStore>,
        index_dir: &Path,
    ) -> Result<Self, RagError> {
        let embedder_fingerprint = embedder.fingerprint();
        let store_fingerprint = store.metric().to_string();

        let mut registry = ChunkRegistry::default();
        let mut manifest = FileManifest::new(&embedder_fingerprint, &store_fingerprint);
        let mut pending_rebuild = Some(RebuildReason::Missing);

        if Self::is_present(index_dir) {
            let existing = FileManifest::load(index_dir);
            if existing.is_compatible_with(&embedder_fingerprint, &store_fingerprint) {
                store.load(index_dir)?;
                registry.load(index_dir)?;
                manifest = existing;
                pending_rebuild = None;
            } else {
                pending_rebuild = Some(RebuildReason::Incompatible {
                    was: format!("{}/{}", existing.embedder(), existing.store()),
                    now: format!("{embedder_fingerprint}/{store_fingerprint}"),
                });
            }
        }

        let searcher = Self::build_searcher(&embedder, &store, &registry);

        Ok(Self {
            embedder,
            store,
            registry,
            manifest,
            searcher,
            index_dir: index_dir.to_path_buf(),
            pending_rebuild,
        })
    }

    /// Like [`open`](Self::open), but refuses to start without a usable index.
    /// Use it on read-only paths, where an empty result would otherwise be
    /// indistinguishable from a codebase that was never indexed.
    pub fn open_existing(
        embedder: Arc<dyn Embedder>,
        store: Arc<dyn VectorStore>,
        index_dir: &Path,
    ) -> Result<Self, RagError> {
        let session = Self::open(embedder, store, index_dir)?;
        match &session.pending_rebuild {
            None => Ok(session),
            Some(RebuildReason::Missing) => Err(RagError::NoIndex(index_dir.to_path_buf())),
            Some(reason) => Err(RagError::Filter(format!(
                "index at {} is unusable: {reason}; re-index to rebuild it",
                index_dir.display()
            ))),
        }
    }

    /// An index counts as present only when both halves of it are on disk: the
    /// manifest alone says nothing about the chunks it describes.
    fn is_present(index_dir: &Path) -> bool {
        index_dir.join("manifest.json").exists() && index_dir.join("chunks.json").exists()
    }

    fn build_searcher(
        embedder: &Arc<dyn Embedder>,
        store: &Arc<dyn VectorStore>,
        registry: &ChunkRegistry,
    ) -> HybridSearcher {
        let engine = SearchEngine::new(embedder.clone(), store.clone());
        HybridSearcher::new(engine, LexicalIndex::from_chunks(registry.iter()))
    }

    /// Why the next [`index`](Self::index) call will rebuild from scratch, if
    /// it will. `None` means the index on disk is usable as-is.
    pub fn pending_rebuild(&self) -> Option<&RebuildReason> {
        self.pending_rebuild.as_ref()
    }

    /// Index `source_root`, skipping files whose content has not changed.
    ///
    /// Set `full` to force a rebuild; otherwise a rebuild still happens when
    /// [`pending_rebuild`](Self::pending_rebuild) says the existing index
    /// cannot be trusted. Progress is reported as `(files_done, files_total)`.
    ///
    /// Nothing is written to disk until [`save`](Self::save).
    pub async fn index(
        &mut self,
        source_root: &Path,
        full: bool,
        on_progress: impl Fn(usize, usize),
    ) -> Result<IndexReport, RagError> {
        // Prefer the recorded reason over the caller's flag: "the embedder
        // changed" is more useful to report back than "you asked for it".
        let rebuilt = self
            .pending_rebuild
            .take()
            .or_else(|| full.then_some(RebuildReason::Requested));

        if rebuilt.is_some() {
            self.store.clear().await?;
            self.registry = ChunkRegistry::default();
            self.manifest =
                FileManifest::new(self.embedder.fingerprint(), self.store.metric().to_string());
        }

        let indexer = Indexer::new(self.embedder.clone(), self.store.clone());
        let chunks_indexed = indexer
            .index_dir(
                source_root,
                &mut self.registry,
                Some(&mut self.manifest),
                on_progress,
            )
            .await?;

        // The lexical index is derived from the registry, so it goes stale the
        // moment indexing touches a chunk.
        self.searcher = Self::build_searcher(&self.embedder, &self.store, &self.registry);

        Ok(IndexReport {
            chunks_indexed,
            rebuilt,
        })
    }

    /// Persist the vectors, the chunk metadata and the manifest together. They
    /// are only consistent as a set, so they are always written as one.
    pub fn save(&self) -> Result<(), RagError> {
        std::fs::create_dir_all(&self.index_dir)?;
        self.store.save(&self.index_dir)?;
        self.registry.save(&self.index_dir)?;
        self.manifest.save(&self.index_dir)?;
        Ok(())
    }

    /// Search the index, returning at most `top_k` results that pass `filters`.
    ///
    /// Filters are metadata narrowing applied after ranking, so the ranker is
    /// asked for extra candidates to compensate. Hits the registry cannot
    /// resolve are dropped: without a chunk there is no path, name or code to
    /// hand back.
    pub async fn search(
        &self,
        query: &str,
        top_k: usize,
        mode: SearchMode,
        filters: &SearchFilters,
    ) -> Result<Vec<SearchResult<'_>>, RagError> {
        let hits = self
            .searcher
            .search(query, filters.fetch_size(top_k), mode)
            .await?;

        Ok(hits
            .iter()
            .filter_map(|hit| {
                self.registry.get(&hit.id).map(|chunk| SearchResult {
                    score: hit.score,
                    chunk,
                })
            })
            .filter(|result| filters.matches(result.chunk))
            .take(top_k)
            .collect())
    }

    /// Build a prompt-ready context window for `query` from the top `top_k`
    /// chunks under `mode`.
    pub async fn context(
        &self,
        query: &str,
        top_k: usize,
        mode: SearchMode,
    ) -> Result<ContextWindow, RagError> {
        let hits = self.searcher.search(query, top_k, mode).await?;
        Ok(ContextWindow::from_hits(query, &hits, &self.registry))
    }

    /// What the index currently holds, including the live vector count.
    pub fn stats(&self) -> IndexStats {
        IndexStats::summarize(
            &self.registry,
            self.embedder.fingerprint(),
            self.store.metric().to_string(),
            Some(self.store.len()),
        )
    }

    /// Where this session reads and writes its index.
    pub fn index_dir(&self) -> &Path {
        &self.index_dir
    }

    /// The chunk metadata backing the index.
    pub fn registry(&self) -> &ChunkRegistry {
        &self.registry
    }
}

/// Most frequent first, then by key so equal counts do not shuffle between runs.
fn sorted_by_count<K: Ord + Copy>(counts: HashMap<K, usize>) -> Vec<(K, usize)> {
    let mut sorted: Vec<(K, usize)> = counts.into_iter().collect();
    sorted
        .sort_by(|(a_key, a_count), (b_key, b_count)| b_count.cmp(a_count).then(a_key.cmp(b_key)));
    sorted
}