refactor: rebuild desktop app with Tauri

This commit is contained in:
leookun
2026-08-24 02:49:00 +08:00
parent a3ec2a0dfc
commit 4053a7fb20
808 changed files with 87359 additions and 181491 deletions
+65
View File
@@ -0,0 +1,65 @@
//! On-disk paths derived from source identity and index configuration.
use std::path::{Path, PathBuf};
use sha2::{Digest, Sha256};
use crate::{config::INDEX_FORMAT_VERSION, ContentType};
#[derive(Clone, Debug)]
pub struct CacheLayout {
pub snapshot: PathBuf,
pub runtime: PathBuf,
}
impl CacheLayout {
pub fn new(
root: &Path,
source_identity: &str,
model: &str,
content: &[ContentType],
chunk_bytes: usize,
) -> Self {
let source = hex::encode(Sha256::digest(source_identity.as_bytes()));
let mut signature = format!("{model}:{chunk_bytes}:");
let mut content = content.to_vec();
content.sort_by_key(|item| *item as u8);
for item in content {
signature.push_str(&format!("{item:?},"));
}
let signature = hex::encode(Sha256::digest(signature.as_bytes()));
let directory = root
.join(format!("indexes/v{INDEX_FORMAT_VERSION}"))
.join(source)
.join(signature);
Self {
snapshot: directory.join("index.bin"),
runtime: directory.join("runtime.bin"),
}
}
}
#[cfg(test)]
mod tests {
use std::ffi::OsStr;
use super::*;
#[test]
fn snapshot_path_tracks_the_index_format_version() {
let layout = CacheLayout::new(
Path::new("cache"),
"source",
"model",
&[ContentType::Code],
750,
);
let version = format!("v{INDEX_FORMAT_VERSION}");
assert!(layout
.snapshot
.components()
.any(|component| component.as_os_str() == OsStr::new(&version)));
assert_eq!(layout.runtime.file_name(), Some(OsStr::new("runtime.bin")));
assert_eq!(layout.snapshot.parent(), layout.runtime.parent());
}
}
+7
View File
@@ -0,0 +1,7 @@
//! Stable cache keys plus atomic, versioned index snapshot persistence.
mod layout;
mod store;
pub use layout::CacheLayout;
pub use store::{load_runtime, load_snapshot, save_runtime, save_snapshot};
+144
View File
@@ -0,0 +1,144 @@
//! Corruption-resistant bincode snapshot loading and atomic replacement.
use std::{fs, io::Write, path::Path};
use crate::{
index::{IndexSnapshot, LoadedIndex},
Error, Result,
};
pub fn load_snapshot(path: &Path) -> Result<Option<IndexSnapshot>> {
let Some(snapshot) = load::<IndexSnapshot>(path)? else {
return Ok(None);
};
snapshot.validate()?;
Ok(Some(snapshot))
}
pub fn save_snapshot(path: &Path, snapshot: &IndexSnapshot) -> Result<()> {
snapshot.validate()?;
save(path, snapshot)
}
pub fn load_runtime(path: &Path) -> Result<Option<LoadedIndex>> {
let Some(runtime) = load::<LoadedIndex>(path)? else {
return Ok(None);
};
runtime.validate()?;
Ok(Some(runtime))
}
pub fn save_runtime(path: &Path, runtime: &LoadedIndex) -> Result<()> {
runtime.validate()?;
save(path, runtime)
}
fn load<T>(path: &Path) -> Result<Option<T>>
where
T: serde::de::DeserializeOwned,
{
if !path.exists() {
return Ok(None);
}
let bytes = fs::read(path).map_err(|error| Error::io(path, error))?;
let (value, consumed): (T, usize) =
bincode::serde::decode_from_slice(&bytes, bincode::config::standard())
.map_err(|error| Error::CorruptIndex(error.to_string()))?;
if consumed != bytes.len() {
return Err(Error::CorruptIndex("cache has trailing bytes".into()));
}
Ok(Some(value))
}
fn save<T>(path: &Path, value: &T) -> Result<()>
where
T: serde::Serialize,
{
let parent = path
.parent()
.ok_or_else(|| Error::InvalidRequest("snapshot path has no parent".into()))?;
fs::create_dir_all(parent).map_err(|error| Error::io(parent, error))?;
let temporary = parent.join(format!(".cache-{}-{}.tmp", std::process::id(), now_nanos()));
let bytes = bincode::serde::encode_to_vec(value, bincode::config::standard())
.map_err(|error| Error::Serialization(error.to_string()))?;
let mut file = fs::File::create(&temporary).map_err(|error| Error::io(&temporary, error))?;
file.write_all(&bytes)
.and_then(|_| file.sync_all())
.map_err(|error| Error::io(&temporary, error))?;
if path.exists() {
fs::remove_file(path).map_err(|error| Error::io(path, error))?;
}
fs::rename(&temporary, path).map_err(|error| Error::io(path, error))
}
fn now_nanos() -> u128 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |value| value.as_nanos())
}
#[cfg(test)]
mod tests {
use crate::{
config::INDEX_FORMAT_VERSION, index::IndexedFile, source::FileStamp, Chunk, ContentType,
};
use super::*;
fn snapshot() -> IndexSnapshot {
IndexSnapshot {
format_version: INDEX_FORMAT_VERSION,
source_identity: "source".into(),
model_id: "test-model".into(),
dimensions: 2,
desired_chunk_bytes: 128,
content: vec![ContentType::Code],
files: vec![IndexedFile {
path: "src/lib.rs".into(),
stamp: FileStamp {
modified_ns: 1,
size: 4,
},
content_type: ContentType::Code,
chunks: vec![Chunk {
file_path: "src/lib.rs".into(),
start_line: 1,
end_line: 1,
language: Some("rust".into()),
content: "code".into(),
}],
definitions: Vec::new(),
lexical_documents: vec![crate::index::lexical_document("code")],
vectors: vec![127, 0],
}],
}
}
#[test]
fn snapshots_round_trip_and_reject_trailing_corruption() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("nested/index.bin");
save_snapshot(&path, &snapshot()).unwrap();
let loaded = load_snapshot(&path).unwrap().unwrap();
assert!(loaded.files[0].chunks[0].content.is_empty());
let mut bytes = fs::read(&path).unwrap();
bytes.push(0xff);
fs::write(&path, bytes).unwrap();
assert!(matches!(load_snapshot(&path), Err(Error::CorruptIndex(_))));
}
#[test]
fn runtime_indexes_round_trip() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("nested/runtime.bin");
let runtime = LoadedIndex::from_snapshot(snapshot()).unwrap();
save_runtime(&path, &runtime).unwrap();
let loaded = load_runtime(&path).unwrap().unwrap();
assert_eq!(loaded.chunks.len(), 1);
assert_eq!(loaded.vectors, vec![127, 0]);
assert!(loaded
.lexical
.exact_symbol("missing", &loaded.chunks, 1)
.is_empty());
}
}
+25
View File
@@ -0,0 +1,25 @@
//! Byte boundaries used while partitioning source files into retrievable chunks.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ChunkBoundary {
pub start: usize,
pub end: usize,
}
pub(crate) fn merge_adjacent(boundaries: Vec<ChunkBoundary>, desired: usize) -> Vec<ChunkBoundary> {
let mut input = boundaries.into_iter();
let Some(mut current) = input.next() else {
return Vec::new();
};
let mut output = Vec::new();
for next in input {
if next.end.saturating_sub(current.start) <= desired {
current.end = next.end;
} else {
output.push(current);
current = next;
}
}
output.push(current);
output
}
@@ -0,0 +1,19 @@
//! Language-agnostic chunking for unsupported or unparsable source files.
use super::{boundary::merge_adjacent, ChunkBoundary};
pub fn line_boundaries(source: &str, desired: usize) -> Vec<ChunkBoundary> {
if source.trim().is_empty() {
return Vec::new();
}
let mut offset = 0;
let boundaries = source
.split_inclusive('\n')
.map(|line| {
let start = offset;
offset += line.len();
ChunkBoundary { start, end: offset }
})
.collect();
merge_adjacent(boundaries, desired)
}
+9
View File
@@ -0,0 +1,9 @@
//! Code-aware Tree-sitter chunking with a deterministic line-based fallback.
mod boundary;
mod line_fallback;
mod tree_sitter;
pub use boundary::ChunkBoundary;
pub use line_fallback::line_boundaries;
pub use tree_sitter::chunk_source;
+115
View File
@@ -0,0 +1,115 @@
//! Recursive Tree-sitter boundary selection and public chunk construction.
use tree_sitter::Node;
use crate::{language::parser_for, types::Chunk};
use super::{boundary::merge_adjacent, line_boundaries, ChunkBoundary};
const MAX_DEPTH: usize = 500;
const MIN_CHUNK_BYTES: usize = 50;
pub fn chunk_source(
source: &str,
file_path: &str,
language: Option<&str>,
desired: usize,
) -> Vec<Chunk> {
if source.trim().is_empty() {
return Vec::new();
}
let boundaries = language
.and_then(parser_for)
.and_then(|mut parser| parser.parse(source, None))
.map(|tree| merge_adjacent(split_node(tree.root_node(), desired, 0), desired))
.filter(|items| !items.is_empty())
.unwrap_or_else(|| line_boundaries(source, desired));
boundaries
.into_iter()
.filter_map(|boundary| {
let content = source.get(boundary.start..boundary.end)?.to_owned();
let start_line = source[..boundary.start]
.bytes()
.filter(|byte| *byte == b'\n')
.count()
+ 1;
let end_line = start_line + content.lines().count().saturating_sub(1);
Some(Chunk {
file_path: file_path.to_owned(),
start_line,
end_line,
language: language.map(str::to_owned),
content,
})
})
.collect()
}
fn split_node(node: Node<'_>, desired: usize, depth: usize) -> Vec<ChunkBoundary> {
if node.child_count() == 0
|| depth > MAX_DEPTH
|| node.end_byte().saturating_sub(node.start_byte()) < MIN_CHUNK_BYTES
{
return vec![ChunkBoundary {
start: node.start_byte(),
end: node.end_byte(),
}];
}
let mut cursor = node.walk();
let children = node.children(&mut cursor).collect::<Vec<_>>();
let mut output = Vec::new();
let mut index = 0;
while index < children.len() {
let child = children[index];
index += 1;
if child.end_byte().saturating_sub(child.start_byte()) > desired {
output.extend(split_node(child, desired, depth + 1));
continue;
}
let start = child.start_byte();
let mut end = child.end_byte();
while index < children.len() && children[index].end_byte().saturating_sub(start) <= desired
{
end = children[index].end_byte();
index += 1;
}
output.push(ChunkBoundary { start, end });
}
output
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rust_chunks_keep_paths_lines_and_source_text() {
let source =
"fn first() {\n println!(\"one\");\n}\n\nfn second() {\n println!(\"two\");\n}\n";
let chunks = chunk_source(source, "src/lib.rs", Some("rust"), 48);
assert!(chunks.len() >= 2);
assert_eq!(chunks[0].file_path, "src/lib.rs");
assert_eq!(chunks[0].start_line, 1);
assert!(chunks.iter().any(|chunk| chunk.content.contains("second")));
assert!(chunks.iter().all(|chunk| !chunk.content.is_empty()));
}
#[test]
fn unsupported_languages_fall_back_to_line_chunks() {
let source = "alpha beta gamma\ndelta epsilon\nzeta eta theta\n";
let chunks = chunk_source(source, "notes.unknown", None, 24);
assert!(chunks.len() >= 2);
assert_eq!(chunks[0].start_line, 1);
assert_eq!(chunks.last().unwrap().end_line, 3);
}
#[test]
fn multibyte_character_at_a_chunk_boundary_does_not_panic() {
let source = "// 中文说明。";
let chunks = chunk_source(source, "src/lib.rs", Some("rust"), 750);
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].content, source);
assert_eq!(chunks[0].start_line, 1);
assert_eq!(chunks[0].end_line, 1);
}
}
+35
View File
@@ -0,0 +1,35 @@
//! User-independent configuration and fixed index format parameters.
use std::path::PathBuf;
/// Current persisted snapshot format. Incompatible changes must bump this value.
pub const INDEX_FORMAT_VERSION: u32 = 7;
/// Search and persistence settings shared by every indexed repository.
#[derive(Clone, Debug)]
pub struct SembleConfig {
pub cache_dir: PathBuf,
pub desired_chunk_bytes: usize,
pub max_file_bytes: u64,
}
impl SembleConfig {
pub fn new(cache_dir: impl Into<PathBuf>) -> Self {
Self {
cache_dir: cache_dir.into(),
desired_chunk_bytes: 750,
max_file_bytes: 1_000_000,
}
}
}
impl Default for SembleConfig {
fn default() -> Self {
let root = std::env::var_os("SEMBLE_CACHE_LOCATION")
.map(PathBuf::from)
.filter(|path| path.is_absolute())
.or_else(|| dirs::home_dir().map(|home| home.join(".cursor-byok-v3/cache/semble")))
.unwrap_or_else(|| PathBuf::from(".cursor-byok-v3/cache/semble"));
Self::new(root)
}
}
@@ -0,0 +1,87 @@
//! Verified, atomic download and reuse of the tokenizer and embedding matrix.
use std::{
fs,
io::{Read, Write},
path::{Path, PathBuf},
};
use sha2::{Digest, Sha256};
use crate::{Error, Result};
const MODEL_URL: &str =
"https://huggingface.co/minishlab/potion-code-16M-v2/resolve/main/model.safetensors";
const TOKENIZER_URL: &str =
"https://huggingface.co/minishlab/potion-code-16M-v2/resolve/main/tokenizer.json";
const MODEL_SHA256: &str = "75cf7a6c2171b230ad19b1e7d8e0b1aee86da5a02af8e7cacedd9921d227623c";
const TOKENIZER_SHA256: &str = "107bbdcbad4bff1d299b7a4c3a2fb17c52890688b7dd0e4c9deab79d3c4f3d45";
#[derive(Clone, Debug)]
pub struct ModelAssets {
pub model: PathBuf,
pub tokenizer: PathBuf,
}
impl ModelAssets {
pub fn model_path(cache_root: &Path) -> PathBuf {
cache_root.join("models/potion-code-16M-v2/model.safetensors")
}
pub fn ensure(cache_root: &Path) -> Result<Self> {
let directory = cache_root.join("models/potion-code-16M-v2");
fs::create_dir_all(&directory).map_err(|error| Error::io(&directory, error))?;
let model = Self::model_path(cache_root);
let tokenizer = directory.join("tokenizer.json");
ensure_asset(&model, MODEL_URL, MODEL_SHA256)?;
ensure_asset(&tokenizer, TOKENIZER_URL, TOKENIZER_SHA256)?;
Ok(Self { model, tokenizer })
}
}
fn ensure_asset(path: &Path, url: &str, expected: &str) -> Result<()> {
if path.is_file() && digest(path)? == expected {
return Ok(());
}
let temporary = path.with_extension(format!("tmp-{}", std::process::id()));
let response = reqwest::blocking::Client::builder()
.build()
.map_err(|error| Error::ModelAsset(error.to_string()))?
.get(url)
.send()
.and_then(reqwest::blocking::Response::error_for_status)
.map_err(|error| Error::ModelAsset(format!("download {url}: {error}")))?;
let bytes = response
.bytes()
.map_err(|error| Error::ModelAsset(error.to_string()))?;
let actual = hex::encode(Sha256::digest(&bytes));
if actual != expected {
return Err(Error::ModelAsset(format!(
"checksum mismatch for {url}: expected {expected}, got {actual}"
)));
}
let mut file = fs::File::create(&temporary).map_err(|error| Error::io(&temporary, error))?;
file.write_all(&bytes)
.and_then(|_| file.sync_all())
.map_err(|error| Error::io(&temporary, error))?;
if path.exists() {
fs::remove_file(path).map_err(|error| Error::io(path, error))?;
}
fs::rename(&temporary, path).map_err(|error| Error::io(path, error))
}
fn digest(path: &Path) -> Result<String> {
let mut file = fs::File::open(path).map_err(|error| Error::io(path, error))?;
let mut hash = Sha256::new();
let mut buffer = [0_u8; 64 * 1024];
loop {
let read = file
.read(&mut buffer)
.map_err(|error| Error::io(path, error))?;
if read == 0 {
break;
}
hash.update(&buffer[..read]);
}
Ok(hex::encode(hash.finalize()))
}
+7
View File
@@ -0,0 +1,7 @@
//! Model2Vec-compatible static token embeddings and managed model assets.
mod assets;
mod model;
pub use assets::ModelAssets;
pub use model::{Embedder, StaticEmbedder};
+150
View File
@@ -0,0 +1,150 @@
//! Safetensors-backed mean-pooled and L2-normalized static embedding inference.
use std::{fs, path::Path};
use half::f16;
use rayon::prelude::*;
use safetensors::{tensor::Dtype, SafeTensors};
use tokenizers::Tokenizer;
use crate::{Error, Result};
pub trait Embedder: Send + Sync {
fn id(&self) -> &str;
fn dimensions(&self) -> usize;
fn encode(&self, texts: &[String]) -> Result<Vec<Vec<f32>>>;
}
pub struct StaticEmbedder {
tokenizer: Tokenizer,
embeddings: Vec<f16>,
rows: usize,
dimensions: usize,
}
impl StaticEmbedder {
pub fn load(model: &Path, tokenizer: &Path) -> Result<Self> {
let bytes = fs::read(model).map_err(|error| Error::io(model, error))?;
let tensors =
SafeTensors::deserialize(&bytes).map_err(|error| Error::Model(error.to_string()))?;
let tensor = tensors
.tensor("embeddings")
.map_err(|error| Error::Model(error.to_string()))?;
if tensor.dtype() != Dtype::F16 || tensor.shape().len() != 2 {
return Err(Error::Model(
"embeddings must be a rank-2 F16 tensor".into(),
));
}
let rows = tensor.shape()[0];
let dimensions = tensor.shape()[1];
let embeddings = tensor
.data()
.chunks_exact(2)
.map(|value| f16::from_le_bytes([value[0], value[1]]))
.collect();
let tokenizer =
Tokenizer::from_file(tokenizer).map_err(|error| Error::Model(error.to_string()))?;
Ok(Self {
tokenizer,
embeddings,
rows,
dimensions,
})
}
fn encode_one(&self, text: &str) -> Result<Vec<f32>> {
let encoding = self
.tokenizer
.encode(text, false)
.map_err(|error| Error::Model(error.to_string()))?;
let ids = encoding
.get_ids()
.iter()
.copied()
.filter(|id| *id != 1)
.take(512)
.collect::<Vec<_>>();
let mut output = vec![0.0; self.dimensions];
let mut count = 0_f32;
for id in ids {
let row = id as usize;
if row >= self.rows {
continue;
}
let start = row * self.dimensions;
for (slot, value) in output
.iter_mut()
.zip(&self.embeddings[start..start + self.dimensions])
{
*slot += value.to_f32();
}
count += 1.0;
}
if count > 0.0 {
for value in &mut output {
*value /= count;
}
}
normalize(&mut output);
Ok(output)
}
}
impl Embedder for StaticEmbedder {
fn id(&self) -> &str {
"minishlab/potion-code-16M-v2"
}
fn dimensions(&self) -> usize {
self.dimensions
}
fn encode(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
texts.par_iter().map(|text| self.encode_one(text)).collect()
}
}
pub(crate) fn normalize(vector: &mut [f32]) {
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
if norm > f32::EPSILON {
for value in vector {
*value /= norm;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn normalization_produces_unit_vectors_and_keeps_zero_stable() {
let mut vector = [3.0, 4.0];
normalize(&mut vector);
assert!((vector[0] - 0.6).abs() < 0.0001);
assert!((vector[1] - 0.8).abs() < 0.0001);
let mut zero = [0.0, 0.0];
normalize(&mut zero);
assert_eq!(zero, [0.0, 0.0]);
}
#[test]
fn loads_and_encodes_the_pinned_model_fixture_when_available() {
let (Some(model), Some(tokenizer)) = (
std::env::var_os("SEMBLE_TEST_MODEL"),
std::env::var_os("SEMBLE_TEST_TOKENIZER"),
) else {
return;
};
let model = StaticEmbedder::load(Path::new(&model), Path::new(&tokenizer)).unwrap();
assert_eq!(model.dimensions(), 256);
let vectors = model
.encode(&["parse an HTTP response".into(), "create an invoice".into()])
.unwrap();
assert_eq!(vectors.len(), 2);
assert!(vectors.iter().all(|vector| vector.len() == 256));
assert!(vectors.iter().all(|vector| {
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
(norm - 1.0).abs() < 0.001
}));
assert_ne!(vectors[0], vectors[1]);
}
}
+48
View File
@@ -0,0 +1,48 @@
//! Errors shared by indexing, persistence, model loading, and retrieval.
use std::path::PathBuf;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("source does not exist: {0}")]
SourceMissing(PathBuf),
#[error("source is not a directory: {0}")]
SourceNotDirectory(PathBuf),
#[error("source path escapes the allowed root: {0}")]
UnsafePath(PathBuf),
#[error("unsupported repository URL: {0}")]
UnsupportedUrl(String),
#[error("git operation failed: {0}")]
Git(String),
#[error("model asset error: {0}")]
ModelAsset(String),
#[error("model inference error: {0}")]
Model(String),
#[error("index is empty: {0}")]
EmptyIndex(PathBuf),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("persisted index is incompatible or corrupt: {0}")]
CorruptIndex(String),
#[error("I/O error at {path}: {source}")]
Io {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error(transparent)]
Json(#[from] serde_json::Error),
#[error("serialization error: {0}")]
Serialization(String),
}
impl Error {
pub(crate) fn io(path: impl Into<PathBuf>, source: std::io::Error) -> Self {
Self::Io {
path: path.into(),
source,
}
}
}
pub type Result<T> = std::result::Result<T, Error>;
+219
View File
@@ -0,0 +1,219 @@
//! Incremental BM25 document data and a compact runtime inverted index.
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use super::{IndexedFile, IndexedLexicalDocument, IndexedTermFrequency};
use crate::index::tokenize;
const K1: f32 = 1.5;
const B: f32 = 0.75;
#[derive(Clone, Copy, Deserialize, Serialize)]
struct Posting {
document: usize,
frequency: u32,
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) struct Bm25Match {
pub document: usize,
pub score: f32,
pub coverage: f32,
}
#[derive(Deserialize, Serialize)]
pub struct Bm25Index {
document_lengths: Vec<u32>,
average_document_length: f32,
postings: HashMap<String, Vec<Posting>>,
}
impl Bm25Index {
pub fn from_files(files: &mut [IndexedFile]) -> Self {
let document_count = files
.iter()
.map(|file| file.lexical_documents.len())
.sum::<usize>();
let mut document_lengths = Vec::with_capacity(document_count);
let mut postings = HashMap::<String, Vec<Posting>>::new();
for file in files {
for document in std::mem::take(&mut file.lexical_documents) {
let position = document_lengths.len();
document_lengths.push(document.length);
for term in document.terms {
let posting = Posting {
document: position,
frequency: term.frequency,
};
postings.entry(term.term).or_default().push(posting);
}
}
}
let total_length = document_lengths
.iter()
.map(|length| u64::from(*length))
.sum::<u64>();
let average_document_length = if document_lengths.is_empty() {
0.0
} else {
total_length as f32 / document_lengths.len() as f32
};
Self {
document_lengths,
average_document_length,
postings,
}
}
pub(crate) fn search(&self, query: &str, limit: usize) -> Vec<Bm25Match> {
if limit == 0 || self.document_lengths.is_empty() {
return Vec::new();
}
let mut query_terms = HashMap::<String, u32>::new();
for term in tokenize(query) {
*query_terms.entry(term).or_default() += 1;
}
let corpus_size = self.document_lengths.len() as f32;
let average_length = self.average_document_length.max(f32::EPSILON);
let mut scores = HashMap::<usize, f32>::new();
let mut matched_weight = HashMap::<usize, f32>::new();
let mut query_weight = 0.0;
for (term, query_frequency) in query_terms {
let postings = self.postings.get(&term);
let document_frequency = postings.map_or(0.0, |values| values.len() as f32);
let inverse_document_frequency =
(1.0 + (corpus_size - document_frequency + 0.5) / (document_frequency + 0.5)).ln();
query_weight += inverse_document_frequency;
let Some(postings) = postings else {
continue;
};
for posting in postings {
let term_frequency = posting.frequency as f32;
let document_length = self.document_lengths[posting.document] as f32;
let normalized_frequency = term_frequency
/ (K1 * (1.0 - B + B * document_length / average_length) + term_frequency);
*scores.entry(posting.document).or_default() +=
query_frequency as f32 * inverse_document_frequency * normalized_frequency;
*matched_weight.entry(posting.document).or_default() += inverse_document_frequency;
}
}
let mut ranked = scores
.into_iter()
.filter(|(_, score)| *score > 0.0)
.map(|(document, score)| Bm25Match {
document,
score,
coverage: matched_weight.get(&document).copied().unwrap_or_default()
/ query_weight.max(f32::EPSILON),
})
.collect::<Vec<_>>();
ranked.sort_by(|left, right| {
right
.score
.total_cmp(&left.score)
.then_with(|| left.document.cmp(&right.document))
});
ranked.truncate(limit.min(ranked.len()));
ranked
}
pub(crate) fn validate(&self, chunks: usize) -> bool {
self.document_lengths.len() == chunks
&& self
.postings
.values()
.flatten()
.all(|posting| posting.document < chunks && posting.frequency > 0)
}
}
pub fn lexical_document(content: &str) -> IndexedLexicalDocument {
let tokens = tokenize(content);
let length = u32::try_from(tokens.len()).unwrap_or(u32::MAX);
let mut counts = HashMap::<String, u32>::new();
for token in tokens {
*counts.entry(token).or_default() += 1;
}
let mut terms = counts
.into_iter()
.map(|(term, frequency)| IndexedTermFrequency { term, frequency })
.collect::<Vec<_>>();
terms.sort_by(|left, right| left.term.cmp(&right.term));
IndexedLexicalDocument { length, terms }
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{source::FileStamp, Chunk, ContentType};
fn indexed_file(contents: &[&str]) -> IndexedFile {
let chunks = contents
.iter()
.enumerate()
.map(|(index, content)| Chunk {
file_path: "src/lib.rs".into(),
start_line: index + 1,
end_line: index + 1,
language: Some("rust".into()),
content: (*content).into(),
})
.collect::<Vec<_>>();
IndexedFile {
path: "src/lib.rs".into(),
stamp: FileStamp {
modified_ns: 1,
size: 1,
},
content_type: ContentType::Code,
lexical_documents: chunks
.iter()
.map(|chunk| lexical_document(&chunk.content))
.collect(),
chunks,
definitions: Vec::new(),
vectors: Vec::new(),
}
}
#[test]
fn ranks_identifier_aware_lexical_matches_and_excludes_zero_scores() {
let file = indexed_file(&[
"fn parseConfig() { parse_config(); }",
"fn authenticate() { verify_token(); }",
"fn parse_document() {}",
]);
let index = Bm25Index::from_files(&mut [file]);
let matches = index.search("parse config", 3);
assert_eq!(
matches
.iter()
.map(|matched| matched.document)
.collect::<Vec<_>>(),
vec![0, 2]
);
assert_eq!(matches[0].coverage, 1.0);
assert!(matches[1].coverage < matches[0].coverage);
assert!(index.search("missing_identifier", 3).is_empty());
}
#[test]
fn lexical_documents_store_sorted_term_frequencies() {
let document = lexical_document("parseConfig parse_config");
assert_eq!(document.length, 6);
assert!(document
.terms
.windows(2)
.all(|terms| terms[0].term < terms[1].term));
assert_eq!(
document
.terms
.iter()
.find(|term| term.term == "parse")
.map(|term| term.frequency),
Some(2)
);
}
}
+407
View File
@@ -0,0 +1,407 @@
//! Incremental index builder that reuses unchanged chunks and embeddings.
use std::{
collections::HashMap,
fs,
num::NonZeroUsize,
path::Path,
sync::Arc,
time::{Duration, Instant},
};
use lru::LruCache;
use parking_lot::Mutex;
use crate::{
cache::{load_runtime, load_snapshot, save_runtime, save_snapshot, CacheLayout},
chunk::chunk_source,
config::{SembleConfig, INDEX_FORMAT_VERSION},
embedding::Embedder,
index::extract_definitions,
language::detect_language,
source::discover_files,
ContentType, Error, Result,
};
use super::{lexical_document, IndexSnapshot, IndexedFile, LoadedIndex};
pub struct IndexRepository {
config: SembleConfig,
embedder: Arc<dyn Embedder>,
build_lock: Mutex<()>,
loaded: Mutex<LruCache<std::path::PathBuf, CachedIndex>>,
}
const SEARCH_REFRESH_INTERVAL: Duration = Duration::from_secs(1);
struct CachedIndex {
index: Arc<LoadedIndex>,
checked_at: Instant,
}
impl IndexRepository {
pub fn new(config: SembleConfig, embedder: Arc<dyn Embedder>) -> Self {
Self {
config,
embedder,
build_lock: Mutex::new(()),
loaded: Mutex::new(LruCache::new(
NonZeroUsize::new(10).expect("non-zero cache size"),
)),
}
}
pub fn load_or_build(
&self,
root: &Path,
identity: &str,
content: &[ContentType],
) -> Result<Arc<LoadedIndex>> {
let layout = CacheLayout::new(
&self.config.cache_dir,
identity,
self.embedder.id(),
content,
self.config.desired_chunk_bytes,
);
let files = discover_files(root, content, self.config.max_file_bytes)?;
if files.is_empty() {
return Err(Error::EmptyIndex(root.to_path_buf()));
}
if let Some(index) = self.checked_if_unchanged(&layout.snapshot, &files) {
return Ok(index);
}
let _guard = self.build_lock.lock();
if let Some(index) = self.checked_if_unchanged(&layout.snapshot, &files) {
return Ok(index);
}
let runtime = load_runtime(&layout.runtime)
.ok()
.flatten()
.filter(|runtime| {
runtime.metadata.source_identity == identity
&& runtime.metadata.model_id == self.embedder.id()
&& runtime.metadata.desired_chunk_bytes == self.config.desired_chunk_bytes
&& runtime.metadata.content == content
&& unchanged_loaded(&files, &runtime.metadata.files)
});
if let Some(runtime) = runtime {
let loaded = Arc::new(runtime);
self.store_loaded(layout.snapshot, loaded.clone());
return Ok(loaded);
}
let previous = load_snapshot(&layout.snapshot)
.ok()
.flatten()
.filter(|snapshot| {
snapshot.source_identity == identity
&& snapshot.model_id == self.embedder.id()
&& snapshot.desired_chunk_bytes == self.config.desired_chunk_bytes
&& snapshot.content == content
});
let snapshot_is_current = previous.as_ref().is_some_and(|snapshot| {
let stamps = snapshot
.files
.iter()
.map(|file| (file.path.as_str(), file.stamp))
.collect::<HashMap<_, _>>();
files.len() == snapshot.files.len()
&& files
.iter()
.all(|file| stamps.get(file.relative_path.as_str()) == Some(&file.stamp))
});
if snapshot_is_current {
let snapshot = previous.expect("current snapshot exists");
let loaded = Arc::new(LoadedIndex::from_snapshot(snapshot)?);
save_runtime(&layout.runtime, &loaded)?;
self.store_loaded(layout.snapshot, loaded.clone());
return Ok(loaded);
}
let mut old = previous
.map(|snapshot| {
snapshot
.files
.into_iter()
.map(|file| (file.path.clone(), file))
.collect::<HashMap<_, _>>()
})
.unwrap_or_default();
let mut indexed = Vec::new();
for file in files {
if let Some(cached) = old
.remove(&file.relative_path)
.filter(|cached| cached.stamp == file.stamp)
{
indexed.push(cached);
continue;
}
let bytes = fs::read(&file.absolute_path)
.map_err(|error| Error::io(&file.absolute_path, error))?;
if bytes.iter().take(8192).any(|byte| *byte == 0) {
continue;
}
let source = String::from_utf8_lossy(&bytes);
let language = detect_language(&file.absolute_path);
let chunks = chunk_source(
&source,
&file.relative_path,
language,
self.config.desired_chunk_bytes,
);
if chunks.is_empty() {
continue;
}
let texts = chunks
.iter()
.map(|chunk| chunk.content.clone())
.collect::<Vec<_>>();
let vectors =
quantize_vectors(self.embedder.encode(&texts)?, self.embedder.dimensions())?;
let lexical_documents = chunks
.iter()
.map(|chunk| lexical_document(&chunk.content))
.collect();
indexed.push(IndexedFile {
path: file.relative_path,
stamp: file.stamp,
content_type: file.content_type,
definitions: extract_definitions(&chunks),
lexical_documents,
chunks,
vectors,
});
}
if indexed.is_empty() {
return Err(Error::EmptyIndex(root.to_path_buf()));
}
let snapshot = IndexSnapshot {
format_version: INDEX_FORMAT_VERSION,
source_identity: identity.to_owned(),
model_id: self.embedder.id().to_owned(),
dimensions: self.embedder.dimensions(),
desired_chunk_bytes: self.config.desired_chunk_bytes,
content: content.to_vec(),
files: indexed,
};
save_snapshot(&layout.snapshot, &snapshot)?;
let loaded = Arc::new(LoadedIndex::from_snapshot(snapshot)?);
save_runtime(&layout.runtime, &loaded)?;
self.store_loaded(layout.snapshot, loaded.clone());
Ok(loaded)
}
/// Returns a recently checked in-memory index, refreshing stale entries first.
pub fn load_for_search(
&self,
root: &Path,
identity: &str,
content: &[ContentType],
) -> Result<Arc<LoadedIndex>> {
let layout = CacheLayout::new(
&self.config.cache_dir,
identity,
self.embedder.id(),
content,
self.config.desired_chunk_bytes,
);
if let Some(index) = self
.loaded
.lock()
.get(&layout.snapshot)
.filter(|cached| cached.checked_at.elapsed() < SEARCH_REFRESH_INTERVAL)
.map(|cached| cached.index.clone())
{
return Ok(index);
}
self.load_or_build(root, identity, content)
}
fn store_loaded(&self, snapshot: std::path::PathBuf, index: Arc<LoadedIndex>) {
self.loaded.lock().put(
snapshot,
CachedIndex {
index,
checked_at: Instant::now(),
},
);
}
fn checked_if_unchanged(
&self,
snapshot: &Path,
files: &[crate::source::SourceFile],
) -> Option<Arc<LoadedIndex>> {
let mut loaded = self.loaded.lock();
let cached = loaded.get_mut(snapshot)?;
if unchanged_loaded(files, &cached.index.metadata.files) {
cached.checked_at = Instant::now();
Some(cached.index.clone())
} else {
None
}
}
pub(crate) fn encode(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
self.embedder.encode(texts)
}
}
fn quantize_vectors(vectors: Vec<Vec<f32>>, dimensions: usize) -> Result<Vec<i8>> {
if vectors.iter().any(|vector| vector.len() != dimensions) {
return Err(Error::CorruptIndex(
"embedder returned an invalid vector shape".into(),
));
}
Ok(vectors
.into_iter()
.flatten()
.map(|value| (value.clamp(-1.0, 1.0) * 127.0).round() as i8)
.collect())
}
fn unchanged_loaded(files: &[crate::source::SourceFile], indexed: &[super::LoadedFile]) -> bool {
if files.len() != indexed.len() {
return false;
}
let stamps = indexed
.iter()
.map(|file| (file.path.as_str(), file.stamp))
.collect::<HashMap<_, _>>();
files
.iter()
.all(|file| stamps.get(file.relative_path.as_str()) == Some(&file.stamp))
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use super::*;
struct CountingEmbedder(AtomicUsize);
impl Embedder for CountingEmbedder {
fn id(&self) -> &str {
"counting-v1"
}
fn dimensions(&self) -> usize {
2
}
fn encode(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
self.0.fetch_add(texts.len(), Ordering::SeqCst);
Ok(texts
.iter()
.map(|text| vec![text.len() as f32, 1.0])
.collect())
}
}
#[test]
fn unchanged_files_reuse_persisted_chunks_and_vectors() {
let source = tempfile::tempdir().unwrap();
let cache = tempfile::tempdir().unwrap();
fs::write(
source.path().join("lib.rs"),
"pub fn one() { println!(\"one\"); }\n",
)
.unwrap();
let embedder = Arc::new(CountingEmbedder(AtomicUsize::new(0)));
let repository = IndexRepository::new(SembleConfig::new(cache.path()), embedder.clone());
let first = repository
.load_or_build(source.path(), "fixture", &[ContentType::Code])
.unwrap();
let encoded = embedder.0.load(Ordering::SeqCst);
assert!(encoded > 0);
let second = repository
.load_or_build(source.path(), "fixture", &[ContentType::Code])
.unwrap();
assert_eq!(embedder.0.load(Ordering::SeqCst), encoded);
assert_eq!(first.chunks.len(), second.chunks.len());
fs::write(
source.path().join("lib.rs"),
"pub fn two() { println!(\"changed and longer\"); }\n",
)
.unwrap();
let changed = repository
.load_or_build(source.path(), "fixture", &[ContentType::Code])
.unwrap();
assert!(embedder.0.load(Ordering::SeqCst) > encoded);
assert!(changed
.chunks
.iter()
.any(|chunk| chunk.content.contains("two")));
}
#[test]
fn a_new_repository_instance_loads_the_same_disk_snapshot() {
let source = tempfile::tempdir().unwrap();
let cache = tempfile::tempdir().unwrap();
fs::write(source.path().join("lib.rs"), "pub fn persisted() {}\n").unwrap();
let first_embedder = Arc::new(CountingEmbedder(AtomicUsize::new(0)));
IndexRepository::new(SembleConfig::new(cache.path()), first_embedder)
.load_or_build(source.path(), "fixture", &[ContentType::Code])
.unwrap();
let second_embedder = Arc::new(CountingEmbedder(AtomicUsize::new(0)));
let loaded = IndexRepository::new(SembleConfig::new(cache.path()), second_embedder.clone())
.load_or_build(source.path(), "fixture", &[ContentType::Code])
.unwrap();
assert_eq!(second_embedder.0.load(Ordering::SeqCst), 0);
assert_eq!(loaded.chunks[0].file_path, "lib.rs");
}
#[test]
fn a_new_repository_refreshes_changed_source_instead_of_using_stale_runtime_cache() {
let source = tempfile::tempdir().unwrap();
let cache = tempfile::tempdir().unwrap();
fs::write(source.path().join("lib.rs"), "pub fn before() {}\n").unwrap();
IndexRepository::new(
SembleConfig::new(cache.path()),
Arc::new(CountingEmbedder(AtomicUsize::new(0))),
)
.load_or_build(source.path(), "fixture", &[ContentType::Code])
.unwrap();
fs::write(
source.path().join("lib.rs"),
"pub fn after() { println!(\"changed and longer\"); }\n",
)
.unwrap();
let embedder = Arc::new(CountingEmbedder(AtomicUsize::new(0)));
let loaded = IndexRepository::new(SembleConfig::new(cache.path()), embedder.clone())
.load_or_build(source.path(), "fixture", &[ContentType::Code])
.unwrap();
assert!(embedder.0.load(Ordering::SeqCst) > 0);
assert!(loaded.chunks[0].content.contains("after"));
}
#[test]
fn search_cache_refreshes_after_the_fixed_interval() {
let source = tempfile::tempdir().unwrap();
let cache = tempfile::tempdir().unwrap();
fs::write(source.path().join("lib.rs"), "pub fn before() {}\n").unwrap();
let repository = IndexRepository::new(
SembleConfig::new(cache.path()),
Arc::new(CountingEmbedder(AtomicUsize::new(0))),
);
repository
.load_or_build(source.path(), "fixture", &[ContentType::Code])
.unwrap();
fs::write(source.path().join("lib.rs"), "pub fn after() {}\n").unwrap();
let cached = repository
.load_for_search(source.path(), "fixture", &[ContentType::Code])
.unwrap();
assert!(cached.chunks[0].content.contains("before"));
for (_, cached) in repository.loaded.lock().iter_mut() {
cached.checked_at = Instant::now() - SEARCH_REFRESH_INTERVAL;
}
let refreshed = repository
.load_for_search(source.path(), "fixture", &[ContentType::Code])
.unwrap();
assert!(refreshed.chunks[0].content.contains("after"));
}
}
+539
View File
@@ -0,0 +1,539 @@
//! Compact inverted lexical index and definition-aware symbol lookup.
use std::{
collections::{HashMap, HashSet},
sync::LazyLock,
};
use regex::Regex;
use serde::{Deserialize, Serialize};
use crate::Chunk;
use super::{IndexedDefinition, IndexedFile};
static DECLARATION: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?m)^\s*(?:(?:export|default|async|declare|public|private|protected|static|pub(?:\([^)]*\))?)\s+)*(?:function|fn|def|class|struct|enum|trait|interface|type|record|module|namespace|protocol)\s+([A-Za-z_][A-Za-z0-9_]*)",
)
.expect("static declaration regex")
});
static BINDING: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?m)^\s*(?:(?:export|default|declare|public|private|protected|static|pub(?:\([^)]*\))?)\s+)*(?:const|let|var|static)\s+([A-Za-z_][A-Za-z0-9_]*)\s*(?::[^=\n]+)?=",
)
.expect("static binding regex")
});
pub fn tokenize(text: &str) -> Vec<String> {
let bytes = text.as_bytes();
let mut output = Vec::new();
let mut index = 0;
while index < bytes.len() {
if !bytes[index].is_ascii_alphabetic() && bytes[index] != b'_' {
index += 1;
continue;
}
let start = index;
index += 1;
while index < bytes.len() && (bytes[index].is_ascii_alphanumeric() || bytes[index] == b'_')
{
index += 1;
}
output.extend(split_identifier(&text[start..index]));
}
output
}
fn split_identifier(token: &str) -> Vec<String> {
let lower = token.to_ascii_lowercase();
let mut parts = Vec::new();
if token.contains('_') {
parts.extend(
lower
.split('_')
.filter(|part| !part.is_empty())
.map(str::to_owned),
);
} else {
let chars = token.char_indices().collect::<Vec<_>>();
let mut start = 0;
for index in 1..chars.len() {
let previous = chars[index - 1].1;
let current = chars[index].1;
let next = chars.get(index + 1).map(|item| item.1);
if (previous.is_ascii_lowercase() && current.is_ascii_uppercase())
|| (previous.is_ascii_uppercase()
&& current.is_ascii_uppercase()
&& next.is_some_and(|value| value.is_ascii_lowercase()))
{
parts.push(token[start..chars[index].0].to_ascii_lowercase());
start = chars[index].0;
}
}
if start > 0 {
parts.push(token[start..].to_ascii_lowercase());
}
}
if parts.len() >= 2 {
std::iter::once(lower).chain(parts).collect()
} else {
vec![lower]
}
}
#[derive(Clone, Debug, Deserialize, Serialize)]
struct Definition {
symbol: String,
parts: Vec<String>,
chunk: usize,
implementation: bool,
}
#[derive(Deserialize, Serialize)]
pub struct SymbolIndex {
definitions: Vec<Definition>,
definitions_by_symbol: HashMap<String, Vec<usize>>,
}
impl SymbolIndex {
pub fn from_files(files: &mut [IndexedFile]) -> Self {
let mut definitions = Vec::new();
let mut definitions_by_symbol = HashMap::<String, Vec<usize>>::new();
let mut chunk_offset = 0;
for file in files {
for indexed in std::mem::take(&mut file.definitions) {
let parts = symbol_parts(&indexed.symbol);
let mut symbol = indexed.symbol;
symbol.make_ascii_lowercase();
let position = definitions.len();
definitions.push(Definition {
symbol: symbol.clone(),
parts,
chunk: chunk_offset + indexed.chunk,
implementation: indexed.implementation,
});
definitions_by_symbol
.entry(symbol)
.or_default()
.push(position);
}
chunk_offset += file.chunks.len();
}
Self {
definitions,
definitions_by_symbol,
}
}
pub fn exact_symbol(&self, query: &str, chunks: &[Chunk], limit: usize) -> Vec<(usize, f32)> {
let symbol = query
.rsplit([':', '.', '>'])
.find(|part| !part.is_empty())
.unwrap_or(query)
.trim()
.to_ascii_lowercase();
self.rank_symbol(&symbol, chunks, limit)
}
pub fn inferred_symbols(
&self,
query: &str,
chunks: &[Chunk],
limit: usize,
) -> Vec<(usize, f32)> {
let raw_terms = tokenize(query).into_iter().collect::<HashSet<_>>();
let mut query_terms = raw_terms
.iter()
.map(|term| normalize_term(term))
.collect::<HashSet<_>>();
if query_terms.contains("update") {
query_terms.insert("set".into());
}
let mut symbols = HashMap::<&str, f32>::new();
for definition in &self.definitions {
let meaningful = definition
.parts
.iter()
.filter(|part| !is_generic_symbol_part(part))
.collect::<Vec<_>>();
let matched = meaningful
.iter()
.filter(|part| query_terms.contains(&normalize_term(part)))
.count();
let exact = definition.parts.len() > 1 && raw_terms.contains(&definition.symbol);
let qualifies = exact
|| (meaningful.len() == 1 && matched == 1)
|| (meaningful.len() >= 2 && matched == meaningful.len())
|| (meaningful.len() >= 3 && matched >= 2);
if !qualifies {
continue;
}
let coverage = matched as f32 / meaningful.len().max(1) as f32;
let unmatched = meaningful.len().saturating_sub(matched) as f32;
let action_match = meaningful.iter().any(|part| {
let normalized = normalize_term(part);
query_terms.contains(&normalized) && is_action_symbol_part(&normalized)
});
let score = if exact { 1.0 } else { 0.0 } + coverage * 5.0 + matched as f32 * 2.0
- unmatched * 3.0
- meaningful.len() as f32 * 0.1
+ if action_match { 2.0 } else { 0.0 };
symbols
.entry(&definition.symbol)
.and_modify(|current| *current = current.max(score))
.or_insert(score);
}
let mut symbols = symbols.into_iter().collect::<Vec<_>>();
symbols.sort_by(|left, right| right.1.total_cmp(&left.1).then_with(|| left.0.cmp(right.0)));
symbols.truncate(limit.min(6));
let mut output = Vec::new();
for (symbol, symbol_score) in symbols {
for (chunk, source_score) in self.rank_symbol(symbol, chunks, 1) {
let path_matches = tokenize(&chunks[chunk].file_path)
.into_iter()
.map(|term| normalize_term(&term))
.collect::<HashSet<_>>()
.intersection(&query_terms)
.count();
let source_weight = source_priority(&chunks[chunk].file_path, symbol).min(1.0);
if source_weight < 0.5 {
continue;
}
output.push((
chunk,
(symbol_score * 10.0 + path_matches as f32 * 2.0) * source_weight
+ source_score,
));
}
}
output.sort_by(|left, right| {
right
.1
.total_cmp(&left.1)
.then_with(|| left.0.cmp(&right.0))
});
output.dedup_by_key(|item| item.0);
output.truncate(limit);
output
}
fn rank_symbol(&self, symbol: &str, chunks: &[Chunk], limit: usize) -> Vec<(usize, f32)> {
let mut output = self
.definitions_by_symbol
.get(symbol)
.into_iter()
.flatten()
.map(|&position| {
let definition = &self.definitions[position];
let chunk = &chunks[definition.chunk];
let score = (1.0 + if definition.implementation { 0.5 } else { 0.0 })
* source_priority(&chunk.file_path, symbol);
(definition.chunk, score)
})
.collect::<Vec<_>>();
output.sort_by(|left, right| {
right
.1
.total_cmp(&left.1)
.then_with(|| chunks[left.0].file_path.cmp(&chunks[right.0].file_path))
.then_with(|| chunks[right.0].start_line.cmp(&chunks[left.0].start_line))
});
output.dedup_by_key(|item| item.0);
output.truncate(limit);
output
}
pub(crate) fn validate(&self, chunks: usize) -> bool {
if self
.definitions
.iter()
.any(|definition| definition.chunk >= chunks)
{
return false;
}
let mut seen = vec![false; self.definitions.len()];
for (symbol, positions) in &self.definitions_by_symbol {
for &position in positions {
let Some(definition) = self.definitions.get(position) else {
return false;
};
if definition.symbol != *symbol || std::mem::replace(&mut seen[position], true) {
return false;
}
}
}
seen.into_iter().all(|value| value)
}
}
pub fn extract_definitions(chunks: &[Chunk]) -> Vec<IndexedDefinition> {
let mut definitions = Vec::new();
for (index, chunk) in chunks.iter().enumerate() {
for captures in DECLARATION
.captures_iter(&chunk.content)
.chain(BINDING.captures_iter(&chunk.content))
{
let Some(matched) = captures.get(1) else {
continue;
};
definitions.push(IndexedDefinition {
symbol: matched.as_str().to_owned(),
chunk: index,
implementation: chunk.content[matched.end()..]
.lines()
.take(12)
.any(|line| line.contains('{') || line.trim_end().ends_with("=>")),
});
}
}
definitions
}
fn symbol_parts(symbol: &str) -> Vec<String> {
let mut values = split_identifier(symbol);
if values.len() > 1 {
values.remove(0);
}
values
.into_iter()
.map(|part| normalize_term(&part))
.collect()
}
fn normalize_term(term: &str) -> String {
let lower = term.to_ascii_lowercase();
match lower.as_str() {
"application" | "applications" => "app".into(),
"listener" | "listeners" | "listening" => "listen".into(),
"rendered" | "rendering" => "render".into(),
"hydrated" | "hydrating" => "hydrate".into(),
"compiled" | "compiling" => "compile".into(),
"parsed" | "parsing" => "parse".into(),
"called" | "calling" => "call".into(),
"cancelled" | "cancelling" | "cancellation" => "cancel".into(),
"dispatched" | "dispatching" => "dispatch".into(),
"interrupted" | "interrupting" | "interruption" => "interrupt".into(),
"started" | "starting" => "start".into(),
_ if lower.len() > 4 && lower.ends_with('s') => lower[..lower.len() - 1].into(),
_ => lower,
}
}
fn is_action_symbol_part(part: &str) -> bool {
matches!(
part,
"build"
| "call"
| "cancel"
| "commit"
| "compile"
| "create"
| "define"
| "dispatch"
| "find"
| "handle"
| "hydrate"
| "interrupt"
| "listen"
| "load"
| "mount"
| "parse"
| "patch"
| "persist"
| "read"
| "reconcile"
| "record"
| "register"
| "render"
| "save"
| "schedule"
| "search"
| "send"
| "set"
| "start"
| "store"
| "update"
| "watch"
| "write"
)
}
fn is_generic_symbol_part(part: &str) -> bool {
matches!(
part,
"all" | "api" | "base" | "impl" | "internal" | "of" | "on" | "the" | "to"
)
}
fn source_priority(path: &str, symbol: &str) -> f32 {
let normalized = path.replace('\\', "/").to_ascii_lowercase();
let mut score = 1.0;
if normalized.contains("/src/") {
score += 0.35;
}
let parts = normalized.split('/').collect::<Vec<_>>();
if let Some(packages) = parts.iter().position(|part| *part == "packages") {
if parts
.get(packages + 1)
.is_some_and(|package| !package.contains('-'))
{
score += 1.0;
}
}
if normalized
.rsplit('/')
.next()
.and_then(|file| file.rsplit_once('.').map(|(stem, _)| stem))
.is_some_and(|stem| stem == symbol || stem.ends_with(symbol))
{
score += 0.4;
}
if normalized.split('/').any(|part| {
matches!(
part,
"test" | "tests" | "__tests__" | "fixtures" | "examples" | "benchmarks"
)
}) || normalized.contains(".test.")
|| normalized.contains(".spec.")
|| normalized.contains("dts-test")
{
score *= 0.08;
}
if normalized
.split('/')
.any(|part| matches!(part, "playground" | "demo" | "demos"))
{
score *= 0.2;
}
score
}
#[cfg(test)]
mod tests {
use super::*;
fn chunk(path: &str, content: &str) -> Chunk {
Chunk {
file_path: path.into(),
start_line: 1,
end_line: 1,
language: Some("typescript".into()),
content: content.into(),
}
}
#[test]
fn tokenizer_preserves_and_splits_common_identifiers() {
assert_eq!(
tokenize("HTTPResponse parse_request"),
vec![
"httpresponse",
"http",
"response",
"parse_request",
"parse",
"request"
]
);
}
#[test]
fn exact_symbol_prefers_implementations_and_non_test_sources() {
let chunks = [
chunk("tests/useState.test.ts", "function useState() {}"),
chunk("src/hooks.ts", "export function useState<T>(value: T)\n"),
chunk(
"src/hooks.ts",
"export function useState(value: unknown) { return value }",
),
];
let index = SymbolIndex::from_files(&mut [IndexedFile {
path: "fixture.ts".into(),
stamp: crate::source::FileStamp {
modified_ns: 1,
size: 1,
},
content_type: crate::ContentType::Code,
definitions: extract_definitions(&chunks),
lexical_documents: chunks
.iter()
.map(|chunk| super::super::lexical_document(&chunk.content))
.collect(),
chunks: chunks.to_vec(),
vectors: vec![0; chunks.len()],
}]);
assert_eq!(index.exact_symbol("useState", &chunks, 3)[0].0, 2);
}
#[test]
fn natural_queries_identify_composite_definitions() {
let chunks = [
chunk(
"src/renderer.ts",
"const patchKeyedChildren = () => { longestIncreasingSubsequence() }",
),
chunk("src/other.ts", "function unrelated() {}"),
];
let index = SymbolIndex::from_files(&mut [IndexedFile {
path: "fixture.ts".into(),
stamp: crate::source::FileStamp {
modified_ns: 1,
size: 1,
},
content_type: crate::ContentType::Code,
definitions: extract_definitions(&chunks),
lexical_documents: chunks
.iter()
.map(|chunk| super::super::lexical_document(&chunk.content))
.collect(),
chunks: chunks.to_vec(),
vectors: vec![0; chunks.len()],
}]);
assert_eq!(
index.inferred_symbols(
"diff keyed children using the longest increasing subsequence",
&chunks,
5,
)[0]
.0,
0
);
}
#[test]
fn natural_queries_prefer_symbols_that_match_behavior_words() {
let chunks = [
chunk("src/runtime.rs", "fn mcp_tool() {}"),
chunk("src/runtime.rs", "fn built_in() {}"),
chunk("src/dispatch.rs", "fn call_tool() {}"),
chunk("tests/dispatch.rs", "fn dispatch_call_tool() {}"),
];
let index = SymbolIndex::from_files(&mut [IndexedFile {
path: "fixture.ts".into(),
stamp: crate::source::FileStamp {
modified_ns: 1,
size: 1,
},
content_type: crate::ContentType::Code,
definitions: extract_definitions(&chunks),
lexical_documents: chunks
.iter()
.map(|chunk| super::super::lexical_document(&chunk.content))
.collect(),
chunks: chunks.to_vec(),
vectors: vec![0; chunks.len()],
}]);
assert_eq!(
index.inferred_symbols(
"dispatch built-in MCP tool calls to the registered server implementation",
&chunks,
5,
)[0]
.0,
2
);
}
}
+14
View File
@@ -0,0 +1,14 @@
//! Incremental persisted index construction and in-memory searchable snapshots.
mod bm25;
mod builder;
mod lexical;
mod snapshot;
pub(crate) use bm25::{lexical_document, Bm25Index, Bm25Match};
pub use builder::IndexRepository;
pub use lexical::{extract_definitions, tokenize, SymbolIndex};
pub use snapshot::{
IndexSnapshot, IndexedDefinition, IndexedFile, IndexedLexicalDocument, IndexedTermFrequency,
LoadedFile, LoadedIndex, LoadedMetadata,
};
+219
View File
@@ -0,0 +1,219 @@
//! Versioned binary snapshot format and flattened runtime representation.
use serde::{Deserialize, Serialize};
use crate::{config::INDEX_FORMAT_VERSION, source::FileStamp, Chunk, ContentType, Error, Result};
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct IndexedFile {
pub path: String,
pub stamp: FileStamp,
pub content_type: ContentType,
pub chunks: Vec<Chunk>,
pub definitions: Vec<IndexedDefinition>,
pub lexical_documents: Vec<IndexedLexicalDocument>,
/// Row-major signed-byte unit vectors, one row per chunk.
pub vectors: Vec<i8>,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct IndexedLexicalDocument {
pub length: u32,
pub terms: Vec<IndexedTermFrequency>,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct IndexedTermFrequency {
pub term: String,
pub frequency: u32,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct IndexedDefinition {
pub symbol: String,
pub chunk: usize,
pub implementation: bool,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct IndexSnapshot {
pub format_version: u32,
pub source_identity: String,
pub model_id: String,
pub dimensions: usize,
pub desired_chunk_bytes: usize,
pub content: Vec<ContentType>,
pub files: Vec<IndexedFile>,
}
impl IndexSnapshot {
pub fn validate(&self) -> Result<()> {
if self.format_version != INDEX_FORMAT_VERSION {
return Err(Error::CorruptIndex("format version mismatch".into()));
}
for file in &self.files {
if file
.definitions
.iter()
.any(|definition| definition.chunk >= file.chunks.len())
{
return Err(Error::CorruptIndex(format!(
"definition chunk is out of bounds in {}",
file.path
)));
}
if file.vectors.len() != file.chunks.len() * self.dimensions {
return Err(Error::CorruptIndex(format!(
"chunk/vector count mismatch in {}",
file.path
)));
}
if file.lexical_documents.len() != file.chunks.len() {
return Err(Error::CorruptIndex(format!(
"chunk/lexical document count mismatch in {}",
file.path
)));
}
for document in &file.lexical_documents {
let mut previous = None;
let mut length = 0_u64;
for term in &document.terms {
if term.term.is_empty()
|| term.frequency == 0
|| previous.is_some_and(|value| value >= term.term.as_str())
{
return Err(Error::CorruptIndex(format!(
"invalid lexical document in {}",
file.path
)));
}
previous = Some(term.term.as_str());
length += u64::from(term.frequency);
}
if length != u64::from(document.length) {
return Err(Error::CorruptIndex(format!(
"lexical document length mismatch in {}",
file.path
)));
}
}
}
Ok(())
}
}
#[derive(Deserialize, Serialize)]
pub struct LoadedFile {
pub path: String,
pub stamp: FileStamp,
pub content_type: ContentType,
}
#[derive(Deserialize, Serialize)]
pub struct LoadedMetadata {
pub source_identity: String,
pub model_id: String,
pub dimensions: usize,
pub desired_chunk_bytes: usize,
pub content: Vec<ContentType>,
pub files: Vec<LoadedFile>,
}
#[derive(Deserialize, Serialize)]
pub struct LoadedIndex {
pub metadata: LoadedMetadata,
pub chunks: Vec<Chunk>,
pub vectors: Vec<i8>,
pub lexical: super::SymbolIndex,
pub bm25: super::Bm25Index,
}
impl LoadedIndex {
pub fn from_snapshot(mut snapshot: IndexSnapshot) -> Result<Self> {
snapshot.validate()?;
let lexical = super::SymbolIndex::from_files(&mut snapshot.files);
let bm25 = super::Bm25Index::from_files(&mut snapshot.files);
let metadata = LoadedMetadata {
source_identity: std::mem::take(&mut snapshot.source_identity),
model_id: std::mem::take(&mut snapshot.model_id),
dimensions: snapshot.dimensions,
desired_chunk_bytes: snapshot.desired_chunk_bytes,
content: std::mem::take(&mut snapshot.content),
files: snapshot
.files
.iter_mut()
.map(|file| LoadedFile {
path: std::mem::take(&mut file.path),
stamp: file.stamp,
content_type: file.content_type,
})
.collect(),
};
let chunks = snapshot
.files
.iter_mut()
.flat_map(|file| std::mem::take(&mut file.chunks))
.collect::<Vec<_>>();
let vectors = snapshot
.files
.iter_mut()
.flat_map(|file| std::mem::take(&mut file.vectors))
.collect::<Vec<_>>();
Ok(Self {
metadata,
chunks,
vectors,
lexical,
bm25,
})
}
pub fn validate(&self) -> Result<()> {
if self.metadata.dimensions == 0
|| self.vectors.len() != self.chunks.len() * self.metadata.dimensions
|| !self.lexical.validate(self.chunks.len())
|| !self.bm25.validate(self.chunks.len())
{
return Err(Error::CorruptIndex("invalid runtime index".into()));
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use crate::source::FileStamp;
use super::*;
#[test]
fn validation_rejects_vector_shape_mismatches() {
let snapshot = IndexSnapshot {
format_version: INDEX_FORMAT_VERSION,
source_identity: "source".into(),
model_id: "model".into(),
dimensions: 2,
desired_chunk_bytes: 100,
content: vec![ContentType::Code],
files: vec![IndexedFile {
path: "lib.rs".into(),
stamp: FileStamp {
modified_ns: 1,
size: 1,
},
content_type: ContentType::Code,
chunks: vec![Chunk {
file_path: "lib.rs".into(),
start_line: 1,
end_line: 1,
language: Some("rust".into()),
content: "x".into(),
}],
definitions: Vec::new(),
lexical_documents: vec![super::super::lexical_document("x")],
vectors: vec![1],
}],
};
assert!(matches!(snapshot.validate(), Err(Error::CorruptIndex(_))));
}
}
@@ -0,0 +1,84 @@
//! Extension and special-filename catalog for code, documentation, and configuration.
use std::path::Path;
use crate::ContentType;
pub fn detect_language(path: &Path) -> Option<&'static str> {
let extension = path.extension()?.to_str()?.to_ascii_lowercase();
Some(match extension.as_str() {
"rs" => "rust",
"py" | "pyi" | "pyw" => "python",
"js" | "jsx" | "mjs" | "cjs" => "javascript",
"ts" | "tsx" | "mts" | "cts" => "typescript",
"go" => "go",
"java" => "java",
"c" | "h" => "c",
"cc" | "cpp" | "cxx" | "hpp" | "hxx" => "cpp",
"cs" => "csharp",
"rb" => "ruby",
"php" => "php",
"swift" => "swift",
"kt" | "kts" => "kotlin",
"scala" => "scala",
"sh" | "bash" => "bash",
"lua" => "lua",
"ex" | "exs" => "elixir",
"dart" => "dart",
"sql" => "sql",
"html" | "htm" => "html",
"css" | "scss" | "less" => "css",
"json" | "json5" => "json",
"yaml" | "yml" => "yaml",
"toml" => "toml",
"md" | "markdown" | "rst" | "adoc" => "markdown",
_ => return None,
})
}
pub fn content_type_for_path(path: &Path) -> Option<ContentType> {
let name = path.file_name()?.to_str()?.to_ascii_lowercase();
if matches!(
name.as_str(),
"cargo.toml"
| "pyproject.toml"
| "package.json"
| "tsconfig.json"
| "dockerfile"
| "makefile"
| ".gitignore"
| ".sembleignore"
) {
return Some(ContentType::Config);
}
let language = detect_language(path)?;
if language == "markdown" {
Some(ContentType::Docs)
} else if matches!(language, "json" | "yaml" | "toml") {
Some(ContentType::Config)
} else {
Some(ContentType::Code)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn classifies_code_docs_and_configuration() {
assert_eq!(
content_type_for_path(Path::new("src/main.rs")),
Some(ContentType::Code)
);
assert_eq!(
content_type_for_path(Path::new("README.md")),
Some(ContentType::Docs)
);
assert_eq!(
content_type_for_path(Path::new("Cargo.toml")),
Some(ContentType::Config)
);
assert_eq!(content_type_for_path(Path::new("image.png")), None);
}
}
+7
View File
@@ -0,0 +1,7 @@
//! Language detection and Tree-sitter parser selection.
mod catalog;
mod parser;
pub use catalog::{content_type_for_path, detect_language};
pub use parser::parser_for;
+20
View File
@@ -0,0 +1,20 @@
//! Tree-sitter parser construction for bundled high-value languages.
use tree_sitter::Parser;
pub fn parser_for(language: &str) -> Option<Parser> {
let language = match language {
"rust" => tree_sitter_rust::LANGUAGE.into(),
"python" => tree_sitter_python::LANGUAGE.into(),
"javascript" => tree_sitter_javascript::LANGUAGE.into(),
"typescript" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
"go" => tree_sitter_go::LANGUAGE.into(),
"java" => tree_sitter_java::LANGUAGE.into(),
"c" => tree_sitter_c::LANGUAGE.into(),
"cpp" => tree_sitter_cpp::LANGUAGE.into(),
_ => return None,
};
let mut parser = Parser::new();
parser.set_language(&language).ok()?;
Some(parser)
}
+20
View File
@@ -0,0 +1,20 @@
//! Transport-independent code indexing, persistence, and hybrid retrieval.
pub mod cache;
pub mod chunk;
pub mod config;
pub mod embedding;
pub mod error;
pub mod index;
pub mod language;
pub mod search;
pub mod source;
pub mod types;
pub use config::SembleConfig;
pub use embedding::{Embedder, StaticEmbedder};
pub use error::{Error, Result};
pub use search::SearchEngine;
pub use types::{
Chunk, ContentType, FindRelatedRequest, IndexStats, SearchRequest, SearchResponse, SearchResult,
};
+810
View File
@@ -0,0 +1,810 @@
//! Cached index lifecycle and semantic, lexical, and related-code queries.
use std::{
collections::{HashMap, HashSet},
path::{Path, PathBuf},
sync::{Arc, LazyLock},
};
use parking_lot::Mutex;
use crate::{
embedding::{Embedder, ModelAssets, StaticEmbedder},
index::{tokenize, Bm25Match, IndexRepository, LoadedIndex},
source::{canonical_source_root, RemoteRepository},
types::snippet,
ContentType, Error, FindRelatedRequest, IndexStats, Result, SearchRequest, SearchResponse,
SearchResult, SembleConfig,
};
use super::{
rerank::{is_symbol_query, path_penalty, rerank},
rrf,
};
const NATURAL_LANGUAGE_SEMANTIC_WEIGHT: f32 = 0.5;
const SYMBOL_SEMANTIC_WEIGHT: f32 = 0.3;
pub struct SearchEngine {
config: SembleConfig,
repository: IndexRepository,
}
static EMBEDDERS: LazyLock<Mutex<HashMap<PathBuf, Arc<StaticEmbedder>>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
impl SearchEngine {
pub fn load_default(config: SembleConfig) -> Result<Self> {
let model_path = ModelAssets::model_path(&config.cache_dir);
let cached = { EMBEDDERS.lock().get(&model_path).cloned() };
let embedder = if let Some(embedder) = cached {
embedder
} else {
let assets = ModelAssets::ensure(&config.cache_dir)?;
let embedder = Arc::new(StaticEmbedder::load(&assets.model, &assets.tokenizer)?);
EMBEDDERS
.lock()
.insert(assets.model.clone(), embedder.clone());
embedder
};
Ok(Self::with_embedder(config, embedder))
}
pub fn with_embedder(config: SembleConfig, embedder: Arc<dyn Embedder>) -> Self {
let repository = IndexRepository::new(config.clone(), embedder);
Self { config, repository }
}
/// Opens a prepared index and incrementally refreshes changed source files.
pub fn prepare(&self, repo: &Path, content: &[ContentType]) -> Result<IndexStats> {
let (root, identity) = self.resolve_source(repo)?;
let content = normalize_content(content);
let index = self.repository.load_or_build(&root, &identity, &content)?;
Ok(index_stats(&index))
}
/// Rescans source stamps and incrementally refreshes a prepared index.
pub fn refresh(&self, repo: &Path, content: &[ContentType]) -> Result<IndexStats> {
let (root, identity) = self.resolve_source(repo)?;
let content = normalize_content(content);
let index = self.repository.load_or_build(&root, &identity, &content)?;
Ok(index_stats(&index))
}
pub fn search(&self, request: SearchRequest) -> Result<SearchResponse> {
validate_query(&request.query, request.top_k)?;
let (root, identity) = self.resolve_source(&request.repo)?;
let content = normalize_content(&request.content);
let index = self
.repository
.load_for_search(&root, &identity, &content)?;
if !request.query.contains(char::is_whitespace) {
let exact = index
.lexical
.exact_symbol(&request.query, &index.chunks, request.top_k);
if !exact.is_empty() {
return response(
&request.query,
exact,
&index,
&root,
request.max_snippet_lines,
);
}
}
let query_vector = quantize_query(self.encode_query(&index, &request.query)?);
let candidate_count = request.top_k.saturating_mul(5).max(request.top_k);
let semantic = rank_semantic(
&query_vector,
&index.vectors,
index.metadata.dimensions,
candidate_count,
);
let lexical = index.bm25.search(&request.query, candidate_count);
let lexical_ranking = lexical
.iter()
.map(|matched| matched.document)
.collect::<Vec<_>>();
let symbol_query = is_symbol_query(&request.query);
let semantic_weight = if symbol_query {
SYMBOL_SEMANTIC_WEIGHT
} else {
NATURAL_LANGUAGE_SEMANTIC_WEIGHT
};
let fused = rrf::fuse(&semantic, &lexical_ranking, semantic_weight);
let definitions = if !symbol_query {
index
.lexical
.inferred_symbols(&request.query, &index.chunks, request.top_k)
} else {
Vec::new()
};
let strong_lexical = if symbol_query || !is_literal_query(&request.query) {
Vec::new()
} else {
strong_lexical_matches(
&lexical,
&index.chunks,
&request.query,
&root,
request.top_k.min(3),
)?
};
let ranked = prioritize_evidence(
strong_lexical,
definitions,
rerank(fused, &index.chunks, &request.query, request.top_k),
request.top_k,
);
response(
&request.query,
ranked,
&index,
&root,
request.max_snippet_lines,
)
}
pub fn find_related(&self, request: FindRelatedRequest) -> Result<SearchResponse> {
if request.line == 0 || request.top_k == 0 {
return Err(Error::InvalidRequest(
"line and top_k must be greater than zero".into(),
));
}
let (root, identity) = self.resolve_source(&request.repo)?;
let content = normalize_content(&request.content);
let index = self
.repository
.load_for_search(&root, &identity, &content)?;
let normalized = request.file_path.replace('\\', "/");
let source = index
.chunks
.iter()
.position(|chunk| {
chunk.file_path == normalized
&& chunk.start_line <= request.line
&& request.line <= chunk.end_line
})
.ok_or_else(|| {
Error::InvalidRequest(format!(
"no indexed chunk contains {}:{}",
request.file_path, request.line
))
})?;
let dimensions = index.metadata.dimensions;
let source_start = source * dimensions;
let source_vector = &index.vectors[source_start..source_start + dimensions];
let mut ranked =
rank_semantic(source_vector, &index.vectors, dimensions, request.top_k + 1)
.into_iter()
.filter(|position| *position != source)
.take(request.top_k)
.map(|position| {
(
position,
quantized_cosine(
source_vector,
&index.vectors[position * dimensions..(position + 1) * dimensions],
),
)
})
.collect::<Vec<_>>();
ranked.sort_by(|left, right| right.1.total_cmp(&left.1));
response(
&format!("Chunks related to {}:{}", request.file_path, request.line),
ranked,
&index,
&root,
request.max_snippet_lines,
)
}
fn resolve_source(&self, source: &Path) -> Result<(PathBuf, String)> {
let value = source.to_string_lossy();
if value.starts_with("https://") || value.starts_with("http://") {
let remote = RemoteRepository::acquire(&value, &self.config.cache_dir)?;
Ok((remote.path, remote.identity))
} else {
let root = canonical_source_root(source)?;
let identity = root.to_string_lossy().into_owned();
Ok((root, identity))
}
}
fn encode_query(&self, _index: &LoadedIndex, query: &str) -> Result<Vec<f32>> {
self.repository
.encode(&[query.to_owned()])
.map(|mut values| values.remove(0))
}
}
fn validate_query(query: &str, top_k: usize) -> Result<()> {
if query.trim().is_empty() {
return Err(Error::InvalidRequest("query must not be empty".into()));
}
if top_k == 0 || top_k > 100 {
return Err(Error::InvalidRequest(
"top_k must be between 1 and 100".into(),
));
}
Ok(())
}
fn index_stats(index: &LoadedIndex) -> IndexStats {
IndexStats {
file_count: index.metadata.files.len(),
chunk_count: index.chunks.len(),
source_bytes: index
.metadata
.files
.iter()
.map(|file| file.stamp.size)
.sum(),
dimensions: index.metadata.dimensions,
}
}
fn strong_lexical_matches(
lexical: &[Bm25Match],
chunks: &[crate::Chunk],
query: &str,
root: &Path,
limit: usize,
) -> Result<Vec<(usize, f32)>> {
let query_terms = tokenize(query).into_iter().collect::<HashSet<_>>();
if query_terms.len() < 2 || lexical.is_empty() || limit == 0 {
return Ok(Vec::new());
}
let top_score = lexical[0].score.max(f32::EPSILON);
let mut files = HashMap::<String, String>::new();
let mut strong = Vec::new();
for matched in lexical {
let score_ratio = matched.score / top_score;
let phrase = if matched.coverage >= 0.5 {
let content = indexed_chunk_text(&chunks[matched.document], root, &mut files)?;
if path_penalty(&chunks[matched.document].file_path) < 0.5
|| content.trim_start().starts_with("#[cfg(test)]")
{
continue;
}
contains_normalized_phrase(query, &content)
} else {
false
};
if !phrase && (matched.coverage < 0.8 || score_ratio < 0.7) {
continue;
}
let strength = (score_ratio + matched.coverage + if phrase { 2.0 } else { 0.0 })
* path_penalty(&chunks[matched.document].file_path);
strong.push((matched.document, strength));
}
strong.sort_by(|left, right| {
right
.1
.total_cmp(&left.1)
.then_with(|| chunks[left.0].file_path.cmp(&chunks[right.0].file_path))
.then_with(|| chunks[left.0].start_line.cmp(&chunks[right.0].start_line))
});
strong.truncate(limit);
Ok(strong)
}
fn is_literal_query(query: &str) -> bool {
let lower = query.trim().to_ascii_lowercase();
if lower.is_empty() {
return false;
}
if [
"not found",
"permission denied",
"access denied",
"unauthorized",
"forbidden",
"invalid argument",
"missing ",
"failed to",
"panic",
]
.iter()
.any(|marker| lower.contains(marker))
|| lower
.chars()
.any(|character| matches!(character, '`' | '"' | '\'' | '{' | '}' | '='))
{
return true;
}
let terms = lower.split_whitespace().collect::<Vec<_>>();
if !(2..=8).contains(&terms.len()) {
return false;
}
let natural_connectors = [
"a", "an", "and", "as", "at", "by", "for", "from", "how", "in", "into", "is", "of", "on",
"the", "to", "using", "when", "where", "while", "with",
];
let behavior_verbs = [
"build",
"call",
"compile",
"create",
"define",
"dispatch",
"find",
"handle",
"hydrate",
"implement",
"load",
"parse",
"persist",
"read",
"record",
"register",
"render",
"save",
"schedule",
"search",
"send",
"store",
"update",
"watch",
"write",
];
!terms.iter().any(|term| natural_connectors.contains(term))
&& !behavior_verbs.contains(&terms[0])
}
fn indexed_chunk_text(
chunk: &crate::Chunk,
root: &Path,
files: &mut HashMap<String, String>,
) -> Result<String> {
if !chunk.content.is_empty() {
return Ok(chunk.content.clone());
}
if !files.contains_key(&chunk.file_path) {
let path = root.join(&chunk.file_path);
let source = std::fs::read_to_string(&path).map_err(|error| Error::io(&path, error))?;
files.insert(chunk.file_path.clone(), source);
}
Ok(files[&chunk.file_path]
.lines()
.skip(chunk.start_line.saturating_sub(1))
.take(chunk.end_line.saturating_sub(chunk.start_line) + 1)
.collect::<Vec<_>>()
.join("\n"))
}
fn contains_normalized_phrase(query: &str, content: &str) -> bool {
let phrase = query
.split_whitespace()
.collect::<Vec<_>>()
.join(" ")
.to_ascii_lowercase();
let source = content
.split_whitespace()
.collect::<Vec<_>>()
.join(" ")
.to_ascii_lowercase();
!phrase.is_empty() && source.contains(&phrase)
}
fn prioritize_evidence(
strong_lexical: Vec<(usize, f32)>,
definitions: Vec<(usize, f32)>,
ranked: Vec<(usize, f32)>,
top_k: usize,
) -> Vec<(usize, f32)> {
let mut output = strong_lexical;
for candidate in definitions.into_iter().chain(ranked) {
if !output.iter().any(|existing| existing.0 == candidate.0) {
output.push(candidate);
}
}
output.truncate(top_k);
output
}
fn normalize_content(content: &[ContentType]) -> Vec<ContentType> {
let mut output = if content.is_empty() {
vec![ContentType::Code]
} else {
content.to_vec()
};
output.sort_by_key(|item| *item as u8);
output.dedup();
output
}
fn rank_semantic(query: &[i8], vectors: &[i8], dimensions: usize, limit: usize) -> Vec<usize> {
let scores = vectors
.chunks_exact(dimensions)
.map(|vector| quantized_dot(query, vector))
.collect::<Vec<_>>();
rank_scores(&scores, limit, false)
}
fn rank_scores(scores: &[f32], limit: usize, exclude_zero: bool) -> Vec<usize> {
let mut indices = (0..scores.len())
.filter(|index| !exclude_zero || scores[*index] > 0.0)
.collect::<Vec<_>>();
indices.sort_by(|left, right| {
scores[*right]
.total_cmp(&scores[*left])
.then_with(|| left.cmp(right))
});
indices.truncate(limit.min(indices.len()));
indices
}
fn quantized_dot(left: &[i8], right: &[i8]) -> f32 {
left.iter()
.zip(right)
.map(|(left, right)| i32::from(*left) * i32::from(*right))
.sum::<i32>() as f32
}
fn quantized_cosine(left: &[i8], right: &[i8]) -> f32 {
quantized_dot(left, right) / (127.0 * 127.0)
}
fn quantize_query(vector: Vec<f32>) -> Vec<i8> {
vector
.into_iter()
.map(|value| (value.clamp(-1.0, 1.0) * 127.0).round() as i8)
.collect()
}
fn response(
query: &str,
ranked: Vec<(usize, f32)>,
index: &LoadedIndex,
root: &Path,
max_lines: Option<usize>,
) -> Result<SearchResponse> {
Ok(SearchResponse {
query: query.to_owned(),
results: ranked
.into_iter()
.map(|(position, score)| {
let chunk = &index.chunks[position];
Ok(SearchResult {
file_path: chunk.file_path.clone(),
start_line: chunk.start_line,
end_line: chunk.end_line,
score,
content: chunk_content(root, chunk, max_lines)?,
})
})
.collect::<Result<Vec<_>>>()?,
})
}
fn chunk_content(
root: &Path,
chunk: &crate::Chunk,
lines: Option<usize>,
) -> Result<Option<String>> {
if lines == Some(0) {
return Ok(None);
}
if !chunk.content.is_empty() {
return Ok(snippet(&chunk.content, lines));
}
let path = root.join(&chunk.file_path);
let source = std::fs::read_to_string(&path).map_err(|error| Error::io(&path, error))?;
let available = chunk.end_line.saturating_sub(chunk.start_line) + 1;
let limit = lines.unwrap_or(available).min(available);
Ok(Some(
source
.lines()
.skip(chunk.start_line.saturating_sub(1))
.take(limit)
.collect::<Vec<_>>()
.join("\n"),
))
}
#[cfg(test)]
mod tests {
use std::fs;
use super::*;
struct KeywordEmbedder;
impl Embedder for KeywordEmbedder {
fn id(&self) -> &str {
"keyword-v1"
}
fn dimensions(&self) -> usize {
3
}
fn encode(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
Ok(texts
.iter()
.map(|text| {
let text = text.to_ascii_lowercase();
let mut vector = vec![
usize::from(text.contains("auth")) as f32,
usize::from(text.contains("invoice")) as f32,
usize::from(text.contains("parse")) as f32,
];
let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
if norm > 0.0 {
vector.iter_mut().for_each(|value| *value /= norm);
}
vector
})
.collect())
}
}
fn fixture() -> (tempfile::TempDir, tempfile::TempDir, SearchEngine) {
let source = tempfile::tempdir().unwrap();
let cache = tempfile::tempdir().unwrap();
fs::create_dir_all(source.path().join("src")).unwrap();
fs::write(
source.path().join("src/auth.rs"),
"pub fn authenticate_request() {\n verify_token();\n}\n",
)
.unwrap();
fs::write(
source.path().join("src/billing.rs"),
"pub fn create_invoice() {\n charge_customer();\n}\n",
)
.unwrap();
let engine =
SearchEngine::with_embedder(SembleConfig::new(cache.path()), Arc::new(KeywordEmbedder));
(source, cache, engine)
}
#[test]
fn hybrid_search_returns_locations_and_bounded_snippets() {
let (source, _cache, engine) = fixture();
let response = engine
.search(SearchRequest {
query: "authenticate request".into(),
repo: source.path().into(),
top_k: 1,
max_snippet_lines: Some(1),
content: vec![ContentType::Code],
})
.unwrap();
assert_eq!(response.results[0].file_path, "src/auth.rs");
assert_eq!(
response.results[0].content.as_deref(),
Some("pub fn authenticate_request() {")
);
}
#[test]
fn bm25_recovers_lexical_matches_when_semantic_scores_are_tied() {
let (source, _cache, engine) = fixture();
fs::write(
source.path().join("src/tracing.rs"),
"pub fn write_record() {\n let description = \"quasar chronicle telemetry durable\";\n}\n",
)
.unwrap();
let response = engine
.search(SearchRequest {
query: "quasar chronicle telemetry durable".into(),
repo: source.path().into(),
top_k: 1,
max_snippet_lines: Some(0),
content: vec![ContentType::Code],
})
.unwrap();
assert_eq!(response.results[0].file_path, "src/tracing.rs");
}
#[test]
fn complete_multi_term_matches_rank_ahead_of_inferred_definitions() {
let (source, cache, engine) = fixture();
fs::write(
source.path().join("src/definitions.rs"),
"pub fn quasar_chronicle() {}\npub fn telemetry_durable() {}\n",
)
.unwrap();
fs::write(
source.path().join("src/literal.rs"),
"pub fn diagnostic() {\n let message = \"quasar chronicle telemetry durable\";\n}\n",
)
.unwrap();
engine.prepare(source.path(), &[ContentType::Code]).unwrap();
drop(engine);
let engine =
SearchEngine::with_embedder(SembleConfig::new(cache.path()), Arc::new(KeywordEmbedder));
let response = engine
.search(SearchRequest {
query: "quasar chronicle telemetry durable".into(),
repo: source.path().into(),
top_k: 5,
max_snippet_lines: Some(0),
content: vec![ContentType::Code],
})
.unwrap();
assert_eq!(response.results[0].file_path, "src/literal.rs");
assert!(response
.results
.iter()
.any(|result| result.file_path == "src/definitions.rs"));
}
#[test]
fn prepare_reports_the_persisted_index_shape() {
let (source, _cache, engine) = fixture();
let stats = engine.prepare(source.path(), &[ContentType::Code]).unwrap();
assert_eq!(stats.file_count, 2);
assert_eq!(stats.chunk_count, 2);
assert!(stats.source_bytes > 0);
assert_eq!(stats.dimensions, 3);
}
#[test]
fn disk_loaded_indexes_read_result_snippets_from_source() {
let (source, cache, engine) = fixture();
engine.prepare(source.path(), &[ContentType::Code]).unwrap();
drop(engine);
let reloaded =
SearchEngine::with_embedder(SembleConfig::new(cache.path()), Arc::new(KeywordEmbedder));
let response = reloaded
.search(SearchRequest {
query: "authenticate request".into(),
repo: source.path().into(),
top_k: 1,
max_snippet_lines: Some(1),
content: vec![ContentType::Code],
})
.unwrap();
assert_eq!(
response.results[0].content.as_deref(),
Some("pub fn authenticate_request() {")
);
}
#[test]
fn disk_loaded_indexes_retain_bm25_search() {
let (source, cache, engine) = fixture();
fs::write(
source.path().join("src/tracing.rs"),
"pub fn write_record() {\n let description = \"quasar chronicle telemetry durable\";\n}\n",
)
.unwrap();
engine.prepare(source.path(), &[ContentType::Code]).unwrap();
drop(engine);
let reloaded =
SearchEngine::with_embedder(SembleConfig::new(cache.path()), Arc::new(KeywordEmbedder));
let response = reloaded
.search(SearchRequest {
query: "quasar chronicle telemetry durable".into(),
repo: source.path().into(),
top_k: 1,
max_snippet_lines: Some(0),
content: vec![ContentType::Code],
})
.unwrap();
assert_eq!(response.results[0].file_path, "src/tracing.rs");
}
#[test]
fn explicit_refresh_makes_modified_files_searchable_during_the_cache_window() {
let (source, _cache, engine) = fixture();
engine.prepare(source.path(), &[ContentType::Code]).unwrap();
fs::write(
source.path().join("src/auth.rs"),
"pub fn parse_request() {\n parse_payload();\n validate_fields();\n}\n",
)
.unwrap();
let stale = engine
.search(SearchRequest {
query: "parse_request".into(),
repo: source.path().into(),
top_k: 1,
max_snippet_lines: Some(3),
content: vec![ContentType::Code],
})
.unwrap();
assert!(!stale.results[0]
.content
.as_deref()
.unwrap_or_default()
.contains("parse_request"));
engine.refresh(source.path(), &[ContentType::Code]).unwrap();
let response = engine
.search(SearchRequest {
query: "parse_request".into(),
repo: source.path().into(),
top_k: 1,
max_snippet_lines: Some(0),
content: vec![ContentType::Code],
})
.unwrap();
assert_eq!(response.results[0].file_path, "src/auth.rs");
assert!(response.results[0].score >= 1.0);
}
#[test]
fn explicit_refresh_makes_new_files_available_to_related_search() {
let (source, _cache, engine) = fixture();
engine.prepare(source.path(), &[ContentType::Code]).unwrap();
fs::write(
source.path().join("src/parser.rs"),
"pub fn parse_document() {\n parse_payload();\n}\n",
)
.unwrap();
engine.refresh(source.path(), &[ContentType::Code]).unwrap();
let response = engine
.find_related(FindRelatedRequest {
repo: source.path().into(),
file_path: "src/parser.rs".into(),
line: 1,
top_k: 1,
max_snippet_lines: Some(0),
content: vec![ContentType::Code],
})
.unwrap();
assert_eq!(response.results.len(), 1);
assert_ne!(response.results[0].file_path, "src/parser.rs");
}
#[test]
fn prepare_removes_deleted_files_from_the_index() {
let (source, _cache, engine) = fixture();
engine.prepare(source.path(), &[ContentType::Code]).unwrap();
fs::remove_file(source.path().join("src/auth.rs")).unwrap();
let stats = engine.prepare(source.path(), &[ContentType::Code]).unwrap();
assert_eq!(stats.file_count, 1);
assert_eq!(stats.chunk_count, 1);
}
#[test]
fn related_search_excludes_the_source_chunk() {
let (source, _cache, engine) = fixture();
let response = engine
.find_related(FindRelatedRequest {
repo: source.path().into(),
file_path: "src/auth.rs".into(),
line: 1,
top_k: 1,
max_snippet_lines: Some(2),
content: vec![ContentType::Code],
})
.unwrap();
assert_eq!(response.results.len(), 1);
assert_ne!(response.results[0].file_path, "src/auth.rs");
}
#[test]
fn invalid_queries_fail_before_indexing() {
let (_source, _cache, engine) = fixture();
let request = SearchRequest::new("", "/missing");
assert!(matches!(
engine.search(request),
Err(Error::InvalidRequest(_))
));
}
#[test]
fn distinguishes_literal_fragments_from_behavior_queries() {
assert!(is_literal_query("Semble Code Search"));
assert!(is_literal_query(
"MCP tool not found server not found permission denied"
));
assert!(!is_literal_query(
"dispatch built-in MCP tool calls to the registered server implementation"
));
assert!(!is_literal_query("create a concurrent React DOM root"));
}
}
+7
View File
@@ -0,0 +1,7 @@
//! Public search engine, Reciprocal Rank Fusion, and code-aware reranking.
mod engine;
mod rerank;
mod rrf;
pub use engine::SearchEngine;
+143
View File
@@ -0,0 +1,143 @@
//! Definition boosts, path priors, file coherence, and result saturation.
use std::collections::HashMap;
use regex::Regex;
use crate::Chunk;
pub fn is_symbol_query(query: &str) -> bool {
let value = query.trim();
!value.contains(char::is_whitespace)
&& (value.contains("::")
|| value.contains("->")
|| value.contains('.')
|| value.starts_with('_')
|| value
.chars()
.any(|character| character == '_' || character.is_ascii_uppercase()))
}
pub fn rerank(
scores: HashMap<usize, f32>,
chunks: &[Chunk],
query: &str,
top_k: usize,
) -> Vec<(usize, f32)> {
if scores.is_empty() {
return Vec::new();
}
let mut scores = scores;
let max = scores.values().copied().fold(0.0_f32, f32::max);
let mut file_sum = HashMap::<&str, f32>::new();
let mut best = HashMap::<&str, usize>::new();
for (&index, &score) in &scores {
*file_sum.entry(&chunks[index].file_path).or_default() += score;
best.entry(&chunks[index].file_path)
.and_modify(|current| {
if score > scores[current] {
*current = index;
}
})
.or_insert(index);
}
let max_sum = file_sum
.values()
.copied()
.fold(0.0_f32, f32::max)
.max(f32::EPSILON);
for (path, index) in best {
*scores.entry(index).or_default() += max * 0.2 * file_sum[path] / max_sum;
}
if is_symbol_query(query) {
let symbol = query
.rsplit([':', '.', '>'])
.find(|part| !part.is_empty())
.unwrap_or(query)
.trim();
let pattern = Regex::new(&format!(r"(?m)(?:^|\s)(?:class|def|fn|func|function|struct|enum|trait|interface|type|record|module|namespace|protocol)\s+{}(?:\s|[<({{:\[;]|$)", regex::escape(symbol))).ok();
if let Some(pattern) = pattern {
for (&index, score) in &mut scores {
if pattern.is_match(&chunks[index].content) {
*score += max * 3.0;
}
}
}
}
let mut ranked = scores
.into_iter()
.map(|(index, score)| (index, score * path_penalty(&chunks[index].file_path)))
.collect::<Vec<_>>();
ranked.sort_by(|left, right| {
right
.1
.total_cmp(&left.1)
.then_with(|| chunks[left.0].file_path.cmp(&chunks[right.0].file_path))
.then_with(|| chunks[left.0].start_line.cmp(&chunks[right.0].start_line))
.then_with(|| left.0.cmp(&right.0))
});
let mut selected = Vec::new();
let mut per_file = HashMap::<&str, usize>::new();
for (index, mut score) in ranked {
let count = per_file.entry(&chunks[index].file_path).or_default();
if *count > 0 {
score *= 0.5_f32.powi(*count as i32);
}
*count += 1;
selected.push((index, score));
}
selected.sort_by(|left, right| {
right
.1
.total_cmp(&left.1)
.then_with(|| chunks[left.0].file_path.cmp(&chunks[right.0].file_path))
.then_with(|| chunks[left.0].start_line.cmp(&chunks[right.0].start_line))
.then_with(|| left.0.cmp(&right.0))
});
selected.truncate(top_k);
selected
}
pub(super) fn path_penalty(path: &str) -> f32 {
let value = path.replace('\\', "/").to_ascii_lowercase();
let mut penalty = 1.0;
if value
.split('/')
.any(|part| matches!(part, "test" | "tests" | "__tests__" | "spec" | "testing"))
|| value.contains(".test.")
|| value.contains(".spec.")
|| value.ends_with("_test.rs")
|| value.ends_with("_test.go")
{
penalty *= 0.3;
}
if value.split('/').any(|part| {
matches!(
part,
"compat" | "_compat" | "legacy" | "example" | "examples"
)
}) {
penalty *= 0.3;
}
if value.ends_with(".d.ts") {
penalty *= 0.7;
}
if value.ends_with("/__init__.py") || value.ends_with("/package-info.java") {
penalty *= 0.5;
}
penalty
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn recognizes_symbols_and_penalizes_tests() {
assert!(is_symbol_query("SearchEngine"));
assert!(!is_symbol_query("find the search engine"));
assert!(path_penalty("tests/search_test.rs") < path_penalty("src/search.rs"));
}
}
+28
View File
@@ -0,0 +1,28 @@
//! Reciprocal Rank Fusion of independently ranked semantic and lexical candidates.
use std::collections::HashMap;
const RRF_K: f32 = 60.0;
pub fn fuse(semantic: &[usize], lexical: &[usize], alpha: f32) -> HashMap<usize, f32> {
let mut scores = HashMap::new();
for (rank, index) in semantic.iter().enumerate() {
*scores.entry(*index).or_insert(0.0) += alpha / (RRF_K + rank as f32 + 1.0);
}
for (rank, index) in lexical.iter().enumerate() {
*scores.entry(*index).or_insert(0.0) += (1.0 - alpha) / (RRF_K + rank as f32 + 1.0);
}
scores
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rewards_items_present_in_both_lists() {
let scores = fuse(&[1, 2], &[2, 3], 0.5);
assert!(scores[&2] > scores[&1]);
assert!(scores[&2] > scores[&3]);
}
}
+173
View File
@@ -0,0 +1,173 @@
//! Deterministic local file discovery with gitignore and Semble exclusions.
use std::{
fs,
path::{Path, PathBuf},
sync::mpsc,
time::UNIX_EPOCH,
};
use ignore::{WalkBuilder, WalkState};
use serde::{Deserialize, Serialize};
use crate::{language::content_type_for_path, ContentType, Error, Result};
const IGNORED_DIRS: &[&str] = &[
".git",
".hg",
".svn",
"node_modules",
"target",
".venv",
"venv",
".tox",
"__pycache__",
".next",
"dist",
"build",
".cache",
".semble",
];
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct FileStamp {
pub modified_ns: u128,
pub size: u64,
}
#[derive(Clone, Debug)]
pub struct SourceFile {
pub absolute_path: PathBuf,
pub relative_path: String,
pub stamp: FileStamp,
pub content_type: ContentType,
}
pub fn discover_files(
root: &Path,
selected: &[ContentType],
max_bytes: u64,
) -> Result<Vec<SourceFile>> {
let mut builder = WalkBuilder::new(root);
builder
.hidden(false)
.follow_links(false)
.git_ignore(true)
.git_exclude(true)
.parents(true)
.add_custom_ignore_filename(".sembleignore")
.filter_entry(|entry| {
entry
.file_name()
.to_str()
.map(|name| {
!entry.file_type().is_some_and(|kind| kind.is_dir())
|| !IGNORED_DIRS.contains(&name)
})
.unwrap_or(true)
});
let (sender, receiver) = mpsc::channel();
builder.build_parallel().run(|| {
let sender = sender.clone();
Box::new(move |entry| {
let result = entry
.map_err(|error| Error::InvalidRequest(error.to_string()))
.and_then(|entry| source_file(root, selected, max_bytes, entry));
let _ = sender.send(result);
WalkState::Continue
})
});
drop(sender);
let mut files = Vec::new();
for result in receiver {
if let Some(file) = result? {
files.push(file);
}
}
files.sort_by(|left, right| left.relative_path.cmp(&right.relative_path));
Ok(files)
}
fn source_file(
root: &Path,
selected: &[ContentType],
max_bytes: u64,
entry: ignore::DirEntry,
) -> Result<Option<SourceFile>> {
if !entry.file_type().is_some_and(|kind| kind.is_file()) {
return Ok(None);
}
let path = entry.path();
let Some(content_type) = content_type_for_path(path) else {
return Ok(None);
};
if !selected.contains(&content_type) {
return Ok(None);
}
let metadata = fs::metadata(path).map_err(|error| Error::io(path, error))?;
if metadata.len() == 0 || metadata.len() > max_bytes {
return Ok(None);
}
let relative_path = path
.strip_prefix(root)
.map_err(|_| Error::UnsafePath(path.to_path_buf()))?;
let modified_ns = metadata
.modified()
.ok()
.and_then(|time| time.duration_since(UNIX_EPOCH).ok())
.map_or(0, |duration| duration.as_nanos());
Ok(Some(SourceFile {
absolute_path: path.to_path_buf(),
relative_path: relative_path.to_string_lossy().replace('\\', "/"),
stamp: FileStamp {
modified_ns,
size: metadata.len(),
},
content_type,
}))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn discovery_obeys_content_scope_and_ignore_files() {
let directory = tempfile::tempdir().unwrap();
fs::create_dir_all(directory.path().join("src")).unwrap();
fs::create_dir_all(directory.path().join("target")).unwrap();
fs::write(directory.path().join("src/lib.rs"), "pub fn visible() {}\n").unwrap();
fs::write(
directory.path().join("src/ignored.rs"),
"pub fn ignored() {}\n",
)
.unwrap();
fs::write(directory.path().join("README.md"), "# docs\n").unwrap();
fs::write(
directory.path().join("target/generated.rs"),
"fn generated() {}\n",
)
.unwrap();
fs::write(directory.path().join(".sembleignore"), "src/ignored.rs\n").unwrap();
let code = discover_files(directory.path(), &[ContentType::Code], 1024).unwrap();
assert_eq!(
code.iter()
.map(|file| file.relative_path.as_str())
.collect::<Vec<_>>(),
vec!["src/lib.rs"]
);
let docs = discover_files(directory.path(), &[ContentType::Docs], 1024).unwrap();
assert_eq!(docs[0].relative_path, "README.md");
}
#[test]
fn discovery_skips_empty_and_oversized_files() {
let directory = tempfile::tempdir().unwrap();
fs::write(directory.path().join("empty.rs"), "").unwrap();
fs::write(directory.path().join("large.rs"), "0123456789").unwrap();
assert!(discover_files(directory.path(), &[ContentType::Code], 5)
.unwrap()
.is_empty());
}
}
+9
View File
@@ -0,0 +1,9 @@
//! Source discovery, ignore handling, path safety, and remote Git acquisition.
mod local;
mod remote_git;
mod security;
pub use local::{discover_files, FileStamp, SourceFile};
pub use remote_git::RemoteRepository;
pub use security::canonical_source_root;
@@ -0,0 +1,87 @@
//! Restricted HTTPS Git acquisition for explicitly requested remote repositories.
use std::{
path::{Path, PathBuf},
process::{Command, Stdio},
};
use sha2::{Digest, Sha256};
use url::Url;
use crate::{Error, Result};
pub struct RemoteRepository {
pub path: PathBuf,
pub identity: String,
}
impl RemoteRepository {
pub fn acquire(url: &str, cache_root: &Path) -> Result<Self> {
let parsed = Url::parse(url).map_err(|_| Error::UnsupportedUrl(url.to_owned()))?;
if !matches!(parsed.scheme(), "https" | "http") || parsed.host_str().is_none() {
return Err(Error::UnsupportedUrl(url.to_owned()));
}
let key = hex::encode(Sha256::digest(url.as_bytes()));
let repos = cache_root.join("repos");
std::fs::create_dir_all(&repos).map_err(|error| Error::io(&repos, error))?;
let path = repos.join(key);
if !path.join(".git").is_dir() {
let temporary =
repos.join(format!(".clone-{}-{}", std::process::id(), random_suffix()));
let result = Command::new("git")
.args(["clone", "--depth", "1", "--", url])
.arg(&temporary)
.stdin(Stdio::null())
.output()
.map_err(|error| Error::Git(error.to_string()))?;
if !result.status.success() {
let _ = std::fs::remove_dir_all(&temporary);
return Err(Error::Git(
String::from_utf8_lossy(&result.stderr).trim().to_owned(),
));
}
std::fs::rename(&temporary, &path).map_err(|error| Error::io(&path, error))?;
}
let output = Command::new("git")
.args(["-C"])
.arg(&path)
.args(["rev-parse", "HEAD"])
.stdin(Stdio::null())
.output()
.map_err(|error| Error::Git(error.to_string()))?;
if !output.status.success() {
return Err(Error::Git(
String::from_utf8_lossy(&output.stderr).trim().to_owned(),
));
}
let revision = String::from_utf8_lossy(&output.stdout).trim().to_owned();
Ok(Self {
path,
identity: format!("{url}@{revision}"),
})
}
}
fn random_suffix() -> u128 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |value| value.as_nanos())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn remote_sources_reject_non_http_urls_before_git_runs() {
let directory = tempfile::tempdir().unwrap();
assert!(matches!(
RemoteRepository::acquire("file:///tmp/repository", directory.path()),
Err(Error::UnsupportedUrl(_))
));
assert!(matches!(
RemoteRepository::acquire("git@github.com:owner/repo.git", directory.path()),
Err(Error::UnsupportedUrl(_))
));
}
}
+15
View File
@@ -0,0 +1,15 @@
//! Canonical path validation that keeps indexing inside an allowed source root.
use std::path::{Path, PathBuf};
use crate::{Error, Result};
pub fn canonical_source_root(path: &Path) -> Result<PathBuf> {
if !path.exists() {
return Err(Error::SourceMissing(path.to_path_buf()));
}
if !path.is_dir() {
return Err(Error::SourceNotDirectory(path.to_path_buf()));
}
path.canonicalize().map_err(|error| Error::io(path, error))
}
+89
View File
@@ -0,0 +1,89 @@
//! Public request, result, source, chunk, and persisted-index data types.
use std::path::PathBuf;
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ContentType {
#[default]
Code,
Docs,
Config,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct Chunk {
pub file_path: String,
pub start_line: usize,
pub end_line: usize,
pub language: Option<String>,
/// Source text is runtime-only; persisted indexes reload it on demand for returned snippets.
#[serde(skip, default)]
pub content: String,
}
#[derive(Clone, Debug)]
pub struct SearchRequest {
pub query: String,
pub repo: PathBuf,
pub top_k: usize,
pub max_snippet_lines: Option<usize>,
pub content: Vec<ContentType>,
}
impl SearchRequest {
pub fn new(query: impl Into<String>, repo: impl Into<PathBuf>) -> Self {
Self {
query: query.into(),
repo: repo.into(),
top_k: 5,
max_snippet_lines: Some(10),
content: vec![ContentType::Code],
}
}
}
#[derive(Clone, Debug)]
pub struct FindRelatedRequest {
pub repo: PathBuf,
pub file_path: String,
pub line: usize,
pub top_k: usize,
pub max_snippet_lines: Option<usize>,
pub content: Vec<ContentType>,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct SearchResult {
pub file_path: String,
pub start_line: usize,
pub end_line: usize,
pub score: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct SearchResponse {
pub query: String,
pub results: Vec<SearchResult>,
}
/// Summary of a prepared repository index, independent of its storage format.
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct IndexStats {
pub file_count: usize,
pub chunk_count: usize,
pub source_bytes: u64,
pub dimensions: usize,
}
pub(crate) fn snippet(content: &str, lines: Option<usize>) -> Option<String> {
match lines {
Some(0) => None,
Some(limit) => Some(content.lines().take(limit).collect::<Vec<_>>().join("\n")),
None => Some(content.to_owned()),
}
}