use crate::error::{PathErrorExt, SsgError};
use crate::plugin::{Plugin, PluginContext};
use crate::search::SearchIndex;
use ssg_search::{
paths::{EMBEDDINGS_FILE, MANIFEST_FILE, MODEL_FILE, TOKENIZER_FILE},
ArtifactsBuilder,
};
use std::fs;
const EXCERPT_CHARS: usize = 160;
#[derive(Debug, Clone, Copy, Default)]
pub struct VectorSearchPlugin;
impl Plugin for VectorSearchPlugin {
fn name(&self) -> &'static str {
"vector-search"
}
fn after_compile(&self, ctx: &PluginContext) -> Result<(), SsgError> {
if !ctx.site_dir.exists() {
return Ok(());
}
if ctx.dry_run {
return Ok(());
}
let index = SearchIndex::build(&ctx.site_dir)?;
if index.is_empty() {
return Ok(());
}
let mut builder = ArtifactsBuilder::default();
for e in &index.entries {
let excerpt = truncate_chars(&e.content, EXCERPT_CHARS);
let _ = builder.add_doc(ssg_search::artifacts::InputDoc {
url: e.url.clone(),
title: e.title.clone(),
body: format!("{}\n{}", e.title, e.content),
excerpt,
});
}
let artifacts = builder.build();
let dir = ctx.site_dir.join("search");
fs::create_dir_all(&dir).with_path(&dir)?;
let emb_path = dir.join(EMBEDDINGS_FILE);
fs::write(&emb_path, &artifacts.embeddings).with_path(&emb_path)?;
let man_path = dir.join(MANIFEST_FILE);
fs::write(&man_path, &artifacts.manifest_json).with_path(&man_path)?;
let model_path = dir.join(MODEL_FILE);
fs::write(&model_path, &artifacts.model).with_path(&model_path)?;
let tok_path = dir.join(TOKENIZER_FILE);
fs::write(&tok_path, &artifacts.tokenizer).with_path(&tok_path)?;
log::info!(
"[vector-search] Wrote {} docs ({} bytes embeddings) to {}",
artifacts.count(),
artifacts.embeddings.len(),
dir.display()
);
Ok(())
}
}
fn truncate_chars(s: &str, max_chars: usize) -> String {
if s.chars().count() <= max_chars {
return s.to_string();
}
s.chars().take(max_chars).collect()
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use std::path::Path;
use tempfile::tempdir;
fn ctx(p: &Path) -> PluginContext {
PluginContext::new(p, p, p, p)
}
fn write_html(dir: &Path, name: &str, body: &str) {
let p = dir.join(name);
fs::write(
&p,
format!(
"<!DOCTYPE html><html><head><title>{name}</title></head><body>{body}</body></html>"
),
)
.unwrap();
}
#[test]
fn plugin_name() {
assert_eq!(VectorSearchPlugin.name(), "vector-search");
}
#[test]
fn plugin_is_noop_on_missing_dir() {
let tmp = tempdir().unwrap();
let missing = tmp.path().join("not-here");
let c = ctx(&missing);
VectorSearchPlugin.after_compile(&c).unwrap();
assert!(!missing.exists());
}
#[test]
fn plugin_is_noop_in_dry_run() {
let tmp = tempdir().unwrap();
write_html(tmp.path(), "page.html", "<p>hello</p>");
let c = ctx(tmp.path()).with_dry_run(true);
VectorSearchPlugin.after_compile(&c).unwrap();
assert!(!tmp.path().join("search").exists());
}
#[test]
fn plugin_is_noop_on_empty_corpus() {
let tmp = tempdir().unwrap();
let c = ctx(tmp.path());
VectorSearchPlugin.after_compile(&c).unwrap();
assert!(!tmp.path().join("search").exists());
}
#[test]
fn plugin_emits_four_artifacts() {
let tmp = tempdir().unwrap();
write_html(
tmp.path(),
"a.html",
"<p>rust webassembly compiles to portable browser modules</p>",
);
write_html(
tmp.path(),
"b.html",
"<p>baking sourdough bread starter flour</p>",
);
let c = ctx(tmp.path());
VectorSearchPlugin.after_compile(&c).unwrap();
let dir = tmp.path().join("search");
assert!(dir.exists());
assert!(dir.join("embeddings.bin").exists());
assert!(dir.join("manifest.json").exists());
assert!(dir.join("model.bin").exists());
assert!(dir.join("tokenizer.bin").exists());
}
#[test]
fn embeddings_size_is_n_x_d_x_4() {
let tmp = tempdir().unwrap();
write_html(tmp.path(), "a.html", "<p>foo bar baz</p>");
write_html(tmp.path(), "b.html", "<p>quux corge grault</p>");
let c = ctx(tmp.path());
VectorSearchPlugin.after_compile(&c).unwrap();
let emb = fs::read(tmp.path().join("search/embeddings.bin")).unwrap();
assert_eq!(emb.len(), 2 * 256 * 4);
}
#[test]
fn manifest_has_one_entry_per_html() {
let tmp = tempdir().unwrap();
write_html(tmp.path(), "x.html", "<p>alpha</p>");
write_html(tmp.path(), "y.html", "<p>beta</p>");
let c = ctx(tmp.path());
VectorSearchPlugin.after_compile(&c).unwrap();
let json = fs::read_to_string(tmp.path().join("search/manifest.json"))
.unwrap();
let m: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(m["count"].as_u64().unwrap(), 2);
assert_eq!(m["entries"].as_array().unwrap().len(), 2);
}
#[test]
fn truncate_chars_respects_unicode_boundaries() {
let s = "résumé café";
let t = truncate_chars(s, 6);
assert_eq!(t, "résumé");
}
#[test]
fn truncate_chars_returns_full_string_when_short() {
let s = "short";
assert_eq!(truncate_chars(s, 100), "short");
}
#[test]
fn embeddings_byte_identical_across_runs() {
let tmp = tempdir().unwrap();
write_html(
tmp.path(),
"a.html",
"<p>identical content produces identical embeddings</p>",
);
let c = ctx(tmp.path());
VectorSearchPlugin.after_compile(&c).unwrap();
let first = fs::read(tmp.path().join("search/embeddings.bin")).unwrap();
VectorSearchPlugin.after_compile(&c).unwrap();
let second =
fs::read(tmp.path().join("search/embeddings.bin")).unwrap();
assert_eq!(first, second);
}
}