Commit
Message
Changed Files (13)
-
modified Cargo.lock
diff --git a/Cargo.lock b/Cargo.lock index e988f80..d0456a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1025,6 +1025,12 @@ dependencies = [ "dtoa", ] +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "ego-tree" version = "0.6.3" @@ -3362,6 +3368,30 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "schemars" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615" +dependencies = [ + "dyn-clone", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -3403,6 +3433,7 @@ dependencies = [ "reqwest", "robotstxt", "rusqlite", + "schemars", "scraper", "serde", "serde_json", @@ -3493,6 +3524,17 @@ dependencies = [ "syn", ] +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "serde_json" version = "1.0.150" -
modified Cargo.toml
diff --git a/Cargo.toml b/Cargo.toml index 1074520..01278f0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ indicatif = "0.17" reqwest = { version = "0.12", default-features = false, features = ["rustls-tls"] } robotstxt = "0.3" rusqlite = { version = "0.32", features = ["bundled", "chrono"] } +schemars = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" tempfile = "3" -
modified README.md
diff --git a/README.md b/README.md index 4b825f3..2f3f860 100644 --- a/README.md +++ b/README.md @@ -1,16 +1,12 @@ # SearchHub -A local search engine for your browser bookmarks and history. Import bookmarks/history from Firefox, Zen, Chrome, or Chromium, search them with full-text queries, and optionally forward searches to external engines like wikipedia or SearXNG (aggregates results from dozens of backends). Content can be automatically tagged via local ONNX embeddings (opt-in; set `tagging_enabled = true` in config). +A local search engine for your browser bookmarks and history. Import bookmarks/history from Firefox, Zen, Chrome, or Chromium, search them with full-text queries, and optionally forward searches to external engines like Wikipedia or SearXNG (aggregates results from dozens of backends). Content can be automatically tagged via a two-stage ONNX pipeline — a fast bi-encoder scores all tags, then a cross-encoder reranker refines the top candidates (opt-in; see auto-tagging section). -The ONNX tagging model is embedded via the optional `tagging` feature (included by default). If you don't need auto-tagging, you can build a ~30% smaller binary by excluding it. +The ONNX models are embedded via the optional `tagging` feature (included by default). If you don't need auto-tagging, you can build a ~30% smaller binary by excluding it. -## Install +## Quick start -**Binaries** are available at [vit.am/~ololduck/search_hub/latest](https://vit.am/~ololduck/search_hub/latest/). Download the binary for your architecture, extract, and run. - -**Source:** Clone the [repository](https://vit.am/~ololduck/search_hub/repository.git) and build with Rust: - -**Prerequisites:** Rust (install via [rustup](https://rustup.rs/)). +**Prerequisites:** [Rust](https://rustup.rs/) (or grab a [pre-built binary](https://vit.am/~ololduck/search_hub/latest/)). ```sh git clone https://vit.am/~ololduck/search_hub/repository.git search_hub @@ -18,46 +14,19 @@ cd search_hub cargo install --path . ``` -This installs the `search_hub` binary to `~/.cargo/bin/search_hub`. - -**Smaller binary (no ONNX tagging):** If you don't need auto-tagging, skip the ~12 MiB ONNX Runtime dependency: +Then import your browser bookmarks and start the web UI: ```sh -cargo install --path . --no-default-features +search_hub import firefox # or "chrome", "zen" +search_hub serve # open http://127.0.0.1:8080 ``` -The resulting binary still supports all features except semantic tagging. Any `search_hub retag` or auto-tagging operation will advise you to rebuild with `--features tagging`. - -To update later, pull the latest code and reinstall. - -## First steps +**Smaller build (no ONNX tagging):** ```sh -# Import bookmarks from Firefox (auto-discovers your profile) -search_hub import firefox - -# Import from Chrome -search_hub import chrome - -# Start the web UI -search_hub serve +cargo install --path . --no-default-features ``` -Open http://127.0.0.1:8080 in your browser. You can now search your bookmarks. - -Search queries are also forwarded to external engines: Wikipedia, [crates.io](https://crates.io) via its public JSON API, and optionally [SearXNG](https://searx.space) (which aggregates Google, Bing, DDG, and dozens more) if `[[engines]]` is configured. Works as a custom search provider in Firefox/Zen via the OpenSearch protocol (your browser should auto-discover it at `/opensearch.xml`). - -When JavaScript is available, search results are progressively loaded: local bookmark results appear instantly, then external engine results stream in as each engine responds. - -The server supports a `?js=true` query parameter: when present, it skips external engines during page render (returning only local results in ~50ms) and lets the client fetch them via a streaming connection. The inline JavaScript automatically adds `?js=true` to the URL, so the fast path is active by default once JS runs. - -For the fastest initial page load from your browser's search bar (before any JS runs), configure your browser's search engine URL with `?js=true` appended: - -- **Firefox:** Settings → Search → Manage Search Engines → Add, use `http://localhost:8080/search?q=%s&js=true` -- **OpenSearch template:** `http://localhost:8080/search?q={searchTerms}&js=true` - -Without `?js=true`, the server waits for all external engines before sending the page — a fallback path for browsers without JavaScript (which still works, just slower on first load). - ## CLI reference | Command | What it does | @@ -71,90 +40,45 @@ Without `?js=true`, the server waits for all external engines before sending the | `search_hub list` | List all bookmarks | | `search_hub insert "Title" https://..."` | Add a bookmark (fetches content, auto-tags if enabled) | | `search_hub remove --id 1` | Delete a bookmark by ID | -| `search_hub retag --all` | Re-run auto-tagging (requires `tagging_enabled = true` in config) | +| `search_hub retag --all` | Re-run auto-tagging (requires `[tagging] enabled = true` in config) | | `search_hub init-config` | Create a default config file at `~/.config/search_hub/config.toml` | +| `search_hub config-schema` | Print the JSON Schema for the config file format | | `search_hub self-update` | Check abbaye Atom feed and update to the latest release | | `search_hub self-update --dry-run` | Check for updates without downloading | | `search_hub self-update --target x86_64-unknown-linux-gnu` | Override the target triple | All commands use `~/.local/share/search_hub/bookmarks.db` by default. Override with `--db-path` or set `db_path` in the config file. -The first time you use a tagging-enabled command (`search_hub insert`, `search_hub retag`, or import with `tagging_enabled = true`), SearchHub downloads an ONNX embedding model to `$XDG_CACHE_DIR` (defaults to `~/.cache/search_hub`) (about 127 MB). This only applies if the binary was built with the `tagging` feature (default: on). +## Web UI -## Configuration +Search queries are forwarded to external engines: Wikipedia, [crates.io](https://crates.io), and optionally [SearXNG](https://searx.space) if configured. Works as a custom search provider in Firefox/Zen via OpenSearch (auto-discovered at `/opensearch.xml`). -Run `search_hub init-config` to create `~/.config/search_hub/config.toml` with all available options commented out. Or create it manually: +When JavaScript is available, results load progressively: local bookmark results appear instantly, then external engine results stream in as each engine responds. The server supports a `?js=true` parameter that skips external engines during the initial page render (~50ms instead of ~1700ms), letting the client fetch them via a streaming connection instead. The inline JavaScript automatically adds `?js=true` to the URL after the first load. -```toml -# Bookmark database path (default: platform data directory) -# db_path = "/home/you/.local/share/search_hub/bookmarks.db" +For the fastest load from your browser's search bar (before JS runs), configure the search engine URL with `?js=true`: + +- **Firefox:** Settings → Search → Manage Search Engines → Add, use `http://localhost:8080/search?q=%s&js=true` +- **OpenSearch template:** `http://localhost:8080/search?q={searchTerms}&js=true` -# Custom tags override the built-in defaults -# [[tags]] -# name = "my-custom-tag" -# examples = ["example text one", "example text two"] +Without `?js=true`, the server waits for all external engines before sending the page — a fallback for browsers without JavaScript (still works, just slower on first load). -# Whether auto-tagging is enabled (default: false, requires ONNX model download on first use) -# tagging_enabled = true +## Configuration -# Minimum confidence for auto-tagging (0.0 to 1.0, default: 0.6) -# tagging_threshold = 0.6 +Run `search_hub init-config` to create `~/.config/search_hub/config.toml` with all available options commented out. The config supports: -# Hosts to skip when fetching content for bookmarking (default: local addresses) -# exclude_urls = ["localhost", "127.0.0.1", "::1"] +- `[tagging]` — auto-tagging settings (enable, threshold, model selection, reranker) +- `[[engines]]` — external search engines (SearXNG, Wikipedia, MDN, crates.io, generic) +- `[[bangs]]` — custom search shortcuts +- `[[tagging.tags]]` — custom tag definitions with optional negatives +- `[web]` — server bind address, page size, worker threads +- `exclude_urls` — hosts to skip when fetching content +- `db_path` — bookmark database location -# Per-engine configuration (optional) -# Multiple instances supported (e.g., public + private crates.io registries) -[[engines]] -type = "searxng" -instance = "https://search.kael.ink" -# timeout_secs = 10.0 # optional per-engine timeout -# Best: use an existing public instance (see https://searx.space). -# Also possible: run your own with Docker: -# docker run -d --name searxng -p 8888:8080 searxng/searxng - -# Custom crates.io registry (optional) -# [[engines]] -# type = "crates_io" -# url = "https://registry.example.com/api/v1/crates?q={}&per_page=10" -# timeout_secs = 5.0 - -# Wikipedia search (optional, defaults to English) -# [[engines]] -# type = "wikipedia" -# lang = "fr" -# timeout_secs = 5.0 - -# MDN Web Docs search (optional, defaults to en-US) -# [[engines]] -# type = "mdn" -# locale = "fr" -# timeout_secs = 5.0 - -# Generic HTML-scraped engine (use with any search site) -# Provide a URL template with `{}` for the query and a CSS selector -# targeting the result container. Results are extracted from `<a>` links -# inside that container (deduplicated, up to 10, http/https only). -# -# Note: most commercial search engines (Google, Bing, DuckDuckGo, etc.) -# block automated requests. This engine works best with small/niche sites -# that don't enforce bot detection. To find the right selector, view the -# page source or use browser dev tools on the search results page. -# [[engines]] -# type = "generic" -# name = "DuckDuckGo" -# url = "https://html.duckduckgo.com/html/?q={}" -# selector = "div.results" -# timeout_secs = 10.0 -# shortcode = "ddg" # optional: override auto-generated shortcode -# bang_enabled = true # optional: disable ! redirect but keep @ -# bang_url = "..." # optional: custom redirect URL (keeps shortcode) -``` +See `search_hub config-schema` for the complete JSON Schema of the config file format. ## Search shortcuts -SearchHub supports two query prefixes that use **shortcodes** — compact aliases -auto-generated from your configured `[[engines]]`. +Shortcodes are compact aliases auto-generated from configured `[[engines]]`. | Prefix | Example | Behavior | |--------|---------|----------| @@ -163,8 +87,6 @@ auto-generated from your configured `[[engines]]`. ### Auto-generated shortcodes -Each engine type gets a sensible default shortcode: - | Engine | Shortcode | Bang URL | |--------|-----------|----------| | Wikipedia (lang=en) | `w` | `https://en.wikipedia.org/w/index.php?search={}` | @@ -175,9 +97,7 @@ Each engine type gets a sensible default shortcode: | SearXNG | `sx` | `{instance}/search?q={}` | | Generic | slugified name | the engine's own URL template | -### Overriding shortcodes per engine - -Set `shortcode`, `bang_url`, or `bang_enabled` directly on the engine: +Override `shortcode`, `bang_url`, or `bang_enabled` per engine: ```toml [[engines]] @@ -185,108 +105,132 @@ type = "wikipedia" lang = "fr" shortcode = "wikifr" # overrides "wfr" bang_enabled = false # disable ! redirect (still searchable via @) -bang_url = "https://..." # custom redirect URL ``` -### Custom bangs (standalone shortcuts — no @ support) +Custom bangs (standalone `!` redirects without `@` support): ```toml [[bangs]] trigger = "gh" url = "https://github.com/search?q={}" name = "GitHub" +``` -# Suppress an auto-generated shortcut -[[bangs]] -trigger = "crates" -enabled = false +If two engines share a shortcode, SearchHub panics at startup naming both. Set `shortcode` on one to resolve. + +## Auto-tagging + +When `[tagging] enabled = true`, SearchHub automatically tags bookmarks during insert and import using a two-stage ONNX pipeline: + +1. **Bi-encoder** — `BAAI/bge-small-en-v1.5` (~33 MB) embeds both the bookmark content and each tag's example set. Scores use **mean-pooling** (average similarity across all tag examples), requiring consensus rather than a single match. + +2. **Cross-encoder reranker** — `BAAI/bge-reranker-v2-m3` (~2 GB, optional) re-ranks the top 10 candidates for higher precision. Falls back to bi-encoder scores on failure. Disable with `reranker_enabled = false` to save disk space. + +Tags can include **negative examples**. The score is `pos_mean - neg_mean * 0.3`, so negative matches penalize a tag without dominating. Tags without negatives are unaffected. A **confidence threshold** (`[tagging] threshold`, default: 0.6) controls which tags are applied. + +### Tag taxonomy + +SearchHub ships with 21 built-in tags. Add or override them with `[[tagging.tags]]` in config. + +| Tag | Domain | +|-----|--------| +| `rust` | Rust language projects and tooling | +| `python` | Python language projects and tooling | +| `javascript` | JS/TS, Node.js, browser APIs | +| `web` | Web frameworks, HTML, CSS, frontend | +| `api` | REST, GraphQL, RPC protocols | +| `database` | SQL, NoSQL, ORMs, migrations | +| `testing` | Unit, integration, E2E, property-based | +| `linux` | Kernel, systemd, shell, config | +| `security` | Auth, crypto, CVE, hardening | +| `ai` | ML, LLMs, computer vision, NLP | +| `data` | ETL, pipelines, formats (CSV, Parquet) | +| `audio` | Sound processing, codecs, DSP | +| `video` | Video codecs, streaming, rendering | +| `gaming` | Game engines, graphics, physics | +| `hardware` | CPU, GPU, memory, peripherals | +| `mobile` | iOS, Android, PWA, responsive | +| `design` | UI/UX, typography, color, icons | +| `documentation` | API refs, ADRs, READMEs, docstrings | +| `news` | Release announcements, changelogs, events | +| `social` | Forums, chat, federated platforms | +| `productivity` | Note-taking, task management, tools | + +### Custom tags with negatives + +```toml +[[tagging.tags]] +name = "my-special-tag" +examples = ["example text one", "example text two"] +negatives = ["something similar but different", "another close miss"] ``` -### Collisions +### Disk space + +The bi-encoder model is ~33 MB. Enabling the reranker adds ~2 GB (~2.1 GB total in cache). Set `reranker_enabled = false` to minimize disk usage. The first tagging command downloads models to `~/.cache/search_hub/`. -If two engines produce the same shortcode, SearchHub panics at startup with -a message naming both engines. Set `shortcode` on one of them to resolve it. +## Deployment -## Run the web server as a systemd user service +### Systemd user services -Keeps the web UI running in the background, starts automatically on login. +**Web server** (auto-start on login): ```sh -VERSION=(search_hub --version | cut -d\ -f2) +VERSION=$(search_hub --version | cut -d\ -f2) mkdir -p ~/.config/systemd/user -wget -O ~/.config/systemd/user/search-hub-web.service https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-web.service +wget -O ~/.config/systemd/user/search-hub-web.service \ + https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-web.service systemctl --user daemon-reload systemctl --user enable --now search-hub-web.service ``` -Check status with `systemctl --user status search-hub-web`. View logs with `journalctl --user -u search-hub-web -f`. - -## Auto-import with systemd +**Auto-import** (daily Firefox/Zen bookmark import): ```sh -VERSION=(search_hub --version | cut -d\ -f2) -mkdir -p ~/.config/systemd/user -wget -O ~/.config/systemd/user/search-hub-import.service https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-import.service -wget -O ~/.config/systemd/user/search-hub-import.timer https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-import.timer +wget -O ~/.config/systemd/user/search-hub-import.service \ + https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-import.service +wget -O ~/.config/systemd/user/search-hub-import.timer \ + https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-import.timer systemctl --user daemon-reload systemctl --user enable --now search-hub-import.timer ``` -This imports bookmarks from Zen Browser daily. Edit the file to import from another browser. - -## Auto-update with systemd +**Auto-update** (weekly release check): ```sh -VERSION=(search_hub --version | cut -d\ -f2) -mkdir -p ~/.config/systemd/user -wget -O ~/.config/systemd/user/search-hub-update.service https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-update.service -wget -O ~/.config/systemd/user/search-hub-update.timer https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-update.timer +wget -O ~/.config/systemd/user/search-hub-update.service \ + https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-update.service +wget -O ~/.config/systemd/user/search-hub-update.timer \ + https://vit.am/~ololduck/search_hub/repository/browse/v$VERSION/contrib/search-hub-update.timer systemctl --user daemon-reload systemctl --user enable --now search-hub-self-update.timer ``` -This checks for new releases weekly and updates the binary automatically. - -## Run with Podman / Docker +Check status with `systemctl --user status search-hub-web`. View logs with `journalctl --user -u search-hub-web -f`. -A container image is available at `oci.vit.am/search-hub:latest`. It serves on port 8080 as the `search_hub` user and expects: +### Podman / Docker -- **Config** mounted at `/home/search_hub/.config/search_hub/config.toml` -- **Database** directory mounted at `/home/search_hub/.local/share/search_hub/` +A container image is available at `oci.vit.am/search-hub:latest` (port 8080, runs as `search_hub` user). ```sh -# Pull and run podman run -d --name search-hub \ -p 8080:8080 \ -v ~/.config/search_hub:/home/search_hub/.config/search_hub:ro \ -v ~/.local/share/search_hub:/home/search_hub/.local/share/search_hub \ oci.vit.am/search-hub:latest serve - -# SIGHUP reload (re-reads config) -podman kill -s HUP search-hub - -# Build locally from the Containerfile -podman build -t search-hub:latest -f Containerfile . ``` -### docker-compose +SIGHUP reloads the config without restarting: `podman kill -s HUP search-hub` -```sh -docker compose up -d -``` +Or use `docker compose up -d` (see `docker-compose.yaml` at the project root, includes a commented SearXNG service). For systemd-native container management, see the Quadlet file at `contrib/search-hub.container`. -See `docker-compose.yaml` at the project root. A SearXNG service is included as a commented-out example. +## Contributing -### Podman Quadlet (systemd-native) +Contributions are welcome! This is a personal project, but bug reports and feature requests are genuinely appreciated -- they help make search\_hub better for everyone. -```sh -mkdir -p ~/.config/containers/systemd -wget -O ~/.config/containers/systemd/search-hub.container https://vit.am/~ololduck/search_hub/repository/browse/main/contrib/search-hub.container -systemctl --user daemon-reload -systemctl --user enable --now search-hub -``` +The best way to contribute is to report a bug or describe an unsupported use case. Pull requests are welcome too, though for larger changes it's a good idea to reach out first to discuss the approach. The project uses a `develop` branch for active work, with `main` reserved for releases. -The Quadlet file uses `%h` (your home directory) for volume source paths. +I'm reachable via [mail](mailto:ololduck@vit.am), IRC (`ololduck` on Libera.Chat), or the Fediverse ([@ololduck@vit.am](https://vit.am/@ololduck)). ## Resources -
modified abbaye.toml
diff --git a/abbaye.toml b/abbaye.toml index cf2496e..b29a588 100644 --- a/abbaye.toml +++ b/abbaye.toml @@ -20,6 +20,7 @@ type = "archive" ignore_patterns = [".git/", "*.tar.gz"] [[builders]] +id = "build default" type = "cargo" targets = ["x86_64-unknown-linux-gnu"] comment = "with ONNX tagging (default features)" @@ -29,3 +30,15 @@ type = "cargo" targets = ["x86_64-unknown-linux-gnu"] no_default_features = true comment = "without ONNX tagging (~30% smaller binary)" + +[[builders]] +id = "generate schema" +category = "Other" +name = "Config JSON Schema" +comment = "Auto-generated JSON Schema for search_hub config.toml" +type = "script" +script = [ + "cargo run --release -- config-schema > target/search_hub.schema.json", +] +outputs = ["target/search_hub.schema.json"] +depends_on = ["build default"] -
modified contrib/release.sh
diff --git a/contrib/release.sh b/contrib/release.sh index 80c4307..0a56b2e 100755 --- a/contrib/release.sh +++ b/contrib/release.sh @@ -59,6 +59,9 @@ cargo test 2>&1 echo "==> Running cargo build --release" cargo build --release 2>&1 +echo "==> Generating config JSON schema" +./target/release/search_hub config-schema > search_hub.schema.json + echo "==> Updating CHANGELOG.md for v$VERSION" RELEASE_DATE=$(date +%Y-%m-%d) awk -v ver="$VERSION" -v date="$RELEASE_DATE" ' @@ -80,7 +83,7 @@ awk -v ver="$VERSION" -v date="$RELEASE_DATE" ' ' CHANGELOG.md > CHANGELOG.md.tmp && mv CHANGELOG.md.tmp CHANGELOG.md echo "==> Committing release: v$VERSION" -git add Cargo.toml Cargo.lock CHANGELOG.md +git add Cargo.toml Cargo.lock CHANGELOG.md search_hub.schema.json git commit -m "release: v$VERSION" echo "==> Tagging v$VERSION" -
modified src/config.rs
diff --git a/src/config.rs b/src/config.rs index b3d1ab0..a2628c8 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,5 +1,6 @@ use figment::Figment; use figment::providers::{Format, Toml}; +use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::PathBuf; @@ -9,7 +10,7 @@ use async_trait::async_trait; use crate::search_engines::{EngineError, ResultEntry, SearchEngine}; /// Configuration for a single search engine instance. -#[derive(Debug, Deserialize, Clone)] +#[derive(Debug, Deserialize, Clone, JsonSchema)] #[serde(tag = "type", rename_all = "lowercase")] pub enum EngineConfig { /// crates.io registry (public or private) @@ -196,7 +197,7 @@ pub struct Shortcut { } /// User-defined override or addition to auto-generated shortcuts. -#[derive(Debug, Deserialize, Clone)] +#[derive(Debug, Deserialize, Clone, JsonSchema)] pub struct BangOverride { /// The trigger keyword (e.g. "w" for !w). pub trigger: String, @@ -295,6 +296,9 @@ impl EngineConfig { /// Application configuration loaded from the TOML config file. /// +/// All config structs derive `schemars::JsonSchema` so the full +/// schema can be printed via `search_hub config-schema`. +/// /// # Example /// /// ```ignore @@ -302,17 +306,11 @@ impl EngineConfig { /// let engines = cfg.engines.clone(); /// println!("{} engines enabled", engines.len()); /// ``` -#[derive(Debug, Default, Deserialize)] +#[derive(Debug, Default, Deserialize, JsonSchema)] pub struct Config { - /// Custom tag definitions. If non-empty, these replace the hardcoded defaults. - #[serde(default)] - pub tags: Vec<crate::tagging::TagDef>, - /// Whether auto-tagging is enabled. Defaults to `false` if not set. - #[serde(default)] - pub tagging_enabled: Option<bool>, - /// Tagging threshold (0.0 to 1.0). Defaults to 0.60 if not set. + /// Tagging-specific configuration (always present; defaults apply when `[tagging]` is absent). #[serde(default)] - pub tagging_threshold: Option<f64>, + pub tagging: TaggingConfig, /// Hostnames to exclude from content fetching. #[serde(default)] pub exclude_urls: Option<Vec<String>>, @@ -326,6 +324,14 @@ pub struct Config { #[serde(default)] pub db_path: Option<String>, + /// Web server configuration (always present; defaults apply when `[web]` is absent). + #[serde(default)] + pub web: WebConfig, +} + +/// Web server configuration options. +#[derive(Debug, Clone, Deserialize, JsonSchema)] +pub struct WebConfig { /// Server bind address (default: "127.0.0.1"). #[serde(default)] pub bind_address: Option<String>, @@ -335,16 +341,85 @@ pub struct Config { /// Actix worker threads (default: 2). #[serde(default)] pub workers: Option<usize>, +} +impl Default for WebConfig { + fn default() -> Self { + Self { + bind_address: Some("127.0.0.1".into()), + page_size: Some(20), + workers: Some(2), + } + } +} + +impl WebConfig { + pub fn bind_address(&self) -> String { + self.bind_address.clone().unwrap_or_else(|| "127.0.0.1".into()) + } + pub fn page_size(&self) -> usize { + self.page_size.unwrap_or(20) + } + pub fn workers(&self) -> usize { + self.workers.unwrap_or(2) + } +} + +/// Tagging-specific configuration options. +#[derive(Debug, Clone, Deserialize, JsonSchema)] +#[serde(default)] +pub struct TaggingConfig { + /// Whether auto-tagging is enabled (default: false). + pub enabled: Option<bool>, + /// Tagging threshold (0.0 to 1.0, default: 0.6). + pub threshold: Option<f64>, /// ONNX embedding model name (default: "BGESmallENV15"). - #[serde(default)] pub onnx_model: Option<String>, - /// Max characters to use from page content for tagging (default: 2000). - #[serde(default)] + /// Max characters to use from page content for tagging (default: 4000). pub truncation: Option<usize>, /// Max tags to assign per bookmark (default: 5). - #[serde(default)] pub max_tags: Option<usize>, + /// Whether to use the cross-encoder reranker (default: true). + pub reranker_enabled: Option<bool>, + /// Cross-encoder reranker model code (default: "rozgo/bge-reranker-v2-m3"). + pub reranker_model: Option<String>, + /// Weight multiplier for the tags column in FTS5 search ranking (default: 3.0). + pub tag_weight: Option<f64>, + /// If true, use only user-provided [[tags]] and ignore built-in tag defaults. + pub override_tags: Option<bool>, + /// Custom tag definitions. If non-empty and `override_tags` is false, + /// these replace the hardcoded defaults (nested under `[tagging]` in TOML). + #[serde(default)] + pub tags: Vec<crate::tagging::TagDef>, +} + +impl Default for TaggingConfig { + fn default() -> Self { + Self { + enabled: Some(false), + threshold: Some(0.35), + onnx_model: Some("BGESmallENV15".into()), + truncation: Some(4000), + max_tags: Some(5), + reranker_enabled: Some(false), + reranker_model: Some("rozgo/bge-reranker-v2-m3".into()), + tag_weight: Some(3.0), + override_tags: Some(false), + tags: Vec::new(), + } + } +} + +impl TaggingConfig { + pub fn enabled(&self) -> bool { self.enabled.unwrap_or(false) } + pub fn threshold(&self) -> f64 { self.threshold.unwrap_or(0.35) } + pub fn onnx_model(&self) -> String { self.onnx_model.clone().unwrap_or_else(|| "BGESmallENV15".into()) } + pub fn truncation(&self) -> usize { self.truncation.unwrap_or(4000) } + pub fn max_tags(&self) -> usize { self.max_tags.unwrap_or(5) } + pub fn reranker_enabled(&self) -> bool { self.reranker_enabled.unwrap_or(false) } + pub fn reranker_model(&self) -> String { self.reranker_model.clone().unwrap_or_else(|| "rozgo/bge-reranker-v2-m3".into()) } + pub fn tag_weight(&self) -> f64 { self.tag_weight.unwrap_or(3.0) } + pub fn override_tags(&self) -> bool { self.override_tags.unwrap_or(false) } } impl Config { @@ -485,7 +560,7 @@ mod tests { #[test] fn load_from_missing_file_returns_default() { let cfg = Config::load_from(&PathBuf::from("/nonexistent/path.toml")); - assert!(cfg.tags.is_empty()); + assert!(cfg.tagging.tags.is_empty()); assert!(cfg.engines.is_empty()); } @@ -578,30 +653,32 @@ instance = "https://search.example.com" let mut file = NamedTempFile::new().unwrap(); write!(file, "invalid toml [[[").unwrap(); let cfg = Config::load_from(&file.path().to_path_buf()); - assert!(cfg.tags.is_empty()); + assert!(cfg.tagging.tags.is_empty()); assert!(cfg.engines.is_empty()); } #[test] fn tagging_enabled_defaults_to_false() { let cfg = Config::default(); - assert!(!cfg.tagging_enabled.unwrap_or(false)); + assert!(!cfg.tagging.enabled()); } #[test] fn tagging_enabled_can_be_false() { let mut file = NamedTempFile::new().unwrap(); - write!(file, r#"tagging_enabled = false"#).unwrap(); + write!(file, r#"[tagging] +enabled = false"#).unwrap(); let cfg = Config::load_from(&file.path().to_path_buf()); - assert_eq!(cfg.tagging_enabled, Some(false)); + assert!(!cfg.tagging.enabled()); } #[test] fn tagging_enabled_can_be_true() { let mut file = NamedTempFile::new().unwrap(); - write!(file, r#"tagging_enabled = true"#).unwrap(); + write!(file, r#"[tagging] +enabled = true"#).unwrap(); let cfg = Config::load_from(&file.path().to_path_buf()); - assert_eq!(cfg.tagging_enabled, Some(true)); + assert!(cfg.tagging.enabled()); } // -- shortcut (bang/@) tests -- -
modified src/main.rs
diff --git a/src/main.rs b/src/main.rs index c1e55b3..dc359c1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -199,6 +199,8 @@ Import { }, /// Create a default config file at the default config path InitConfig, +/// Print the JSON schema for the config file format +ConfigSchema, /// Check for updates and apply them automatically SelfUpdate { #[arg(long)] @@ -316,22 +318,28 @@ async fn main() { None => Config::load(), }; let engines = config.engines.clone(); - let _tagging_enabled = config.tagging_enabled.unwrap_or(false); - let _tag_threshold: f32 = config.tagging_threshold.map(|t| t as f32).unwrap_or(0.60); + let _tagging_enabled = config.tagging.enabled(); + let _override_tags = config.tagging.override_tags(); + let _tag_threshold: f32 = config.tagging.threshold() as f32; let exclude_hosts: Vec<String> = config.exclude_urls.clone().unwrap_or_else(|| { vec!["localhost".into(), "127.0.0.1".into(), "::1".into()] }); - let _tags: Vec<TagDef> = if config.tags.is_empty() { + let _tags: Vec<TagDef> = if _override_tags { + config.tagging.tags.clone() + } else if config.tagging.tags.is_empty() { default_tags() } else { - config.tags.clone() + config.tagging.tags.clone() }; - let bind_address = config.bind_address.clone().unwrap_or_else(|| "127.0.0.1".into()); - let page_size = config.page_size.unwrap_or(20); - let workers = config.workers.unwrap_or(2); - let _onnx_model = config.onnx_model.clone().unwrap_or_else(|| "BGESmallENV15".into()); - let _truncation = config.truncation.unwrap_or(2000); - let _max_tags = config.max_tags.unwrap_or(5); + let bind_address = config.web.bind_address(); + let page_size = config.web.page_size(); + let workers = config.web.workers(); + let _onnx_model = config.tagging.onnx_model(); + let _truncation = config.tagging.truncation(); + let _max_tags = config.tagging.max_tags(); + let _reranker_enabled = config.tagging.reranker_enabled(); + let _reranker_model = config.tagging.reranker_model(); + let _tag_weight = config.tagging.tag_weight(); let cache_dir = directories::ProjectDirs::from("com", "search_hub", "search_hub") .map(|d| d.cache_dir().to_path_buf()) .unwrap_or_else(|| { @@ -346,7 +354,7 @@ async fn main() { let db_path = resolve_db_path(db_path, config.db_path.as_deref()); let engines = Arc::new(RwLock::new(engines)); info!("Starting server on {}:{}", bind_address, port); - let srv_cfg = web::ServerConfig { port, bind_address, page_size, workers }; + let srv_cfg = web::ServerConfig { port, bind_address, page_size, workers, tag_weight: _tag_weight }; let reload_engines = engines.clone(); let config_path = config_path.clone(); @@ -414,7 +422,7 @@ async fn main() { return None; } info!("tagging content..."); - match TaggingEngine::new(&_tags, _tag_threshold, _max_tags, _truncation, &_onnx_model) { + match TaggingEngine::new(&_tags, _tag_threshold, _max_tags, _truncation, &_onnx_model, _reranker_enabled, &_reranker_model) { Ok(mut engine) => { let t = engine.tags_for(c).unwrap_or_default(); if t.is_empty() { @@ -435,7 +443,7 @@ async fn main() { #[cfg(not(feature = "tagging"))] let tags_val: Option<String> = { if _tagging_enabled { - eprintln!("Warning: tagging not available in this build. Rebuild with 'cargo install --features tagging' to enable it."); + eprintln!("Warning: tagging is enabled in config but this build does not support it (compile with --features tagging)."); } None }; @@ -470,7 +478,7 @@ async fn main() { Command::Search { query, db_path } => { let db_path = resolve_db_path(db_path, config.db_path.as_deref()); let conn = storage::init_db(&db_path.to_string_lossy()).expect("Failed to open database"); - let bookmarks = storage::search_bookmarks(&conn, &query, 1, 10000).expect("Failed to search"); + let bookmarks = storage::search_bookmarks(&conn, &query, 1, 10000, _tag_weight).expect("Failed to search"); if bookmarks.is_empty() { println!("No results for \"{query}\""); @@ -486,14 +494,14 @@ async fn main() { #[cfg(feature = "tagging")] Command::Retag { id, all, db_path } => { if !_tagging_enabled { - println!("Tagging is disabled in config. Enable it with tagging_enabled = true."); + println!("Tagging is disabled in config. Set `[tagging] enabled = true` to enable."); return; } let db_path = resolve_db_path(db_path, config.db_path.as_deref()); let conn = storage::init_db(&db_path.to_string_lossy()).expect("Failed to open database"); - let mut engine = match TaggingEngine::new(&_tags, _tag_threshold, _max_tags, _truncation, &_onnx_model) { + let mut engine = match TaggingEngine::new(&_tags, _tag_threshold, _max_tags, _truncation, &_onnx_model, _reranker_enabled, &_reranker_model) { Ok(e) => e, Err(e) => { eprintln!("Warning: failed to initialize tagger: {e}"); @@ -523,7 +531,7 @@ async fn main() { println!(" {}. #{} {} {} [{}]", i + 1, b.id, b.title, time, tags); } println!(); - print!("Enter IDs (comma-separated) or 'all': "); + print!("Enter IDs (space-separated), ranges (x-y), or 'all': "); std::io::Write::flush(&mut std::io::stdout()).unwrap(); let mut input = String::new(); std::io::stdin().read_line(&mut input).unwrap(); @@ -538,9 +546,7 @@ async fn main() { .map(|b| b.id as i64) .collect() } else { - input.split(',') - .filter_map(|s| s.trim().parse::<i64>().ok()) - .collect() + parse_retag_ids(input) } }; @@ -563,36 +569,38 @@ async fn main() { for rowid in &ids { match storage::get_bookmark(&conn, *rowid) { Ok(Some(b)) => { + eprintln!("→ {}", b.url); match b.content { Some(ref content) => { match engine.tags_for(content) { Ok(tags) => { let tags_str = if tags.is_empty() { None } else { Some(tags.join(", ")) }; + eprintln!(" ✓ tags: {}", tags_str.as_deref().unwrap_or("none")); storage::update_bookmark_tags(&conn, *rowid, tags_str.as_deref()) - .unwrap_or_else(|e| eprintln!("Warning: failed to update tags: {e}")); + .unwrap_or_else(|e| eprintln!(" ✘ failed to update tags: {e}")); tagged += 1; pb.set_message(format!("{} tagged", "✓")); } Err(e) => { - eprintln!("Warning: tagging failed for #{rowid}: {e}"); + eprintln!(" ✘ tagging failed: {e}"); skipped += 1; pb.set_message(format!("{} failed", "✘")); } } } None => { - eprintln!("Warning: #{rowid} has no content, skipping"); + eprintln!(" ✘ no content, skipping"); skipped += 1; pb.set_message(format!("{} no content", "✘")); } } } Ok(None) => { - eprintln!("Warning: bookmark #{rowid} not found"); + eprintln!(" ✘ bookmark #{rowid} not found"); skipped += 1; } Err(e) => { - eprintln!("Warning: failed to read bookmark #{rowid}: {e}"); + eprintln!(" ✘ failed to read bookmark #{rowid}: {e}"); skipped += 1; } } @@ -604,7 +612,7 @@ async fn main() { #[cfg(not(feature = "tagging"))] Command::Retag { .. } => { - println!("Tagging is not available in this build. Rebuild with 'cargo install --features tagging' to enable it."); + println!("Tagging is not available in this build. Compile with --features tagging to enable it."); } Command::InitConfig => { let path = config_path.clone().unwrap_or_else(config_file_path); @@ -622,21 +630,54 @@ async fn main() { }; let content = format!( "# SearchHub configuration\n\ + \"$schema\" = \"https://vit.am/~ololduck/search_hub/latest/dist/search_hub.schema.json\"\n\ \n\ # Bookmark database path (default: platform data directory)\n\ # db_path = \"{default_display}\"\n\ \n\ # Custom tags override the built-in defaults\n\ - # [[tags]]\n\ + # [[tagging.tags]]\n\ # name = \"my-custom-tag\"\n\ # examples = [\"example text one\", \"example text two\"]\n\ + # negatives = [\"close miss one\", \"close miss two\"]\n\ \n\ \n\ + [tagging]\n\ # Whether auto-tagging is enabled (default: false)\n\ - # tagging_enabled = true\n\ + # enabled = true\n\ + \n\ + # Minimum confidence for auto-tagging (0.0 to 1.0, default: 0.35)\n\ + # threshold = 0.35\n\ + \n\ + # ONNX embedding model name (default: \"BGESmallENV15\")\n\ + # onnx_model = \"BGESmallENV15\"\n\ + \n\ + # Max characters from page content for tagging (default: 4000)\n\ + # truncation = 4000\n\ + \n\ + # Max tags to assign per bookmark (default: 5)\n\ + # max_tags = 5\n\ + \n\ + # Cross-encoder reranker for improved tagging (default: false)\n\ + # reranker_enabled = false\n\ + # Cross-encoder model (default: \"rozgo/bge-reranker-v2-m3\")\n\ + # reranker_model = \"rozgo/bge-reranker-v2-m3\"\n\ + \n\ + # Weight multiplier for the tags column in search ranking (default: 3.0)\n\ + # tag_weight = 3.0\n\ + \n\ + # Use only your [[tagging.tags]] and ignore built-in defaults (default: false)\n\ + # override_tags = true\n\ + \n\ + [web]\n\ + # Server bind address (default: \"127.0.0.1\")\n\ + # bind_address = \"127.0.0.1\"\n\ \n\ - # Minimum confidence for auto-tagging (0.0 to 1.0, default: 0.6)\n\ - # tagging_threshold = 0.6\n\ + # Results per page in the web UI (default: 20)\n\ + # page_size = 20\n\ + \n\ + # Actix worker threads (default: 2)\n\ + # workers = 2\n\ \n\ # Hosts to skip when fetching content for bookmarking (default: localhost addresses)\n\ # exclude_urls = [\"localhost\", \"127.0.0.1\", \"::1\"]\n\ @@ -655,6 +696,10 @@ async fn main() { tokio::fs::write(&path, content).await.expect("Failed to write config file"); println!("Default config created at {path:?}"); } + Command::ConfigSchema => { + let schema = schemars::schema_for!(search_hub::config::Config); + println!("{}", serde_json::to_string_pretty(&schema).expect("serialize schema")); + } Command::SelfUpdate { feed_url, target, dry_run } => { let mut updater = search_hub::self_update::SelfUpdate::new() .dry_run(dry_run); @@ -741,15 +786,20 @@ fn resolve_profiles(importer: &(impl Importer + ?Sized), profile: Option<String> } async fn run_import(source: &str, profile: Option<String>, db_path: &str, config: &Config, kind: ImportKind, cache_path: PathBuf) { - let _tagging_enabled = config.tagging_enabled.unwrap_or(false); - let _tag_threshold: f32 = config.tagging_threshold.map(|t| t as f32).unwrap_or(0.60); - let _tags: Vec<TagDef> = if config.tags.is_empty() { default_tags() } else { config.tags.clone() }; + let _tagging_enabled = config.tagging.enabled(); + let _override_tags = config.tagging.override_tags(); + let _tag_threshold: f32 = config.tagging.threshold() as f32; + let _tags: Vec<TagDef> = if _override_tags { + config.tagging.tags.clone() + } else if config.tagging.tags.is_empty() { default_tags() } else { config.tagging.tags.clone() }; let exclude_hosts: Vec<String> = config.exclude_urls.clone().unwrap_or_else(|| { vec!["localhost".into(), "127.0.0.1".into(), "::1".into()] }); - let _max_tags = config.max_tags.unwrap_or(5); - let _truncation = config.truncation.unwrap_or(2000); - let _onnx_model = config.onnx_model.clone().unwrap_or_else(|| "BGESmallENV15".into()); + let _max_tags = config.tagging.max_tags(); + let _truncation = config.tagging.truncation(); + let _onnx_model = config.tagging.onnx_model(); + let _reranker_enabled = config.tagging.reranker_enabled(); + let _reranker_model = config.tagging.reranker_model(); let importer: Box<dyn Importer> = match source { "firefox" => Box::new(FirefoxImporter), @@ -877,6 +927,10 @@ async fn run_import(source: &str, profile: Option<String>, db_path: &str, config let task_tagging_enabled = _tagging_enabled; #[cfg(feature = "tagging")] let task_onnx = _onnx_model.clone(); + #[cfg(feature = "tagging")] + let task_reranker_enabled = _reranker_enabled; + #[cfg(feature = "tagging")] + let task_reranker_model = _reranker_model.clone(); tokio::task::spawn_blocking(move || { let mut fetcher = match Fetcher::new(task_cache) { Ok(f) => f, @@ -887,7 +941,7 @@ async fn run_import(source: &str, profile: Option<String>, db_path: &str, config }; #[cfg(feature = "tagging")] let mut tagger = if task_tagging_enabled { - TaggingEngine::new(&task_tags, task_threshold, _max_tags, _truncation, &task_onnx).ok() + TaggingEngine::new(&task_tags, task_threshold, _max_tags, _truncation, &task_onnx, task_reranker_enabled, &task_reranker_model).ok() } else { None }; @@ -899,18 +953,25 @@ async fn run_import(source: &str, profile: Option<String>, db_path: &str, config bar.inc(1); continue; } + let _ = tx.send(format!("[{task_id}] → {url}")); if let Some(md) = fetch_and_convert(&mut fetcher, url, Some(task_id)) { #[cfg(feature = "tagging")] let entry_tags = tagger.as_mut() .and_then(|e| e.tags_for(&md).ok()) .unwrap_or_default(); #[cfg(not(feature = "tagging"))] - let entry_tags: Vec<String> = Vec::new(); + let entry_tags: Vec<String> = { + if _tagging_enabled { + let _ = tx.send(format!("[{task_id}] Warning: tagging is enabled in config but this build does not support it (compile with --features tagging).")); + } + Vec::new() + }; let tags_str = if entry_tags.is_empty() { None } else { Some(entry_tags.join(", ")) }; + let _ = tx.send(format!(" ✓ tags: {}", tags_str.as_deref().unwrap_or("none"))); storage::update_bookmark_content_tags( &conn, *rowid, Some(&md), tags_str.as_deref(), ).unwrap_or_else(|e| { - let _ = tx.send(format!("failed to update bookmark {rowid}: {e}")); + let _ = tx.send(format!(" ✘ failed to update bookmark {rowid}: {e}")); }); } bar.inc(1); @@ -918,17 +979,15 @@ async fn run_import(source: &str, profile: Option<String>, db_path: &str, config }) }).collect(); - // Read errors from channel, keep ringbuffer of last 10, print live + // Read progress + errors from channel, keep rolling buffer, print live let error_handle = tokio::task::spawn_blocking(move || { - let mut error_buffer: VecDeque<String> = VecDeque::with_capacity(10); - while let Ok(err) = rx.recv() { - if error_buffer.len() < 10 { - error_buffer.push_back(err.clone()); - eprintln!("{err}"); - } else { - error_buffer.pop_front(); - error_buffer.push_back(err); + let mut buffer: VecDeque<String> = VecDeque::with_capacity(10); + while let Ok(msg) = rx.recv() { + if buffer.len() >= 10 { + buffer.pop_front(); } + buffer.push_back(msg.clone()); + eprintln!("{msg}"); } }); @@ -981,6 +1040,32 @@ fn fetch_and_convert(fetcher: &mut Fetcher, url: &str, task_id: Option<usize>) - } } +fn parse_retag_ids(input: &str) -> Vec<i64> { + if input.trim() == "all" { + return Vec::new(); // caller handles "all" separately + } + let mut ids: Vec<i64> = Vec::new(); + for token in input.split_whitespace() { + if let Some((a, b)) = token.split_once('-') { + let lo = a.parse::<i64>().ok(); + let hi = b.parse::<i64>().ok(); + match (lo, hi) { + (Some(lo), Some(hi)) if lo <= hi => { + for id in lo..=hi { ids.push(id); } + } + _ => { + eprintln!("Warning: invalid range '{token}', skipping"); + } + } + } else if let Ok(id) = token.parse::<i64>() { + ids.push(id); + } else { + eprintln!("Warning: invalid input '{token}', skipping"); + } + } + ids +} + #[cfg(test)] mod tests { use super::*; @@ -1151,4 +1236,88 @@ mod tests { assert!(is_excluded_url("http://my-internal.dev/api", &hosts)); assert!(!is_excluded_url("https://example.com", &hosts)); } + + #[test] + fn parse_retag_ids_single_ids() { + let result = parse_retag_ids("3 5 7"); + assert_eq!(result, vec![3, 5, 7]); + } + + #[test] + fn parse_retag_ids_range() { + let result = parse_retag_ids("1-5"); + assert_eq!(result, vec![1, 2, 3, 4, 5]); + } + + #[test] + fn parse_retag_ids_mixed() { + let result = parse_retag_ids("1-3 7 10-12"); + assert_eq!(result, vec![1, 2, 3, 7, 10, 11, 12]); + } + + #[test] + fn parse_retag_ids_empty() { + let result = parse_retag_ids(""); + assert!(result.is_empty()); + } + + #[test] + fn parse_retag_ids_garbage_skipped() { + let result = parse_retag_ids("1 foo 3- 5-1 bar"); + assert_eq!(result, vec![1]); + } + + #[test] + fn parse_retag_ids_whitespace() { + let result = parse_retag_ids(" 1 2 3 "); + assert_eq!(result, vec![1, 2, 3]); + } + + #[test] + fn parse_retag_ids_tabs_and_newlines() { + let result = parse_retag_ids("1\t2\n3"); + assert_eq!(result, vec![1, 2, 3]); + } + + #[test] + fn parse_retag_ids_single_id() { + let result = parse_retag_ids("42"); + assert_eq!(result, vec![42]); + } + + #[test] + fn parse_retag_ids_reversed_range() { + let result = parse_retag_ids("5-1"); + assert!(result.is_empty()); + } + + #[test] + fn parse_retag_ids_single_item_range() { + let result = parse_retag_ids("3-3"); + assert_eq!(result, vec![3]); + } + + #[test] + fn parse_retag_ids_zero() { + let result = parse_retag_ids("0"); + assert_eq!(result, vec![0]); + } + + #[test] + fn parse_retag_ids_negative_numbers_skipped() { + let result = parse_retag_ids("-1 3"); + assert_eq!(result, vec![3]); + } + + #[test] + fn parse_retag_ids_negative_range_skipped() { + let result = parse_retag_ids("-3--1"); + assert!(result.is_empty()); + } + + #[test] + fn parse_retag_ids_all_returns_empty() { + let result = parse_retag_ids("all"); + assert!(result.is_empty()); + } } -
modified src/storage.rs
diff --git a/src/storage.rs b/src/storage.rs index 2268831..27def0f 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -234,12 +234,14 @@ pub fn list_bookmarks(conn: &Connection, page: usize, page_size: usize) -> Resul Ok(bookmarks) } -/// Full-text search the bookmarks table using the FTS5 MATCH syntax. +/// Full-text search the bookmarks table using the FTS5 MATCH syntax, with +/// optional tag-weight boosting applied in Rust for compatibility with all +/// bundled SQLite versions. /// /// # Example /// /// ```ignore -/// let results = search_hub::storage::search_bookmarks(&conn, "rust tutorial", 1, 20) +/// let results = search_hub::storage::search_bookmarks(&conn, "rust tutorial", 1, 20, 3.0) /// .expect("search"); /// ``` /// @@ -249,31 +251,50 @@ pub fn list_bookmarks(conn: &Connection, page: usize, page_size: usize) -> Resul /// * `query` - An FTS5 search query string. /// * `page` - 1-indexed page number. /// * `page_size` - Results per page. +/// * `tag_weight`- Weight multiplier for the tags column (applied in Rust). /// /// # Returns /// /// A `Vec<Bookmark>` matching the query (empty if none match). -pub fn search_bookmarks(conn: &Connection, query: &str, page: usize, page_size: usize) -> Result<Vec<Bookmark>> { +pub fn search_bookmarks(conn: &Connection, query: &str, page: usize, page_size: usize, tag_weight: f64) -> Result<Vec<Bookmark>> { let offset = (page.saturating_sub(1)) * page_size; - let mut stmt = conn.prepare("SELECT rowid, title, url, description, source, content, tags, created_at FROM bookmarks WHERE bookmarks MATCH ? ORDER BY rank LIMIT ? OFFSET ?")?; - let book_iter = stmt.query_map(params![query, page_size as i64, offset as i64], |row| { - Ok(Bookmark { - id: row.get(0)?, - title: row.get(1)?, - url: row.get(2)?, - description: row.get(3)?, - source: row.get(4)?, - content: row.get(5)?, - tags: row.get(6)?, - created_at: row.get(7)?, - }) - })?; - - let mut bookmarks = Vec::new(); - for book in book_iter { - bookmarks.push(book?); + // rank is an FTS5 built-in auxiliary function. We select it alongside + // bookmark columns, then boost results whose tags column matches the query. + let mut stmt = conn.prepare("SELECT rowid, title, url, description, source, content, tags, created_at, rank FROM bookmarks WHERE bookmarks MATCH ? ORDER BY rank LIMIT ? OFFSET ?")?; + let mut bookmarks: Vec<(Bookmark, f64)> = stmt.query_map(params![query, page_size as i64, offset as i64], |row| { + Ok(( + Bookmark { + id: row.get(0)?, + title: row.get(1)?, + url: row.get(2)?, + description: row.get(3)?, + source: row.get(4)?, + content: row.get(5)?, + tags: row.get(6)?, + created_at: row.get(7)?, + }, + row.get::<_, f64>(8)?, + )) + })? + .filter_map(|r| r.ok()) + .collect(); + + if tag_weight != 1.0 && tag_weight > 0.0 { + let lowered = query.to_lowercase(); + for (bm, rank) in &mut bookmarks { + if let Some(ref tags) = bm.tags { + if tags.to_lowercase().contains(&lowered) { + // rank is negative (lower = better match). Multiply by + // tag_weight to make a matching tag result rank even + // better (more negative). + *rank *= tag_weight; + } + } + } + bookmarks.sort_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); } - Ok(bookmarks) + + Ok(bookmarks.into_iter().map(|(bm, _)| bm).collect()) } /// Count bookmarks matching an FTS5 query. @@ -405,7 +426,7 @@ mod tests { fn insert_and_search() { let conn = test_db(); insert_bookmark(&conn, &sample_bm()).expect("insert"); - let results = search_bookmarks(&conn, "Test", 1, 100).expect("search"); + let results = search_bookmarks(&conn, "Test", 1, 100, 3.0).expect("search"); assert_eq!(results.len(), 1); assert_eq!(results[0].title, "Test Page"); } @@ -414,7 +435,7 @@ mod tests { fn search_no_match() { let conn = test_db(); insert_bookmark(&conn, &sample_bm()).expect("insert"); - let results = search_bookmarks(&conn, "zzznotfound", 1, 100).expect("search"); + let results = search_bookmarks(&conn, "zzznotfound", 1, 100, 3.0).expect("search"); assert!(results.is_empty()); } -
modified src/tagging.rs
diff --git a/src/tagging.rs b/src/tagging.rs index 6778a09..14a129f 100644 --- a/src/tagging.rs +++ b/src/tagging.rs @@ -1,9 +1,13 @@ +use schemars::JsonSchema; use serde::Deserialize; #[cfg(feature = "tagging")] use std::path::PathBuf; #[cfg(feature = "tagging")] -use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions}; +use fastembed::{ + EmbeddingModel, TextEmbedding, TextInitOptions, + TextRerank, RerankInitOptions, RerankerModel, +}; #[cfg(feature = "tagging")] use fastembed::similarity::cosine_similarity; @@ -25,7 +29,7 @@ fn parse_model(name: &str) -> anyhow::Result<EmbeddingModel> { }) } -/// A named tag with example texts used for semantic similarity scoring. +/// A named tag with positive and negative example texts for semantic similarity. /// /// # Example /// @@ -34,19 +38,24 @@ fn parse_model(name: &str) -> anyhow::Result<EmbeddingModel> { /// /// let tag = TagDef { /// name: "rust".into(), -/// examples: vec!["Rust ownership".into(), "cargo build system".into()], +/// examples: vec!["serde serialization".into(), "tokio async runtime".into()], +/// negatives: vec![], /// }; /// assert_eq!(tag.name, "rust"); /// ``` -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Deserialize, JsonSchema)] pub struct TagDef { /// The tag label (e.g. "rust", "web"). pub name: String, - /// Example phrases that exemplify this tag for embedding comparison. + /// Positive example phrases that exemplify this tag. pub examples: Vec<String>, + /// Negative example phrases that should NOT match this tag. + #[serde(default)] + pub negatives: Vec<String>, } -/// Return the hardcoded default set of 25 tags with 3 example texts each. +/// Return the hardcoded default set of 25 tags with 3 positive and 2 negative +/// examples each. /// /// # Example /// @@ -54,204 +63,418 @@ pub struct TagDef { /// use search_hub::tagging::default_tags; /// /// let tags = default_tags(); -/// assert_eq!(tags.len(), 25); +/// assert_eq!(tags.len(), 21); /// assert_eq!(tags[0].name, "rust"); /// ``` pub fn default_tags() -> Vec<TagDef> { vec![ - TagDef { name: "rust".into(), examples: vec![ - "Rust ownership and borrow checker enforcing memory safety at compile time".into(), - "pattern matching with enums and the Result type for error handling".into(), - "cargo build system, crates.io ecosystem, and procedural macros".into(), - ]}, - TagDef { name: "python".into(), examples: vec![ - "Python indentation-based syntax, list comprehensions, and generator expressions".into(), - "dynamic typing, duck typing, and Python's data model protocols".into(), - "pip packaging, virtual environments, and Python import system".into(), - ]}, - TagDef { name: "web".into(), examples: vec![ - "HTML semantic markup, accessibility attributes, and document structure".into(), - "CSS layout with flexbox and grid, responsive design with media queries".into(), - "DOM manipulation, event bubbling, and Web API interfaces in the browser".into(), - ]}, - TagDef { name: "audio".into(), examples: vec![ - "music streaming, albums, playlists, and artist discovery".into(), - "podcast episodes, RSS feeds, and audio content distribution".into(), - "radio stations, live broadcasts, and audio programming".into(), - ]}, - TagDef { name: "backend".into(), examples: vec![ - "HTTP server routing, request handling, and response middleware chains".into(), - "connection pooling, ORM patterns, and server-side template rendering".into(), - "backend service architecture, message queues, and inter-service communication".into(), - ]}, - TagDef { name: "devops".into(), examples: vec![ - "container images, Dockerfiles, and Kubernetes pod orchestration".into(), - "infrastructure provisioning with Terraform and configuration management".into(), - "CI/CD build pipelines, artifact management, and deployment strategies".into(), - ]}, - TagDef { name: "data".into(), examples: vec![ - "data frame operations, statistical analysis, and numerical computing".into(), - "data visualization with plotting libraries and charting techniques".into(), - "ETL workflows, data cleaning, and batch processing pipelines".into(), - ]}, - TagDef { name: "ai".into(), examples: vec![ - "transformer attention mechanisms, tokenization, and embedding layers".into(), - "gradient descent, backpropagation, and neural network loss functions".into(), - "model quantization, fine-tuning strategies, and inference optimization".into(), - ]}, - TagDef { name: "linux".into(), examples: vec![ - "file permission bits, process management, and signal handling".into(), - "piping stdout, redirecting file descriptors, and shell expansion rules".into(), - "package managers, init systems, and systemd unit files".into(), - ]}, - TagDef { name: "security".into(), examples: vec![ - "authentication tokens, OAuth flows, and JWT session handling".into(), - "input sanitization, parameterized queries, and XSS/CSRF prevention".into(), - "certificate authorities, TLS handshakes, and mTLS configurations".into(), - ]}, - TagDef { name: "design".into(), examples: vec![ - "design tokens, component libraries, and design system consistency".into(), - "typographic scale, whitespace rhythm, and visual hierarchy principles".into(), - "color contrast, WCAG accessibility ratios, and responsive breakpoints".into(), - ]}, - TagDef { name: "mobile".into(), examples: vec![ - "touch gesture handling, viewport sizing, and responsive mobile layouts".into(), - "app lifecycle, push notifications, and background task management".into(), - "native platform APIs, mobile sensors, and cross-platform mobile frameworks".into(), - ]}, - TagDef { name: "gaming".into(), examples: vec![ - "game loop architecture, frame-rate independence, and delta time".into(), - "physics simulation, collision detection, and spatial partitioning".into(), - "shader programs, GPU rendering pipeline, and 3D transformations".into(), - ]}, - TagDef { name: "tutorial".into(), examples: vec![ - "beginner-friendly walkthroughs with code examples and expected output".into(), - "learning objectives, prerequisite knowledge, and progressive skill building".into(), - "interactive code playgrounds, exercises, and quiz-based reinforcement".into(), - ]}, - TagDef { name: "news".into(), examples: vec![ - "version bumps, deprecation timelines, and migration announcements".into(), - "community announcements, conference talks, and ecosystem updates".into(), - "release notes, changelogs, and feature release highlights".into(), - ]}, - TagDef { name: "video".into(), examples: vec![ - "video streaming platforms, channels, and content creation".into(), - "video editing, encoding formats, and transcoding workflows".into(), - "live streaming, video on demand, and media playback".into(), - ]}, - TagDef { name: "tools".into(), examples: vec![ - "text editor configuration, IDE plugins, and developer workflow tooling".into(), - "version control workflows, git branching strategies, and merge patterns".into(), - "debugger breakpoints, profiling tools, and performance tracing utilities".into(), - ]}, - TagDef { name: "database".into(), examples: vec![ - "SQL table schemas, foreign key relationships, and constraint design".into(), - "index structures, query plan analysis, and query performance tuning".into(), - "ACID transactions, isolation levels, and connection pool configuration".into(), - ]}, - TagDef { name: "cli".into(), examples: vec![ - "command argument parsing, subcommand patterns, and flag conventions".into(), - "terminal output formatting, colored logging, and progress indicators".into(), - "stdin/stdout pipes, exit codes, and shell completion scripts".into(), - ]}, - TagDef { name: "social".into(), examples: vec![ - "social media platforms, feeds, and community discussions".into(), - "user profiles, followers, and content sharing features".into(), - "messaging systems, real-time chat, and social networking APIs".into(), - ]}, - TagDef { name: "testing".into(), examples: vec![ - "unit test assertions, test fixtures, and parametrized test cases".into(), - "mocking external dependencies, test doubles, and fake implementations".into(), - "integration tests, end-to-end testing, and continuous testing in CI".into(), - ]}, - TagDef { name: "javascript".into(), examples: vec![ - "JavaScript closures, prototypal inheritance, and the event loop".into(), - "async/await patterns, Promise chaining, and callback conventions".into(), - "ES modules, npm packages, and JavaScript bundler tooling".into(), - ]}, - TagDef { name: "api".into(), examples: vec![ - "RESTful resource design, URL patterns, and HTTP method semantics".into(), - "request validation, error response formatting, and status code conventions".into(), - "API versioning, rate limiting, and OpenAPI specification documents".into(), - ]}, - TagDef { name: "documentation".into(), examples: vec![ - "API reference docs, docstrings, and inline code annotations".into(), - "architecture decision records and design documentation practices".into(), - "README writing, project wikis, and onboarding guides for contributors".into(), - ]}, - TagDef { name: "productivity".into(), examples: vec![ - "habit tracking, time management, and personal workflow optimization".into(), - "note-taking systems, knowledge base management, and personal wikis".into(), - "task organization, prioritization frameworks, and automation of repetitive work".into(), - ]}, + TagDef { + name: "rust".into(), + examples: vec![ + "serde derive macros for serialization and deserialization of structs".into(), + "tokio async runtime for concurrent network services and I/O".into(), + "cargo build scripts, feature flags, and workspace dependency management".into(), + ], + negatives: vec![ + "garbage collected memory management with tracing and compacting".into(), + "Python virtual environment and pip freeze dependency locking".into(), + ], + }, + TagDef { + name: "python".into(), + examples: vec![ + "pandas DataFrame groupby and merge operations for data analysis".into(), + "Django REST framework serializers and viewset routing".into(), + "pip install and requirements.txt for package dependency management".into(), + ], + negatives: vec![ + "Cargo.toml manifest with crate type and edition fields".into(), + "JavaScript npm install and package-lock.json dependency resolution".into(), + ], + }, + TagDef { + name: "javascript".into(), + examples: vec![ + "React useEffect dependency array and component re-render optimization".into(), + "async await with fetch API for HTTP requests and Promise error handling".into(), + "npm install and package.json scripts for build toolchain management".into(), + ], + negatives: vec![ + "Python decorator syntax for function wrapping and metaprogramming".into(), + "Rust macro_rules for compile-time code generation".into(), + ], + }, + TagDef { + name: "web".into(), + examples: vec![ + "HTML semantic landmarks and ARIA roles for accessible document structure".into(), + "CSS Grid and Flexbox for responsive multi-column page layouts".into(), + "HTTP caching headers ETag and Cache-Control for performance optimization".into(), + "Vue.js computed properties and watchers for reactive data binding".into(), + "DOM event delegation and custom event dispatch in the browser".into(), + ], + negatives: vec![ + "native iOS SwiftUI view modifiers and state property wrappers".into(), + "kernel module compilation and device driver development".into(), + ], + }, + TagDef { + name: "api".into(), + examples: vec![ + "RESTful endpoint design with proper HTTP method and status code semantics".into(), + "OpenAPI specification generation and API client code generation".into(), + "GraphQL schema type definitions and resolver data fetching patterns".into(), + ], + negatives: vec![ + "CSS animations and visual transitions for user interface effects".into(), + "HTML table structure and responsive styling across viewports".into(), + ], + }, + TagDef { + name: "database".into(), + examples: vec![ + "SQL JOIN query optimization with EXPLAIN ANALYZE and index scans".into(), + "database migration files for incremental schema versioning and rollback".into(), + "connection pool sizing and transaction isolation level configuration".into(), + ], + negatives: vec![ + "HTTP request body parsing and validation middleware".into(), + "frontend state management with Redux or Zustand".into(), + ], + }, + TagDef { + name: "linux".into(), + examples: vec![ + "systemd service unit with ExecStart, Restart, and journalctl log inspection".into(), + "file permission octal modes umask and access control list management".into(), + "process signal handling SIGTERM vs SIGKILL and systemd kill mode".into(), + ], + negatives: vec![ + "Kubernetes pod lifecycle probes and container restart policy".into(), + "Docker image layer caching and build context optimization".into(), + ], + }, + TagDef { + name: "security".into(), + examples: vec![ + "OAuth 2.0 Authorization Code flow with PKCE for public clients".into(), + "SQL injection prevention via parameterized queries and ORM escaping".into(), + "TLS certificate chain validation and mutual TLS for service-to-service auth".into(), + ], + negatives: vec![ + "database indexing strategies for query performance".into(), + "CSS custom properties for theme color management".into(), + ], + }, + TagDef { + name: "ai".into(), + examples: vec![ + "transformer self-attention mechanism for sequence-to-sequence modeling".into(), + "fine-tuning large language models with LoRA adapters and quantization".into(), + "vector embeddings for semantic search and retrieval-augmented generation".into(), + ], + negatives: vec![ + "SQL aggregate functions and window partition operations".into(), + "Linux process scheduling policies and nice values".into(), + ], + }, + TagDef { + name: "data".into(), + examples: vec![ + "pandas DataFrame groupby aggregation and pivot table transformation".into(), + "Matplotlib and Seaborn for statistical data visualization and charting".into(), + "ETL pipeline with data validation and incremental load strategies".into(), + ], + negatives: vec![ + "machine learning model training and hyperparameter optimization".into(), + "RESTful API endpoint design patterns and conventions".into(), + ], + }, + TagDef { + name: "audio".into(), + examples: vec![ + "Spotify Web API for playlist management and track metadata retrieval".into(), + "audio encoding formats MP3 AAC FLAC and bitrate versus quality tradeoffs".into(), + "podcast RSS feed generation with enclosure tags and episode metadata".into(), + ], + negatives: vec![ + "video codec configuration and resolution scaling parameters".into(), + "image EXIF metadata extraction and format conversion".into(), + ], + }, + TagDef { + name: "video".into(), + examples: vec![ + "HLS adaptive bitrate streaming with m3u8 playlist and TS segment files".into(), + "FFmpeg transcoding pipeline with codec selection and filter graphs".into(), + "video editing timeline with multi-track compositing and keyframe animation".into(), + ], + negatives: vec![ + "audio-only podcast feed generation and MP3 tag metadata".into(), + "static image optimization with WebP and AVIF format conversion".into(), + ], + }, + TagDef { + name: "gaming".into(), + examples: vec![ + "game loop with fixed timestep update and variable render interpolation".into(), + "entity component system architecture for game object management".into(), + "3D rendering pipeline with vertex shaders and fragment shader programs".into(), + ], + negatives: vec![ + "web application MVC framework routing and controller patterns".into(), + "CSS animation keyframes and transition timing functions".into(), + ], + }, + TagDef { + name: "hardware".into(), + examples: vec![ + "CPU cache hierarchy and instruction pipeline architecture".into(), + "PCIe lane allocation for GPU compute and NVMe storage bandwidth".into(), + "memory bandwidth versus latency tradeoffs for database workload".into(), + "mechanical keyboard switch types actuation force and keycap profile".into(), + "mouse sensor DPI polling rate and lift-off distance configuration".into(), + "display panel technology IPS VA OLED with refresh rate and color gamut".into(), + ], + negatives: vec![ + "software rendering pipeline and fragment shader algorithm implementation".into(), + "virtual machine resource allocation in cloud hypervisor environment".into(), + ], + }, + TagDef { + name: "mobile".into(), + examples: vec![ + "SwiftUI view modifier chaining and @State property wrapper patterns".into(), + "Android Kotlin coroutine scope for background task lifecycle management".into(), + "responsive viewport meta tag and touch event handler implementation".into(), + ], + negatives: vec![ + "desktop browser extension manifest and content script APIs".into(), + "server-side rendering and template engine integration".into(), + ], + }, + TagDef { + name: "testing".into(), + examples: vec![ + "unit test assertions with parametrized test cases for edge case coverage".into(), + "mock HTTP server for external API integration test isolation".into(), + "end-to-end browser test with Playwright page object model pattern".into(), + ], + negatives: vec![ + "production deployment configuration and environment variable management".into(), + "application feature implementation and business logic".into(), + ], + }, + TagDef { + name: "design".into(), + examples: vec![ + "typographic scale modular ratio and font pairing for readable text hierarchy".into(), + "color contrast ratio WCAG AA compliance and accessible color palette generation".into(), + "design system component tokens for spacing color and typography consistency".into(), + ], + negatives: vec![ + "backend API endpoint implementation and data validation".into(), + "command-line tool argument parsing and output formatting".into(), + ], + }, + TagDef { + name: "documentation".into(), + examples: vec![ + "Rustdoc documentation comments with code examples and crate-level docs".into(), + "architecture decision record for tracking design rationale and tradeoffs".into(), + "README structure with installation usage and contributing guide sections".into(), + "API reference page organized by endpoint with request response examples".into(), + ], + negatives: vec![ + "automated test assertions with mock setup and fixture teardown".into(), + ], + }, + TagDef { + name: "news".into(), + examples: vec![ + "framework major version release announcements and migration guides".into(), + "conference keynote summaries and community event calendar updates".into(), + "deprecation notices with timeline and replacement API migration path".into(), + ], + negatives: vec![ + "in-depth technical deep-dive with implementation walkthrough".into(), + "personal productivity workflows and note-taking methodology".into(), + ], + }, + TagDef { + name: "social".into(), + examples: vec![ + "forum thread discussion with community Q&A and answer ranking".into(), + "social media feed algorithm for content ranking and recommendation".into(), + "real-time messaging WebSocket connection for chat application".into(), + ], + negatives: vec![ + "technical documentation reference material".into(), + "version control branching strategy and code review workflow".into(), + ], + }, + TagDef { + name: "productivity".into(), + examples: vec![ + "personal knowledge management with Zettelkasten note-taking methodology".into(), + "task organization using kanban board with WIP limit and priority labels".into(), + "keyboard shortcut automation and text expansion for repetitive typing tasks".into(), + ], + negatives: vec![ + "programming language syntax and compiler error messages".into(), + "specific library API usage and function parameter documentation".into(), + ], + }, ] } #[cfg(feature = "tagging")] pub struct TaggingEngine { model: TextEmbedding, - tag_examples: Vec<(String, Vec<Vec<f32>>)>, + pos_embeddings: Vec<(String, Vec<Vec<f32>>)>, + neg_embeddings: Vec<(String, Vec<Vec<f32>>)>, threshold: f32, max_tags: usize, truncation: usize, + reranker: Option<TextRerank>, + reranker_docs: Vec<String>, } #[cfg(feature = "tagging")] impl TaggingEngine { /// Create a new tagging engine. /// - /// Downloads the ONNX embedding model on first run (cached afterwards). + /// Downloads the ONNX embedding model and optionally the cross-encoder + /// reranker model on first run (cached afterwards). /// /// # Parameters /// - /// * `tags` - Slice of `TagDef` entries. - /// * `threshold` - Minimum cosine-similarity score (0.0 to 1.0). - /// * `max_tags` - Default max tags to assign. - /// * `truncation` - Max characters to use from page content. - /// * `model_name` - ONNX model name (e.g. "BGESmallENV15"). + /// * `tags` - Slice of `TagDef` entries. + /// * `threshold` - Minimum similarity score (0.0 to 1.0). + /// * `max_tags` - Default max tags to assign. + /// * `truncation` - Max characters to use from page content. + /// * `model_name` - ONNX bi-encoder model name (e.g. "BGESmallENV15"). + /// * `reranker_enabled`- Whether to load and use the cross-encoder reranker. + /// * `reranker_model` - Cross-encoder model code (e.g. "rozgo/bge-reranker-v2-m3"). /// /// # Errors /// - /// Returns an error if the model is unknown or fails to load. + /// Returns an error if the bi-encoder model is unknown or fails to load. + /// Reranker load failures are logged as warnings and degrade gracefully. /// /// # Example /// /// ```ignore /// let tags = search_hub::tagging::default_tags(); - /// let mut engine = search_hub::tagging::TaggingEngine::new(&tags, 0.60, 5, 2000, "BGESmallENV15") - /// .expect("model init"); + /// let mut engine = search_hub::tagging::TaggingEngine::new( + /// &tags, 0.60, 5, 4000, "BGESmallENV15", true, "rozgo/bge-reranker-v2-m3", + /// ).expect("model init"); /// ``` - pub fn new(tags: &[TagDef], threshold: f32, max_tags: usize, truncation: usize, model_name: &str) -> anyhow::Result<Self> { + pub fn new( + tags: &[TagDef], + threshold: f32, + max_tags: usize, + truncation: usize, + model_name: &str, + reranker_enabled: bool, + reranker_model: &str, + ) -> anyhow::Result<Self> { let cache_dir = cache_dir_for_models(); let model = parse_model(model_name)?; let mut embedder = TextEmbedding::try_new( TextInitOptions::new(model) - .with_cache_dir(cache_dir) + .with_cache_dir(cache_dir.clone()) .with_show_download_progress(true), )?; - let mut all_examples: Vec<String> = Vec::new(); - let mut tag_indices: Vec<(usize, &str)> = Vec::new(); + let mut pos_texts: Vec<String> = Vec::new(); + let mut neg_texts: Vec<String> = Vec::new(); + let mut pos_indices: Vec<(usize, &str)> = Vec::new(); + let mut neg_indices: Vec<(usize, &str)> = Vec::new(); for (ti, tag) in tags.iter().enumerate() { for example in &tag.examples { - tag_indices.push((ti, &tag.name)); - all_examples.push(format!("passage: {example}")); + pos_indices.push((ti, &tag.name)); + pos_texts.push(format!("passage: {example}")); + } + for neg in &tag.negatives { + neg_indices.push((ti, &tag.name)); + neg_texts.push(format!("passage: {neg}")); } } - let embeddings = embedder.embed(all_examples, None)?; + let pos_embeddings_raw = embedder.embed(pos_texts, None)?; + let neg_embeddings_raw = if neg_texts.is_empty() { + Vec::new() + } else { + embedder.embed(neg_texts, None)? + }; + + let mut pos_embeddings: Vec<(String, Vec<Vec<f32>>)> = tags + .iter() + .map(|t| (t.name.clone(), Vec::new())) + .collect(); - let mut tag_examples: Vec<(String, Vec<Vec<f32>>)> = tags + let mut neg_embeddings: Vec<(String, Vec<Vec<f32>>)> = tags .iter() .map(|t| (t.name.clone(), Vec::new())) .collect(); - for ((ti, _name), emb) in tag_indices.iter().zip(embeddings.iter()) { - tag_examples[*ti].1.push(emb.clone()); + for ((ti, _), emb) in pos_indices.iter().zip(pos_embeddings_raw.iter()) { + pos_embeddings[*ti].1.push(emb.clone()); + } + for ((ti, _), emb) in neg_indices.iter().zip(neg_embeddings_raw.iter()) { + neg_embeddings[*ti].1.push(emb.clone()); } - Ok(Self { model: embedder, tag_examples, threshold, max_tags, truncation }) + let reranker_docs: Vec<String> = tags + .iter() + .map(|t| { + let negs = if t.negatives.is_empty() { + String::new() + } else { + format!(" (not: {})", t.negatives.join("; ")) + }; + format!("{}{}", t.name, negs) + }) + .collect(); + + let reranker = if reranker_enabled { + match reranker_model.parse::<RerankerModel>() { + Ok(model) => { + match TextRerank::try_new( + RerankInitOptions::new(model) + .with_cache_dir(cache_dir) + ) { + Ok(r) => { + eprintln!("info: cross-encoder reranker loaded"); + Some(r) + } + Err(e) => { + eprintln!("warning: failed to load reranker model '{reranker_model}': {e}. Tagging will use bi-encoder only."); + None + } + } + } + Err(e) => { + eprintln!("warning: invalid reranker model '{reranker_model}': {e}"); + None + } + } + } else { + None + }; + + Ok(Self { + model: embedder, + pos_embeddings, + neg_embeddings, + threshold, + max_tags, + truncation, + reranker, + reranker_docs, + }) } fn truncate<'a>(&self, content: &'a str) -> &'a str { @@ -274,33 +497,81 @@ impl TaggingEngine { } let query_emb = &emb[0]; - let mut scores: Vec<(usize, f32)> = self.tag_examples + let mut scores: Vec<(usize, f32)> = self.pos_embeddings .iter() .enumerate() - .map(|(i, (_, examples))| { - let max_sim = examples - .iter() - .map(|proto| cosine_similarity(query_emb, proto)) - .fold(f32::NEG_INFINITY, f32::max); - (i, max_sim) + .map(|(i, (_, pos_examples))| { + let pos_mean = if pos_examples.is_empty() { + 0.0 + } else { + let sum: f32 = pos_examples.iter() + .map(|proto| cosine_similarity(query_emb, proto)) + .sum(); + sum / pos_examples.len() as f32 + }; + let neg_examples = &self.neg_embeddings[i].1; + let neg_mean = if neg_examples.is_empty() { + 0.0 + } else { + let sum: f32 = neg_examples.iter() + .map(|proto| cosine_similarity(query_emb, proto)) + .sum(); + sum / neg_examples.len() as f32 + }; + let score = if neg_examples.is_empty() { + pos_mean + } else { + (pos_mean - neg_mean * 0.3).max(0.0) + }; + (i, score) }) .collect(); scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - Ok(scores + let mut result: Vec<(String, f32)> = scores .into_iter() - .map(|(i, score)| (self.tag_examples[i].0.clone(), score)) - .collect()) + .map(|(i, score)| (self.pos_embeddings[i].0.clone(), score)) + .collect(); + + let truncated = self.truncate(content); + let rerank_count = result.len().min(10); + let candidates: Vec<&str> = self.reranker_docs + .iter() + .map(|d| d.as_str()) + .collect(); + if let Some(reranker) = self.reranker.as_mut() { + let rerank_slice = &candidates[..rerank_count.min(candidates.len())]; + + if !rerank_slice.is_empty() { + match reranker.rerank(truncated, rerank_slice, false, None) { + Ok(reranked) => { + let mut reranked_scores: Vec<(usize, f32)> = reranked + .into_iter() + .map(|r| (r.index, r.score)) + .collect(); + reranked_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + let reranked_result: Vec<(String, f32)> = reranked_scores + .into_iter() + .map(|(i, score)| (self.pos_embeddings[i].0.clone(), score)) + .collect(); + result = reranked_result; + } + Err(e) => { + eprintln!("warning: reranker inference failed: {e}. Falling back to bi-encoder scores."); + } + } + } + } + + Ok(result) } - /// Score `content` against all tag prototypes and return tags above the - /// configured threshold. + /// Score `content` against all tags and return names above the threshold. /// /// # Parameters /// - /// * `content` - The text to tag (e.g. page body converted to Markdown). - /// * `max_tags` - Maximum number of tags to return. + /// * `content` - The text to tag (e.g. page body converted to Markdown). /// /// # Returns /// @@ -315,8 +586,9 @@ impl TaggingEngine { /// /// ```ignore /// let tags = search_hub::tagging::default_tags(); - /// let mut engine = search_hub::tagging::TaggingEngine::new(&tags, 0.40) - /// .expect("model init"); + /// let mut engine = search_hub::tagging::TaggingEngine::new( + /// &tags, 0.60, 5, 4000, "BGESmallENV15", false, "", + /// ).expect("model init"); /// let matched = engine.tags_for("the rust programming language") /// .expect("tagging failed"); /// println!("{:?}", matched); @@ -335,7 +607,7 @@ impl TaggingEngine { /// /// * `content` - The text to tag. /// * `max_tags` - Maximum number of tags to return. - /// * `threshold` - Minimum cosine-similarity score (0.0 to 1.0). + /// * `threshold` - Minimum similarity score (0.0 to 1.0). /// /// # Returns /// @@ -350,8 +622,9 @@ impl TaggingEngine { /// /// ```ignore /// let tags = search_hub::tagging::default_tags(); - /// let mut engine = search_hub::tagging::TaggingEngine::new(&tags) - /// .expect("model init"); + /// let mut engine = search_hub::tagging::TaggingEngine::new( + /// &tags, 0.60, 5, 4000, "BGESmallENV15", false, "", + /// ).expect("model init"); /// let matched = engine.tags_for_with_threshold("rust programming", 5, 0.30) /// .expect("tagging failed"); /// for (tag, score) in &matched { -
modified src/web/handlers.rs
diff --git a/src/web/handlers.rs b/src/web/handlers.rs index 007a5b2..3d412ef 100644 --- a/src/web/handlers.rs +++ b/src/web/handlers.rs @@ -217,7 +217,7 @@ async fn search( let total_pages = total_results.div_ceil(page_size); let bookmarks = if has_query { - storage::search_bookmarks(&db_pool.conn(), &q, page, page_size).unwrap_or_default() + storage::search_bookmarks(&db_pool.conn(), &q, page, page_size, cfg.tag_weight).unwrap_or_default() } else { storage::list_bookmarks(&db_pool.conn(), page, page_size).unwrap_or_default() }; @@ -333,7 +333,7 @@ async fn api_search( let total_pages = total_results.div_ceil(page_size); let bookmarks = if has_query { - storage::search_bookmarks(&db_pool.conn(), &q, page, page_size).unwrap_or_default() + storage::search_bookmarks(&db_pool.conn(), &q, page, page_size, cfg.tag_weight).unwrap_or_default() } else { storage::list_bookmarks(&db_pool.conn(), page, page_size).unwrap_or_default() }; @@ -483,7 +483,7 @@ async fn search_stream( let total_pages = total_results.div_ceil(page_size); let bookmarks = - storage::search_bookmarks(&db_pool.conn(), &q, page, page_size).unwrap_or_default(); + storage::search_bookmarks(&db_pool.conn(), &q, page, page_size, cfg.tag_weight).unwrap_or_default(); let api_bookmarks: Vec<SearchApiResult> = bookmarks .into_iter() -
modified src/web/mod.rs
diff --git a/src/web/mod.rs b/src/web/mod.rs index a0e1ac9..9e3eeab 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -40,6 +40,7 @@ pub struct ServerConfig { pub bind_address: String, pub page_size: usize, pub workers: usize, + pub tag_weight: f64, } /// Simplified shortcut info for rendering in Tera templates. -
modified tests/search_integration.rs
diff --git a/tests/search_integration.rs b/tests/search_integration.rs index 6b020d1..354fb19 100644 --- a/tests/search_integration.rs +++ b/tests/search_integration.rs @@ -18,6 +18,7 @@ macro_rules! setup { bind_address: "127.0.0.1".into(), page_size: 20, workers: 1, + tag_weight: 3.0, }); let shortcuts = web::Data::new(HashMap::<String, Shortcut>::new()); let mut tera_ = Tera::default(); @@ -376,6 +377,7 @@ async fn sse_bang_returns_bang_event() { bind_address: "127.0.0.1".into(), page_size: 20, workers: 1, + tag_weight: 3.0, }); let mut tera_ = Tera::default(); tera_ -
modified tests/tagging_thresholds.rs
diff --git a/tests/tagging_thresholds.rs b/tests/tagging_thresholds.rs index 6f6b6a0..290d7ce 100644 --- a/tests/tagging_thresholds.rs +++ b/tests/tagging_thresholds.rs @@ -398,7 +398,7 @@ so users can challenge moderation decisions they disagree with. #[test] fn explore_tagging_thresholds() { let tags = default_tags(); - let mut engine = TaggingEngine::new(&tags, 0.40, 5, 2000, "BGESmallENV15").expect("failed to init tagging engine"); + let mut engine = TaggingEngine::new(&tags, 0.40, 5, 2000, "BGESmallENV15", false, "").expect("failed to init tagging engine"); let thresholds = [ 0.30, 0.35, 0.40, 0.45, 0.50, 0.55, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90,