80 lines
2.9 KiB
Rust
80 lines
2.9 KiB
Rust
|
use anyhow::Result;
|
||
|
|
use searchi_core::SearchIndex;
|
||
|
|
use searchi_service::{HybridSearch, SearchOptions};
|
||
|
|
use std::fs;
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
#[ignore = "Requires explicitly provisioned Google model and CPU worker"]
|
||
|
|
fn real_google_search_persists_and_suppresses_changed_deleted_sources() -> Result<()> {
|
||
|
|
assert_eq!(
|
||
|
|
std::env::var("SEARCHI_FORCE_CPU").unwrap(),
|
||
|
|
"1",
|
||
|
|
"VM tests require CPU only"
|
||
|
|
);
|
||
|
|
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 temp = tempfile::tempdir()?;
|
||
|
|
let root = temp.path().join("documents");
|
||
|
|
fs::create_dir(&root)?;
|
||
|
|
let db = temp.path().join("index.db");
|
||
|
|
fs::write(
|
||
|
|
root.join("agreement.txt"),
|
||
|
|
"The supplier partnership may be ended with sixty days written notice.",
|
||
|
|
)?;
|
||
|
|
fs::write(
|
||
|
|
root.join("garden.txt"),
|
||
|
|
"Tomatoes need regular watering and lots of sunlight.",
|
||
|
|
)?;
|
||
|
|
let mut lexical = SearchIndex::open(&db)?;
|
||
|
|
lexical.index_folder(&root)?;
|
||
|
|
let mut hybrid = HybridSearch::open(&db, &model, &worker, 1)?;
|
||
|
|
let report = hybrid.index_embeddings()?;
|
||
|
|
assert_eq!(report.passages, 2);
|
||
|
|
let hits = hybrid.search(
|
||
|
|
"Wie beenden wir die Zusammenarbeit mit dem Lieferanten?",
|
||
|
|
10,
|
||
|
|
&SearchOptions::default(),
|
||
|
|
)?;
|
||
|
|
assert_eq!(hits[0].filename, "agreement.txt");
|
||
|
|
assert!(hits[0].source.is_some());
|
||
|
|
assert!(hits[0].snippet.contains("sixty days"));
|
||
|
|
let exact = hybrid.search("garden.txt", 10, &SearchOptions::default())?;
|
||
|
|
assert_eq!(exact[0].match_kind, "exact_filename");
|
||
|
|
let restricted = hybrid.search(
|
||
|
|
"Wie beenden wir die Zusammenarbeit mit dem Lieferanten?",
|
||
|
|
10,
|
||
|
|
&SearchOptions {
|
||
|
|
extension: Some("pdf".into()),
|
||
|
|
..Default::default()
|
||
|
|
},
|
||
|
|
)?;
|
||
|
|
assert!(restricted.is_empty());
|
||
|
|
drop(hybrid);
|
||
|
|
let mut hybrid = HybridSearch::open(&db, &model, &worker, 1)?;
|
||
|
|
assert_eq!(
|
||
|
|
hybrid.search(
|
||
|
|
"How can we stop working with the vendor?",
|
||
|
|
10,
|
||
|
|
&SearchOptions::default()
|
||
|
|
)?[0]
|
||
|
|
.filename,
|
||
|
|
"agreement.txt"
|
||
|
|
);
|
||
|
|
fs::write(root.join("agreement.txt"),"Entirely unrelated replacement content about repainting a kitchen and changing cupboard handles.")?;
|
||
|
|
lexical.index_folder(&root)?;
|
||
|
|
assert!(hybrid
|
||
|
|
.search_semantic(
|
||
|
|
"How can we stop working with the vendor?",
|
||
|
|
10,
|
||
|
|
&SearchOptions::default()
|
||
|
|
)?
|
||
|
|
.iter()
|
||
|
|
.all(|hit| hit.filename != "agreement.txt"));
|
||
|
|
fs::remove_file(root.join("garden.txt"))?;
|
||
|
|
lexical.index_folder(&root)?;
|
||
|
|
assert!(hybrid
|
||
|
|
.search_semantic("How do we water vegetables?", 10, &SearchOptions::default())?
|
||
|
|
.iter()
|
||
|
|
.all(|hit| hit.filename != "garden.txt"));
|
||
|
|
Ok(())
|
||
|
|
}
|