diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..8140821 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,24 @@ + +# Memory Context + +# [Recall/hamburg] recent context, 2026-05-07 10:45am PDT + +Legend: ๐ŸŽฏsession ๐Ÿ”ดbugfix ๐ŸŸฃfeature ๐Ÿ”„refactor โœ…change ๐Ÿ”ตdiscovery โš–๏ธdecision ๐Ÿšจsecurity_alert ๐Ÿ”security_note +Format: ID TIME TYPE TITLE +Fetch details: get_observations([IDs]) | Search: mem-search skill + +Stats: 9 obs (3,300t read) | 133,379t work | 98% savings + +### Apr 20, 2026 +7 10:24a โš–๏ธ Recall/Mumbai App โ€” Full-Solution Expansion Prompt Requested +8 " ๐Ÿ”ต Recall/Mumbai โ€” Existing Architecture Bottlenecks Confirmed +9 " โœ… Plan File Repurposed as Full Product Vision Master Prompt +12 11:54a โš–๏ธ Recall โ€” Full Product Vision Defined as Agent Build Prompt +13 " ๐ŸŸฃ Phase 1A โ€” Local AI Captioning Module Designed (captioner.py) +14 " ๐ŸŸฃ Phase 1Bโ€“1C โ€” Concurrent Ingest + FastAPI Daemon Architecture Designed +15 " ๐ŸŸฃ Phase 2 โ€” Six Source Connectors Designed with BaseConnector Interface +16 " ๐ŸŸฃ Phase 1D + Phase 3 โ€” Raycast HTTP Client Rewrite and Source-Filter UX Designed +17 " โœ… New Dependencies Added to pyproject.toml for All Phases + +Access 133k tokens of past work via get_observations([IDs]) or mem-search skill. + \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 4195fe8..6b88ddf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "vector-embedded-finder" version = "0.2.0" -description = "Local multimodal memory with semantic search, powered by Gemini Embedding 2 and ChromaDB" +description = "Local-first multimodal memory with hybrid semantic search, sqlite durability, and a persistent Recall daemon" requires-python = ">=3.11" license = "MIT" dependencies = [ @@ -26,6 +26,9 @@ dependencies = [ "requests-oauthlib>=1.3.0", "beautifulsoup4>=4.12.0", "pypdf>=4.0.0", + "hnswlib>=0.8.0", + "sentence-transformers>=3.0.0", + "Pillow>=10.0.0", ] [project.optional-dependencies] diff --git a/raycast/assets/command-icon.png b/raycast/assets/command-icon.png index 2563eed..3c27f18 100644 Binary files a/raycast/assets/command-icon.png and b/raycast/assets/command-icon.png differ diff --git a/raycast/package.json b/raycast/package.json index 4a550f6..2e68ca7 100644 --- a/raycast/package.json +++ b/raycast/package.json @@ -4,7 +4,7 @@ "title": "Recall", "description": "Search local files and connected apps with semantic memory", "icon": "command-icon.png", - "author": "hughp", + "author": "ayush", "categories": [ "Productivity" ], diff --git a/raycast/src/lib/runner.ts b/raycast/src/lib/runner.ts index dc84e7d..ccf751d 100644 --- a/raycast/src/lib/runner.ts +++ b/raycast/src/lib/runner.ts @@ -1,32 +1,27 @@ /** - * VEF runner โ€” all searches go via the persistent daemon over HTTP. + * Recall runner โ€” all daemon traffic goes over the Unix domain socket. * Falls back to auto-starting the daemon if it is not running. */ +import http from "http"; +import os from "os"; import { spawnSync, spawn } from "child_process"; import { getPreferenceValues } from "@raycast/api"; -// โ”€โ”€ Preferences โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - export interface Preferences { - /** Absolute path to the directory that contains vector_embedded_finder/ */ pythonPackagePath: string; - /** Path to python3 binary. Defaults to "python3". */ pythonPath?: string; - /** Gemini API key โ€” stored in macOS Keychain via Raycast password preference */ geminiApiKey?: string; } -// โ”€โ”€ Error types โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - export type VefErrorCode = - | "NOT_INSTALLED" // package not importable from the given path - | "PYTHON_NOT_FOUND" // python binary not found - | "AUTH_ERROR" // GEMINI_API_KEY missing or invalid - | "RATE_LIMIT" // 429 / quota exceeded - | "TIMEOUT" // request timed out - | "DAEMON_ERROR" // daemon returned a non-2xx response - | "UNKNOWN"; // any other failure + | "NOT_INSTALLED" + | "PYTHON_NOT_FOUND" + | "AUTH_ERROR" + | "RATE_LIMIT" + | "TIMEOUT" + | "DAEMON_ERROR" + | "UNKNOWN"; export class VefRunnerError extends Error { constructor( @@ -38,8 +33,6 @@ export class VefRunnerError extends Error { } } -// โ”€โ”€ Types โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - export interface SearchResult { id: string; similarity: number; @@ -73,20 +66,7 @@ export interface ProgressInfo { total_indexed: number; } -// โ”€โ”€ Daemon coordinates โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - -const DAEMON_HOST = "127.0.0.1"; -const DAEMON_PORT = 19847; -const BASE_URL = `http://${DAEMON_HOST}:${DAEMON_PORT}`; - -// โ”€โ”€ Helpers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - -function abortAfter(ms: number): AbortSignal { - const controller = new AbortController(); - const id = setTimeout(() => controller.abort(), ms); - controller.signal.addEventListener("abort", () => clearTimeout(id), { once: true }); - return controller.signal; -} +const SOCKET_PATH = process.env.RECALL_SOCKET_PATH?.trim() || `${os.homedir()}/.recall/recall.sock`; function resolvePrefs(): Preferences { return getPreferenceValues(); @@ -99,66 +79,115 @@ function resolvePython(prefs: Preferences): string { return "python3"; } -async function _pollHealth(timeoutMs: number): Promise { +function delay(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +function formatDetail(detail: unknown): string { + if (!detail) return ""; + if (typeof detail === "string") return detail; + if (Array.isArray(detail)) return detail.map((item) => formatDetail(item)).filter(Boolean).join("; "); + if (typeof detail === "object") return JSON.stringify(detail); + return String(detail); +} + +function requestJson( + path: string, + options: { method?: string; body?: unknown; timeoutMs?: number } = {}, +): Promise<{ statusCode: number; body: T }> { + const method = options.method ?? "GET"; + const timeoutMs = options.timeoutMs ?? 5000; + const bodyText = options.body === undefined ? undefined : JSON.stringify(options.body); + + return new Promise((resolve, reject) => { + const req = http.request( + { + socketPath: SOCKET_PATH, + path, + method, + headers: bodyText + ? { + "Content-Type": "application/json", + "Content-Length": Buffer.byteLength(bodyText), + } + : undefined, + }, + (res) => { + const chunks: Buffer[] = []; + res.on("data", (chunk) => chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk))); + res.on("end", () => { + const raw = Buffer.concat(chunks).toString("utf-8"); + let parsed: unknown = {}; + if (raw.trim()) { + try { + parsed = JSON.parse(raw); + } catch { + parsed = raw; + } + } + resolve({ statusCode: res.statusCode ?? 0, body: parsed as T }); + }); + }, + ); + + const timer = setTimeout(() => { + req.destroy(new Error(`Request timed out after ${timeoutMs} ms`)); + }, timeoutMs); + + req.on("error", (err) => { + clearTimeout(timer); + reject(err); + }); + req.on("close", () => clearTimeout(timer)); + + if (bodyText) req.write(bodyText); + req.end(); + }); +} + +async function pollHealth(timeoutMs: number): Promise { const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { try { - // /health is constant-time on the daemon, so a generous per-probe - // abort is safe and avoids false negatives under macOS CPU contention. - const resp = await fetch(`${BASE_URL}/health`, { signal: abortAfter(2000) }); - if (resp.ok) return true; + const resp = await requestJson<{ status?: string }>("/health", { timeoutMs: 2000 }); + if (resp.statusCode >= 200 && resp.statusCode < 300) return true; } catch { - // not ready yet + // keep polling } - await new Promise((r) => setTimeout(r, 400)); + await delay(400); } return false; } -/** - * Ensure the daemon is running, starting it if necessary. - * Polls /health for up to 10 seconds after spawning. - */ async function ensureDaemon(): Promise { - // Fast path: already up try { - const resp = await fetch(`${BASE_URL}/health`, { signal: abortAfter(3000) }); - if (resp.ok) return; + const resp = await requestJson<{ status?: string }>("/health", { timeoutMs: 3000 }); + if (resp.statusCode >= 200 && resp.statusCode < 300) return; } catch { // fall through to start } - // Race guard: if another process is warming up, avoid double-spawn. - if (await _pollHealth(2000)) return; + if (await pollHealth(2000)) return; const prefs = resolvePrefs(); const python = resolvePython(prefs); const packagePath = prefs.pythonPackagePath?.trim() || ""; - const env: Record = { ...process.env } as Record; - if (prefs.geminiApiKey?.trim()) { - env["GEMINI_API_KEY"] = prefs.geminiApiKey.trim(); - } + if (prefs.geminiApiKey?.trim()) env["GEMINI_API_KEY"] = prefs.geminiApiKey.trim(); if (packagePath) { const existing = env["PYTHONPATH"] || ""; env["PYTHONPATH"] = existing ? `${packagePath}:${existing}` : packagePath; } - // Spawn daemon detached so it outlives Raycast. - // Use async spawn (not spawnSync) to avoid blocking Raycast's main thread. let spawnError: string | null = null; try { - const proc = spawn( - python, - ["-m", "vector_embedded_finder.daemon", "_serve"], - { - detached: true, - stdio: "ignore", - env, - cwd: packagePath || undefined, - }, - ); - proc.unref(); // allow parent (Raycast) to exit independently + const proc = spawn(python, ["-m", "vector_embedded_finder.daemon", "_serve"], { + detached: true, + stdio: "ignore", + env, + cwd: packagePath || undefined, + }); + proc.unref(); proc.on("error", (err) => { spawnError = err.message; }); @@ -166,256 +195,143 @@ async function ensureDaemon(): Promise { spawnError = err instanceof Error ? err.message : String(err); } - // Poll /health for up to 10 seconds (daemon takes ~1-2s to warm up) - const ready = await _pollHealth(10000); + const ready = await pollHealth(15000); if (!ready) { const hint = spawnError ? ` (${spawnError})` : ""; throw new VefRunnerError( "DAEMON_ERROR", - `Daemon failed to start within 10 s${hint}. Run "vef-daemon start" manually.`, + `Daemon failed to start within 15 s${hint}. Run "vef-daemon start" manually.`, ); } } -// โ”€โ”€ Public API โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - export interface RunSearchOptions { nResults?: number; sources?: string[] | null; } -function formatDetail(detail: unknown): string { - if (!detail) return ""; - if (typeof detail === "string") return detail; - if (Array.isArray(detail)) { - const msgs = detail - .map((item) => { - if (typeof item === "string") return item; - if (item && typeof item === "object" && "msg" in item) { - return String((item as { msg?: unknown }).msg ?? ""); - } - return JSON.stringify(item); - }) - .filter(Boolean); - return msgs.join("; "); - } - if (typeof detail === "object") return JSON.stringify(detail); - return String(detail); -} - -/** - * Run a semantic search via the VEF daemon. - * - * @throws {VefRunnerError} with appropriate code on any failure - */ export async function runSearch( query: string, nResultsOrOptions: number | RunSearchOptions = 20, ): Promise { if (!query.trim()) return []; - const opts: RunSearchOptions = - typeof nResultsOrOptions === "number" - ? { nResults: nResultsOrOptions } - : nResultsOrOptions; - - const nResults = opts.nResults ?? 20; - const sources = opts.sources ?? null; - + typeof nResultsOrOptions === "number" ? { nResults: nResultsOrOptions } : nResultsOrOptions; await ensureDaemon(); - - const body: Record = { query, n_results: nResults }; - if (sources && sources.length > 0) body.sources = sources; - - let resp: Response; - try { - resp = await fetch(`${BASE_URL}/search`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(body), - signal: abortAfter(30000), - }); - } catch (err: unknown) { - const msg = err instanceof Error ? err.message : String(err); - if (msg.includes("TimeoutError") || msg.includes("timed out") || msg.includes("ETIMEDOUT")) { - throw new VefRunnerError("TIMEOUT", "Search timed out after 30 s"); - } - throw new VefRunnerError("UNKNOWN", `Fetch failed: ${msg}`); - } - - if (!resp.ok) { - let detail = ""; - try { - const payload = (await resp.json()) as { detail?: unknown }; - detail = formatDetail(payload.detail); - } catch { - // ignore parse error - } - const err = `Daemon returned HTTP ${resp.status}${detail ? `: ${detail}` : ""}`; - if (resp.status === 401 || resp.status === 403) { - throw new VefRunnerError("AUTH_ERROR", err); - } - if (resp.status === 429) { - throw new VefRunnerError("RATE_LIMIT", err); - } + const resp = await requestJson("/search", { + method: "POST", + body: { + query, + n_results: opts.nResults ?? 20, + ...(opts.sources && opts.sources.length > 0 ? { sources: opts.sources } : {}), + }, + timeoutMs: 30000, + }); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + const detail = formatDetail((resp.body as { detail?: unknown }).detail); + const err = `Daemon returned HTTP ${resp.statusCode}${detail ? `: ${detail}` : ""}`; + if (resp.statusCode === 401 || resp.statusCode === 403) throw new VefRunnerError("AUTH_ERROR", err); + if (resp.statusCode === 429) throw new VefRunnerError("RATE_LIMIT", err); throw new VefRunnerError("DAEMON_ERROR", err); } - - const results = (await resp.json()) as SearchResult[]; - return results; + return resp.body as SearchResult[]; } -/** - * Validate that the daemon is reachable and the DB has items. - * Returns the item count in the vector store. - * - * @throws {VefRunnerError} with appropriate code on any failure - */ export async function validateSetup(): Promise<{ count: number }> { await ensureDaemon(); - - // Use /stats (which reports count) rather than /health (liveness-only). - // /stats may be slow when chromadb is backlogged, so give it room. - let resp: Response; - try { - resp = await fetch(`${BASE_URL}/stats`, { signal: abortAfter(5000) }); - } catch (err: unknown) { - const msg = err instanceof Error ? err.message : String(err); - throw new VefRunnerError("UNKNOWN", `Stats check failed: ${msg}`); - } - - if (!resp.ok) { - throw new VefRunnerError("DAEMON_ERROR", `Daemon stats returned HTTP ${resp.status}`); + const resp = await requestJson<{ count?: number }>("/stats", { timeoutMs: 5000 }); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Daemon stats returned HTTP ${resp.statusCode}`); } - - const data = (await resp.json()) as { status?: string; count?: number }; - return { count: typeof data.count === "number" ? data.count : 0 }; + return { count: typeof resp.body.count === "number" ? resp.body.count : 0 }; } -/** - * Fetch the list of indexed sources from the daemon. - */ export async function fetchSources(): Promise { try { await ensureDaemon(); - const resp = await fetch(`${BASE_URL}/sources`, { signal: abortAfter(2000) }); - if (!resp.ok) return []; - const data = (await resp.json()) as { sources: string[] }; - return data.sources; + const resp = await requestJson<{ sources?: string[] }>("/sources", { timeoutMs: 3000 }); + return resp.statusCode >= 200 && resp.statusCode < 300 ? resp.body.sources ?? [] : []; } catch { return []; } } -/** - * Fetch per-connector status from the daemon. - */ export async function fetchConnectorStatus(): Promise { await ensureDaemon(); - const resp = await fetch(`${BASE_URL}/connector-status`, { signal: abortAfter(3000) }); - if (!resp.ok) { - throw new VefRunnerError("DAEMON_ERROR", `Connector status request failed: HTTP ${resp.status}`); + const resp = await requestJson("/connector-status", { timeoutMs: 3000 }); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Connector status request failed: HTTP ${resp.statusCode}`); } - return (await resp.json()) as ConnectorStatusMap; + return resp.body; } -/** - * Fetch indexing progress from the daemon. - */ export async function fetchProgress(): Promise { await ensureDaemon(); - const resp = await fetch(`${BASE_URL}/progress`, { signal: abortAfter(3000) }); - if (!resp.ok) { - throw new VefRunnerError("DAEMON_ERROR", `Progress request failed: HTTP ${resp.status}`); + const resp = await requestJson("/progress", { timeoutMs: 3000 }); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Progress request failed: HTTP ${resp.statusCode}`); } - return (await resp.json()) as ProgressInfo; + return resp.body; } -/** - * Trigger immediate connector sync. Returns immediately (daemon handles in_progress). - */ export async function triggerSync(source?: string): Promise<{ status: string; last_sync: Record }> { await ensureDaemon(); - const payload = source ? { source } : {}; - const resp = await fetch(`${BASE_URL}/sync`, { + const resp = await requestJson<{ status: string; last_sync: Record; detail?: unknown }>("/sync", { method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(payload), - signal: abortAfter(30000), + body: source ? { source } : {}, + timeoutMs: 30000, }); - if (!resp.ok) { - let detail = ""; - try { - const payloadErr = (await resp.json()) as { detail?: unknown }; - detail = formatDetail(payloadErr.detail); - } catch { - // ignore - } - throw new VefRunnerError("DAEMON_ERROR", `Sync request failed: HTTP ${resp.status}${detail ? `: ${detail}` : ""}`); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + const detail = formatDetail(resp.body.detail); + throw new VefRunnerError( + "DAEMON_ERROR", + `Sync request failed: HTTP ${resp.statusCode}${detail ? `: ${detail}` : ""}`, + ); } - return (await resp.json()) as { status: string; last_sync: Record }; + return resp.body as { status: string; last_sync: Record }; } -/** - * Check whether a connector sync is actively running. - */ export async function fetchSyncRunning(): Promise { try { - const resp = await fetch(`${BASE_URL}/sync-running`, { signal: abortAfter(3000) }); - if (!resp.ok) return false; - const data = (await resp.json()) as { running: boolean }; - return data.running; + const resp = await requestJson<{ running?: boolean }>("/sync-running", { timeoutMs: 3000 }); + return !!resp.body.running; } catch { return false; } } -/** - * Fetch list of watched directories. - */ export async function fetchWatchedDirs(): Promise { await ensureDaemon(); - const resp = await fetch(`${BASE_URL}/watched-dirs`, { signal: abortAfter(3000) }); - if (!resp.ok) return []; - const data = (await resp.json()) as { dirs: string[] }; - return data.dirs; + const resp = await requestJson<{ dirs?: string[] }>("/watched-dirs", { timeoutMs: 3000 }); + return resp.statusCode >= 200 && resp.statusCode < 300 ? resp.body.dirs ?? [] : []; } -/** - * Add a directory to the watched list. - */ export async function addWatchedDir(path: string): Promise { await ensureDaemon(); - const resp = await fetch(`${BASE_URL}/watched-dirs`, { + const resp = await requestJson<{ dirs?: string[] }>("/watched-dirs", { method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ path }), - signal: abortAfter(5000), + body: { path }, + timeoutMs: 5000, }); - if (!resp.ok) throw new VefRunnerError("DAEMON_ERROR", `Failed to add directory: HTTP ${resp.status}`); - const data = (await resp.json()) as { dirs: string[] }; - return data.dirs; + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Failed to add directory: HTTP ${resp.statusCode}`); + } + return resp.body.dirs ?? []; } -/** - * Remove a directory from the watched list. - */ export async function removeWatchedDir(path: string): Promise { await ensureDaemon(); - const resp = await fetch(`${BASE_URL}/watched-dirs`, { + const resp = await requestJson<{ dirs?: string[] }>("/watched-dirs", { method: "DELETE", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ path }), - signal: abortAfter(5000), + body: { path }, + timeoutMs: 5000, }); - if (!resp.ok) throw new VefRunnerError("DAEMON_ERROR", `Failed to remove directory: HTTP ${resp.status}`); - const data = (await resp.json()) as { dirs: string[] }; - return data.dirs; + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Failed to remove directory: HTTP ${resp.statusCode}`); + } + return resp.body.dirs ?? []; } -/** - * Save API keys / config to daemon (persists to ~/.vef/.env). - */ export async function saveConfigure(cfg: { gemini_api_key?: string; canvas_api_key?: string; @@ -424,18 +340,16 @@ export async function saveConfigure(cfg: { schoology_consumer_secret?: string; }): Promise { await ensureDaemon(); - const resp = await fetch(`${BASE_URL}/configure`, { + const resp = await requestJson<{ ok?: boolean }>("/configure", { method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(cfg), - signal: abortAfter(5000), + body: cfg, + timeoutMs: 5000, }); - if (!resp.ok) throw new VefRunnerError("DAEMON_ERROR", `Configure failed: HTTP ${resp.status}`); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Configure failed: HTTP ${resp.statusCode}`); + } } -/** - * Run trayce connect in Terminal via AppleScript. - */ export async function connectInTerminal(source: string): Promise { const { execFile } = await import("node:child_process"); await new Promise((resolve, reject) => { @@ -447,9 +361,6 @@ export async function connectInTerminal(source: string): Promise { }); } -/** - * Trigger index of a local directory path in Terminal. - */ export async function indexFolderInTerminal(path: string): Promise { const { execFile } = await import("node:child_process"); const escaped = path.replace(/"/g, '\\"'); @@ -461,3 +372,43 @@ export async function indexFolderInTerminal(path: string): Promise { ); }); } + +export async function fetchModelStatus(): Promise> { + await ensureDaemon(); + const resp = await requestJson>("/model-status", { timeoutMs: 5000 }); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Model status failed: HTTP ${resp.statusCode}`); + } + return resp.body; +} + +export async function fetchIndexStatus(): Promise> { + await ensureDaemon(); + const resp = await requestJson>("/index-status", { timeoutMs: 5000 }); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Index status failed: HTTP ${resp.statusCode}`); + } + return resp.body; +} + +export async function fetchMigrationStatus(): Promise> { + await ensureDaemon(); + const resp = await requestJson>("/migration-status", { timeoutMs: 5000 }); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Migration status failed: HTTP ${resp.statusCode}`); + } + return resp.body; +} + +export async function rebuildIndex(): Promise> { + await ensureDaemon(); + const resp = await requestJson>("/rebuild-index", { + method: "POST", + body: {}, + timeoutMs: 30000, + }); + if (resp.statusCode < 200 || resp.statusCode >= 300) { + throw new VefRunnerError("DAEMON_ERROR", `Rebuild index failed: HTTP ${resp.statusCode}`); + } + return resp.body; +} diff --git a/raycast/src/manage.tsx b/raycast/src/manage.tsx index 4a2b99e..8001eda 100644 --- a/raycast/src/manage.tsx +++ b/raycast/src/manage.tsx @@ -18,10 +18,14 @@ import { addWatchedDir, connectInTerminal, fetchConnectorStatus, + fetchIndexStatus, + fetchMigrationStatus, + fetchModelStatus, fetchProgress, fetchSyncRunning, fetchWatchedDirs, indexFolderInTerminal, + rebuildIndex, removeWatchedDir, saveConfigure, triggerSync, @@ -102,7 +106,7 @@ function ConfigureForm() { } > - + ({ indexing: false, queued: 0, total_indexed: 0 }); const [connectors, setConnectors] = useState({}); const [watchedDirs, setWatchedDirs] = useState([]); + const [modelStatus, setModelStatus] = useState>({}); + const [indexStatus, setIndexStatus] = useState>({}); + const [migrationStatus, setMigrationStatus] = useState>({}); const prevDocs = useRef(0); const loadAll = useCallback(async () => { try { - const [health, connStatus, prog, dirs, syncRunning] = await Promise.all([ + const [health, connStatus, prog, dirs, syncRunning, models, indexInfo, migrationInfo] = await Promise.all([ validateSetup(), fetchConnectorStatus(), fetchProgress(), fetchWatchedDirs(), fetchSyncRunning(), + fetchModelStatus(), + fetchIndexStatus(), + fetchMigrationStatus(), ]); setDaemonOk(true); setDocCount(health.count ?? 0); @@ -259,6 +269,9 @@ export default function ManageRecall() { setProgress(prog); setWatchedDirs(dirs); setSyncing(syncRunning); + setModelStatus(models); + setIndexStatus(indexInfo); + setMigrationStatus(migrationInfo); } catch { setDaemonOk(false); } finally { @@ -344,6 +357,21 @@ export default function ManageRecall() { } } + async function handleRebuildIndex() { + const toast = await showToast({ style: Toast.Style.Animated, title: "Rebuilding search indexโ€ฆ" }); + try { + const result = await rebuildIndex(); + toast.style = Toast.Style.Success; + toast.title = "Search index rebuilt"; + toast.message = `${result["count"] ?? 0} items`; + await loadAll(); + } catch (err) { + toast.style = Toast.Style.Failure; + toast.title = "Rebuild failed"; + toast.message = err instanceof Error ? err.message : String(err); + } + } + const connectorIcon = (name: string): Icon => { const icons: Record = { gmail: Icon.Envelope, @@ -357,6 +385,11 @@ export default function ManageRecall() { return icons[name] ?? Icon.Plug; }; + const modelRows = (modelStatus["models"] as Record> | undefined) ?? {}; + const runtime = (modelStatus["runtime"] as Record | undefined) ?? {}; + const indexBackend = String(indexStatus["backend"] ?? "unknown"); + const migrationState = String(migrationStatus["status"] ?? "unknown"); + return ( {/* โ”€โ”€ System โ”€โ”€ */} @@ -378,11 +411,46 @@ export default function ManageRecall() { actions={ + void handleRebuildIndex()} /> push()} /> void loadAll()} /> } /> + + void handleRebuildIndex()} /> + void loadAll()} /> + + } + /> + `${name}:${String(row["status"] ?? "unknown")}`) + .join(" ยท ") || "No model status"} + accessories={[ + { + text: runtime["apple_silicon"] ? "Apple Silicon" : "Compatibility mode", + }, + ]} + actions={ + + void loadAll()} /> + + } + /> {/* โ”€โ”€ Connectors โ”€โ”€ */} @@ -479,7 +547,7 @@ export default function ManageRecall() { push()} /> diff --git a/raycast/src/open-memory.tsx b/raycast/src/open-memory.tsx index 926c899..f39bc6a 100644 --- a/raycast/src/open-memory.tsx +++ b/raycast/src/open-memory.tsx @@ -56,7 +56,7 @@ export default async function OpenMemory(props: { arguments: { query: string } } } catch (err: unknown) { const msg = err instanceof Error ? err.message : String(err); if (msg.includes("GEMINI_API_KEY") || msg.includes("AUTH_ERROR")) { - await showHUD("โš  API key missing โ€” check extension preferences"); + await showHUD("โš  Optional enrichment key invalid โ€” check extension preferences"); await openExtensionPreferences(); } else if (msg.includes("daemon") || msg.includes("DAEMON")) { await showHUD("โš  Daemon not running โ€” run 'vef-daemon start' in terminal"); diff --git a/raycast/src/search-memory.tsx b/raycast/src/search-memory.tsx index 64af65a..ca3d70a 100644 --- a/raycast/src/search-memory.tsx +++ b/raycast/src/search-memory.tsx @@ -104,7 +104,7 @@ export default function SearchMemory() { showToast({ style: Toast.Style.Failure, title: "Auth Error", - message: "GEMINI_API_KEY missing or invalid โ€” check extension preferences", + message: "Optional cloud enrichment key is invalid โ€” local search still works once configuration is fixed", }); } else if (e.code === "RATE_LIMIT") { showToast({ style: Toast.Style.Failure, title: "Rate Limited", message: e.message }); diff --git a/setup_wizard.py b/setup_wizard.py index 8a6cfe5..75009d2 100644 --- a/setup_wizard.py +++ b/setup_wizard.py @@ -167,7 +167,7 @@ def _detect_embedding_provider() -> str: ) except Exception: provider = "" - return provider or "gemini" + return provider or "local" # โ”€โ”€ Screen 1: Splash โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ @@ -176,7 +176,7 @@ def _detect_embedding_provider() -> str: def screen_splash() -> None: CONSOLE.clear() CONSOLE.print() - CONSOLE.print(f" [{C['dim']}]โ•ญโ”€ vector-embedded-finder setup โ”€โ•ฎ[/]") + CONSOLE.print(f" [{C['dim']}]โ•ญโ”€ recall setup โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ•ฎ[/]") CONSOLE.print(f" [{C['dim']}]โ”‚ local multimodal memory โ”‚[/]") CONSOLE.print(f" [{C['dim']}]โ•ฐโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ•ฏ[/]") CONSOLE.print() @@ -211,11 +211,13 @@ def screen_detect() -> dict: # Required packages time.sleep(0.12) - pkg_checks = [("chromadb", "chromadb"), ("python-dotenv", "dotenv")] + pkg_checks = [("sqlite runtime", "sqlite3"), ("python-dotenv", "dotenv")] if provider == "gemini": pkg_checks.append(("google-genai", "google.genai")) elif provider in {"ollama", "nim"}: pkg_checks.append(("httpx", "httpx")) + else: + pkg_checks.append(("sentence-transformers (optional)", "sentence_transformers")) for display, import_name in pkg_checks: try: @@ -241,7 +243,7 @@ def screen_detect() -> dict: result["existing_key"] = existing_key CONSOLE.print(f" {warn('API key found [dim]โ†’ you can update or keep it[/dim]')}") else: - CONSOLE.print(f" {warn('no API key found [dim]โ†’ we will set it up[/dim]')}") + CONSOLE.print(f" {warn('no API key found [dim]โ†’ optional cloud enrichment will stay disabled[/dim]')}") else: CONSOLE.print( f" {ok('API key setup skipped [dim](non-Gemini provider selected)[/dim]')}" @@ -256,7 +258,7 @@ def screen_detect() -> dict: def screen_api_key(detected: dict) -> str: - provider = detected.get("embedding_provider", "gemini") + provider = detected.get("embedding_provider", "local") if provider != "gemini": CONSOLE.clear() step_header(1, 3, "EMBEDDING PROVIDER") @@ -268,7 +270,7 @@ def screen_api_key(detected: dict) -> str: return "" CONSOLE.clear() - step_header(1, 3, "GEMINI API KEY") + step_header(1, 3, "OPTIONAL GEMINI ENRICHMENT") existing = detected.get("existing_key", "") if existing: @@ -282,7 +284,7 @@ def screen_api_key(detected: dict) -> str: CONSOLE.print() return existing - CONSOLE.print(f" [{C['dim']}]You need a free Gemini key to generate embeddings.[/]") + CONSOLE.print(f" [{C['dim']}]Recall now embeds and searches locally. Gemini is optional for richer cloud captions.[/]") CONSOLE.print() CONSOLE.print(f" [{C['accent']}]โ†’ Get yours free at:[/] [underline]https://aistudio.google.com/apikey[/underline]") CONSOLE.print() @@ -912,7 +914,7 @@ def screen_done(detected: dict, stats: dict[str, int]) -> None: ), ( "Gemini API Key", - f"(set โ€” see {ENV_FILE.name} in repo root)" if provider == "gemini" else "(not required)", + f"(optional โ€” see {ENV_FILE.name} in repo root)" if provider == "gemini" else "(optional enrichment only)", C["warn"] if provider == "gemini" else C["dim"], ), ] diff --git a/tests/test_connectors.py b/tests/test_connectors.py index f10b680..81a174f 100644 --- a/tests/test_connectors.py +++ b/tests/test_connectors.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any +from vector_embedded_finder import keychain from vector_embedded_finder.connectors.canvas import CanvasConnector from vector_embedded_finder.connectors.gmail import GmailConnector @@ -13,6 +14,7 @@ def test_gmail_not_authenticated_without_token(monkeypatch: Any, tmp_path: Path) token_path = tmp_path / "gmail.json" monkeypatch.setattr(config, "GMAIL_CREDENTIALS_FILE", token_path, raising=False) + monkeypatch.setattr(keychain, "_security_available", lambda: False) assert GmailConnector().is_authenticated() is False @@ -21,6 +23,7 @@ def test_canvas_not_authenticated_without_creds(monkeypatch: Any, tmp_path: Path creds_path = tmp_path / "canvas.json" monkeypatch.setattr(config, "CANVAS_CREDENTIALS_FILE", creds_path, raising=False) + monkeypatch.setattr(keychain, "_security_available", lambda: False) assert CanvasConnector().is_authenticated() is False @@ -31,4 +34,10 @@ def test_gmail_is_authenticated_with_token(monkeypatch: Any, tmp_path: Path) -> token_path.parent.mkdir(parents=True, exist_ok=True) token_path.write_text(json.dumps({"token": "abc"})) monkeypatch.setattr(config, "GMAIL_CREDENTIALS_FILE", token_path, raising=False) + monkeypatch.setattr(keychain, "_security_available", lambda: False) + assert GmailConnector().is_authenticated() is True + + +def test_gmail_is_authenticated_with_keychain_secret(monkeypatch: Any) -> None: + monkeypatch.setattr(keychain, "load_json", lambda *args, **kwargs: {"token": "abc"}) assert GmailConnector().is_authenticated() is True diff --git a/tests/test_daemon_http.py b/tests/test_daemon_http.py index c2d9ac8..2e1c329 100644 --- a/tests/test_daemon_http.py +++ b/tests/test_daemon_http.py @@ -10,22 +10,21 @@ def _client(monkeypatch: Any, tmp_path: Path) -> TestClient: - from vector_embedded_finder import config, embedder, store + from vector_embedded_finder import config, embedder, migration, model_manager, store search_mod = importlib.import_module("vector_embedded_finder.search") monkeypatch.setattr(config, "EMBEDDING_PROVIDER", "ollama", raising=False) monkeypatch.setattr(config, "WATCHED_DIRS_FILE", tmp_path / "watched_dirs.json", raising=False) monkeypatch.setattr(config, "ensure_vef_dirs", lambda: None, raising=False) + monkeypatch.setattr(config, "ensure_runtime_dirs", lambda: None, raising=False) monkeypatch.setattr(embedder, "warmup_provider", lambda: None) + monkeypatch.setattr(model_manager, "warmup", lambda: None) + monkeypatch.setattr(migration, "ensure_migrated", lambda: {"status": "complete"}) + monkeypatch.setattr(migration, "status", lambda: {"status": "complete"}) + monkeypatch.setattr(store, "initialize", lambda: None) + monkeypatch.setattr(store, "index_status", lambda: {"backend": "memory", "count": 0, "ready": True, "dirty": False}) + monkeypatch.setattr(store, "get_sources", lambda: []) - class _Coll: - def count(self) -> int: - return 0 - - def get(self, limit: int = 0, include: list[str] | None = None) -> dict[str, list[dict]]: - return {"metadatas": []} - - monkeypatch.setattr(store, "_get_collection", lambda: _Coll()) monkeypatch.setattr(store, "count", lambda: 0) monkeypatch.setattr(search_mod, "search", lambda *_a, **_k: []) monkeypatch.setattr(daemon, "_run_connector_sync_once", lambda **_k: {}) diff --git a/tests/test_ingest.py b/tests/test_ingest.py index 0b6e29f..271ef43 100644 --- a/tests/test_ingest.py +++ b/tests/test_ingest.py @@ -28,15 +28,70 @@ def test_text_file_ingested(monkeypatch: Any, tmp_path: Path, fake_embedding: li monkeypatch.setattr(ingest.store, "exists", lambda _doc_id: False) monkeypatch.setattr(ingest.embedder, "embed_text", lambda _text: fake_embedding) - def fake_add(doc_id: str, embedding: list[float], metadata: dict[str, Any], document: str = "") -> None: + def fake_add( + doc_id: str, + embedding: list[float], + metadata: dict[str, Any], + document: str = "", + enrichment: dict[str, Any] | None = None, + ) -> None: captured["doc_id"] = doc_id captured["embedding"] = embedding captured["metadata"] = metadata captured["document"] = document + captured["enrichment"] = enrichment or {} monkeypatch.setattr(ingest.store, "add", fake_add) + monkeypatch.setattr( + ingest.store, + "retire_path_versions", + lambda path, *, keep_doc_id: captured.__setitem__("retired", (str(path), keep_doc_id)), + ) result = ingest.ingest_file(p, source="manual") assert result["status"] == "embedded" assert result["category"] == "text" assert captured["metadata"]["file_name"] == "note.txt" + assert captured["enrichment"] == { + "caption": "", + "ocr_text": "", + "gps_city": "", + "face_count": 0, + "exif_date": "", + "exif_camera": "", + } + assert captured["retired"] == (str(p.resolve()), captured["doc_id"]) + + +def test_ingest_failure_does_not_retire_path_versions( + monkeypatch: Any, + tmp_path: Path, +) -> None: + p = tmp_path / "note.txt" + p.write_text("hello world") + + calls = {"add": 0, "retire": 0} + monkeypatch.setattr(ingest.store, "exists", lambda _doc_id: False) + monkeypatch.setattr( + ingest.store, + "add", + lambda *_a, **_k: calls.__setitem__("add", calls["add"] + 1), + ) + monkeypatch.setattr( + ingest.store, + "retire_path_versions", + lambda *_a, **_k: calls.__setitem__("retire", calls["retire"] + 1), + ) + monkeypatch.setattr( + ingest.embedder, + "embed_text", + lambda _text: (_ for _ in ()).throw(RuntimeError("embed fail")), + ) + + try: + ingest.ingest_file(p, source="manual") + except RuntimeError: + pass + + assert calls["add"] == 0 + assert calls["retire"] == 0 diff --git a/tests/test_keychain.py b/tests/test_keychain.py new file mode 100644 index 0000000..eb90d7c --- /dev/null +++ b/tests/test_keychain.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from vector_embedded_finder import keychain + + +def test_load_json_migrates_legacy_file_to_keychain(monkeypatch: Any, tmp_path: Path) -> None: + legacy = tmp_path / "gmail.json" + legacy.write_text(json.dumps({"token": "abc"})) + + calls: list[list[str]] = [] + + def fake_run(args: list[str]) -> Any: + calls.append(args) + if args[0] == "find-generic-password": + return type("Result", (), {"returncode": 44, "stdout": "", "stderr": "not found"})() + if args[0] == "add-generic-password": + return type("Result", (), {"returncode": 0, "stdout": "", "stderr": ""})() + raise AssertionError(f"Unexpected command: {args}") + + monkeypatch.setattr(keychain, "_security_available", lambda: True) + monkeypatch.setattr(keychain, "_run_security", fake_run) + + payload = keychain.load_json("gmail", legacy_path=legacy) + + assert payload == {"token": "abc"} + assert not legacy.exists() + assert calls == [ + ["find-generic-password", "-a", keychain.DEFAULT_ACCOUNT, "-s", "com.recall.credentials.gmail", "-w"], + [ + "add-generic-password", + "-U", + "-a", + keychain.DEFAULT_ACCOUNT, + "-s", + "com.recall.credentials.gmail", + "-w", + json.dumps({"token": "abc"}, sort_keys=True), + ], + ] + + +def test_load_json_prefers_existing_keychain_secret(monkeypatch: Any, tmp_path: Path) -> None: + legacy = tmp_path / "gmail.json" + legacy.write_text(json.dumps({"token": "legacy"})) + + def fake_run(args: list[str]) -> Any: + if args[0] == "find-generic-password": + return type("Result", (), {"returncode": 0, "stdout": json.dumps({"token": "secure"}), "stderr": ""})() + raise AssertionError(f"Unexpected command: {args}") + + monkeypatch.setattr(keychain, "_security_available", lambda: True) + monkeypatch.setattr(keychain, "_run_security", fake_run) + + payload = keychain.load_json("gmail", legacy_path=legacy) + + assert payload == {"token": "secure"} + assert legacy.exists() diff --git a/tests/test_search.py b/tests/test_search.py index 4bd53d6..5c66391 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -3,37 +3,42 @@ import importlib from typing import Any +from vector_embedded_finder.store import Candidate + search_mod = importlib.import_module("vector_embedded_finder.search") -def _raw_result( +def _candidate( *, doc_id: str = "id-1", - distance: float = 0.6, + score: float = 0.6, file_name: str = "draft.txt", description: str = "some text", media_category: str = "text", source: str = "manual", -) -> dict[str, Any]: - return { - "ids": [[doc_id]], - "metadatas": [[{ +) -> Candidate: + return Candidate( + file_id=1, + doc_id=doc_id, + distance=max(0.0, 1.0 - score), + score=score, + metadata={ "file_path": "/tmp/draft.txt", "file_name": file_name, "media_category": media_category, "timestamp": "2025-01-01T00:00:00+00:00", "description": description, "source": source, - }]], - "documents": [["body text"]], - "distances": [[distance]], - } + "preview": "body text", + }, + ) def test_similarity_threshold_filters_low_scores(monkeypatch: Any, fake_embedding: list[float]) -> None: monkeypatch.setattr(search_mod.embedder, "embed_query", lambda _q: fake_embedding) monkeypatch.setattr(search_mod, "MIN_SIMILARITY", 0.45) - monkeypatch.setattr(search_mod.store, "search", lambda *_a, **_k: _raw_result(distance=0.7)) + monkeypatch.setattr(search_mod.store, "dense_search", lambda *_a, **_k: [_candidate(score=0.3)]) + monkeypatch.setattr(search_mod.store, "keyword_search", lambda *_a, **_k: []) results = search_mod.search("anything", n_results=5) assert results == [] @@ -42,21 +47,28 @@ def test_similarity_threshold_filters_low_scores(monkeypatch: Any, fake_embeddin def test_keyword_boost_raises_score(monkeypatch: Any, fake_embedding: list[float]) -> None: monkeypatch.setattr(search_mod.embedder, "embed_query", lambda _q: fake_embedding) monkeypatch.setattr(search_mod, "MIN_SIMILARITY", 0.0) - - def fake_search(*_a: Any, **kwargs: Any) -> dict[str, Any]: - if kwargs.get("where_document"): - return _raw_result( - distance=0.6, + monkeypatch.setattr( + search_mod.store, + "dense_search", + lambda *_a, **_k: [ + _candidate( + score=0.4, file_name="plasma-wound-report.txt", description="plasma treatment notes", ) - return _raw_result( - distance=0.6, - file_name="plasma-wound-report.txt", - description="plasma treatment notes", - ) - - monkeypatch.setattr(search_mod.store, "search", fake_search) + ], + ) + monkeypatch.setattr( + search_mod.store, + "keyword_search", + lambda *_a, **_k: [ + _candidate( + score=0.41, + file_name="plasma-wound-report.txt", + description="plasma treatment notes", + ) + ], + ) results = search_mod.search("plasma report", n_results=5) assert results assert float(results[0]["similarity"]) > 0.4 @@ -66,43 +78,80 @@ def test_intent_filter_image(monkeypatch: Any, fake_embedding: list[float]) -> N monkeypatch.setattr(search_mod.embedder, "embed_query", lambda _q: fake_embedding) calls: list[dict[str, Any]] = [] - def fake_search(_embedding: list[float], **kwargs: Any) -> dict[str, Any]: + def fake_dense(_embedding: list[float], **kwargs: Any): calls.append(kwargs) - return {"ids": [[]], "metadatas": [[]], "documents": [[]], "distances": [[]]} + return [] - monkeypatch.setattr(search_mod.store, "search", fake_search) + monkeypatch.setattr(search_mod.store, "dense_search", fake_dense) + monkeypatch.setattr(search_mod.store, "keyword_search", lambda *_a, **_k: []) _ = search_mod.search("photo of sunset") assert calls - where = calls[0].get("where") - assert where == {"media_category": {"$eq": "image"}} + filters = calls[0].get("filters") + assert filters == {"media_category": "image"} def test_intent_filter_email(monkeypatch: Any, fake_embedding: list[float]) -> None: monkeypatch.setattr(search_mod.embedder, "embed_query", lambda _q: fake_embedding) calls: list[dict[str, Any]] = [] - def fake_search(_embedding: list[float], **kwargs: Any) -> dict[str, Any]: + def fake_dense(_embedding: list[float], **kwargs: Any): calls.append(kwargs) - return {"ids": [[]], "metadatas": [[]], "documents": [[]], "distances": [[]]} + return [] - monkeypatch.setattr(search_mod.store, "search", fake_search) + monkeypatch.setattr(search_mod.store, "dense_search", fake_dense) + monkeypatch.setattr(search_mod.store, "keyword_search", lambda *_a, **_k: []) _ = search_mod.search("email from john") assert calls - where = calls[0].get("where") - assert isinstance(where, dict) - assert "$and" in where - terms = where["$and"] - assert {"media_category": {"$eq": "email"}} in terms - assert {"source": {"$eq": "gmail"}} in terms + filters = calls[0].get("filters") or {} + assert filters.get("media_category") == "email" + assert filters.get("sources") == ["gmail"] def test_no_results_empty_db(monkeypatch: Any, fake_embedding: list[float]) -> None: + monkeypatch.setattr(search_mod.embedder, "embed_query", lambda _q: fake_embedding) + monkeypatch.setattr(search_mod.store, "dense_search", lambda *_a, **_k: []) + monkeypatch.setattr(search_mod.store, "keyword_search", lambda *_a, **_k: []) + assert search_mod.search("nothing here") == [] + + +def test_query_embedding_cache_skips_second_embed(monkeypatch: Any, fake_embedding: list[float]) -> None: + calls: list[str] = [] + search_mod._embed_query_cached.cache_clear() + + def fake_embed(query: str) -> list[float]: + calls.append(query) + return fake_embedding + + monkeypatch.setattr(search_mod.embedder, "embed_query", fake_embed) + monkeypatch.setattr(search_mod.store, "dense_search", lambda *_a, **_k: []) + monkeypatch.setattr(search_mod.store, "keyword_search", lambda *_a, **_k: []) + + assert search_mod.search("repeat query") == [] + assert search_mod.search("repeat query") == [] + assert calls == ["repeat query"] + + +def test_result_cache_invalidates_on_store_epoch(monkeypatch: Any, fake_embedding: list[float]) -> None: + search_mod._embed_query_cached.cache_clear() + search_mod._RESULT_CACHE.clear() + + epoch = {"value": 1} + state = {"doc_id": "doc-v1"} + monkeypatch.setattr(search_mod.store, "cache_epoch", lambda: int(epoch["value"])) monkeypatch.setattr(search_mod.embedder, "embed_query", lambda _q: fake_embedding) monkeypatch.setattr( search_mod.store, - "search", - lambda *_a, **_k: {"ids": [[]], "metadatas": [[]], "documents": [[]], "distances": [[]]}, + "dense_search", + lambda *_a, **_k: [_candidate(doc_id=str(state["doc_id"]), score=0.9)], ) - assert search_mod.search("nothing here") == [] + monkeypatch.setattr(search_mod.store, "keyword_search", lambda *_a, **_k: []) + + first = search_mod.search("epoch test") + state["doc_id"] = "doc-v2" + epoch["value"] = 2 + second = search_mod.search("epoch test") + + assert first[0]["id"] == "doc-v1" + assert second[0]["id"] == "doc-v2" diff --git a/vector_embedded_finder/cli.py b/vector_embedded_finder/cli.py index 1e2cb90..265c81c 100644 --- a/vector_embedded_finder/cli.py +++ b/vector_embedded_finder/cli.py @@ -42,7 +42,7 @@ def _console() -> Console | None: def _daemon_base_url() -> str: - return f"http://{config.DAEMON_HOST}:{config.DAEMON_PORT}" + return config.RECALL_SOCKET_BASE_URL def _render_results(results: list[dict[str, Any]]) -> None: @@ -62,11 +62,12 @@ def _run_daemon_command(args: list[str]) -> int: def _fetch_json(path: str, method: str = "GET", payload: dict[str, Any] | None = None, timeout: float = 30.0) -> dict[str, Any]: url = f"{_daemon_base_url()}{path}" - with httpx.Client(timeout=timeout) as client: + transport = httpx.HTTPTransport(uds=str(config.SOCKET_PATH)) + with httpx.Client(transport=transport, base_url=_daemon_base_url(), timeout=timeout) as client: if method == "GET": - resp = client.get(url) + resp = client.get(path) else: - resp = client.post(url, json=payload or {}) + resp = client.post(path, json=payload or {}) resp.raise_for_status() data = resp.json() if not isinstance(data, dict): diff --git a/vector_embedded_finder/config.py b/vector_embedded_finder/config.py index a14b354..cc27814 100644 --- a/vector_embedded_finder/config.py +++ b/vector_embedded_finder/config.py @@ -1,85 +1,151 @@ from __future__ import annotations import os +import platform from pathlib import Path from dotenv import load_dotenv PROJECT_DIR = Path(__file__).parent.parent load_dotenv(PROJECT_DIR / ".env") -DEFAULT_VEF_DIR = Path(os.environ.get("VEF_DIR", str(Path.home() / ".vef"))) -load_dotenv(DEFAULT_VEF_DIR / ".env", override=True) -DATA_DIR = Path(os.environ.get("VEF_DATA_DIR", str(PROJECT_DIR / "data"))) -CHROMA_DIR = DATA_DIR / "chromadb" -EMBEDDING_PROVIDER = os.environ.get("VEF_EMBEDDING_PROVIDER", "gemini").strip().lower() -EMBEDDING_MODEL = os.environ.get("VEF_EMBEDDING_MODEL", "gemini-embedding-2-preview") -EMBEDDING_DIMENSIONS = int(os.environ.get("VEF_EMBEDDING_DIMENSIONS", "768")) +def _env_bool(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +LEGACY_VEF_DIR = Path(os.environ.get("VEF_DIR", str(Path.home() / ".vef"))) +RECALL_HOME = Path( + os.environ.get("RECALL_HOME") + or os.environ.get("VEF_DIR") + or str(Path.home() / ".recall") +) + +load_dotenv(LEGACY_VEF_DIR / ".env", override=False) +load_dotenv(RECALL_HOME / ".env", override=True) + +DB_DIR = RECALL_HOME / "db" +HNSW_DIR = RECALL_HOME / "hnsw" +MODELS_DIR = RECALL_HOME / "models" +CACHE_DIR = RECALL_HOME / "cache" +LOG_DIR = RECALL_HOME / "logs" +CREDENTIALS_DIR = RECALL_HOME / "credentials" + +SQLITE_PATH = DB_DIR / "recall.db" +MODEL_MANIFEST_PATH = MODELS_DIR / "manifest.json" +MIGRATION_STATUS_PATH = RECALL_HOME / "migration_status.json" +WATCHED_DIRS_FILE = RECALL_HOME / "watched_dirs.json" +PID_FILE = RECALL_HOME / "daemon.pid" +SOCKET_PATH = Path( + os.environ.get("RECALL_SOCKET_PATH", str(RECALL_HOME / "recall.sock")) +) + +# Backward-compatible aliases used throughout the existing repo. +VEF_DIR = RECALL_HOME +DEFAULT_VEF_DIR = LEGACY_VEF_DIR + +LEGACY_DATA_DIR = Path(os.environ.get("VEF_DATA_DIR", str(PROJECT_DIR / "data"))) +CHROMA_DIR = LEGACY_DATA_DIR / "chromadb" + +EMBEDDING_PROVIDER = os.environ.get("RECALL_EMBEDDING_PROVIDER", "").strip().lower() +if not EMBEDDING_PROVIDER: + EMBEDDING_PROVIDER = os.environ.get("VEF_EMBEDDING_PROVIDER", "local").strip().lower() +if EMBEDDING_PROVIDER not in {"local", "gemini", "ollama", "nim"}: + EMBEDDING_PROVIDER = "local" + +EMBEDDING_MODEL = os.environ.get( + "RECALL_EMBEDDING_MODEL", + os.environ.get("VEF_EMBEDDING_MODEL", "nomic-ai/nomic-embed-text-v1.5"), +) +VISION_EMBEDDING_MODEL = os.environ.get( + "RECALL_VISION_EMBEDDING_MODEL", + "nomic-ai/nomic-embed-vision-v1.5", +) +RERANKER_MODEL = os.environ.get( + "RECALL_RERANKER_MODEL", + "cross-encoder/ms-marco-MiniLM-L6-v2", +) +EMBEDDING_DIMENSIONS = int( + os.environ.get( + "RECALL_EMBEDDING_DIMENSIONS", + os.environ.get("VEF_EMBEDDING_DIMENSIONS", "768"), + ) +) MAX_TEXT_TOKENS = 8192 -OLLAMA_BASE_URL = os.environ.get("VEF_OLLAMA_BASE_URL", "http://127.0.0.1:11434").rstrip("/") +OLLAMA_BASE_URL = os.environ.get( + "VEF_OLLAMA_BASE_URL", + os.environ.get("RECALL_OLLAMA_BASE_URL", "http://127.0.0.1:11434"), +).rstrip("/") OLLAMA_EMBED_MODEL = os.environ.get("VEF_OLLAMA_EMBED_MODEL", "nomic-embed-text") -OLLAMA_EMBED_URL = os.environ.get("VEF_OLLAMA_EMBED_URL", f"{OLLAMA_BASE_URL}/api/embeddings") +OLLAMA_EMBED_URL = os.environ.get( + "VEF_OLLAMA_EMBED_URL", + f"{OLLAMA_BASE_URL}/api/embeddings", +) NIM_EMBED_URL = os.environ.get("VEF_NIM_EMBED_URL", "").strip() NIM_EMBED_MODEL = os.environ.get("VEF_NIM_EMBED_MODEL", "nvidia/nv-embedqa-e5-v5") -if "VEF_COLLECTION_NAME" in os.environ: - COLLECTION_NAME = os.environ["VEF_COLLECTION_NAME"] -elif EMBEDDING_PROVIDER == "gemini": - COLLECTION_NAME = "vector_embedded_finder" -elif EMBEDDING_PROVIDER == "ollama": - COLLECTION_NAME = "vector_embedded_finder_ollama" -elif EMBEDDING_PROVIDER == "nim": - COLLECTION_NAME = "vector_embedded_finder_nim" -else: - COLLECTION_NAME = f"vector_embedded_finder_{EMBEDDING_PROVIDER}" - -# Credentials and runtime state directory -VEF_DIR = Path(os.environ.get("VEF_DIR", str(DEFAULT_VEF_DIR))) -CREDENTIALS_DIR = VEF_DIR / "credentials" -PID_FILE = VEF_DIR / "daemon.pid" +COLLECTION_NAME = os.environ.get("VEF_COLLECTION_NAME", "vector_embedded_finder") SUPPORTED_EXTENSIONS = { "image": {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff"}, "audio": {".mp3", ".wav", ".m4a", ".ogg", ".flac", ".aac"}, "video": {".mp4", ".mov", ".avi", ".mkv", ".webm"}, "document": {".pdf"}, - "text": {".txt", ".md", ".csv", ".json", ".yaml", ".yml", ".toml", ".xml", ".html", ".py", ".js", ".ts", ".go", ".rs", ".sh"}, + "text": { + ".txt", + ".md", + ".csv", + ".json", + ".yaml", + ".yml", + ".toml", + ".xml", + ".html", + ".py", + ".js", + ".ts", + ".go", + ".rs", + ".sh", + }, } -ALL_EXTENSIONS = set() +ALL_EXTENSIONS: set[str] = set() for exts in SUPPORTED_EXTENSIONS.values(): ALL_EXTENSIONS.update(exts) -# โ”€โ”€ Resource limits โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - -# Number of concurrent ingest workers in the thread pool +# Resource limits MAX_CONCURRENT_INGEST = int(os.environ.get("VEF_CONCURRENCY", "10")) - -# Daemon HTTP port -DAEMON_PORT = int(os.environ.get("VEF_PORT", "19847")) -DAEMON_HOST = "127.0.0.1" - -# CPU guard: back off if sustained CPU exceeds this percentage -CPU_GUARD_PERCENT = 30 - -# Memory guard: require at least this many bytes free before loading AI models -MIN_FREE_RAM_BYTES = 8 * 1024 ** 3 # 8 GB - -# โ”€โ”€ Connector sync intervals (seconds) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - -GMAIL_POLL_INTERVAL = 15 * 60 # 15 minutes -GCAL_POLL_INTERVAL = 30 * 60 # 30 minutes -CALAI_POLL_INTERVAL = 30 * 60 # 30 minutes -LMS_POLL_INTERVAL = 60 * 60 # 60 minutes -GDRIVE_POLL_INTERVAL = 30 * 60 # 30 minutes -NOTION_POLL_INTERVAL = 30 * 60 # 30 minutes +CPU_GUARD_PERCENT = int(os.environ.get("RECALL_CPU_GUARD_PERCENT", "30")) +MIN_FREE_RAM_BYTES = 8 * 1024**3 + +# Daemon transport +DAEMON_HOST = os.environ.get("RECALL_COMPAT_HOST", "127.0.0.1") +DAEMON_PORT = int(os.environ.get("VEF_PORT", os.environ.get("RECALL_COMPAT_PORT", "19847"))) +RECALL_ENABLE_COMPAT_HTTP = _env_bool("RECALL_ENABLE_COMPAT_HTTP", default=False) +RECALL_SOCKET_BASE_URL = "http://recall.local" + +# Search/index flags +DUAL_WRITE_CHROMA = _env_bool("RECALL_DUAL_WRITE_CHROMA", default=True) +READ_FROM_CHROMA = _env_bool("RECALL_READ_FROM_CHROMA", default=False) +ENABLE_CLOUD_ENRICHMENT = _env_bool("RECALL_ENABLE_CLOUD_ENRICHMENT", default=False) +ENABLE_OPTIONAL_CAPTIONING = _env_bool("RECALL_ENABLE_OPTIONAL_CAPTIONING", default=True) +INDEX_REBUILD_BATCH = int(os.environ.get("RECALL_INDEX_REBUILD_BATCH", "500")) + +# Connector sync intervals (seconds) +GMAIL_POLL_INTERVAL = 15 * 60 +GCAL_POLL_INTERVAL = 30 * 60 +CALAI_POLL_INTERVAL = 30 * 60 +LMS_POLL_INTERVAL = 60 * 60 +GDRIVE_POLL_INTERVAL = 30 * 60 +NOTION_POLL_INTERVAL = 30 * 60 CONNECTOR_SYNC_BUDGET_S = float(os.environ.get("VEF_CONNECTOR_SYNC_BUDGET_S", "600")) -# โ”€โ”€ Connector credential files โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - GMAIL_CREDENTIALS_FILE = CREDENTIALS_DIR / "gmail.json" CANVAS_CREDENTIALS_FILE = CREDENTIALS_DIR / "canvas.json" CALAI_CREDENTIALS_FILE = CREDENTIALS_DIR / "calai.json" @@ -87,16 +153,12 @@ GDRIVE_CREDENTIALS_FILE = CREDENTIALS_DIR / "gdrive.json" NOTION_CREDENTIALS_FILE = CREDENTIALS_DIR / "notion.json" -# โ”€โ”€ Watched directories (populated by setup wizard) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - -WATCHED_DIRS_FILE = VEF_DIR / "watched_dirs.json" - def get_api_key() -> str: key = os.environ.get("GEMINI_API_KEY", "") if not key: raise ValueError( - "GEMINI_API_KEY not set. Add it to .env or set as environment variable." + "GEMINI_API_KEY not set. Add it to ~/.recall/.env or configure it in Recall." ) return key @@ -104,9 +166,7 @@ def get_api_key() -> str: def get_nim_api_key() -> str: key = os.environ.get("NIM_API_KEY", "") if not key: - raise ValueError( - "NIM_API_KEY not set. Add it to .env or set as environment variable." - ) + raise ValueError("NIM_API_KEY not set. Add it to ~/.recall/.env.") return key @@ -119,6 +179,25 @@ def get_media_category(ext: str) -> str | None: def ensure_vef_dirs() -> None: - """Create ~/.vef and subdirectories if they don't exist.""" - VEF_DIR.mkdir(parents=True, exist_ok=True) + ensure_runtime_dirs() + + +def ensure_runtime_dirs() -> None: + RECALL_HOME.mkdir(parents=True, exist_ok=True) + DB_DIR.mkdir(parents=True, exist_ok=True) + HNSW_DIR.mkdir(parents=True, exist_ok=True) + MODELS_DIR.mkdir(parents=True, exist_ok=True) + CACHE_DIR.mkdir(parents=True, exist_ok=True) + LOG_DIR.mkdir(parents=True, exist_ok=True) CREDENTIALS_DIR.mkdir(parents=True, exist_ok=True) + if SOCKET_PATH.exists() and SOCKET_PATH.is_dir(): + raise RuntimeError(f"Socket path points to a directory: {SOCKET_PATH}") + + +def is_apple_silicon() -> bool: + machine = platform.machine().lower() + return machine in {"arm64", "aarch64"} and sys_platform_is_macos() + + +def sys_platform_is_macos() -> bool: + return platform.system().lower() == "darwin" diff --git a/vector_embedded_finder/connectors/calai.py b/vector_embedded_finder/connectors/calai.py index b0d0320..14ad8f6 100644 --- a/vector_embedded_finder/connectors/calai.py +++ b/vector_embedded_finder/connectors/calai.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import Callable -from .. import config, embedder, store, utils +from .. import config, embedder, keychain, store, utils from .base import BaseConnector logger = logging.getLogger(__name__) @@ -21,16 +21,14 @@ def _creds_path() -> Path: def _load_api_key() -> str | None: - p = _creds_path() - if p.exists(): - data = json.loads(p.read_text()) + data = keychain.load_json("calai", legacy_path=_creds_path()) + if data: return data.get("api_key") return None def _save_api_key(key: str) -> None: - config.ensure_vef_dirs() - _creds_path().write_text(json.dumps({"api_key": key}, indent=2)) + keychain.save_json("calai", {"api_key": key}, legacy_path=_creds_path()) class CalAIConnector(BaseConnector): diff --git a/vector_embedded_finder/connectors/canvas.py b/vector_embedded_finder/connectors/canvas.py index d3fe51e..e62893a 100644 --- a/vector_embedded_finder/connectors/canvas.py +++ b/vector_embedded_finder/connectors/canvas.py @@ -10,7 +10,7 @@ from pathlib import Path from typing import Callable -from .. import config, embedder, store, utils +from .. import config, embedder, keychain, store, utils from .base import BaseConnector logger = logging.getLogger(__name__) @@ -21,15 +21,15 @@ def _creds_path() -> Path: def _load_creds() -> dict | None: - p = _creds_path() - if p.exists(): - return json.loads(p.read_text()) - return None + return keychain.load_json("canvas", legacy_path=_creds_path()) def _save_creds(token: str, base_url: str) -> None: - config.ensure_vef_dirs() - _creds_path().write_text(json.dumps({"token": token, "base_url": base_url}, indent=2)) + keychain.save_json( + "canvas", + {"token": token, "base_url": base_url}, + legacy_path=_creds_path(), + ) def _strip_html(html: str) -> str: diff --git a/vector_embedded_finder/connectors/gcal.py b/vector_embedded_finder/connectors/gcal.py index 9f0c2f1..6b42848 100644 --- a/vector_embedded_finder/connectors/gcal.py +++ b/vector_embedded_finder/connectors/gcal.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import Callable -from .. import config, embedder, store, utils +from .. import config, embedder, keychain, store, utils from .base import BaseConnector logger = logging.getLogger(__name__) @@ -43,15 +43,11 @@ def _oauth_client_path() -> Path: def _load_token() -> dict | None: - p = _creds_path() - if p.exists(): - return json.loads(p.read_text()) - return None + return keychain.load_json("gmail", legacy_path=_creds_path()) def _save_token(data: dict) -> None: - config.ensure_vef_dirs() - _creds_path().write_text(json.dumps(data, indent=2)) + keychain.save_json("gmail", data, legacy_path=_creds_path()) class GCalConnector(BaseConnector): diff --git a/vector_embedded_finder/connectors/gdrive.py b/vector_embedded_finder/connectors/gdrive.py index 630224b..6105026 100644 --- a/vector_embedded_finder/connectors/gdrive.py +++ b/vector_embedded_finder/connectors/gdrive.py @@ -11,7 +11,7 @@ import pypdf -from .. import config, embedder, store, utils +from .. import config, embedder, keychain, store, utils from .base import BaseConnector logger = logging.getLogger(__name__) @@ -49,15 +49,11 @@ def _oauth_client_path() -> Path: def _load_token() -> dict | None: - p = _token_path() - if p.exists(): - return json.loads(p.read_text()) - return None + return keychain.load_json("gmail", legacy_path=_token_path()) def _save_token(data: dict) -> None: - config.ensure_vef_dirs() - _token_path().write_text(json.dumps(data, indent=2)) + keychain.save_json("gmail", data, legacy_path=_token_path()) class GDriveConnector(BaseConnector): diff --git a/vector_embedded_finder/connectors/gmail.py b/vector_embedded_finder/connectors/gmail.py index 1dbb63d..991903b 100644 --- a/vector_embedded_finder/connectors/gmail.py +++ b/vector_embedded_finder/connectors/gmail.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Callable -from .. import config, embedder, store, utils +from .. import config, embedder, keychain, store, utils from .base import BaseConnector logger = logging.getLogger(__name__) @@ -42,15 +42,11 @@ def _oauth_client_path() -> Path: def _load_token() -> dict | None: - p = _creds_path() - if p.exists(): - return json.loads(p.read_text()) - return None + return keychain.load_json("gmail", legacy_path=_creds_path()) def _save_token(data: dict) -> None: - config.ensure_vef_dirs() - _creds_path().write_text(json.dumps(data, indent=2)) + keychain.save_json("gmail", data, legacy_path=_creds_path()) class GmailConnector(BaseConnector): diff --git a/vector_embedded_finder/connectors/notion.py b/vector_embedded_finder/connectors/notion.py index 9ccdc5a..2e05f37 100644 --- a/vector_embedded_finder/connectors/notion.py +++ b/vector_embedded_finder/connectors/notion.py @@ -10,7 +10,7 @@ import httpx -from .. import config, embedder, store, utils +from .. import config, embedder, keychain, store, utils from .base import BaseConnector logger = logging.getLogger(__name__) @@ -25,16 +25,14 @@ def _creds_path() -> Path: def _load_api_key() -> str | None: - p = _creds_path() - if p.exists(): - data = json.loads(p.read_text()) + data = keychain.load_json("notion", legacy_path=_creds_path()) + if data: return data.get("api_key") return None def _save_api_key(key: str) -> None: - config.ensure_vef_dirs() - _creds_path().write_text(json.dumps({"api_key": key}, indent=2)) + keychain.save_json("notion", {"api_key": key}, legacy_path=_creds_path()) def _headers(api_key: str) -> dict[str, str]: diff --git a/vector_embedded_finder/connectors/schoology.py b/vector_embedded_finder/connectors/schoology.py index 6fc41bc..e2c2f7c 100644 --- a/vector_embedded_finder/connectors/schoology.py +++ b/vector_embedded_finder/connectors/schoology.py @@ -10,7 +10,7 @@ from pathlib import Path from typing import Callable -from .. import config, embedder, store, utils +from .. import config, embedder, keychain, store, utils from .base import BaseConnector logger = logging.getLogger(__name__) @@ -23,19 +23,19 @@ def _creds_path() -> Path: def _load_creds() -> dict | None: - p = _creds_path() - if p.exists(): - return json.loads(p.read_text()) - return None + return keychain.load_json("schoology", legacy_path=_creds_path()) def _save_creds(consumer_key: str, consumer_secret: str, base_url: str) -> None: - config.ensure_vef_dirs() - _creds_path().write_text(json.dumps({ - "consumer_key": consumer_key, - "consumer_secret": consumer_secret, - "base_url": base_url.rstrip("/"), - }, indent=2)) + keychain.save_json( + "schoology", + { + "consumer_key": consumer_key, + "consumer_secret": consumer_secret, + "base_url": base_url.rstrip("/"), + }, + legacy_path=_creds_path(), + ) def _strip_html(html: str) -> str: diff --git a/vector_embedded_finder/daemon.py b/vector_embedded_finder/daemon.py index fff7dcc..4776e9e 100644 --- a/vector_embedded_finder/daemon.py +++ b/vector_embedded_finder/daemon.py @@ -1,18 +1,8 @@ -"""Persistent search daemon โ€” FastAPI server on 127.0.0.1:19847. - -CLI: vef-daemon start | stop | status | sync [source] | check-embed - -Design notes: -- Uses FastAPI lifespan (not the deprecated on_event) for startup/shutdown. -- Sources are tracked in an in-memory set (_known_sources), populated at - startup from a sampled scan and updated on every /ingest call. This avoids - the O(n) full-collection scan that the naive implementation would require. -- Connector syncs run in a dedicated background thread (_connector_sync_loop). - The thread checks every 60s and only syncs when the daemon has been idle - (no /search requests) for at least 30 seconds. -- cmd_start polls /health for up to 6 seconds so "Daemon started" only - prints when the daemon is actually healthy, not just when the process spawns. -- stderr is redirected to ~/.vef/daemon.log so crashes are diagnosable. +"""Persistent Recall daemon. + +Primary transport is a Unix domain socket under ~/.recall/recall.sock. +Optional localhost HTTP compatibility can be enabled with +RECALL_ENABLE_COMPAT_HTTP=1 for migration and diagnostics. """ from __future__ import annotations @@ -20,9 +10,7 @@ import json import logging import os -import errno import signal -import socket import sys import threading import time as _time @@ -31,21 +19,24 @@ from pathlib import Path from typing import Callable, List, Optional -logger = logging.getLogger(__name__) +import httpx -# โ”€โ”€ Shared state โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ +logger = logging.getLogger(__name__) -_last_search_time: float = 0.0 # updated by /search; read by sync thread -_known_sources: set[str] = set() # populated at startup; updated on ingest +_last_search_time: float = 0.0 _sync_lock = threading.Lock() -_sync_done = threading.Event() # set whenever a sync run completes -_last_sync_result: dict[str, dict] = {} # result of the most recent sync run +_sync_done = threading.Event() +_last_sync_result: dict[str, dict] = {} _last_connector_sync: dict[str, float] = {} _ingest_lock = threading.Lock() _ingest_in_flight = 0 +_watcher = None +_watcher_queue_depth: Callable[[], int] | None = None +_watcher_lock = threading.Lock() from . import config as _config -SYNC_STATE_FILE = _config.VEF_DIR / "sync_state.json" + +SYNC_STATE_FILE = _config.RECALL_HOME / "sync_state.json" def _load_sync_state() -> dict[str, float]: @@ -62,7 +53,7 @@ def _load_sync_state() -> dict[str, float]: def _save_sync_state() -> None: try: - _config.ensure_vef_dirs() + _config.ensure_runtime_dirs() payload = {name: float(ts) for name, ts in _last_connector_sync.items()} SYNC_STATE_FILE.write_text(json.dumps(payload, indent=2, sort_keys=True)) except Exception as exc: @@ -70,7 +61,6 @@ def _save_sync_state() -> None: def _is_idle() -> bool: - """True when no search has happened in the last 30 seconds.""" return _time.time() - _last_search_time > 30 @@ -80,11 +70,7 @@ def _track_ingest(delta: int) -> None: _ingest_in_flight = max(0, _ingest_in_flight + delta) -# โ”€โ”€ Connector sync background thread โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - - def _connector_specs() -> dict[str, tuple[float, str, str]]: - """Connector config mapping: name -> (interval_s, module, class).""" from . import config return { @@ -103,7 +89,6 @@ def _run_connector_sync_once( force: bool, only_sources: set[str] | None = None, ) -> dict[str, dict]: - """Run one connector sync pass and return per-connector status.""" import importlib specs = _connector_specs() @@ -114,8 +99,6 @@ def _run_connector_sync_once( global _last_sync_result status: dict[str, dict] = {} - # Non-blocking try. If lock is busy, return in_progress immediately so - # callers can switch to polling mode rather than blocking indefinitely. acquired = _sync_lock.acquire(blocking=False) if not acquired: return {name: {"status": "skipped", "reason": "sync_in_progress"} for name in specs} @@ -125,11 +108,9 @@ def _run_connector_sync_once( for name, (interval, module_path, class_name) in specs.items(): if only_sources and name not in only_sources: continue - if not force and now - _last_connector_sync[name] < interval: status[name] = {"status": "skipped", "reason": "interval_not_reached"} continue - try: module = importlib.import_module(module_path) conn_class = getattr(module, class_name) @@ -137,7 +118,6 @@ def _run_connector_sync_once( if not conn.is_authenticated(): status[name] = {"status": "skipped", "reason": "not_authenticated"} continue - since: datetime | None = None if _last_connector_sync[name] > 0: since = datetime.fromtimestamp(_last_connector_sync[name], tz=timezone.utc) @@ -150,17 +130,15 @@ def _should_pause() -> bool: should_pause=_should_pause, budget_s=_config.CONNECTOR_SYNC_BUDGET_S, ) - - embedded = sum(1 for r in results if r.get("status") == "embedded") - errors = sum(1 for r in results if r.get("status") == "error") - skipped = sum(1 for r in results if r.get("status") == "skipped") - had_partial = errors > 0 or (embedded == 0 and skipped > 0) or (embedded > 0 and (errors > 0 or skipped > 0)) + embedded = sum(1 for row in results if row.get("status") == "embedded") + errors = sum(1 for row in results if row.get("status") == "error") + skipped = sum(1 for row in results if row.get("status") == "skipped") + had_partial = errors > 0 or (embedded == 0 and skipped > 0) or ( + embedded > 0 and (errors > 0 or skipped > 0) + ) if embedded > 0: _last_connector_sync[name] = _time.time() _save_sync_state() - if embedded > 0: - _known_sources.add(name) - status[name] = { "status": "partial" if had_partial else "ok", "embedded": embedded, @@ -168,7 +146,6 @@ def _should_pause() -> bool: } if errors > 0: status[name]["error_count"] = errors - logger.info("Connector sync %s: %d new items (of %d)", name, embedded, len(results)) except Exception as exc: status[name] = {"status": "error", "error": str(exc)} logger.warning("Connector sync failed for %s: %s", name, exc) @@ -180,106 +157,104 @@ def _should_pause() -> bool: def _connector_sync_loop() -> None: - """Run all authenticated connectors at their configured intervals.""" - - # Give the daemon a 10s head-start before the first sync attempt. _time.sleep(10) - while True: - _time.sleep(60) # check every minute - + _time.sleep(60) if not _is_idle(): - logger.debug("Connector sync skipped โ€” daemon not idle") continue - _run_connector_sync_once(force=False) -# โ”€โ”€ FastAPI app โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - - def _build_app(): from fastapi import Body, FastAPI, HTTPException from pydantic import BaseModel, ValidationError - # โ”€โ”€ Lifespan (replaces deprecated on_event) โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + def _load_watched_dirs() -> list[str]: + if not _config.WATCHED_DIRS_FILE.exists(): + return [] + try: + payload = json.loads(_config.WATCHED_DIRS_FILE.read_text()) + if isinstance(payload, list): + return [str(row) for row in payload if row] + except Exception: + pass + return [] + + def _persist_watched_dirs(dirs: list[str]) -> None: + _config.ensure_runtime_dirs() + _config.WATCHED_DIRS_FILE.write_text(json.dumps(dirs, indent=2)) + + def _on_new_file(path: Path) -> None: + from .ingest import ingest_file + + try: + if path.exists(): + _track_ingest(1) + ingest_file(path, source="files") + except Exception as exc: + logger.debug("Auto-index failed for %s: %s", path, exc) + finally: + _track_ingest(-1) + + def _on_delete_file(path: Path) -> None: + from . import store + + try: + store.delete_by_path(path) + except Exception as exc: + logger.debug("Delete handling failed for %s: %s", path, exc) + + def _restart_watcher(dirs: list[str]) -> None: + from .watcher import FileWatcher + + resolved_dirs = [Path(row).expanduser().resolve() for row in dirs] + existing_dirs = [row for row in resolved_dirs if row.exists()] + with _watcher_lock: + global _watcher_queue_depth, _watcher + if _watcher is not None: + _watcher.stop() + _watcher = None + _watcher_queue_depth = None + if not existing_dirs: + return + watcher = FileWatcher() + watcher.start(existing_dirs, _on_new_file, delete_callback=_on_delete_file) + _watcher = watcher + _watcher_queue_depth = watcher.queued + logger.info("File watcher started on %d directories", len(existing_dirs)) @asynccontextmanager async def lifespan(app: FastAPI): - # โ”€โ”€ Startup โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - from . import store, embedder, config + from . import embedder, keychain, migration, store - config.ensure_vef_dirs() - if config.EMBEDDING_PROVIDER == "gemini": - try: - config.get_api_key() - except ValueError as exc: - logger.error("Startup failed: %s", exc) - raise RuntimeError(str(exc)) from exc - store._get_collection() + _config.ensure_runtime_dirs() + migration.ensure_migrated() + keychain.migrate_legacy_credentials() + store.initialize() embedder.warmup_provider() _last_connector_sync.update(_load_sync_state()) - # Pre-populate sources cache without a full scan. - try: - coll = store._get_collection() - total = coll.count() - sample = coll.get(limit=min(total, 5000), include=["metadatas"]) - for m in sample.get("metadatas") or []: - if m and m.get("source"): - _known_sources.add(m["source"]) - logger.debug("Sources cache pre-populated: %s", sorted(_known_sources)) - except Exception as exc: - logger.warning("Could not pre-populate sources cache: %s", exc) - - # Start filesystem watcher if directories are configured. - watcher = None try: - if config.WATCHED_DIRS_FILE.exists(): - dirs_data = json.loads(config.WATCHED_DIRS_FILE.read_text()) - dirs = [Path(d) for d in dirs_data if d] - if dirs: - from .watcher import FileWatcher - from .ingest import ingest_file - - def _on_new_file(p: Path) -> None: - try: - if p.exists(): - _track_ingest(1) - result = ingest_file(p, source="files") - if result.get("status") == "embedded": - _known_sources.add("files") - logger.info("Auto-indexed: %s", p) - except Exception as exc: - logger.debug("Auto-index failed for %s: %s", p, exc) - finally: - _track_ingest(-1) - - watcher = FileWatcher() - watcher.start(dirs, _on_new_file) - logger.info("File watcher started on %d directories", len(dirs)) + _restart_watcher(_load_watched_dirs()) except Exception as exc: logger.warning("Could not start file watcher: %s", exc) - # Start connector sync background thread. sync_thread = threading.Thread( target=_connector_sync_loop, daemon=True, - name="vef-connector-sync", + name="recall-connector-sync", ) sync_thread.start() + logger.info("Recall daemon ready at %s", _config.SOCKET_PATH) + yield + with _watcher_lock: + global _watcher_queue_depth, _watcher + if _watcher is not None: + _watcher.stop() + _watcher = None + _watcher_queue_depth = None - logger.info("VEF daemon ready โ€” port %d", config.DAEMON_PORT) - - yield # โ”€โ”€ server is now running โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - - # โ”€โ”€ Shutdown โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - if watcher: - watcher.stop() - - app = FastAPI(title="VEF Daemon", version="1.1.0", lifespan=lifespan) - - # โ”€โ”€ Models โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + app = FastAPI(title="Recall Daemon", version="2.0.0", lifespan=lifespan) class SearchRequest(BaseModel): query: str @@ -306,67 +281,71 @@ class IngestRequest(BaseModel): class SyncRequest(BaseModel): source: Optional[str] = None - # โ”€โ”€ Routes โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - - @app.post("/search", response_model=List[SearchResult]) - async def search(req: dict = Body(...)): - global _last_search_time - _last_search_time = _time.time() - - try: - parsed = SearchRequest.model_validate(req) - except ValidationError as exc: - raise HTTPException(status_code=422, detail=exc.errors()) from exc - - if not parsed.query.strip(): - return [] - - from .search import search as vef_search - try: - results = vef_search( - parsed.query, - n_results=parsed.n_results, - sources=parsed.sources, - ) - return results - except Exception as exc: - raise HTTPException(status_code=500, detail=str(exc)) from exc - @app.get("/health") async def health(): - # Liveness check must be constant-time. Do NOT touch the DB here โ€” - # chromadb's count() serialises and can take hundreds of ms under load, - # which would make the daemon look dead whenever Raycast polls. return {"status": "ok"} + @app.get("/ready") + async def ready(): + from . import migration, store + + return { + "status": "ok", + "migration": migration.status().get("status", "not_started"), + "index": store.index_status(), + } + @app.get("/stats") async def stats(): - # Potentially slow โ€” calls chromadb count(). Separate from /health - # so liveness probes don't pay this cost. from . import store + return {"status": "ok", "count": store.count()} @app.get("/sources") async def sources(): - return {"sources": sorted(_known_sources)} + from . import store + + return {"sources": store.get_sources()} @app.get("/progress") async def progress(): from . import store + with _ingest_lock: in_flight = _ingest_in_flight - return {"indexing": in_flight > 0, "queued": in_flight, "total_indexed": store.count()} + queued = _watcher_queue_depth() if _watcher_queue_depth else 0 + return { + "indexing": (in_flight > 0 or queued > 0), + "queued": queued + in_flight, + "total_indexed": store.count(), + } + + @app.post("/search", response_model=List[SearchResult]) + async def search(req: dict = Body(...)): + global _last_search_time + _last_search_time = _time.time() + try: + parsed = SearchRequest.model_validate(req) + except ValidationError as exc: + raise HTTPException(status_code=422, detail=exc.errors()) from exc + if not parsed.query.strip(): + return [] + from .search import search as recall_search + + try: + return recall_search(parsed.query, n_results=parsed.n_results, sources=parsed.sources) + except Exception as exc: + raise HTTPException(status_code=500, detail=str(exc)) from exc @app.post("/ingest") async def ingest(req: dict = Body(...)): - from .ingest import ingest_file import asyncio + from .ingest import ingest_file try: parsed = IngestRequest.model_validate(req) except ValidationError as exc: raise HTTPException(status_code=422, detail=exc.errors()) from exc - loop = asyncio.get_running_loop() _track_ingest(1) try: @@ -376,16 +355,14 @@ async def ingest(req: dict = Body(...)): ) finally: _track_ingest(-1) - if result.get("status") == "embedded": - _known_sources.add(parsed.source) return result @app.get("/connector-status") async def connector_status(): import importlib - specs = _connector_specs() + result: dict[str, dict] = {} - for name, (interval, module_path, class_name) in specs.items(): + for name, (interval, module_path, class_name) in _connector_specs().items(): last = _last_connector_sync.get(name, 0.0) try: module = importlib.import_module(module_path) @@ -415,22 +392,14 @@ async def sync(req: dict | None = Body(default=None)): if parsed.source: parsed_source = parsed.source.strip().lower() if parsed_source not in _connector_specs(): - raise HTTPException( - status_code=400, - detail=f"Unknown source '{parsed_source}'. Valid sources: {', '.join(sorted(_connector_specs()))}", - ) + raise HTTPException(status_code=400, detail=f"Unknown source '{parsed_source}'") - # Always return instantly โ€” start sync in background if not already running. if _sync_lock.locked(): return {"status": "in_progress", "last_sync": _last_sync_result} - # Fire and forget: run sync in daemon thread, don't await it. only = {parsed_source} if parsed_source else None loop = asyncio.get_running_loop() - loop.run_in_executor( - None, - lambda: _run_connector_sync_once(force=True, only_sources=only), - ) + loop.run_in_executor(None, lambda: _run_connector_sync_once(force=True, only_sources=only)) return {"status": "started", "last_sync": _last_sync_result} @app.get("/sync-running") @@ -439,146 +408,221 @@ async def sync_running(): @app.get("/watched-dirs") async def get_watched_dirs(): - from . import config - dirs: list[str] = [] - if config.WATCHED_DIRS_FILE.exists(): - try: - dirs = json.loads(config.WATCHED_DIRS_FILE.read_text()) - except Exception: - pass - return {"dirs": dirs} + return {"dirs": _load_watched_dirs()} @app.post("/watched-dirs") async def add_watched_dir(req: dict = Body(default={})): - from . import config path = str(req.get("path", "")).strip() if not path: raise HTTPException(status_code=400, detail="path required") - from pathlib import Path as _Path - resolved = str(_Path(path).expanduser().resolve()) - dirs: list[str] = [] - if config.WATCHED_DIRS_FILE.exists(): - try: - dirs = json.loads(config.WATCHED_DIRS_FILE.read_text()) - except Exception: - pass + resolved = str(Path(path).expanduser().resolve()) + dirs = _load_watched_dirs() if resolved not in dirs: dirs.append(resolved) - config.WATCHED_DIRS_FILE.write_text(json.dumps(dirs, indent=2)) + _persist_watched_dirs(dirs) + try: + _restart_watcher(dirs) + except Exception as exc: + raise HTTPException(status_code=500, detail=f"Watcher reload failed: {exc}") from exc return {"dirs": dirs} @app.delete("/watched-dirs") async def remove_watched_dir(req: dict = Body(default={})): - from . import config path = str(req.get("path", "")).strip() - dirs: list[str] = [] - if config.WATCHED_DIRS_FILE.exists(): - try: - dirs = json.loads(config.WATCHED_DIRS_FILE.read_text()) - except Exception: - pass - dirs = [d for d in dirs if d != path] - config.WATCHED_DIRS_FILE.write_text(json.dumps(dirs, indent=2)) + resolved = str(Path(path).expanduser().resolve()) if path else "" + dirs = [row for row in _load_watched_dirs() if row != path and row != resolved] + _persist_watched_dirs(dirs) + try: + _restart_watcher(dirs) + except Exception as exc: + raise HTTPException(status_code=500, detail=f"Watcher reload failed: {exc}") from exc return {"dirs": dirs} @app.post("/configure") async def configure(req: dict = Body(default={})): - """Write API keys to ~/.vef/.env so they persist across daemon restarts.""" - from pathlib import Path as _Path - env_file = _Path.home() / ".vef" / ".env" + env_file = _config.RECALL_HOME / ".env" env_file.parent.mkdir(parents=True, exist_ok=True) - lines: list[str] = [] - if env_file.exists(): - lines = env_file.read_text().splitlines() + lines: list[str] = env_file.read_text().splitlines() if env_file.exists() else [] def _set(key: str, val: str) -> None: nonlocal lines - lines = [l for l in lines if not l.startswith(f"{key}=")] + lines = [line for line in lines if not line.startswith(f"{key}=")] lines.append(f"{key}={val}") - if req.get("gemini_api_key"): - _set("GEMINI_API_KEY", str(req["gemini_api_key"])) - os.environ["GEMINI_API_KEY"] = str(req["gemini_api_key"]) - if req.get("canvas_api_key"): - _set("CANVAS_API_KEY", str(req["canvas_api_key"])) - os.environ["CANVAS_API_KEY"] = str(req["canvas_api_key"]) - if req.get("canvas_base_url"): - _set("CANVAS_BASE_URL", str(req["canvas_base_url"])) - os.environ["CANVAS_BASE_URL"] = str(req["canvas_base_url"]) - if req.get("schoology_consumer_key"): - _set("SCHOOLOGY_CONSUMER_KEY", str(req["schoology_consumer_key"])) - os.environ["SCHOOLOGY_CONSUMER_KEY"] = str(req["schoology_consumer_key"]) - if req.get("schoology_consumer_secret"): - _set("SCHOOLOGY_CONSUMER_SECRET", str(req["schoology_consumer_secret"])) - os.environ["SCHOOLOGY_CONSUMER_SECRET"] = str(req["schoology_consumer_secret"]) + for src_key, env_key in ( + ("gemini_api_key", "GEMINI_API_KEY"), + ("canvas_api_key", "CANVAS_API_KEY"), + ("canvas_base_url", "CANVAS_BASE_URL"), + ("schoology_consumer_key", "SCHOOLOGY_CONSUMER_KEY"), + ("schoology_consumer_secret", "SCHOOLOGY_CONSUMER_SECRET"), + ): + value = req.get(src_key) + if value: + _set(env_key, str(value)) + os.environ[env_key] = str(value) env_file.write_text("\n".join(lines) + "\n") return {"ok": True} - return app + @app.get("/model-status") + async def model_status(): + from . import model_manager + + return model_manager.model_status() + @app.get("/index-status") + async def index_status(): + from . import store -# โ”€โ”€ Server start โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + return store.index_status() + @app.get("/migration-status") + async def migration_status(): + from . import migration -def _configure_logging() -> None: - """Install a rotating file handler for ~/.vef/daemon.log. + return migration.status() - Prevents unbounded growth from repetitive warnings (e.g. chromadb compactor - errors spamming _safe_count). Keeps the last 2 MB * 3 rotations. - """ + @app.post("/rebuild-index") + async def rebuild_index(): + from . import store + + return store.rebuild_hot_index() + + return app + + +def _configure_logging() -> None: from logging.handlers import RotatingFileHandler - from . import config root = logging.getLogger() - # If already configured (hot reload), skip. if any(isinstance(h, RotatingFileHandler) for h in root.handlers): return - config.ensure_vef_dirs() - log_path = config.VEF_DIR / "daemon.log" + _config.ensure_runtime_dirs() + log_path = _config.LOG_DIR / "daemon.log" handler = RotatingFileHandler( log_path, maxBytes=2 * 1024 * 1024, backupCount=3, encoding="utf-8", ) - handler.setFormatter(logging.Formatter( - "%(asctime)s %(levelname)s %(name)s: %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - )) + handler.setFormatter( + logging.Formatter( + "%(asctime)s %(levelname)s %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + ) root.addHandler(handler) root.setLevel(logging.INFO) +def _uds_client(timeout: float = 2.0) -> httpx.Client: + transport = httpx.HTTPTransport(uds=str(_config.SOCKET_PATH)) + return httpx.Client(transport=transport, base_url=_config.RECALL_SOCKET_BASE_URL, timeout=timeout) + + +def _poll_health_socket(timeout_s: float) -> bool: + deadline = _time.time() + timeout_s + while _time.time() < deadline: + try: + with _uds_client(timeout=2.0) as client: + resp = client.get("/health") + if resp.is_success: + return True + except Exception: + pass + _time.sleep(0.3) + return False + + +def _poll_health_http(host: str, port: int, timeout_s: float) -> bool: + deadline = _time.time() + timeout_s + while _time.time() < deadline: + try: + resp = httpx.get(f"http://{host}:{port}/health", timeout=2.0) + if resp.is_success: + return True + except Exception: + pass + _time.sleep(0.3) + return False + + +def _build_compat_proxy_app(): + from fastapi import FastAPI, Request, Response + + app = FastAPI(title="Recall Compat Proxy", version="1.0.0") + + async def _forward(path: str, request: Request) -> Response: + upstream = f"/{path}" if path else "/" + body = await request.body() + headers = { + key: value + for key, value in request.headers.items() + if key.lower() not in {"host", "content-length", "connection"} + } + transport = httpx.AsyncHTTPTransport(uds=str(_config.SOCKET_PATH)) + async with httpx.AsyncClient( + transport=transport, + base_url=_config.RECALL_SOCKET_BASE_URL, + timeout=60.0, + ) as client: + upstream_resp = await client.request( + request.method, + upstream, + params=request.query_params, + content=body, + headers=headers, + ) + response_headers = { + key: value + for key, value in upstream_resp.headers.items() + if key.lower() not in {"content-length", "transfer-encoding", "connection"} + } + return Response( + content=upstream_resp.content, + status_code=upstream_resp.status_code, + headers=response_headers, + media_type=upstream_resp.headers.get("content-type"), + ) + + @app.api_route( + "/", + methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"], + ) + async def proxy_root(request: Request) -> Response: + return await _forward("", request) + + @app.api_route( + "/{path:path}", + methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"], + ) + async def proxy_path(path: str, request: Request) -> Response: + return await _forward(path, request) + + return app + + def _run_server() -> None: import uvicorn - from . import config - config.ensure_vef_dirs() + _config.ensure_runtime_dirs() _configure_logging() - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - try: - sock.bind((config.DAEMON_HOST, config.DAEMON_PORT)) - except OSError as exc: - if exc.errno == errno.EADDRINUSE: - if _poll_health(config.DAEMON_HOST, config.DAEMON_PORT, timeout_s=2.0): - print(f"another daemon already serving on {config.DAEMON_PORT}") - sys.exit(0) - print(f"port {config.DAEMON_PORT} already in use; try `vef-daemon stop`", file=sys.stderr) - sys.exit(1) - raise - finally: + if _config.SOCKET_PATH.exists(): try: - sock.close() + _config.SOCKET_PATH.unlink() except Exception: pass - config.PID_FILE.write_text(str(os.getpid())) + _config.PID_FILE.write_text(str(os.getpid())) def _cleanup(signum=None, frame=None): try: - config.PID_FILE.unlink(missing_ok=True) + _config.PID_FILE.unlink(missing_ok=True) + except Exception: + pass + try: + _config.SOCKET_PATH.unlink(missing_ok=True) except Exception: pass sys.exit(0) @@ -587,31 +631,34 @@ def _cleanup(signum=None, frame=None): signal.signal(signal.SIGINT, _cleanup) app = _build_app() + compat_thread = None + if _config.RECALL_ENABLE_COMPAT_HTTP: + compat_app = _build_compat_proxy_app() + compat_thread = threading.Thread( + target=uvicorn.run, + kwargs={ + "app": compat_app, + "host": _config.DAEMON_HOST, + "port": _config.DAEMON_PORT, + "log_level": "warning", + "access_log": False, + }, + daemon=True, + name="recall-compat-http", + ) + compat_thread.start() try: - try: - uvicorn.run( - app, - host=config.DAEMON_HOST, - port=config.DAEMON_PORT, - log_level="warning", - access_log=False, - ) - except OSError as exc: - if exc.errno == errno.EADDRINUSE: - if _poll_health(config.DAEMON_HOST, config.DAEMON_PORT, timeout_s=2.0): - print(f"another daemon already serving on {config.DAEMON_PORT}") - sys.exit(0) - print(f"port {config.DAEMON_PORT} already in use; try `vef-daemon stop`", file=sys.stderr) - sys.exit(1) - raise + uvicorn.run( + app, + uds=str(_config.SOCKET_PATH), + log_level="warning", + access_log=False, + ) finally: _cleanup() -# โ”€โ”€ CLI commands โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ - - def _pid_running(pid: int) -> bool: try: os.kill(pid, 0) @@ -621,96 +668,45 @@ def _pid_running(pid: int) -> bool: def _read_pid() -> int | None: - from . import config try: - return int(config.PID_FILE.read_text().strip()) + return int(_config.PID_FILE.read_text().strip()) except Exception: return None -def _poll_health(host: str, port: int, timeout_s: float) -> bool: - """Poll /health until it responds OK or timeout expires. - - /health is a constant-time liveness check (see app.health). We give httpx - a generous timeout so a momentary CPU spike does not fail the probe. - """ - import socket - deadline = _time.time() + timeout_s - while _time.time() < deadline: - try: - with socket.create_connection((host, port), timeout=0.3): - pass - import httpx - resp = httpx.get(f"http://{host}:{port}/health", timeout=2.0) - if resp.is_success: - return True - except Exception: - pass - _time.sleep(0.3) - return False - - -def _port_in_use(host: str, port: int) -> bool: - """True if something is already bound to host:port (regardless of pid file).""" - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(0.3) - try: - sock.bind((host, port)) - return False - except OSError as exc: - return exc.errno == errno.EADDRINUSE - finally: - try: - sock.close() - except Exception: - pass - - def cmd_start() -> None: - from . import config - config.ensure_vef_dirs() - - # Case 1: PID file points at a live process โ€” verify it's actually serving. + _config.ensure_runtime_dirs() pid = _read_pid() if pid and _pid_running(pid): - if _poll_health(config.DAEMON_HOST, config.DAEMON_PORT, timeout_s=3.0): + if _poll_health_socket(timeout_s=3.0): print(f"Daemon already running (pid {pid})") return - # Live pid but not healthy โ€” kill it so we can start fresh. try: os.kill(pid, signal.SIGTERM) - for _ in range(10): - _time.sleep(0.2) - if not _pid_running(pid): - break - except Exception: + except PermissionError: + print(f"Daemon process {pid} is unhealthy and cannot be terminated (permission denied).") + sys.exit(1) + except ProcessLookupError: pass - config.PID_FILE.unlink(missing_ok=True) - - # Case 2: PID file stale โ€” clear it. - if pid and not _pid_running(pid): - config.PID_FILE.unlink(missing_ok=True) - - # Case 3: Port already in use by something that's not in our PID file. - # Check /health โ€” if it's our daemon, treat as already running. - if _port_in_use(config.DAEMON_HOST, config.DAEMON_PORT): - if _poll_health(config.DAEMON_HOST, config.DAEMON_PORT, timeout_s=3.0): - print(f"Daemon already running on port {config.DAEMON_PORT} (pid file missing)") - return - print( - f"Port {config.DAEMON_PORT} is in use but /health is not responding.", - file=sys.stderr, - ) - print("Run `vef-daemon stop` or `lsof -i :19847` to investigate.", file=sys.stderr) - sys.exit(1) + for _ in range(20): + _time.sleep(0.2) + if not _pid_running(pid): + break + if _pid_running(pid): + print(f"Daemon process {pid} is unhealthy and could not be stopped.") + print("Stop it manually and retry.") + sys.exit(1) + _config.PID_FILE.unlink(missing_ok=True) + elif pid: + _config.PID_FILE.unlink(missing_ok=True) - # Case 4: Fresh start. Spawn the daemon. - log_path = config.VEF_DIR / "daemon.log" + log_path = _config.LOG_DIR / "daemon.log" project_root = str(Path(__file__).parent.parent) env = os.environ.copy() existing_py_path = env.get("PYTHONPATH", "") env["PYTHONPATH"] = f"{project_root}:{existing_py_path}" if existing_py_path else project_root import subprocess + with open(log_path, "a") as log_fh: proc = subprocess.Popen( [sys.executable, "-m", "vector_embedded_finder.daemon", "_serve"], @@ -721,19 +717,10 @@ def cmd_start() -> None: start_new_session=True, ) - if _poll_health(config.DAEMON_HOST, config.DAEMON_PORT, timeout_s=15.0): - # Verify the child is still alive (it may have exited after health came up - # if another daemon races us โ€” though we guarded against that above). - if proc.poll() is not None: - pid_from_file = _read_pid() - if pid_from_file: - print(f"Daemon running (pid {pid_from_file}); spawned helper pid {proc.pid} exited.") - else: - print("Daemon responding but spawned pid exited โ€” investigate daemon.log.") - return + if _poll_health_socket(timeout_s=20.0): print(f"Daemon started (pid {proc.pid})") else: - print(f"Daemon process spawned (pid {proc.pid}) but did not respond within 15s.") + print(f"Daemon process spawned (pid {proc.pid}) but did not respond within 20s.") print(f"Check logs: {log_path}") @@ -743,12 +730,10 @@ def cmd_stop() -> None: print("Daemon not running (no PID file)") return if not _pid_running(pid): - from . import config - config.PID_FILE.unlink(missing_ok=True) + _config.PID_FILE.unlink(missing_ok=True) print("Daemon not running (stale PID cleaned up)") return os.kill(pid, signal.SIGTERM) - # Wait briefly so subsequent commands don't see a stale state for _ in range(10): _time.sleep(0.3) if not _pid_running(pid): @@ -757,69 +742,43 @@ def cmd_stop() -> None: def cmd_status() -> None: - from . import config pid = _read_pid() if not pid or not _pid_running(pid): - log_path = config.VEF_DIR / "daemon.log" + log_path = _config.LOG_DIR / "daemon.log" print("Daemon: stopped") if log_path.exists(): - # Show last 5 lines of log to help diagnose crashes lines = log_path.read_text().splitlines() tail = lines[-5:] if len(lines) >= 5 else lines if any(tail): print(f"Last log lines ({log_path}):") - for ln in tail: - print(f" {ln}") + for line in tail: + print(f" {line}") return - try: - import httpx - # /stats may be slow if chromadb is backlogged โ€” keep it short. - resp = httpx.get( - f"http://{config.DAEMON_HOST}:{config.DAEMON_PORT}/stats", - timeout=5.0, - ) - data = resp.json() - print(f"Daemon: running (pid {pid}, {data.get('count', '?')} documents indexed)") + with _uds_client(timeout=5.0) as client: + stats = client.get("/stats").json() + ready = client.get("/ready").json() + print(f"Daemon: running (pid {pid}, {stats.get('count', '?')} documents indexed)") + print(f"Migration: {ready.get('migration', 'unknown')}") except Exception: - # Stats slow โ†’ fall back to /health for liveness. - try: - import httpx - resp = httpx.get( - f"http://{config.DAEMON_HOST}:{config.DAEMON_PORT}/health", - timeout=2.0, - ) - if resp.is_success: - print(f"Daemon: running (pid {pid}, index stats unavailable)") - return - except Exception: - pass - print(f"Daemon: running (pid {pid}, health check failed โ€” may still be starting up)") + print(f"Daemon: running (pid {pid}, health check failed)") def cmd_sync(source: str | None = None) -> None: - from . import config - import httpx - payload = {"source": source} if source else {} try: - resp = httpx.post( - f"http://{config.DAEMON_HOST}:{config.DAEMON_PORT}/sync", - json=payload, - timeout=120.0, - ) - resp.raise_for_status() + with _uds_client(timeout=120.0) as client: + resp = client.post("/sync", json=payload) + resp.raise_for_status() + data = resp.json() except Exception as exc: print(f"Sync failed: {exc}") print("Start daemon first: vef-daemon start") sys.exit(1) - - data = resp.json() sync_data = data.get("last_sync", {}) if not sync_data: print("No connectors synced.") return - for name in sorted(sync_data): result = sync_data[name] if result.get("status") == "ok": @@ -838,11 +797,9 @@ def cmd_check_embed() -> None: except Exception as exc: print(f"Embedding provider check failed: {exc}") sys.exit(1) - if not vec: print("Embedding provider check failed: empty embedding vector") sys.exit(1) - print(f"Embedding provider OK (dimension {len(vec)})") diff --git a/vector_embedded_finder/embedder.py b/vector_embedded_finder/embedder.py index 75ee589..44ea675 100644 --- a/vector_embedded_finder/embedder.py +++ b/vector_embedded_finder/embedder.py @@ -1,14 +1,10 @@ -"""Embedding provider wrapper with retry/backoff. - -Supported providers: -- gemini (default) -- ollama (local) -- nim (OpenAI-compatible embeddings endpoint) -""" +"""Embedding provider wrapper with local-first defaults.""" from __future__ import annotations +import hashlib import logging +import math import random import time from pathlib import Path @@ -17,14 +13,14 @@ from google import genai from google.genai import types -from . import config +from . import config, model_manager logger = logging.getLogger(__name__) _client: genai.Client | None = None _MAX_RETRIES = 5 -_BASE_BACKOFF = 2.0 # seconds +_BASE_BACKOFF = 2.0 def _get_client() -> genai.Client: @@ -40,7 +36,6 @@ def _is_rate_limit(exc: Exception) -> bool: def _call_with_retry(fn, provider: str): - """Call fn(), retrying on rate-limit errors with exponential back-off.""" wait = _BASE_BACKOFF for attempt in range(_MAX_RETRIES): try: @@ -50,9 +45,11 @@ def _call_with_retry(fn, provider: str): jitter = random.uniform(0, wait * 0.2) actual = wait + jitter logger.warning( - "%s rate-limited โ€” retrying in %.1fs (attempt %d/%d)", + "%s rate-limited, retrying in %.1fs (attempt %d/%d)", provider, - actual, attempt + 1, _MAX_RETRIES, + actual, + attempt + 1, + _MAX_RETRIES, ) time.sleep(actual) wait = min(wait * 2, 32.0) @@ -61,6 +58,51 @@ def _call_with_retry(fn, provider: str): raise RuntimeError("unreachable") +def _normalize(vec: list[float]) -> list[float]: + norm = math.sqrt(sum(v * v for v in vec)) + if norm <= 0: + return vec + return [float(v / norm) for v in vec] + + +def _hash_embedding(text: str) -> list[float]: + dims = config.EMBEDDING_DIMENSIONS + vec = [0.0] * dims + for token in text.lower().split(): + digest = hashlib.sha256(token.encode("utf-8")).digest() + for i in range(0, min(16, dims)): + idx = (digest[i] + (i * 31)) % dims + sign = -1.0 if digest[(i + 1) % len(digest)] % 2 else 1.0 + vec[idx] += sign * ((digest[(i + 2) % len(digest)] / 255.0) + 0.5) + if not any(vec): + vec[0] = 1.0 + return _normalize(vec) + + +def _embed_text_local(text: str) -> list[float]: + model = model_manager.get_text_model() + if model is None: + return _hash_embedding(text) + values = model.encode([text], normalize_embeddings=True)[0] + return [float(x) for x in values[: config.EMBEDDING_DIMENSIONS]] + + +def _embed_path_local(path: Path, fallback_text: str) -> list[float]: + model = model_manager.get_vision_model() + if model is None: + return _embed_text_local(fallback_text) + try: + from PIL import Image + except Exception: + return _embed_text_local(fallback_text) + try: + image = Image.open(path).convert("RGB") + values = model.encode([image], normalize_embeddings=True)[0] + return [float(x) for x in values[: config.EMBEDDING_DIMENSIONS]] + except Exception: + return _embed_text_local(fallback_text) + + def _embed_text_gemini(text: str, task: str) -> list[float]: client = _get_client() @@ -73,9 +115,9 @@ def _call(): output_dimensionality=config.EMBEDDING_DIMENSIONS, ), ) - return result.embeddings[0].values + return [float(v) for v in result.embeddings[0].values] - return _call_with_retry(_call, provider="gemini") + return _normalize(_call_with_retry(_call, provider="gemini")) def _embed_text_ollama(text: str) -> list[float]: @@ -92,7 +134,7 @@ def _call(): raise ValueError("Ollama embedding response missing 'embedding'") return [float(x) for x in emb] - return _call_with_retry(_call, provider="ollama") + return _normalize(_call_with_retry(_call, provider="ollama")) def _embed_text_nim(text: str) -> list[float]: @@ -113,35 +155,42 @@ def _call(): raise ValueError("NIM embedding response missing data[0].embedding") return [float(x) for x in rows[0]["embedding"]] - return _call_with_retry(_call, provider="nim") + return _normalize(_call_with_retry(_call, provider="nim")) def warmup_provider() -> None: provider = config.EMBEDDING_PROVIDER + if provider == "local": + model_manager.warmup() + return if provider == "gemini": _get_client() + model_manager.warmup() return if provider == "ollama": - # Quick local health check. httpx.get(f"{config.OLLAMA_BASE_URL}/api/tags", timeout=5.0).raise_for_status() + model_manager.warmup() return if provider == "nim": if not config.NIM_EMBED_URL: raise ValueError("VEF_NIM_EMBED_URL is not configured") _ = config.get_nim_api_key() + model_manager.warmup() return - raise ValueError(f"Unsupported VEF_EMBEDDING_PROVIDER: {provider}") + raise ValueError(f"Unsupported RECALL_EMBEDDING_PROVIDER: {provider}") def embed_text(text: str, task: str = "RETRIEVAL_DOCUMENT") -> list[float]: provider = config.EMBEDDING_PROVIDER + if provider == "local": + return _embed_text_local(text) if provider == "gemini": return _embed_text_gemini(text, task=task) if provider == "ollama": return _embed_text_ollama(text) if provider == "nim": return _embed_text_nim(text) - raise ValueError(f"Unsupported VEF_EMBEDDING_PROVIDER: {provider}") + raise ValueError(f"Unsupported provider: {provider}") def embed_query(query: str) -> list[float]: @@ -149,12 +198,12 @@ def embed_query(query: str) -> list[float]: def embed_image(path: Path) -> list[float]: + if config.EMBEDDING_PROVIDER == "local": + return _embed_path_local(path, fallback_text=f"Image file {path.name}") if config.EMBEDDING_PROVIDER != "gemini": - return embed_text(f"Image file: {path.name}") + return embed_text(f"Image file {path.name}") client = _get_client() - with open(path, "rb") as f: - image_bytes = f.read() - + image_bytes = path.read_bytes() from . import utils mt = utils.mime_type(path) @@ -169,81 +218,18 @@ def _call(): output_dimensionality=config.EMBEDDING_DIMENSIONS, ), ) - return result.embeddings[0].values + return [float(v) for v in result.embeddings[0].values] - return _call_with_retry(_call, provider="gemini") + return _normalize(_call_with_retry(_call, provider="gemini")) def embed_audio(path: Path) -> list[float]: - if config.EMBEDDING_PROVIDER != "gemini": - return embed_text(f"Audio file: {path.name}") - client = _get_client() - with open(path, "rb") as f: - audio_bytes = f.read() - - from . import utils - mt = utils.mime_type(path) - - def _call(): - result = client.models.embed_content( - model=config.EMBEDDING_MODEL, - contents=types.Content( - parts=[types.Part(inline_data=types.Blob(mime_type=mt, data=audio_bytes))] - ), - config=types.EmbedContentConfig( - task_type="RETRIEVAL_DOCUMENT", - output_dimensionality=config.EMBEDDING_DIMENSIONS, - ), - ) - return result.embeddings[0].values - - return _call_with_retry(_call, provider="gemini") + return embed_text(f"Audio file {path.name}") def embed_video(path: Path) -> list[float]: - if config.EMBEDDING_PROVIDER != "gemini": - return embed_text(f"Video file: {path.name}") - client = _get_client() - with open(path, "rb") as f: - video_bytes = f.read() - - from . import utils - mt = utils.mime_type(path) - - def _call(): - result = client.models.embed_content( - model=config.EMBEDDING_MODEL, - contents=types.Content( - parts=[types.Part(inline_data=types.Blob(mime_type=mt, data=video_bytes))] - ), - config=types.EmbedContentConfig( - task_type="RETRIEVAL_DOCUMENT", - output_dimensionality=config.EMBEDDING_DIMENSIONS, - ), - ) - return result.embeddings[0].values - - return _call_with_retry(_call, provider="gemini") + return embed_text(f"Video file {path.name}") def embed_pdf(path: Path) -> list[float]: - if config.EMBEDDING_PROVIDER != "gemini": - return embed_text(f"PDF file: {path.name}") - client = _get_client() - with open(path, "rb") as f: - pdf_bytes = f.read() - - def _call(): - result = client.models.embed_content( - model=config.EMBEDDING_MODEL, - contents=types.Content( - parts=[types.Part(inline_data=types.Blob(mime_type="application/pdf", data=pdf_bytes))] - ), - config=types.EmbedContentConfig( - task_type="RETRIEVAL_DOCUMENT", - output_dimensionality=config.EMBEDDING_DIMENSIONS, - ), - ) - return result.embeddings[0].values - - return _call_with_retry(_call, provider="gemini") + return embed_text(f"PDF file {path.name}") diff --git a/vector_embedded_finder/ingest.py b/vector_embedded_finder/ingest.py index ad8450a..7a26b2c 100644 --- a/vector_embedded_finder/ingest.py +++ b/vector_embedded_finder/ingest.py @@ -1,4 +1,4 @@ -"""File ingestion pipeline โ€” detect type, caption, embed, store.""" +"""File ingestion pipeline โ€” local-first enrichment, embedding, and storage.""" from __future__ import annotations @@ -7,121 +7,153 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -from typing import Callable +from typing import Any, Callable import psutil -from . import config, embedder, store, utils +from . import captioner, config, embedder, store, utils logger = logging.getLogger(__name__) def _extract_pdf_text(path: Path) -> str: - """Extract plain text from a PDF using pypdf. Returns empty string on failure.""" try: from pypdf import PdfReader + reader = PdfReader(str(path)) pages: list[str] = [] for page in reader.pages: - t = page.extract_text() or "" - stripped = t.strip() - if stripped: - pages.append(stripped) + text = (page.extract_text() or "").strip() + if text: + pages.append(text) return "\n".join(pages) - except Exception as e: - logger.debug("pypdf extraction failed for %s: %s", path, e) + except Exception as exc: + logger.debug("pypdf extraction failed for %s: %s", path, exc) return "" def _cpu_guard() -> None: - """Back off until CPU usage is below the guard threshold.""" while True: usage = psutil.cpu_percent(interval=0.5) if usage <= config.CPU_GUARD_PERCENT: - break + return logger.debug("CPU %.0f%% > %d%%, backing off 5s", usage, config.CPU_GUARD_PERCENT) time.sleep(5) def _set_low_priority() -> None: - """Lower OS scheduling priority for the current thread.""" try: os.nice(10) except (AttributeError, PermissionError): pass +def _image_enrichment(path: Path, category: str) -> dict[str, Any]: + enrichment: dict[str, Any] = { + "caption": "", + "ocr_text": "", + "gps_city": "", + "face_count": 0, + "exif_date": "", + "exif_camera": "", + } + if category != "image": + return enrichment + + try: + from PIL import Image + + image = Image.open(path) + exif = getattr(image, "getexif", lambda: None)() or {} + make = exif.get(271) or "" + model = exif.get(272) or "" + when = exif.get(36867) or exif.get(306) or "" + enrichment["exif_camera"] = " ".join(part for part in (str(make).strip(), str(model).strip()) if part) + enrichment["exif_date"] = str(when) + except Exception as exc: + logger.debug("EXIF extraction failed for %s: %s", path, exc) + + if config.ENABLE_OPTIONAL_CAPTIONING: + try: + caption = captioner.caption_file(path) + if caption: + enrichment["caption"] = caption + except Exception as exc: + logger.debug("Caption enrichment failed for %s: %s", path, exc) + return enrichment + + +def _build_text_payload( + path: Path, + category: str, + description: str, + caption: str | None, +) -> tuple[list[float], str]: + if category == "text": + text = path.read_text(errors="replace") + if len(text) > 32000: + text = text[:32000] + return embedder.embed_text(text), text[:5000] + + if category == "document": + extracted_text = _extract_pdf_text(path) + if extracted_text: + if len(extracted_text) > 32000: + extracted_text = extracted_text[:32000] + return embedder.embed_text(extracted_text), extracted_text[:5000] + fallback = description or f"PDF file {path.name}" + return embedder.embed_pdf(path), fallback + + if category == "image": + if caption: + return embedder.embed_text(caption), caption + return embedder.embed_image(path), description or f"Image file {path.name}" + + if category == "audio": + if caption: + return embedder.embed_text(caption), caption + return embedder.embed_audio(path), description or f"Audio file {path.name}" + + if category == "video": + if caption: + return embedder.embed_text(caption), caption + return embedder.embed_video(path), description or f"Video file {path.name}" + + raise ValueError(f"Unknown category: {category}") + + def ingest_file( path: str | Path, source: str = "manual", description: str = "", -) -> dict: - path = Path(path).resolve() +) -> dict[str, Any]: + path = Path(path).expanduser().resolve() if not path.exists(): raise FileNotFoundError(f"File not found: {path}") - if not utils.is_supported(path): raise ValueError(f"Unsupported file type: {path.suffix}") category = config.get_media_category(path.suffix.lower()) - doc_id = utils.file_hash(path) + if category is None: + raise ValueError(f"Unknown category for file: {path}") + stat = path.stat() + doc_id = utils.file_hash(path) if store.exists(doc_id): return {"status": "skipped", "reason": "already embedded", "id": doc_id, "path": str(path)} caption: str | None = None - - if category in ("image", "audio", "video"): + if category in {"image", "audio", "video"} and config.ENABLE_OPTIONAL_CAPTIONING: _cpu_guard() try: - from . import captioner caption = captioner.caption_file(path) - except Exception as e: - logger.debug("captioner import/call failed for %s: %s", path, e) + except Exception as exc: + logger.debug("Caption generation failed for %s: %s", path, exc) - if category == "text": - text = path.read_text(errors="replace") - if len(text) > 32000: - text = text[:32000] - embedding = embedder.embed_text(text) - doc_text = text[:500] - elif category == "image": - if caption: - embedding = embedder.embed_text(caption) - doc_text = caption[:500] - else: - embedding = embedder.embed_image(path) - doc_text = description or f"Image: {path.name}" - elif category == "audio": - if caption: - embedding = embedder.embed_text(caption) - doc_text = caption[:500] - else: - embedding = embedder.embed_audio(path) - doc_text = description or f"Audio: {path.name}" - elif category == "video": - if caption: - embedding = embedder.embed_text(caption) - doc_text = caption[:500] - else: - embedding = embedder.embed_video(path) - doc_text = description or f"Video: {path.name}" - elif category == "document": - # Prefer text extraction: semantically searchable + far smaller API payload. - extracted_text = _extract_pdf_text(path) - if extracted_text: - if len(extracted_text) > 32000: - extracted_text = extracted_text[:32000] - embedding = embedder.embed_text(extracted_text) - doc_text = extracted_text[:500] - else: - # Fallback: binary PDF embedding (scanned documents, image-only PDFs) - embedding = embedder.embed_pdf(path) - doc_text = description or f"PDF: {path.name}" - else: - raise ValueError(f"Unknown category: {category}") - - effective_description = caption or description + embedding, document = _build_text_payload(path, category, description, caption) + enrichment = _image_enrichment(path, category) + if caption and not enrichment.get("caption"): + enrichment["caption"] = caption metadata = { "file_path": str(path), @@ -130,11 +162,14 @@ def ingest_file( "media_category": category, "timestamp": utils.now_iso(), "source": source, - "description": effective_description, - "file_size": path.stat().st_size, + "description": caption or description, + "file_size": stat.st_size, + "mtime": stat.st_mtime, + "sha256": doc_id, } - store.add(doc_id, embedding, metadata, document=doc_text) + store.add(doc_id, embedding, metadata, document=document, enrichment=enrichment) + store.retire_path_versions(path, keep_doc_id=doc_id) return {"status": "embedded", "id": doc_id, "path": str(path), "category": category} @@ -143,17 +178,14 @@ def ingest_text( description: str = "", source: str = "manual", tags: str = "", -) -> dict: +) -> dict[str, Any]: doc_id = utils.text_hash(text) - if store.exists(doc_id): return {"status": "skipped", "reason": "already embedded", "id": doc_id} - embedding = embedder.embed_text(text) - metadata = { "file_path": "", - "file_name": "", + "file_name": description or "text snippet", "file_type": "text/plain", "media_category": "text", "timestamp": utils.now_iso(), @@ -161,54 +193,47 @@ def ingest_text( "description": description, "tags": tags, "file_size": len(text.encode()), + "sha256": doc_id, + "mtime": 0.0, } - - store.add(doc_id, embedding, metadata, document=text[:500]) + store.add(doc_id, embedder.embed_text(text), metadata, document=text[:5000], enrichment=None) return {"status": "embedded", "id": doc_id, "category": "text"} -def _ingest_worker(file_path: Path, source: str) -> dict: - """Worker function run inside the thread pool.""" +def _ingest_worker(file_path: Path, source: str) -> dict[str, Any]: _set_low_priority() try: return ingest_file(file_path, source=source) - except Exception as e: - return {"status": "error", "path": str(file_path), "error": str(e)} + except Exception as exc: + return {"status": "error", "path": str(file_path), "error": str(exc)} def ingest_directory( path: str | Path, source: str = "manual", recursive: bool = True, - progress_callback: Callable[[int, int, dict], None] | None = None, -) -> list[dict]: - path = Path(path).resolve() + progress_callback: Callable[[int, int, dict[str, Any]], None] | None = None, +) -> list[dict[str, Any]]: + path = Path(path).expanduser().resolve() pattern = "**/*" if recursive else "*" - files = [f for f in sorted(path.glob(pattern)) if f.is_file() and utils.is_supported(f)] total = len(files) - results: list[dict] = [None] * total # type: ignore[list-item] + results: list[dict[str, Any]] = [None] * total # type: ignore[list-item] completed = 0 with ThreadPoolExecutor(max_workers=config.MAX_CONCURRENT_INGEST) as pool: - future_to_idx = { - pool.submit(_ingest_worker, fp, source): i - for i, fp in enumerate(files) - } + future_to_idx = {pool.submit(_ingest_worker, fp, source): i for i, fp in enumerate(files)} for future in as_completed(future_to_idx): idx = future_to_idx[future] try: result = future.result() - except Exception as e: - result = {"status": "error", "path": str(files[idx]), "error": str(e)} - + except Exception as exc: + result = {"status": "error", "path": str(files[idx]), "error": str(exc)} results[idx] = result completed += 1 - if progress_callback: try: progress_callback(completed, total, result) except Exception: pass - return results diff --git a/vector_embedded_finder/keychain.py b/vector_embedded_finder/keychain.py new file mode 100644 index 0000000..4d31cb5 --- /dev/null +++ b/vector_embedded_finder/keychain.py @@ -0,0 +1,157 @@ +"""Keychain-backed storage for connector credentials.""" + +from __future__ import annotations + +import getpass +import json +import logging +import shutil +import subprocess +from pathlib import Path +from typing import Any, Iterable + +from . import config + +logger = logging.getLogger(__name__) + +SERVICE_PREFIX = "com.recall.credentials" +DEFAULT_ACCOUNT = getpass.getuser() + + +def service_name(source: str) -> str: + return f"{SERVICE_PREFIX}.{source}" + + +def _security_available() -> bool: + return config.sys_platform_is_macos() and shutil.which("security") is not None + + +def _run_security(args: list[str]) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["security", *args], + capture_output=True, + text=True, + check=False, + ) + + +def _iter_service_names(source: str, aliases: Iterable[str]) -> list[str]: + names = [source, *aliases] + return [service_name(name) for name in names] + + +def get_secret( + source: str, + *, + aliases: Iterable[str] = (), + account: str = DEFAULT_ACCOUNT, +) -> str | None: + if not _security_available(): + return None + + for name in _iter_service_names(source, aliases): + result = _run_security(["find-generic-password", "-a", account, "-s", name, "-w"]) + if result.returncode == 0: + return result.stdout.rstrip("\n") + return None + + +def set_secret(source: str, secret: str, *, account: str = DEFAULT_ACCOUNT) -> None: + if not _security_available(): + raise RuntimeError("macOS Keychain unavailable") + + result = _run_security(["add-generic-password", "-U", "-a", account, "-s", service_name(source), "-w", secret]) + if result.returncode != 0: + raise RuntimeError(result.stderr.strip() or f"Failed to store credentials for {source}") + + +def delete_secret(source: str, *, account: str = DEFAULT_ACCOUNT) -> None: + if not _security_available(): + return + _run_security(["delete-generic-password", "-a", account, "-s", service_name(source)]) + + +def _load_legacy_json(path: Path | None) -> dict[str, Any] | None: + if path is None or not path.exists(): + return None + payload = json.loads(path.read_text()) + return payload if isinstance(payload, dict) else None + + +def _delete_legacy_file(path: Path | None) -> None: + if path is None or not path.exists(): + return + try: + path.unlink() + except OSError as exc: + logger.warning("Could not delete legacy credential file %s: %s", path, exc) + + +def load_json( + source: str, + *, + legacy_path: Path | None = None, + aliases: Iterable[str] = (), + account: str = DEFAULT_ACCOUNT, +) -> dict[str, Any] | None: + raw = get_secret(source, aliases=aliases, account=account) + if raw: + payload = json.loads(raw) + return payload if isinstance(payload, dict) else None + + payload = _load_legacy_json(legacy_path) + if payload is None: + return None + + if _security_available(): + set_secret(source, json.dumps(payload, sort_keys=True), account=account) + _delete_legacy_file(legacy_path) + + return payload + + +def save_json( + source: str, + payload: dict[str, Any], + *, + legacy_path: Path | None = None, + account: str = DEFAULT_ACCOUNT, +) -> None: + if _security_available(): + set_secret(source, json.dumps(payload, sort_keys=True), account=account) + _delete_legacy_file(legacy_path) + return + + if legacy_path is None: + raise RuntimeError("No fallback credential path available") + config.ensure_runtime_dirs() + legacy_path.write_text(json.dumps(payload, indent=2, sort_keys=True)) + + +def migrate_legacy_credentials() -> dict[str, Any]: + migrated: list[str] = [] + skipped: list[str] = [] + + specs = { + "gmail": config.GMAIL_CREDENTIALS_FILE, + "canvas": config.CANVAS_CREDENTIALS_FILE, + "calai": config.CALAI_CREDENTIALS_FILE, + "schoology": config.SCHOOLOGY_CREDENTIALS_FILE, + "notion": config.NOTION_CREDENTIALS_FILE, + } + + for source, path in specs.items(): + if not path.exists(): + continue + try: + payload = _load_legacy_json(path) + if payload is None: + skipped.append(source) + continue + save_json(source, payload, legacy_path=path) + migrated.append(source) + except Exception as exc: + skipped.append(source) + logger.warning("Credential migration failed for %s: %s", source, exc) + + return {"migrated": migrated, "skipped": skipped} diff --git a/vector_embedded_finder/migration.py b/vector_embedded_finder/migration.py new file mode 100644 index 0000000..46e02a9 --- /dev/null +++ b/vector_embedded_finder/migration.py @@ -0,0 +1,127 @@ +"""Runtime migration from legacy ~/.vef and Chroma-backed state.""" + +from __future__ import annotations + +import json +import logging +import shutil +import time +from pathlib import Path +from typing import Any + +from . import config + +logger = logging.getLogger(__name__) + + +def _read_status() -> dict[str, Any]: + if not config.MIGRATION_STATUS_PATH.exists(): + return {"status": "not_started"} + try: + payload = json.loads(config.MIGRATION_STATUS_PATH.read_text()) + if isinstance(payload, dict): + return payload + except Exception: + pass + return {"status": "not_started"} + + +def _write_status(status: str, **extra: Any) -> dict[str, Any]: + payload = {"status": status, **extra, "updated_at": time.time()} + config.ensure_runtime_dirs() + config.MIGRATION_STATUS_PATH.write_text(json.dumps(payload, indent=2, sort_keys=True)) + return payload + + +def status() -> dict[str, Any]: + return _read_status() + + +def _copy_if_exists(src: Path, dst: Path) -> None: + if not src.exists() or dst.exists(): + return + dst.parent.mkdir(parents=True, exist_ok=True) + if src.is_dir(): + shutil.copytree(src, dst, dirs_exist_ok=True) + else: + shutil.copy2(src, dst) + + +def _migrate_filesystem_state() -> None: + legacy = config.LEGACY_VEF_DIR + if not legacy.exists(): + return + for name in ("credentials", "watched_dirs.json", "sync_state.json", ".env"): + _copy_if_exists(legacy / name, config.RECALL_HOME / name) + + +def _import_chroma() -> dict[str, Any]: + try: + import chromadb + except Exception as exc: + return {"imported": 0, "skipped": 0, "error": f"chromadb unavailable: {exc}"} + + if not config.CHROMA_DIR.exists(): + return {"imported": 0, "skipped": 0} + + from . import store + + client = chromadb.PersistentClient(path=str(config.CHROMA_DIR)) + coll = client.get_or_create_collection( + name=config.COLLECTION_NAME, + metadata={"hnsw:space": "cosine"}, + ) + total = int(coll.count()) + if total <= 0: + return {"imported": 0, "skipped": 0} + + imported = 0 + skipped = 0 + batch_size = max(1, int(config.INDEX_REBUILD_BATCH)) + for offset in range(0, total, batch_size): + rows = coll.get( + include=["embeddings", "metadatas", "documents"], + limit=batch_size, + offset=offset, + ) + ids = rows.get("ids", []) + embeddings = rows.get("embeddings") or [] + metadatas = rows.get("metadatas") or [] + documents = rows.get("documents") or [] + for idx, doc_id in enumerate(ids): + if store.exists(doc_id): + skipped += 1 + continue + embedding = embeddings[idx] if idx < len(embeddings) else None + metadata = metadatas[idx] if idx < len(metadatas) else {} + document = documents[idx] if idx < len(documents) else "" + if not embedding: + skipped += 1 + continue + store.add( + str(doc_id), + [float(v) for v in embedding], + dict(metadata or {}), + document=str(document or ""), + ) + imported += 1 + return {"imported": imported, "skipped": skipped} + + +def ensure_migrated() -> dict[str, Any]: + existing = _read_status() + if existing.get("status") == "complete": + return existing + + _write_status("running") + try: + _migrate_filesystem_state() + from . import store + + store.initialize() + chroma_result = _import_chroma() + result = _write_status("complete", chroma=chroma_result) + return result + except Exception as exc: + logger.exception("Migration failed") + return _write_status("failed", error=str(exc)) diff --git a/vector_embedded_finder/model_manager.py b/vector_embedded_finder/model_manager.py new file mode 100644 index 0000000..dfb719b --- /dev/null +++ b/vector_embedded_finder/model_manager.py @@ -0,0 +1,176 @@ +"""Local model lifecycle for Recall. + +The implementation is intentionally resilient: +- prefers local sentence-transformers models cached under ~/.recall/models +- records model state in a manifest for UI and daemon status +- falls back to deterministic in-process embeddings when heavyweight deps are + not installed, so the rest of the product still functions +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass +from pathlib import Path +from threading import Lock +from typing import Any + +from . import config + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ModelSpec: + key: str + model_name: str + local_dir: Path + kind: str + + +_MANIFEST_LOCK = Lock() +_TEXT_MODEL = None +_VISION_MODEL = None + + +def _model_specs() -> list[ModelSpec]: + return [ + ModelSpec( + key="text", + model_name=config.EMBEDDING_MODEL, + local_dir=config.MODELS_DIR / "text", + kind="embedding", + ), + ModelSpec( + key="vision", + model_name=config.VISION_EMBEDDING_MODEL, + local_dir=config.MODELS_DIR / "vision", + kind="embedding", + ), + ModelSpec( + key="reranker", + model_name=config.RERANKER_MODEL, + local_dir=config.MODELS_DIR / "reranker", + kind="reranker", + ), + ] + + +def _read_manifest() -> dict[str, Any]: + if not config.MODEL_MANIFEST_PATH.exists(): + return {} + try: + payload = json.loads(config.MODEL_MANIFEST_PATH.read_text()) + if isinstance(payload, dict): + return payload + except Exception as exc: + logger.debug("Could not read model manifest: %s", exc) + return {} + + +def _write_manifest(payload: dict[str, Any]) -> None: + config.ensure_runtime_dirs() + config.MODEL_MANIFEST_PATH.write_text(json.dumps(payload, indent=2, sort_keys=True)) + + +def ensure_manifest() -> dict[str, Any]: + with _MANIFEST_LOCK: + payload = _read_manifest() + models = payload.setdefault("models", {}) + for spec in _model_specs(): + row = models.setdefault( + spec.key, + { + "name": spec.model_name, + "kind": spec.kind, + "path": str(spec.local_dir), + "status": "pending", + "backend": "sentence-transformers", + }, + ) + row["name"] = spec.model_name + row["path"] = str(spec.local_dir) + _write_manifest(payload) + return payload + + +def _mark_status(key: str, *, status: str, backend: str, detail: str = "") -> None: + payload = ensure_manifest() + models = payload.setdefault("models", {}) + row = models.setdefault(key, {}) + row["status"] = status + row["backend"] = backend + if detail: + row["detail"] = detail + elif "detail" in row: + del row["detail"] + _write_manifest(payload) + + +def _load_sentence_transformer(model_name: str, cache_dir: Path): + from sentence_transformers import SentenceTransformer + + cache_dir.mkdir(parents=True, exist_ok=True) + return SentenceTransformer(model_name, cache_folder=str(cache_dir)) + + +def get_text_model(): + global _TEXT_MODEL + if _TEXT_MODEL is not None: + return _TEXT_MODEL + spec = next(s for s in _model_specs() if s.key == "text") + try: + _TEXT_MODEL = _load_sentence_transformer(spec.model_name, spec.local_dir) + _mark_status("text", status="ready", backend="sentence-transformers") + return _TEXT_MODEL + except Exception as exc: + logger.info("Local text model unavailable, using hash fallback: %s", exc) + _mark_status("text", status="fallback", backend="hash", detail=str(exc)) + _TEXT_MODEL = False + return None + + +def get_vision_model(): + global _VISION_MODEL + if _VISION_MODEL is not None: + return _VISION_MODEL + spec = next(s for s in _model_specs() if s.key == "vision") + try: + _VISION_MODEL = _load_sentence_transformer(spec.model_name, spec.local_dir) + _mark_status("vision", status="ready", backend="sentence-transformers") + return _VISION_MODEL + except Exception as exc: + logger.info("Local vision model unavailable, using text fallback: %s", exc) + _mark_status("vision", status="fallback", backend="text-proxy", detail=str(exc)) + _VISION_MODEL = False + return None + + +def warmup() -> None: + ensure_manifest() + if config.EMBEDDING_PROVIDER == "local": + get_text_model() + get_vision_model() + elif config.EMBEDDING_PROVIDER == "gemini": + _mark_status("text", status="external", backend="gemini") + _mark_status("vision", status="external", backend="gemini") + elif config.EMBEDDING_PROVIDER == "ollama": + _mark_status("text", status="external", backend="ollama") + _mark_status("vision", status="external", backend="ollama") + elif config.EMBEDDING_PROVIDER == "nim": + _mark_status("text", status="external", backend="nim") + _mark_status("vision", status="external", backend="nim") + + +def model_status() -> dict[str, Any]: + payload = ensure_manifest() + payload.setdefault("runtime", {}) + payload["runtime"].update( + { + "provider": config.EMBEDDING_PROVIDER, + "apple_silicon": config.is_apple_silicon(), + "models_dir": str(config.MODELS_DIR), + } + ) + return payload diff --git a/vector_embedded_finder/reranker.py b/vector_embedded_finder/reranker.py index d486cb2..a5e5942 100644 --- a/vector_embedded_finder/reranker.py +++ b/vector_embedded_finder/reranker.py @@ -1,16 +1,69 @@ -"""Ranking utilities for blending multiple retrieval signals.""" +"""Ranking utilities for blending retrieval signals.""" from __future__ import annotations +import logging +from collections import OrderedDict +from typing import Any + +logger = logging.getLogger(__name__) + def reciprocal_rank_fusion( ranked_lists: list[list[str]], *, k: int = 60, ) -> list[str]: - """Merge ranked id lists using Reciprocal Rank Fusion (RRF).""" scores: dict[str, float] = {} for ranked in ranked_lists: for rank, doc_id in enumerate(ranked, start=1): scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank) return [doc_id for doc_id, _ in sorted(scores.items(), key=lambda kv: kv[1], reverse=True)] + + +_RERANK_CACHE: OrderedDict[tuple[str, tuple[str, ...]], list[str]] = OrderedDict() +_RERANK_CACHE_SIZE = 128 + + +def maybe_rerank( + query: str, + rows: list[dict[str, Any]], + *, + low_confidence_threshold: float = 0.15, +) -> list[dict[str, Any]]: + if len(rows) < 2: + return rows + + top_scores = [float(row.get("similarity", 0.0)) for row in rows[:5]] + if not top_scores or (max(top_scores) - min(top_scores)) >= low_confidence_threshold: + return rows + + key = (query, tuple(str(row.get("id", "")) for row in rows[:20])) + cached = _RERANK_CACHE.get(key) + if cached is not None: + index = {str(row.get("id", "")): row for row in rows} + return [index[doc_id] for doc_id in cached if doc_id in index] + + # Lightweight local heuristic rerank until a heavyweight cross-encoder is + # available: prefer rows whose caption/ocr/preview contain more query terms. + terms = [part for part in query.lower().split() if part] + + def score(row: dict[str, Any]) -> tuple[float, float]: + haystack = " ".join( + str(row.get(field, "")) + for field in ("file_name", "description", "preview") + ).lower() + meta = row.get("metadata", {}) + if isinstance(meta, dict): + haystack += " " + " ".join( + str(meta.get(field, "")) + for field in ("caption", "ocr_text", "gps_city", "exif_camera") + ).lower() + overlap = sum(1 for term in terms if term in haystack) + return (float(overlap), float(row.get("similarity", 0.0))) + + reranked = sorted(rows, key=score, reverse=True) + _RERANK_CACHE[key] = [str(row.get("id", "")) for row in reranked[:20]] + while len(_RERANK_CACHE) > _RERANK_CACHE_SIZE: + _RERANK_CACHE.popitem(last=False) + return reranked diff --git a/vector_embedded_finder/search.py b/vector_embedded_finder/search.py index 28a2f45..3f9a1a4 100644 --- a/vector_embedded_finder/search.py +++ b/vector_embedded_finder/search.py @@ -4,16 +4,20 @@ import os import re -import logging +from collections import OrderedDict from datetime import datetime, timedelta, timezone +from functools import lru_cache +from typing import Any -from . import embedder, store -from .reranker import reciprocal_rank_fusion +from . import config, embedder, store +from .reranker import maybe_rerank, reciprocal_rank_fusion -logger = logging.getLogger(__name__) - -MIN_SIMILARITY = float(os.environ.get("VEF_MIN_SIMILARITY", "0.45")) +MIN_SIMILARITY = float(os.environ.get("VEF_MIN_SIMILARITY", "0.35")) RRF_K = int(os.environ.get("VEF_RRF_K", "60")) +EMBED_CACHE_SIZE = int(os.environ.get("VEF_SEARCH_EMBED_CACHE_SIZE", "256")) +RESULT_CACHE_SIZE = int(os.environ.get("RECALL_RESULT_CACHE_SIZE", "128")) + +_RESULT_CACHE: OrderedDict[tuple[Any, ...], list[dict[str, Any]]] = OrderedDict() _MEDIA_KEYWORDS = { "image": {"image", "photo", "picture", "screenshot"}, @@ -25,22 +29,28 @@ def _tokenize_query(query: str) -> set[str]: - tokens = {t for t in re.findall(r"[a-z0-9]+", query.lower()) if len(t) > 1} - return tokens + return {t for t in re.findall(r"[a-z0-9]+", query.lower()) if len(t) > 1} -def _keyword_boost(result: dict, query: str) -> float: - media_category = str(result.get("media_category", "")).lower() - if media_category not in {"text", "document"}: - return 0.0 - +def _keyword_boost(result: dict[str, Any], query: str) -> float: words = _tokenize_query(query) if not words: return 0.0 - - text = f"{result.get('file_name', '')} {result.get('description', '')}".lower() - matches = sum(1 for w in words if w in text) - return min(0.15 * matches / max(len(words), 1), 0.15) + text = " ".join( + [ + str(result.get("file_name", "")), + str(result.get("description", "")), + str(result.get("preview", "")), + ] + ).lower() + meta = result.get("metadata", {}) + if isinstance(meta, dict): + text += " " + " ".join( + str(meta.get(key, "")) + for key in ("caption", "ocr_text", "gps_city", "exif_camera") + ).lower() + matches = sum(1 for word in words if word in text) + return min(0.2 * matches / max(len(words), 1), 0.2) def _detect_media_intent(query: str) -> str | None: @@ -75,145 +85,183 @@ def _detect_time_cutoff(query: str) -> str | None: return None -def _build_results(raw: dict) -> list[dict]: - results: list[dict] = [] - if not raw.get("ids") or not raw["ids"] or not raw["ids"][0]: - return results - - for i in range(len(raw["ids"][0])): - meta = raw["metadatas"][0][i] - distance = raw["distances"][0][i] - similarity = 1 - distance - results.append( - { - "id": raw["ids"][0][i], - "similarity": round(similarity, 4), - "file_path": meta.get("file_path", ""), - "file_name": meta.get("file_name", ""), - "media_category": meta.get("media_category", ""), - "timestamp": meta.get("timestamp", ""), - "description": meta.get("description", ""), - "source": meta.get("source", ""), - "preview": raw["documents"][0][i][:200] if raw["documents"][0][i] else "", - "metadata": {k: v for k, v in meta.items()}, - } - ) - return results - - -def search( +def _build_filters( query: str, - n_results: int = 20, + *, media_type: str | None = None, sources: list[str] | None = None, -) -> list[dict]: - try: - query_embedding = embedder.embed_query(query) - except Exception as exc: - logger.warning("Search embedding failed for query %r: %s", query, exc) - return [] - - where: dict | None = None - filters: list[dict] = [] - +) -> dict[str, Any]: + filters: dict[str, Any] = {} inferred_media = _detect_media_intent(query) if not media_type else None if media_type: - filters.append({"media_category": {"$eq": media_type}}) + filters["media_category"] = media_type elif inferred_media: - filters.append({"media_category": {"$eq": inferred_media}}) + filters["media_category"] = inferred_media inferred_sources = _detect_source_intent(query) if not sources else None source_filters = sources or inferred_sources if source_filters: - if len(source_filters) == 1: - filters.append({"source": {"$eq": source_filters[0]}}) - else: - filters.append({"source": {"$in": source_filters}}) + filters["sources"] = list(source_filters) since_cutoff = _detect_time_cutoff(query) if since_cutoff: - filters.append({"timestamp": {"$gte": since_cutoff}}) - - if len(filters) == 1: - where = filters[0] - elif len(filters) > 1: - where = {"$and": filters} - - vector_raw = store.search(query_embedding, n_results=n_results, where=where) - vector_results = _build_results(vector_raw) - - keyword_token = "" - query_tokens = sorted(_tokenize_query(query), key=len, reverse=True) - for token in query_tokens: - if len(token) >= 3: - keyword_token = token - break - - keyword_results: list[dict] = [] - if keyword_token: - keyword_raw = store.search( - query_embedding, - n_results=n_results, - where=where, - where_document={"$contains": keyword_token}, + filters["since"] = since_cutoff + return filters + + +def _candidate_to_result(candidate) -> dict[str, Any]: + meta = dict(candidate.metadata) + preview = str(meta.get("preview", ""))[:200] + return { + "id": candidate.doc_id, + "similarity": round(float(candidate.score), 4), + "file_path": meta.get("file_path", ""), + "file_name": meta.get("file_name", ""), + "media_category": meta.get("media_category", ""), + "timestamp": meta.get("timestamp", ""), + "description": meta.get("description", ""), + "source": meta.get("source", ""), + "preview": preview, + "metadata": meta, + } + + +@lru_cache(maxsize=EMBED_CACHE_SIZE) +def _embed_query_cached( + query: str, + provider: str, + model: str, + dimensions: int, +) -> tuple[float, ...]: + return tuple(float(v) for v in embedder.embed_query(query)) + + +def _query_embedding(query: str) -> list[float]: + return list( + _embed_query_cached( + query, + config.EMBEDDING_PROVIDER, + config.EMBEDDING_MODEL, + config.EMBEDDING_DIMENSIONS, ) - keyword_results = _build_results(keyword_raw) + ) + - by_id: dict[str, dict] = {r["id"]: r for r in vector_results} - for row in keyword_results: +def _result_cache_key( + query: str, + n_results: int, + media_type: str | None, + sources: list[str] | None, +) -> tuple[Any, ...]: + return ( + query.strip().lower(), + int(n_results), + media_type or "", + tuple(sorted(sources or [])), + config.EMBEDDING_PROVIDER, + config.EMBEDDING_MODEL, + store.cache_epoch(), + ) + + +def _result_cache_get(key: tuple[Any, ...]) -> list[dict[str, Any]] | None: + cached = _RESULT_CACHE.get(key) + if cached is None: + return None + _RESULT_CACHE.move_to_end(key) + return cached + + +def _result_cache_put(key: tuple[Any, ...], value: list[dict[str, Any]]) -> None: + _RESULT_CACHE[key] = value + _RESULT_CACHE.move_to_end(key) + while len(_RESULT_CACHE) > RESULT_CACHE_SIZE: + _RESULT_CACHE.popitem(last=False) + + +def search( + query: str, + n_results: int = 20, + media_type: str | None = None, + sources: list[str] | None = None, +) -> list[dict[str, Any]]: + query = query.strip() + if not query: + return [] + + cache_key = _result_cache_key(query, n_results, media_type, sources) + cached = _result_cache_get(cache_key) + if cached is not None: + return cached + + filters = _build_filters(query, media_type=media_type, sources=sources) + query_embedding = _query_embedding(query) + + dense = store.dense_search(query_embedding, n_results=n_results, filters=filters) + try: + keyword = store.keyword_search(query, n_results=n_results, filters=filters) + except Exception: + keyword = [] + + by_id: dict[str, dict[str, Any]] = {} + for candidate in dense: + row = _candidate_to_result(candidate) + by_id[row["id"]] = row + for candidate in keyword: + row = _candidate_to_result(candidate) existing = by_id.get(row["id"]) if existing is None or float(row["similarity"]) > float(existing["similarity"]): by_id[row["id"]] = row - if keyword_results: + if keyword: fused_ids = reciprocal_rank_fusion( - [ - [r["id"] for r in vector_results], - [r["id"] for r in keyword_results], - ], + [[row.doc_id for row in dense], [row.doc_id for row in keyword]], k=RRF_K, ) else: - fused_ids = [r["id"] for r in vector_results] + fused_ids = [row.doc_id for row in dense] fused_len = max(len(fused_ids), 1) + results: list[dict[str, Any]] = [] for idx, doc_id in enumerate(fused_ids): row = by_id.get(doc_id) - if not row: + if row is None: continue rrf_bonus = max(0.0, 0.08 * (1 - (idx / fused_len))) - boosted = min( - 1.0, - float(row["similarity"]) + _keyword_boost(row, query) + rrf_bonus, + row["similarity"] = round( + min(1.0, float(row["similarity"]) + _keyword_boost(row, query) + rrf_bonus), + 4, ) - row["similarity"] = round(boosted, 4) + if float(row["similarity"]) >= MIN_SIMILARITY: + results.append(row) + + if not results and dense: + fallback = [_candidate_to_result(candidate) for candidate in dense[:n_results]] + results = [row for row in fallback if float(row["similarity"]) >= MIN_SIMILARITY] - results = [by_id[doc_id] for doc_id in fused_ids if doc_id in by_id] - results = [r for r in results if float(r["similarity"]) >= MIN_SIMILARITY] - results.sort(key=lambda r: float(r["similarity"]), reverse=True) - return results + results.sort(key=lambda item: float(item["similarity"]), reverse=True) + results = maybe_rerank(query, results[:n_results]) + _result_cache_put(cache_key, results[:n_results]) + return results[:n_results] -def format_results(results: list[dict]) -> str: +def format_results(results: list[dict[str, Any]]) -> str: if not results: return "No results found." lines = [] - for i, r in enumerate(results, 1): - score_pct = f"{r['similarity'] * 100:.1f}%" - path = r["file_path"] or "(text snippet)" - category = r["media_category"] - ts = r["timestamp"][:10] if r["timestamp"] else "unknown" - - lines.append(f"**{i}. [{category}] {r['file_name'] or 'text'}** โ€” {score_pct} match") - if path: - lines.append(f" Path: `{path}`") - lines.append(f" Date: {ts} | Source: {r['source']}") - if r["preview"]: - preview = r["preview"][:150].replace("\n", " ") + for i, row in enumerate(results, 1): + score_pct = f"{float(row['similarity']) * 100:.1f}%" + path = row["file_path"] or "(text snippet)" + category = row["media_category"] + ts = str(row["timestamp"])[:10] if row["timestamp"] else "unknown" + lines.append(f"**{i}. [{category}] {row['file_name'] or 'text'}** โ€” {score_pct} match") + lines.append(f" Path: `{path}`") + lines.append(f" Date: {ts} | Source: {row['source']}") + if row["preview"]: + preview = str(row["preview"])[:150].replace("\n", " ") lines.append(f" Preview: {preview}") - if r["description"]: - lines.append(f" Description: {r['description']}") + if row["description"]: + lines.append(f" Description: {row['description']}") lines.append("") - return "\n".join(lines) diff --git a/vector_embedded_finder/store.py b/vector_embedded_finder/store.py index ea547c5..2e2bc96 100644 --- a/vector_embedded_finder/store.py +++ b/vector_embedded_finder/store.py @@ -1,122 +1,881 @@ -"""ChromaDB vector store interface.""" +"""Recall storage layer. + +sqlite is the durable source of truth. A hot vector index is layered on top of +it for fast dense search, with a best-effort Chroma dual-write during +migration. The public surface preserves the small subset of helpers the rest of +the repo already uses. +""" from __future__ import annotations -import chromadb +import json import logging +import math +import sqlite3 +import threading import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable from . import config logger = logging.getLogger(__name__) +_CONN: sqlite3.Connection | None = None +_LOCK = threading.RLock() +_HOT_INDEX = None +_CACHE_EPOCH = 0 -_client: chromadb.PersistentClient | None = None -_collection: chromadb.Collection | None = None -# Throttle repeated count()/query() error logs. chromadb can repeatedly fail -# the HNSW compactor backfill and spam the log on every call otherwise. -_last_count_warn_ts: float = 0.0 -_last_count_warn_msg: str = "" -_WARN_THROTTLE_S: float = 60.0 +def _bump_cache_epoch() -> None: + global _CACHE_EPOCH + _CACHE_EPOCH += 1 -def _get_collection() -> chromadb.Collection: - global _client, _collection - if _collection is None: - config.CHROMA_DIR.mkdir(parents=True, exist_ok=True) - _client = chromadb.PersistentClient(path=str(config.CHROMA_DIR)) - _collection = _client.get_or_create_collection( - name=config.COLLECTION_NAME, - metadata={"hnsw:space": "cosine"}, +def cache_epoch() -> int: + return _CACHE_EPOCH + + +def _connect() -> sqlite3.Connection: + global _CONN + if _CONN is None: + config.ensure_runtime_dirs() + conn = sqlite3.connect(str(config.SQLITE_PATH), check_same_thread=False) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA foreign_keys=ON") + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute("PRAGMA temp_store=MEMORY") + _init_schema(conn) + _CONN = conn + return _CONN + + +def _init_schema(conn: sqlite3.Connection) -> None: + conn.executescript( + """ + CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS manifest ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + doc_id TEXT UNIQUE NOT NULL, + path TEXT NOT NULL DEFAULT '', + file_name TEXT NOT NULL DEFAULT '', + source TEXT NOT NULL DEFAULT '', + file_type TEXT NOT NULL DEFAULT '', + media_category TEXT NOT NULL DEFAULT '', + timestamp TEXT NOT NULL DEFAULT '', + mtime REAL NOT NULL DEFAULT 0, + sha256 TEXT NOT NULL DEFAULT '', + state TEXT NOT NULL DEFAULT 'active', + description TEXT NOT NULL DEFAULT '', + file_size INTEGER NOT NULL DEFAULT 0, + preview TEXT NOT NULL DEFAULT '', + metadata_json TEXT NOT NULL DEFAULT '{}' + ); + + CREATE INDEX IF NOT EXISTS manifest_state_idx ON manifest(state); + CREATE INDEX IF NOT EXISTS manifest_path_idx ON manifest(path); + CREATE INDEX IF NOT EXISTS manifest_source_idx ON manifest(source); + CREATE INDEX IF NOT EXISTS manifest_media_idx ON manifest(media_category); + + CREATE TABLE IF NOT EXISTS vectors ( + file_id INTEGER PRIMARY KEY REFERENCES manifest(id) ON DELETE CASCADE, + embedding_json TEXT NOT NULL, + updated_at REAL NOT NULL + ); + + CREATE TABLE IF NOT EXISTS enrichment ( + file_id INTEGER PRIMARY KEY REFERENCES manifest(id) ON DELETE CASCADE, + caption TEXT NOT NULL DEFAULT '', + ocr_text TEXT NOT NULL DEFAULT '', + gps_lat REAL, + gps_lon REAL, + gps_city TEXT NOT NULL DEFAULT '', + face_count INTEGER NOT NULL DEFAULT 0, + exif_date TEXT NOT NULL DEFAULT '', + exif_camera TEXT NOT NULL DEFAULT '' + ); + + CREATE VIRTUAL TABLE IF NOT EXISTS fts_content USING fts5( + file_id UNINDEXED, + path, + title, + body, + metadata_text, + tokenize='porter unicode61' + ); + """ + ) + + +def _meta_get(key: str, default: str = "") -> str: + row = _connect().execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone() + if row is None: + return default + return str(row["value"]) + + +def _meta_set(key: str, value: str) -> None: + _connect().execute( + "INSERT INTO meta(key, value) VALUES(?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (key, value), + ) + _connect().commit() + + +def _json_dumps(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":")) + + +def _json_loads(value: str | None) -> dict[str, Any]: + if not value: + return {} + try: + payload = json.loads(value) + if isinstance(payload, dict): + return payload + except Exception: + pass + return {} + + +def _normalize(vec: list[float]) -> list[float]: + norm = math.sqrt(sum(v * v for v in vec)) + if norm <= 0: + return vec + return [float(v / norm) for v in vec] + + +def _cosine_distance(left: list[float], right: list[float]) -> float: + if not left or not right: + return 1.0 + size = min(len(left), len(right)) + dot = sum(float(left[i]) * float(right[i]) for i in range(size)) + return max(0.0, min(2.0, 1.0 - dot)) + + +@dataclass +class Candidate: + file_id: int + doc_id: str + distance: float + score: float + metadata: dict[str, Any] + + +class _MemoryHotIndex: + def __init__(self) -> None: + self._vectors: dict[int, list[float]] = {} + self._dirty = False + self.backend = "memory" + self.path = config.HNSW_DIR / "hot.index" + + def load_from_rows(self, rows: Iterable[tuple[int, list[float]]]) -> None: + self._vectors = {file_id: vec for file_id, vec in rows} + self._dirty = False + + def add_or_update(self, file_id: int, embedding: list[float]) -> None: + self._vectors[file_id] = embedding + + def delete(self, file_id: int) -> None: + self._vectors.pop(file_id, None) + + def search(self, embedding: list[float], limit: int) -> list[tuple[int, float]]: + rows = [ + (file_id, _cosine_distance(embedding, candidate)) + for file_id, candidate in self._vectors.items() + ] + rows.sort(key=lambda item: item[1]) + return rows[:limit] + + def mark_dirty(self) -> None: + self._dirty = True + + def status(self) -> dict[str, Any]: + return { + "backend": self.backend, + "path": str(self.path), + "ready": True, + "dirty": self._dirty, + "count": len(self._vectors), + } + + +class _HnswHotIndex(_MemoryHotIndex): + def __init__(self) -> None: + super().__init__() + import hnswlib + + self._hnswlib = hnswlib + self.backend = "hnswlib" + self.path = config.HNSW_DIR / "hot.bin" + self._index = None + self._labels: set[int] = set() + + def load_from_rows(self, rows: Iterable[tuple[int, list[float]]]) -> None: + tuples = list(rows) + self._vectors = {file_id: vec for file_id, vec in tuples} + max_elements = max(len(self._vectors) + 64, 256) + index = self._hnswlib.Index(space="cosine", dim=config.EMBEDDING_DIMENSIONS) + index.init_index( + max_elements=max_elements, + ef_construction=200, + M=32, + allow_replace_deleted=True, ) - return _collection + if tuples: + embeddings = [vec for _, vec in tuples] + labels = [file_id for file_id, _ in tuples] + index.add_items(embeddings, labels) + index.set_ef(max(50, min(200, len(tuples) + 10))) + self._index = index + self._labels = set(self._vectors) + self._dirty = False + try: + config.HNSW_DIR.mkdir(parents=True, exist_ok=True) + index.save_index(str(self.path)) + except Exception as exc: + logger.debug("Could not save hnsw index: %s", exc) + + def add_or_update(self, file_id: int, embedding: list[float]) -> None: + if self._index is None: + self.load_from_rows([(file_id, embedding)]) + return + if file_id in self._labels: + # hnswlib update semantics are awkward with stable labels; keep the + # in-memory view correct and request a rebuild. + self._vectors[file_id] = embedding + self.mark_dirty() + return + try: + self._index.add_items([embedding], [file_id]) + self._labels.add(file_id) + self._vectors[file_id] = embedding + self._index.save_index(str(self.path)) + except Exception as exc: + logger.debug("Incremental hnsw update failed, will rebuild: %s", exc) + self._vectors[file_id] = embedding + self.mark_dirty() + + def delete(self, file_id: int) -> None: + self._vectors.pop(file_id, None) + if self._index is None: + return + if file_id in self._labels: + try: + self._index.mark_deleted(file_id) + self._labels.discard(file_id) + self._index.save_index(str(self.path)) + except Exception: + self.mark_dirty() + + def search(self, embedding: list[float], limit: int) -> list[tuple[int, float]]: + if self._dirty or self._index is None or not self._vectors: + return super().search(embedding, limit) + labels, distances = self._index.knn_query(embedding, k=min(limit, len(self._vectors))) + results: list[tuple[int, float]] = [] + for label, distance in zip(labels[0], distances[0]): + results.append((int(label), float(distance))) + return results + + +def _hot_index(): + global _HOT_INDEX + if _HOT_INDEX is None: + try: + _HOT_INDEX = _HnswHotIndex() + except Exception as exc: + logger.info("hnswlib unavailable, using in-memory hot index: %s", exc) + _HOT_INDEX = _MemoryHotIndex() + return _HOT_INDEX + + +def initialize() -> None: + _connect() + rebuild_hot_index() + + +def _active_vector_rows() -> list[tuple[int, list[float]]]: + conn = _connect() + rows = conn.execute( + """ + SELECT m.id, v.embedding_json + FROM manifest m + JOIN vectors v ON v.file_id = m.id + WHERE m.state = 'active' + """ + ).fetchall() + result: list[tuple[int, list[float]]] = [] + for row in rows: + try: + embedding = [float(x) for x in json.loads(row["embedding_json"])] + except Exception: + continue + result.append((int(row["id"]), embedding)) + return result + + +def rebuild_hot_index() -> dict[str, Any]: + with _LOCK: + start = time.time() + rows = _active_vector_rows() + _hot_index().load_from_rows(rows) + _meta_set("index_last_rebuild_at", str(time.time())) + _bump_cache_epoch() + return { + "status": "ok", + "count": len(rows), + "duration_s": round(time.time() - start, 4), + **_hot_index().status(), + } + + +def _metadata_text(metadata: dict[str, Any], document: str, enrichment: dict[str, Any] | None) -> str: + parts = [ + str(metadata.get("file_name", "")), + str(metadata.get("description", "")), + str(metadata.get("source", "")), + str(metadata.get("media_category", "")), + document, + ] + if enrichment: + for key in ("caption", "ocr_text", "gps_city", "exif_date", "exif_camera"): + value = enrichment.get(key) + if value: + parts.append(str(value)) + return "\n".join(p for p in parts if p).strip() -def _throttled_warn(label: str, exc: Exception) -> None: - global _last_count_warn_ts, _last_count_warn_msg - msg = f"{label}: {exc}" - now = time.time() - if msg != _last_count_warn_msg or (now - _last_count_warn_ts) > _WARN_THROTTLE_S: - logger.warning("%s", msg) - _last_count_warn_ts = now - _last_count_warn_msg = msg +def _row_to_metadata(row: sqlite3.Row) -> dict[str, Any]: + metadata = _json_loads(row["metadata_json"]) + metadata.update( + { + "file_path": row["path"], + "file_name": row["file_name"], + "file_type": row["file_type"], + "media_category": row["media_category"], + "timestamp": row["timestamp"], + "source": row["source"], + "description": row["description"], + "file_size": row["file_size"], + "preview": row["preview"], + "path": row["path"], + "state": row["state"], + } + ) + return metadata -def _safe_count(coll: chromadb.Collection) -> int: +def _maybe_dual_write_chroma( + doc_id: str, + embedding: list[float], + metadata: dict[str, Any], + document: str, +) -> None: + if not config.DUAL_WRITE_CHROMA: + return try: - return int(coll.count()) + import chromadb + + config.CHROMA_DIR.mkdir(parents=True, exist_ok=True) + client = chromadb.PersistentClient(path=str(config.CHROMA_DIR)) + coll = client.get_or_create_collection( + name=config.COLLECTION_NAME, + metadata={"hnsw:space": "cosine"}, + ) + coll.upsert( + ids=[doc_id], + embeddings=[embedding], + metadatas=[metadata], + documents=[document], + ) except Exception as exc: - _throttled_warn("Chroma count failed", exc) - return 0 + logger.debug("Legacy chroma dual-write failed: %s", exc) + + +def _ensure_manifest_row( + conn: sqlite3.Connection, + doc_id: str, + metadata: dict[str, Any], + document: str, +) -> int: + path = str(metadata.get("file_path", "")) + file_name = str(metadata.get("file_name", "")) + file_type = str(metadata.get("file_type", "")) + media_category = str(metadata.get("media_category", "")) + timestamp = str(metadata.get("timestamp", "")) + source = str(metadata.get("source", "")) + description = str(metadata.get("description", "")) + file_size = int(metadata.get("file_size", 0) or 0) + sha256 = str(metadata.get("sha256", doc_id)) + mtime = float(metadata.get("mtime", 0.0) or 0.0) + preview = (document or "")[:500] + conn.execute( + """ + INSERT INTO manifest( + doc_id, path, file_name, source, file_type, media_category, + timestamp, mtime, sha256, state, description, file_size, + preview, metadata_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'active', ?, ?, ?, ?) + ON CONFLICT(doc_id) DO UPDATE SET + path = excluded.path, + file_name = excluded.file_name, + source = excluded.source, + file_type = excluded.file_type, + media_category = excluded.media_category, + timestamp = excluded.timestamp, + mtime = excluded.mtime, + sha256 = excluded.sha256, + state = 'active', + description = excluded.description, + file_size = excluded.file_size, + preview = excluded.preview, + metadata_json = excluded.metadata_json + """, + ( + doc_id, + path, + file_name, + source, + file_type, + media_category, + timestamp, + mtime, + sha256, + description, + file_size, + preview, + _json_dumps(metadata), + ), + ) + row = conn.execute("SELECT id FROM manifest WHERE doc_id = ?", (doc_id,)).fetchone() + return int(row["id"]) def add( doc_id: str, embedding: list[float], - metadata: dict, + metadata: dict[str, Any], document: str = "", + enrichment: dict[str, Any] | None = None, ) -> None: - coll = _get_collection() - coll.upsert( - ids=[doc_id], - embeddings=[embedding], - metadatas=[metadata], - documents=[document], - ) + embedding = _normalize([float(v) for v in embedding[: config.EMBEDDING_DIMENSIONS]]) + with _LOCK: + conn = _connect() + file_id = _ensure_manifest_row(conn, doc_id, metadata, document) + conn.execute( + "INSERT OR REPLACE INTO vectors(file_id, embedding_json, updated_at) VALUES (?, ?, ?)", + (file_id, _json_dumps(embedding), time.time()), + ) + if enrichment is not None: + conn.execute( + """ + INSERT INTO enrichment( + file_id, caption, ocr_text, gps_lat, gps_lon, gps_city, + face_count, exif_date, exif_camera + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(file_id) DO UPDATE SET + caption = excluded.caption, + ocr_text = excluded.ocr_text, + gps_lat = excluded.gps_lat, + gps_lon = excluded.gps_lon, + gps_city = excluded.gps_city, + face_count = excluded.face_count, + exif_date = excluded.exif_date, + exif_camera = excluded.exif_camera + """, + ( + file_id, + str(enrichment.get("caption", "")), + str(enrichment.get("ocr_text", "")), + enrichment.get("gps_lat"), + enrichment.get("gps_lon"), + str(enrichment.get("gps_city", "")), + int(enrichment.get("face_count", 0) or 0), + str(enrichment.get("exif_date", "")), + str(enrichment.get("exif_camera", "")), + ), + ) + conn.execute("DELETE FROM fts_content WHERE file_id = ?", (str(file_id),)) + conn.execute( + "INSERT INTO fts_content(file_id, path, title, body, metadata_text) VALUES (?, ?, ?, ?, ?)", + ( + str(file_id), + str(metadata.get("file_path", "")), + str(metadata.get("file_name", "")), + document, + _metadata_text(metadata, document, enrichment), + ), + ) + conn.commit() + _hot_index().add_or_update(file_id, embedding) + _maybe_dual_write_chroma(doc_id, embedding, metadata, document) + _bump_cache_epoch() -def search( - query_embedding: list[float], - n_results: int = 5, - where: dict | None = None, - where_document: dict | None = None, -) -> dict: - coll = _get_collection() - total = _safe_count(coll) - if total <= 0: - return {"ids": [[]], "metadatas": [[]], "documents": [[]], "distances": [[]]} - kwargs = { - "query_embeddings": [query_embedding], - "n_results": min(n_results, total), - "include": ["metadatas", "documents", "distances"], - } - if where: - kwargs["where"] = where - if where_document: - kwargs["where_document"] = where_document - try: - return coll.query(**kwargs) - except Exception as exc: - _throttled_warn("Chroma query failed", exc) - return {"ids": [[]], "metadatas": [[]], "documents": [[]], "distances": [[]]} +def _find_doc_id_by_path(path: str) -> str | None: + row = _connect().execute( + "SELECT doc_id FROM manifest WHERE path = ? AND state = 'active'", + (path,), + ).fetchone() + return str(row["doc_id"]) if row else None -def exists(doc_id: str) -> bool: - coll = _get_collection() - result = coll.get(ids=[doc_id]) - return len(result["ids"]) > 0 +def delete(doc_id: str) -> None: + with _LOCK: + conn = _connect() + row = conn.execute("SELECT id FROM manifest WHERE doc_id = ?", (doc_id,)).fetchone() + if row is None: + return + file_id = int(row["id"]) + conn.execute("UPDATE manifest SET state = 'deleted' WHERE id = ?", (file_id,)) + conn.execute("DELETE FROM vectors WHERE file_id = ?", (file_id,)) + conn.execute("DELETE FROM enrichment WHERE file_id = ?", (file_id,)) + conn.execute("DELETE FROM fts_content WHERE file_id = ?", (str(file_id),)) + conn.commit() + _hot_index().delete(file_id) + _bump_cache_epoch() -def delete(doc_id: str) -> None: - coll = _get_collection() - coll.delete(ids=[doc_id]) +def delete_by_path(path: str | Path) -> None: + doc_id = _find_doc_id_by_path(str(path)) + if doc_id: + delete(doc_id) + + +def retire_path_versions(path: str | Path, *, keep_doc_id: str) -> int: + with _LOCK: + conn = _connect() + rows = conn.execute( + """ + SELECT id + FROM manifest + WHERE path = ? AND state = 'active' AND doc_id != ? + """, + (str(path), keep_doc_id), + ).fetchall() + if not rows: + return 0 + file_ids = [int(row["id"]) for row in rows] + for file_id in file_ids: + conn.execute("UPDATE manifest SET state = 'deleted' WHERE id = ?", (file_id,)) + conn.execute("DELETE FROM vectors WHERE file_id = ?", (file_id,)) + conn.execute("DELETE FROM enrichment WHERE file_id = ?", (file_id,)) + conn.execute("DELETE FROM fts_content WHERE file_id = ?", (str(file_id),)) + conn.commit() + for file_id in file_ids: + _hot_index().delete(file_id) + _bump_cache_epoch() + return len(file_ids) + + +def exists(doc_id: str) -> bool: + row = _connect().execute( + "SELECT 1 FROM manifest WHERE doc_id = ? AND state = 'active'", + (doc_id,), + ).fetchone() + return row is not None def count() -> int: - return _safe_count(_get_collection()) + row = _connect().execute( + "SELECT COUNT(*) AS n FROM manifest WHERE state = 'active'" + ).fetchone() + return int(row["n"]) if row else 0 -def list_all(limit: int = 100, offset: int = 0) -> dict: - coll = _get_collection() - return coll.get( - limit=limit, - offset=offset, - include=["metadatas", "documents"], - ) +def list_all(limit: int = 100, offset: int = 0) -> dict[str, Any]: + rows = _connect().execute( + """ + SELECT * + FROM manifest + WHERE state = 'active' + ORDER BY timestamp DESC, id DESC + LIMIT ? OFFSET ? + """, + (limit, offset), + ).fetchall() + ids: list[str] = [] + metadatas: list[dict[str, Any]] = [] + documents: list[str] = [] + for row in rows: + ids.append(str(row["doc_id"])) + metadatas.append(_row_to_metadata(row)) + documents.append(str(row["preview"])) + return {"ids": ids, "metadatas": metadatas, "documents": documents} + + +def update_metadata(doc_id: str, metadata: dict[str, Any]) -> None: + with _LOCK: + conn = _connect() + row = conn.execute("SELECT id FROM manifest WHERE doc_id = ?", (doc_id,)).fetchone() + if row is None: + return + file_id = int(row["id"]) + existing = get_by_doc_ids([doc_id]).get(doc_id, {}) + merged = {**existing.get("metadata", {}), **metadata} + _ensure_manifest_row(conn, doc_id, merged, str(merged.get("preview", ""))) + conn.execute("DELETE FROM fts_content WHERE file_id = ?", (str(file_id),)) + conn.execute( + "INSERT INTO fts_content(file_id, path, title, body, metadata_text) VALUES (?, ?, ?, ?, ?)", + ( + str(file_id), + str(merged.get("file_path", "")), + str(merged.get("file_name", "")), + str(merged.get("preview", "")), + _metadata_text(merged, str(merged.get("preview", "")), None), + ), + ) + conn.commit() + _bump_cache_epoch() + + +def get_sources() -> list[str]: + rows = _connect().execute( + "SELECT DISTINCT source FROM manifest WHERE state = 'active' AND source != '' ORDER BY source" + ).fetchall() + return [str(row["source"]) for row in rows] -def update_metadata(doc_id: str, metadata: dict) -> None: - """Update metadata for an existing entry (no re-embedding).""" - coll = _get_collection() - coll.update(ids=[doc_id], metadatas=[metadata]) +def get_by_doc_ids(doc_ids: list[str]) -> dict[str, dict[str, Any]]: + if not doc_ids: + return {} + placeholders = ",".join("?" for _ in doc_ids) + rows = _connect().execute( + f""" + SELECT m.*, e.caption, e.ocr_text, e.gps_lat, e.gps_lon, e.gps_city, + e.face_count, e.exif_date, e.exif_camera + FROM manifest m + LEFT JOIN enrichment e ON e.file_id = m.id + WHERE m.doc_id IN ({placeholders}) + """, + tuple(doc_ids), + ).fetchall() + payload: dict[str, dict[str, Any]] = {} + for row in rows: + metadata = _row_to_metadata(row) + metadata.update( + { + "caption": row["caption"] if "caption" in row.keys() else "", + "ocr_text": row["ocr_text"] if "ocr_text" in row.keys() else "", + "gps_lat": row["gps_lat"] if "gps_lat" in row.keys() else None, + "gps_lon": row["gps_lon"] if "gps_lon" in row.keys() else None, + "gps_city": row["gps_city"] if "gps_city" in row.keys() else "", + "face_count": row["face_count"] if "face_count" in row.keys() else 0, + "exif_date": row["exif_date"] if "exif_date" in row.keys() else "", + "exif_camera": row["exif_camera"] if "exif_camera" in row.keys() else "", + } + ) + payload[str(row["doc_id"])] = { + "file_id": int(row["id"]), + "doc_id": str(row["doc_id"]), + "metadata": metadata, + } + return payload + + +def _match_filters(metadata: dict[str, Any], filters: dict[str, Any] | None) -> bool: + if not filters: + return True + media_category = filters.get("media_category") + if media_category and metadata.get("media_category") != media_category: + return False + sources = filters.get("sources") + if sources and metadata.get("source") not in set(sources): + return False + since = filters.get("since") + if since and str(metadata.get("timestamp", "")) < str(since): + return False + must_exist = filters.get("path_exists") + if must_exist and metadata.get("file_path") and not Path(str(metadata["file_path"])).exists(): + return False + extra = filters.get("metadata") or {} + for key, value in extra.items(): + if metadata.get(key) != value: + return False + return True + + +def dense_search( + query_embedding: list[float], + n_results: int = 20, + filters: dict[str, Any] | None = None, + oversample: int | None = None, +) -> list[Candidate]: + limit = oversample or max(50, n_results * 5) + rows = _hot_index().search(_normalize(query_embedding), limit) + doc_rows: list[Candidate] = [] + file_ids = [file_id for file_id, _ in rows] + if not file_ids: + return [] + placeholders = ",".join("?" for _ in file_ids) + data = _connect().execute( + f""" + SELECT m.*, e.caption, e.ocr_text, e.gps_lat, e.gps_lon, e.gps_city, + e.face_count, e.exif_date, e.exif_camera + FROM manifest m + LEFT JOIN enrichment e ON e.file_id = m.id + WHERE m.id IN ({placeholders}) AND m.state = 'active' + """, + tuple(file_ids), + ).fetchall() + by_file_id = {int(row["id"]): row for row in data} + for file_id, distance in rows: + row = by_file_id.get(file_id) + if row is None: + continue + metadata = _row_to_metadata(row) + metadata.update( + { + "caption": row["caption"] if "caption" in row.keys() else "", + "ocr_text": row["ocr_text"] if "ocr_text" in row.keys() else "", + "gps_lat": row["gps_lat"] if "gps_lat" in row.keys() else None, + "gps_lon": row["gps_lon"] if "gps_lon" in row.keys() else None, + "gps_city": row["gps_city"] if "gps_city" in row.keys() else "", + "face_count": row["face_count"] if "face_count" in row.keys() else 0, + "exif_date": row["exif_date"] if "exif_date" in row.keys() else "", + "exif_camera": row["exif_camera"] if "exif_camera" in row.keys() else "", + } + ) + if not _match_filters(metadata, filters): + continue + doc_rows.append( + Candidate( + file_id=file_id, + doc_id=str(row["doc_id"]), + distance=float(distance), + score=max(0.0, 1.0 - float(distance)), + metadata=metadata, + ) + ) + if len(doc_rows) >= n_results: + break + if len(doc_rows) >= n_results: + return doc_rows + return _linear_dense_search(query_embedding, n_results=n_results, filters=filters, skip={c.doc_id for c in doc_rows}, append=doc_rows) + + +def _linear_dense_search( + query_embedding: list[float], + *, + n_results: int, + filters: dict[str, Any] | None, + skip: set[str], + append: list[Candidate] | None = None, +) -> list[Candidate]: + rows = _connect().execute( + """ + SELECT m.*, v.embedding_json, e.caption, e.ocr_text, e.gps_lat, e.gps_lon, + e.gps_city, e.face_count, e.exif_date, e.exif_camera + FROM manifest m + JOIN vectors v ON v.file_id = m.id + LEFT JOIN enrichment e ON e.file_id = m.id + WHERE m.state = 'active' + """ + ).fetchall() + ranked: list[Candidate] = list(append or []) + query_embedding = _normalize(query_embedding) + scored: list[Candidate] = [] + for row in rows: + doc_id = str(row["doc_id"]) + if doc_id in skip: + continue + metadata = _row_to_metadata(row) + metadata.update( + { + "caption": row["caption"] if "caption" in row.keys() else "", + "ocr_text": row["ocr_text"] if "ocr_text" in row.keys() else "", + "gps_lat": row["gps_lat"] if "gps_lat" in row.keys() else None, + "gps_lon": row["gps_lon"] if "gps_lon" in row.keys() else None, + "gps_city": row["gps_city"] if "gps_city" in row.keys() else "", + "face_count": row["face_count"] if "face_count" in row.keys() else 0, + "exif_date": row["exif_date"] if "exif_date" in row.keys() else "", + "exif_camera": row["exif_camera"] if "exif_camera" in row.keys() else "", + } + ) + if not _match_filters(metadata, filters): + continue + embedding = [float(x) for x in json.loads(row["embedding_json"])] + distance = _cosine_distance(query_embedding, embedding) + scored.append( + Candidate( + file_id=int(row["id"]), + doc_id=doc_id, + distance=distance, + score=max(0.0, 1.0 - distance), + metadata=metadata, + ) + ) + scored.sort(key=lambda item: item.distance) + ranked.extend(scored[: max(0, n_results - len(ranked))]) + return ranked[:n_results] + + +def keyword_search( + query: str, + n_results: int = 20, + filters: dict[str, Any] | None = None, +) -> list[Candidate]: + tokens = [token for token in query.lower().split() if token.strip()] + if not tokens: + return [] + match_expr = " OR ".join(f'"{token.replace("\"", "")}"' for token in tokens[:10]) + rows = _connect().execute( + """ + SELECT m.*, bm25(fts_content) AS rank, + e.caption, e.ocr_text, e.gps_lat, e.gps_lon, e.gps_city, + e.face_count, e.exif_date, e.exif_camera + FROM fts_content + JOIN manifest m ON m.id = CAST(fts_content.file_id AS INTEGER) + LEFT JOIN enrichment e ON e.file_id = m.id + WHERE fts_content MATCH ? AND m.state = 'active' + ORDER BY rank + LIMIT ? + """, + (match_expr, max(50, n_results * 5)), + ).fetchall() + results: list[Candidate] = [] + for row in rows: + metadata = _row_to_metadata(row) + metadata.update( + { + "caption": row["caption"] if "caption" in row.keys() else "", + "ocr_text": row["ocr_text"] if "ocr_text" in row.keys() else "", + "gps_lat": row["gps_lat"] if "gps_lat" in row.keys() else None, + "gps_lon": row["gps_lon"] if "gps_lon" in row.keys() else None, + "gps_city": row["gps_city"] if "gps_city" in row.keys() else "", + "face_count": row["face_count"] if "face_count" in row.keys() else 0, + "exif_date": row["exif_date"] if "exif_date" in row.keys() else "", + "exif_camera": row["exif_camera"] if "exif_camera" in row.keys() else "", + } + ) + if not _match_filters(metadata, filters): + continue + rank = float(row["rank"]) if row["rank"] is not None else 0.0 + score = 1.0 / (1.0 + max(rank, 0.0)) + results.append( + Candidate( + file_id=int(row["id"]), + doc_id=str(row["doc_id"]), + distance=max(0.0, 1.0 - score), + score=score, + metadata=metadata, + ) + ) + if len(results) >= n_results: + break + return results + + +def index_status() -> dict[str, Any]: + return { + "count": count(), + "sqlite_path": str(config.SQLITE_PATH), + "last_rebuild_at": _meta_get("index_last_rebuild_at", ""), + **_hot_index().status(), + } diff --git a/vector_embedded_finder/watcher.py b/vector_embedded_finder/watcher.py index 75faa95..d42bc06 100644 --- a/vector_embedded_finder/watcher.py +++ b/vector_embedded_finder/watcher.py @@ -1,15 +1,12 @@ -"""Filesystem watcher โ€” indexes new/modified files within ~10 seconds. - -Uses the watchdog library to monitor configured directories. A 2-second -debounce prevents re-indexing partially-written files. -""" +"""Filesystem watcher with debounce, dedupe, and delete handling.""" from __future__ import annotations import logging +import queue import threading import time -from collections import defaultdict +from dataclasses import dataclass from pathlib import Path from typing import Callable @@ -20,103 +17,91 @@ logger = logging.getLogger(__name__) DEBOUNCE_SECONDS = 2.0 +MAX_QUEUE_SIZE = 500 -class _DebounceTimer: - """Fires callback once after no new events for `delay` seconds.""" - - def __init__(self, delay: float, callback: Callable[[Path], None], path: Path): - self._delay = delay - self._callback = callback - self._path = path - self._timer: threading.Timer | None = None - self._lock = threading.Lock() - - def touch(self) -> None: - with self._lock: - if self._timer: - self._timer.cancel() - self._timer = threading.Timer(self._delay, self._fire) - self._timer.daemon = True - self._timer.start() - - def _fire(self) -> None: - try: - self._callback(self._path) - except Exception as e: - logger.error("watcher callback error for %s: %s", self._path, e) +@dataclass +class _PendingEvent: + op: str + path: Path + ready_at: float class _FileEventHandler: - """Watchdog event handler that debounces and queues ingest calls.""" - - def __init__(self, callback: Callable[[Path], None]): - self._callback = callback - self._timers: dict[str, _DebounceTimer] = {} + def __init__(self, event_queue: "queue.Queue[_PendingEvent]"): + self._queue = event_queue + self._pending: dict[str, _PendingEvent] = {} self._lock = threading.Lock() - # watchdog calls these methods def on_created(self, event) -> None: - self._handle(event) + self._handle("upsert", event) def on_modified(self, event) -> None: - self._handle(event) + self._handle("upsert", event) def on_moved(self, event) -> None: - # Index the destination path (file was renamed/moved to a new location) if event.is_directory: return - dest = Path(event.dest_path) - if dest.name.startswith("._") or not utils.is_supported(dest): - return - key = str(dest) - with self._lock: - if key not in self._timers: - self._timers[key] = _DebounceTimer(DEBOUNCE_SECONDS, self._on_ready, dest) - self._timers[key].touch() + self._enqueue("delete", Path(event.src_path)) + self._enqueue("upsert", Path(event.dest_path)) + + def on_deleted(self, event) -> None: + self._handle("delete", event) - def _handle(self, event) -> None: + def _handle(self, op: str, event) -> None: if event.is_directory: return - path = Path(event.src_path) + self._enqueue(op, Path(event.src_path)) + + def _enqueue(self, op: str, path: Path) -> None: if path.name.startswith("._"): return - if not utils.is_supported(path): + if op == "upsert" and not utils.is_supported(path): return - - key = str(path) + ready_at = time.time() + DEBOUNCE_SECONDS + key = f"{op}:{path}" with self._lock: - if key not in self._timers: - self._timers[key] = _DebounceTimer(DEBOUNCE_SECONDS, self._on_ready, path) - self._timers[key].touch() + self._pending[key] = _PendingEvent(op=op, path=path, ready_at=ready_at) - def _on_ready(self, path: Path) -> None: - # CPU guard before ingesting - while psutil.cpu_percent(interval=0.5) > config.CPU_GUARD_PERCENT: - logger.debug("CPU busy, deferring ingest of %s by 5s", path) - time.sleep(5) - self._callback(path) + def flush_ready(self) -> None: + now = time.time() + ready: list[_PendingEvent] = [] + with self._lock: + for key, pending in list(self._pending.items()): + if pending.ready_at <= now: + ready.append(pending) + del self._pending[key] + for pending in ready: + try: + self._queue.put_nowait(pending) + except queue.Full: + logger.debug("Watcher queue full, dropping event for %s", pending.path) class FileWatcher: - """Start / stop watching a list of directories.""" - def __init__(self): self._observer = None self._handler: _FileEventHandler | None = None + self._queue: queue.Queue[_PendingEvent] = queue.Queue(maxsize=MAX_QUEUE_SIZE) + self._worker: threading.Thread | None = None + self._flusher: threading.Thread | None = None + self._stopped = threading.Event() def start( self, directories: list[Path], callback: Callable[[Path], None], + delete_callback: Callable[[Path], None] | None = None, ) -> None: - """Begin watching `directories`. `callback(path)` called for each new/modified file.""" - from watchdog.observers import Observer from watchdog.events import FileSystemEventHandler + try: + from watchdog.observers.fsevents import FSEventsObserver as Observer + except Exception: + from watchdog.observers import Observer - self._handler = _FileEventHandler(callback) + self._stopped.clear() + self._handler = _FileEventHandler(self._queue) - # Wrap our handler in a watchdog-compatible shim class WatchdogShim(FileSystemEventHandler): def __init__(self, inner: _FileEventHandler): self._inner = inner @@ -130,24 +115,63 @@ def on_modified(self, event): def on_moved(self, event): self._inner.on_moved(event) - shim = WatchdogShim(self._handler) + def on_deleted(self, event): + self._inner.on_deleted(event) self._observer = Observer() - for d in directories: - expanded = d.expanduser().resolve() - if expanded.exists(): - self._observer.schedule(shim, str(expanded), recursive=True) - logger.info("Watching directory: %s", expanded) - else: - logger.warning("Watch directory does not exist, skipping: %s", expanded) - + shim = WatchdogShim(self._handler) + for directory in directories: + resolved = directory.expanduser().resolve() + if resolved.exists(): + self._observer.schedule(shim, str(resolved), recursive=True) + logger.info("Watching directory: %s", resolved) self._observer.start() + def _flush_loop() -> None: + while not self._stopped.is_set(): + if self._handler is not None: + self._handler.flush_ready() + time.sleep(0.5) + + def _worker_loop() -> None: + while not self._stopped.is_set(): + try: + pending = self._queue.get(timeout=0.5) + except queue.Empty: + continue + while psutil.cpu_percent(interval=0.2) > config.CPU_GUARD_PERCENT: + time.sleep(1.0) + try: + if pending.op == "delete": + if delete_callback is not None: + delete_callback(pending.path) + elif pending.path.exists(): + callback(pending.path) + except Exception as exc: + logger.error("Watcher callback error for %s: %s", pending.path, exc) + finally: + self._queue.task_done() + + self._flusher = threading.Thread(target=_flush_loop, daemon=True, name="recall-watch-flusher") + self._worker = threading.Thread(target=_worker_loop, daemon=True, name="recall-watch-worker") + self._flusher.start() + self._worker.start() + def stop(self) -> None: + self._stopped.set() if self._observer: self._observer.stop() self._observer.join() self._observer = None + if self._worker: + self._worker.join(timeout=1.0) + self._worker = None + if self._flusher: + self._flusher.join(timeout=1.0) + self._flusher = None def is_alive(self) -> bool: return self._observer is not None and self._observer.is_alive() + + def queued(self) -> int: + return self._queue.qsize()