mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-06 13:44:21 +08:00
refactor: rebuild desktop app with Tauri
This commit is contained in:
+65
@@ -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());
|
||||
}
|
||||
}
|
||||
Vendored
+7
@@ -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
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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()))
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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]);
|
||||
}
|
||||
}
|
||||
@@ -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>;
|
||||
@@ -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)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
};
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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,
|
||||
};
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//! Public search engine, Reciprocal Rank Fusion, and code-aware reranking.
|
||||
|
||||
mod engine;
|
||||
mod rerank;
|
||||
mod rrf;
|
||||
|
||||
pub use engine::SearchEngine;
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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]);
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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(_))
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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()),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user