Files
searchi/crates/searchi-service/src/lib.rs
T

818 lines
29 KiB
Rust
Raw Normal View History

//! Shared offline hybrid search using real Google embeddings and persisted,
//! filtered HNSW retrieval. Production-scale validation remains a release gate.
mod coordination;
mod media;
pub use coordination::{
BudgetedProvider, InferenceBudget, InferencePermit, InferenceRole, PendingForeground,
};
mod prepared;
pub use prepared::PreparedGraphs;
mod pipeline;
mod routing;
pub use routing::SemanticOrigin;
mod store;
pub use searchi_media::{FfmpegTools, MediaLimits};
use anyhow::{ensure, Result};
use searchi_core::metrics::{MetricsMode, SpanToken, TraceCollector};
use searchi_core::SearchIndex;
pub use searchi_core::{SearchMetrics, SearchOptions, SearchResponse, SearchResult};
pub use searchi_semantic::WorkerInfo;
use searchi_semantic::{GemmaEmbedder, TextWindow, GEMMA_DIMENSIONS, GEMMA_FINGERPRINT};
use serde::Serialize;
use std::{
collections::{HashMap, VecDeque},
path::Path,
};
use store::{Store, MAX_SCAN_CHUNKS};
const MAX_BATCH: usize = 8;
const QUERY_CACHE_ENTRIES: usize = 32;
const BYTE_SLICE: usize = 64 * 1024;
const BYTE_OVERLAP: usize = 512;
const MAX_PASSAGE_WINDOWS: usize = 2048;
const CHUNK_POLICY: &str = "byte64k-overlap512-title512-v1";
fn generation_fingerprint(model: &str) -> String {
format!(
"{model}:{CHUNK_POLICY}:{}:image280-frame70-audio30s-overlap2-distributed240-scene120-v1",
searchi_semantic::GEMMA_MEDIA_FINGERPRINT
)
}
fn reusable_policy(expected: &str, active: &str) -> bool {
if active == expected || active == format!("{expected}:text-only") {
return true;
}
let Some(policy) = active.strip_prefix(&format!("{expected}:custom-")) else {
return false;
};
let values = policy
.split('-')
.map(str::parse::<u64>)
.collect::<std::result::Result<Vec<_>, _>>();
matches!(values,Ok(v) if v.len()==8 && v[0]>=2 && (5000..=30000).contains(&v[3]) && v[4]<v[3] && v[5]>0 && v[6]<=u32::MAX as u64)
}
/// Model-specific query/document prompts and actual tokenizer windows belong
/// to the provider. Storage and ranking never assume a particular dimension.
pub trait EmbeddingProvider: Send {
/// Maximum document batch admitted as one inference call. Coordinated
/// background sessions yield at every document boundary.
fn document_batch_limit(&self) -> usize {
MAX_BATCH
}
fn fingerprint(&self) -> &str;
fn dimensions(&self) -> usize;
fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>>;
fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>>;
fn embed_query(&mut self, query: &str) -> Result<Vec<f32>>;
fn embed_code_query(&mut self, _query: &str) -> Result<Vec<f32>> {
anyhow::bail!("Provider does not support Google code retrieval")
}
fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
self.embed_code_query(query)
}
fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
self.embed_query(query)
}
fn embed_query_traced(
&mut self,
query: &str,
debug: bool,
_trace: &mut TraceCollector,
_parent: Option<SpanToken>,
) -> Result<Vec<f32>> {
if debug {
self.embed_query_debug(query)
} else {
self.embed_query(query)
}
}
fn embed_code_query_traced(
&mut self,
query: &str,
debug: bool,
_trace: &mut TraceCollector,
_parent: Option<SpanToken>,
) -> Result<Vec<f32>> {
if debug {
self.embed_code_query_debug(query)
} else {
self.embed_code_query(query)
}
}
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
Default::default()
}
fn embed_images(&mut self, _paths: &[&Path], _budget: u32) -> Result<Vec<Vec<f32>>> {
anyhow::bail!("Provider does not support real images")
}
fn embed_audio(&mut self, _pcm: &[f32]) -> Result<Vec<f32>> {
anyhow::bail!("Provider does not support real audio")
}
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
None
}
}
impl EmbeddingProvider for GemmaEmbedder {
fn fingerprint(&self) -> &str {
GEMMA_FINGERPRINT
}
fn dimensions(&self) -> usize {
GEMMA_DIMENSIONS
}
fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
self.split_windows(text, title)
}
fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
self.embed_documents(texts, title)
}
fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
self.embed_query(query)
}
fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
GemmaEmbedder::embed_query_debug(self, query)
}
fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
GemmaEmbedder::embed_code_query(self, query)
}
fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
GemmaEmbedder::embed_code_query_debug(self, query)
}
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
GemmaEmbedder::query_diagnostics(self)
}
fn embed_images(&mut self, paths: &[&Path], budget: u32) -> Result<Vec<Vec<f32>>> {
self.embed_images_with_budget(paths, budget)
}
fn embed_audio(&mut self, pcm: &[f32]) -> Result<Vec<f32>> {
self.embed_audio_pcm(pcm)
}
fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
Some(self.info.clone())
}
}
#[derive(Debug, Serialize)]
pub struct MediaFailure {
pub path: String,
pub message: String,
}
#[derive(Debug, Serialize)]
pub struct EmbeddingReport {
pub media_files: usize,
pub media_failures: usize,
pub failures: Vec<MediaFailure>,
pub passages: usize,
pub chunks: usize,
pub generation: i64,
pub model_fingerprint: String,
/// Completed groups copied or resumed without document/media inference.
pub reused_sources: usize,
pub reused_chunks: usize,
pub encoded_sources: usize,
pub document_inference_calls: usize,
pub image_inference_calls: usize,
pub audio_inference_calls: usize,
pub resumed: bool,
}
#[derive(Debug, Serialize)]
pub struct EmbeddingStatus {
pub failures: Vec<MediaFailure>,
pub generation: Option<i64>,
pub chunks: usize,
pub model_fingerprint: Option<String>,
pub stale: bool,
/// Rebuild is generation-based: all current passages are pending when stale.
pub pending: usize,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum QueryTask {
Search,
Code,
}
/// Shared local hybrid service. Google EmbeddingGemma 2 is the default and
/// only shipped provider; no alternative model is silently substituted.
pub struct HybridSearch {
lexical: SearchIndex,
embedder: Box<dyn EmbeddingProvider>,
store: Store,
generation_fingerprint: String,
// Real provider outputs only, scoped to this model session. Queries never
// persist and results/filters are never cached here.
query_cache_fingerprint: String,
query_cache: VecDeque<(QueryTask, String, Vec<f32>)>,
}
impl HybridSearch {
/// Attach shared admission without replacing the warmed model or query LRU.
pub fn with_budget(
mut self,
budget: &InferenceBudget,
role: InferenceRole,
actual_worker_threads: usize,
) -> Self {
self.embedder = Box::new(BudgetedProvider::new(
self.embedder,
budget,
role,
actual_worker_threads,
));
self
}
/// Connect background publication to an existing warmed reader.
pub fn share_prepared_graphs(&mut self, graphs: &PreparedGraphs) -> Result<()> {
self.store.share_prepared_graphs(graphs)
}
/// Actual model/runtime/device reported by the persistent native worker.
pub fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
self.embedder.worker_info()
}
/// Cheap local status check. No model, runtime worker, or network is loaded.
/// Corpus revision triggers detect file inserts, updates, and deletions.
pub fn embedding_status(db: impl AsRef<Path>) -> Result<EmbeddingStatus> {
let connection =
rusqlite::Connection::open_with_flags(db, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?;
connection.busy_timeout(std::time::Duration::from_secs(5))?;
let exists: bool = connection.query_row(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE name='semantic_meta')",
[],
|row| row.get(0),
)?;
let pending: usize =
connection.query_row("SELECT (SELECT count(*) FROM passages)+(SELECT count(*) FROM files WHERE extension IN ('png','jpg','jpeg','webp','bmp','mp4','mkv','mov','avi','webm','m4v','mpeg','mpg','mp3','flac','wav','m4a','ogg','opus','aac','aiff','aif'))", [], |row| row.get(0))?;
if !exists {
return Ok(EmbeddingStatus {
failures: Vec::new(),
generation: None,
chunks: 0,
model_fingerprint: None,
stale: true,
pending,
});
}
let version: i64 = connection.query_row(
"SELECT schema_version FROM semantic_meta WHERE singleton=1",
[],
|row| row.get(0),
)?;
if version < 6 {
return Ok(EmbeddingStatus {
failures: Vec::new(),
generation: None,
chunks: 0,
model_fingerprint: None,
stale: true,
pending,
});
}
ensure!(
(6..=7).contains(&version),
"Unsupported semantic schema version {version}"
);
let (generation,fingerprint,revision,current):(Option<i64>,Option<String>,Option<i64>,i64)=connection.query_row(
"SELECT g.id,g.fingerprint,CASE WHEN g.ann_blob IS NOT NULL AND g.dimensions=768 THEN g.corpus_revision END,m.corpus_revision FROM semantic_meta m LEFT JOIN semantic_generations g ON g.id=m.active_generation WHERE m.singleton=1",[],
|row| Ok((row.get(0)?,row.get(1)?,row.get(2)?,row.get(3)?)))?;
let chunks: usize = connection.query_row(
"SELECT count(*) FROM semantic_chunks WHERE generation=?1",
[generation],
|row| row.get(0),
)?;
let mut failures_query=connection.prepare("SELECT path,message FROM semantic_media_errors WHERE generation=?1 ORDER BY path LIMIT 1000")?;
let failures = failures_query
.query_map([generation], |r| {
Ok(MediaFailure {
path: r.get(0)?,
message: r.get(1)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
let stale = revision != Some(current)
|| fingerprint.as_deref() != Some(generation_fingerprint(GEMMA_FINGERPRINT).as_str());
Ok(EmbeddingStatus {
failures,
generation,
chunks,
model_fingerprint: fingerprint,
stale,
pending: if stale { pending } else { 0 },
})
}
pub fn open(
db: impl AsRef<Path>,
model_dir: impl AsRef<Path>,
runtime_library: impl AsRef<Path>,
threads: usize,
) -> Result<Self> {
let embedder = GemmaEmbedder::load(model_dir, runtime_library, threads)?;
Self::with_provider(db, Box::new(embedder))
}
/// Explicit provider injection for future model generations, never fallback.
pub fn with_provider(
db: impl AsRef<Path>,
embedder: Box<dyn EmbeddingProvider>,
) -> Result<Self> {
let lexical = SearchIndex::open(&db)?;
let store = Store::open(db.as_ref(), embedder.dimensions())?;
let expected = generation_fingerprint(embedder.fingerprint());
let active = store.active_fingerprint()?;
let generation_fingerprint = match active {
Some(active) if reusable_policy(&expected, &active) => active,
_ => expected,
};
Ok(Self {
query_cache_fingerprint: embedder.fingerprint().to_owned(),
query_cache: VecDeque::new(),
lexical,
embedder,
store,
generation_fingerprint,
})
}
/// Execute real query inference and load the active ANN generation before
/// foreground use. Call on a background worker after assets are verified.
pub fn warmup(&mut self) -> Result<()> {
self.lexical.prewarm()?;
let vector = self.embedder.embed_query("warm up local file search")?;
let _ = self.store.search(
&vector,
&self.generation_fingerprint,
1,
&SearchOptions::default(),
)?;
Ok(())
}
/// Build or resume a replacement generation in bounded batches. Completed
/// groups can be reused; failures retain staging work and the ready reader.
pub fn index_embeddings(&mut self) -> Result<EmbeddingReport> {
let fingerprint = format!(
"{}:text-only",
generation_fingerprint(self.embedder.fingerprint())
);
self.rebuild_generation(fingerprint, None, |service, generation, work| {
service.build(generation, work)
})
}
/// Filter each retrieval lane before ranking, preserve filename priority,
/// then fuse keyword and semantic ranks with reciprocal-rank fusion (k=60).
pub fn search(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
let mut trace = TraceCollector::new(MetricsMode::Basic);
self.search_traced(query, limit, options, &mut trace, None)
}
pub fn search_measured(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
debug: bool,
) -> Result<SearchResponse<SearchResult>> {
let mut trace = TraceCollector::new(if debug {
MetricsMode::Debug
} else {
MetricsMode::Basic
});
let results = self.search_traced(query, limit, options, &mut trace, None)?;
let metrics = trace.finish(results.len());
Ok(SearchResponse { results, metrics })
}
pub fn search_traced(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
trace: &mut TraceCollector,
parent: Option<SpanToken>,
) -> Result<Vec<SearchResult>> {
trace.timed("hybrid.search", parent, |trace, parent| {
self.search_inner(query, limit, options, trace, parent)
})
}
fn search_inner(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
trace: &mut TraceCollector,
parent: Option<SpanToken>,
) -> Result<Vec<SearchResult>> {
let candidates = limit.saturating_mul(4).clamp(20, 1000);
// Core validates filters and query budgets even for empty requests.
let mut lexical = self
.lexical
.search_with_options_traced(query, candidates, options, trace, parent)?;
if query.trim().is_empty() || limit == 0 {
return Ok(Vec::new());
}
trace.timed("hybrid.identifier_priority", parent, |_, _| {
promote_identifiers(query.trim(), &mut lexical)
});
let vector = self.query_vector(query, options.code_only, trace, parent)?;
let semantic =
self.routed_semantic_traced(query, &vector, candidates, options, trace, parent)?;
let results = trace.timed("hybrid.rank_fusion", parent, |trace, span| {
trace.counter(span, "lexical_files", lexical.len() as u64);
trace.counter(span, "semantic_files", semantic.len() as u64);
fuse(lexical, semantic, limit.min(1000))
});
Ok(results)
}
/// Semantic lane for reproducible relevance evaluation and diagnostics.
pub fn search_semantic(
&mut self,
query: &str,
limit: usize,
options: &SearchOptions,
) -> Result<Vec<SearchResult>> {
let _ = self.lexical.search_with_options(query, 0, options)?;
if query.trim().is_empty() || limit == 0 {
return Ok(Vec::new());
}
let mut trace = TraceCollector::new(MetricsMode::Basic);
let vector = self.query_vector(query, options.code_only, &mut trace, None)?;
self.routed_semantic(query, &vector, limit.min(1000), options)
}
fn query_vector(
&mut self,
query: &str,
code: bool,
trace: &mut TraceCollector,
parent: Option<SpanToken>,
) -> Result<Vec<f32>> {
let query = query.trim();
ensure!(
!query.is_empty() && query.len() <= 4096,
"Query must be nonempty and at most 4096 UTF-8 bytes"
);
let span = trace.begin("semantic.query_embedding", parent);
let task = if code {
QueryTask::Code
} else {
QueryTask::Search
};
trace.counter(span, "code_task", u64::from(code));
if self.query_cache_fingerprint != self.embedder.fingerprint() {
self.query_cache.clear();
self.query_cache_fingerprint = self.embedder.fingerprint().to_owned();
}
if let Some(position) = self
.query_cache
.iter()
.position(|(cached_task, cached, _)| *cached_task == task && cached == query)
{
let entry = self
.query_cache
.remove(position)
.expect("Located cache entry");
let vector = entry.2.clone();
self.query_cache.push_back(entry);
trace.counter(span, "query_cache_hit", 1);
trace.counter(span, "query_cache_miss", 0);
trace.counter(span, "failed", 0);
// No worker was called: never attach its previous diagnostics.
trace.end(span, Default::default());
return Ok(vector);
}
trace.counter(span, "query_cache_hit", 0);
trace.counter(span, "query_cache_miss", 1);
let vector = (|| {
let vector = if code {
self.embedder
.embed_code_query_traced(query, trace.is_debug(), trace, span)?
} else {
self.embedder
.embed_query_traced(query, trace.is_debug(), trace, span)?
};
ensure!(
vector.len() == self.embedder.dimensions()
&& vector.iter().all(|value| value.is_finite())
&& vector.iter().any(|value| *value != 0.0),
"Provider returned an invalid query vector"
);
Ok(vector)
})();
trace.counter(span, "failed", u64::from(vector.is_err()));
trace.end(
span,
if trace.is_debug() {
self.embedder.query_diagnostics()
} else {
Default::default()
},
);
let vector = vector?;
if self.query_cache.len() == QUERY_CACHE_ENTRIES {
self.query_cache.pop_front();
}
self.query_cache
.push_back((task, query.to_owned(), vector.clone()));
Ok(vector)
}
}
fn bounded_title(title: &str) -> &str {
let mut end = title.len().min(512);
while !title.is_char_boundary(end) {
end -= 1;
}
&title[..end]
}
fn byte_slices(text: &str) -> Result<Vec<(usize, usize)>> {
ensure!(
text.len() <= 2 * 1024 * 1024,
"Passage exceeds 2 MiB service byte budget"
);
let mut slices = Vec::new();
let mut start = 0;
while start < text.len() {
let mut end = (start + BYTE_SLICE).min(text.len());
while !text.is_char_boundary(end) {
end -= 1;
}
ensure!(end > start, "UTF-8 slice made no progress");
slices.push((start, end));
if end == text.len() {
break;
}
start = end.saturating_sub(BYTE_OVERLAP);
while !text.is_char_boundary(start) {
start -= 1;
}
}
Ok(slices)
}
fn passage_windows(
provider: &mut dyn EmbeddingProvider,
text: &str,
title: &str,
) -> Result<Vec<TextWindow>> {
let mut windows = Vec::new();
for (start, end) in byte_slices(text)? {
let fragment = &text[start..end];
if fragment.trim().is_empty() {
continue;
}
let mut slice = provider.split_windows(fragment, title)?;
for window in &mut slice {
window.start_byte += start;
window.end_byte += start;
}
windows.extend(slice);
ensure!(
windows.len() <= MAX_PASSAGE_WINDOWS,
"Passage exceeds {MAX_PASSAGE_WINDOWS} semantic window budget"
);
}
Ok(windows)
}
fn promote_identifiers(query: &str, results: &mut [SearchResult]) {
if query.chars().any(char::is_whitespace)
|| !query.chars().any(|c| c.is_ascii_digit() || c == '_')
{
return;
}
for result in results {
if result.match_kind == "content" && literal_identifier_match(&result.snippet, query) {
result.match_kind = "exact_identifier".into();
}
}
}
fn literal_identifier_match(text: &str, query: &str) -> bool {
text.match_indices(query).any(|(start, _)| {
let end = start + query.len();
let is_identifier = |c: char| c.is_alphanumeric() || c == '_' || c == '-';
!text[..start].chars().next_back().is_some_and(is_identifier)
&& !text[end..].chars().next().is_some_and(is_identifier)
})
}
fn fuse(
lexical: Vec<SearchResult>,
semantic: Vec<SearchResult>,
limit: usize,
) -> Vec<SearchResult> {
let mut results = HashMap::<String, SearchResult>::new();
for (rank, mut result) in lexical.into_iter().enumerate() {
let weight = if result.match_kind == "filename" {
1.15
} else {
1.0
};
result.score = weight / (61.0 + rank as f64);
results.insert(result.path.clone(), result);
}
for (rank, mut result) in semantic.into_iter().enumerate() {
let contribution = 1.0 / (61.0 + rank as f64);
if let Some(existing) = results.get_mut(&result.path) {
existing.score += contribution;
if existing.match_kind == "content" {
existing.match_kind = "hybrid".into();
}
} else {
result.score = contribution;
results.insert(result.path.clone(), result);
}
}
let tier = |result: &SearchResult| match result.match_kind.as_str() {
"exact_filename" => 0,
"exact_identifier" => 1,
_ => 2,
};
let mut results: Vec<_> = results.into_values().collect();
results.sort_by(|a, b| {
tier(a)
.cmp(&tier(b))
.then_with(|| b.score.total_cmp(&a.score))
.then_with(|| a.path.cmp(&b.path))
});
results.truncate(limit);
results
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn failed_query_inference_retains_diagnostics_in_shared_trace() -> Result<()> {
// Failure-only provider: no fabricated embeddings are generated.
struct FailingProvider;
impl EmbeddingProvider for FailingProvider {
fn fingerprint(&self) -> &str {
"failure-injection"
}
fn dimensions(&self) -> usize {
768
}
fn split_windows(&mut self, _: &str, _: &str) -> Result<Vec<TextWindow>> {
anyhow::bail!("unused")
}
fn embed_documents(&mut self, _: &[&str], _: &str) -> Result<Vec<Vec<f32>>> {
anyhow::bail!("unused")
}
fn embed_query(&mut self, _: &str) -> Result<Vec<f32>> {
anyhow::bail!("injected worker failure")
}
fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
std::collections::BTreeMap::from([("request_write_us".into(), 7)])
}
}
let directory = tempfile::tempdir()?;
let db = directory.path().join("index.db");
let mut service = HybridSearch::with_provider(&db, Box::new(FailingProvider))?;
let mut trace = TraceCollector::new(MetricsMode::Debug);
assert!(service
.search_traced(
"ordinary query",
5,
&SearchOptions::default(),
&mut trace,
None
)
.is_err());
let metrics = trace.finish(0);
let spans = metrics.trace.unwrap();
let embedding = spans
.iter()
.find(|s| s.name == "semantic.query_embedding")
.unwrap();
assert_eq!(embedding.counters["failed"], 1);
assert_eq!(embedding.counters["request_write_us"], 7);
let parent = &spans[embedding.parent.unwrap()];
assert!(
embedding.start_ms + embedding.duration_ms
<= parent.start_ms + parent.duration_ms + 1e-6
);
Ok(())
}
fn result(path: &str, kind: &str) -> SearchResult {
SearchResult {
path: path.into(),
filename: path.into(),
snippet: "evidence".into(),
score: 1.0,
match_kind: kind.into(),
source: None,
extraction_error: None,
}
}
#[test]
fn exact_filename_wins_even_when_another_result_matches_both_lanes() {
let results = fuse(
vec![
result("invoice.pdf", "exact_filename"),
result("notes.txt", "content"),
],
vec![
result("notes.txt", "semantic"),
result("another.txt", "semantic"),
],
10,
);
assert_eq!(results[0].path, "invoice.pdf");
assert_eq!(results[1].match_kind, "hybrid");
assert_eq!(results.len(), 3);
}
#[test]
fn identifiers_use_literal_boundaries_and_survive_semantic_competition() {
let mut lexical = vec![result("support.txt", "content")];
lexical[0].snippet = "Error E_SEARCH_071 means upgrade the app".into();
promote_identifiers("E_SEARCH_071", &mut lexical);
let results = fuse(lexical, vec![result("wrong.txt", "semantic")], 10);
assert_eq!(results[0].match_kind, "exact_identifier");
assert!(!literal_identifier_match(
"E_SEARCH_071_EXTRA",
"E_SEARCH_071"
));
}
#[test]
fn keyword_semantic_consensus_beats_weak_partial_filename() {
let results = fuse(
vec![
result("meeting-cancellation.txt", "filename"),
result("supplier-agreement.txt", "content"),
],
vec![result("supplier-agreement.txt", "semantic")],
10,
);
assert_eq!(results[0].path, "supplier-agreement.txt");
assert_eq!(results[0].match_kind, "hybrid");
assert_eq!(results[1].match_kind, "filename");
}
#[test]
fn byte_slices_preserve_utf8_overlap_offsets_and_complete_tail() {
let text = format!("{}TAIL", "Ä🐈中文 ".repeat(10000));
let slices = byte_slices(&text).unwrap();
assert!(slices.len() > 1);
assert_eq!(slices[0].0, 0);
assert_eq!(slices.last().unwrap().1, text.len());
for (start, end) in &slices {
assert!(text.is_char_boundary(*start) && text.is_char_boundary(*end));
assert!(end - start <= BYTE_SLICE);
}
for pair in slices.windows(2) {
assert!(pair[1].0 < pair[0].1 && pair[1].1 > pair[0].1);
}
assert!(text[slices.last().unwrap().0..].ends_with("TAIL"));
assert!(byte_slices(&"x".repeat(2 * 1024 * 1024 + 1)).is_err());
let title = "中文".repeat(120);
assert!(bounded_title(&title).len() <= 512 && title.starts_with(bounded_title(&title)));
}
#[test]
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
fn google_cpu_preserves_and_retrieves_content_after_64k_byte_boundary() -> Result<()> {
let model =
std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Set SEARCHI_GEMMA_MODEL_DIR");
let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Set SEARCHI_GEMMA_WORKER");
let mut provider =
GemmaEmbedder::load_with_device(model, worker, 1, searchi_semantic::Device::Cpu)?;
let text=format!("{}\nÜbernachtungskosten auf Dienstreisen werden gegen Vorlage der Hotelrechnung erstattet.","Routine system checks were completed successfully. ".repeat(1500));
assert!(text.len() > BYTE_SLICE);
let windows = passage_windows(&mut provider, &text, "routine-log.txt")?;
let last = windows.last().unwrap();
assert_eq!(last.end_byte, text.len());
assert!(last.start_byte >= BYTE_SLICE - BYTE_OVERLAP);
assert!(text.find("Übernachtungskosten").unwrap() > BYTE_SLICE);
assert!(last.text.contains("Hotelrechnung"));
assert_eq!(
text.get(last.start_byte..last.end_byte),
Some(last.text.as_str())
);
let vectors = provider.embed_documents(
&[windows[0].text.as_str(), last.text.as_str()],
"routine-log.txt",
)?;
let query = provider.embed_query(
"Where does it say that hotel stays during work travel can be reimbursed?",
)?;
let cosine = |vector: &[f32]| vector.iter().zip(&query).map(|(a, b)| a * b).sum::<f32>();
assert!(
cosine(&vectors[1]) > cosine(&vectors[0]),
"Real model did not retrieve the tail window"
);
Ok(())
}
}