diff --git a/.gitignore b/.gitignore index ec75915..df56f48 100644 --- a/.gitignore +++ b/.gitignore @@ -43,3 +43,9 @@ Desktop.ini # Scratchpad scratchpad/ + +# Local-embedding assets (downloaded by `make assets`, never committed). +# Runtime assets that get go:embed-ed into the localembed build: +internal/embeddings/local/embedded/ +# Link-time static tokenizer lib (pulled in via CGO_LDFLAGS): +internal/embeddings/local/lib/ diff --git a/CLAUDE.md b/CLAUDE.md index 5edf196..2038451 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,16 +1,24 @@ # Agent Memory -Local-first desktop app + MCP server for semantic file search. Watches directories, embeds file contents via OpenAI, stores vectors in SQLite (sqlite-vec), and provides KNN semantic search. +Local-first desktop app + MCP server for semantic file search. Watches directories, embeds file contents (a bundled local model by default; OpenAI opt-in), stores vectors in SQLite (sqlite-vec), and provides KNN semantic search. ## Build & Run ```bash -make build # wails build -skipbindings (DO NOT use plain `wails build` — hangs on binding generation due to CGo) +make assets # download pinned local model/tokenizer/ORT libs into assets/embedded/ (~150 MB, checksummed via assets/manifest.json) +make build # wails build -skipbindings, builds with -tags localembed (DO NOT use plain `wails build` — hangs on binding generation due to CGo) +make build-darwin-amd64 # cross-build the Intel-mac app from an arm64 Mac (swaps in darwin-amd64 assets) make dev # hot-reload dev mode -make test # go test ./... +make test # go test ./... (default build, no localembed tag — stays native-lib-free) make clean # rm -rf build/bin ``` +> `make build` depends on `make assets`, which downloads ~150 MB of native artifacts (model, +> tokenizer, ONNX Runtime lib) the first time — needs network access. The build sets +> `-tags localembed` plus `CGO_LDFLAGS` to link the tokenizer library; the resulting binary is +> ~180 MB (it bundles the model). The native code is build-tag-isolated, so plain +> `go test ./...` and CI need no assets. + > **macOS 26+ gotcha:** Go ≤ 1.24 produces CGo binaries the kernel kills instantly on launch > (exit 137, `dyld: missing LC_UUID`) — this breaks `make build`, the `wails` CLI, and the built > app. Fix: `go env -w GOTOOLCHAIN=go1.26.4`, then `go install github.com/wailsapp/wails/v2/cmd/wails@v2.11.0`. @@ -44,7 +52,8 @@ Delivery (main.go, mcp/, tray.go, frontend/) |---------|---------| | `internal/engine` | Core orchestrator: scan → extract → chunk → embed → store | | `internal/store` | SQLite + sqlite-vec persistence (config, directories, files, chunks, vectors) | -| `internal/embeddings` | OpenAI embedding client with batching (2048/req) and retry | +| `internal/embeddings` | Embedder interface + OpenAI client (opt-in) with batching (2048/req) and retry | +| `internal/embeddings/local` | Bundled in-process local embedder (default): ONNX Runtime + HF tokenizer, `multilingual-e5-small` (384-dim), offline | | `internal/chunker` | Token-based text splitting (tiktoken, cl100k_base) | | `internal/watcher` | fsnotify recursive directory watcher with 500ms debounce | | `internal/extractor` | Multi-format content extraction (text, docx, xlsx, pptx, pdf, images) | diff --git a/Documentation/ARCHITECTURE.md b/Documentation/ARCHITECTURE.md index ca4c7d9..a6f97ad 100644 --- a/Documentation/ARCHITECTURE.md +++ b/Documentation/ARCHITECTURE.md @@ -17,7 +17,7 @@ Delivery -> Service (Engine) -> Domain (interfaces) -> Infrastructure (imp **Domain** — each package defines its own interface in `iface.go`. No cross-domain imports. -**Infrastructure** — concrete implementations that satisfy domain interfaces: SQLite, OpenAI API, fsnotify, file extractor. +**Infrastructure** — concrete implementations that satisfy domain interfaces: SQLite, the embedding providers (local ONNX model and OpenAI API), fsnotify, file extractor. ## Composition Root @@ -30,7 +30,7 @@ Six interfaces, each in its own package: | Interface | Package | Defined in | Implemented by | |-----------|---------|------------|----------------| | `Store` | `internal/store` | `iface.go` | `sqlite.go` (SQLite + sqlite-vec) | -| `Embedder` | `internal/embeddings` | `iface.go` | `openai.go` (OpenAI API) | +| `Embedder` | `internal/embeddings` | `iface.go` | `local/` (bundled ONNX model, **default**) and `openai.go` (OpenAI API, opt-in) | | `Chunker` | `internal/chunker` | `iface.go` | `chunker.go` (tiktoken, cl100k_base) | | `Watcher` | `internal/watcher` | `iface.go` | `fswatcher.go` (fsnotify) | | `FileEventHandler` | `internal/watcher` | `iface.go` | `engine.go` (Engine implements OnCreate/OnModify/OnDelete) | @@ -63,7 +63,8 @@ tray.go System tray (macOS CGo) internal/ domain/types.go Shared types: Directory, File, Chunk, SearchResult, IndexStats, ActivityLogEntry store/ Store interface + SQLite implementation - embeddings/ Embedder interface + OpenAI implementation + embeddings/ Embedder interface + OpenAI implementation (opt-in) + local/ Bundled ONNX local embedder (default provider) chunker/ Chunker interface + token-based splitter watcher/ Watcher + FileEventHandler interfaces + fsnotify implementation extractor/ Extractor interface + multi-format file extraction @@ -96,17 +97,36 @@ One binary, two modes controlled by the `--mcp` flag: Both share the same SQLite database. WAL mode supports concurrent readers with one writer. The GUI writes, the stdio process reads. -## Embedding Model Change Flow +## Embedding Providers -When the user changes the embedding model in Settings: +Two infrastructure implementations satisfy the same `Embedder` interface, so the engine, store, watcher, and MCP layers stay provider-agnostic: -1. `app.SetConfig("embedding_model", newModel)` detects the change -2. Persists the new model to the config table -3. Creates a new embedder via the injected `EmbedderFactory` (so `app.go` never imports concrete embeddings) -4. Calls `engine.SetEmbedder()` to swap it +- **Local (default)** — `internal/embeddings/local/`. A bundled `multilingual-e5-small` ONNX model (384 dimensions) runs in-process on the CPU via ONNX Runtime (`yalue/onnxruntime_go`), with a Hugging Face tokenizer (`daulet/tokenizers`). The session is lazy-initialized on first use; inference tokenizes, runs the model, mean-pools over the attention mask, and L2-normalizes. In stdio (`--mcp`) mode the model loads lazily on the first search so the MCP handshake stays instant. Works fully offline. +- **OpenAI (opt-in)** — `openai.go`. Used only when the user supplies an API key and selects the OpenAI provider in Settings. + +`main.go` picks the implementation from the `embedding_provider` config key; the local branch also wires the model-matched tokenizer into the chunker (via `chunker.WithTokenizer`) so chunk boundaries are measured in the model's own tokens. + +### Asset Bundling + +The native artifacts (ONNX model weights, tokenizer, ONNX Runtime shared library) are **not** committed to git. Instead: + +- `assets/manifest.json` (in git) pins each artifact's URL and SHA-256 checksum. +- `make assets` downloads and checksum-verifies them into `assets/embedded/` (gitignored). +- The `localembed`-tagged build embeds them into the binary via `go:embed`; on first use they are extracted (atomically, checksum-verified) to `~/.agent-memory/runtime/` because the ORT library must be `dlopen`-ed from a real file path. + +The native-lib-dependent code lives behind the `//go:build localembed` tag, so the default `go build`/`go test ./...` (and CI) stays lib-free and green; the pure-Go embedding pipeline is unit-testable without the tag. `make build` sets `-tags localembed` plus the `CGO_LDFLAGS` to link the tokenizer library. + +## Embedding Provider / Model Change Flow + +When the user changes the embedding provider or model in Settings: + +1. `app.SetConfig` detects the change (`embedding_provider` or `embedding_model`) +2. Persists the new value to the config table +3. Creates a new embedder via the injected `EmbedderFactory` (so `app.go` never imports concrete embeddings) — the factory maps the provider to `local.New(...)` or `NewOpenAIEmbedder(...)` +4. Calls `engine.SetEmbedder()` (and, for a provider switch, swaps the matching chunker/tokenizer) to swap it 5. Calls `engine.Reset()` which stops the watcher, drops and recreates the vector table with the new dimension, and restarts -The same factory pattern applies when the API key changes — the embedder is swapped so new requests use the updated key immediately. +The same factory pattern applies when the OpenAI API key changes — the embedder is swapped so new requests use the updated key immediately. ## Testing diff --git a/Documentation/Bugs/fswatcher-create-event-swallowed-linux.md b/Documentation/Bugs/fswatcher-create-event-swallowed-linux.md new file mode 100644 index 0000000..9e48921 --- /dev/null +++ b/Documentation/Bugs/fswatcher-create-event-swallowed-linux.md @@ -0,0 +1,54 @@ +# Bug: watcher swallows Create events on Linux (debounce replaces instead of merges) + +**Date found:** 2026-07-16 +**Found during:** local-embeddings epic, Phase 5 Linux verification (first-ever test run on Linux) +**Status:** Fix submitted — PR #5 (`fix/watcher-create-swallowed-linux`, based on `main`, independent of the epic PR stack) +**Severity:** Low today (no user-visible breakage), latent correctness risk + +## Symptom + +`go test ./internal/watcher/` fails deterministically on Linux (3/3 runs, Docker +`golang:1.26-bookworm`, arm64): + +``` +--- FAIL: TestOnCreate (2.01s) + fswatcher_test.go:121: expected OnCreate to be called +``` + +`TestOnModify` and `TestOnDelete` pass. The full suite is green on macOS. + +## Root cause + +Writing a new file on Linux (inotify) emits **two events** for the same path within +milliseconds: `CREATE`, then `WRITE`. In `internal/watcher/fswatcher.go`, +`handleEvent` → `debounce(path, fn)` stores **one pending closure per path and +replaces it** on each new event (`fswatcher.go:114-128`): the `WRITE` event cancels +the `CREATE` timer and substitutes a closure that only sees `WRITE`. When the +debounce fires, the handler classifies the event as a modification — +`handler.OnCreate` is never called for newly written files. + +macOS (fsnotify kqueue/FSEvents backend) delivers/coalesces these events differently, +so the swallow never manifests there. + +## Impact + +- **Today: effectively none for users.** `engine.OnCreate` and `engine.OnModify` + are identical (`engine.go:717-747` — both call `IndexFile`), so new files on + Linux are still indexed, merely misclassified as modifications. +- **Latent risk:** any future divergence between create and modify handling + (e.g. create-only bookkeeping, activity-log semantics, per-event UX) silently + breaks on Linux only. +- Blocks a fully green `go test ./...` on Linux (Phase 5 verify gate) until fixed. + +## Fix (applied in PR #5) + +In the debouncer, **accumulate the fsnotify op bits per path** instead of replacing +the closure — e.g. keep `pendingOps map[string]fsnotify.Op`, OR-ing each event's op; +when the timer fires, classify with precedence Remove/Rename > Create > Write/Chmod +from the merged bits, then clear the entry. Semantics on macOS are unchanged +(single-op case degenerates to today's behavior); Linux create+write merges to +Create. `TestOnCreate` then passes on both platforms. + +Fix delivered as its own small PR per working rules (the watcher is outside the +local-embeddings epic's scope): **PR #5**. Verified there: watcher suite 3/3 pass +on Linux arm64 (was 3/3 fail) and 3/3 pass on macOS. diff --git a/Documentation/Bugs/local-embed-token-budget-overflow.md b/Documentation/Bugs/local-embed-token-budget-overflow.md new file mode 100644 index 0000000..18fe244 --- /dev/null +++ b/Documentation/Bugs/local-embed-token-budget-overflow.md @@ -0,0 +1,54 @@ +# Bug: local embedder overflows the model context — every multi-chunk file silently dropped + +**Date found:** 2026-07-17 +**Found during:** Phase 5 Windows GUI field test (first real-world vault ever indexed) +**Status:** Fixed — branch `fix/local-embed-token-budget` (stacks on the epic PR chain) +**Severity:** Critical for the local provider — affects ALL platforms, latent since Phase 3 + +## Symptom + +Indexing a real vault (109 files, D:\Brain) silently indexed only 26 files. No errors +anywhere user-visible (see the companion visibility bug, PR #8). The 26 survivors were +exactly the files ≤ 1.6 KB; every file ≥ 1.7 KB — i.e., anything needing more than one +chunk — was dropped. All previous platform verifications used tiny single-chunk test +files, so five platforms' smokes all passed while the pipeline was broken for real +workloads everywhere. + +## Root cause (the "512 by 516" error) + +ONNX Runtime fails hard when a sequence exceeds mE5-small's 512-token context: + +``` +BroadcastIterator::Append ... Attempting to broadcast an axis by a dimension other than 1. 512 by 516 +``` + +`buildChunker` passed `embedder.MaxInputTokens()` (512) straight through as the chunk +budget. But the embedder prepends the E5 instruction prefix ("passage: ") and the +tokenizer adds special tokens **after** chunking, so a chunk cut at exactly 512 tokens +reached the model at 516. The epic specified the defense precisely ("effective chunk +size = min(chunk_size, MaxInputTokens − prefix − special) → default 480"); the +implementation skipped the reservation. + +The unchunked **query path** was also exposed: `EmbedQuery` has no chunker, so a search +query longer than ~500 tokens crashed inference on every platform. + +## Fix (two layers) + +1. `app.go buildChunker`: chunk budget = `MaxInputTokens() − local.EmbedTokenReserve` + (32 → effective 480, matching the epic). +2. `LocalEmbedder.embedBatch`: defensively truncate any tokenized input to the model + limit, preserving the trailing EOS token — protects queries and any future caller + regardless of chunking correctness. + +## Verification + +- Unit: truncation table test + fake-session proof the model never receives >512 tokens. +- Integration (permanent): `TestIntegrationLargeDocument` — real pipeline, multi-chunk + document + over-long query. Failed before the fix on macOS AND Windows; passes on both. +- Field: clean re-index of the same vault on Windows → **105 files / 932 chunks, all 5 + dossier files present, 79 multi-chunk files, zero errors** (was 26/26/0). + +## Lesson recorded + +Platform verifications used single-chunk corpora only; a large-document case is now a +permanent integration test, and the Phase 6 checklist gained a real-vault field test. diff --git a/Documentation/Bugs/pdf-extraction-passthrough.md b/Documentation/Bugs/pdf-extraction-passthrough.md new file mode 100644 index 0000000..65facfa --- /dev/null +++ b/Documentation/Bugs/pdf-extraction-passthrough.md @@ -0,0 +1,43 @@ +# Bug: PDF "extraction" embeds raw bytes, not text + +**Status:** Open · **Severity:** Medium (silent quality loss) · **Reported:** 2026-07-06 +**Component:** `internal/extractor` · **Track:** separate from the local-embeddings epic (pre-existing) + +## Summary + +PDF files are treated as "supported," but their text is never actually extracted. +`extractPDFPassthrough` (`internal/extractor/extractor.go:177`) does `os.ReadFile(path)` and returns +the **raw PDF bytes** as the `Text` to chunk and embed. PDFs are binary (compressed content streams, +xref tables, object dictionaries), so what gets embedded is mostly non-text noise, not the +document's readable content. + +## Impact + +- Semantic search over PDFs is effectively broken — matches are against binary noise, not content. +- The README claims "PDF — text content extraction," which the code does not do. +- Under the local-embeddings epic, the same garbage would burn **local CPU** instead of API dollars. + +## Evidence + +`internal/extractor/extractor.go`: + +```go +func extractPDFPassthrough(path string) (Result, error) { + data, err := os.ReadFile(path) + if err != nil { return Result{}, err } + return Result{Text: string(data)}, nil // raw PDF bytes, not extracted text +} +``` + +`.pdf` is registered in `binaryExtractors`, so the file is reported as supported and indexed. + +## Proposed fix (separate task) + +Replace the passthrough with real PDF text extraction (a pure-Go PDF text library returning +concatenated page text), and add a unit test with a small sample PDF. Keep this **out of the +local-embeddings epic scope** — it is an independent defect. + +## Notes + +Found during the Phase-0 "get familiar / try to break it" pass. Logged per the bug-report workflow; +not a blocker for the epic. diff --git a/Documentation/Epics/background-presence.md b/Documentation/Epics/background-presence.md new file mode 100644 index 0000000..84ee1e8 --- /dev/null +++ b/Documentation/Epics/background-presence.md @@ -0,0 +1,82 @@ +# Epic: Background Presence (cross-platform tray, start-on-login, single instance) + +**Date:** 2026-08-12 +**Status:** Proposed — definition requested in PR #2 review; first in the post-merge epic queue +**Owner:** Bo Motlagh (definition scoped by Bo; drafted by Nestor per review) + +## Goal + +Close the product gap the close=quit triage left on Windows/Linux: **the app only watches +and indexes while its window is open.** On macOS the status-bar tray keeps the app resident +in the background; on Windows/Linux, closing the window now stops watching entirely (stdio +search keeps working against the existing index, but the index goes stale until the next +launch). Windows has a first-class notification-area pattern for exactly this — the current +limitation is only that `tray.go` is darwin-only Objective-C. + +After this epic: on every platform that supports it, closing the window hides the app into +a tray with a Show/Quit menu and watching continues; the app can optionally start on login; +and running two instances against one DB is structurally impossible. + +## User-set constraints (settled in the PR #2 review — do not relitigate) + +- Cross-platform tray/notification-area icon with **Show / Quit**, restoring hide-on-close + on platforms that have a tray. +- **Start-on-login option**, per-OS mechanism (macOS Login Items / Windows registry Run key / + Linux XDG autostart). +- **Linux reality check is part of the epic**: tray support varies by desktop environment + (appindicator vs legacy tray protocols). Document what we target and what degrades to + close=quit — degradation is acceptable, silence about it is not. +- **Single-instance guard** so the zombie-stacking class of bug is structurally impossible + regardless of tray state. +- Current **close=quit stays as the documented fallback** wherever a tray isn't available. + +## Execution Notes (read first if you are the implementing session) + +- Follow the conventions proven by `local-embeddings.md`: phase gates with explicit + **Verify** steps, record deviations in this file, never start a phase before the prior + gate passes. +- The `hasTray` seam already exists: `tray.go` (darwin) / `tray_stub.go` (!darwin) each own + the constant, and `main.go` keys `HideWindowOnClose` off it. This epic's tray + implementations replace the stub per platform and flip `hasTray` there — `main.go` should + need no changes for the hide-on-close behavior. +- Shutdown correctness is already handled (engine.Stop-before-Close + atomic index writes, + PR #2 revision round). Tray-Quit must go through the same shutdown path. + +## Research the epic must settle (Phase 0) + +1. **Library evaluation — likely `fyne-io/systray`**: the known hard question is message-loop + integration with Wails v2 (both want the main thread on some platforms). Spike: tray icon + + Show/Quit alongside a running Wails window on Windows first, then Linux. Record findings + here the way local-embeddings recorded its Phase 0 numbers. Fallback candidates if it + fails: platform-native minimal implementations (Win32 Shell_NotifyIcon via CGo; keep the + existing darwin ObjC). +2. **Linux tray matrix**: GNOME (needs appindicator extension), KDE, XFCE, Pop!_OS — what + works out of the box, what needs a package, what doesn't work at all. Output: a support + table in this doc + the degrade-to-close=quit list. +3. **Single-instance mechanism**: evaluate a lock file beside the DB (flock/LockFileEx) vs a + local socket. Must handle stale locks after crashes; second launch should surface the + existing instance's window when possible rather than just erroring. +4. **Start-on-login mechanics** per OS, including uninstall/cleanup behavior (removing the + app must not leave a broken login entry). + +## Phases (outline — implementation plan written after Phase 0, per workflow) + +- **Phase 0 — Spike (throwaway):** systray×Wails coexistence on Windows + one Linux DE. + **Gate:** tray icon + working Show/Quit alongside a live Wails window, findings recorded here. +- **Phase 1 — Single-instance guard** (independent of tray; kills the zombie class on its + own). **Verify:** second launch on each OS surfaces/exits cleanly; crash + relaunch does + not deadlock on a stale lock. +- **Phase 2 — Windows tray** (flip `hasTray` on windows, hide-on-close restored, Quit goes + through full shutdown). **Verify:** close → still watching (index a file while hidden); + Quit → process exits, no zombies. +- **Phase 3 — Linux tray** per the support matrix, with documented degradation. +- **Phase 4 — Start-on-login** (all platforms, off by default, Settings toggle). +- **Phase 5 — Docs + epic close-out** (README/FEATURES/ARCHITECTURE; record measured + behavior per platform). + +## Out of scope + +- Any change to indexing/search behavior (this epic is presence/lifecycle only). +- macOS tray rework (existing ObjC tray stays; only refactor if the Phase 0 library choice + makes unification free). +- Auto-update, menubar richness beyond Show/Quit (parking lot). diff --git a/Documentation/Epics/local-embeddings.md b/Documentation/Epics/local-embeddings.md index 51ce4b3..7b83fe0 100644 --- a/Documentation/Epics/local-embeddings.md +++ b/Documentation/Epics/local-embeddings.md @@ -1,7 +1,13 @@ # Epic: Local Embeddings as the Primary Provider **Date:** 2026-07-02 -**Status:** Proposed +**Status:** Complete pending review — all phases (0–6) executed and verified. Every shipping +target (macOS arm64/x86_64, Linux arm64/x64, Windows x64) built, tested, and — for Windows and +Linux — GUI field-tested on real hardware with real data. The Windows field test surfaced and +fixed a critical latent cross-platform bug (multi-chunk embed overflow, PR #10 + bug report). +Delivered as a stacked PR chain #2→#3→#4→#6→#7→#10 plus independent fixes #5/#8/#9, all draft, +awaiting Bo's review. Flips to Complete when the stack merges (one post-merge chore: re-run +the Windows suite after #5 lands). **Owner:** Bo Motlagh ## Goal @@ -122,6 +128,49 @@ newer-headers/older-runtime mismatch). The artifact is hosted in this repo's Git pinned by SHA-256 in `assets/manifest.json` like every other artifact — `make assets` treats it identically to officially-published binaries. +## Phase 0 Results + +**Date:** 2026-07-06 · **Outcome: GATE PASSED** — local embedding stack proven end-to-end on +macOS arm64 (throwaway spike, scratchpad only; no repo code touched). + +**Verified stack (all pinned versions worked):** ONNX Runtime 1.26.0 CPU via `onnxruntime_go` +v1.31.0; tokenizer via `daulet/tokenizers` v1.27.0 (prebuilt `libtokenizers.darwin-arm64`); +models `Xenova/multilingual-e5-small` int8 (118 MB) and +`ibm-granite/granite-embedding-97m-multilingual-r2` int8 (98 MB). Pipeline +tokenizer → ORT → mean-pool → L2-normalize produced a valid **384-dim, L2-norm = 1.0** vector; +multilingual tokenization (EN/ES/ZH/AR) sane. + +| Metric | mE5-small | Granite-97m | +|---|---|---| +| Cold load | 129 ms | 145 ms | +| Per-chunk @ batch 16 (~300 tok) | ~34 ms (~30/s) | ~32 ms | +| Peak RSS (upper bound) | ~1.3 GB | ~1.5 GB | +| Disk | 118 MB | 98 MB | + +**Quality (cosine; distance = 1 − cos):** ordering correct — related 0.13–0.18 < +cross-lingual 0.17–0.22 < unrelated ~0.29 (mE5). Cross-lingual (EN↔ES) retrieval works well. + +**Model decision:** **`multilingual-e5-small` is the committed default.** **Granite-97m logged as +a future upgrade candidate** — it loaded and ran at full speed on arm64 (the AVX2-int8 concern did +not materialize), smaller file, crisper related/unrelated separation, but marginally weaker +cross-lingual; needs a broader recall eval + prefix-convention decision before promotion. Upside, +not a dependency (model = data file). + +**Findings that adjust the plan:** + +1. **mE5 requires `token_type_ids`.** The Xenova export takes **three** INT64 inputs + (`input_ids`, `attention_mask`, `token_type_ids`), not two — pass a zero tensor for + `token_type_ids` or ORT errors. Affects the Phase 2 `LocalEmbedder` inference code. (Granite + takes the two-input form.) +2. **Default search threshold `1.5` is miscalibrated for local vectors** (real related-vs-unrelated + separation is ~0.25 cosine distance; 1.5 admits nearly everything). Make the default + **provider-aware**, and confirm which distance metric sqlite-vec's `vec0` table is configured for + before tuning. +3. **Peak RSS ~1 GB** (ORT memory arena + batch activations) — constrain ORT arena / batch size in + `LocalEmbedder` for the desktop memory budget. +4. **Build flags:** minimal `CGO_LDFLAGS="-L -ltokenizers"` suffices on darwin arm64; no + `-framework` flags needed (the lib embeds `-ldl -lm`). + ## Architecture ### Design principles (unchanged) @@ -297,17 +346,35 @@ and the indexing-progress UX all already exist and are the extension points. | File | Purpose | |---|---| -| `internal/embeddings/local/local.go` | `LocalEmbedder` (lazy ONNX session, tokenize→infer→pool→normalize, prefixes) | -| `internal/embeddings/local/assets.go` | `go:embed` + extract-to-`~/.agent-memory/runtime/` with checksums | -| `internal/embeddings/local/local_test.go` | Unit tests + build-tagged integration test | -| `internal/chunker/hf_tokenizer.go` (or similar) | HF tokenizer adapter satisfying `chunker.Tokenizer` | -| `assets/manifest.json` | Pinned artifact URLs + SHA-256 (in git) | +| `internal/embeddings/local/local.go` | `LocalEmbedder` (implements `Embedder`): prefixes, sub-batching, tensor assembly (incl. zero `token_type_ids`), mean-pool, L2-normalize, lazy-session orchestration. **Pure Go, no build tag** — unit-tested via fake session/tokenizer seams. | +| `internal/embeddings/local/session_ort.go` | **(`//go:build localembed`)** real ONNX Runtime session via `onnxruntime_go` (thread cap + arena limit). | +| `internal/embeddings/local/tokenizer_hf.go` | **(`//go:build localembed`)** HF tokenizer via `daulet/tokenizers` + the `chunker.Tokenizer` adapter. **Relocated here from `internal/chunker/` (Phase 2 decision — see below)** to keep the pure-Go `chunker` package CGo-free. | +| `internal/embeddings/local/assets.go` | asset resolution + checksummed atomic extraction to `~/.agent-memory/runtime//`. Source = dev env path (`AGENT_MEMORY_LOCAL_ASSETS`) in Phase 2; swapped to `go:embed` in Phase 3. | +| `internal/embeddings/local/local_test.go` | unit tests (fake session+tokenizer, no tag, always run) + `//go:build localembed` integration test. | +| `assets/manifest.json` | Pinned artifact URLs + SHA-256 (in git) — populated in Phase 3. | + +**Phase 2 implementation decisions (recorded 2026-07-07, per the "update the tables, don't silently diverge" rule):** + +1. **CGo isolation via a `//go:build localembed` tag + interface seams.** The ONNX/tokenizer + infrastructure (which requires native libs) lives in tagged files behind `onnxSession` / + `tokenizer` interfaces; the pure pipeline logic and unit tests carry no tag. This keeps the + default `go build ./...` / `go test ./...` (and CI) green and native-lib-free through Phase 2; + `make build` gains the tag + `CGO_LDFLAGS` in Phase 3. Reinforces the architecture rules + (interfaces at boundaries, mocks for all seams). +2. **HF tokenizer adapter moved from `internal/chunker/` to `internal/embeddings/local/`.** Placing + the CGo-dependent adapter in the `chunker` package would force that pure-Go domain package (and + its always-run tests) to require the native tokenizer lib. `main.go` injects it via the + `chunker.WithTokenizer` seam. Better honors separation of concerns than the original placement. +3. **Phase 2 sources assets from a dev path** (env var), not `go:embed`; weights/libs are never + committed. `go:embed` + `make assets` + manifest land in Phase 3. ### Explicitly unchanged `internal/watcher/`, `internal/extractor/`, `internal/mcp/dispatch.go`, `internal/mcp/stdio.go` -(interface types unchanged — `ReadOnlyEngineService` signature is stable), `tray.go`, +(interface types unchanged — `ReadOnlyEngineService` signature is stable), `internal/store/sqlite_readonly.go`, all search/KNN logic in `store.Search`. +(~~`tray.go`~~ — removed from this list in Phase 5: it required a `//go:build darwin` guard ++ non-darwin stub to compile anywhere but macOS; see Phase 5 Linux findings.) ### Known dead-code risks to check at the end @@ -367,6 +434,94 @@ and the indexing-progress UX all already exist and are the extension points. ### Phase 5 — Cross-platform builds +**Execution decisions (recorded 2026-07-16, per the "update the tables, don't silently diverge" rule):** + +- **Delivered as small stacked PRs** (Bo's small-PRs rule; the epic doesn't mandate one PR): + foundation → Linux → macOS x86_64 → Windows. Windows moved last purely for convenience + (only target needing a second machine for the Rust tokenizer build); no dependency reason. +- **Foundation PR (`feat/xplat-foundation`)** implements Open Concern #8: the ORT shared + library's `go:embed` moves from `assets_embed.go` into per-platform + `assets_embed__.go` files (each defines `embeddedORTLib` + `ortLibFile`; + darwin-arm64 first). Model + tokenizer stay in the shared embed (platform-independent + bytes). Each later platform PR adds one sibling file + manifest entries only. +- **Makefile checksum portability**: `shasum -a 256` (macOS-only) replaced by a `SHA256` + variable that picks `sha256sum` on Linux — prerequisite for `make assets` inside Docker/CI. +- `dylibCandidates` gains the versioned Linux name (`libonnxruntime.so.1.26.0`) that the + official Linux tarball actually ships. + +**Linux findings (recorded 2026-07-16, branch `feat/xplat-linux`):** + +- **`tray.go` was never darwin-guarded** — its Cocoa CGo preamble compiled on every platform, + making *any* non-macOS build impossible. Fixed here (`//go:build darwin` + no-op + `tray_stub.go`); this removes tray.go from the "Explicitly unchanged" list (deviation + recorded per the rules — CLAUDE.md already declared it macOS-only in intent). +- **Pre-existing watcher bug surfaced by first-ever Linux test run:** inotify emits + CREATE+WRITE for a new file; the per-path debouncer replaces rather than merges events, so + OnCreate is swallowed (OnModify fires instead — no user impact today since both index the + file). Documented in `Documentation/Bugs/fswatcher-create-event-swallowed-linux.md`; + fix deliberately kept out of this epic's PRs (watcher is out of scope). +- **linux-arm64 verified end-to-end** in Docker (golang:1.26-bookworm): `make assets` + checksums, test suite (watcher known-fail excepted), real int8 inference (cosine ordering + 0.140 < 0.171 < 0.286, matching Phase 0), full Wails build (webkit2gtk-4.1 via the + `webkit2_41` tag), and an offline (`--network none`) MCP search that retrieved correct + semantic matches from a macOS-indexed DB — cross-platform vector compatibility confirmed. +- **linux-amd64**: artifacts pinned and checksum-verified; runtime smoke still pending + (needs an x64 Linux env — Docker on the arm64 dev Mac would run it under slow emulation). + +**macOS x86_64 findings (recorded 2026-07-17, branch `feat/xplat-macos-intel`):** + +- **ORT 1.26.0 source-built for mac-Intel** exactly as planned (official prebuilts stopped at + 1.23): cross-compiled from the arm64 dev Mac (`./build.sh --config Release --osx_arch x86_64 + --build_shared_lib --skip_tests`; needs CMake ≥4 ok, Python ≥3.10 — system 3.9 fails on + `match` syntax). 27.5 MB dylib, published as repo release **`ort-1.26.0-darwin-x64`** + (prerelease, reproducible recipe in its notes), pinned in the manifest like every artifact. +- **Darwin embed file consolidated**: `assets_embed_darwin_arm64.go` → + `assets_embed_darwin.go` (`localembed && darwin`) — both mac arches share the dylib file + name, mirroring the Linux single-file pattern; the arch difference is which artifact + `make assets` fetches. +- **New `make build-darwin-amd64`** target cross-builds the Intel app from an arm64 Mac + (GOARCH=amd64 assets fetch + `wails build -platform darwin/amd64`). +- **Arch-collision bug found and fixed before ship:** the runtime extraction dir + (`~/.agent-memory/runtime//`) was not arch-namespaced, so an Intel build (or a + home dir migrated from an Intel Mac) poisoned the arm64 build's dlopen with an + incompatible-architecture dylib. Extraction dirs are now `--`. + No user impact (the extraction scheme never shipped — it exists only in this PR stack). +- **Verified under Rosetta 2 on the arm64 dev Mac:** integration test passes as an x86_64 + binary (cosine ordering 0.140 < 0.171 < 0.288, matching all other platforms); + `GOARCH=amd64 make assets` downloads + checksum-verifies from the repo release; full + x86_64 Wails app builds, extracts its assets to its own arch dir, and answers an MCP + search correctly with both arch dirs coexisting. Native arm64 re-verified after. + +**Windows x64 findings (recorded 2026-07-17, branch `feat/xplat-windows`):** + +- **Built on-device over SSH** (Tailscale) to the Windows PC, driven from the dev Mac. + Toolchain installed via scoop: Go 1.26.5, MinGW gcc 16.1.0, rustup-gnu + cargo 1.97.1, + make. Recipe + provenance in `Documentation/windows-build.md`. +- **libtokenizers.a source-built** (no upstream Windows prebuilt): `daulet/tokenizers` v1.27.0 + tag, **GNU** Rust toolchain (must match MinGW gcc; MSVC would produce an unlinkable `.lib`). + The v1.27.0 layout emits `libtokenizers_ffi.a` — renamed to `libtokenizers.a` on packaging. + Published as repo release **`tokenizers-1.27.0-windows-x64`**, SHA-256 pinned. ONNX Runtime + DLL is the **official** Microsoft `onnxruntime-win-x64-1.26.0.zip` (member-pinned). +- **`assets_embed_windows.go`** embeds `onnxruntime.dll`; **Makefile** learned to extract + `.zip` archives (unzip → Windows `tar.exe` fallback) so `make assets` runs on Windows. +- **Two Windows-only build gaps found and fixed in the Makefile** (both invisible on + macOS/Linux, guarded by `GOOS=windows`): + 1. The Rust static lib needs NT/Winsock/crypto syscall libs MinGW doesn't link by default + (`undefined reference to Nt*/Rtl*`) → `LINK_LIBS` appends `-lntdll -lws2_32 -lbcrypt + -luserenv -ladvapi32 -lkernel32 -lncrypt`. + 2. sqlite-vec's cgo `#include "sqlite3.h"` has no system header on Windows (macOS SDK / + Linux libsqlite3 supplied it) → new `winhdr` step stages the headers mattn/go-sqlite3 + bundles (its `sqlite3-binding.h` IS the amalgamation `sqlite3.h`) into `build/winhdr`, + version-matched to the SQLite mattn compiles in. +- **Verified natively on Windows x64:** `make assets` downloads + checksum-verifies all four + artifacts (incl. `.zip` extraction and our published tokenizer release); integration test + links and passes (cosine ordering 0.140 < 0.171 < 0.286, matching every platform); full + `make build` produces `agent-memory.exe` (193 MB PE32+ GUI); the exe extracts its embedded + assets to `windows-amd64-/` and answers an MCP search on a Mac-indexed DB + ("how do I cook italian pasta" → carbonara-recipe.md) — cross-platform vector compatibility + confirmed. (Network not forcibly disabled — SSH session — but the local provider makes no + outbound calls by design; the `--network none` Linux/Intel runs already proved that property.) + 1. Linux x64/arm64: assets manifest entries, CI build, smoke test. 2. Windows x64: CI job builds `libtokenizers.a` with Rust toolchain (no published binary); ORT DLL from official release; smoke test. @@ -383,7 +538,135 @@ and the indexing-progress UX all already exist and are the extension points. 2. Architecture rule check per `Documentation/ARCHITECTURE.md` (inward deps, wiring in main.go, mocks for all interfaces, thin delivery layer). 3. `go vet ./...`, `go test ./...`, `make build`, full manual test both modes. -4. Update this epic's Status to Complete; record the chosen default model and measured numbers. +4. **Verification gaps carried from Phase 5** (recorded 2026-07-17 — coverage gaps, not known + defects; each Phase 5 platform's core inference + search path IS verified): + - [x] **Windows GUI smoke** — DONE 2026-07-17/18, and it earned its keep: the field test + (keyless onboarding ✓, real 109-file vault indexed, dossier search via Claude Desktop + returning correct contextual results ✓) surfaced **three real bugs**, all fixed: + 1. **Token-budget overflow** — every multi-chunk file (>~1.7 KB) silently failed to + embed on EVERY platform, latent since Phase 3 (all prior verifications used tiny + single-chunk corpora). Fixed + permanent large-doc integration test: **PR #10**; + report `Documentation/Bugs/local-embed-token-budget-overflow.md`. Field-verified: + 105 files / 932 chunks / 79 multi-chunk / 0 errors (was 26/26/0). + 2. **Index failures invisible** (stderr only, nothing in the Log page): **PR #8**. + 3. **Close-window zombies on Windows** (HideWindowOnClose without the mac-only tray; + six concurrent instances accumulated): **PR #9**. + Observations for Bo (not fixed): onboarding registers the folder but indexing starts + only via the Dashboard button (intentional per app.go comment — confirm UX intent); + Microsoft-Store-installed Claude Desktop reads its config from the MSIX sandbox + (`%LOCALAPPDATA%\Packages\Claude_*\LocalCache\Roaming\Claude\`), so the Install button + can't reach it — documented limitation, needs a decision on Store-install detection. + - [x] **Linux GUI field test** — DONE 2026-07-20 (Surface, Pop!_OS 24.04, fix #10 build): + first-ever Linux GUI run — keyless onboarding ✓, indexed 4 large docs (76 chunks, all + multi-chunk, 0 errors) ✓, MCP semantic search correct ✓. Throughput ~4 chunks/s on the + older dual-core i7 (vs ~30/s on the M-series dev Mac) — expected CPU scaling, no + pathology. **Packaging lesson:** mid-test the machine upgraded 22.04→24.04, which + REMOVES webkit2gtk-4.0 — the 4.0-linked binary stopped loading. Linux release builds + must target **webkit2gtk-4.1** (`-tags webkit2_41`; present on 22.04 AND 24.04; the + linux-arm64 build already does — only the ad-hoc x64 test build used 4.0). + - [x] **Full untagged `go test ./...` on Windows** — DONE 2026-07-17: every package green + except `TestOnCreate`, which fails exactly as on Linux (assumption CONFIRMED — Windows + delivers the same create+write double-event). PR #5's fix branch verified on Windows: + watcher suite 3/3 pass. So the sole Windows failure is the known bug with a proven fix; + re-run once #5 merges into the stack for the final green checkmark. + - [x] **linux-amd64 runtime smoke** — DONE 2026-07-17 on real x64 hardware (Surface Book + i7, Pop!_OS 22.04): binary built in an Ubuntu 22.04 amd64 container (glibc-matched; + `go build` with wails production tags — the full wails-CLI build path was already proven + on linux-arm64), integration test passed in-container (0.140 < 0.171 < 0.286), then the + binary ran natively on the Surface with **zero library installs** (`ldd` clean against + stock webkit2gtk-4.0/gtk-3), extracted assets to `linux-amd64-/`, and + answered the MCP search correctly against a macOS-indexed DB. +5. Update this epic's Status to Complete; record the chosen default model and measured numbers. + +**Review revision round (2026-08-12, addressing Bo's PR #2 review):** the shipped +dimension-only read-only guard was an unrecorded simplification of this epic's specified +`embedding_fingerprint` — now resolved per the review: the engine records the full +`provider:model:dimensions` fingerprint on every index run/reset (bare dimension kept for +pre-fingerprint DBs), and the read-only guard compares fingerprints first, so a +same-dimension model swap (e.g. granite-97m, also 384-dim) is caught instead of silently +returning mixed-vector garbage. Also landed in the round: crash-safe shutdown (atomic +per-file index transaction + Stop-waits-for-file-boundary), activity-log retention +(30-day TTL / 5000-row cap) with per-path error upsert (no more identical rows per +launch), error logging consolidated inside IndexFile, deterministic cross-platform +debounce-merge tests, and `[]` (not `null`) for empty search results. + +**Phase 6 executed 2026-07-20 — results:** + +- **Orphan hunt: clean.** Zero callers of the old `Embed()` name anywhere. Every + `"text-embedding-3-small"` / `1536` occurrence lives in a legitimate OpenAI-path site + (defaults.go, openai.go, OpenAI UI pickers, tests asserting OpenAI resolution). One + deliberate fallback kept: `store.defaultVecDimension = 1536` (dim ≤ 0 fallback, documented + as historical-schema preservation; the single production caller always passes the resolved + dimension). Fixed one stale comment (domain.SearchOptions.Threshold now points at the + provider-aware `embeddings.DefaultThreshold`: openai 1.5 / local 0.6). +- **Architecture audit: passes.** Inward-only imports verified package-by-package + (store→domain only — never imports embeddings; embeddings/chunker/watcher/extractor + import no siblings; engine depends on interfaces of all domains; mcp defines its own + service interfaces and imports only domain). All wiring in main.go. All six domain + interfaces have mocks (`chunker.Tokenizer`, a two-method seam, uses package-local fakes + in its consumers — acceptable). Delivery layer thin (app.go pass-through + the sanctioned + provider-swap flow). +- **Formal pass: green.** `go vet` clean; untagged suite 8/8 packages; tagged integration + suite (incl. the large-document regression) passes; `make build` (arm64, fix included); + stdio mode answers a semantic search correctly with the fixed binary; GUI mode + Phase-4-verified and field-verified on Windows + Linux with the fix. +- **Post-review items only:** re-run the Windows suite once PR #5 merges (expect all-green); + flip Status from "Complete pending review" to "Complete" when the PR stack lands. + +**Final measured numbers (recorded per this checklist):** default model +`multilingual-e5-small` int8 (384-dim, 512-token ctx, effective chunk budget 480); +cosine-distance ordering related ≈0.13–0.14 < cross-lingual ≈0.17 < unrelated ≈0.28–0.29, +reproduced identically on macOS arm64, macOS x86_64 (Rosetta), Linux arm64/x64, Windows x64; +throughput ~30 chunks/s (M-series) to ~4 chunks/s (2016 dual-core i7); field scale: +105 files / 932 chunks real vault on Windows, offline, zero errors. + +## Phase 3 — Detailed Plan (DRAFT, pending Bo review) + +Drafted 2026-07-07 as the method-level plan for the Phase 3 outline above. **Not yet built** — +Phase 3 is the biggest, behavior-changing phase and its architecture decisions want Bo's sign-off +first. **Recommend delivering as 3 smaller PRs** (3a build/bundling → 3b config/switching → 3c +safety) to fit the "small PRs, never break anything" model. + +### 3a — Build & distribution (no running-app behavior change yet) + +| File | Change | +|---|---| +| `Makefile` | New `assets` target: download pinned model/tokenizer/ORT-lib into `assets/embedded/` (gitignored), verify SHA-256 vs `assets/manifest.json`. `build` depends on `assets`; passes `-tags localembed` + `CGO_LDFLAGS` for `libtokenizers.a`. | +| `assets/manifest.json` (new) | Pinned URLs + SHA-256 (in git). | +| `.gitignore` | `assets/embedded/`. | +| `internal/embeddings/local/assets_embed.go` (new, `//go:build localembed`) | `go:embed` model+tokenizer+dylib; extract via existing `extractAndVerify` to `~/.agent-memory/runtime//`; `resolveAssets` falls back to this when no dev dir/env is set. | +| `internal/embeddings/local/assets_embed_stub.go` (new, `//go:build !localembed`) | "no embedded assets" → keeps default build small/green. | + +**Verify:** `make assets && make build` runs the local model offline; plain `go test ./...` stays green/lib-free. + +### 3b — Config, factory & provider switching (local becomes default) + +| File | Change (method-level) | +|---|---| +| `app.go` | `EmbedderFactory` → `func(provider, apiKey, model) (embeddings.Embedder, error)`. `SetConfig` gains `embedding_provider` case: swap embedder **and chunker** → `engine.Reset()`. Provider-aware `embedding_model`/`openai_api_key` handling. | +| `internal/engine/engine.go` | **Add `SetChunker(c chunker.Chunker)`** (Open Concern #1) — provider switch swaps tokenizer/chunker atomically with the embedder. | +| `main.go` | Both branches read `embedding_provider`, build via factory (`local`→`local.New`, `openai`→`NewOpenAIEmbedder`); GUI wires embedder-matched tokenizer into chunker (`WithTokenizer` + `WithMaxInputTokens`); **stdio stays lazy**. | +| `internal/embeddings/defaults.go` (new) | Centralize `defaultModel(provider)` / `defaultDimension(provider, model)` (Open Concern #4) — remove scattered `"text-embedding-3-small"`/`1536` hardcodes. | + +**Verify:** fresh DB indexes a test folder fully offline; provider switch local↔openai triggers reset + re-index. + +### 3c — Safety, store & Stats + +| File | Change | +|---|---| +| `internal/store/sqlite.go` | `migrate()`/`Reset()` create `chunk_embeddings` at the active provider's dimension (passed in — see Decision 2), not hardcoded. `Stats()` reports `embedding_provider`+`embedding_model` from config (layering fix). | +| `app.go` | Write `embedding_fingerprint` (`provider:model:dim`) on each index run; check at GUI startup → mismatch surfaces "re-index required", not garbage. | +| `internal/engine/readonly.go` | Read-only dimension guard (Open Concern #2): compare embedder `Dimensions()`/fingerprint to stored table; mismatch → actionable error, not raw sqlite-vec failure. | +| `internal/mcp/server.go` | `index_status` gains a `provider` field (from `Stats`). | +| threshold | Confirm sqlite-vec metric; provider-aware default threshold from Phase 0 numbers (~0.25 local), centralized (Open Concern #5). | + +**Verify:** existing OpenAI-vectored DB hits the fingerprint-mismatch path (not garbage); `--mcp` returns a clear error on dimension mismatch. + +### Decisions needing Bo's sign-off + +1. **`make build` now requires native libs + `-tags localembed`** (via `make assets`, ~150 MB); binary ~26 MB → ~180 MB. Epic accepts this — confirm. +2. **Store dimension: reorder the composition root** so `main.go` resolves provider/model/dimension up front and **passes the dimension into the store**, rather than `migrate()` hardcoding it — avoids a new hardcode *and* keeps `store` from importing `embeddings`. Touches wiring order. +3. **Existing-user handling:** when `embedding_provider` is unset on an upgraded DB with OpenAI vectors + a key → default to **openai** (preserve their setup), not local. Full onboarding migration is Phase 4; the resolution rule starts here. ## Risks @@ -399,6 +682,23 @@ and the indexing-progress UX all already exist and are the extension points. | Binary size ~180 MB | Accepted | Decision made 2026-07-02; the trade for zero API cost | | int8 quantization quality drop vs fp32 | Low | Phase 0 sanity comparison; fp32 fallback possible at 449 MB if unacceptable | +## Open Concerns & Plan Adjustments (from review + Phase 0) + +Captured 2026-07-06 during pre-implementation review and the Phase 0 spike. Each item is tagged to +the phase that must address it, so nothing is lost mid-epic. + +| # | Concern | Fix in | Note | +|---|---|---|---| +| 1 | **Chunker not swapped on runtime provider switch.** `engine` has `SetEmbedder` but no `SetChunker`; `app.SetConfig` swaps only the embedder — switching OpenAI↔local at runtime would leave the wrong tokenizer/clamp. | **Phase 3** | Add `engine.SetChunker` (or a combined provider swap); `SetConfig` provider case swaps embedder + chunker atomically. | +| 2 | **Read-only stdio dimension mismatch unhandled.** `readonly.Search` embeds the query with a config-derived embedder; if its dim disagrees with the stored vectors, sqlite-vec errors and the read-only process cannot re-index. | **Phase 3** | stdio compares embedder `Dimensions()`/fingerprint to the stored table; return an actionable error ("rebuild in the GUI"), not a raw failure. | +| 3 | **Existing-user onboarding regression.** Switching the gate to `onboarding_complete` re-onboards current users and defaults them to local, mismatching their OpenAI vectors. | **Phase 4** | Migration: if the DB has watched dirs or a saved key, back-fill `onboarding_complete=true` and preserve `provider=openai` for existing OpenAI DBs. | +| 4 | **Scattered default-model/dimension hardcodes** (`"text-embedding-3-small"` ×5, `1536` ×2). | **Phase 3** (seed in Phase 1) | Centralize `defaultModel(provider)` / `defaultDimension(provider,model)`; stop trading a `1536`→`384` hardcode. | +| 5 | **Threshold `1.5` duplicated (~4 sites) and miscalibrated** for local (Phase 0: real separation ~0.25 cosine distance). | **Phase 3** | Centralize the default; make it provider-aware; confirm sqlite-vec's configured distance metric. | +| 6 | **mE5 needs `token_type_ids`** (3 INT64 inputs, not 2). | **Phase 2** | Pass a zero tensor in `LocalEmbedder`. | +| 7 | **Peak RSS ~1 GB** (ORT arena + batch activations). | **Phase 2** | Constrain ORT arena / batch size. | +| 8 | **`go:embed` must be build-tagged per platform** (one binary can embed only one platform's ORT lib). | **Phase 5** | Build constraints on the embed directives. | +| 9 | **`OpenAIEmbedder.MaxInputTokens()==0` ("no limit")** — footgun if `chunk_size` ever exceeds OpenAI's 8191. | Low priority | Documented trade-off; add a comment. | + ## Out of Scope (this epic) - Ollama / external-runtime provider option diff --git a/Documentation/Epics/mcp-search-filters.md b/Documentation/Epics/mcp-search-filters.md new file mode 100644 index 0000000..38573db --- /dev/null +++ b/Documentation/Epics/mcp-search-filters.md @@ -0,0 +1,87 @@ +# Epic: MCP Search Filters (filesystem-tier metadata) + +**Date:** 2026-08-12 +**Status:** Proposed — definition requested in PR #2 review; queued after `background-presence` +**Owner:** Bo Motlagh (definition scoped by Bo; drafted by Nestor per review) + +## Goal + +The MCP `search` tool (stdio and HTTP) gains **optional filter parameters** that constrain +results by file metadata, so a caller can ask "search only my markdown notes," "only files +under this directory," or "only files modified in the last month" — instead of over-fetching +and filtering client-side after the fact. + +## User-set constraints (settled in the PR #2 review — do not relitigate) + +- **Filesystem-tier metadata only**: path/glob, file type/extension, watched directory, + file modified time, file size, indexed-at time. **Document-embedded properties** + (title/author/created from docx/pdf/etc.) are explicitly **out of scope** for this epic — + parking-lot note only. +- **MCP surface only**: both transports' `search` tool schemas. **No GUI search work.** +- Definition merged in PR #2; implementation sequenced after `background-presence` unless + Bo reorders. + +## Current state (verified against the PR #2 branch) + +- The `files` table carries only `path`, `hash`, `indexed_at`, `directory_id` — **no mtime, + size, or type column**. Filters on modified-time/size need columns captured at index time: + a schema migration plus a backfill decision. Backfill is cheap — an `os.Stat` per file at + the next scan populates the columns **without re-embedding** (the content hash is + unchanged, so the pipeline is skipped; only the file row updates). +- Filter params must flow through `domain.SearchParams` so engine / readonly / store stay + interface-clean; the tool schemas in `internal/mcp/server.go` are the only delivery-layer + touch (both transports dispatch through the same search path). +- `store.Search` today over-fetches (`limit+offset`, ×2 under a threshold) and filters + distance in Go — there is already a post-KNN filtering stage to build on. + +## The key design question (the epic's real research content) + +**How filtering composes with the KNN query** in `store.Search`: sqlite-vec `vec0` KNN does +not trivially JOIN-filter. Research and **measure** both options, recording numbers in this +doc the way local-embeddings recorded its Phase 0 results: + +1. **Over-fetch then filter** (in SQL over the joined rows, or in Go): simple, no schema + coupling to sqlite-vec — but `limit`/`offset`/`threshold` semantics must hold + **post-filter**. A filter that excludes 90% of neighbors must not return 1 result because + the over-fetch was too small: the over-fetch factor needs to adapt (e.g. iterative + re-query with growing K until `limit` post-filter results or the index is exhausted). +2. **sqlite-vec metadata columns / partition keys** in the `vec0` table (newer sqlite-vec + feature): filters pushed into the KNN itself. Verify version availability in our pinned + sqlite-vec, migration cost for existing tables, and which of our filter types it can + express (globs almost certainly not — likely a hybrid: partition/metadata for cheap + equality filters, post-filter for globs/ranges). + +**Design-once note (from the review):** result post-processing (over-fetch → dedupe/filter → +limit/offset/threshold) is the **same machinery `multi-representation-indexing` needs** for +its per-chunk best-score dedupe. Whichever epic lands first should design `store.Search`'s +post-KNN stage to serve both. + +## Execution Notes (read first if you are the implementing session) + +- Epic conventions per `local-embeddings.md`: phase gates with Verify, record deviations here. +- Schema migration must be additive (`ALTER TABLE ADD COLUMN`, nullable) — existing DBs keep + working un-backfilled; NULL metadata simply doesn't match metadata filters until the next + scan backfills (document this in the tool description). +- Tool schema: filters land as optional properties (e.g. `path_glob`, `extensions[]`, + `directory`, `modified_after/before`, `min/max_size`, `indexed_after/before`) — absent + filters must behave byte-for-byte like today's search (regression tests on that). + +## Phases (outline) + +- **Phase 0 — Measure the composition options** on a realistic index (thousands of chunks): + over-fetch-adaptive vs metadata-column KNN, latency + correctness under selective filters. + **Gate:** numbers recorded here; approach chosen. +- **Phase 1 — Schema + backfill:** mtime/size (+type derived from path) columns, migration, + os.Stat backfill on scan without re-embed. **Verify:** upgraded DB backfills on next scan; + no re-embedding occurs (chunk counts and vectors unchanged). +- **Phase 2 — `domain.SearchParams` + store filtering** per the chosen approach, with the + post-filter limit/offset/threshold semantics pinned by tests (including the + filter-eats-90%-of-neighbors case). +- **Phase 3 — MCP tool schemas** (both transports) + docs. **Verify:** filtered searches via + a live MCP client; unfiltered search unchanged. + +## Out of scope / parking lot + +- Document-embedded properties (docx/pdf titles, authors, created dates) — future epic. +- GUI search/filter UI. +- Saved filters, filter persistence, per-directory defaults. diff --git a/Documentation/Epics/multi-representation-indexing.md b/Documentation/Epics/multi-representation-indexing.md new file mode 100644 index 0000000..4686a6a --- /dev/null +++ b/Documentation/Epics/multi-representation-indexing.md @@ -0,0 +1,98 @@ +# Epic: Multi-Representation Indexing (dual-vector chunks) + +**Date:** 2026-08-12 +**Status:** Proposed — definition requested in PR #2 review; queued after `mcp-search-filters` +**Owner:** Bo Motlagh (definition + design input scoped by Bo; drafted by Nestor per review) + +## Problem (verified against the PR #2 branch) + +A chunk's single `content` currently serves two conflicting jobs: **the text we embed** and +**the text we return**. `.html`/`.htm` are in the extractor's `textExtensions` (raw +passthrough), so an HTML chunk's vector substantially describes class names and tags +(`text-slate-800`, `border-collapse`) rather than what the page says — findable only by +markup-shaped queries. Stripping tags before embedding is wrong the other way: for source +files the markup *is* the artifact, and a code-shaped query should match raw and get raw +back. Any per-file "document or code?" classification has no right answer. + +## Proposed shape (Bo's design input — to validate in this epic, not settled) + +Decouple embedded text from stored text: **up to two vectors per chunk**, both pointing at +the same stored raw payload. + +- **Vector A** = embed(raw chunk) — matches code/markup-shaped queries. +- **Vector B** = embed(extracted text of that chunk) — matches prose queries about the + rendered content. +- **Emission is measurable, not taxonomic** — at index time compare extracted length to raw + length: + - Ratio ≈ 1 (prose: md/txt/docx): one vector from raw — exactly today's behavior, no + regression. + - Extracted ≪ raw (html/svg/ipynb): both vectors. + - Extraction ≈ nothing (minified bundles, base64 blobs): one raw vector **plus a + file-level summary chunk** (path, type, title/headings outline) — the same shape as the + ZIP metadata stub the extractor already emits. +- Extraction stays **deterministic and free**: for HTML, walk text nodes + structural + landmarks (``, headings, `alt`, `aria-label`, link/button labels) via + `golang.org/x/net/html`. No LLM, no per-file cost. + +## Binding constraints (settled — do not relitigate) + +- Files on disk are **never modified**. +- `search` returns **raw `content` by default** — existing callers unchanged; an optional + `include_extracted` response field is fine. +- **No LLM** anywhere in the pipeline. +- The extra embed/vector cost lands on local compute and is accepted (~10–15% vector growth + in a mostly-prose corpus). +- **Hard prerequisite: the PDF extraction fix** (see + `Documentation/Bugs/pdf-extraction-passthrough.md` + its scheduled fix) — same mechanism, + and PDFs become "embed extracted, store extracted" with **no raw vector** (nobody queries + for `/Filter /FlateDecode`). + +## Design questions this epic must work out (the real content) + +1. **Schema.** `chunk_embeddings` is `vec0(chunk_id INTEGER PRIMARY KEY, embedding + FLOAT[dim])` — strictly one vector per chunk today. Two vectors per chunk means a + synthetic embedding id with a (chunk_id, representation) mapping — via an aux table or + sqlite-vec's metadata/auxiliary columns — plus migration and `Reset()` changes. +2. **Search semantics.** KNN now returns *embedding* rows, not chunks — results must dedupe + to one row per chunk keeping the best score, which forces over-fetch before + `limit`/`offset`/`threshold` apply. This is the **same post-KNN machinery + `mcp-search-filters` needs**; whichever epic lands second inherits the first's design — + design `store.Search`'s post-processing once for both. +3. **Pairing granularity — the subtlest question.** Chunk boundaries are computed on raw + text, so "the extracted text of chunk 42" means extracting from a raw HTML *fragment* + (`x/net/html` tolerates fragments, but landmark quality mid-document needs verification) + — versus the alternative of two parallel chunkings of the file with looser file-level + pairing. **Spike this first; it decides the data model.** +4. **Threshold interplay.** Raw-vector and extracted-vector distances for the same chunk + will distribute differently; confirm the provider-aware default threshold behaves + sensibly when both representations compete in one result set (measure, per the + local-embeddings Phase 0 pattern). + +## Execution Notes (read first if you are the implementing session) + +- Epic conventions per `local-embeddings.md`: phase gates with Verify, record deviations and + measured numbers in this doc. +- Sequenced **after** `mcp-search-filters` (unless Bo reorders) — and note design question 2 + is shared with it either way. +- The emission ratio thresholds (what counts as "≈1", "≪") must be recorded here once + measured against a real corpus — they are product behavior, not implementation detail. + +## Phases (outline) + +- **Phase 0 — Pairing spike (throwaway):** fragment-extraction quality on real HTML + mid-document vs parallel-chunking alternative. **Gate:** data model decided, recorded here. +- **Phase 1 — Schema migration** ((chunk_id, representation) mapping, Reset, fingerprint + interaction). **Verify:** existing single-vector DBs migrate or trigger the re-index path + cleanly — never silently mixed. +- **Phase 2 — Extraction + emission rules** (ratio-based; html first, svg/ipynb after). + **Verify:** prose corpus produces byte-identical behavior to today (no regression class). +- **Phase 3 — Search dedupe + threshold calibration** (shared machinery with + mcp-search-filters). **Verify:** prose-shaped query finds the HTML page's content; code- + shaped query still finds raw markup; measured distances recorded. +- **Phase 4 — MCP `include_extracted` + docs.** + +## Out of scope / parking lot + +- LLM-generated summaries or descriptions of any kind. +- Per-file-type user configuration of representations. +- OCR / image content extraction. diff --git a/Documentation/FEATURES.md b/Documentation/FEATURES.md index a326a76..f782134 100644 --- a/Documentation/FEATURES.md +++ b/Documentation/FEATURES.md @@ -2,7 +2,7 @@ ## Core Pipeline -**Directory watching and indexing.** Add directories via the GUI or MCP. The engine walks each directory, filters by ignore patterns and supported file types, then runs the pipeline: extract text -> chunk -> embed via OpenAI -> store vectors in SQLite. File content is hashed (SHA-256) so unchanged files are skipped on subsequent scans. +**Directory watching and indexing.** Add directories via the GUI or MCP. The engine walks each directory, filters by ignore patterns and supported file types, then runs the pipeline: extract text -> chunk -> embed (via the active provider — local by default) -> store vectors in SQLite. File content is hashed (SHA-256) so unchanged files are skipped on subsequent scans. **Semantic search.** Query text is embedded, then matched against stored vectors using KNN (sqlite-vec cosine distance). Returns ranked results with file path, chunk text, and similarity score. Available via MCP `search` tool. @@ -24,11 +24,15 @@ Token-based splitting using tiktoken (cl100k_base encoding). Configurable chunk ## Embedding -OpenAI API client supporting `text-embedding-3-small` (1536 dimensions) and `text-embedding-3-large` (3072 dimensions). Batches up to 2048 inputs per API call. Retries with exponential backoff on rate limits (max 3 retries). The engine further batches by token count (250K tokens per batch) to stay within API limits. +Two providers, selectable in Settings. Both satisfy the same `Embedder` interface, so the rest of the pipeline is provider-agnostic. -**Model switching.** Changing the embedding model via Settings drops the vector table, recreates it with the new dimension, and re-indexes everything. The embedder is hot-swapped so no restart is needed. +**Local (default).** A bundled `multilingual-e5-small` model (384 dimensions, ~100 languages) runs in-process on the CPU via ONNX Runtime. Works fully offline — no API key, no network calls. The model, tokenizer, and runtime library are bundled into the app; onboarding requires no key. -**API key changes.** Updating the OpenAI API key immediately swaps the embedder so new requests use the updated key. +**OpenAI (opt-in).** OpenAI API client supporting `text-embedding-3-small` (1536 dimensions) and `text-embedding-3-large` (3072 dimensions). Enabled by adding an API key in Settings. Batches up to 2048 inputs per API call. Retries with exponential backoff on rate limits (max 3 retries). The engine further batches by token count (250K tokens per batch) to stay within API limits. + +**Provider & model switching.** Changing the provider or model via Settings drops the vector table, recreates it with the new dimension, and re-indexes everything (vectors from different models are incompatible). The embedder is hot-swapped so no restart is needed. + +**API key changes.** Updating the OpenAI API key (while the OpenAI provider is active) immediately swaps the embedder so new requests use the updated key. ## Ignore Patterns @@ -68,7 +72,7 @@ Both implement the MCP protocol: `initialize` handshake, `tools/list`, and `tool Wails v2 with React frontend. Native webview, no browser required. -**Onboarding** — first-launch flow: API key entry, directory picker, then dashboard. +**Onboarding** — first-launch flow: welcome, directory picker, then dashboard. No API key required — the local provider works out of the box. Opting into OpenAI is done later in Settings. **Dashboard** — total files, chunks, last indexed time, indexing progress bar, embedding model display. @@ -76,7 +80,7 @@ Wails v2 with React frontend. Native webview, no browser required. **Controls** — Start/Stop/Restart/Reset buttons. Reset clears all vectors and re-indexes (with confirmation dialog). -**Settings** — API key (masked), model picker, chunk size/overlap, MCP port, auth token display with rotate button. +**Settings** — embedding provider selector (local default vs OpenAI); when OpenAI is selected, an API key field (masked) and model picker appear. Plus chunk size/overlap, MCP port, and auth token display with rotate button. Switching provider re-indexes everything. **Activity Log** — paginated log of indexing events (indexed, ignored, deleted, errors). @@ -97,7 +101,7 @@ The config path is customizable for non-standard installs. Claude Desktop launch - HTTP server binds to `127.0.0.1` only — no network exposure. - Bearer token auth on all MCP HTTP requests. Token auto-generated, stored in SQLite, rotatable via GUI. -- Only outbound traffic is HTTPS to `api.openai.com`. +- No outbound traffic by default — the local embedding provider runs in-process. Outbound HTTPS to `api.openai.com` occurs only if the OpenAI provider is opted into. - GUI uses native webview IPC — no localhost web server for the UI. - SQLite file uses standard filesystem permissions. @@ -111,7 +115,7 @@ All settings are stored in the SQLite `config` table and manageable from the GUI **MCP server port:** Default `9847`. Binds to `127.0.0.1` only. -**Embedding model:** `text-embedding-3-small` (1536 dimensions) by default. Switch to `text-embedding-3-large` (3072 dimensions) in Settings. Changing the model re-indexes everything. +**Embedding provider & model:** local `multilingual-e5-small` (384 dimensions) by default. Switch to the OpenAI provider in Settings to use `text-embedding-3-small` (1536 dimensions) or `text-embedding-3-large` (3072 dimensions), which require an API key. Changing the provider or model re-indexes everything. **Chunk size:** 512 tokens default, configurable. Overlap: 50 tokens default, configurable. diff --git a/Documentation/ROADMAP.md b/Documentation/ROADMAP.md index 2c9cc9c..e846008 100644 --- a/Documentation/ROADMAP.md +++ b/Documentation/ROADMAP.md @@ -2,9 +2,10 @@ ## Local Embeddings (Epic) -Replace OpenAI as the default embedding provider with a bundled in-process local model -(ONNX Runtime + multilingual-e5-small); OpenAI becomes an explicit opt-in backup. Full design -and implementation plan: [Epics/local-embeddings.md](Epics/local-embeddings.md). +**Implemented** — a bundled in-process local model (ONNX Runtime + `multilingual-e5-small`, 384-dim) +is now the default embedding provider, working fully offline with no API key; OpenAI is an explicit +opt-in in Settings. Working on macOS arm64; cross-platform builds (Linux, Windows, macOS x86_64) +are the remaining phase. Full design and status: [Epics/local-embeddings.md](Epics/local-embeddings.md). ## Search Configuration UI diff --git a/Documentation/windows-build.md b/Documentation/windows-build.md new file mode 100644 index 0000000..1df3cca --- /dev/null +++ b/Documentation/windows-build.md @@ -0,0 +1,90 @@ +# Windows x64 Build & Artifact Provenance (local-embeddings, Phase 5) + +How the Windows local-embedding build is produced and how to reproduce its one +source-built artifact. Records what was actually run (versions below are what +shipped), so a future maintainer can rebuild after a dependency bump. + +## Why Windows needs a build step at all + +Every other platform gets both native libraries as official prebuilts. Windows +is the exception on one of them: + +| Artifact | Windows source | +|---|---| +| ONNX Runtime DLL | **Official** Microsoft release (`onnxruntime-win-x64-1.26.0.zip`) — pinned directly in `assets/manifest.json`. | +| `libtokenizers.a` (static) | **No upstream Windows prebuilt** — `daulet/tokenizers` publishes macOS/Linux only. We build it from source and host it in this repo's releases (`tokenizers-1.27.0-windows-x64`). | + +## Toolchain (installed via scoop) + +``` +scoop install go mingw rustup-gnu make +rustup default stable-x86_64-pc-windows-gnu +``` + +Verified versions: Go 1.26.5, MinGW gcc 16.1.0, rustc/cargo 1.97.1 (GNU), +GNU make 4.4.1. + +**The GNU pairing is load-bearing.** CGo links the Go binary with MinGW `gcc`, +so the Rust static lib must come from the `*-windows-gnu` toolchain. Building it +with the default MSVC toolchain yields a `.lib` in a format MinGW cannot link. + +## Reproducing `libtokenizers.a` + +``` +git clone --depth 1 --branch v1.27.0 https://github.com/daulet/tokenizers +cd tokenizers +cargo build --release +# Output: target/release/libtokenizers_ffi.a (~40 MB) +``` + +Note the **name change**: this repo layout (v1.27.0) emits `libtokenizers_ffi.a` +from the `tokenizers-ffi` crate. The macOS/Linux prebuilts — and therefore the +`-ltokenizers` link flag and the manifest's `member: libtokenizers.a` — expect +`libtokenizers.a`. Same archive, so we simply rename on packaging: + +``` +cp target/release/libtokenizers_ffi.a libtokenizers.a +tar czf libtokenizers.windows-x86_64.tar.gz libtokenizers.a +# publish as a repo release, pin archive + member SHA-256 in the manifest +``` + +## Building the app + +``` +git clone https://github.com/unitedeffectslabs/agent-memory +cd agent-memory && git checkout feat/xplat-windows +cd frontend && npm install && cd .. +make assets # downloads DLL + model + tokenizer + our published static lib; verifies all checksums +make build # wails build -skipbindings -tags localembed → build/bin/agent-memory.exe +``` + +`make assets` extracts the `.zip` ORT archive via `unzip` if present, else falls +back to Windows' bundled `C:\Windows\System32\tar.exe` (Git Bash's GNU tar +cannot read zips). Checksums use `sha256sum` if present, else `shasum -a 256` +(both provided by Git for Windows). + +## Verification (the Phase 5 gate) + +1. Integration test — real tokenizer + ORT + model, as a native Windows binary: + ``` + CGO_LDFLAGS="-L$PWD/internal/embeddings/local/lib -ltokenizers" \ + go test -tags localembed -count=1 -run TestIntegrationEmbed -v ./internal/embeddings/local/ + ``` + Expect PASS, cosine ordering ≈ related 0.14 < cross-lingual 0.17 < unrelated + 0.29 (matches macOS arm64/x86_64 and Linux). +2. Offline app run: disconnect network, `agent-memory.exe --db test.db`, onboard + with **no API key**, index a small folder, confirm search returns matches. + +## Troubleshooting (observed / anticipated) + +- **Missing Windows system symbols at link** (`undefined reference to Nt*` / + `Rtl*`, `ws2_32`, `bcrypt`, `userenv`): the Rust static lib pulls in NT/Winsock/ + crypto syscalls MinGW doesn't link by default. **Handled** — the Makefile's + `LINK_LIBS` appends `-lntdll -lws2_32 -lbcrypt -luserenv -ladvapi32 -lkernel32 + -lncrypt` when `GOOS=windows` (empty elsewhere). Observed and fixed during the + first Windows build; listed here in case a tokenizers/Rust bump adds more. +- **`onnxruntime_providers_shared.dll` load error**: the official zip ships this + second DLL; CPU-only use normally doesn't need it, but if ORT fails to load, + pin it as an extra `embedded/` artifact beside the main DLL. +- **Defender quarantines the fresh unsigned exe**: add a temp exclusion for the + repo folder while testing. diff --git a/Makefile b/Makefile index 524818f..1dc8e30 100644 --- a/Makefile +++ b/Makefile @@ -1,13 +1,114 @@ -.PHONY: build dev test clean +.PHONY: build build-darwin-amd64 dev test clean assets winhdr -build: - wails build -skipbindings +# --- Local-embedding asset bundling ----------------------------------------- +# Artifacts (model, tokenizer, ONNX Runtime dylib, static tokenizer lib) are +# pinned by URL + SHA-256 in assets/manifest.json and downloaded by `make +# assets` into the local embedding package (gitignored). They are never +# committed. `make build` links the tokenizer static lib via CGO_LDFLAGS and +# compiles with the `localembed` tag (which go:embed's the runtime assets). +MANIFEST := assets/manifest.json +LOCAL_DIR := internal/embeddings/local +EMBED_DIR := $(LOCAL_DIR)/embedded +LIB_DIR := $(LOCAL_DIR)/lib +PLATFORM := $(shell go env GOOS)-$(shell go env GOARCH) +# Portable SHA-256: macOS ships shasum, Linux ships sha256sum. Both print +# "<hash> <file>", so the awk '{print $1}' callers work with either. +SHA256 := $(shell command -v sha256sum >/dev/null 2>&1 && echo sha256sum || echo shasum -a 256) -dev: - wails dev +# Windows-only: the source-built libtokenizers.a (Rust std, GNU toolchain) pulls +# in low-level NT/Winsock/crypto syscalls that MinGW does not link by default. +# Naming them here resolves "undefined reference to Nt*/Rtl*" at link time. +# Empty on macOS/Linux, so those builds are unaffected. +LINK_LIBS := -ltokenizers +CGO_EXTRA_CFLAGS := +WIN_PREREQ := +ifeq ($(shell go env GOOS),windows) +LINK_LIBS += -lntdll -lws2_32 -lbcrypt -luserenv -ladvapi32 -lkernel32 -lncrypt +# sqlite-vec's cgo build #includes sqlite3.h / sqlite3ext.h, which macOS and +# Linux supply from the system but Windows does not. Stage the exact headers +# mattn/go-sqlite3 bundles (its sqlite3-binding.h IS the amalgamation sqlite3.h), +# so they match the SQLite that mattn compiles in — see the winhdr target. +CGO_EXTRA_CFLAGS := -I$(PWD)/build/winhdr +WIN_PREREQ := winhdr +endif + +# Stage sqlite headers for the Windows build from the mattn/go-sqlite3 module +# (Windows has no system sqlite3.h). No-op / unused on macOS and Linux. +winhdr: + @mkdir -p "$(PWD)/build/winhdr" + @d=$$(go list -m -f '{{.Dir}}' github.com/mattn/go-sqlite3); \ + d=$$(cygpath -u "$$d" 2>/dev/null || echo "$$d"); \ + cp "$$d/sqlite3-binding.h" "$(PWD)/build/winhdr/sqlite3.h"; \ + cp "$$d/sqlite3ext.h" "$(PWD)/build/winhdr/sqlite3ext.h"; \ + echo ">> staged Windows sqlite headers from $$d" + +build: assets $(WIN_PREREQ) + CGO_CFLAGS="$(CGO_EXTRA_CFLAGS)" CGO_LDFLAGS="-L$(PWD)/$(LIB_DIR) $(LINK_LIBS)" wails build -skipbindings -tags localembed + +# Cross-build the Intel-mac app from an arm64 Mac. GOARCH=amd64 makes the +# assets target fetch the darwin-amd64 artifacts (the embedded/ and lib/ dirs +# hold ONE platform at a time — the checksum check refetches on arch switch, +# so alternating with `make build` is safe, just re-downloads). +build-darwin-amd64: + GOARCH=amd64 $(MAKE) assets + CGO_LDFLAGS="-L$(PWD)/$(LIB_DIR) -ltokenizers" wails build -skipbindings -tags localembed -platform darwin/amd64 + +dev: assets $(WIN_PREREQ) + CGO_CFLAGS="$(CGO_EXTRA_CFLAGS)" CGO_LDFLAGS="-L$(PWD)/$(LIB_DIR) $(LINK_LIBS)" wails dev -tags localembed test: go test ./... clean: rm -rf build/bin + +# Download + checksum-verify each manifest artifact for the current platform. +# Idempotent: files already present with a matching checksum are skipped. +assets: + @echo ">> fetching local-embedding assets for $(PLATFORM)" + @command -v jq >/dev/null || { echo "ERROR: jq is required"; exit 1; } + @test -f $(MANIFEST) || { echo "ERROR: $(MANIFEST) not found"; exit 1; } + @jq -e '.platforms["$(PLATFORM)"]' $(MANIFEST) >/dev/null 2>&1 || \ + { echo "ERROR: no manifest entry for platform $(PLATFORM)"; exit 1; } + @mkdir -p $(EMBED_DIR) $(LIB_DIR) + @set -e; \ + for key in $$(jq -r '.platforms["$(PLATFORM)"] | keys[]' $(MANIFEST)); do \ + sel() { jq -r ".platforms[\"$(PLATFORM)\"].$$key.$$1 // empty" $(MANIFEST); }; \ + url=$$(sel url); sha=$$(sel sha256); dest=$$(sel dest); \ + member=$$(sel member); membersha=$$(sel member_sha256); \ + destpath=$(LOCAL_DIR)/$$dest; \ + want=$$sha; [ -n "$$membersha" ] && want=$$membersha; \ + if [ -f "$$destpath" ]; then \ + have=$$($(SHA256) "$$destpath" | awk '{print $$1}'); \ + if [ "$$have" = "$$want" ]; then echo " ok (cached) $$dest"; continue; fi; \ + echo " stale, refetching $$dest"; \ + fi; \ + echo " downloading $$key -> $$dest"; \ + tmp=$$(mktemp); \ + curl -fsSL -o "$$tmp" "$$url"; \ + got=$$($(SHA256) "$$tmp" | awk '{print $$1}'); \ + if [ "$$got" != "$$sha" ]; then \ + echo "ERROR: archive checksum mismatch for $$key: got $$got want $$sha"; rm -f "$$tmp"; exit 1; \ + fi; \ + mkdir -p "$$(dirname "$$destpath")"; \ + if [ -n "$$member" ]; then \ + xd=$$(mktemp -d); \ + case "$$url" in \ + *.zip) if command -v unzip >/dev/null 2>&1; then unzip -q "$$tmp" "$$member" -d "$$xd"; \ + elif [ -x /c/Windows/System32/tar.exe ]; then /c/Windows/System32/tar.exe xf "$$tmp" -C "$$xd" "$$member"; \ + else tar xf "$$tmp" -C "$$xd" "$$member"; fi ;; \ + *) tar xzf "$$tmp" -C "$$xd" "$$member" ;; \ + esac; \ + cp "$$xd/$$member" "$$destpath"; \ + rm -rf "$$xd"; \ + got2=$$($(SHA256) "$$destpath" | awk '{print $$1}'); \ + if [ "$$got2" != "$$membersha" ]; then \ + echo "ERROR: member checksum mismatch for $$key: got $$got2 want $$membersha"; rm -f "$$tmp"; exit 1; \ + fi; \ + else \ + cp "$$tmp" "$$destpath"; \ + fi; \ + rm -f "$$tmp"; \ + echo " verified $$dest"; \ + done + @echo ">> assets ready in $(EMBED_DIR) and $(LIB_DIR)" diff --git a/README.md b/README.md index cbe2b52..8773da5 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ Works as a standalone file tracker/indexer and as an MCP server that gives Claud ## How It Works -1. **Launch the app** and enter your OpenAI API key in the onboarding screen (used for generating embeddings — your files never leave your machine except as embedding API calls). +1. **Launch the app.** It works out of the box with a bundled local embedding model — no API key, no account, no network calls. Onboarding is just: welcome → add folders → done. (You *can* opt into OpenAI embeddings later in Settings if you prefer; see below.) 2. **Add directories** you want indexed. The app watches them in real time — new files, edits, and deletions are picked up automatically. 3. **Search your files**, connect Claude Desktop so it can search your files during conversations. @@ -23,19 +23,26 @@ Under the hood, Claude Desktop launches the app as a subprocess (`--mcp` flag) t - **PDFs** — text content extraction - **Media & archives** — metadata only (file name, size, dimensions, contents list) -Files are chunked into ~512-token segments (configurable), embedded via OpenAI, and stored as vectors. Unchanged files (matched by SHA-256 hash) are skipped on re-scan. +Files are chunked into ~512-token segments (configurable), embedded by the active provider (local by default), and stored as vectors. Unchanged files (matched by SHA-256 hash) are skipped on re-scan. ### Supported Embedding Models -- `text-embedding-3-small` (1536 dimensions) — default, faster, cheaper +**Local (default, opt-out):** + +- `multilingual-e5-small` (384 dimensions) — bundled into the app, runs in-process on the CPU, ~100 languages. Works fully offline: no API key, no account, no network calls. + +**OpenAI (opt-in):** enable in Settings by adding an API key. Higher retrieval quality, at a per-request cost; sends text to `api.openai.com`. + +- `text-embedding-3-small` (1536 dimensions) — faster, cheaper - `text-embedding-3-large` (3072 dimensions) — higher quality, more expensive -Switching models in Settings re-indexes everything (vectors from different models are incompatible). +Switching provider or model in Settings re-indexes everything (vectors from different models are incompatible). ### Privacy & Security +- **No outbound network calls by default** — the local embedding model runs entirely on your machine, so with the default provider nothing ever leaves your computer. - All data stays local in `~/.agent-memory/agent-memory.db` -- The only outbound network call is to `api.openai.com` for embeddings +- The only outbound network call is to `api.openai.com` for embeddings — and only if you opt into the OpenAI provider in Settings. - The MCP HTTP server binds to `127.0.0.1` only (never exposed to the network) - The stdio transport (Claude Desktop) uses process-level isolation — no network at all @@ -48,9 +55,10 @@ Switching models in Settings re-indexes everything (vectors from different model - Go 1.24+ - Wails v2 CLI (`go install github.com/wailsapp/wails/v2/cmd/wails@latest`) - Node.js 18+ and npm (for the frontend build) -- OpenAI API key +- Network access on the first build — `make build` runs `make assets`, which downloads the pinned local model, tokenizer, and ONNX Runtime library (~150 MB, checksummed against `assets/manifest.json`) into the gitignored `assets/embedded/`. Cached after the first run. +- OpenAI API key — optional, only needed if you opt into the OpenAI provider at runtime - macOS, Linux, or Windows -- CGo enabled (required for SQLite + sqlite-vec; default on macOS/Linux, may need `CC=gcc` on Windows) +- CGo enabled (required for SQLite + sqlite-vec, and for the local embedder's native libraries; default on macOS/Linux, may need `CC=gcc` on Windows) ### Local Development @@ -74,6 +82,8 @@ make test # run all Go tests make clean # remove build artifacts ``` +`make build` depends on `make assets` (downloads the local model/tokenizer/ORT libraries, ~150 MB, on the first run) and builds with the `localembed` build tag plus the `CGO_LDFLAGS` needed to link the native libraries. The resulting binary is ~180 MB (it bundles the embedding model). `make test` runs the default, native-lib-free build — the `localembed`-tagged code is isolated so plain `go test ./...` needs no assets. + Or manually: ```bash diff --git a/app.go b/app.go index fd0d8b9..80d2163 100644 --- a/app.go +++ b/app.go @@ -10,9 +10,13 @@ import ( "log" "os" "path/filepath" + goruntime "runtime" + "strconv" + "github.com/borzou/vecstore/internal/chunker" "github.com/borzou/vecstore/internal/domain" "github.com/borzou/vecstore/internal/embeddings" + "github.com/borzou/vecstore/internal/embeddings/local" "github.com/borzou/vecstore/internal/engine" "github.com/borzou/vecstore/internal/mcp" "github.com/borzou/vecstore/internal/store" @@ -23,9 +27,59 @@ import ( // appInstance is a package-level reference so CGo tray callbacks can reach the app. var appInstance *App -// EmbedderFactory creates an Embedder for the given API key and model. -// Injected by main.go so app.go doesn't depend on concrete embedder packages. -type EmbedderFactory func(apiKey, model string) embeddings.Embedder +// EmbedderFactory creates an Embedder for the given provider, API key and model. +// Injected by main.go so app.go doesn't own provider-to-implementation wiring. +type EmbedderFactory func(provider, apiKey, model string) (embeddings.Embedder, error) + +// resolveProvider delegates to the single shared policy in the embeddings +// package (an explicit provider wins; else an OpenAI key implies openai; else +// the default) — the engine and read-only search resolve through the same +// function, so the rule cannot drift between processes. +func resolveProvider(configuredProvider, apiKey string) string { + return embeddings.ResolveProvider(configuredProvider, apiKey) +} + +// resolveModel returns the stored model if set, else the provider's default. +func resolveModel(provider, storedModel string) string { + if storedModel != "" { + return storedModel + } + return embeddings.DefaultModel(provider) +} + +// buildChunker constructs a chunker matched to the embedding provider, reusing +// the stored chunk_size/chunk_overlap config. For the local provider it injects +// the model's own tokenizer and clamps chunk size to the model's context window +// so chunks are measured in the embedding model's tokens. +func buildChunker(cfg *store.SQLiteStore, provider string, embedder embeddings.Embedder) (chunker.Chunker, error) { + var opts []chunker.Option + if sizeStr, _ := cfg.GetConfig("chunk_size"); sizeStr != "" { + if size, err := strconv.Atoi(sizeStr); err == nil { + opts = append(opts, chunker.WithChunkSize(size)) + } + } + if overlapStr, _ := cfg.GetConfig("chunk_overlap"); overlapStr != "" { + if overlap, err := strconv.Atoi(overlapStr); err == nil { + opts = append(opts, chunker.WithOverlap(overlap)) + } + } + if provider == embeddings.ProviderLocal { + tok, err := local.NewChunkerTokenizer(local.Config{}) + if err != nil { + return nil, fmt.Errorf("local chunker tokenizer: %w", err) + } + // Budget below the model's hard limit: the embedder prepends the E5 + // prefix and the tokenizer adds special tokens AFTER chunking, so a + // chunk cut exactly at MaxInputTokens overflows the model (512+prefix + // → the "512 by 516" ORT failure that silently dropped every + // multi-chunk file). Epic: effective chunk size ≈ 480. + opts = append(opts, + chunker.WithTokenizer(tok), + chunker.WithMaxInputTokens(embedder.MaxInputTokens()-local.EmbedTokenReserve), + ) + } + return chunker.New(opts...) +} // App exposes methods to the Wails frontend. type App struct { @@ -63,6 +117,12 @@ func (a *App) shutdown(ctx context.Context) { if err := a.mcpServer.Stop(); err != nil { log.Printf("mcp server stop: %v", err) } + // Stop before Close: cancel in-flight indexing and wait for it to reach a + // file boundary so the store never shuts down under an active write. With + // close=quit on Windows/Linux, quitting mid-scan is a routine event. + if err := a.engine.Stop(); err != nil { + log.Printf("engine stop: %v", err) + } if err := a.engine.Close(); err != nil { log.Printf("engine close: %v", err) } @@ -153,36 +213,84 @@ func (a *App) GetConfig(key string) string { return val } +// currentProvider resolves the active embedding provider from config, applying +// the same policy as the composition root. +func (a *App) currentProvider() string { + provider, _ := a.store.GetConfig("embedding_provider") + apiKey, _ := a.store.GetConfig("openai_api_key") + return resolveProvider(provider, apiKey) +} + // SetConfig writes a configuration value to the store. -// If the embedding model changes, the embedder is swapped and the index is -// reset — the virtual table must be recreated with the new dimension. -// If the API key changes, the embedder is swapped so new requests use it. +// - embedding_provider: persist, rebuild the embedder AND the matched chunker, +// then reset the index (the vector table is recreated at the new dimension). +// - embedding_model: only meaningful for the active provider; rebuild the +// embedder and reset (dimension may change, e.g. OpenAI small↔large). +// - openai_api_key: only re-swap the embedder when openai is the active provider. func (a *App) SetConfig(key, value string) error { switch key { + case "embedding_provider": + oldProvider := a.currentProvider() + if oldProvider == value { + return a.store.SetConfig(key, value) + } + log.Printf("[SetConfig] embedding provider changed %s → %s — swapping embedder+chunker and resetting index", oldProvider, value) + if err := a.store.SetConfig(key, value); err != nil { + return err + } + // Provider switch uses the provider's default model; the local model is + // fixed, and OpenAI users can adjust embedding_model afterward. + model := embeddings.DefaultModel(value) + apiKey, _ := a.store.GetConfig("openai_api_key") + embedder, err := a.newEmbedder(value, apiKey, model) + if err != nil { + return fmt.Errorf("build embedder for provider %s: %w", value, err) + } + newChunker, err := buildChunker(a.store, value, embedder) + if err != nil { + return fmt.Errorf("build chunker for provider %s: %w", value, err) + } + if err := a.store.SetConfig("embedding_model", model); err != nil { + return err + } + a.engine.SetEmbedder(embedder) + a.engine.SetChunker(newChunker) + return a.engine.Reset() + case "embedding_model": + provider := a.currentProvider() oldModel, _ := a.store.GetConfig("embedding_model") if oldModel == "" { - oldModel = "text-embedding-3-small" + oldModel = embeddings.DefaultModel(provider) } - if oldModel != value { - log.Printf("[SetConfig] embedding model changed %s → %s — resetting index", oldModel, value) - if err := a.store.SetConfig(key, value); err != nil { - return err - } - apiKey, _ := a.store.GetConfig("openai_api_key") - a.engine.SetEmbedder(a.newEmbedder(apiKey, value)) - return a.engine.Reset() + if oldModel == value { + return nil } + log.Printf("[SetConfig] embedding model changed %s → %s — resetting index", oldModel, value) + if err := a.store.SetConfig(key, value); err != nil { + return err + } + apiKey, _ := a.store.GetConfig("openai_api_key") + embedder, err := a.newEmbedder(provider, apiKey, value) + if err != nil { + return fmt.Errorf("build embedder: %w", err) + } + a.engine.SetEmbedder(embedder) + return a.engine.Reset() case "openai_api_key": if err := a.store.SetConfig(key, value); err != nil { return err } - model, _ := a.store.GetConfig("embedding_model") - if model == "" { - model = "text-embedding-3-small" + // A key change only affects requests when openai is the active provider. + if a.currentProvider() == embeddings.ProviderOpenAI { + model := resolveModel(embeddings.ProviderOpenAI, a.GetConfig("embedding_model")) + embedder, err := a.newEmbedder(embeddings.ProviderOpenAI, value, model) + if err != nil { + return fmt.Errorf("build embedder: %w", err) + } + a.engine.SetEmbedder(embedder) } - a.engine.SetEmbedder(a.newEmbedder(value, model)) return nil } @@ -282,10 +390,28 @@ func generateToken() (string, error) { return hex.EncodeToString(b), nil } -// defaultClaudeDesktopConfigPath returns the standard Claude Desktop config location. +// defaultClaudeDesktopConfigPath returns the standard Claude Desktop config +// location for the current OS. Claude Desktop reads a different path per +// platform; writing the macOS path on Windows produces a config it never sees +// (while still reporting success — the Phase 5 Windows smoke caught exactly that). func defaultClaudeDesktopConfigPath() string { - home, _ := os.UserHomeDir() - return filepath.Join(home, "Library", "Application Support", "Claude", "claude_desktop_config.json") + switch goruntime.GOOS { + case "windows": + if appData := os.Getenv("APPDATA"); appData != "" { + return filepath.Join(appData, "Claude", "claude_desktop_config.json") + } + home, _ := os.UserHomeDir() + return filepath.Join(home, "AppData", "Roaming", "Claude", "claude_desktop_config.json") + case "darwin": + home, _ := os.UserHomeDir() + return filepath.Join(home, "Library", "Application Support", "Claude", "claude_desktop_config.json") + default: // linux + if xdg := os.Getenv("XDG_CONFIG_HOME"); xdg != "" { + return filepath.Join(xdg, "Claude", "claude_desktop_config.json") + } + home, _ := os.UserHomeDir() + return filepath.Join(home, ".config", "Claude", "claude_desktop_config.json") + } } // GetClaudeDesktopConfigPath returns the current config path (custom or default). diff --git a/assets/manifest.json b/assets/manifest.json new file mode 100644 index 0000000..50532cf --- /dev/null +++ b/assets/manifest.json @@ -0,0 +1,136 @@ +{ + "schema": 1, + "comment": "Pinned local-embedding artifacts (URLs + SHA-256). Weights/libs are NEVER committed; `make assets` downloads and checksum-verifies these into internal/embeddings/local/embedded/ (go:embed runtime assets) and internal/embeddings/local/lib/ (link-time static lib). For tarball artifacts, `sha256` is the checksum of the downloaded archive and `member_sha256` is the checksum of the extracted member. `dest` is relative to internal/embeddings/local/. Model + tokenizer entries are identical bytes on every platform (per-platform keyed for a uniform `make assets`). darwin-amd64's onnxruntime is our own source build (official mac-Intel prebuilts stopped at ORT 1.23) hosted in this repo's GitHub releases — see the ort-1.26.0-darwin-x64 release notes for the reproducible recipe. windows-amd64's tokenizers_static has no upstream prebuilt: it is built from Rust source per Documentation/windows-build.md and hosted in this repo's releases (tokenizers-1.27.0-windows-x64).", + "platforms": { + "darwin-arm64": { + "model": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/onnx/model_quantized.onnx", + "sha256": "f80102d3f2a1229f387d3c81909990d8945513e347b0eab049f7de3c6f98c193", + "dest": "embedded/model_quantized.onnx" + }, + "tokenizer": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/tokenizer.json", + "sha256": "0b44a9d7b51c3c62626640cda0e2c2f70fdacdc25bbbd68038369d14ebdf4c39", + "dest": "embedded/tokenizer.json" + }, + "onnxruntime": { + "url": "https://github.com/microsoft/onnxruntime/releases/download/v1.26.0/onnxruntime-osx-arm64-1.26.0.tgz", + "sha256": "7a1280bbb1701ea514f71828765237e7896e0f2e1cd332f1f70dbd5c3e33aca3", + "member": "onnxruntime-osx-arm64-1.26.0/lib/libonnxruntime.1.26.0.dylib", + "member_sha256": "30afadcfc3c704f7671f8430d6252956651c1972373901d2be629da2e6a4d8ee", + "dest": "embedded/libonnxruntime.1.26.0.dylib" + }, + "tokenizers_static": { + "url": "https://github.com/daulet/tokenizers/releases/download/v1.27.0/libtokenizers.darwin-arm64.tar.gz", + "sha256": "fb84b8b2e349a5952767ffe80ccd862fc44084de47f3b0cc3f0b7c9d4e649cf7", + "member": "libtokenizers.a", + "member_sha256": "c91ae814afb8fe4f000099972208f60c3aa2d13899a7cd5a31fee2e9e2efbac7", + "dest": "lib/libtokenizers.a" + } + }, + "windows-amd64": { + "model": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/onnx/model_quantized.onnx", + "sha256": "f80102d3f2a1229f387d3c81909990d8945513e347b0eab049f7de3c6f98c193", + "dest": "embedded/model_quantized.onnx" + }, + "tokenizer": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/tokenizer.json", + "sha256": "0b44a9d7b51c3c62626640cda0e2c2f70fdacdc25bbbd68038369d14ebdf4c39", + "dest": "embedded/tokenizer.json" + }, + "onnxruntime": { + "url": "https://github.com/microsoft/onnxruntime/releases/download/v1.26.0/onnxruntime-win-x64-1.26.0.zip", + "sha256": "6ebe99b5564bf4d029b6e93eac9ff423682b6212eade769e9ca3f685eaf500b4", + "member": "onnxruntime-win-x64-1.26.0/lib/onnxruntime.dll", + "member_sha256": "b2ba7ca16e0e4fe71ad5148744ab885a2f5809e52a0c3de4d9ba3853a03977f9", + "dest": "embedded/onnxruntime.dll" + }, + "tokenizers_static": { + "url": "https://github.com/unitedeffectslabs/agent-memory/releases/download/tokenizers-1.27.0-windows-x64/libtokenizers.windows-x86_64.tar.gz", + "sha256": "882f520174a6cb14dcf4dca559375d65915a69a075bd85c22e40463bc0b466a8", + "member": "libtokenizers.a", + "member_sha256": "5e0815434a9d9eea40638ae9303f239c85059868673dc7e2a5a1eee72a9ad692", + "dest": "lib/libtokenizers.a" + } + }, + "darwin-amd64": { + "model": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/onnx/model_quantized.onnx", + "sha256": "f80102d3f2a1229f387d3c81909990d8945513e347b0eab049f7de3c6f98c193", + "dest": "embedded/model_quantized.onnx" + }, + "tokenizer": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/tokenizer.json", + "sha256": "0b44a9d7b51c3c62626640cda0e2c2f70fdacdc25bbbd68038369d14ebdf4c39", + "dest": "embedded/tokenizer.json" + }, + "onnxruntime": { + "url": "https://github.com/unitedeffectslabs/agent-memory/releases/download/ort-1.26.0-darwin-x64/onnxruntime-osx-x86_64-1.26.0-uelabs.tgz", + "sha256": "a78e1d8c0183833bf21f23440b0b36ae3cd4f99dbfac470249ec41768493f2b8", + "member": "libonnxruntime.1.26.0.dylib", + "member_sha256": "e29a80a09a41def826f3b31519978944a6ab5c3dc0613124a1863ff49eefce9b", + "dest": "embedded/libonnxruntime.1.26.0.dylib" + }, + "tokenizers_static": { + "url": "https://github.com/daulet/tokenizers/releases/download/v1.27.0/libtokenizers.darwin-x86_64.tar.gz", + "sha256": "6239efe5a81fde8089ef2df8ae710366542b4e5deab6d8ecb74d7d1862db2ddb", + "member": "libtokenizers.a", + "member_sha256": "c8b1f9ac68193f185a0d09ffc2e8d8892ebcfc71aad6318ffdbd6dce11d05460", + "dest": "lib/libtokenizers.a" + } + }, + "linux-arm64": { + "model": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/onnx/model_quantized.onnx", + "sha256": "f80102d3f2a1229f387d3c81909990d8945513e347b0eab049f7de3c6f98c193", + "dest": "embedded/model_quantized.onnx" + }, + "tokenizer": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/tokenizer.json", + "sha256": "0b44a9d7b51c3c62626640cda0e2c2f70fdacdc25bbbd68038369d14ebdf4c39", + "dest": "embedded/tokenizer.json" + }, + "onnxruntime": { + "url": "https://github.com/microsoft/onnxruntime/releases/download/v1.26.0/onnxruntime-linux-aarch64-1.26.0.tgz", + "sha256": "34ff1c2d0f12e2cf3d33a0c5f82e39792e1d581fbd6968fd7c30d173654be01a", + "member": "onnxruntime-linux-aarch64-1.26.0/lib/libonnxruntime.so.1.26.0", + "member_sha256": "115ecb838e703d390262b8b4d07d5248e6693c67658d4c98c48f94905ab27af4", + "dest": "embedded/libonnxruntime.so.1.26.0" + }, + "tokenizers_static": { + "url": "https://github.com/daulet/tokenizers/releases/download/v1.27.0/libtokenizers.linux-arm64.tar.gz", + "sha256": "e96545ad05930c26f51f63d932ee6d3bbd32bbed149e102c5290d587a2293067", + "member": "libtokenizers.a", + "member_sha256": "aa4d3a7d28023e439fa7a645fb606195acafb5c33c8cd01928ed1aacfc4c074e", + "dest": "lib/libtokenizers.a" + } + }, + "linux-amd64": { + "model": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/onnx/model_quantized.onnx", + "sha256": "f80102d3f2a1229f387d3c81909990d8945513e347b0eab049f7de3c6f98c193", + "dest": "embedded/model_quantized.onnx" + }, + "tokenizer": { + "url": "https://huggingface.co/Xenova/multilingual-e5-small/resolve/main/tokenizer.json", + "sha256": "0b44a9d7b51c3c62626640cda0e2c2f70fdacdc25bbbd68038369d14ebdf4c39", + "dest": "embedded/tokenizer.json" + }, + "onnxruntime": { + "url": "https://github.com/microsoft/onnxruntime/releases/download/v1.26.0/onnxruntime-linux-x64-1.26.0.tgz", + "sha256": "1254da24fb389cf39dc0ff3451ab48301740ffbfcbaf646849df92f80ee92c57", + "member": "onnxruntime-linux-x64-1.26.0/lib/libonnxruntime.so.1.26.0", + "member_sha256": "5bd5bedf736fc501692435d0ec4f6e8b2bdf48cd30af8e6d00d61b3ddc9a7ab8", + "dest": "embedded/libonnxruntime.so.1.26.0" + }, + "tokenizers_static": { + "url": "https://github.com/daulet/tokenizers/releases/download/v1.27.0/libtokenizers.linux-amd64.tar.gz", + "sha256": "72556cdca798dd4ea7cdaba308e5f0d68a8cb93b67c96edf485b7a0edd7b07f4", + "member": "libtokenizers.a", + "member_sha256": "e6862b31745bb7d07980fcee70e49cd3b4318097609180f5d2d3fb394f305d50", + "dest": "lib/libtokenizers.a" + } + } + } +} diff --git a/frontend/src/App.jsx b/frontend/src/App.jsx index c39b3fa..cd9433e 100644 --- a/frontend/src/App.jsx +++ b/frontend/src/App.jsx @@ -22,8 +22,8 @@ export default function App() { async function checkOnboarding() { try { - const apiKey = await window.go.main.App.GetConfig('openai_api_key') - setNeedsOnboarding(!apiKey || apiKey.trim() === '') + const complete = await window.go.main.App.GetConfig('onboarding_complete') + setNeedsOnboarding(complete !== 'true') } catch { setNeedsOnboarding(true) } diff --git a/frontend/src/pages/Dashboard.jsx b/frontend/src/pages/Dashboard.jsx index 5a6a829..1c0405f 100644 --- a/frontend/src/pages/Dashboard.jsx +++ b/frontend/src/pages/Dashboard.jsx @@ -84,7 +84,11 @@ export default function Dashboard() { <div style={st.grid}> <StatCard label="Total Files" value={stats?.TotalFiles ?? '-'} /> <StatCard label="Total Chunks" value={stats?.TotalChunks ?? '-'} /> - <StatCard label="Embedding Model" value={stats?.EmbeddingModel || '-'} small /> + <StatCard + label="Embedding Model" + value={stats?.EmbeddingModel ? `${stats.Provider || '-'} · ${stats.EmbeddingModel}` : '-'} + small + /> <StatCard label="Last Indexed" value={formatTime(stats?.LastIndexedAt)} small /> </div> diff --git a/frontend/src/pages/Onboarding.jsx b/frontend/src/pages/Onboarding.jsx index e0aa40d..ca24233 100644 --- a/frontend/src/pages/Onboarding.jsx +++ b/frontend/src/pages/Onboarding.jsx @@ -1,21 +1,35 @@ import React, { useState } from 'react' +const OPENAI_MODELS = ['text-embedding-3-small', 'text-embedding-3-large'] + export default function Onboarding({ onComplete }) { const [step, setStep] = useState(0) + const [useOpenAI, setUseOpenAI] = useState(false) const [apiKey, setApiKey] = useState('') + const [openaiModel, setOpenaiModel] = useState('text-embedding-3-small') const [dirs, setDirs] = useState([]) const [error, setError] = useState('') const [saving, setSaving] = useState(false) - async function handleSaveApiKey() { - if (!apiKey.trim()) { setError('API key is required.'); return } - setSaving(true); setError('') - try { - await window.go.main.App.SetConfig('openai_api_key', apiKey.trim()) - setStep(2) - } catch (e) { - setError('Failed to save: ' + (e?.message || String(e))) - } finally { setSaving(false) } + // Welcome → directories. If the user opted into OpenAI, persist that provider + // choice here; otherwise the backend keeps the default (local) provider. + async function handleStartFromWelcome() { + setError('') + if (useOpenAI) { + if (!apiKey.trim()) { setError('API key is required to use OpenAI.'); return } + setSaving(true) + try { + await window.go.main.App.SetConfig('embedding_provider', 'openai') + await window.go.main.App.SetConfig('openai_api_key', apiKey.trim()) + await window.go.main.App.SetConfig('embedding_model', openaiModel) + } catch (e) { + setError('Failed to save OpenAI settings: ' + (e?.message || String(e))) + setSaving(false) + return + } + setSaving(false) + } + setStep(1) } async function handleChooseFolder() { @@ -41,17 +55,24 @@ export default function Onboarding({ onComplete }) { for (const path of dirs) { await window.go.main.App.RegisterDirectory(path) } - setStep(3) + setStep(2) } catch (e) { setError('Failed to save directories: ' + (e?.message || String(e))) } } - function handleFinish() { + async function handleFinish() { + setError('') + try { + await window.go.main.App.SetConfig('onboarding_complete', 'true') + } catch (e) { + setError('Failed to finish setup: ' + (e?.message || String(e))) + return + } onComplete() } - const totalSteps = 4 + const totalSteps = 3 return ( <div style={s.container}> @@ -67,42 +88,66 @@ export default function Onboarding({ onComplete }) { <> <h2 style={s.heading}>Welcome to Agent Memory</h2> <p style={s.text}> - Agent Memory watches your directories, creates embeddings with OpenAI, - and stores vectors locally. It provides semantic search via MCP for - Claude and other AI assistants. + Agent Memory watches your directories, creates embeddings, and stores + vectors locally. It provides semantic search via MCP for Claude and + other AI assistants. </p> - <div style={s.btnRow}> - <button style={s.btnPrimary} onClick={() => setStep(1)}>Get Started</button> - </div> - </> - )} - - {step === 1 && ( - <> - <h2 style={s.heading}>OpenAI API Key</h2> <p style={s.text}> - Enter your OpenAI API key to enable embeddings. - It's stored locally in the database — never sent anywhere except OpenAI. + By default it runs a local, on-device embedding model — private, no + API key, and fully offline. You can switch to OpenAI anytime in Settings. </p> - <input - type="password" - style={s.input} - placeholder="sk-..." - value={apiKey} - onChange={(e) => setApiKey(e.target.value)} - onKeyDown={(e) => e.key === 'Enter' && handleSaveApiKey()} - /> - {error && <div style={s.error}>{error}</div>} - <div style={s.btnRow}> - <button style={s.btn} onClick={() => setStep(0)}>Back</button> - <button style={s.btnPrimary} onClick={handleSaveApiKey} disabled={saving}> - {saving ? 'Saving...' : 'Continue'} + + {/* Optional: opt into OpenAI instead of the local default */} + {!useOpenAI ? ( + <button + style={s.linkBtn} + onClick={() => { setError(''); setUseOpenAI(true) }} + > + Use OpenAI instead + </button> + ) : ( + <div style={s.openaiBox}> + <div style={s.openaiHeader}> + <span style={s.openaiTitle}>Use OpenAI</span> + <button + style={s.linkBtn} + onClick={() => { setError(''); setUseOpenAI(false) }} + > + Use local model instead + </button> + </div> + <p style={{ ...s.text, marginBottom: 12 }}> + Higher quality, but requires an API key and sends text to OpenAI. + The key is stored locally in the database. + </p> + <input + type="password" + style={s.input} + placeholder="sk-..." + value={apiKey} + onChange={(e) => setApiKey(e.target.value)} + /> + <select + style={s.select} + value={openaiModel} + onChange={(e) => setOpenaiModel(e.target.value)} + > + {OPENAI_MODELS.map((m) => <option key={m} value={m}>{m}</option>)} + </select> + </div> + )} + + {error && <div style={{ ...s.error, marginTop: 12 }}>{error}</div>} + + <div style={{ ...s.btnRow, marginTop: 20 }}> + <button style={s.btnPrimary} onClick={handleStartFromWelcome} disabled={saving}> + {saving ? 'Saving...' : 'Get Started'} </button> </div> </> )} - {step === 2 && ( + {step === 1 && ( <> <h2 style={s.heading}>Add Directories</h2> <p style={s.text}> @@ -137,7 +182,7 @@ export default function Onboarding({ onComplete }) { {error && <div style={{ ...s.error, marginTop: 12 }}>{error}</div>} <div style={{ ...s.btnRow, marginTop: 20 }}> - <button style={s.btn} onClick={() => setStep(1)}>Back</button> + <button style={s.btn} onClick={() => setStep(0)}>Back</button> <button style={{ ...s.btnPrimary, ...(dirs.length === 0 ? s.btnDisabled : {}) }} onClick={handleSaveDirsAndContinue} @@ -149,7 +194,7 @@ export default function Onboarding({ onComplete }) { </> )} - {step === 3 && ( + {step === 2 && ( <> <h2 style={s.heading}>All Set</h2> <p style={s.text}> @@ -221,8 +266,45 @@ const s = { borderRadius: 8, color: '#e5e5e7', fontSize: 14, - marginBottom: 16, + marginBottom: 12, + outline: 'none', + }, + select: { + width: '100%', + padding: '10px 12px', + background: 'rgba(255,255,255,0.06)', + border: '1px solid rgba(255,255,255,0.1)', + borderRadius: 8, + color: '#e5e5e7', + fontSize: 14, outline: 'none', + cursor: 'pointer', + }, + linkBtn: { + background: 'none', + border: 'none', + color: '#0a84ff', + fontSize: 13, + cursor: 'pointer', + padding: 0, + fontWeight: 500, + }, + openaiBox: { + background: 'rgba(255,255,255,0.03)', + border: '1px solid rgba(255,255,255,0.08)', + borderRadius: 10, + padding: 16, + }, + openaiHeader: { + display: 'flex', + alignItems: 'center', + justifyContent: 'space-between', + marginBottom: 10, + }, + openaiTitle: { + fontSize: 14, + fontWeight: 600, + color: '#fff', }, dirList: { display: 'flex', diff --git a/frontend/src/pages/Settings.jsx b/frontend/src/pages/Settings.jsx index 44a6bda..8646105 100644 --- a/frontend/src/pages/Settings.jsx +++ b/frontend/src/pages/Settings.jsx @@ -3,6 +3,7 @@ import React, { useState, useEffect } from 'react' const MODELS = ['text-embedding-3-small', 'text-embedding-3-large'] export default function Settings() { + const [provider, setProvider] = useState('local') const [apiKey, setApiKey] = useState('') const [model, setModel] = useState('') const [chunkSize, setChunkSize] = useState('') @@ -14,12 +15,14 @@ export default function Settings() { const [error, setError] = useState('') const [showDeleteConfirm, setShowDeleteConfirm] = useState(false) const [pendingModel, setPendingModel] = useState(null) // model change awaiting confirmation + const [pendingProvider, setPendingProvider] = useState(null) // provider change awaiting confirmation useEffect(() => { loadSettings() }, []) async function loadSettings() { try { - const [key, m, cs, co, p, token] = await Promise.all([ + const [prov, key, m, cs, co, p, token] = await Promise.all([ + window.go.main.App.GetConfig('embedding_provider'), window.go.main.App.GetConfig('openai_api_key'), window.go.main.App.GetConfig('embedding_model'), window.go.main.App.GetConfig('chunk_size'), @@ -27,6 +30,7 @@ export default function Settings() { window.go.main.App.GetConfig('mcp_port'), window.go.main.App.GetConfig('auth_token'), ]) + setProvider(prov || 'local') setApiKey(key || '') setModel(m || 'text-embedding-3-small') setChunkSize(cs || '512') @@ -40,6 +44,27 @@ export default function Settings() { } } + // Provider change is destructive when there are existing embeddings (the + // backend clears + re-indexes because vectors across providers/dimensions are + // incompatible), so gate it behind a confirmation when chunks already exist. + async function requestProviderChange(next) { + if (next === provider) return + setError('') + try { + const stats = await window.go.main.App.GetStats() + if (stats.TotalChunks > 0) { + setPendingProvider(next) + return + } + } catch { /* proceed anyway */ } + applyProviderChange(next) + } + + function applyProviderChange(next) { + setProvider(next) + saveConfig('embedding_provider', next, 'provider') + } + async function saveConfig(key, value, label) { setError('') try { @@ -91,41 +116,63 @@ export default function Settings() { {error && <div style={s.error}>{error}</div>} - {/* API Configuration */} - <Section title="API Configuration"> - <Row label="OpenAI API Key"> - <input - type="password" - style={s.input} - value={apiKey} - onChange={(e) => setApiKey(e.target.value)} - placeholder="sk-..." - /> - <SaveBtn onClick={() => saveConfig('openai_api_key', apiKey, 'apiKey')} saved={saved.apiKey} /> - </Row> - <Row label="Embedding Model"> - <select - style={s.select} - value={model} - onChange={async (e) => { - const newModel = e.target.value - if (newModel === model) return - // Check if there are existing embeddings that would be wiped - try { - const stats = await window.go.main.App.GetStats() - if (stats.TotalChunks > 0) { - setPendingModel(newModel) // show confirmation dialog - return - } - } catch { /* proceed anyway */ } - setModel(newModel) - saveConfig('embedding_model', newModel, 'model') - }} - > - {MODELS.map((m) => <option key={m} value={m}>{m}</option>)} - </select> - {saved.model && <span style={s.saved}>Saved</span>} + {/* Embedding Provider */} + <Section title="Embedding Provider"> + <Row label="Provider"> + <div style={s.providerChoices}> + <ProviderOption + selected={provider === 'local'} + onClick={() => requestProviderChange('local')} + title="Local (default)" + desc="Private, on-device, no API key" + /> + <ProviderOption + selected={provider === 'openai'} + onClick={() => requestProviderChange('openai')} + title="OpenAI" + desc="Higher quality, requires a key and sends text to OpenAI" + /> + </div> + {saved.provider && <span style={s.saved}>Saved</span>} </Row> + + {provider === 'openai' && ( + <> + <Row label="OpenAI API Key"> + <input + type="password" + style={s.input} + value={apiKey} + onChange={(e) => setApiKey(e.target.value)} + placeholder="sk-..." + /> + <SaveBtn onClick={() => saveConfig('openai_api_key', apiKey, 'apiKey')} saved={saved.apiKey} /> + </Row> + <Row label="Embedding Model"> + <select + style={s.select} + value={model} + onChange={async (e) => { + const newModel = e.target.value + if (newModel === model) return + // Check if there are existing embeddings that would be wiped + try { + const stats = await window.go.main.App.GetStats() + if (stats.TotalChunks > 0) { + setPendingModel(newModel) // show confirmation dialog + return + } + } catch { /* proceed anyway */ } + setModel(newModel) + saveConfig('embedding_model', newModel, 'model') + }} + > + {MODELS.map((m) => <option key={m} value={m}>{m}</option>)} + </select> + {saved.model && <span style={s.saved}>Saved</span>} + </Row> + </> + )} </Section> {/* Chunking */} @@ -177,7 +224,9 @@ export default function Settings() { <span style={s.readOnly}>127.0.0.1:{port || '9847'}</span> </Row> <Row label="Outbound"> - <span style={s.readOnly}>api.openai.com only</span> + <span style={s.readOnly}> + {provider === 'openai' ? 'api.openai.com' : 'none (fully offline)'} + </span> </Row> </Section> @@ -228,6 +277,33 @@ export default function Settings() { </div> )} + {/* Provider change confirmation dialog */} + {pendingProvider && ( + <div style={s.overlay} onClick={() => setPendingProvider(null)}> + <div style={s.dialog} onClick={(e) => e.stopPropagation()}> + <div style={s.dialogTitle}>Change Embedding Provider?</div> + <div style={s.dialogText}> + Switching from <strong>{provider}</strong> to <strong>{pendingProvider}</strong> will + reset the index and re-embed all files. Existing embeddings will be deleted because + vectors from different providers are incompatible. + </div> + <div style={s.dialogButtons}> + <button style={s.btn} onClick={() => setPendingProvider(null)}>Cancel</button> + <button + style={{ ...s.btn, background: '#0a84ff', color: '#fff', borderColor: 'transparent' }} + onClick={() => { + const next = pendingProvider + setPendingProvider(null) + applyProviderChange(next) + }} + > + Change & Reset Index + </button> + </div> + </div> + </div> + )} + {/* Danger Zone */} <div style={s.danger}> <div style={s.dangerTitle}>Danger Zone</div> @@ -288,6 +364,23 @@ function Row({ label, children }) { ) } +function ProviderOption({ selected, onClick, title, desc }) { + return ( + <button + onClick={onClick} + style={{ ...s.providerOption, ...(selected ? s.providerOptionActive : {}) }} + > + <span style={s.providerRadio}> + <span style={{ ...s.providerRadioDot, ...(selected ? s.providerRadioDotActive : {}) }} /> + </span> + <span style={s.providerOptionText}> + <span style={s.providerOptionTitle}>{title}</span> + <span style={s.providerOptionDesc}>{desc}</span> + </span> + </button> + ) +} + function SaveBtn({ onClick, saved }) { return ( <> @@ -398,6 +491,62 @@ const s = { fontSize: 13, color: '#0a84ff', }, + providerChoices: { + display: 'flex', + flexDirection: 'column', + gap: 8, + flex: 1, + }, + providerOption: { + display: 'flex', + alignItems: 'flex-start', + gap: 10, + padding: '10px 12px', + background: 'rgba(255,255,255,0.03)', + border: '1px solid rgba(255,255,255,0.1)', + borderRadius: 8, + cursor: 'pointer', + textAlign: 'left', + transition: 'all 0.15s', + }, + providerOptionActive: { + background: 'rgba(10,132,255,0.12)', + borderColor: '#0a84ff', + }, + providerRadio: { + width: 16, + height: 16, + borderRadius: '50%', + border: '1.5px solid rgba(255,255,255,0.3)', + display: 'flex', + alignItems: 'center', + justifyContent: 'center', + flexShrink: 0, + marginTop: 1, + }, + providerRadioDot: { + width: 8, + height: 8, + borderRadius: '50%', + background: 'transparent', + }, + providerRadioDotActive: { + background: '#0a84ff', + }, + providerOptionText: { + display: 'flex', + flexDirection: 'column', + gap: 2, + }, + providerOptionTitle: { + fontSize: 13, + fontWeight: 600, + color: '#e5e5e7', + }, + providerOptionDesc: { + fontSize: 12, + color: 'rgba(255,255,255,0.45)', + }, saved: { fontSize: 12, color: '#34c759', diff --git a/go.mod b/go.mod index b859cb9..acdf8ee 100644 --- a/go.mod +++ b/go.mod @@ -5,12 +5,14 @@ go 1.24.0 require ( github.com/asg017/sqlite-vec-go-bindings v0.1.6 github.com/bmatcuk/doublestar/v4 v4.10.0 + github.com/daulet/tokenizers v1.27.0 github.com/fsnotify/fsnotify v1.9.0 github.com/mattn/go-sqlite3 v1.14.22 github.com/nguyenthenguyen/docx v0.0.0-20230621112118-9c8e795a11db github.com/tiktoken-go/tokenizer v0.2.0 github.com/wailsapp/wails/v2 v2.11.0 github.com/xuri/excelize/v2 v2.10.1 + github.com/yalue/onnxruntime_go v1.31.0 ) require ( diff --git a/go.sum b/go.sum index 3a9441e..3252c4b 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/bep/debounce v1.2.1 h1:v67fRdBA9UQu2NhLFXrSg0Brw7CexQekrBwDMM8bzeY= github.com/bep/debounce v1.2.1/go.mod h1:H8yggRPQKLUhUoqrJC1bO2xNya7vanpDl7xR3ISbCJ0= github.com/bmatcuk/doublestar/v4 v4.10.0 h1:zU9WiOla1YA122oLM6i4EXvGW62DvKZVxIe6TYWexEs= github.com/bmatcuk/doublestar/v4 v4.10.0/go.mod h1:xBQ8jztBU6kakFMg+8WGxn0c6z1fTSPVIjEY1Wr7jzc= +github.com/daulet/tokenizers v1.27.0 h1:MmFYAEDFz69s/nNQfHg59DWqHz3v94m99kEZ/JbL+s4= +github.com/daulet/tokenizers v1.27.0/go.mod h1:YjFY1o1HGMyWkQgbXJDghhvke/yFDp2vGdIO2hYs4MQ= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI= @@ -85,6 +87,8 @@ github.com/xuri/excelize/v2 v2.10.1 h1:V62UlqopMqha3kOpnlHy2CcRVw1V8E63jFoWUmMzx github.com/xuri/excelize/v2 v2.10.1/go.mod h1:iG5tARpgaEeIhTqt3/fgXCGoBRt4hNXgCp3tfXKoOIc= github.com/xuri/nfp v0.0.2-0.20250530014748-2ddeb826f9a9 h1:+C0TIdyyYmzadGaL/HBLbf3WdLgC29pgyhTjAT/0nuE= github.com/xuri/nfp v0.0.2-0.20250530014748-2ddeb826f9a9/go.mod h1:WwHg+CVyzlv/TX9xqBFXEZAuxOPxn2k1GNHwG41IIUQ= +github.com/yalue/onnxruntime_go v1.31.0 h1:1ln4YW1SFOFfGJZXe3jNOb2JUSt+l2pEneZfV8HdtFA= +github.com/yalue/onnxruntime_go v1.31.0/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4= golang.org/x/crypto v0.48.0 h1:/VRzVqiRSggnhY7gNRxPauEQ5Drw9haKdM0jqfcCFts= golang.org/x/crypto v0.48.0/go.mod h1:r0kV5h3qnFPlQnBSrULhlsRfryS2pmewsg+XfMgkVos= golang.org/x/image v0.25.0 h1:Y6uW6rH1y5y/LK1J8BPWZtr6yZ7hrsy6hFrXjgsc2fQ= diff --git a/internal/chunker/chunker.go b/internal/chunker/chunker.go index be139dc..d5b37f6 100644 --- a/internal/chunker/chunker.go +++ b/internal/chunker/chunker.go @@ -12,11 +12,38 @@ const DefaultChunkSize = 512 // DefaultChunkOverlap is the default number of overlapping tokens between consecutive chunks. const DefaultChunkOverlap = 50 +// tiktokenAdapter wraps a tokenizer.Codec to satisfy the Tokenizer interface, +// converting between the codec's []uint tokens and the interface's []uint32. +type tiktokenAdapter struct { + codec tokenizer.Codec +} + +func (a tiktokenAdapter) Encode(text string) ([]uint32, error) { + ids, _, err := a.codec.Encode(text) + if err != nil { + return nil, err + } + out := make([]uint32, len(ids)) + for i, id := range ids { + out[i] = uint32(id) + } + return out, nil +} + +func (a tiktokenAdapter) Decode(tokens []uint32) (string, error) { + ids := make([]uint, len(tokens)) + for i, t := range tokens { + ids[i] = uint(t) + } + return a.codec.Decode(ids) +} + // TokenChunker implements the Chunker interface using tiktoken-based token counting. type TokenChunker struct { chunkSize int overlap int - codec tokenizer.Codec + tokenizer Tokenizer + maxTokens int } // Option configures a TokenChunker. @@ -40,6 +67,25 @@ func WithOverlap(overlap int) Option { } } +// WithTokenizer sets a custom tokenizer, overriding the default tiktoken codec. +func WithTokenizer(t Tokenizer) Option { + return func(c *TokenChunker) { + if t != nil { + c.tokenizer = t + } + } +} + +// WithMaxInputTokens clamps the effective chunk size so no chunk exceeds the +// tokenizer/model's maximum input length. A value of 0 disables clamping. +func WithMaxInputTokens(n int) Option { + return func(c *TokenChunker) { + if n > 0 { + c.maxTokens = n + } + } +} + // New creates a new TokenChunker with the given options. // It uses the cl100k_base encoding for token counting. func New(opts ...Option) (*TokenChunker, error) { @@ -51,7 +97,7 @@ func New(opts ...Option) (*TokenChunker, error) { c := &TokenChunker{ chunkSize: DefaultChunkSize, overlap: DefaultChunkOverlap, - codec: codec, + tokenizer: tiktokenAdapter{codec: codec}, } for _, opt := range opts { opt(c) @@ -67,7 +113,7 @@ func (c *TokenChunker) ChunkText(content string) ([]ChunkResult, error) { return nil, nil } - tokens, _, err := c.codec.Encode(content) + tokens, err := c.tokenizer.Encode(content) if err != nil { return nil, err } @@ -77,18 +123,25 @@ func (c *TokenChunker) ChunkText(content string) ([]ChunkResult, error) { return nil, nil } + // Effective chunk size: clamp to maxTokens when it is smaller than the + // configured chunk size. With the default (maxTokens==0) this is a no-op. + size := c.chunkSize + if c.maxTokens > 0 && c.maxTokens < size { + size = c.maxTokens + } + var results []ChunkResult idx := 0 start := 0 for start < totalTokens { - end := start + c.chunkSize + end := start + size if end > totalTokens { end = totalTokens } chunkTokens := tokens[start:end] - chunkText, err := c.codec.Decode(chunkTokens) + chunkText, err := c.tokenizer.Decode(chunkTokens) if err != nil { return nil, err } @@ -101,7 +154,7 @@ func (c *TokenChunker) ChunkText(content string) ([]ChunkResult, error) { idx++ - step := c.chunkSize - c.overlap + step := size - c.overlap if step < 1 { step = 1 } diff --git a/internal/chunker/chunker_test.go b/internal/chunker/chunker_test.go index c0ae5a0..0c17fa1 100644 --- a/internal/chunker/chunker_test.go +++ b/internal/chunker/chunker_test.go @@ -213,3 +213,80 @@ func TestChunkTextSimple(t *testing.T) { t.Error("expected non-empty result") } } + +// fakeTokenizer is a trivial whitespace tokenizer used to verify that a custom +// tokenizer injected via WithTokenizer is actually used. Each rune becomes a +// token, so token counts are deterministic and independent of tiktoken. +type fakeTokenizer struct { + encodeCalls int +} + +func (f *fakeTokenizer) Encode(text string) ([]uint32, error) { + f.encodeCalls++ + runes := []rune(text) + out := make([]uint32, len(runes)) + for i, r := range runes { + out[i] = uint32(r) + } + return out, nil +} + +func (f *fakeTokenizer) Decode(tokens []uint32) (string, error) { + runes := make([]rune, len(tokens)) + for i, t := range tokens { + runes[i] = rune(t) + } + return string(runes), nil +} + +func TestWithTokenizerIsHonored(t *testing.T) { + ft := &fakeTokenizer{} + c := mustNew(t, WithTokenizer(ft), WithChunkSize(3), WithOverlap(0)) + + // "abcdef" -> 6 rune-tokens; with size 3 and no overlap -> 2 chunks. + results, err := c.ChunkText("abcdef") + if err != nil { + t.Fatalf("ChunkText() error: %v", err) + } + if ft.encodeCalls == 0 { + t.Fatal("custom tokenizer Encode was never called") + } + if len(results) != 2 { + t.Fatalf("expected 2 chunks, got %d", len(results)) + } + if results[0].Content != "abc" || results[1].Content != "def" { + t.Errorf("unexpected chunk contents: %q, %q", results[0].Content, results[1].Content) + } +} + +func TestWithMaxInputTokensClampsChunkSize(t *testing.T) { + // Deterministic tokenizer so chunk boundaries are exact. + content := "abcdefghij" // 10 rune-tokens + + // Without clamping: chunkSize 8 -> a single chunk covering all 10? No, + // 8 < 10 so two chunks (8 + 2). Establish the baseline first. + base := mustNew(t, WithTokenizer(&fakeTokenizer{}), WithChunkSize(8), WithOverlap(0)) + baseResults, err := base.ChunkText(content) + if err != nil { + t.Fatalf("ChunkText() error: %v", err) + } + if len(baseResults) == 0 || baseResults[0].TokenCount != 8 { + t.Fatalf("baseline expected first chunk of 8 tokens, got %+v", baseResults) + } + + // With clamping to 4: chunks become smaller (max 4 tokens each). + clamped := mustNew(t, WithTokenizer(&fakeTokenizer{}), WithChunkSize(8), WithOverlap(0), WithMaxInputTokens(4)) + clampedResults, err := clamped.ChunkText(content) + if err != nil { + t.Fatalf("ChunkText() error: %v", err) + } + for i, r := range clampedResults { + if r.TokenCount > 4 { + t.Errorf("chunk %d has %d tokens, expected <= 4 after clamping", i, r.TokenCount) + } + } + if len(clampedResults) <= len(baseResults) { + t.Errorf("clamping should produce more (smaller) chunks: got %d clamped vs %d base", + len(clampedResults), len(baseResults)) + } +} diff --git a/internal/chunker/iface.go b/internal/chunker/iface.go index 4dd20bd..049e1f5 100644 --- a/internal/chunker/iface.go +++ b/internal/chunker/iface.go @@ -11,3 +11,10 @@ type ChunkResult struct { type Chunker interface { ChunkText(content string) ([]ChunkResult, error) } + +// Tokenizer encodes text to tokens and back. It abstracts the underlying +// tokenization backend (e.g. tiktoken) so alternate tokenizers can be injected. +type Tokenizer interface { + Encode(text string) ([]uint32, error) + Decode(tokens []uint32) (string, error) +} diff --git a/internal/domain/types.go b/internal/domain/types.go index 724c3d7..4857d0b 100644 --- a/internal/domain/types.go +++ b/internal/domain/types.go @@ -30,7 +30,7 @@ type SearchParams struct { Query string Limit int // Max results. Default: 10. Offset int // Skip first N results (pagination). Default: 0. - Threshold float32 // Max distance; results farther than this are excluded. Default: 1.5 (cosine). 0 means no threshold. + Threshold float32 // Max distance; results farther than this are excluded. 0 means "use the provider-aware default" (openai 1.5, local 0.6 — see embeddings.DefaultThreshold). } type SearchResult struct { @@ -54,6 +54,7 @@ type IndexStats struct { TotalChunks int LastIndexedAt time.Time IsIndexing bool + Provider string EmbeddingModel string // Progress tracking during indexing IndexedFiles int // files processed so far in current run diff --git a/internal/embeddings/defaults.go b/internal/embeddings/defaults.go new file mode 100644 index 0000000..763f917 --- /dev/null +++ b/internal/embeddings/defaults.go @@ -0,0 +1,102 @@ +package embeddings + +import "fmt" + +// Provider identifiers for the supported embedding backends. +const ( + ProviderLocal = "local" + ProviderOpenAI = "openai" +) + +// ResolveProvider applies the single provider-resolution policy: an explicit +// configured provider wins; otherwise an existing OpenAI key implies the +// openai provider (preserving pre-provider-config users); otherwise the +// default. Every consumer of a stored provider string (composition root, +// engine, read-only search) must resolve through here — divergent copies of +// this rule are how a legacy DB gets a 'local' threshold and fingerprint +// applied to an OpenAI index. +func ResolveProvider(configuredProvider, apiKey string) string { + if configuredProvider != "" { + return configuredProvider + } + if apiKey != "" { + return ProviderOpenAI + } + return DefaultProvider() +} + +// Fingerprint returns the canonical index-identity string +// (provider:model:dimensions) recorded when an index is built and compared +// before read-only searches. Writer and checker must both use this +// constructor — a hand-built copy that drifts makes every valid index look +// mismatched, or a real mismatch look valid. +func Fingerprint(provider string, e Embedder) string { + return fmt.Sprintf("%s:%s:%d", provider, e.ModelName(), e.Dimensions()) +} + +// Default model identifiers per provider. +const ( + defaultLocalModel = "multilingual-e5-small" + defaultOpenAIModel = "text-embedding-3-small" +) + +// Embedding dimensions per model. +const ( + localDimension = 384 + openAISmallDimension = 1536 + openAILargeDimension = 3072 +) + +// DefaultProvider returns the provider used when none is configured. The +// local, in-process model is the default so the app works offline out of the box. +func DefaultProvider() string { return ProviderLocal } + +// DefaultModel returns the default model identifier for a provider. Any +// unrecognized provider falls back to the local model (the default provider). +func DefaultModel(provider string) string { + if provider == ProviderOpenAI { + return defaultOpenAIModel + } + return defaultLocalModel +} + +// DefaultDimension returns the embedding vector dimension for a provider/model +// pair. It centralizes the dimension knowledge that would otherwise be scattered +// as magic numbers across the store and composition root. +func DefaultDimension(provider, model string) int { + switch provider { + case ProviderOpenAI: + if model == "text-embedding-3-large" { + return openAILargeDimension + } + return openAISmallDimension + default: // ProviderLocal and any unrecognized provider + return localDimension + } +} + +// DefaultThreshold returns the default maximum distance for search results for a +// given provider. Results farther than this are excluded when a caller does not +// supply an explicit threshold. +// +// METRIC NOTE (flagged in the PR #2 review round): the chunk_embeddings vec0 +// table is declared without distance_metric, so sqlite-vec returns EUCLIDEAN +// (L2) distance, not cosine. Because every vector we store is L2-normalized, +// the two are monotonically equivalent (L2 = sqrt(2·cosine_distance)), so +// ranking is identical either way — but these threshold values are therefore +// L2-scale cutoffs, not the cosine values earlier comments claimed. On the L2 +// scale the Phase 0 mE5 ranges map to: related ≈0.51–0.60, cross-lingual +// ≈0.58–0.66, unrelated ≈0.76. The local 0.6 cutoff (field-validated for +// same-language search) truncates part of the cross-lingual band; whether to +// declare distance_metric=cosine (table rebuild) or retune the L2 value +// (~0.66–0.70) is an owner decision recorded in the PR review notes. +func DefaultThreshold(provider string) float32 { + switch provider { + case ProviderOpenAI: + return 1.5 + default: + // ProviderLocal (and any unrecognized provider). L2-scale cutoff, + // see METRIC NOTE. Tunable pending real-corpus evaluation. + return 0.6 + } +} diff --git a/internal/embeddings/defaults_test.go b/internal/embeddings/defaults_test.go new file mode 100644 index 0000000..a5af178 --- /dev/null +++ b/internal/embeddings/defaults_test.go @@ -0,0 +1,20 @@ +package embeddings + +import "testing" + +func TestDefaultThreshold(t *testing.T) { + cases := []struct { + provider string + want float32 + }{ + {ProviderOpenAI, 1.5}, + {ProviderLocal, 0.6}, + {"", 0.6}, // unrecognized → local default + {"something", 0.6}, // unrecognized → local default + } + for _, c := range cases { + if got := DefaultThreshold(c.provider); got != c.want { + t.Errorf("DefaultThreshold(%q) = %v, want %v", c.provider, got, c.want) + } + } +} diff --git a/internal/embeddings/iface.go b/internal/embeddings/iface.go index 118f5da..1ea3f12 100644 --- a/internal/embeddings/iface.go +++ b/internal/embeddings/iface.go @@ -2,7 +2,9 @@ package embeddings // Embedder produces vector embeddings for text. type Embedder interface { - Embed(texts []string) ([][]float32, error) + EmbedDocuments(texts []string) ([][]float32, error) // indexing path (batched) + EmbedQuery(text string) ([]float32, error) // search path Dimensions() int ModelName() string + MaxInputTokens() int // 0 = no practical limit (OpenAI); model ctx for local } diff --git a/internal/embeddings/local/assets.go b/internal/embeddings/local/assets.go new file mode 100644 index 0000000..506324f --- /dev/null +++ b/internal/embeddings/local/assets.go @@ -0,0 +1,182 @@ +package local + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "os" + "path/filepath" +) + +// Asset file names expected inside the assets directory. +const ( + assetModelFile = "model_quantized.onnx" + assetTokenizerFile = "tokenizer.json" + + // assetsDirEnv is the environment variable naming the directory that holds + // the model, tokenizer and ONNX Runtime shared library during development + // (Phase 2). Phase 3 replaces this with go:embed + checksummed extraction. + assetsDirEnv = "AGENT_MEMORY_LOCAL_ASSETS" +) + +// dylibCandidates are the ONNX Runtime shared-library file names we look for, +// most specific first. Used only for developer-override directories +// (Config.AssetsDir / AGENT_MEMORY_LOCAL_ASSETS); the bundled path resolves the +// per-platform ortLibFile directly. The Linux release tarball ships the +// versioned name (libonnxruntime.so.1.26.0), hence both .so forms. +var dylibCandidates = []string{ + "libonnxruntime.1.26.0.dylib", + "libonnxruntime.dylib", + "libonnxruntime.so.1.26.0", + "libonnxruntime.so", + "onnxruntime.dll", +} + +// resolveAssets locates the model, tokenizer and ONNX Runtime shared library. +// +// Priority: +// 1. A developer override — cfg.AssetsDir, falling back to the +// AGENT_MEMORY_LOCAL_ASSETS environment variable — points at a directory +// that already holds the three files. +// 2. Otherwise, the assets bundled into the binary via go:embed (localembed +// build only) are extracted to ~/.agent-memory/runtime/<fingerprint>/ and +// resolved from there. In the default (untagged) build no assets are +// embedded, so extractEmbeddedAssets returns an error and resolveAssets +// fails with an actionable message. +func resolveAssets(cfg Config) (modelPath, tokPath, dylibPath string, err error) { + base := cfg.AssetsDir + if base == "" { + base = os.Getenv(assetsDirEnv) + } + if base == "" { + // No dev override: fall back to the go:embed-ed, extracted assets. + extracted, eerr := extractEmbeddedAssets() + if eerr != nil { + return "", "", "", eerr + } + base = extracted + } + return resolveFromDir(base) +} + +// resolveFromDir resolves the three asset paths from a directory, requiring the +// model and tokenizer files to be present and at least one ONNX Runtime shared +// library to be found. +func resolveFromDir(base string) (modelPath, tokPath, dylibPath string, err error) { + modelPath = filepath.Join(base, assetModelFile) + tokPath = filepath.Join(base, assetTokenizerFile) + + if err := mustExist(modelPath); err != nil { + return "", "", "", err + } + if err := mustExist(tokPath); err != nil { + return "", "", "", err + } + + dylibPath, err = findDylib(base) + if err != nil { + return "", "", "", err + } + return modelPath, tokPath, dylibPath, nil +} + +// findDylib returns the first ONNX Runtime shared library found under base. +func findDylib(base string) (string, error) { + for _, name := range dylibCandidates { + p := filepath.Join(base, name) + if _, err := os.Stat(p); err == nil { + return p, nil + } + } + return "", fmt.Errorf("local: no ONNX Runtime shared library found in %s (looked for %v)", + base, dylibCandidates) +} + +func mustExist(path string) error { + if _, err := os.Stat(path); err != nil { + return fmt.Errorf("local: required asset missing: %s: %w", path, err) + } + return nil +} + +// fingerprint returns a short, stable identifier for a provider/model/dimension +// triple, used to namespace the on-disk runtime extraction directory. +func fingerprint(provider, model string, dim int) string { + sum := sha256.Sum256([]byte(fmt.Sprintf("%s:%s:%d", provider, model, dim))) + return hex.EncodeToString(sum[:])[:16] +} + +// extractAndVerify copies srcPath into destDir atomically (write to a temp file +// then rename), verifying the SHA-256 checksum. If wantSHA is empty the checksum +// is computed and not enforced. The operation is idempotent: if the destination +// already exists with a matching checksum it is left untouched. Returns the +// final destination path. +// +// Phase 2 sources assets from a developer directory and does not require +// extraction; this helper exists for Phase 3's go:embed-backed extraction and is +// unit-tested here so the behavior is pinned early. +func extractAndVerify(srcPath, destDir, wantSHA string) (string, error) { + if err := os.MkdirAll(destDir, 0o755); err != nil { + return "", fmt.Errorf("local: create runtime dir: %w", err) + } + destPath := filepath.Join(destDir, filepath.Base(srcPath)) + + // Idempotent fast path: destination already present and (if requested) + // matches the wanted checksum. + if existing, err := sha256File(destPath); err == nil { + if wantSHA == "" || existing == wantSHA { + return destPath, nil + } + } + + srcSum, err := sha256File(srcPath) + if err != nil { + return "", fmt.Errorf("local: hash source: %w", err) + } + if wantSHA != "" && srcSum != wantSHA { + return "", fmt.Errorf("local: checksum mismatch for %s: got %s want %s", + srcPath, srcSum, wantSHA) + } + + tmp, err := os.CreateTemp(destDir, ".tmp-*") + if err != nil { + return "", fmt.Errorf("local: temp file: %w", err) + } + tmpName := tmp.Name() + defer os.Remove(tmpName) // no-op after successful rename + + src, err := os.Open(srcPath) + if err != nil { + tmp.Close() + return "", fmt.Errorf("local: open source: %w", err) + } + if _, err := io.Copy(tmp, src); err != nil { + src.Close() + tmp.Close() + return "", fmt.Errorf("local: copy asset: %w", err) + } + src.Close() + if err := tmp.Close(); err != nil { + return "", fmt.Errorf("local: close temp: %w", err) + } + + if err := os.Rename(tmpName, destPath); err != nil { + return "", fmt.Errorf("local: rename into place: %w", err) + } + return destPath, nil +} + +// sha256File returns the hex-encoded SHA-256 of the file at path. +func sha256File(path string) (string, error) { + f, err := os.Open(path) + if err != nil { + return "", err + } + defer f.Close() + h := sha256.New() + if _, err := io.Copy(h, f); err != nil { + return "", err + } + return hex.EncodeToString(h.Sum(nil)), nil +} diff --git a/internal/embeddings/local/assets_embed.go b/internal/embeddings/local/assets_embed.go new file mode 100644 index 0000000..7f138c0 --- /dev/null +++ b/internal/embeddings/local/assets_embed.go @@ -0,0 +1,101 @@ +//go:build localembed + +package local + +import ( + "embed" + "fmt" + "os" + "path/filepath" + "runtime" +) + +// embeddedAssets carries the platform-independent runtime assets compiled into +// the binary (model + tokenizer — identical bytes on every platform). They are +// downloaded and checksum-verified by `make assets` into ./embedded/ before the +// tagged build compiles (the go:embed directive requires the files to exist). +// The platform-specific ONNX Runtime shared library is embedded separately in +// the per-platform assets_embed_<GOOS>_<GOARCH>.go file (embeddedORTLib / +// ortLibFile). +// +// NOTE: go:embed can only reach files inside this package's own directory tree, +// so the runtime assets live under internal/embeddings/local/embedded/ — NOT the +// repo-root assets/ directory. This is a deliberate deviation from the epic's +// documented assets/embedded/ path. +// +// The static tokenizer library (libtokenizers.a) is linked at build time via +// CGO_LDFLAGS and is deliberately NOT embedded here. +// +//go:embed embedded/model_quantized.onnx +//go:embed embedded/tokenizer.json +var embeddedAssets embed.FS + +// assetsEmbedded reports whether bundled assets are compiled into this build. +// True here (localembed); the stub sets it false. Tests use it to distinguish +// the "no override configured" outcome across the two builds. +const assetsEmbedded = true + +// embeddedAssetSources pairs each runtime asset's embed.FS with its path, in +// extraction order. Each is written to disk under its base name. Model and +// tokenizer come from the shared embeddedAssets; the ONNX Runtime library comes +// from the per-platform embeddedORTLib. +var embeddedAssetSources = []struct { + fs *embed.FS + path string +}{ + {&embeddedAssets, "embedded/" + assetModelFile}, + {&embeddedAssets, "embedded/" + assetTokenizerFile}, + {&embeddedORTLib, "embedded/" + ortLibFile}, +} + +// extractEmbeddedAssets materializes the go:embed-ed runtime assets into +// ~/.agent-memory/runtime/<fingerprint>/ and returns that directory. Extraction +// is atomic (write-temp-then-rename), checksum-computed and idempotent — reusing +// extractAndVerify from assets.go — so concurrent GUI/stdio processes are safe +// and a warm run costs only stat+hash. +func extractEmbeddedAssets() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("local: locate home dir: %w", err) + } + // The directory is namespaced by GOOS-GOARCH in addition to the model + // fingerprint: the ORT shared library is architecture-specific, and two + // builds sharing one home directory is a real scenario (an Intel-mac build + // under Rosetta, or a home directory migrated from an Intel Mac). Without + // the arch in the path, the first build's extraction poisons the second's + // dlopen with an incompatible-architecture error. + destDir := filepath.Join(home, ".agent-memory", "runtime", + fmt.Sprintf("%s-%s-%s", runtime.GOOS, runtime.GOARCH, + fingerprint("local", modelName, modelDim))) + + for _, src := range embeddedAssetSources { + embedPath := src.path + data, err := src.fs.ReadFile(embedPath) + if err != nil { + return "", fmt.Errorf("local: read embedded asset %s: %w", embedPath, err) + } + + // extractAndVerify copies from a source *file* into destDir under the + // same base name, so stage the embedded bytes to a temp file named after + // the asset first. + base := filepath.Base(embedPath) + stageDir, err := os.MkdirTemp("", "agent-memory-embed-*") + if err != nil { + return "", fmt.Errorf("local: stage dir: %w", err) + } + stagePath := filepath.Join(stageDir, base) + if err := os.WriteFile(stagePath, data, 0o644); err != nil { + os.RemoveAll(stageDir) + return "", fmt.Errorf("local: stage embedded asset %s: %w", base, err) + } + + // wantSHA is empty: the bytes are already trusted (compiled into the + // binary); extractAndVerify still hashes for its idempotent fast path. + _, err = extractAndVerify(stagePath, destDir, "") + os.RemoveAll(stageDir) + if err != nil { + return "", err + } + } + return destDir, nil +} diff --git a/internal/embeddings/local/assets_embed_darwin.go b/internal/embeddings/local/assets_embed_darwin.go new file mode 100644 index 0000000..bcdc733 --- /dev/null +++ b/internal/embeddings/local/assets_embed_darwin.go @@ -0,0 +1,21 @@ +//go:build localembed && darwin + +package local + +import "embed" + +// Per-platform ONNX Runtime shared library (macOS). One binary can embed only +// one platform's ORT lib (epic Open Concern #8), so each Phase 5 target adds a +// sibling of this file — the embed directive, the FS var, and the ortLibFile +// constant are the ONLY platform-specific pieces; everything else in this +// package is shared. Both mac arches ship the same dylib file name — arm64 from +// the official ORT release, x86_64 from our own source build (official mac-Intel +// prebuilts stopped at 1.23) hosted in this repo's GitHub releases; the arch +// difference is which artifact `make assets` downloads (assets/manifest.json). +// +//go:embed embedded/libonnxruntime.1.26.0.dylib +var embeddedORTLib embed.FS + +// ortLibFile is the base name of this platform's embedded ONNX Runtime shared +// library, both inside embedded/ and after extraction to the runtime dir. +const ortLibFile = "libonnxruntime.1.26.0.dylib" diff --git a/internal/embeddings/local/assets_embed_linux.go b/internal/embeddings/local/assets_embed_linux.go new file mode 100644 index 0000000..bbffb69 --- /dev/null +++ b/internal/embeddings/local/assets_embed_linux.go @@ -0,0 +1,18 @@ +//go:build localembed && linux + +package local + +import "embed" + +// Per-platform ONNX Runtime shared library (Linux). One binary can embed only +// one platform's ORT lib (epic Open Concern #8). Both Linux arches ship the +// same versioned file name — the arm64/x64 distinction is which artifact +// `make assets` downloads into embedded/ (see assets/manifest.json), so a +// single build-constrained file covers linux/arm64 and linux/amd64. +// +//go:embed embedded/libonnxruntime.so.1.26.0 +var embeddedORTLib embed.FS + +// ortLibFile is the base name of this platform's embedded ONNX Runtime shared +// library, both inside embedded/ and after extraction to the runtime dir. +const ortLibFile = "libonnxruntime.so.1.26.0" diff --git a/internal/embeddings/local/assets_embed_stub.go b/internal/embeddings/local/assets_embed_stub.go new file mode 100644 index 0000000..543c090 --- /dev/null +++ b/internal/embeddings/local/assets_embed_stub.go @@ -0,0 +1,19 @@ +//go:build !localembed + +package local + +import "errors" + +// assetsEmbedded reports whether bundled assets are compiled into this build. +// False here (default/untagged build); the localembed build sets it true. +const assetsEmbedded = false + +// extractEmbeddedAssets is the default-build stub. No assets are embedded in the +// untagged build (no go:embed, so `go build ./...` needs no downloaded files and +// stays native-lib-free), so resolving assets without a developer override +// (Config.AssetsDir / AGENT_MEMORY_LOCAL_ASSETS) is an error here. +func extractEmbeddedAssets() (string, error) { + return "", errors.New( + "local: bundled assets available only in the 'localembed' build " + + "(build with -tags localembed, or set " + assetsDirEnv + ")") +} diff --git a/internal/embeddings/local/assets_embed_windows.go b/internal/embeddings/local/assets_embed_windows.go new file mode 100644 index 0000000..7d47555 --- /dev/null +++ b/internal/embeddings/local/assets_embed_windows.go @@ -0,0 +1,19 @@ +//go:build localembed && windows + +package local + +import "embed" + +// Per-platform ONNX Runtime shared library (Windows x64). One binary can embed +// only one platform's ORT lib (epic Open Concern #8); the official Microsoft +// release ships the DLL — see assets/manifest.json (windows-amd64) and +// Documentation/windows-build.md for the full Windows build procedure +// (the static tokenizer lib has no upstream prebuilt and is built from Rust +// source there). +// +//go:embed embedded/onnxruntime.dll +var embeddedORTLib embed.FS + +// ortLibFile is the base name of this platform's embedded ONNX Runtime shared +// library, both inside embedded/ and after extraction to the runtime dir. +const ortLibFile = "onnxruntime.dll" diff --git a/internal/embeddings/local/chunker_tokenizer.go b/internal/embeddings/local/chunker_tokenizer.go new file mode 100644 index 0000000..200c2cb --- /dev/null +++ b/internal/embeddings/local/chunker_tokenizer.go @@ -0,0 +1,17 @@ +//go:build localembed + +package local + +import "github.com/borzou/vecstore/internal/chunker" + +// NewChunkerTokenizer builds a chunker.Tokenizer backed by the embedding model's +// own HuggingFace tokenizer, resolving tokenizer.json from the configured or +// bundled assets. This lets main.go (untagged) wire a provider-matched chunker +// without importing the tagged tokenizer implementation directly. +func NewChunkerTokenizer(cfg Config) (chunker.Tokenizer, error) { + _, tokPath, _, err := resolveAssets(cfg) + if err != nil { + return nil, err + } + return NewHFChunkerTokenizer(tokPath) +} diff --git a/internal/embeddings/local/chunker_tokenizer_stub.go b/internal/embeddings/local/chunker_tokenizer_stub.go new file mode 100644 index 0000000..5e71415 --- /dev/null +++ b/internal/embeddings/local/chunker_tokenizer_stub.go @@ -0,0 +1,15 @@ +//go:build !localembed + +package local + +import ( + "errors" + + "github.com/borzou/vecstore/internal/chunker" +) + +// NewChunkerTokenizer is the default-build stub. The real, tokenizer-backed +// implementation is provided by the `localembed` build. +func NewChunkerTokenizer(cfg Config) (chunker.Tokenizer, error) { + return nil, errors.New("local tokenizer requires the localembed build") +} diff --git a/internal/embeddings/local/integration_test.go b/internal/embeddings/local/integration_test.go new file mode 100644 index 0000000..9f7a2f1 --- /dev/null +++ b/internal/embeddings/local/integration_test.go @@ -0,0 +1,156 @@ +//go:build localembed + +package local + +import ( + "math" + "strings" + "testing" + + "github.com/borzou/vecstore/internal/chunker" +) + +// TestIntegrationEmbed exercises the real ONNX Runtime + HF tokenizer pipeline +// against the mE5-small assets. It requires the native libraries and the assets +// directory (AGENT_MEMORY_LOCAL_ASSETS), so it only builds under the +// `localembed` tag and skips if assets are absent. +func TestIntegrationEmbed(t *testing.T) { + // Assets come from either the dev override (AGENT_MEMORY_LOCAL_ASSETS) or the + // go:embed-ed set (present in every localembed build). Skip only if neither + // resolves — proving the embed->extract path when run with no override set. + if _, _, _, err := resolveAssets(Config{}); err != nil { + t.Skipf("no local assets available: %v", err) + } + + e := New(Config{Threads: 2, BatchSize: 8}) + + // dimensionality + unit norm + q, err := e.EmbedQuery("hello world") + if err != nil { + t.Fatalf("EmbedQuery: %v", err) + } + if len(q) != 384 { + t.Fatalf("dim = %d, want 384", len(q)) + } + if n := vecNorm(q); math.Abs(n-1) > 1e-4 { + t.Fatalf("query vector norm = %v, want ~1", n) + } + + // Phase 0 sanity ordering: + // related < cross-lingual < unrelated (cosine distance) + query := "how do I reset my password" + related := "steps to recover a forgotten account password" + crossES := "pasos para recuperar una contraseña de cuenta olvidada" + unrelated := "the recipe calls for two cups of flour" + + qv, err := e.EmbedQuery(query) + if err != nil { + t.Fatal(err) + } + docs, err := e.EmbedDocuments([]string{related, crossES, unrelated}) + if err != nil { + t.Fatal(err) + } + for i, d := range docs { + if len(d) != 384 { + t.Fatalf("doc %d dim = %d, want 384", i, len(d)) + } + if n := vecNorm(d); math.Abs(n-1) > 1e-4 { + t.Fatalf("doc %d norm = %v, want ~1", i, n) + } + } + + distRelated := 1 - cosine(qv, docs[0]) + distCross := 1 - cosine(qv, docs[1]) + distUnrelated := 1 - cosine(qv, docs[2]) + + t.Logf("cosine distances: related=%.4f cross-lingual=%.4f unrelated=%.4f", + distRelated, distCross, distUnrelated) + + if !(distRelated < distCross) { + t.Errorf("expected related (%.4f) < cross-lingual (%.4f)", distRelated, distCross) + } + if !(distCross < distUnrelated) { + t.Errorf("expected cross-lingual (%.4f) < unrelated (%.4f)", distCross, distUnrelated) + } +} + +func cosine(a, b []float32) float64 { + var dot float64 + for i := range a { + dot += float64(a[i]) * float64(b[i]) + } + return dot +} + +// TestOrtLibFileIsKnownCandidate is a tripwire for future platform files: the +// per-platform ortLibFile must be a name findDylib recognizes, so a developer +// override directory populated with the same artifacts always resolves. +func TestOrtLibFileIsKnownCandidate(t *testing.T) { + for _, name := range dylibCandidates { + if name == ortLibFile { + return + } + } + t.Fatalf("ortLibFile %q is not in dylibCandidates %v", ortLibFile, dylibCandidates) +} + +// TestIntegrationLargeDocument covers the scenario every platform smoke missed: +// a document big enough to need multiple chunks, through the REAL wired +// pipeline (HF tokenizer + reserved chunk budget + ONNX inference), plus an +// over-long query through the unchunked query path. Regression for the +// "512 by 516" failure that silently dropped every >1-chunk file. +func TestIntegrationLargeDocument(t *testing.T) { + if _, _, _, err := resolveAssets(Config{}); err != nil { + t.Skipf("no local assets available: %v", err) + } + e := New(Config{Threads: 2, BatchSize: 8}) + tok, err := NewChunkerTokenizer(Config{}) + if err != nil { + t.Fatalf("chunker tokenizer: %v", err) + } + c, err := chunker.New( + chunker.WithTokenizer(tok), + chunker.WithMaxInputTokens(e.MaxInputTokens()-EmbedTokenReserve), + ) + if err != nil { + t.Fatalf("chunker: %v", err) + } + + doc := strings.Repeat("Session notes: the demo plan needs a gap execution review and an architecture pivot before the milestone. ", 300) // well beyond one chunk + chunks, err := c.ChunkText(doc) + if err != nil { + t.Fatalf("chunk: %v", err) + } + if len(chunks) < 2 { + t.Fatalf("test needs a multi-chunk doc, got %d chunks", len(chunks)) + } + texts := make([]string, len(chunks)) + for i, ch := range chunks { + texts[i] = ch.Content + } + vecs, err := e.EmbedDocuments(texts) + if err != nil { + t.Fatalf("EmbedDocuments over %d real chunks: %v", len(chunks), err) + } + if len(vecs) != len(chunks) { + t.Fatalf("got %d vectors for %d chunks", len(vecs), len(chunks)) + } + for i, v := range vecs { + if len(v) != 384 { + t.Fatalf("chunk %d dim = %d, want 384", i, len(v)) + } + } + + // The unchunked query path: a query far beyond the context window must + // embed (truncated) rather than crash inference. + longQuery := strings.Repeat("what was the demo plan and architecture pivot for the dossier project ", 60) + qv, err := e.EmbedQuery(longQuery) + if err != nil { + t.Fatalf("EmbedQuery over-long query: %v", err) + } + if len(qv) != 384 { + t.Fatalf("query dim = %d, want 384", len(qv)) + } + t.Logf("large-doc pipeline OK: %d chunks embedded, over-long query embedded", len(chunks)) +} diff --git a/internal/embeddings/local/local.go b/internal/embeddings/local/local.go new file mode 100644 index 0000000..35475fe --- /dev/null +++ b/internal/embeddings/local/local.go @@ -0,0 +1,318 @@ +// Package local implements an in-process, CPU-based embedding provider backed +// by ONNX Runtime and a HuggingFace tokenizer (multilingual-e5-small). +// +// This file (and assets.go) are pure Go with NO CGo / native-library imports, +// so the default `go build`/`go test` stays green without ONNX Runtime or the +// tokenizer static library present. The native infrastructure lives in files +// guarded by the `//go:build localembed` tag (session_ort.go, tokenizer_hf.go), +// with stubs (session_stub.go, tokenizer_stub.go) for the default build. +// +// The pipeline is: prefix -> tokenize -> pad/build tensors -> ONNX forward -> +// mean-pool over the attention mask -> L2-normalize. The E5 family requires a +// "query: " prefix for search text and a "passage: " prefix for indexed text. +package local + +import ( + "errors" + "fmt" + "math" + "sync" + + "github.com/borzou/vecstore/internal/embeddings" +) + +// LocalEmbedder implements the embeddings.Embedder seam. +var _ embeddings.Embedder = (*LocalEmbedder)(nil) + +// errLocalTagRequired is returned by the native constructors in the default +// build (no `localembed` tag). Defined here in a single untagged file so both +// the stubs and any error-comparison logic can reference it. +var errLocalTagRequired = errors.New("local inference requires the 'localembed' build tag") + +// Model constants for multilingual-e5-small. +const ( + modelName = "multilingual-e5-small" + modelDim = 384 + modelMaxTokens = 512 + + queryPrefix = "query: " + passagePrefix = "passage: " + + defaultBatchSize = 16 +) + +// EmbedTokenReserve is the token headroom the embedder consumes around each +// input at embed time: the E5 instruction prefix ("passage: " / "query: ") +// plus the tokenizer's special tokens, rounded up generously. Chunkers must +// budget chunks at MaxInputTokens() − EmbedTokenReserve so the prefixed, +// tokenized sequence never exceeds the model's hard limit — the epic's +// "effective chunk size ≈ 480" (512 − 32). Passing MaxInputTokens() straight +// through as the chunk budget overflows the model by the prefix width (the +// "512 by 516" ORT crash that silently dropped every multi-chunk file). +const EmbedTokenReserve = 32 + +// onnxSession is the seam over the ONNX Runtime session. Implementations take +// padded, batched int64 input tensors (input_ids, attention_mask, +// token_type_ids) and return one mean-pooled vector per input row. Injecting a +// fake here lets the pipeline be unit-tested without native libraries. +type onnxSession interface { + run(inputIDs, attnMask, typeIDs [][]int64) ([][]float32, error) + close() error +} + +// tokenizerBackend is the seam over the HuggingFace tokenizer. It encodes a +// single string to token IDs (including the model's special tokens). +type tokenizerBackend interface { + encode(text string) ([]uint32, error) +} + +// Config configures a LocalEmbedder. +type Config struct { + // AssetsDir is the directory holding the model, tokenizer and ONNX Runtime + // shared library. If empty, the AGENT_MEMORY_LOCAL_ASSETS env var is used. + AssetsDir string + // Threads caps ONNX Runtime intra-op parallelism. <= 0 lets the runtime + // pick a sensible default. + Threads int + // BatchSize is the internal sub-batch size for EmbedDocuments. <= 0 uses + // defaultBatchSize. + BatchSize int +} + +// LocalEmbedder is an in-process ONNX embedder implementing embeddings.Embedder. +type LocalEmbedder struct { + cfg Config + once sync.Once + session onnxSession + tok tokenizerBackend + initErr error +} + +// New constructs a LocalEmbedder. It is cheap: no model is loaded and no native +// library is touched until the first embed call (lazy session init). +func New(cfg Config) *LocalEmbedder { + return &LocalEmbedder{cfg: cfg} +} + +// Dimensions returns the embedding vector dimensionality. +func (e *LocalEmbedder) Dimensions() int { return modelDim } + +// ModelName returns the model identifier. +func (e *LocalEmbedder) ModelName() string { return modelName } + +// MaxInputTokens returns the model's context window in tokens. +func (e *LocalEmbedder) MaxInputTokens() int { return modelMaxTokens } + +// EmbedDocuments embeds indexed content. Each text is prefixed with "passage: " +// then embedded in sub-batches of cfg.BatchSize. +func (e *LocalEmbedder) EmbedDocuments(texts []string) ([][]float32, error) { + if len(texts) == 0 { + return [][]float32{}, nil + } + if err := e.ensureSession(); err != nil { + return nil, err + } + prefixed := withPrefix(passagePrefix, texts) + + batchSize := e.cfg.BatchSize + if batchSize <= 0 { + batchSize = defaultBatchSize + } + + out := make([][]float32, 0, len(prefixed)) + for start := 0; start < len(prefixed); start += batchSize { + end := start + batchSize + if end > len(prefixed) { + end = len(prefixed) + } + vecs, err := e.embedBatch(prefixed[start:end]) + if err != nil { + return nil, err + } + out = append(out, vecs...) + } + return out, nil +} + +// EmbedQuery embeds a single search query, prefixed with "query: ". +func (e *LocalEmbedder) EmbedQuery(text string) ([]float32, error) { + if err := e.ensureSession(); err != nil { + return nil, err + } + vecs, err := e.embedBatch(withPrefix(queryPrefix, []string{text})) + if err != nil { + return nil, err + } + if len(vecs) != 1 { + return nil, fmt.Errorf("local: expected 1 vector, got %d", len(vecs)) + } + return vecs[0], nil +} + +// embedBatch tokenizes, builds padded tensors, runs the session, and +// L2-normalizes each pooled vector. +func (e *LocalEmbedder) embedBatch(texts []string) ([][]float32, error) { + tokenIDs := make([][]uint32, len(texts)) + for i, t := range texts { + ids, err := e.tok.encode(t) + if err != nil { + return nil, fmt.Errorf("local: tokenize: %w", err) + } + // Defense-in-depth: never hand the model more than its context window. + // The chunker budgets indexed chunks below the limit, but queries reach + // here unchunked and any budgeting bug would otherwise crash inference. + tokenIDs[i] = truncateTokens(ids, modelMaxTokens) + } + + ids, mask, types := buildInputs(tokenIDs) + pooled, err := e.session.run(ids, mask, types) + if err != nil { + return nil, fmt.Errorf("local: inference: %w", err) + } + if len(pooled) != len(texts) { + return nil, fmt.Errorf("local: expected %d vectors, got %d", len(texts), len(pooled)) + } + + out := make([][]float32, len(pooled)) + for i, v := range pooled { + out[i] = l2Normalize(v) + } + return out, nil +} + +// ensureSession lazily constructs the ONNX session and tokenizer exactly once. +// If a session and tokenizer are already set (tests inject fakes), real +// construction is skipped. +func (e *LocalEmbedder) ensureSession() error { + e.once.Do(func() { + if e.session != nil && e.tok != nil { + return + } + modelPath, tokPath, dylibPath, err := resolveAssets(e.cfg) + if err != nil { + e.initErr = err + return + } + sess, err := newORTSession(dylibPath, modelPath, e.cfg.Threads) + if err != nil { + e.initErr = err + return + } + tok, err := newHFTokenizer(tokPath) + if err != nil { + _ = sess.close() + e.initErr = err + return + } + e.session = sess + e.tok = tok + }) + return e.initErr +} + +// --------------------------------------------------------------------------- +// Pure helpers (unit-tested, no native deps) +// --------------------------------------------------------------------------- + +// withPrefix returns a new slice with prefix prepended to each text. +func withPrefix(prefix string, texts []string) []string { + out := make([]string, len(texts)) + for i, t := range texts { + out[i] = prefix + t + } + return out +} + +// truncateTokens caps a token sequence at max tokens. The tokenizer emits the +// model's end-of-sequence special token last; truncation preserves it so an +// over-long sequence stays well-formed (<s> … </s>) instead of ending +// mid-stream. +func truncateTokens(ids []uint32, max int) []uint32 { + if max <= 0 || len(ids) <= max { + return ids + } + out := make([]uint32, max) + copy(out, ids[:max-1]) + out[max-1] = ids[len(ids)-1] + return out +} + +// buildInputs converts per-sequence token IDs into padded, batched int64 +// tensors. All sequences are right-padded to the batch's max length. The +// attention mask is 1 for real tokens and 0 for padding; token_type_ids are all +// zero (required by the mE5 Xenova export, which takes three INT64 inputs). +func buildInputs(tokenIDs [][]uint32) (ids, mask, types [][]int64) { + n := len(tokenIDs) + ids = make([][]int64, n) + mask = make([][]int64, n) + types = make([][]int64, n) + + maxLen := 0 + for _, seq := range tokenIDs { + if len(seq) > maxLen { + maxLen = len(seq) + } + } + + for i, seq := range tokenIDs { + rowIDs := make([]int64, maxLen) + rowMask := make([]int64, maxLen) + rowTypes := make([]int64, maxLen) // all zero + for j, id := range seq { + rowIDs[j] = int64(id) + rowMask[j] = 1 + } + ids[i] = rowIDs + mask[i] = rowMask + types[i] = rowTypes + } + return ids, mask, types +} + +// meanPool averages the token hidden states of a single sequence, weighted by +// the attention mask (padding positions are ignored). hidden is [seqLen][dim]; +// mask is [seqLen]. Returns a [dim] vector. If no positions are active it +// returns a zero vector of the input's dimensionality. +func meanPool(hidden [][]float32, mask []int64) []float32 { + if len(hidden) == 0 { + return nil + } + dim := len(hidden[0]) + out := make([]float32, dim) + var count float64 + for i, row := range hidden { + if i < len(mask) && mask[i] == 0 { + continue + } + count++ + for k := 0; k < dim && k < len(row); k++ { + out[k] += row[k] + } + } + if count == 0 { + return out + } + for k := range out { + out[k] = float32(float64(out[k]) / count) + } + return out +} + +// l2Normalize returns a unit-norm copy of v. A zero vector is returned +// unchanged (as a copy) to avoid division by zero. +func l2Normalize(v []float32) []float32 { + out := make([]float32, len(v)) + var norm float64 + for _, x := range v { + norm += float64(x) * float64(x) + } + if norm == 0 { + copy(out, v) + return out + } + norm = math.Sqrt(norm) + for i, x := range v { + out[i] = float32(float64(x) / norm) + } + return out +} diff --git a/internal/embeddings/local/local_test.go b/internal/embeddings/local/local_test.go new file mode 100644 index 0000000..14d41bc --- /dev/null +++ b/internal/embeddings/local/local_test.go @@ -0,0 +1,409 @@ +package local + +import ( + "math" + "os" + "path/filepath" + "reflect" + "strings" + "testing" +) + +// --------------------------------------------------------------------------- +// Fakes (no native libraries) +// --------------------------------------------------------------------------- + +// fakeTokenizer maps a string to token IDs by length so different inputs pad +// differently. Deterministic and native-free. +type fakeTokenizer struct { + calls []string +} + +func (f *fakeTokenizer) encode(text string) ([]uint32, error) { + f.calls = append(f.calls, text) + ids := make([]uint32, len(text)) + for i := range ids { + ids[i] = uint32(i + 1) + } + return ids, nil +} + +// fakeSession records the tensors it received and returns a fixed pooled vector +// per input row (un-normalized, so l2Normalize is observable downstream). +type fakeSession struct { + lastIDs [][]int64 + lastMask [][]int64 + lastTypes [][]int64 + vec []float32 // returned for every row + closed bool +} + +func (s *fakeSession) run(inputIDs, attnMask, typeIDs [][]int64) ([][]float32, error) { + s.lastIDs = inputIDs + s.lastMask = attnMask + s.lastTypes = typeIDs + out := make([][]float32, len(inputIDs)) + for i := range out { + cp := make([]float32, len(s.vec)) + copy(cp, s.vec) + out[i] = cp + } + return out, nil +} + +func (s *fakeSession) close() error { s.closed = true; return nil } + +func newFakeEmbedder(vec []float32) (*LocalEmbedder, *fakeSession, *fakeTokenizer) { + sess := &fakeSession{vec: vec} + tk := &fakeTokenizer{} + e := New(Config{}) + e.session = sess + e.tok = tk + return e, sess, tk +} + +// --------------------------------------------------------------------------- +// Pure helpers +// --------------------------------------------------------------------------- + +func TestWithPrefix(t *testing.T) { + got := withPrefix("passage: ", []string{"a", "b"}) + want := []string{"passage: a", "passage: b"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("withPrefix = %v, want %v", got, want) + } + // original inputs unchanged + orig := []string{"x"} + _ = withPrefix("query: ", orig) + if orig[0] != "x" { + t.Fatalf("withPrefix mutated input: %v", orig) + } +} + +func TestBuildInputs(t *testing.T) { + tokenIDs := [][]uint32{ + {10, 11, 12}, // len 3 + {20, 21}, // len 2 -> padded to 3 + } + ids, mask, types := buildInputs(tokenIDs) + + // three parallel inputs, padded to max len 3 + wantIDs := [][]int64{{10, 11, 12}, {20, 21, 0}} + wantMask := [][]int64{{1, 1, 1}, {1, 1, 0}} + wantTypes := [][]int64{{0, 0, 0}, {0, 0, 0}} + + if !reflect.DeepEqual(ids, wantIDs) { + t.Errorf("ids = %v, want %v", ids, wantIDs) + } + if !reflect.DeepEqual(mask, wantMask) { + t.Errorf("mask = %v, want %v", mask, wantMask) + } + if !reflect.DeepEqual(types, wantTypes) { + t.Errorf("token_type_ids = %v, want %v (must be all zero)", types, wantTypes) + } + + // exactly three input tensors are produced, all same shape + if len(ids) != len(mask) || len(mask) != len(types) { + t.Fatalf("input row counts differ: ids=%d mask=%d types=%d", len(ids), len(mask), len(types)) + } + for i := range ids { + if len(ids[i]) != 3 || len(mask[i]) != 3 || len(types[i]) != 3 { + t.Fatalf("row %d not padded to 3: ids=%v mask=%v types=%v", i, ids[i], mask[i], types[i]) + } + } +} + +func TestBuildInputsTokenTypeIDsAllZero(t *testing.T) { + _, _, types := buildInputs([][]uint32{{1, 2, 3, 4}, {5}}) + for i, row := range types { + for j, v := range row { + if v != 0 { + t.Fatalf("token_type_ids[%d][%d] = %d, want 0", i, j, v) + } + } + } +} + +func TestMeanPool(t *testing.T) { + // two real tokens, one padded (mask 0). Padded row must be ignored. + hidden := [][]float32{ + {1, 2}, + {3, 4}, + {100, 100}, // padding — ignored + } + mask := []int64{1, 1, 0} + got := meanPool(hidden, mask) + want := []float32{2, 3} // (1+3)/2, (2+4)/2 + if !reflect.DeepEqual(got, want) { + t.Fatalf("meanPool = %v, want %v", got, want) + } +} + +func TestMeanPoolAllMasked(t *testing.T) { + got := meanPool([][]float32{{5, 5}}, []int64{0}) + want := []float32{0, 0} + if !reflect.DeepEqual(got, want) { + t.Fatalf("meanPool all-masked = %v, want %v", got, want) + } +} + +func TestL2Normalize(t *testing.T) { + got := l2Normalize([]float32{3, 4}) + want := []float32{0.6, 0.8} + for i := range want { + if math.Abs(float64(got[i]-want[i])) > 1e-6 { + t.Fatalf("l2Normalize = %v, want %v", got, want) + } + } + // resulting norm ~= 1 + if n := vecNorm(got); math.Abs(n-1) > 1e-6 { + t.Fatalf("norm = %v, want ~1", n) + } +} + +func TestL2NormalizeZero(t *testing.T) { + got := l2Normalize([]float32{0, 0, 0}) + if vecNorm(got) != 0 { + t.Fatalf("zero vector should stay zero, got %v", got) + } +} + +// --------------------------------------------------------------------------- +// Pipeline via fakes +// --------------------------------------------------------------------------- + +func TestEmbedQueryPrefixAndShape(t *testing.T) { + e, sess, tk := newFakeEmbedder([]float32{3, 4}) // un-normalized + + vec, err := e.EmbedQuery("hello") + if err != nil { + t.Fatal(err) + } + + // query prefix applied to the tokenizer input + if len(tk.calls) != 1 || tk.calls[0] != "query: hello" { + t.Fatalf("tokenizer calls = %v, want [\"query: hello\"]", tk.calls) + } + // result is L2-normalized (3,4 -> 0.6,0.8) + if math.Abs(float64(vec[0]-0.6)) > 1e-6 || math.Abs(float64(vec[1]-0.8)) > 1e-6 { + t.Fatalf("EmbedQuery vec = %v, want normalized [0.6 0.8]", vec) + } + if n := vecNorm(vec); math.Abs(n-1) > 1e-6 { + t.Fatalf("EmbedQuery norm = %v, want ~1", n) + } + // three padded inputs reached the session, token_type_ids all zero + assertThreeZeroTypeInputs(t, sess) +} + +func TestEmbedDocumentsPrefixAndBatching(t *testing.T) { + e, sess, tk := newFakeEmbedder([]float32{0, 3}) + e.cfg.BatchSize = 2 // force multiple sub-batches over 5 inputs + + texts := []string{"a", "b", "c", "d", "e"} + vecs, err := e.EmbedDocuments(texts) + if err != nil { + t.Fatal(err) + } + + if len(vecs) != len(texts) { + t.Fatalf("got %d vectors, want %d", len(vecs), len(texts)) + } + // passage prefix applied to every input + for i, c := range tk.calls { + want := "passage: " + texts[i] + if c != want { + t.Fatalf("tokenizer call %d = %q, want %q", i, c, want) + } + } + // every returned vector is unit-norm and 2-dim + for i, v := range vecs { + if len(v) != 2 { + t.Fatalf("vec %d has dim %d, want 2", i, len(v)) + } + if n := vecNorm(v); math.Abs(n-1) > 1e-6 { + t.Fatalf("vec %d norm = %v, want ~1", i, n) + } + } + // 5 inputs at batch size 2 -> sub-batches of 2,2,1; last one has 1 row + if len(sess.lastIDs) != 1 { + t.Fatalf("last sub-batch rows = %d, want 1", len(sess.lastIDs)) + } + assertThreeZeroTypeInputs(t, sess) +} + +func TestEmbedDocumentsEmpty(t *testing.T) { + e, _, _ := newFakeEmbedder([]float32{1}) + vecs, err := e.EmbedDocuments(nil) + if err != nil { + t.Fatal(err) + } + if len(vecs) != 0 { + t.Fatalf("want empty result, got %v", vecs) + } +} + +func TestMetadata(t *testing.T) { + e := New(Config{}) + if e.Dimensions() != 384 { + t.Errorf("Dimensions = %d, want 384", e.Dimensions()) + } + if e.ModelName() != "multilingual-e5-small" { + t.Errorf("ModelName = %q", e.ModelName()) + } + if e.MaxInputTokens() != 512 { + t.Errorf("MaxInputTokens = %d, want 512", e.MaxInputTokens()) + } +} + +// --------------------------------------------------------------------------- +// Assets helpers +// --------------------------------------------------------------------------- + +func TestFingerprintStable(t *testing.T) { + a := fingerprint("local", "multilingual-e5-small", 384) + b := fingerprint("local", "multilingual-e5-small", 384) + if a != b { + t.Fatalf("fingerprint not deterministic: %q vs %q", a, b) + } + if a == fingerprint("openai", "text-embedding-3-small", 1536) { + t.Fatalf("fingerprint collision across providers") + } + if len(a) != 16 { + t.Fatalf("fingerprint len = %d, want 16", len(a)) + } +} + +func TestResolveAssetsMissing(t *testing.T) { + if _, _, _, err := resolveAssets(Config{AssetsDir: t.TempDir()}); err == nil { + t.Fatal("expected error for empty assets dir") + } + // No dir and no env. Without embedded assets (default build) this is an + // error; in the localembed build the go:embed-ed set is the valid fallback, + // so success is expected there. + t.Setenv(assetsDirEnv, "") + _, _, _, err := resolveAssets(Config{}) + if !assetsEmbedded && err == nil { + t.Fatal("expected error when no assets dir configured and none embedded") + } +} + +func TestExtractAndVerify(t *testing.T) { + src := filepath.Join(t.TempDir(), assetModelFile) + content := []byte("weights") + if err := os.WriteFile(src, content, 0o644); err != nil { + t.Fatal(err) + } + want, err := sha256File(src) + if err != nil { + t.Fatal(err) + } + dest := t.TempDir() + + got, err := extractAndVerify(src, dest, want) + if err != nil { + t.Fatal(err) + } + out, err := os.ReadFile(got) + if err != nil || string(out) != string(content) { + t.Fatalf("extracted content mismatch: %q err=%v", out, err) + } + + // idempotent second call + if _, err := extractAndVerify(src, dest, want); err != nil { + t.Fatalf("second extract failed: %v", err) + } + // checksum mismatch is rejected + if _, err := extractAndVerify(src, t.TempDir(), "deadbeef"); err == nil { + t.Fatal("expected checksum mismatch error") + } +} + +// --------------------------------------------------------------------------- +// helpers +// --------------------------------------------------------------------------- + +func assertThreeZeroTypeInputs(t *testing.T, s *fakeSession) { + t.Helper() + if s.lastIDs == nil || s.lastMask == nil || s.lastTypes == nil { + t.Fatal("expected all three inputs (ids, mask, types) to reach session") + } + if len(s.lastIDs) != len(s.lastMask) || len(s.lastMask) != len(s.lastTypes) { + t.Fatalf("input row counts differ: ids=%d mask=%d types=%d", + len(s.lastIDs), len(s.lastMask), len(s.lastTypes)) + } + for i, row := range s.lastTypes { + for j, v := range row { + if v != 0 { + t.Fatalf("token_type_ids[%d][%d] = %d, want 0", i, j, v) + } + } + } +} + +func vecNorm(v []float32) float64 { + var s float64 + for _, x := range v { + s += float64(x) * float64(x) + } + return math.Sqrt(s) +} + +// --- token-budget regression tests (the "512 by 516" bug) ------------------- + +func TestTruncateTokens(t *testing.T) { + mk := func(n int) []uint32 { + ids := make([]uint32, n) + for i := range ids { + ids[i] = uint32(i + 100) + } + return ids + } + t.Run("under limit unchanged", func(t *testing.T) { + ids := mk(10) + got := truncateTokens(ids, 512) + if len(got) != 10 { + t.Fatalf("len = %d, want 10", len(got)) + } + }) + t.Run("at limit unchanged", func(t *testing.T) { + if got := truncateTokens(mk(512), 512); len(got) != 512 { + t.Fatalf("len = %d, want 512", len(got)) + } + }) + t.Run("over limit capped preserving EOS", func(t *testing.T) { + ids := mk(516) + got := truncateTokens(ids, 512) + if len(got) != 512 { + t.Fatalf("len = %d, want 512", len(got)) + } + if got[511] != ids[515] { + t.Fatalf("last token = %d, want the original final (EOS) token %d", got[511], ids[515]) + } + if got[510] != ids[510] { + t.Fatalf("truncation should keep the first max-1 tokens intact") + } + }) + t.Run("zero max disables", func(t *testing.T) { + if got := truncateTokens(mk(600), 0); len(got) != 600 { + t.Fatalf("max=0 should disable truncation") + } + }) +} + +// TestEmbedBatchTruncatesOversizedInput proves the model never receives more +// than modelMaxTokens even when a caller hands the embedder unchunked text +// (the query path has no chunker; regression for the silent multi-chunk-file +// indexing failure). +func TestEmbedBatchTruncatesOversizedInput(t *testing.T) { + e, sess, _ := newFakeEmbedder([]float32{1, 0}) + long := strings.Repeat("x", modelMaxTokens+300) // fakeTokenizer: 1 token per byte + + if _, err := e.EmbedQuery(long); err != nil { + t.Fatalf("EmbedQuery long input: %v", err) + } + for _, row := range sess.lastIDs { + if len(row) > modelMaxTokens { + t.Fatalf("session received %d tokens, model limit is %d", len(row), modelMaxTokens) + } + } +} diff --git a/internal/embeddings/local/session_ort.go b/internal/embeddings/local/session_ort.go new file mode 100644 index 0000000..c5b892e --- /dev/null +++ b/internal/embeddings/local/session_ort.go @@ -0,0 +1,141 @@ +//go:build localembed + +package local + +import ( + "fmt" + "sync" + + ort "github.com/yalue/onnxruntime_go" +) + +// mE5 (Xenova quantized export) takes three INT64 inputs and produces the +// token-level last_hidden_state. +var ( + ortInputNames = []string{"input_ids", "attention_mask", "token_type_ids"} + ortOutputNames = []string{"last_hidden_state"} +) + +// ortInitOnce guards process-wide ONNX Runtime environment initialization. +var ortInitOnce sync.Once +var ortInitErr error + +// ortSession is the real onnxSession implementation backed by ONNX Runtime. +type ortSession struct { + sess *ort.DynamicAdvancedSession +} + +// newORTSession initializes the ONNX Runtime environment (once per process) and +// creates a session for the given model. threads caps intra-op parallelism so +// background indexing does not peg the machine (<= 0 leaves the runtime default). +func newORTSession(dylibPath, modelPath string, threads int) (onnxSession, error) { + ortInitOnce.Do(func() { + ort.SetSharedLibraryPath(dylibPath) + if !ort.IsInitialized() { + ortInitErr = ort.InitializeEnvironment() + } + }) + if ortInitErr != nil { + return nil, fmt.Errorf("local: init ONNX Runtime: %w", ortInitErr) + } + + opts, err := ort.NewSessionOptions() + if err != nil { + return nil, fmt.Errorf("local: session options: %w", err) + } + defer opts.Destroy() + if threads > 0 { + if err := opts.SetIntraOpNumThreads(threads); err != nil { + return nil, fmt.Errorf("local: set intra-op threads: %w", err) + } + } + + sess, err := ort.NewDynamicAdvancedSession(modelPath, ortInputNames, ortOutputNames, opts) + if err != nil { + return nil, fmt.Errorf("local: create session: %w", err) + } + return &ortSession{sess: sess}, nil +} + +// run builds padded 2-D tensors for the batch, executes the model, and +// mean-pools each sequence's token embeddings over its attention mask. It +// returns one (un-normalized) pooled vector per input row; L2 normalization is +// applied by the caller. +func (s *ortSession) run(inputIDs, attnMask, typeIDs [][]int64) ([][]float32, error) { + n := len(inputIDs) + if n == 0 { + return [][]float32{}, nil + } + maxLen := len(inputIDs[0]) + + flatIDs := flatten(inputIDs, n, maxLen) + flatMask := flatten(attnMask, n, maxLen) + flatType := flatten(typeIDs, n, maxLen) + + shape := ort.NewShape(int64(n), int64(maxLen)) + + tIDs, err := ort.NewTensor(shape, flatIDs) + if err != nil { + return nil, fmt.Errorf("local: input_ids tensor: %w", err) + } + defer tIDs.Destroy() + tMask, err := ort.NewTensor(shape, flatMask) + if err != nil { + return nil, fmt.Errorf("local: attention_mask tensor: %w", err) + } + defer tMask.Destroy() + tType, err := ort.NewTensor(shape, flatType) + if err != nil { + return nil, fmt.Errorf("local: token_type_ids tensor: %w", err) + } + defer tType.Destroy() + + outputs := []ort.Value{nil} + if err := s.sess.Run([]ort.Value{tIDs, tMask, tType}, outputs); err != nil { + return nil, fmt.Errorf("local: run: %w", err) + } + out, ok := outputs[0].(*ort.Tensor[float32]) + if !ok { + return nil, fmt.Errorf("local: unexpected output tensor type %T", outputs[0]) + } + defer out.Destroy() + + data := out.GetData() + total := n * maxLen + if total == 0 { + return nil, fmt.Errorf("local: empty output") + } + dim := len(data) / total + if dim*total != len(data) { + return nil, fmt.Errorf("local: output size %d not divisible by n*maxLen=%d", len(data), total) + } + + // Reshape per sequence into [maxLen][dim] and mean-pool over the mask. + pooled := make([][]float32, n) + for i := 0; i < n; i++ { + hidden := make([][]float32, maxLen) + for j := 0; j < maxLen; j++ { + base := (i*maxLen + j) * dim + hidden[j] = data[base : base+dim] + } + pooled[i] = meanPool(hidden, attnMask[i]) + } + return pooled, nil +} + +func (s *ortSession) close() error { + if s.sess != nil { + s.sess.Destroy() + s.sess = nil + } + return nil +} + +// flatten concatenates n rows of length maxLen into a single row-major buffer. +func flatten(rows [][]int64, n, maxLen int) []int64 { + flat := make([]int64, n*maxLen) + for i := 0; i < n; i++ { + copy(flat[i*maxLen:], rows[i]) + } + return flat +} diff --git a/internal/embeddings/local/session_stub.go b/internal/embeddings/local/session_stub.go new file mode 100644 index 0000000..8dd1a16 --- /dev/null +++ b/internal/embeddings/local/session_stub.go @@ -0,0 +1,9 @@ +//go:build !localembed + +package local + +// newORTSession is the default-build stub. Building with the `localembed` tag +// (and linking ONNX Runtime) provides the real implementation. +func newORTSession(dylibPath, modelPath string, threads int) (onnxSession, error) { + return nil, errLocalTagRequired +} diff --git a/internal/embeddings/local/tokenizer_hf.go b/internal/embeddings/local/tokenizer_hf.go new file mode 100644 index 0000000..aab35ba --- /dev/null +++ b/internal/embeddings/local/tokenizer_hf.go @@ -0,0 +1,83 @@ +//go:build localembed + +package local + +import ( + "fmt" + "os" + + tok "github.com/daulet/tokenizers" + + "github.com/borzou/vecstore/internal/chunker" +) + +// loadTokenizer reads a HuggingFace tokenizer.json and constructs a tokenizer. +func loadTokenizer(path string) (*tok.Tokenizer, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("local: read tokenizer: %w", err) + } + t, err := tok.FromBytes(data) + if err != nil { + return nil, fmt.Errorf("local: load tokenizer: %w", err) + } + return t, nil +} + +// hfTokenizer is the real tokenizerBackend for inference. It encodes WITH the +// model's special tokens (<s> ... </s>), which the ONNX model expects. +type hfTokenizer struct { + t *tok.Tokenizer +} + +// newHFTokenizer constructs the inference tokenizer backend. +func newHFTokenizer(tokPath string) (tokenizerBackend, error) { + t, err := loadTokenizer(tokPath) + if err != nil { + return nil, err + } + return &hfTokenizer{t: t}, nil +} + +func (h *hfTokenizer) encode(text string) ([]uint32, error) { + ids, _ := h.t.Encode(text, true) // addSpecialTokens = true + return ids, nil +} + +// HFChunkerTokenizer adapts the HuggingFace tokenizer to chunker.Tokenizer so +// the chunker can measure boundaries in the embedding model's own tokens. +// main.go injects it via chunker.WithTokenizer when the provider is local +// (wired in Phase 3). It encodes WITHOUT special tokens so token counts reflect +// content only; the chunk-size clamp accounts for special/prefix tokens +// separately. +type HFChunkerTokenizer struct { + t *tok.Tokenizer +} + +// compile-time assertion that the adapter satisfies the chunker seam. +var _ chunker.Tokenizer = (*HFChunkerTokenizer)(nil) + +// NewHFChunkerTokenizer loads a tokenizer.json for use as a chunker.Tokenizer. +func NewHFChunkerTokenizer(tokPath string) (*HFChunkerTokenizer, error) { + t, err := loadTokenizer(tokPath) + if err != nil { + return nil, err + } + return &HFChunkerTokenizer{t: t}, nil +} + +// Encode returns content token IDs (no special tokens). +func (h *HFChunkerTokenizer) Encode(text string) ([]uint32, error) { + ids, _ := h.t.Encode(text, false) + return ids, nil +} + +// Decode reconstructs text from token IDs, skipping special tokens. +func (h *HFChunkerTokenizer) Decode(tokens []uint32) (string, error) { + return h.t.Decode(tokens, true), nil +} + +// Close releases the underlying tokenizer. +func (h *HFChunkerTokenizer) Close() error { + return h.t.Close() +} diff --git a/internal/embeddings/local/tokenizer_stub.go b/internal/embeddings/local/tokenizer_stub.go new file mode 100644 index 0000000..6513b5c --- /dev/null +++ b/internal/embeddings/local/tokenizer_stub.go @@ -0,0 +1,9 @@ +//go:build !localembed + +package local + +// newHFTokenizer is the default-build stub. Building with the `localembed` tag +// (and linking libtokenizers) provides the real implementation. +func newHFTokenizer(tokPath string) (tokenizerBackend, error) { + return nil, errLocalTagRequired +} diff --git a/internal/embeddings/openai.go b/internal/embeddings/openai.go index e3f7925..77a4c4a 100644 --- a/internal/embeddings/openai.go +++ b/internal/embeddings/openai.go @@ -52,8 +52,8 @@ type apiError struct { Type string `json:"type"` } -// Embed generates embeddings for the given texts, batching as needed. -func (e *OpenAIEmbedder) Embed(texts []string) ([][]float32, error) { +// EmbedDocuments generates embeddings for the given texts, batching as needed. +func (e *OpenAIEmbedder) EmbedDocuments(texts []string) ([][]float32, error) { if len(texts) == 0 { return nil, nil } @@ -142,6 +142,22 @@ func (e *OpenAIEmbedder) callAPI(texts []string) ([]embeddingData, error) { return nil, fmt.Errorf("max retries exceeded: %w", lastErr) } +// EmbedQuery generates an embedding for a single query text. +func (e *OpenAIEmbedder) EmbedQuery(text string) ([]float32, error) { + vecs, err := e.EmbedDocuments([]string{text}) + if err != nil { + return nil, err + } + if len(vecs) == 0 { + return nil, nil + } + return vecs[0], nil +} + +// MaxInputTokens returns 0, meaning no practical limit is enforced here. +// OpenAI's per-input limit is 8191 tokens, well above our chunk sizes. +func (e *OpenAIEmbedder) MaxInputTokens() int { return 0 } + // Dimensions returns the embedding dimension for the configured model. func (e *OpenAIEmbedder) Dimensions() int { switch e.model { diff --git a/internal/embeddings/openai_test.go b/internal/embeddings/openai_test.go index 43a8086..f530aa2 100644 --- a/internal/embeddings/openai_test.go +++ b/internal/embeddings/openai_test.go @@ -48,7 +48,7 @@ func TestEmbed_Success(t *testing.T) { }) defer srv.Close() - results, err := embedder.Embed([]string{"hello", "world"}) + results, err := embedder.EmbedDocuments([]string{"hello", "world"}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -92,7 +92,7 @@ func TestEmbed_Batching(t *testing.T) { texts[i] = "text" } - results, err := embedder.Embed(texts) + results, err := embedder.EmbedDocuments(texts) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -117,7 +117,7 @@ func TestEmbed_APIError(t *testing.T) { }) defer srv.Close() - _, err := embedder.Embed([]string{"hello"}) + _, err := embedder.EmbedDocuments([]string{"hello"}) if err == nil { t.Fatal("expected error, got nil") } @@ -150,7 +150,7 @@ func TestEmbed_RetryOn429(t *testing.T) { }) defer srv.Close() - results, err := embedder.Embed([]string{"hello"}) + results, err := embedder.EmbedDocuments([]string{"hello"}) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -164,7 +164,7 @@ func TestEmbed_RetryOn429(t *testing.T) { func TestEmbed_EmptyInput(t *testing.T) { embedder := NewOpenAIEmbedder("key", "text-embedding-3-small") - results, err := embedder.Embed(nil) + results, err := embedder.EmbedDocuments(nil) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -191,3 +191,39 @@ func TestModelName(t *testing.T) { t.Errorf("expected text-embedding-3-small, got %s", e.ModelName()) } } + +func TestEmbedQuery(t *testing.T) { + srv, embedder := newMockServer(t, func(w http.ResponseWriter, r *http.Request) { + var req embeddingRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Fatalf("decode request: %v", err) + } + + resp := embeddingResponse{} + for i := range req.Input { + resp.Data = append(resp.Data, embeddingData{ + Embedding: makeEmbedding(1536), + Index: i, + }) + } + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + }) + defer srv.Close() + + vec, err := embedder.EmbedQuery("hello") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(vec) != 1536 { + t.Errorf("expected 1536 dimensions, got %d", len(vec)) + } +} + +func TestMaxInputTokens(t *testing.T) { + e := NewOpenAIEmbedder("key", "text-embedding-3-small") + if e.MaxInputTokens() != 0 { + t.Errorf("expected 0, got %d", e.MaxInputTokens()) + } +} diff --git a/internal/engine/engine.go b/internal/engine/engine.go index 7eeccaf..a5015ad 100644 --- a/internal/engine/engine.go +++ b/internal/engine/engine.go @@ -7,6 +7,7 @@ import ( "log" "os" "path/filepath" + "strconv" "strings" "sync" "time" @@ -51,7 +52,8 @@ type Engine struct { extractor extractor.Extractor mu sync.Mutex indexing bool - stopCh chan struct{} // closed by Stop() to cancel in-flight indexing + stopCh chan struct{} // closed by Stop() to cancel in-flight indexing + indexWG sync.WaitGroup // tracks in-flight indexing work; Stop() waits on it so shutdown lands on a file boundary // Progress tracking indexedFiles int totalToIndex int @@ -74,6 +76,13 @@ func (eng *Engine) SetEmbedder(e embeddings.Embedder) { eng.embedder = e } +// SetChunker swaps the chunker. A provider switch must swap the chunker +// alongside the embedder so the tokenizer and max-input clamp match the new +// model. Must be called while the engine is stopped. +func (eng *Engine) SetChunker(c chunker.Chunker) { + eng.chunker = c +} + // GetIgnorePatterns returns the current ignore pattern list. If none have been // configured yet, it seeds and persists the defaults. func (eng *Engine) GetIgnorePatterns() ([]string, error) { @@ -196,8 +205,40 @@ func (eng *Engine) ListLogEntries(limit, offset int) ([]domain.ActivityLogEntry, } // IndexFile extracts text from a file, hashes it, and runs the chunk-embed-store -// pipeline if the content has changed since the last index. +// pipeline if the content has changed since the last index. Failures are +// recorded in the activity log exactly once, here — callers must not add their +// own logActivity for the returned error (stderr context lines are fine). +// Error rows are upserted per path (timestamp/detail updated in place), so a +// persistently-failing file keeps one living row instead of appending an +// identical row every launch. func (eng *Engine) IndexFile(path string) error { + err := eng.indexFile(path) + if err != nil { + eng.logIndexError(path, err) + } + return err +} + +// logIndexError records an index failure via the per-path upsert. +func (eng *Engine) logIndexError(path string, err error) { + eng.upsertError(path, fmt.Sprintf("index: %v", err)) +} + +// upsertError records an error for a path, replacing any previous error row +// for that path (latest failure state wins; duplicates never accumulate). +func (eng *Engine) upsertError(path, detail string) { + entry := domain.ActivityLogEntry{ + Timestamp: time.Now(), + Path: path, + Action: "error", + Detail: detail, + } + if err := eng.store.UpsertLogEntry(entry); err != nil { + log.Printf("engine: log error row: %v", err) + } +} + +func (eng *Engine) indexFile(path string) error { // 1. Check if the file type is supported at all. if !eng.extractor.IsSupported(path) { eng.logActivity(path, "ignored", "unsupported file type") @@ -223,7 +264,6 @@ func (eng *Engine) IndexFile(path string) error { // 4. Extract text content (handles text, docx, xlsx, pptx, metadata, etc.) result, err := eng.extractor.Extract(path) if err != nil { - eng.logActivity(path, "error", fmt.Sprintf("extract: %v", err)) return fmt.Errorf("extract %s: %w", path, err) } if result.Text == "" { @@ -262,14 +302,11 @@ func (eng *Engine) IndexFile(path string) error { return fmt.Errorf("find directory for %s: %w", path, err) } - // 9. Remove old chunks if file existed. - if existing != nil { - if err := eng.store.RemoveChunksByFile(existing.ID); err != nil { - return fmt.Errorf("remove old chunks for %s: %w", path, err) - } - } - - // 10. Upsert file record. + // 9. Atomically replace the file's index entry (remove old chunks, upsert + // the file row, insert new chunks) in ONE transaction. A process death + // mid-index must leave the file either fully indexed or untouched — three + // separate writes could record the file at the new hash with zero chunks, + // and the hash short-circuit above would then skip it forever. file := domain.File{ DirectoryID: dirID, Path: path, @@ -279,29 +316,17 @@ func (eng *Engine) IndexFile(path string) error { if existing != nil { file.ID = existing.ID } - if err := eng.store.UpsertFile(file); err != nil { - return fmt.Errorf("upsert file %s: %w", path, err) - } - - // 11. Insert new chunks. Re-fetch the file ID because INSERT OR REPLACE - // may have assigned a new auto-increment ID. - stored, err := eng.store.GetFileByPath(path) - if err != nil { - return fmt.Errorf("get file after upsert %s: %w", path, err) - } - if stored == nil { - return fmt.Errorf("file not found after upsert: %s", path) - } - file.ID = stored.ID - if err := eng.store.InsertChunks(file.ID, domainChunks); err != nil { - return fmt.Errorf("insert chunks for %s: %w", path, err) + if err := eng.store.UpsertFileWithChunks(file, domainChunks); err != nil { + return fmt.Errorf("index write for %s: %w", path, err) } eng.logActivity(path, "indexed", fmt.Sprintf("%d chunks", len(domainChunks))) return nil } -// maxTokensPerBatch is the max tokens per OpenAI embedding API call. +// maxTokensPerBatch is the max tokens per OpenAI embedding API call. This is +// OpenAI-API-specific and harmless for a local embedder (which sub-batches +// internally), since it only bounds how many chunks are sent per call. const maxTokensPerBatch = 250000 // conservative, API limit is 300K // embedBatched sends chunks to the embedder in batches that fit within the API @@ -342,7 +367,7 @@ func (eng *Engine) embedSlice(chunks []chunker.ChunkResult) ([][]float32, error) for i, c := range chunks { texts[i] = c.Content } - return eng.embedder.Embed(texts) + return eng.embedder.EmbedDocuments(texts) } // findDirectoryID returns the directory ID for the watched directory that @@ -384,17 +409,17 @@ func (eng *Engine) Search(params domain.SearchParams) ([]domain.SearchResult, er params.Limit = 10 } if params.Threshold <= 0 { - params.Threshold = 1.5 + params.Threshold = embeddings.DefaultThreshold(eng.resolvedProvider()) } - vectors, err := eng.embedder.Embed([]string{params.Query}) + vector, err := eng.embedder.EmbedQuery(params.Query) if err != nil { return nil, fmt.Errorf("embed query: %w", err) } - if len(vectors) == 0 { + if len(vector) == 0 { return nil, fmt.Errorf("embedder returned no vectors") } - return eng.store.Search(vectors[0], params.Limit, params.Offset, params.Threshold) + return eng.store.Search(vector, params.Limit, params.Offset, params.Threshold) } // AddDirectory adds a directory to the store and watcher, then walks and @@ -443,6 +468,11 @@ func (eng *Engine) AddDirectory(path string) error { log.Printf("engine: walk directory %s: %v", path, walkErr) } + if !eng.tryBeginIndexWork() { + return nil // shutting down + } + defer eng.indexWG.Done() + eng.mu.Lock() eng.indexing = true eng.totalToIndex = len(filePaths) @@ -468,6 +498,11 @@ func (eng *Engine) AddDirectory(path string) error { eng.indexedFiles++ eng.mu.Unlock() } + + // The index identity (fingerprint) must exist as soon as the first index + // run completes — a fresh onboarding session's DB would otherwise carry no + // fingerprint until the next launch, leaving the read-only guard inactive. + eng.recordIndexIdentity() return nil } @@ -584,6 +619,11 @@ func (eng *Engine) initialScan() { return } + if !eng.tryBeginIndexWork() { + return // shutting down + } + defer eng.indexWG.Done() + eng.mu.Lock() eng.indexing = true eng.totalToIndex = len(filePaths) @@ -617,6 +657,54 @@ func (eng *Engine) initialScan() { } } log.Printf("engine: initial scan complete — %d files processed, %d errors", len(filePaths), errored) + + eng.recordIndexIdentity() +} + +// recordIndexIdentity persists what the index was built with so read-only +// consumers can detect a mismatch before querying sqlite-vec. The full +// fingerprint (provider:model:dimensions, per the epic) catches same-dimension +// model swaps that a bare dimension cannot — e.g. the named upgrade candidate +// (granite-97m) is also 384-dim. The bare dimension is kept alongside for +// backward compatibility with DBs written before the fingerprint existed. +func (eng *Engine) recordIndexIdentity() { + fp := embeddings.Fingerprint(eng.resolvedProvider(), eng.embedder) + if err := eng.store.SetConfig("embedding_fingerprint", fp); err != nil { + log.Printf("engine: record fingerprint: %v", err) + } + if err := eng.store.SetConfig("embedding_dimension", strconv.Itoa(eng.embedder.Dimensions())); err != nil { + log.Printf("engine: record dimension: %v", err) + } +} + +// resolvedProvider resolves the active provider through the single shared +// policy (config value, else key-implies-openai, else default) — the same rule +// the composition root and read-only search apply. +func (eng *Engine) resolvedProvider() string { + provider, _ := eng.store.GetConfig("embedding_provider") + apiKey, _ := eng.store.GetConfig("openai_api_key") + return embeddings.ResolveProvider(provider, apiKey) +} + +// tryBeginIndexWork registers in-flight indexing work with the WaitGroup that +// Stop() waits on, refusing if Stop has already been called. The check and the +// Add happen under the same mutex that Stop uses to close stopCh, so an Add +// can never race Stop's Wait — without this, a watcher debounce timer that +// fired just before watcher.Stop cancels timers could Add during/after Wait +// and the store would shut down under an active write (or trip the WaitGroup +// Add-concurrent-with-Wait panic). +func (eng *Engine) tryBeginIndexWork() bool { + eng.mu.Lock() + defer eng.mu.Unlock() + if eng.stopCh != nil { + select { + case <-eng.stopCh: + return false // Stop already called + default: + } + } + eng.indexWG.Add(1) + return true } // stopped reports whether Stop has been called (i.e. stopCh is closed). @@ -636,7 +724,10 @@ func (eng *Engine) stopped() bool { } } -// Stop stops the file watcher and cancels any in-flight indexing. +// Stop stops the file watcher, cancels any in-flight indexing, and WAITS for +// the in-flight work to reach a file boundary before returning. Callers that +// close the store next (shutdown) rely on this: without the wait, an indexing +// goroutine could still be writing while the store shuts down under it. func (eng *Engine) Stop() error { eng.mu.Lock() if eng.stopCh != nil { @@ -651,7 +742,12 @@ func (eng *Engine) Stop() error { eng.store.SetConfig("watcher_running", "false") - return eng.watcher.Stop() + err := eng.watcher.Stop() + // After stopCh is closed the scan loops exit at the next file boundary and + // the stopped watcher delivers no new events; this wait is bounded by one + // file's index time. + eng.indexWG.Wait() + return err } // Restart stops and then starts the file watcher. @@ -671,6 +767,7 @@ func (eng *Engine) Reset() error { if err := eng.store.Reset(eng.embedder.Dimensions()); err != nil { return err } + eng.recordIndexIdentity() return eng.Start() } @@ -703,6 +800,10 @@ func (eng *Engine) OnCreate(path string) { eng.logActivity(path, "ignored", "matched ignore pattern") return } + if !eng.tryBeginIndexWork() { + return // shutting down + } + defer eng.indexWG.Done() if err := eng.IndexFile(path); err != nil { log.Printf("engine: OnCreate %s: %v", path, err) } @@ -719,6 +820,10 @@ func (eng *Engine) OnModify(path string) { eng.logActivity(path, "ignored", "matched ignore pattern") return } + if !eng.tryBeginIndexWork() { + return // shutting down + } + defer eng.indexWG.Done() if err := eng.IndexFile(path); err != nil { log.Printf("engine: OnModify %s: %v", path, err) } @@ -727,7 +832,10 @@ func (eng *Engine) OnModify(path string) { // OnDelete handles file deletion events by removing the file from the index. func (eng *Engine) OnDelete(path string) { if err := eng.RemoveFileFromIndex(path); err != nil { - eng.logActivity(path, "error", fmt.Sprintf("delete: %v", err)) + // Same per-path upsert policy as index errors: one living "error" row + // per path, latest failure state wins (a path either fails to index or + // fails to delete — its most recent error is the relevant one). + eng.upsertError(path, fmt.Sprintf("delete: %v", err)) log.Printf("engine: OnDelete %s: %v", path, err) } else { eng.logActivity(path, "deleted", "") diff --git a/internal/engine/engine_test.go b/internal/engine/engine_test.go index fa1dd77..03e9028 100644 --- a/internal/engine/engine_test.go +++ b/internal/engine/engine_test.go @@ -7,6 +7,8 @@ import ( "os" "path/filepath" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -54,9 +56,9 @@ func TestIndexFile(t *testing.T) { var ( upsertedFile domain.File insertedChunks []domain.Chunk - insertedFileID int64 ) + upsertCalled := false ms := &mocks.MockStore{ GetFileByPathFn: func(path string) (*domain.File, error) { return nil, nil // new file @@ -64,29 +66,15 @@ func TestIndexFile(t *testing.T) { ListDirectoriesFn: func() ([]domain.Directory, error) { return []domain.Directory{{ID: 42, Path: dir}}, nil }, - UpsertFileFn: func(f domain.File) error { + // The atomic replace carries the file row and its chunks in one call. + UpsertFileWithChunksFn: func(f domain.File, chunks []domain.Chunk) error { + upsertCalled = true upsertedFile = f - return nil - }, - InsertChunksFn: func(fileID int64, chunks []domain.Chunk) error { - insertedFileID = fileID insertedChunks = chunks return nil }, } - // After upsert, GetFileByPath should return the file with an ID. - upsertCalled := false - ms.UpsertFileFn = func(f domain.File) error { - upsertedFile = f - upsertCalled = true - // Simulate that after upsert, store returns the file with ID. - ms.GetFileByPathFn = func(path string) (*domain.File, error) { - return &domain.File{ID: 7, DirectoryID: 42, Path: path, Hash: expectedHash, IndexedAt: time.Now()}, nil - } - return nil - } - mc := &mocks.MockChunker{ ChunkTextFn: func(c string) ([]chunker.ChunkResult, error) { return []chunker.ChunkResult{ @@ -97,7 +85,7 @@ func TestIndexFile(t *testing.T) { } me := &mocks.MockEmbedder{ - EmbedFn: func(texts []string) ([][]float32, error) { + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { vecs := make([][]float32, len(texts)) for i := range texts { vecs[i] = []float32{float32(i), 0.5} @@ -115,7 +103,7 @@ func TestIndexFile(t *testing.T) { } if !upsertCalled { - t.Fatal("UpsertFile was not called") + t.Fatal("UpsertFileWithChunks was not called") } if upsertedFile.Hash != expectedHash { t.Errorf("hash = %s, want %s", upsertedFile.Hash, expectedHash) @@ -123,9 +111,6 @@ func TestIndexFile(t *testing.T) { if upsertedFile.DirectoryID != 42 { t.Errorf("directoryID = %d, want 42", upsertedFile.DirectoryID) } - if insertedFileID != 7 { - t.Errorf("insertedFileID = %d, want 7", insertedFileID) - } if len(insertedChunks) != 2 { t.Fatalf("len(chunks) = %d, want 2", len(insertedChunks)) } @@ -188,7 +173,7 @@ func TestSearch(t *testing.T) { } me := &mocks.MockEmbedder{ - EmbedFn: func(texts []string) ([][]float32, error) { + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { if len(texts) != 1 || texts[0] != "test query" { t.Errorf("unexpected texts: %v", texts) } @@ -332,8 +317,8 @@ func TestShouldSkipDir(t *testing.T) { {".git/**", ".git", true}, {"vendor/**", "vendor", true}, {"vendor/**", "src", false}, - {".git", ".git", true}, // exact match - {"*.log", "logs", false}, // file pattern doesn't skip dirs + {".git", ".git", true}, // exact match + {"*.log", "logs", false}, // file pattern doesn't skip dirs } for _, tc := range cases { @@ -473,15 +458,11 @@ func TestAddDirectorySkipsIgnoredFiles(t *testing.T) { GetFileByPathFn: func(path string) (*domain.File, error) { return nil, nil }, - InsertChunksFn: func(fileID int64, chunks []domain.Chunk) error { return nil }, } - // Track indexed paths via UpsertFile, since IndexFile now uses ChunkText (no filePath arg). - ms.UpsertFileFn = func(f domain.File) error { + // Track indexed paths via the atomic replace call. + ms.UpsertFileWithChunksFn = func(f domain.File, chunks []domain.Chunk) error { indexedPaths = append(indexedPaths, f.Path) - ms.GetFileByPathFn = func(path string) (*domain.File, error) { - return &domain.File{ID: 1, Path: path}, nil - } return nil } @@ -492,7 +473,7 @@ func TestAddDirectorySkipsIgnoredFiles(t *testing.T) { } me := &mocks.MockEmbedder{ - EmbedFn: func(texts []string) ([][]float32, error) { + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { vecs := make([][]float32, len(texts)) for i := range texts { vecs[i] = []float32{0.1} @@ -617,3 +598,148 @@ func TestReset(t *testing.T) { t.Errorf("store.Reset dimension = %d, want 1536", resetDim) } } + +// TestIndexErrorsAreLoggedToActivityLog is a regression test for silent +// indexing failures: when IndexFile errors during AddDirectory or a watcher +// event, the failure must land in the activity log (the Log page), not just +// the invisible process stderr. Found when an embedding bug silently dropped +// 79 of 105 real-vault files with zero user-visible signal. +func TestIndexErrorsAreLoggedToActivityLog(t *testing.T) { + dir := t.TempDir() + tempFileInDir(t, dir, "doc.txt", "some content that will fail to embed") + + var logged []domain.ActivityLogEntry + ms := &mocks.MockStore{ + AddDirectoryFn: func(path string) error { return nil }, + GetConfigFn: func(key string) (string, error) { return "", nil }, + GetFileByPathFn: func(path string) (*domain.File, error) { return nil, nil }, + ListDirectoriesFn: func() ([]domain.Directory, error) { + return []domain.Directory{{ID: 1, Path: dir}}, nil + }, + // Error rows flow through the per-path upsert. + UpsertLogEntryFn: func(entry domain.ActivityLogEntry) error { + logged = append(logged, entry) + return nil + }, + } + mc := &mocks.MockChunker{ + ChunkTextFn: func(content string) ([]chunker.ChunkResult, error) { + return []chunker.ChunkResult{{Content: content, TokenCount: 5}}, nil + }, + } + me := &mocks.MockEmbedder{ + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { + return nil, fmt.Errorf("embedder exploded") + }, + } + eng := New(ms, me, mc, &mocks.MockWatcher{}, defaultMockExtractor()) + + assertErrorLogged := func(caller string) { + t.Helper() + found := 0 + for _, e := range logged { + if e.Action == "error" && strings.Contains(e.Detail, "embedder exploded") { + found++ + } + } + if found == 0 { + t.Fatalf("%s: index failure not logged to activity log; got %+v", caller, logged) + } + if found > 1 { + t.Fatalf("%s: index failure logged %d times, want exactly once (caller double-log?)", caller, found) + } + } + + // AddDirectory path. + if err := eng.AddDirectory(dir); err != nil { + t.Fatalf("AddDirectory: %v", err) + } + assertErrorLogged("AddDirectory") + + // initialScan path — the startup rescan where the motivating field + // incident (79 files silently dropped) actually happened. Synchronous + // when called directly. + logged = nil + eng.initialScan() + assertErrorLogged("initialScan") + + // Watcher-event paths. + logged = nil + eng.OnCreate(filepath.Join(dir, "doc.txt")) + assertErrorLogged("OnCreate") + logged = nil + eng.OnModify(filepath.Join(dir, "doc.txt")) + assertErrorLogged("OnModify") +} + +// TestStopWaitsForInflightIndexing pins the shutdown contract: Stop() must not +// return while a file is mid-index, so shutdown (Stop then Close) never closes +// the store under an active write. The embedder blocks until the test releases +// it; Stop() must block with it, then return only after the file's store write +// completed. +func TestStopWaitsForInflightIndexing(t *testing.T) { + dir := t.TempDir() + tempFileInDir(t, dir, "doc.txt", "content to index slowly") + + embedStarted := make(chan struct{}) + releaseEmbed := make(chan struct{}) + var wroteFile atomic.Bool + + ms := &mocks.MockStore{ + AddDirectoryFn: func(path string) error { return nil }, + GetConfigFn: func(key string) (string, error) { return "", nil }, + GetFileByPathFn: func(path string) (*domain.File, error) { return nil, nil }, + ListDirectoriesFn: func() ([]domain.Directory, error) { + return []domain.Directory{{ID: 1, Path: dir}}, nil + }, + UpsertFileWithChunksFn: func(f domain.File, chunks []domain.Chunk) error { + wroteFile.Store(true) + return nil + }, + } + mc := &mocks.MockChunker{ + ChunkTextFn: func(content string) ([]chunker.ChunkResult, error) { + return []chunker.ChunkResult{{Content: content, TokenCount: 3}}, nil + }, + } + var startOnce sync.Once + me := &mocks.MockEmbedder{ + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { + startOnce.Do(func() { close(embedStarted) }) + <-releaseEmbed // hold the file mid-index until the test releases it + return [][]float32{{0.1, 0.2}}, nil + }, + } + // No eng.Start(): AddDirectory indexes on its own, and Stop()'s wait + // contract must hold regardless of whether the watcher was started. + eng := New(ms, me, mc, &mocks.MockWatcher{}, defaultMockExtractor()) + + go func() { _ = eng.AddDirectory(dir) }() + <-embedStarted // the file is now mid-index + + stopReturned := make(chan struct{}) + go func() { + _ = eng.Stop() + close(stopReturned) + }() + + // Stop must NOT return while the file is still embedding. + select { + case <-stopReturned: + t.Fatal("Stop() returned while a file was mid-index") + case <-time.After(100 * time.Millisecond): + // good: still waiting + } + + close(releaseEmbed) // let the in-flight file finish + + select { + case <-stopReturned: + // good: Stop returned once the file boundary was reached + case <-time.After(2 * time.Second): + t.Fatal("Stop() did not return after in-flight indexing completed") + } + if !wroteFile.Load() { + t.Fatal("in-flight file's store write did not complete before Stop returned") + } +} diff --git a/internal/engine/readonly.go b/internal/engine/readonly.go index b37cea6..8b58ec1 100644 --- a/internal/engine/readonly.go +++ b/internal/engine/readonly.go @@ -3,6 +3,7 @@ package engine import ( "encoding/json" "fmt" + "strconv" "github.com/borzou/vecstore/internal/domain" "github.com/borzou/vecstore/internal/embeddings" @@ -26,18 +27,44 @@ func (ro *ReadOnlyEngine) Search(params domain.SearchParams) ([]domain.SearchRes if params.Limit <= 0 { params.Limit = 10 } + + // Resolve the provider through the SAME policy the composition root used + // to build this process's embedder (config, else key-implies-openai, else + // default). Reading the config value alone would diverge on legacy DBs + // where only an API key is set: main.go wires an OpenAI embedder while a + // config-only read here would resolve 'local' — mis-picking the threshold + // and mislabeling the fingerprint. + providerCfg, _ := ro.store.GetConfig("embedding_provider") + apiKey, _ := ro.store.GetConfig("openai_api_key") + provider := embeddings.ResolveProvider(providerCfg, apiKey) if params.Threshold <= 0 { - params.Threshold = 1.5 + params.Threshold = embeddings.DefaultThreshold(provider) + } + + // Guard against querying an index built with a different embedding model + // than the one this read-only process is configured with. The full + // fingerprint (provider:model:dimensions, per the epic) catches + // same-dimension model swaps that the bare dimension cannot; the dimension + // check remains as fallback for DBs written before the fingerprint existed. + ownFP := embeddings.Fingerprint(provider, ro.embedder) + if indexFP, _ := ro.store.GetConfig("embedding_fingerprint"); indexFP != "" { + if indexFP != ownFP { + return nil, fmt.Errorf("index was built with embedding %q but this process is configured for %q — mixed vectors would return garbage-ranked results; reopen the GUI app to rebuild the index", indexFP, ownFP) + } + } else if dimStr, _ := ro.store.GetConfig("embedding_dimension"); dimStr != "" { + if indexDim, convErr := strconv.Atoi(dimStr); convErr == nil && indexDim != ro.embedder.Dimensions() { + return nil, fmt.Errorf("index was built with a different embedding model (dim %d) than the active provider (dim %d) — reopen the GUI app to rebuild the index", indexDim, ro.embedder.Dimensions()) + } } - vectors, err := ro.embedder.Embed([]string{params.Query}) + vector, err := ro.embedder.EmbedQuery(params.Query) if err != nil { return nil, fmt.Errorf("embed query: %w", err) } - if len(vectors) == 0 { + if len(vector) == 0 { return nil, fmt.Errorf("embedder returned no vectors") } - return ro.store.Search(vectors[0], params.Limit, params.Offset, params.Threshold) + return ro.store.Search(vector, params.Limit, params.Offset, params.Threshold) } // ListDirectories returns all watched directories from the store. diff --git a/internal/engine/readonly_test.go b/internal/engine/readonly_test.go index 877de8e..5378b4d 100644 --- a/internal/engine/readonly_test.go +++ b/internal/engine/readonly_test.go @@ -18,15 +18,17 @@ func TestReadOnlySearch(t *testing.T) { if limit != 10 { t.Errorf("expected default limit 10, got %d", limit) } - if threshold != 1.5 { - t.Errorf("expected default threshold 1.5, got %f", threshold) + // No provider configured → falls back to the local default provider, + // whose default threshold is 0.6. + if threshold != 0.6 { + t.Errorf("expected default threshold 0.6, got %f", threshold) } return expected, nil }, } me := &mocks.MockEmbedder{ - EmbedFn: func(texts []string) ([][]float32, error) { + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { return [][]float32{{1.0, 2.0}}, nil }, } @@ -59,7 +61,7 @@ func TestReadOnlySearch_ExplicitParams(t *testing.T) { } me := &mocks.MockEmbedder{ - EmbedFn: func(texts []string) ([][]float32, error) { + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { return [][]float32{{1.0}}, nil }, } @@ -71,10 +73,94 @@ func TestReadOnlySearch_ExplicitParams(t *testing.T) { } } +func TestReadOnlySearch_ProviderAwareThreshold(t *testing.T) { + ms := &mocks.MockStore{ + GetConfigFn: func(key string) (string, error) { + if key == "embedding_provider" { + return "openai", nil + } + return "", nil + }, + SearchFn: func(embedding []float32, limit, offset int, threshold float32) ([]domain.SearchResult, error) { + // OpenAI provider default threshold is 1.5. + if threshold != 1.5 { + t.Errorf("expected openai default threshold 1.5, got %f", threshold) + } + return nil, nil + }, + } + me := &mocks.MockEmbedder{ + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { + return [][]float32{{1.0}}, nil + }, + } + + ro := NewReadOnly(ms, me) + if _, err := ro.Search(domain.SearchParams{Query: "test"}); err != nil { + t.Fatalf("Search: %v", err) + } +} + +func TestReadOnlySearch_DimensionMismatch(t *testing.T) { + searchCalled := false + ms := &mocks.MockStore{ + GetConfigFn: func(key string) (string, error) { + if key == "embedding_dimension" { + return "1536", nil // index built with a 1536-dim model + } + return "", nil + }, + SearchFn: func(embedding []float32, limit, offset int, threshold float32) ([]domain.SearchResult, error) { + searchCalled = true + return nil, nil + }, + } + me := &mocks.MockEmbedder{ + DimensionsFn: func() int { return 384 }, // active provider is 384-dim + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { + return [][]float32{{1.0}}, nil + }, + } + + ro := NewReadOnly(ms, me) + _, err := ro.Search(domain.SearchParams{Query: "test"}) + if err == nil { + t.Fatal("expected dimension-mismatch error, got nil") + } + if searchCalled { + t.Error("store.Search must not be called on a dimension mismatch") + } +} + +func TestReadOnlySearch_DimensionMatch(t *testing.T) { + ms := &mocks.MockStore{ + GetConfigFn: func(key string) (string, error) { + if key == "embedding_dimension" { + return "384", nil + } + return "", nil + }, + SearchFn: func(embedding []float32, limit, offset int, threshold float32) ([]domain.SearchResult, error) { + return nil, nil + }, + } + me := &mocks.MockEmbedder{ + DimensionsFn: func() int { return 384 }, + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { + return [][]float32{{1.0}}, nil + }, + } + + ro := NewReadOnly(ms, me) + if _, err := ro.Search(domain.SearchParams{Query: "test"}); err != nil { + t.Fatalf("Search with matching dimension should succeed: %v", err) + } +} + func TestReadOnlySearch_EmbedderError(t *testing.T) { ms := &mocks.MockStore{} me := &mocks.MockEmbedder{ - EmbedFn: func(texts []string) ([][]float32, error) { + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { return nil, fmt.Errorf("no API key configured") }, } @@ -163,3 +249,71 @@ func TestReadOnlyGetIgnorePatterns_Defaults(t *testing.T) { t.Errorf("expected default patterns, got %d", len(patterns)) } } + +// TestReadOnlySearch_FingerprintMismatchSameDimension pins the case the bare +// dimension guard is blind to: a same-width model swap (e.g. the epic's named +// upgrade candidate granite-97m is also 384-dim). Mixed vectors would return +// garbage-ranked results silently. +func TestReadOnlySearch_FingerprintMismatchSameDimension(t *testing.T) { + searchCalled := false + ms := &mocks.MockStore{ + GetConfigFn: func(key string) (string, error) { + switch key { + case "embedding_fingerprint": + return "local:multilingual-e5-small:384", nil // index identity + case "embedding_provider": + return "local", nil + } + return "", nil + }, + SearchFn: func(embedding []float32, limit, offset int, threshold float32) ([]domain.SearchResult, error) { + searchCalled = true + return nil, nil + }, + } + me := &mocks.MockEmbedder{ + DimensionsFn: func() int { return 384 }, // SAME dimension... + ModelNameFn: func() string { return "granite-embedding-97m" }, // ...different model + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { + return [][]float32{{1.0}}, nil + }, + } + + ro := NewReadOnly(ms, me) + _, err := ro.Search(domain.SearchParams{Query: "test"}) + if err == nil { + t.Fatal("expected fingerprint-mismatch error for same-dimension model swap, got nil") + } + if searchCalled { + t.Error("store.Search must not be called on a fingerprint mismatch") + } +} + +// TestReadOnlySearch_FingerprintMatch: matching fingerprints search normally. +func TestReadOnlySearch_FingerprintMatch(t *testing.T) { + ms := &mocks.MockStore{ + GetConfigFn: func(key string) (string, error) { + switch key { + case "embedding_fingerprint": + return "local:multilingual-e5-small:384", nil + case "embedding_provider": + return "local", nil + } + return "", nil + }, + SearchFn: func(embedding []float32, limit, offset int, threshold float32) ([]domain.SearchResult, error) { + return []domain.SearchResult{}, nil + }, + } + me := &mocks.MockEmbedder{ + DimensionsFn: func() int { return 384 }, + ModelNameFn: func() string { return "multilingual-e5-small" }, + EmbedDocumentsFn: func(texts []string) ([][]float32, error) { + return [][]float32{{1.0}}, nil + }, + } + ro := NewReadOnly(ms, me) + if _, err := ro.Search(domain.SearchParams{Query: "test"}); err != nil { + t.Fatalf("matching fingerprint should search cleanly: %v", err) + } +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 7ff11a9..3c6984f 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -331,7 +331,7 @@ func getToolDefinitions() []toolDefinition { }, "threshold": map[string]interface{}{ "type": "number", - "description": "Maximum cosine distance for results. Lower values mean stricter matching. Default: 1.5. Set to 0 to disable filtering.", + "description": "Maximum cosine distance for results. Lower values mean stricter matching. The default is provider-dependent (tuned per embedding provider). Set to 0 to use the provider default; use a positive value to override it.", }, }, "required": []string{"query"}, @@ -375,7 +375,7 @@ func getToolDefinitions() []toolDefinition { }, { Name: "index_status", - Description: "Returns total files, total chunks, last indexed timestamp, currently indexing flag, and embedding model in use.", + Description: "Returns total files, total chunks, last indexed timestamp, currently indexing flag, and the active embedding provider and model in use.", InputSchema: map[string]interface{}{ "type": "object", "properties": map[string]interface{}{}, @@ -504,7 +504,7 @@ func getReadOnlyToolDefinitions() []toolDefinition { }, "threshold": map[string]interface{}{ "type": "number", - "description": "Maximum cosine distance for results. Lower values mean stricter matching. Default: 1.5. Set to 0 to disable filtering.", + "description": "Maximum cosine distance for results. Lower values mean stricter matching. The default is provider-dependent (tuned per embedding provider). Set to 0 to use the provider default; use a positive value to override it.", }, }, "required": []string{"query"}, @@ -520,7 +520,7 @@ func getReadOnlyToolDefinitions() []toolDefinition { }, { Name: "index_status", - Description: "Returns total files, total chunks, last indexed timestamp, currently indexing flag, and embedding model in use.", + Description: "Returns total files, total chunks, last indexed timestamp, currently indexing flag, and the active embedding provider and model in use.", InputSchema: map[string]interface{}{ "type": "object", "properties": map[string]interface{}{}, diff --git a/internal/mocks/mocks.go b/internal/mocks/mocks.go index 492f3b6..f637bb5 100644 --- a/internal/mocks/mocks.go +++ b/internal/mocks/mocks.go @@ -14,22 +14,22 @@ import ( // --------------------------------------------------------------------------- type MockStore struct { - GetConfigFn func(key string) (string, error) - SetConfigFn func(key, value string) error - AddDirectoryFn func(path string) error - RemoveDirectoryFn func(path string) error - ListDirectoriesFn func() ([]domain.Directory, error) - UpsertFileFn func(f domain.File) error - RemoveFileFn func(path string) error - GetFileByPathFn func(path string) (*domain.File, error) - InsertChunksFn func(fileID int64, chunks []domain.Chunk) error - RemoveChunksByFileFn func(fileID int64) error - SearchFn func(embedding []float32, limit, offset int, threshold float32) ([]domain.SearchResult, error) - StatsFn func() (domain.IndexStats, error) - InsertLogEntryFn func(entry domain.ActivityLogEntry) error - ListLogEntriesFn func(limit, offset int) ([]domain.ActivityLogEntry, int, error) - ResetFn func(embeddingDimension int) error - CloseFn func() error + GetConfigFn func(key string) (string, error) + SetConfigFn func(key, value string) error + AddDirectoryFn func(path string) error + RemoveDirectoryFn func(path string) error + ListDirectoriesFn func() ([]domain.Directory, error) + RemoveFileFn func(path string) error + GetFileByPathFn func(path string) (*domain.File, error) + UpsertFileWithChunksFn func(f domain.File, chunks []domain.Chunk) error + UpsertLogEntryFn func(entry domain.ActivityLogEntry) error + RemoveChunksByFileFn func(fileID int64) error + SearchFn func(embedding []float32, limit, offset int, threshold float32) ([]domain.SearchResult, error) + StatsFn func() (domain.IndexStats, error) + InsertLogEntryFn func(entry domain.ActivityLogEntry) error + ListLogEntriesFn func(limit, offset int) ([]domain.ActivityLogEntry, int, error) + ResetFn func(embeddingDimension int) error + CloseFn func() error } func (m *MockStore) GetConfig(key string) (string, error) { @@ -67,9 +67,16 @@ func (m *MockStore) ListDirectories() ([]domain.Directory, error) { return nil, nil } -func (m *MockStore) UpsertFile(f domain.File) error { - if m.UpsertFileFn != nil { - return m.UpsertFileFn(f) +func (m *MockStore) UpsertFileWithChunks(f domain.File, chunks []domain.Chunk) error { + if m.UpsertFileWithChunksFn != nil { + return m.UpsertFileWithChunksFn(f, chunks) + } + return nil +} + +func (m *MockStore) UpsertLogEntry(entry domain.ActivityLogEntry) error { + if m.UpsertLogEntryFn != nil { + return m.UpsertLogEntryFn(entry) } return nil } @@ -88,13 +95,6 @@ func (m *MockStore) GetFileByPath(path string) (*domain.File, error) { return nil, nil } -func (m *MockStore) InsertChunks(fileID int64, chunks []domain.Chunk) error { - if m.InsertChunksFn != nil { - return m.InsertChunksFn(fileID, chunks) - } - return nil -} - func (m *MockStore) RemoveChunksByFile(fileID int64) error { if m.RemoveChunksByFileFn != nil { return m.RemoveChunksByFileFn(fileID) @@ -149,14 +149,34 @@ func (m *MockStore) Close() error { // --------------------------------------------------------------------------- type MockEmbedder struct { - EmbedFn func(texts []string) ([][]float32, error) - DimensionsFn func() int - ModelNameFn func() string + EmbedDocumentsFn func(texts []string) ([][]float32, error) + EmbedQueryFn func(text string) ([]float32, error) + DimensionsFn func() int + ModelNameFn func() string + MaxInputTokensFn func() int } -func (m *MockEmbedder) Embed(texts []string) ([][]float32, error) { - if m.EmbedFn != nil { - return m.EmbedFn(texts) +func (m *MockEmbedder) EmbedDocuments(texts []string) ([][]float32, error) { + if m.EmbedDocumentsFn != nil { + return m.EmbedDocumentsFn(texts) + } + return nil, nil +} + +func (m *MockEmbedder) EmbedQuery(text string) ([]float32, error) { + if m.EmbedQueryFn != nil { + return m.EmbedQueryFn(text) + } + // Fallback: reuse EmbedDocumentsFn so search tests don't need a separate stub. + if m.EmbedDocumentsFn != nil { + vecs, err := m.EmbedDocumentsFn([]string{text}) + if err != nil { + return nil, err + } + if len(vecs) == 0 { + return nil, nil + } + return vecs[0], nil } return nil, nil } @@ -175,6 +195,13 @@ func (m *MockEmbedder) ModelName() string { return "" } +func (m *MockEmbedder) MaxInputTokens() int { + if m.MaxInputTokensFn != nil { + return m.MaxInputTokensFn() + } + return 0 +} + // --------------------------------------------------------------------------- // MockChunker implements chunker.Chunker // --------------------------------------------------------------------------- diff --git a/internal/store/iface.go b/internal/store/iface.go index aa1d61d..ec57d1f 100644 --- a/internal/store/iface.go +++ b/internal/store/iface.go @@ -9,14 +9,25 @@ type Store interface { AddDirectory(path string) error RemoveDirectory(path string) error ListDirectories() ([]domain.Directory, error) - UpsertFile(f domain.File) error RemoveFile(path string) error GetFileByPath(path string) (*domain.File, error) - InsertChunks(fileID int64, chunks []domain.Chunk) error RemoveChunksByFile(fileID int64) error + // UpsertFileWithChunks atomically replaces a file's index entry: old chunks + // (and their vectors) are removed, the file row is upserted, and the new + // chunks are inserted — all in one transaction, so a crash mid-index leaves + // the file either fully indexed or untouched-and-retryable, never recorded + // at the new hash with missing chunks. (The former separate UpsertFile / + // InsertChunks steps live on as concrete SQLiteStore methods for tests but + // are no longer part of the engine's contract.) + UpsertFileWithChunks(f domain.File, chunks []domain.Chunk) error Search(embedding []float32, limit, offset int, threshold float32) ([]domain.SearchResult, error) Stats() (domain.IndexStats, error) InsertLogEntry(entry domain.ActivityLogEntry) error + // UpsertLogEntry keeps at most one row per (path, action): if one exists + // its timestamp and detail are updated in place, otherwise the entry is + // inserted. Used for error rows so a persistently-failing file yields one + // living row instead of an identical append on every launch. + UpsertLogEntry(entry domain.ActivityLogEntry) error ListLogEntries(limit, offset int) ([]domain.ActivityLogEntry, int, error) Reset(embeddingDimension int) error Close() error diff --git a/internal/store/sqlite.go b/internal/store/sqlite.go index ab46f51..a228e80 100644 --- a/internal/store/sqlite.go +++ b/internal/store/sqlite.go @@ -4,6 +4,7 @@ import ( "database/sql" "encoding/binary" "fmt" + "log" "math" "time" @@ -13,13 +14,26 @@ import ( "github.com/borzou/vecstore/internal/domain" ) +// defaultVecDimension is the fallback embedding dimension used when a caller +// does not supply one (dim <= 0). It matches OpenAI text-embedding-3-small, +// preserving the historical schema for databases created before dimensions +// were provider-driven. +const defaultVecDimension = 1536 + // SQLiteStore implements Store using SQLite + sqlite-vec. type SQLiteStore struct { db *sql.DB } -// NewSQLiteStore opens (or creates) a SQLite database at dbPath and initializes the schema. -func NewSQLiteStore(dbPath string) (*SQLiteStore, error) { +// NewSQLiteStore opens (or creates) a SQLite database at dbPath and initializes +// the schema. dim sets the width of the vector table for a freshly created +// database; a value <= 0 falls back to defaultVecDimension. Existing databases +// are unaffected — the vector table is created with CREATE ... IF NOT EXISTS, so +// the stored dimension always wins for an already-migrated DB. +func NewSQLiteStore(dbPath string, dim int) (*SQLiteStore, error) { + if dim <= 0 { + dim = defaultVecDimension + } sqlite_vec.Auto() db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=on&_busy_timeout=5000") if err != nil { @@ -31,14 +45,42 @@ func NewSQLiteStore(dbPath string) (*SQLiteStore, error) { db.SetMaxOpenConns(1) s := &SQLiteStore{db: db} - if err := s.migrate(); err != nil { + if err := s.migrate(dim); err != nil { db.Close() return nil, fmt.Errorf("migrate: %w", err) } + if err := s.pruneActivityLog(); err != nil { + // Retention is hygiene, not correctness — log-worthy upstream but must + // never block opening the store. + log.Printf("store: prune activity log: %v", err) + } return s, nil } -func (s *SQLiteStore) migrate() error { +// Activity-log retention: entries older than the TTL are dropped, and the +// table is capped to the newest logRetentionMaxRows. Without this the table +// grows without bound (the Log page pays COUNT(*) over it every poll, on the +// single connection, in contention with indexing writes). +const ( + logRetentionDays = 30 + logRetentionMaxRows = 5000 +) + +// pruneActivityLog applies the TTL and row cap. Called at store open. +func (s *SQLiteStore) pruneActivityLog() error { + cutoff := time.Now().UTC().AddDate(0, 0, -logRetentionDays) + if _, err := s.db.Exec(`DELETE FROM activity_log WHERE timestamp < ?`, cutoff); err != nil { + return err + } + _, err := s.db.Exec( + `DELETE FROM activity_log WHERE id NOT IN + (SELECT id FROM activity_log ORDER BY timestamp DESC LIMIT ?)`, + logRetentionMaxRows, + ) + return err +} + +func (s *SQLiteStore) migrate(dim int) error { stmts := []string{ `CREATE TABLE IF NOT EXISTS config ( key TEXT PRIMARY KEY, @@ -68,10 +110,10 @@ func (s *SQLiteStore) migrate() error { token_count INTEGER NOT NULL DEFAULT 0, FOREIGN KEY(file_id) REFERENCES files(id) ON DELETE CASCADE )`, - `CREATE VIRTUAL TABLE IF NOT EXISTS chunk_embeddings USING vec0( + fmt.Sprintf(`CREATE VIRTUAL TABLE IF NOT EXISTS chunk_embeddings USING vec0( chunk_id INTEGER PRIMARY KEY, - embedding FLOAT[1536] - )`, + embedding FLOAT[%d] + )`, dim), `CREATE TABLE IF NOT EXISTS activity_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, timestamp DATETIME NOT NULL, @@ -221,6 +263,79 @@ func (s *SQLiteStore) InsertChunks(fileID int64, chunks []domain.Chunk) error { return tx.Commit() } +// UpsertFileWithChunks atomically replaces a file's index entry in a single +// transaction: delete the file's old chunks + vectors, upsert the file row, +// insert the new chunks + vectors. Atomicity is the crash-safety guarantee for +// indexing: without it, a process death after the file row is written but +// before its chunks land records the file as indexed-at-hash with no content, +// and the hash short-circuit then skips it forever. +func (s *SQLiteStore) UpsertFileWithChunks(f domain.File, chunks []domain.Chunk) error { + tx, err := s.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + // Remove existing chunks + vectors (no-op for a brand-new file). Single + // statements — a per-chunk DELETE loop would hold the write lock for + // hundreds of round-trips on the one connection exactly when the watcher + // is busiest. + if f.ID != 0 { + if _, err := tx.Exec( + `DELETE FROM chunk_embeddings WHERE chunk_id IN (SELECT id FROM chunks WHERE file_id = ?)`, + f.ID, + ); err != nil { + return err + } + if _, err := tx.Exec(`DELETE FROM chunks WHERE file_id = ?`, f.ID); err != nil { + return err + } + } + + // Upsert the file row and resolve its (possibly new) ID within the tx. + res, err := tx.Exec( + `INSERT OR REPLACE INTO files(directory_id, path, hash, indexed_at) VALUES(?, ?, ?, ?)`, + f.DirectoryID, f.Path, f.Hash, f.IndexedAt.UTC(), + ) + if err != nil { + return fmt.Errorf("upsert file: %w", err) + } + fileID, err := res.LastInsertId() + if err != nil { + return err + } + + stmtChunk, err := tx.Prepare(`INSERT INTO chunks(file_id, chunk_index, content, token_count) VALUES(?, ?, ?, ?)`) + if err != nil { + return err + } + defer stmtChunk.Close() + stmtVec, err := tx.Prepare(`INSERT INTO chunk_embeddings(chunk_id, embedding) VALUES(?, ?)`) + if err != nil { + return err + } + defer stmtVec.Close() + + for _, c := range chunks { + res, err := stmtChunk.Exec(fileID, c.Index, c.Content, c.TokenCount) + if err != nil { + return fmt.Errorf("insert chunk: %w", err) + } + chunkID, err := res.LastInsertId() + if err != nil { + return err + } + if len(c.Embedding) > 0 { + blob := float32SliceToBlob(c.Embedding) + if _, err := stmtVec.Exec(chunkID, blob); err != nil { + return fmt.Errorf("insert embedding: %w", err) + } + } + } + + return tx.Commit() +} + func (s *SQLiteStore) RemoveChunksByFile(fileID int64) error { // Get chunk IDs first to remove from vec table. rows, err := s.db.Query(`SELECT id FROM chunks WHERE file_id = ?`, fileID) @@ -288,7 +403,9 @@ func (s *SQLiteStore) Search(embedding []float32, limit, offset int, threshold f } defer rows.Close() - var all []domain.SearchResult + // Non-nil so an empty result set serializes as [] (not null) all the way + // out through the MCP transports — strict JSON clients index into it. + all := make([]domain.SearchResult, 0, fetchLimit) for rows.Next() { var r domain.SearchResult if err := rows.Scan(&r.ChunkIndex, &r.Content, &r.FilePath, &r.Score); err != nil { @@ -307,7 +424,7 @@ func (s *SQLiteStore) Search(embedding []float32, limit, offset int, threshold f if offset > 0 && offset < len(all) { all = all[offset:] } else if offset >= len(all) { - return nil, nil + return []domain.SearchResult{}, nil } // Apply limit. @@ -338,11 +455,11 @@ func (s *SQLiteStore) Stats() (domain.IndexStats, error) { stats.LastIndexedAt = parseTimestamp(lastIndexed.String) } - model, _ := s.GetConfig("embedding_model") - if model == "" { - model = "text-embedding-3-small" - } - stats.EmbeddingModel = model + // Report the active provider/model straight from config. When config is + // empty (never configured) the fields are left empty rather than assuming a + // specific default, since the composition root owns provider selection. + stats.Provider, _ = s.GetConfig("embedding_provider") + stats.EmbeddingModel, _ = s.GetConfig("embedding_model") return stats, nil } @@ -357,6 +474,32 @@ func (s *SQLiteStore) InsertLogEntry(entry domain.ActivityLogEntry) error { return err } +// UpsertLogEntry replaces all existing (path, action) rows with this single +// entry, in one transaction — one living row per failing path instead of an +// identical append per launch. Delete-then-insert (rather than UPDATE) also +// collapses duplicate rows accumulated by pre-upsert builds the first time a +// path fails again after upgrading. +func (s *SQLiteStore) UpsertLogEntry(entry domain.ActivityLogEntry) error { + tx, err := s.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + if _, err := tx.Exec( + `DELETE FROM activity_log WHERE path = ? AND action = ?`, + entry.Path, entry.Action, + ); err != nil { + return err + } + if _, err := tx.Exec( + `INSERT INTO activity_log(timestamp, path, action, detail) VALUES(?, ?, ?, ?)`, + entry.Timestamp.UTC(), entry.Path, entry.Action, entry.Detail, + ); err != nil { + return err + } + return tx.Commit() +} + func (s *SQLiteStore) ListLogEntries(limit, offset int) ([]domain.ActivityLogEntry, int, error) { var total int if err := s.db.QueryRow(`SELECT COUNT(*) FROM activity_log`).Scan(&total); err != nil { @@ -392,7 +535,7 @@ func (s *SQLiteStore) ListLogEntries(limit, offset int) ([]domain.ActivityLogEnt // directories are preserved so the user doesn't have to re-onboard. func (s *SQLiteStore) Reset(embeddingDimension int) error { if embeddingDimension <= 0 { - embeddingDimension = 1536 + embeddingDimension = defaultVecDimension } stmts := []string{ diff --git a/internal/store/sqlite_readonly_test.go b/internal/store/sqlite_readonly_test.go index 378c6dc..f0ff2d6 100644 --- a/internal/store/sqlite_readonly_test.go +++ b/internal/store/sqlite_readonly_test.go @@ -11,7 +11,7 @@ func TestReadOnlyStore_OpensInitializedDB(t *testing.T) { dbPath := filepath.Join(dir, "test.db") // Create and initialize with the read-write constructor. - rw, err := NewSQLiteStore(dbPath) + rw, err := NewSQLiteStore(dbPath, 1536) if err != nil { t.Fatalf("NewSQLiteStore: %v", err) } diff --git a/internal/store/sqlite_test.go b/internal/store/sqlite_test.go index a101242..0a2ec70 100644 --- a/internal/store/sqlite_test.go +++ b/internal/store/sqlite_test.go @@ -1,7 +1,9 @@ package store import ( + "encoding/json" "path/filepath" + "strings" "testing" "time" @@ -11,7 +13,7 @@ import ( func newTestStore(t *testing.T) *SQLiteStore { t.Helper() dir := t.TempDir() - s, err := NewSQLiteStore(filepath.Join(dir, "test.db")) + s, err := NewSQLiteStore(filepath.Join(dir, "test.db"), 1536) if err != nil { t.Fatalf("NewSQLiteStore: %v", err) } @@ -226,6 +228,20 @@ func TestStats(t *testing.T) { if stats.TotalChunks != 1 { t.Fatalf("want 1 chunk, got %d", stats.TotalChunks) } + + // Provider/model are reported straight from config; empty when unset. + if stats.Provider != "" || stats.EmbeddingModel != "" { + t.Fatalf("want empty provider/model when unconfigured, got %q/%q", stats.Provider, stats.EmbeddingModel) + } + s.SetConfig("embedding_provider", "local") + s.SetConfig("embedding_model", "multilingual-e5-small") + stats, _ = s.Stats() + if stats.Provider != "local" { + t.Fatalf("want provider 'local', got %q", stats.Provider) + } + if stats.EmbeddingModel != "multilingual-e5-small" { + t.Fatalf("want model 'multilingual-e5-small', got %q", stats.EmbeddingModel) + } } func TestReset(t *testing.T) { @@ -249,6 +265,48 @@ func TestReset(t *testing.T) { } } +func TestNewSQLiteStoreDimension(t *testing.T) { + dir := t.TempDir() + dbPath := filepath.Join(dir, "dim.db") + + // Fresh DB created at 384 dims should accept a 384-wide embedding. + s, err := NewSQLiteStore(dbPath, 384) + if err != nil { + t.Fatalf("NewSQLiteStore(384): %v", err) + } + s.AddDirectory("/tmp/a") + dirs, _ := s.ListDirectories() + f := domain.File{DirectoryID: dirs[0].ID, Path: "/tmp/a/f.txt", Hash: "h", IndexedAt: time.Now().UTC()} + s.UpsertFile(f) + got, _ := s.GetFileByPath("/tmp/a/f.txt") + + emb := make([]float32, 384) + emb[0] = 1 + if err := s.InsertChunks(got.ID, []domain.Chunk{{Index: 0, Content: "x", TokenCount: 1, Embedding: emb}}); err != nil { + t.Fatalf("InsertChunks(384): %v", err) + } + s.Close() + + // Reopening with a different dim must NOT change the existing table. + s2, err := NewSQLiteStore(dbPath, 1536) + if err != nil { + t.Fatalf("reopen: %v", err) + } + defer s2.Close() + stats, _ := s2.Stats() + if stats.TotalChunks != 1 { + t.Fatalf("existing 384-dim data should survive reopen, got %d chunks", stats.TotalChunks) + } + // A 384-wide query still matches — table width was preserved as 384. + res, err := s2.Search(emb, 5, 0, 0) + if err != nil { + t.Fatalf("Search on preserved 384 table: %v", err) + } + if len(res) == 0 { + t.Fatal("expected a result from the preserved 384-dim table") + } +} + func TestSearch(t *testing.T) { s := newTestStore(t) s.AddDirectory("/tmp/a") @@ -287,3 +345,177 @@ func TestSearch(t *testing.T) { t.Fatalf("unexpected content: %q", results[0].Content) } } + +// TestUpsertFileWithChunksAtomic pins the crash-safety contract: the +// remove-old-chunks / upsert-file / insert-chunks sequence is one transaction, +// so a failure partway through leaves the previous index entry fully intact — +// never the file recorded at the new hash with missing chunks (which the hash +// short-circuit would then skip forever). +func TestUpsertFileWithChunksAtomic(t *testing.T) { + s := newTestStore(t) + if err := s.AddDirectory("/tmp/a"); err != nil { + t.Fatal(err) + } + dirs, _ := s.ListDirectories() + + emb := make([]float32, 1536) + emb[0] = 0.5 + + // Index v1 successfully. + v1 := domain.File{DirectoryID: dirs[0].ID, Path: "/tmp/a/doc.txt", Hash: "hash-v1", IndexedAt: time.Now().UTC()} + if err := s.UpsertFileWithChunks(v1, []domain.Chunk{ + {Index: 0, Content: "v1 chunk zero", TokenCount: 3, Embedding: emb}, + {Index: 1, Content: "v1 chunk one", TokenCount: 3, Embedding: emb}, + }); err != nil { + t.Fatalf("v1 index: %v", err) + } + stored, _ := s.GetFileByPath("/tmp/a/doc.txt") + if stored == nil || stored.Hash != "hash-v1" { + t.Fatalf("v1 not stored correctly: %+v", stored) + } + + // Attempt v2 with a chunk whose embedding has the WRONG dimension — the + // vec0 insert fails mid-transaction. Everything must roll back. + v2 := domain.File{ID: stored.ID, DirectoryID: dirs[0].ID, Path: "/tmp/a/doc.txt", Hash: "hash-v2", IndexedAt: time.Now().UTC()} + badEmb := []float32{1, 2, 3} // store was created with dim 1536 + if err := s.UpsertFileWithChunks(v2, []domain.Chunk{ + {Index: 0, Content: "v2 chunk zero", TokenCount: 3, Embedding: emb}, + {Index: 1, Content: "v2 chunk one", TokenCount: 3, Embedding: badEmb}, + }); err == nil { + t.Fatal("expected wrong-dimension embedding to fail the transaction") + } + + // The file must still be recorded at hash-v1 with BOTH v1 chunks intact. + after, _ := s.GetFileByPath("/tmp/a/doc.txt") + if after == nil { + t.Fatal("file row vanished after failed update") + } + if after.Hash != "hash-v1" { + t.Fatalf("hash = %q after failed update, want hash-v1 (partial write leaked!)", after.Hash) + } + results, err := s.Search(emb, 10, 0, 0) + if err != nil { + t.Fatalf("search after rollback: %v", err) + } + v1Chunks := 0 + for _, r := range results { + if r.FilePath == "/tmp/a/doc.txt" && strings.HasPrefix(r.Content, "v1 ") { + v1Chunks++ + } + } + if v1Chunks != 2 { + t.Fatalf("searchable v1 chunks after rollback = %d, want 2", v1Chunks) + } + + // A good v2 then replaces v1 completely. + if err := s.UpsertFileWithChunks(v2, []domain.Chunk{ + {Index: 0, Content: "v2 only chunk", TokenCount: 3, Embedding: emb}, + }); err != nil { + t.Fatalf("good v2 index: %v", err) + } + final, _ := s.GetFileByPath("/tmp/a/doc.txt") + if final.Hash != "hash-v2" { + t.Fatalf("hash = %q, want hash-v2", final.Hash) + } + results, _ = s.Search(emb, 10, 0, 0) + for _, r := range results { + if r.FilePath == "/tmp/a/doc.txt" && strings.HasPrefix(r.Content, "v1 ") { + t.Fatalf("stale v1 chunk still searchable after replace: %q", r.Content) + } + } +} + +// TestUpsertLogEntryDedupes pins the one-living-row-per-(path,action) +// contract: repeat failures update the existing row instead of appending. +func TestUpsertLogEntryDedupes(t *testing.T) { + s := newTestStore(t) + e := domain.ActivityLogEntry{Timestamp: time.Now(), Path: "/a/b.pdf", Action: "error", Detail: "index: boom v1"} + if err := s.UpsertLogEntry(e); err != nil { + t.Fatal(err) + } + e.Detail = "index: boom v2" + e.Timestamp = time.Now().Add(time.Minute) + if err := s.UpsertLogEntry(e); err != nil { + t.Fatal(err) + } + entries, total, err := s.ListLogEntries(10, 0) + if err != nil { + t.Fatal(err) + } + if total != 1 || len(entries) != 1 { + t.Fatalf("total = %d entries = %d, want exactly 1 row", total, len(entries)) + } + if entries[0].Detail != "index: boom v2" { + t.Fatalf("detail = %q, want the updated v2 detail", entries[0].Detail) + } + // A different path appends normally. + e2 := domain.ActivityLogEntry{Timestamp: time.Now(), Path: "/a/c.pdf", Action: "error", Detail: "index: other"} + if err := s.UpsertLogEntry(e2); err != nil { + t.Fatal(err) + } + if _, total, _ = s.ListLogEntries(10, 0); total != 2 { + t.Fatalf("total = %d, want 2 after a second distinct path", total) + } +} + +// TestActivityLogRetention pins the TTL prune at store open. +func TestActivityLogRetention(t *testing.T) { + dir := t.TempDir() + dbPath := filepath.Join(dir, "retention.db") + s, err := NewSQLiteStore(dbPath, 8) + if err != nil { + t.Fatal(err) + } + old := domain.ActivityLogEntry{Timestamp: time.Now().AddDate(0, 0, -60), Path: "/old.txt", Action: "indexed", Detail: "1 chunks"} + fresh := domain.ActivityLogEntry{Timestamp: time.Now(), Path: "/fresh.txt", Action: "indexed", Detail: "1 chunks"} + if err := s.InsertLogEntry(old); err != nil { + t.Fatal(err) + } + if err := s.InsertLogEntry(fresh); err != nil { + t.Fatal(err) + } + s.Close() + + // Re-open: the 60-day-old row must be pruned, the fresh one kept. + s2, err := NewSQLiteStore(dbPath, 8) + if err != nil { + t.Fatal(err) + } + defer s2.Close() + entries, total, err := s2.ListLogEntries(10, 0) + if err != nil { + t.Fatal(err) + } + if total != 1 || entries[0].Path != "/fresh.txt" { + t.Fatalf("after reopen: total = %d first = %+v, want only /fresh.txt", total, entries) + } +} + +// TestSearchEmptyResultsMarshalsToArray pins the wire contract: an empty +// result set must serialize as [] (not null) — strict JSON clients call +// .length on it. Covers both the no-matches and offset-past-results branches. +func TestSearchEmptyResultsMarshalsToArray(t *testing.T) { + s := newTestStore(t) + q := make([]float32, 1536) + q[0] = 1 + + for name, fn := range map[string]func() ([]domain.SearchResult, error){ + "no matches": func() ([]domain.SearchResult, error) { return s.Search(q, 5, 0, 0) }, + "offset past results": func() ([]domain.SearchResult, error) { return s.Search(q, 5, 100, 0) }, + } { + results, err := fn() + if err != nil { + t.Fatalf("%s: %v", name, err) + } + if results == nil { + t.Fatalf("%s: results is nil, must be an empty slice", name) + } + b, err := json.Marshal(results) + if err != nil { + t.Fatal(err) + } + if string(b) != "[]" { + t.Fatalf("%s: marshals to %s, want []", name, b) + } + } +} diff --git a/internal/watcher/fswatcher.go b/internal/watcher/fswatcher.go index 09ff105..f84f5ce 100644 --- a/internal/watcher/fswatcher.go +++ b/internal/watcher/fswatcher.go @@ -20,6 +20,7 @@ type FSWatcher struct { stopCh chan struct{} handler FileEventHandler timers map[string]*time.Timer + pending map[string]fsnotify.Op timersMu sync.Mutex } @@ -32,6 +33,7 @@ func NewFSWatcher() (*FSWatcher, error) { return &FSWatcher{ watcher: w, timers: make(map[string]*time.Timer), + pending: make(map[string]fsnotify.Op), }, nil } @@ -100,33 +102,60 @@ func (fw *FSWatcher) handleEvent(event fsnotify.Event) { } } - fw.debounce(path, func() { - if event.Has(fsnotify.Remove) || event.Has(fsnotify.Rename) { - fw.handler.OnDelete(path) - } else if event.Has(fsnotify.Create) { - fw.handler.OnCreate(path) - } else if event.Has(fsnotify.Write) || event.Has(fsnotify.Chmod) { - fw.handler.OnModify(path) - } - }) + fw.debounce(path, event.Op) } -func (fw *FSWatcher) debounce(path string, fn func()) { +// debounce coalesces events per path for debounceDuration. Ops are merged +// (OR-ed), not replaced: on Linux, writing a new file emits CREATE then WRITE +// as separate events within milliseconds — replacing the pending event would +// swallow the CREATE and misreport the file as modified (macOS coalesces +// differently, which long masked this). +func (fw *FSWatcher) debounce(path string, op fsnotify.Op) { fw.timersMu.Lock() defer fw.timersMu.Unlock() + fw.pending[path] |= op if t, ok := fw.timers[path]; ok { t.Stop() } fw.timers[path] = time.AfterFunc(debounceDuration, func() { - fn() fw.timersMu.Lock() + merged := fw.pending[path] + delete(fw.pending, path) delete(fw.timers, path) fw.timersMu.Unlock() + fw.dispatch(path, merged) }) } +// dispatch classifies a merged op set. Precedence: a removal ends the story +// regardless of what preceded it; a creation outranks the writes that filled +// the new file with content. Exception: atomic-save editors (vim with +// backupcopy=no, and similar rename-then-recreate patterns) emit RENAME then +// CREATE for the same path within one debounce window — the merged set carries +// Rename, but the file still exists and must not be dropped from the index, so +// a Remove/Rename verdict is confirmed against the filesystem before firing. +func (fw *FSWatcher) dispatch(path string, op fsnotify.Op) { + switch { + case op.Has(fsnotify.Remove) || op.Has(fsnotify.Rename): + if _, err := os.Stat(path); err == nil { + // Path still exists: rename-and-recreate, not a deletion. + if op.Has(fsnotify.Create) { + fw.handler.OnCreate(path) + } else { + fw.handler.OnModify(path) + } + return + } + fw.handler.OnDelete(path) + case op.Has(fsnotify.Create): + fw.handler.OnCreate(path) + case op.Has(fsnotify.Write) || op.Has(fsnotify.Chmod): + fw.handler.OnModify(path) + } +} + // Stop halts the event processing goroutine but keeps the fsnotify watcher // alive so directories can still be added and Start() can resume later. // Use Close() to release the underlying fsnotify resources. @@ -141,12 +170,13 @@ func (fw *FSWatcher) Stop() error { close(fw.stopCh) fw.running = false - // Cancel pending debounce timers. + // Cancel pending debounce timers and drop their merged ops. fw.timersMu.Lock() for _, t := range fw.timers { t.Stop() } fw.timers = make(map[string]*time.Timer) + fw.pending = make(map[string]fsnotify.Op) fw.timersMu.Unlock() return nil diff --git a/internal/watcher/fswatcher_test.go b/internal/watcher/fswatcher_test.go index b2e88bc..c13d405 100644 --- a/internal/watcher/fswatcher_test.go +++ b/internal/watcher/fswatcher_test.go @@ -6,6 +6,8 @@ import ( "sync" "testing" "time" + + "github.com/fsnotify/fsnotify" ) // mockHandler records filesystem events for assertions. @@ -206,3 +208,114 @@ func waitFor(t *testing.T, timeout time.Duration, cond func() bool) bool { } return false } + +// TestDebounceMergesOps drives the debouncer directly — no real filesystem +// events — so the Create+Write merge behavior is pinned deterministically on +// EVERY platform. The end-to-end TestOnCreate only catches a merge regression +// on OSes that emit the CREATE-then-WRITE double event (Linux/Windows); on +// macOS it passes either way, which would let the primary dev environment +// ship a regression. +func TestDebounceMergesOps(t *testing.T) { + newFW := func(t *testing.T) (*FSWatcher, *mockHandler) { + t.Helper() + fw, err := NewFSWatcher() + if err != nil { + t.Fatalf("NewFSWatcher: %v", err) + } + t.Cleanup(func() { fw.Close() }) + h := &mockHandler{} + fw.handler = h + return fw, h + } + wait := func() { time.Sleep(debounceDuration + 200*time.Millisecond) } + + t.Run("create then write fires OnCreate", func(t *testing.T) { + fw, h := newFW(t) + fw.debounce("/p/new.txt", fsnotify.Create) + fw.debounce("/p/new.txt", fsnotify.Write) // the Linux/Windows double event + wait() + if got := h.getCreates(); len(got) != 1 || got[0] != "/p/new.txt" { + t.Fatalf("OnCreate calls = %v, want exactly [/p/new.txt]", got) + } + if got := h.getModifies(); len(got) != 0 { + t.Fatalf("OnModify calls = %v, want none (Create outranks Write)", got) + } + }) + + t.Run("write alone fires OnModify", func(t *testing.T) { + fw, h := newFW(t) + fw.debounce("/p/existing.txt", fsnotify.Write) + wait() + if got := h.getModifies(); len(got) != 1 { + t.Fatalf("OnModify calls = %v, want exactly one", got) + } + if got := h.getCreates(); len(got) != 0 { + t.Fatalf("OnCreate calls = %v, want none", got) + } + }) + + t.Run("delete wins over create and write", func(t *testing.T) { + fw, h := newFW(t) + fw.debounce("/p/gone.txt", fsnotify.Create) + fw.debounce("/p/gone.txt", fsnotify.Write) + fw.debounce("/p/gone.txt", fsnotify.Remove) + wait() + if got := h.getDeletes(); len(got) != 1 { + t.Fatalf("OnDelete calls = %v, want exactly one (removal ends the story)", got) + } + if len(h.getCreates()) != 0 || len(h.getModifies()) != 0 { + t.Fatalf("create/modify fired alongside delete: creates=%v modifies=%v", h.getCreates(), h.getModifies()) + } + }) +} + +// TestDispatchAtomicSaveRename pins the atomic-save pattern: editors like vim +// (backupcopy=no) RENAME the file away then CREATE it fresh within one +// debounce window. The merged ops carry Rename, but the path still exists — +// it must dispatch as a create, not silently vanish from the index. A Rename +// with the path truly gone still dispatches as delete. +func TestDispatchAtomicSaveRename(t *testing.T) { + t.Run("rename then create, file exists -> OnCreate", func(t *testing.T) { + fw, h := func() (*FSWatcher, *mockHandler) { + fw, err := NewFSWatcher() + if err != nil { + t.Fatalf("NewFSWatcher: %v", err) + } + t.Cleanup(func() { fw.Close() }) + h := &mockHandler{} + fw.handler = h + return fw, h + }() + real := tempFileInDirW(t, t.TempDir(), "saved.md", "new content") + fw.dispatch(real, fsnotify.Rename|fsnotify.Create) + if got := h.getCreates(); len(got) != 1 { + t.Fatalf("OnCreate calls = %v, want exactly one", got) + } + if got := h.getDeletes(); len(got) != 0 { + t.Fatalf("OnDelete fired for a file that still exists: %v", got) + } + }) + + t.Run("rename, file gone -> OnDelete", func(t *testing.T) { + fw, err := NewFSWatcher() + if err != nil { + t.Fatalf("NewFSWatcher: %v", err) + } + t.Cleanup(func() { fw.Close() }) + h := &mockHandler{} + fw.handler = h + fw.dispatch("/definitely/not/a/real/path.md", fsnotify.Rename) + if got := h.getDeletes(); len(got) != 1 { + t.Fatalf("OnDelete calls = %v, want exactly one", got) + } + }) +} + +func tempFileInDirW(t *testing.T, dir, name, content string) string { + t.Helper() + p := filepath.Join(dir, name) + if err := os.WriteFile(p, []byte(content), 0644); err != nil { + t.Fatal(err) + } + return p +} diff --git a/main.go b/main.go index d02b191..10d4d9e 100644 --- a/main.go +++ b/main.go @@ -3,13 +3,13 @@ package main import ( "embed" "flag" + "fmt" "log" "os" "path/filepath" - "strconv" - "github.com/borzou/vecstore/internal/chunker" "github.com/borzou/vecstore/internal/embeddings" + "github.com/borzou/vecstore/internal/embeddings/local" "github.com/borzou/vecstore/internal/engine" "github.com/borzou/vecstore/internal/extractor" "github.com/borzou/vecstore/internal/mcp" @@ -21,6 +21,35 @@ import ( "github.com/wailsapp/wails/v2/pkg/options/assetserver" ) +// makeEmbedder is the composition root's EmbedderFactory. Constructors are cheap +// and lazy: local.New does not load the model until the first embed call, so the +// read-only MCP process can build one without forcing native libraries to load. +func makeEmbedder(provider, apiKey, model string) (embeddings.Embedder, error) { + switch provider { + case embeddings.ProviderLocal: + return local.New(local.Config{}), nil + case embeddings.ProviderOpenAI: + return embeddings.NewOpenAIEmbedder(apiKey, model), nil + default: + return nil, fmt.Errorf("unknown embedding provider %q", provider) + } +} + +// peekConfig reads the provider-selection config from an existing database +// without migrating it. A fresh/uninitialized DB yields empty values, which the +// resolver treats as "brand new install → default provider". +func peekConfig(dbPath string) (provider, apiKey, model string) { + ro, err := store.NewReadOnlySQLiteStore(dbPath) + if err != nil { + return "", "", "" + } + defer ro.Close() + provider, _ = ro.GetConfig("embedding_provider") + apiKey, _ = ro.GetConfig("openai_api_key") + model, _ = ro.GetConfig("embedding_model") + return provider, apiKey, model +} + //go:embed all:frontend/dist var assets embed.FS @@ -46,13 +75,16 @@ func main() { } defer s.Close() + providerCfg, _ := s.GetConfig("embedding_provider") apiKey, _ := s.GetConfig("openai_api_key") - model, _ := s.GetConfig("embedding_model") - if model == "" { - model = "text-embedding-3-small" - } + modelCfg, _ := s.GetConfig("embedding_model") + provider := resolveProvider(providerCfg, apiKey) + model := resolveModel(provider, modelCfg) - embedder := embeddings.NewOpenAIEmbedder(apiKey, model) + embedder, err := makeEmbedder(provider, apiKey, model) + if err != nil { + log.Fatalf("create embedder: %v", err) + } roEngine := engine.NewReadOnly(s, embedder) stdio := mcp.NewStdioServer(roEngine) @@ -69,35 +101,49 @@ func main() { log.Fatalf("create data directory: %v", err) } - // 1. Open store. - s, err := store.NewSQLiteStore(*dbPath) + // 1. Peek existing config (if any) to resolve the provider BEFORE opening + // the store, so a fresh DB's vector table is created at the correct + // dimension. A fresh/uninitialized DB peeks empty → default provider. + peekProvider, peekKey, peekModel := peekConfig(*dbPath) + provider := resolveProvider(peekProvider, peekKey) + model := resolveModel(provider, peekModel) + dim := embeddings.DefaultDimension(provider, model) + + // 2. Open store with the resolved dimension. + s, err := store.NewSQLiteStore(*dbPath, dim) if err != nil { log.Fatalf("open store: %v", err) } - // 2. Read config from store. + // 3. Persist the resolved provider/model so Stats and later config changes + // reflect reality (idempotent; harmless on re-launch). apiKey, _ := s.GetConfig("openai_api_key") - model, _ := s.GetConfig("embedding_model") - if model == "" { - model = "text-embedding-3-small" + if err := s.SetConfig("embedding_provider", provider); err != nil { + log.Fatalf("persist embedding_provider: %v", err) + } + if err := s.SetConfig("embedding_model", model); err != nil { + log.Fatalf("persist embedding_model: %v", err) } - // 3. Create embedder (may have empty API key on first run). - embedder := embeddings.NewOpenAIEmbedder(apiKey, model) - - // 4. Create chunker with stored config. - var chunkOpts []chunker.Option - if sizeStr, _ := s.GetConfig("chunk_size"); sizeStr != "" { - if size, err := strconv.Atoi(sizeStr); err == nil { - chunkOpts = append(chunkOpts, chunker.WithChunkSize(size)) + // 3b. Backfill onboarding_complete for pre-existing users. If the flag has + // never been set but the DB already has watched directories or an OpenAI + // key, this is an upgraded install that should skip onboarding. A truly + // fresh DB (no dirs, no key) leaves the flag unset so onboarding runs. + if complete, _ := s.GetConfig("onboarding_complete"); complete == "" { + dirs, _ := s.ListDirectories() + if len(dirs) > 0 || apiKey != "" { + if err := s.SetConfig("onboarding_complete", "true"); err != nil { + log.Fatalf("persist onboarding_complete: %v", err) + } } } - if overlapStr, _ := s.GetConfig("chunk_overlap"); overlapStr != "" { - if overlap, err := strconv.Atoi(overlapStr); err == nil { - chunkOpts = append(chunkOpts, chunker.WithOverlap(overlap)) - } + + // 4. Create the provider-matched embedder and chunker. + embedder, err := makeEmbedder(provider, apiKey, model) + if err != nil { + log.Fatalf("create embedder: %v", err) } - c, err := chunker.New(chunkOpts...) + c, err := buildChunker(s, provider, embedder) if err != nil { log.Fatalf("create chunker: %v", err) } @@ -137,9 +183,7 @@ func main() { store: s, mcpServer: mcpServer, dbPath: *dbPath, - newEmbedder: func(apiKey, model string) embeddings.Embedder { - return embeddings.NewOpenAIEmbedder(apiKey, model) - }, + newEmbedder: makeEmbedder, } appInstance = app @@ -147,7 +191,13 @@ func main() { Title: "Agent Memory", Width: 900, Height: 700, - HideWindowOnClose: true, + // Hide-on-close is only safe where a tray exists to surface/quit the + // hidden app; without one, hiding leaves an invisible process and + // relaunches stack zombie instances contending for the single-writer + // SQLite DB (observed on Windows: six concurrent instances). hasTray + // is owned by the tray build-tag pair (tray.go / tray_stub.go), so + // this can never drift from the actual tray implementation. + HideWindowOnClose: hasTray, AssetServer: &assetserver.Options{ Assets: assets, }, diff --git a/main_test.go b/main_test.go new file mode 100644 index 0000000..c4995c8 --- /dev/null +++ b/main_test.go @@ -0,0 +1,102 @@ +package main + +import ( + "path/filepath" + "testing" + + "github.com/borzou/vecstore/internal/embeddings" + "github.com/borzou/vecstore/internal/engine" + "github.com/borzou/vecstore/internal/mocks" + "github.com/borzou/vecstore/internal/store" +) + +func TestResolveProvider(t *testing.T) { + cases := []struct { + name string + configured string + apiKey string + want string + }{ + {"default when nothing set", "", "", embeddings.ProviderLocal}, + {"openai when key present", "", "sk-abc", embeddings.ProviderOpenAI}, + {"explicit provider wins over key", embeddings.ProviderLocal, "sk-abc", embeddings.ProviderLocal}, + {"explicit openai with no key", embeddings.ProviderOpenAI, "", embeddings.ProviderOpenAI}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := resolveProvider(tc.configured, tc.apiKey); got != tc.want { + t.Fatalf("resolveProvider(%q,%q)=%q want %q", tc.configured, tc.apiKey, got, tc.want) + } + }) + } +} + +func TestResolveModel(t *testing.T) { + if got := resolveModel(embeddings.ProviderLocal, ""); got != "multilingual-e5-small" { + t.Fatalf("local default model = %q", got) + } + if got := resolveModel(embeddings.ProviderOpenAI, ""); got != "text-embedding-3-small" { + t.Fatalf("openai default model = %q", got) + } + if got := resolveModel(embeddings.ProviderOpenAI, "text-embedding-3-large"); got != "text-embedding-3-large" { + t.Fatalf("stored model should win, got %q", got) + } +} + +// TestSetConfigSwitchProvider proves that switching the provider swaps both the +// embedder and the chunker and resets the index. The app's config store is a +// real SQLite DB; the engine is built over mocks so the Reset call is observable. +// Switching to OpenAI (tiktoken chunker) keeps this test native-lib-free. +func TestSetConfigSwitchProvider(t *testing.T) { + dir := t.TempDir() + s, err := store.NewSQLiteStore(filepath.Join(dir, "cfg.db"), 384) + if err != nil { + t.Fatalf("store: %v", err) + } + defer s.Close() + + // The engine runs over mocks so Reset is observable. initialScan (fired by + // Reset→Start) lists directories; the default mock returns none, so it exits. + var resetCalled bool + engStore := &mocks.MockStore{ + ResetFn: func(dim int) error { resetCalled = true; return nil }, + } + + newEmb := &mocks.MockEmbedder{ + DimensionsFn: func() int { return 1536 }, + MaxInputTokensFn: func() int { return 0 }, + } + var factoryProvider, factoryModel string + eng := engine.New(engStore, &mocks.MockEmbedder{}, &mocks.MockChunker{}, &mocks.MockWatcher{}, &mocks.MockExtractor{}) + + app := &App{ + engine: eng, + store: s, + newEmbedder: func(provider, apiKey, model string) (embeddings.Embedder, error) { + factoryProvider = provider + factoryModel = model + return newEmb, nil + }, + } + + // Start from the default (local) provider; switch to openai. + if err := app.SetConfig("embedding_provider", embeddings.ProviderOpenAI); err != nil { + t.Fatalf("SetConfig: %v", err) + } + + if factoryProvider != embeddings.ProviderOpenAI { + t.Fatalf("factory called with provider %q, want openai", factoryProvider) + } + if factoryModel != "text-embedding-3-small" { + t.Fatalf("factory model = %q, want text-embedding-3-small", factoryModel) + } + if !resetCalled { + t.Fatal("expected engine.Reset() to hit the store") + } + if p, _ := s.GetConfig("embedding_provider"); p != embeddings.ProviderOpenAI { + t.Fatalf("persisted provider = %q", p) + } + if m, _ := s.GetConfig("embedding_model"); m != "text-embedding-3-small" { + t.Fatalf("persisted model = %q", m) + } +} diff --git a/tray.go b/tray.go index 467a595..1740ee1 100644 --- a/tray.go +++ b/tray.go @@ -1,3 +1,5 @@ +//go:build darwin + package main /* @@ -80,6 +82,11 @@ import ( //go:embed assets/trayicon.png var trayIcon []byte +// hasTray reports whether this platform has a status-bar tray to hide into. +// Owned by the tray build-tag pair so main.go's hide-on-close behavior can +// never drift from the tray implementation. +const hasTray = true + // setupTray creates a macOS status bar item with Show / Quit menu. // Returns a cleanup function. func (a *App) setupTray() func() { diff --git a/tray_stub.go b/tray_stub.go new file mode 100644 index 0000000..1e0f1c8 --- /dev/null +++ b/tray_stub.go @@ -0,0 +1,15 @@ +//go:build !darwin + +package main + +// hasTray reports whether this platform has a status-bar tray to hide into. +// Owned by the tray build-tag pair so main.go's hide-on-close behavior can +// never drift from the tray implementation (see background-presence epic for +// the plan to bring a tray to the remaining platforms). +const hasTray = false + +// setupTray is a no-op on platforms without the macOS status-bar integration +// (tray.go is darwin-only Objective-C via CGo). Returns a no-op cleanup func. +func (a *App) setupTray() func() { + return func() {} +}