diff --git a/bun.lock b/bun.lock index 867b9632219..4a75a735e53 100644 --- a/bun.lock +++ b/bun.lock @@ -161,14 +161,17 @@ }, "devDependencies": { "@types/bun": "catalog:", + "@zvec/zvec": "catalog:", "fastembed": "catalog:", "onnxruntime-node": "catalog:", }, "peerDependencies": { + "@zvec/zvec": ">=0.1.0", "fastembed": "2.1.0", "onnxruntime-node": "1.21.0", }, "optionalPeers": [ + "@zvec/zvec", "fastembed", "onnxruntime-node", ], @@ -323,6 +326,9 @@ }, }, }, + "trustedDependencies": [ + "@zvec/zvec", + ], "patchedDependencies": { "@ark/schema@0.56.1": "patches/@ark%2Fschema@0.56.1.patch", "puppeteer-core@25.3.0": "patches/puppeteer-core@25.3.0.patch", @@ -371,6 +377,7 @@ "@types/turndown": "5.0.6", "@typescript/native-preview": "7.0.0-dev.20260609.1", "@xterm/headless": "^6.0.0", + "@zvec/zvec": "0.5.0", "arktype": "2.2.2", "chalk": "^5.6.2", "chart.js": "^4.5.1", @@ -932,6 +939,16 @@ "@xterm/headless": ["@xterm/headless@6.0.0", "", {}, "sha512-5Yj1QINYCyzrZtf8OFIHi47iQtI+0qYFPHmouEfG8dHNxbZ9Tb9YGSuLcsEwj9Z+OL75GJqPyJbyoFer80a2Hw=="], + "@zvec/bindings-darwin-arm64": ["@zvec/bindings-darwin-arm64@0.5.0", "", { "os": "darwin", "cpu": "arm64" }, "sha512-NeM1XEvc/M0VTKx7dbr/dFsuPa0GbqNbjjAc+tl6NPO/ZYN1z8ORD2nGyheLFLayPBIVg6EliS7V9TIeV31RkQ=="], + + "@zvec/bindings-linux-arm64": ["@zvec/bindings-linux-arm64@0.5.0", "", { "os": "linux", "cpu": "arm64" }, "sha512-Mb3cd2ZHY4xoucpPWNQ/kZUleIsnbsfk1QP67iUPIKen6n+Lydt+S5urQUkfTs/oGv0aLt5i9+b4mzTpPlgsTQ=="], + + "@zvec/bindings-linux-x64": ["@zvec/bindings-linux-x64@0.5.0", "", { "os": "linux", "cpu": "x64" }, "sha512-08Cj0xJgOmr0GFAWlUFPyJYq6ysLurcA7Enx4WClCBcbtwQaZ9SHc3A2VRqLZ8dgIm5Dei9ZbDnwMwgMO2yV0w=="], + + "@zvec/bindings-win32-x64": ["@zvec/bindings-win32-x64@0.5.0", "", { "os": "win32", "cpu": "x64" }, "sha512-EkpcCO4IEDhXYejrQPIgXO49vdicQl3vCsFn3AV/oRPjLsR4lgRCncveq4ejWRdnIIIzZX6nX0rZQ1V0KPJ62A=="], + + "@zvec/zvec": ["@zvec/zvec@0.5.0", "", { "dependencies": { "bindings": "^1.5.0" }, "optionalDependencies": { "@zvec/bindings-darwin-arm64": "0.5.0", "@zvec/bindings-linux-arm64": "0.5.0", "@zvec/bindings-linux-x64": "0.5.0", "@zvec/bindings-win32-x64": "0.5.0" } }, "sha512-00moPv7j7YByeoEdLbitGfd7GO9+dZrPex91S6htmB4ZYbmvdyxz3VWXYuIPnjiE7YPZoY6rdAEFk7MEfMnzxQ=="], + "adm-zip": ["adm-zip@0.5.18", "", {}, "sha512-ufJnssQGbxzLNS1Ho9bCtX4rQKCCvoVuDLHoJyc3F9dOGDB4BkWs2Ci0kv53lqocAEQ/Cbi+I2XCsNYGqVYqng=="], "ansi-escapes": ["ansi-escapes@7.3.0", "", { "dependencies": { "environment": "^1.0.0" } }, "sha512-BvU8nYgGQBxcmMuEeUEmNTvrMVjJNSH7RgW24vXexN4Ven6qCvy4TntnvlnwnMLTVlcRQQdbRY8NKnaIoeWDNg=="], @@ -962,6 +979,8 @@ "before-after-hook": ["before-after-hook@4.0.0", "", {}, "sha512-q6tR3RPqIB1pMiTRMFcZwuG5T8vwp+vUvEG0vuI6B+Rikh5BfPp2fQ82c925FOs+b0lcFQ8CFrL+KbilfZFhOQ=="], + "bindings": ["bindings@1.5.0", "", { "dependencies": { "file-uri-to-path": "1.0.0" } }, "sha512-p2q/t/mhvuOj/UeLlV6566GD/guowlr0hHxClI0W9m7MWYkL1F0hLo+0Aexs9HSPCtR1SXQ0TD3MMKrXZajbiQ=="], + "bluebird": ["bluebird@3.4.7", "", {}, "sha512-iD3898SR7sWVRHbiQv+sHUtHnMvC1o3nW5rAcqnq3uOn07DSAppZYUkIGslDz6gXC7HfunPe7YVBgoEJASPcHA=="], "boolbase": ["boolbase@2.0.0", "", {}, "sha512-DkVaaQHymRhpYEYo9x1oo7Q7B0Y6KJUsjm3c9eTyFDby4MHLBTwZ6ZDWBel5zrYxj1WsZgC5oLpiz+93MluXeA=="], @@ -1102,6 +1121,8 @@ "file-stream-rotator": ["file-stream-rotator@0.6.1", "", { "dependencies": { "moment": "^2.29.1" } }, "sha512-u+dBid4PvZw17PmDeRcNOtCP9CCK/9lRN2w+r1xIS7yOL9JFrIBKTvrYsxT4P0pGtThYTn++QS5ChHaUov3+zQ=="], + "file-uri-to-path": ["file-uri-to-path@1.0.0", "", {}, "sha512-0Zt+s3L7Vf1biwWZ29aARiVYLx7iMGnEUl9x33fbB/j3jR81u/O2LbqK+Bm1CDSNDKVtJ/YjwY7TUd5SkeLQLw=="], + "flatbuffers": ["flatbuffers@25.9.23", "", {}, "sha512-MI1qs7Lo4Syw0EOzUl0xjs2lsoeqFku44KpngfIduHBYvzm8h2+7K8YMQh1JtVVVrUvhLpNwqVi4DERegUJhPQ=="], "fn.name": ["fn.name@1.1.0", "", {}, "sha512-GRnmB5gPyJpAhTQdSZTSp9uaPSvl09KoYcMQtsB9rQoOmzs9dH6ffeccH+Z+cv6P68Hu5bC6JjRh4Ah/mHSNRw=="], diff --git a/docs/tools/code_search.md b/docs/tools/code_search.md new file mode 100644 index 00000000000..f32e89eea2b --- /dev/null +++ b/docs/tools/code_search.md @@ -0,0 +1,25 @@ +# code_search + +Semantic code search using vector embeddings. + +Finds code snippets by meaning, not just text matching. Use when you need to find code by concept (e.g., "authentication middleware", "database connection pool") rather than exact strings. + +## Prerequisites + +- Enable **Code Search (Zvec)** in `/settings` under Tools → Available Tools. +- Install `@zvec/zvec`: `bun add @zvec/zvec` +- Run `/code-index` to build the index before searching. The tool does **not** auto-index on first use — it only searches the existing index. + +## Parameters + +- `query` (string, required): Natural language search query. +- `pattern` (string, optional): File glob pattern to filter results (e.g. `*.ts`). +- `topK` (number, optional): Max results to return (default 20). + +## Behavior + +Returns matching code chunks with file paths, line numbers, and relevance scores. Supports hybrid (vector + FTS) search when embeddings are available, falling back to FTS-only mode otherwise. + +If the index is empty, the tool returns a message telling you to run `/code-index`. + +Prefer `grep` for exact string/regex matching. Use `code_search` when you need semantic understanding of what the code does. diff --git a/package.json b/package.json index e9c758e51ba..f64e54848ff 100644 --- a/package.json +++ b/package.json @@ -54,6 +54,7 @@ "@types/turndown": "5.0.6", "@typescript/native-preview": "7.0.0-dev.20260609.1", "@xterm/headless": "^6.0.0", + "@zvec/zvec": "0.5.0", "arktype": "2.2.2", "chalk": "^5.6.2", "chart.js": "^4.5.1", @@ -96,6 +97,9 @@ "overrides": { "@ark/schema": "0.56.1" }, + "trustedDependencies": [ + "@zvec/zvec" + ], "scripts": { "setup": "bun install && bun run build:native && bun --cwd=packages/coding-agent link && sh scripts/link-omp.sh", "dev": "bun --cwd=packages/coding-agent src/cli.ts", diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 87cd5f7e25d..0b097dca87a 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -12,6 +12,9 @@ - Added model-keyed fallback management to the /models Roles view: model and `provider/*` chains render as a separate section below the roles (divider + "+ New fallback…" row for creating one by picking the protected model, then keying it by model or provider), with the same replace/remove/reorder editing as role chains; the model strip gains `fallbacks:` and `fallbacks:/*` chips as shortcuts. - Added `/queue ` plus `->` / `=>` composer shorthand for follow-up messages that wait until the agent yields. The shorthand opens a dim `Queueing` header and splits sequential numeric, Roman-numeral, or alphabetic lists into separately highlighted queue entries. - Added per-model TPS/TTFT tracking: every completed assistant turn folds its timing into recency-weighted aggregates in `~/.omp/agent.db`, and the /models browser shows measured speed — a right-aligned `118t/s` column on wide terminals (plus TTFT, e.g. `0.9s 118t/s`, when wider) and `~118t/s · 0.9s ttft` facts in the selection detail line — with no dependency on the `omp stats` session scan. +- Added `code_search` tool for semantic code search using Zvec in-process vector database with HNSW indexing and BM25 full-text search; indexes the workspace on first use and supports hybrid (vector + FTS) search when embeddings are available, falling back to FTS-only mode otherwise. +- Added `tools.codeSearchEnabled` setting (default: off) to enable/disable semantic code search in the /settings panel. +- Added `/code-index` slash command to explicitly build or rebuild the Zvec code search index; the `code_search` tool no longer auto-indexes on first use. ### Changed diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index e07b1404509..f929f903d04 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -3644,6 +3644,18 @@ export const SETTINGS_SCHEMA = { }, }, + "tools.codeSearchEnabled": { + type: "boolean", + default: false, + ui: { + tab: "tools", + group: "Available Tools", + label: "Code Search (Zvec)", + description: + "Enable semantic code search using Zvec in-process vector database. Use /code-index to build the index before searching.", + }, + }, + "browser.enabled": { type: "boolean", default: true, diff --git a/packages/coding-agent/src/prompts/tools/code-search.md b/packages/coding-agent/src/prompts/tools/code-search.md new file mode 100644 index 00000000000..1da7c970d4b --- /dev/null +++ b/packages/coding-agent/src/prompts/tools/code-search.md @@ -0,0 +1,7 @@ +Semantic code search using vector embeddings. + +Finds code snippets by meaning, not just text matching. Use when you need to find code by concept (e.g., "authentication middleware", "database connection pool") rather than exact strings. + +Returns matching code chunks with file paths, line numbers, and relevance scores. The index must be built before searching — if the index is empty, the tool returns a message telling you to run `/code-index` first. Use `/code-index` to build or rebuild the index. + +Prefer `grep` for exact string/regex matching. Use `code_search` when you need semantic understanding of what the code does. diff --git a/packages/coding-agent/src/slash-commands/builtin-registry.ts b/packages/coding-agent/src/slash-commands/builtin-registry.ts index e5a5711b96c..f079d7c86f5 100644 --- a/packages/coding-agent/src/slash-commands/builtin-registry.ts +++ b/packages/coding-agent/src/slash-commands/builtin-registry.ts @@ -24,6 +24,7 @@ import { MarketplaceManager, } from "../extensibility/plugins/marketplace"; import { resolveMemoryBackend } from "../memory-backend"; +import { loadMnemopi } from "../mnemopi/state"; import { runPauseScreen } from "../modes/components/pause-screen"; import { describeLoopLimitRuntime } from "../modes/loop-limit"; import { theme } from "../modes/theme/theme"; @@ -33,6 +34,7 @@ import type { AgentSession, FreshSessionResult } from "../session/agent-session" import { COMPACT_MODES, parseCompactArgs } from "../session/compact-modes"; import { resolveResumableSession } from "../session/session-listing"; import { formatShakeSummary, type ShakeMode } from "../session/shake-types"; +import { indexWorkspace } from "../tools/code-search-index"; import { expandTilde, resolveToCwd } from "../tools/path-utils"; import { urlHyperlinkAlways } from "../tui"; import { @@ -1159,6 +1161,64 @@ const BUILTIN_SLASH_COMMAND_REGISTRY: ReadonlyArray = [ runtime.ctx.editor.setText(""); }, }, + { + name: "code-index", + description: "Index the workspace for semantic code search (Zvec)", + inlineHint: "[rebuild]", + allowArgs: true, + getTuiAutocompleteDescription: runtime => { + if (!runtime.ctx.settings.get("tools.codeSearchEnabled")) { + return "Code Search: disabled in settings"; + } + return "Code Search: index workspace"; + }, + handleTui: async (_command, runtime) => { + const enabled = runtime.ctx.settings.get("tools.codeSearchEnabled"); + if (!enabled) { + runtime.ctx.showWarning("Code search is disabled. Enable it in /settings under Tools → Available Tools."); + runtime.ctx.editor.setText(""); + return; + } + const { ZVecCodeStore } = await loadMnemopi(); + if (!ZVecCodeStore.available()) { + runtime.ctx.showWarning("Code search requires @zvec/zvec. Install it with: bun add @zvec/zvec"); + runtime.ctx.editor.setText(""); + return; + } + runtime.ctx.showStatus("Indexing workspace for code search..."); + runtime.ctx.editor.setText(""); + try { + const result = await indexWorkspace(runtime.ctx.session.sessionManager.getCwd()); + runtime.ctx.showStatus( + `Code search indexed: ${result.filesIndexed} file${result.filesIndexed === 1 ? "" : "s"}, ${result.chunksIndexed} chunk${result.chunksIndexed === 1 ? "" : "s"} in ${result.duration}ms`, + ); + } catch (err) { + runtime.ctx.showError(`Code search indexing failed: ${String(err)}`); + } + }, + handle: async (_command, runtime) => { + const enabled = runtime.settings.get("tools.codeSearchEnabled"); + if (!enabled) { + await runtime.output("Code search is disabled. Enable it in /settings under Tools → Available Tools."); + return commandConsumed(); + } + const { ZVecCodeStore } = await loadMnemopi(); + if (!ZVecCodeStore.available()) { + await runtime.output("Code search requires @zvec/zvec. Install it with: bun add @zvec/zvec"); + return commandConsumed(); + } + await runtime.output("Indexing workspace for code search..."); + try { + const result = await indexWorkspace(runtime.cwd); + await runtime.output( + `Code search indexed: ${result.filesIndexed} file${result.filesIndexed === 1 ? "" : "s"}, ${result.chunksIndexed} chunk${result.chunksIndexed === 1 ? "" : "s"} in ${result.duration}ms`, + ); + } catch (err) { + await runtime.output(`Code search indexing failed: ${String(err)}`); + } + return commandConsumed(); + }, + }, { name: "context", description: "Show estimated context usage breakdown", diff --git a/packages/coding-agent/src/tools/__tests__/code-search-index.test.ts b/packages/coding-agent/src/tools/__tests__/code-search-index.test.ts new file mode 100644 index 00000000000..b0af123be8f --- /dev/null +++ b/packages/coding-agent/src/tools/__tests__/code-search-index.test.ts @@ -0,0 +1,101 @@ +import { afterEach, beforeEach, describe, expect, setDefaultTimeout, test } from "bun:test"; +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import * as path from "node:path"; +import { type CodeChunk, ZVecCodeStore } from "@oh-my-pi/pi-mnemopi"; + +// Zvec native FFI calls (upsert/optimize/searchFts) are slow on CI hardware. +setDefaultTimeout(15_000); + +describe.skipIf(!ZVecCodeStore.available())("ZVecCodeStore edit-then-reindex", () => { + let tempDir: string; + let store: ZVecCodeStore; + + beforeEach(async () => { + tempDir = await mkdtemp(path.join(tmpdir(), "zvec-code-search-")); + // ZVecCreateAndOpen requires a path that does not yet exist. + store = ZVecCodeStore.create(path.join(tempDir, "index"), 384); + }); + + afterEach(async () => { + store.close(); + await rm(tempDir, { recursive: true, force: true }); + }); + + function makeChunk(filePath: string, content: string, fileHash: string): CodeChunk { + return { + id: `${fileHash}-1`, + filePath, + language: "typescript", + lineStart: 1, + lineEnd: 1, + content, + chunkHash: Bun.hash(content).toString(16), + embedding: [], + }; + } + + test("upserting new chunks after removeFile does not leave stale chunks", () => { + const filePath = "src/test.ts"; + const oldContent = "function foo() { return 1; }"; + const newContent = "function bar() { return 2; }"; + + // Step 1: Upsert initial chunk. + const oldHash = Bun.hash(`${filePath}\0${oldContent}`).toString(16); + store.upsertChunks([makeChunk(filePath, oldContent, oldHash)]); + store.optimize(); + expect(store.docCount).toBe(1); + + // Step 2: Verify search finds the old content. + const fooResults = store.searchFts("foo", 10); + expect(fooResults.length).toBe(1); + expect(fooResults[0]!.content).toContain("foo"); + + // Step 3: Simulate an edit — remove old chunks, then upsert new ones. + store.removeFile(filePath); + const newHash = Bun.hash(`${filePath}\0${newContent}`).toString(16); + store.upsertChunks([makeChunk(filePath, newContent, newHash)]); + store.optimize(); + + // Step 4: docCount should be 1 (not 2 — no stale chunks). + expect(store.docCount).toBe(1); + + // Step 5: Search for "foo" should return nothing (old chunk deleted). + const staleResults = store.searchFts("foo", 10); + expect(staleResults.length).toBe(0); + + // Step 6: Search for "bar" should return the new chunk. + const barResults = store.searchFts("bar", 10); + expect(barResults.length).toBe(1); + expect(barResults[0]!.content).toContain("bar"); + expect(barResults[0]!.filePath).toBe(filePath); + }); + + test("getFileHashes reflects current chunks after re-indexing", () => { + const filePath = "src/utils.ts"; + const oldContent = "export function add(a, b) { return a + b; }"; + const newContent = "export function multiply(a, b) { return a * b; }"; + + // Index initial content. + const oldHash = Bun.hash(`${filePath}\0${oldContent}`).toString(16); + store.upsertChunks([makeChunk(filePath, oldContent, oldHash)]); + store.optimize(); + + let hashes = store.getFileHashes(); + expect(hashes.has(filePath)).toBe(true); + expect(hashes.get(filePath)!.size).toBe(1); + + // Re-index with changed content. + store.removeFile(filePath); + const newHash = Bun.hash(`${filePath}\0${newContent}`).toString(16); + store.upsertChunks([makeChunk(filePath, newContent, newHash)]); + store.optimize(); + + // getFileHashes should reflect only the new chunk hash. + hashes = store.getFileHashes(); + expect(hashes.get(filePath)!.size).toBe(1); + const currentHashes = hashes.get(filePath)!; + const newChunkHash = Bun.hash(newContent).toString(16); + expect(currentHashes.has(newChunkHash)).toBe(true); + }); +}); diff --git a/packages/coding-agent/src/tools/builtin-names.ts b/packages/coding-agent/src/tools/builtin-names.ts index 253bd33453e..9199bc19ee3 100644 --- a/packages/coding-agent/src/tools/builtin-names.ts +++ b/packages/coding-agent/src/tools/builtin-names.ts @@ -23,6 +23,7 @@ export const BUILTIN_TOOL_NAMES = [ "web_search", "search_tool_bm25", "write", + "code_search", "memory_edit", "retain", "recall", diff --git a/packages/coding-agent/src/tools/code-search-helpers.ts b/packages/coding-agent/src/tools/code-search-helpers.ts new file mode 100644 index 00000000000..277a01ba440 --- /dev/null +++ b/packages/coding-agent/src/tools/code-search-helpers.ts @@ -0,0 +1,125 @@ +import * as path from "node:path"; + +// ─── Language detection ───────────────────────────────────────────────────── + +const EXTENSION_TO_LANGUAGE: ReadonlyMap = new Map([ + ["ts", "typescript"], + ["tsx", "typescript"], + ["js", "javascript"], + ["jsx", "javascript"], + ["mjs", "javascript"], + ["cjs", "javascript"], + ["py", "python"], + ["rs", "rust"], + ["go", "go"], + ["java", "java"], + ["c", "c"], + ["cpp", "cpp"], + ["cc", "cpp"], + ["h", "c"], + ["hpp", "cpp"], + ["hxx", "cpp"], + ["rb", "ruby"], + ["lua", "lua"], + ["sql", "sql"], + ["sh", "shell"], + ["bash", "shell"], + ["zsh", "shell"], + ["fish", "shell"], + ["ps1", "powershell"], + ["yaml", "yaml"], + ["yml", "yaml"], + ["json", "json"], + ["toml", "toml"], + ["xml", "xml"], + ["html", "html"], + ["css", "css"], + ["scss", "scss"], + ["less", "less"], + ["vue", "vue"], + ["svelte", "svelte"], + ["md", "markdown"], + ["mdx", "markdown"], + ["php", "php"], + ["swift", "swift"], + ["kt", "kotlin"], + ["scala", "scala"], + ["clj", "clojure"], + ["ex", "elixir"], + ["exs", "elixir"], + ["erl", "erlang"], + ["hs", "haskell"], + ["ml", "ocaml"], + ["nim", "nim"], + ["zig", "zig"], + ["v", "verilog"], + ["sv", "systemverilog"], + ["dart", "dart"], + ["gradle", "gradle"], + ["dockerfile", "dockerfile"], + ["makefile", "makefile"], + ["r", "r"], + ["jl", "julia"], +]); + +const INDEXABLE_EXTENSIONS: ReadonlySet = new Set(EXTENSION_TO_LANGUAGE.keys()); + +const SKIP_DIRECTORIES = new Set(["node_modules", ".git", "dist", "build", "target", ".next", "out", "coverage"]); + +const MAX_FILE_SIZE = 1024 * 1024; // 1 MB +const BINARY_CHECK_BYTES = 8192; + +export { BINARY_CHECK_BYTES, EXTENSION_TO_LANGUAGE, INDEXABLE_EXTENSIONS, MAX_FILE_SIZE, SKIP_DIRECTORIES }; + +export function detectLanguage(filePath: string): string { + const ext = path.extname(filePath).slice(1).toLowerCase(); + if (ext === "") { + const basename = path.basename(filePath).toLowerCase(); + if (basename === "dockerfile") return "dockerfile"; + if (basename === "makefile") return "makefile"; + } + return EXTENSION_TO_LANGUAGE.get(ext) ?? "text"; +} + +export function isIndexableFile(filePath: string): boolean { + const ext = path.extname(filePath).slice(1).toLowerCase(); + if (ext !== "") return INDEXABLE_EXTENSIONS.has(ext); + const basename = path.basename(filePath).toLowerCase(); + return basename === "dockerfile" || basename === "makefile"; +} + +// ─── Code chunking ────────────────────────────────────────────────────────── + +export interface Chunk { + lineStart: number; + lineEnd: number; + content: string; +} + +export function chunkFile(content: string, chunkSize: number, overlap: number): Chunk[] { + const lines = content.split("\n"); + if (lines.length === 0) return []; + const chunks: Chunk[] = []; + const step = Math.max(1, chunkSize - overlap); + for (let start = 0; start < lines.length; start += step) { + const end = Math.min(start + chunkSize - 1, lines.length - 1); + const chunkLines = lines.slice(start, end + 1); + chunks.push({ + lineStart: start + 1, + lineEnd: end + 1, + content: chunkLines.join("\n"), + }); + if (end === lines.length - 1) break; + } + return chunks; +} + +// ─── Binary detection ─────────────────────────────────────────────────────── + +export function isBinaryContent(buffer: Buffer): boolean { + const checkLen = Math.min(buffer.length, BINARY_CHECK_BYTES); + for (let i = 0; i < checkLen; i++) { + if (buffer[i] === 0) return true; + } + return false; +} diff --git a/packages/coding-agent/src/tools/code-search-index.ts b/packages/coding-agent/src/tools/code-search-index.ts new file mode 100644 index 00000000000..7b2e494ba89 --- /dev/null +++ b/packages/coding-agent/src/tools/code-search-index.ts @@ -0,0 +1,181 @@ +import * as path from "node:path"; +import type * as MnemopiNs from "@oh-my-pi/pi-mnemopi"; +import type { CodeChunk, ZVecCodeStore } from "@oh-my-pi/pi-mnemopi"; +import { FileType, listWorkspace } from "@oh-my-pi/pi-natives"; +import { logger } from "@oh-my-pi/pi-utils"; +import { loadMnemopi } from "../mnemopi/state"; +import { + chunkFile, + detectLanguage, + isBinaryContent, + isIndexableFile, + MAX_FILE_SIZE, + SKIP_DIRECTORIES, +} from "./code-search-helpers"; +import { throwIfAborted } from "./tool-errors"; + +export interface IndexResult { + totalFiles: number; + filesIndexed: number; + chunksIndexed: number; + duration: number; +} + +export async function indexWorkspace(cwd: string, signal?: AbortSignal): Promise { + const start = Date.now(); + const mnemopi: typeof MnemopiNs = await loadMnemopi(); + + if (!mnemopi.ZVecCodeStore.available()) { + throw new Error("Code search requires @zvec/zvec. Install it with: bun add @zvec/zvec"); + } + + const indexPath = mnemopi.zvecCodeIndexPath(); + const model = mnemopi.currentEmbeddingModel(); + const dimension = mnemopi.embeddingDimFor(model); + + let store: ZVecCodeStore; + try { + store = mnemopi.ZVecCodeStore.open(indexPath); + } catch { + store = mnemopi.ZVecCodeStore.create(indexPath, dimension); + } + + const sourceFiles = await collectSourceFiles(cwd, signal); + if (sourceFiles.length === 0) { + return { totalFiles: 0, filesIndexed: 0, chunksIndexed: 0, duration: Date.now() - start }; + } + + const storeFileHashes = store.getFileHashes(); + const indexedPaths = new Set(storeFileHashes.keys()); + + const chunksToUpsert: CodeChunk[] = []; + const filesToRemove: string[] = []; + const currentFiles = new Set(); + const chunkSize = mnemopi.zvecChunkSize(); + const overlap = mnemopi.zvecChunkOverlap(); + const canEmbed = await mnemopi.available(); + + let filesIndexed = 0; + + for (const filePath of sourceFiles) { + throwIfAborted(signal); + + const file = Bun.file(filePath); + const exists = await file.exists(); + if (!exists) continue; + + const stat = await file.stat(); + if (stat.size > MAX_FILE_SIZE) { + // File grew too large — remove any existing chunks. + if (storeFileHashes.has(filePath)) store.removeFile(filePath); + continue; + } + + const buffer = Buffer.from(await file.arrayBuffer()); + if (isBinaryContent(buffer)) { + // File became binary — remove any existing chunks. + if (storeFileHashes.has(filePath)) store.removeFile(filePath); + continue; + } + + // Only mark as current after confirming it's indexable. + currentFiles.add(filePath); + + const content = buffer.toString("utf-8"); + const fileHash = Bun.hash(`${filePath}\0${content}`).toString(16); + + const chunks = chunkFile(content, chunkSize, overlap); + const newChunkHashes = new Set(chunks.map(c => Bun.hash(c.content).toString(16))); + + const indexedHashes = storeFileHashes.get(filePath); + if (indexedHashes !== undefined) { + if (setsEqual(indexedHashes, newChunkHashes)) continue; + store.removeFile(filePath); + } + + const language = detectLanguage(filePath); + const batchSize = 32; + for (let i = 0; i < chunks.length; i += batchSize) { + throwIfAborted(signal); + const batch = chunks.slice(i, i + batchSize); + let embeddings: Float32Array[] | null = null; + if (canEmbed) { + try { + embeddings = await mnemopi.embed(batch.map(c => c.content)); + } catch (err) { + logger.debug("code_search: embedding batch failed", { error: String(err) }); + } + } + + for (let j = 0; j < batch.length; j++) { + const chunk = batch[j]!; + chunksToUpsert.push({ + id: `${fileHash}-${chunk.lineStart}`, + filePath, + language, + lineStart: chunk.lineStart, + lineEnd: chunk.lineEnd, + content: chunk.content, + chunkHash: Bun.hash(chunk.content).toString(16), + embedding: embeddings?.[j] ? Array.from(embeddings[j]!) : [], + }); + } + } + + filesIndexed++; + } + + for (const indexedPath of indexedPaths) { + if (!currentFiles.has(indexedPath)) { + filesToRemove.push(indexedPath); + } + } + + for (const filePath of filesToRemove) { + store.removeFile(filePath); + } + + if (chunksToUpsert.length > 0) { + store.upsertChunks(chunksToUpsert); + } + + if (chunksToUpsert.length > 0 || filesToRemove.length > 0) { + store.optimize(); + } + + return { + totalFiles: sourceFiles.length, + filesIndexed, + chunksIndexed: chunksToUpsert.length, + duration: Date.now() - start, + }; +} + +async function collectSourceFiles(cwd: string, signal?: AbortSignal): Promise { + const result = await listWorkspace({ + path: cwd, + maxDepth: 20, + gitignore: true, + hidden: false, + signal, + }); + + const files: string[] = []; + for (const entry of result.entries) { + if (entry.fileType !== FileType.File) continue; + if (!isIndexableFile(entry.path)) continue; + if (entry.size !== undefined && entry.size > MAX_FILE_SIZE) continue; + + const parts = entry.path.split("/"); + if (parts.some(p => SKIP_DIRECTORIES.has(p))) continue; + + files.push(path.join(cwd, entry.path)); + } + return files; +} + +function setsEqual(a: Set, b: Set): boolean { + if (a.size !== b.size) return false; + for (const item of a) if (!b.has(item)) return false; + return true; +} diff --git a/packages/coding-agent/src/tools/code-search.ts b/packages/coding-agent/src/tools/code-search.ts new file mode 100644 index 00000000000..f87f8028393 --- /dev/null +++ b/packages/coding-agent/src/tools/code-search.ts @@ -0,0 +1,214 @@ +import * as path from "node:path"; +import type { AgentTool, AgentToolResult } from "@oh-my-pi/pi-agent-core"; +import { + currentEmbeddingModel, + embeddingDimFor, + available as embeddingsAvailable, + embedQuery, + ZVecCodeStore, + type ZVecSearchResult, + zvecCodeIndexPath, + zvecTopK, +} from "@oh-my-pi/pi-mnemopi"; +import { logger, untilAborted } from "@oh-my-pi/pi-utils"; +import { type } from "arktype"; +import codeSearchDescription from "../prompts/tools/code-search.md" with { type: "text" }; +import type { ToolSession } from "."; +import { replaceTabs, shortenPath, TRUNCATE_LENGTHS, truncateToWidth } from "./render-utils"; +import { ToolError } from "./tool-errors"; + +const codeSearchSchema = type({ + query: type("string").describe("natural language search query"), + "pattern?": type("string").describe("optional file glob pattern to filter results (e.g. '*.ts')"), + "topK?": type("number").describe("max results to return (default 20)"), +}); + +export type CodeSearchParams = typeof codeSearchSchema.infer; + +// ─── Result formatting ─────────────────────────────────────────────────────── + +function formatSearchResult( + results: ZVecSearchResult[], + indexedChunkCount: number, + mode: "hybrid" | "fts", + indexEmpty: boolean, +): string { + if (results.length === 0) { + if (indexEmpty) { + return `No matches found. The code index is empty — run /code-index to index your workspace.`; + } + return `No matches found (${indexedChunkCount} indexed chunk${indexedChunkCount === 1 ? "" : "s"}, Zvec index, ${mode} mode).`; + } + + const lines: string[] = [ + `Found ${results.length} match${results.length === 1 ? "" : "es"} (${indexedChunkCount} indexed chunk${indexedChunkCount === 1 ? "" : "s"}, Zvec index, ${mode} mode)`, + "", + ]; + + for (const result of results) { + const displayPath = shortenPath(result.filePath); + const scoreStr = result.score.toFixed(2); + lines.push(`${displayPath}:${result.lineStart}-${result.lineEnd} [distance: ${scoreStr}]`); + + const contentLines = result.content.split("\n"); + for (const line of contentLines) { + const sanitized = replaceTabs(line); + const truncated = truncateToWidth(sanitized, TRUNCATE_LENGTHS.LINE); + lines.push(` ${truncated}`); + } + lines.push(""); + } + + return lines.join("\n").trimEnd(); +} + +// ─── Tool ──────────────────────────────────────────────────────────────────── + +export class CodeSearchTool implements AgentTool { + readonly name = "code_search"; + readonly approval = "read" as const; + readonly label = "Code Search"; + readonly description = codeSearchDescription; + readonly parameters = codeSearchSchema; + readonly strict = true; + readonly loadMode = "discoverable"; + readonly summary = "Semantic code search using vector embeddings"; + + #store: ZVecCodeStore | null = null; + #storePromise: Promise | null = null; + + static createIf(session: ToolSession): CodeSearchTool | null { + if (!session.settings.get("tools.codeSearchEnabled")) return null; + if (!ZVecCodeStore.available()) return null; + return new CodeSearchTool(); + } + + // ─── Store lifecycle ───────────────────────────────────────────────────── + + async #getStore(): Promise { + if (this.#store) return this.#store; + if (this.#storePromise) return this.#storePromise; + + this.#storePromise = this.#createStore(); + try { + this.#store = await this.#storePromise; + return this.#store; + } finally { + this.#storePromise = null; + } + } + + async #createStore(): Promise { + const indexPath = zvecCodeIndexPath(); + const model = currentEmbeddingModel(); + const dimension = embeddingDimFor(model); + + // Try to open existing index, fall back to create. + try { + return ZVecCodeStore.open(indexPath); + } catch { + return ZVecCodeStore.create(indexPath, dimension); + } + } + + // ─── Search ───────────────────────────────────────────────────────────── + + async #search( + query: string, + topK: number, + _signal?: AbortSignal, + ): Promise<{ results: ZVecSearchResult[]; mode: "hybrid" | "fts"; indexEmpty: boolean }> { + const store = await this.#getStore(); + if (store.docCount === 0) { + return { results: [], mode: "fts", indexEmpty: true }; + } + + const canEmbed = await embeddingsAvailable(); + if (canEmbed) { + try { + const queryEmbedding = await embedQuery(query); + if (queryEmbedding) { + const vector = Array.from(queryEmbedding); + const results = store.searchHybrid({ vector, fts: query }, topK); + return { results, mode: "hybrid", indexEmpty: false }; + } + } catch (err) { + logger.debug("code_search: query embedding failed, falling back to FTS", { error: String(err) }); + } + } else { + logger.debug("code_search: embeddings unavailable, using FTS-only mode"); + } + + const results = store.searchFts(query, topK); + return { results, mode: "fts", indexEmpty: false }; + } + + // ─── Execute ──────────────────────────────────────────────────────────── + + async execute(_id: string, params: CodeSearchParams, signal?: AbortSignal): Promise { + return untilAborted(signal, async () => { + const query = params.query.trim(); + if (query.length === 0) { + throw new ToolError("Query is required and must not be empty."); + } + + const topK = params.topK ?? zvecTopK(); + if (!Number.isInteger(topK) || topK <= 0) { + throw new ToolError("topK must be a positive integer."); + } + + // When a pattern filter is provided, overfetch so the filter is reliable. + const searchTopK = params.pattern ? Math.max(topK * 5, 100) : topK; + const { results, mode, indexEmpty } = await this.#search(query, searchTopK, signal); + + // Apply optional glob pattern filter. + let filtered = results; + if (params.pattern) { + const globPattern = params.pattern; + filtered = results.filter(r => matchesGlob(r.filePath, globPattern)); + } + + // Slice to the requested topK after filtering. + filtered = filtered.slice(0, topK); + + const store = await this.#getStore(); + const text = formatSearchResult(filtered, store.docCount, mode, indexEmpty); + + return { + content: [{ type: "text", text }], + details: {}, + }; + }); + } +} + +// ─── Helpers ───────────────────────────────────────────────────────────────── + +function matchesGlob(filePath: string, pattern: string): boolean { + // Simple glob matching: convert pattern to regex. + // Supports * (any chars except /), ** (any chars including /), and ? (single char). + const regex = globToRegex(pattern); + return regex.test(path.basename(filePath)) || regex.test(filePath); +} + +function globToRegex(pattern: string): RegExp { + let regexStr = ""; + for (let i = 0; i < pattern.length; i++) { + const char = pattern[i]!; + if (char === "*") { + if (pattern[i + 1] === "*") { + regexStr += ".*"; + i++; // Skip the next * + } else { + regexStr += "[^/]*"; + } + } else if (char === "?") { + regexStr += "[^/]"; + } else if (".+^$(){}|[]\\".includes(char)) { + regexStr += `\\${char}`; + } else { + regexStr += char; + } + } + return new RegExp(`^${regexStr}$`); +} diff --git a/packages/coding-agent/src/tools/index.ts b/packages/coding-agent/src/tools/index.ts index 287c22ff646..bfbf1befaec 100644 --- a/packages/coding-agent/src/tools/index.ts +++ b/packages/coding-agent/src/tools/index.ts @@ -41,6 +41,7 @@ import { BashTool } from "./bash"; import { BrowserTool } from "./browser"; import { type BuiltinToolName, normalizeToolNames } from "./builtin-names"; import { type CheckpointState, CheckpointTool, type CompletedRewindState, RewindTool } from "./checkpoint"; +import { CodeSearchTool } from "./code-search"; import { DebugTool } from "./debug"; import { EvalTool } from "./eval"; import { resolveEvalBackends } from "./eval-backends"; @@ -79,6 +80,7 @@ export * from "./ast-grep"; export * from "./bash"; export * from "./browser"; export * from "./checkpoint"; +export * from "./code-search"; export * from "./debug"; export * from "./eval"; export * from "./eval-backends"; @@ -453,6 +455,7 @@ export const BUILTIN_TOOLS: Record = { ast_edit: s => new AstEditTool(s), ask: AskTool.createIf, debug: DebugTool.createIf, + code_search: CodeSearchTool.createIf, eval: s => new EvalTool(s), ssh: loadSshTool, github: GithubTool.createIf, diff --git a/packages/coding-agent/test/tool-discovery/initial-tools.test.ts b/packages/coding-agent/test/tool-discovery/initial-tools.test.ts index 303023cc4fe..5ef52ece5bf 100644 --- a/packages/coding-agent/test/tool-discovery/initial-tools.test.ts +++ b/packages/coding-agent/test/tool-discovery/initial-tools.test.ts @@ -27,6 +27,7 @@ const allToolsSettings = Settings.isolated({ "browser.enabled": true, "checkpoint.enabled": true, "todo.enabled": true, + "tools.codeSearchEnabled": true, "memory.backend": "mnemopi", "autolearn.enabled": true, "tools.discoveryMode": "all", diff --git a/packages/mnemopi/CHANGELOG.md b/packages/mnemopi/CHANGELOG.md index 21170cfd9c3..21459b06fbd 100644 --- a/packages/mnemopi/CHANGELOG.md +++ b/packages/mnemopi/CHANGELOG.md @@ -2,6 +2,12 @@ ## [Unreleased] +### Added + +- Added `ZVecCodeStore`, a vector store wrapper for Zvec (`@zvec/zvec`) providing HNSW vector search, BM25 full-text search, and hybrid (RRF) search for code chunks. +- Added `zvec-config.ts` with environment-configurable settings for Zvec paths, chunk sizing, and top-K defaults (`OMP_ZVEC_*` env vars). +- Added `@zvec/zvec` as an optional peer dependency. + ## [16.3.9] - 2026-07-06 ### Fixed diff --git a/packages/mnemopi/package.json b/packages/mnemopi/package.json index d0a7ef18bb2..58a00a46cd4 100644 --- a/packages/mnemopi/package.json +++ b/packages/mnemopi/package.json @@ -45,10 +45,14 @@ "lru-cache": "catalog:" }, "peerDependencies": { + "@zvec/zvec": ">=0.1.0", "fastembed": "2.1.0", "onnxruntime-node": "1.21.0" }, "peerDependenciesMeta": { + "@zvec/zvec": { + "optional": true + }, "fastembed": { "optional": true }, @@ -58,6 +62,7 @@ }, "devDependencies": { "@types/bun": "catalog:", + "@zvec/zvec": "catalog:", "fastembed": "catalog:", "onnxruntime-node": "catalog:" }, diff --git a/packages/mnemopi/src/core/index.ts b/packages/mnemopi/src/core/index.ts index 2c15491538e..0b48d34b414 100644 --- a/packages/mnemopi/src/core/index.ts +++ b/packages/mnemopi/src/core/index.ts @@ -36,3 +36,5 @@ export { storeMemory, update, } from "./memory"; +export * from "./zvec-config"; +export * from "./zvec-store"; diff --git a/packages/mnemopi/src/core/zvec-config.ts b/packages/mnemopi/src/core/zvec-config.ts new file mode 100644 index 00000000000..171bec8abea --- /dev/null +++ b/packages/mnemopi/src/core/zvec-config.ts @@ -0,0 +1,26 @@ +import { homedir } from "node:os"; +import { join } from "node:path"; +import { type Env, envBool, envInt, envString } from "../util/env"; + +export const DEFAULT_ZVEC_DIR = join(homedir(), ".hermes", "zvec"); +export const DEFAULT_ZVEC_CODE_INDEX = join(DEFAULT_ZVEC_DIR, "code-index"); + +export function zvecCodeIndexPath(env: Env = process.env): string { + return envString("OMP_ZVEC_CODE_INDEX", DEFAULT_ZVEC_CODE_INDEX, env); +} + +export function zvecChunkSize(env: Env = process.env): number { + return envInt("OMP_ZVEC_CHUNK_SIZE", 100, env); +} + +export function zvecChunkOverlap(env: Env = process.env): number { + return envInt("OMP_ZVEC_CHUNK_OVERLAP", 10, env); +} + +export function zvecTopK(env: Env = process.env): number { + return envInt("OMP_ZVEC_TOP_K", 20, env); +} + +export function zvecEnabled(env: Env = process.env): boolean { + return envBool("OMP_ZVEC_ENABLED", true, env); +} diff --git a/packages/mnemopi/src/core/zvec-store.ts b/packages/mnemopi/src/core/zvec-store.ts new file mode 100644 index 00000000000..1f1b9eb695b --- /dev/null +++ b/packages/mnemopi/src/core/zvec-store.ts @@ -0,0 +1,352 @@ +import { createRequire } from "node:module"; +import { logger } from "@oh-my-pi/pi-utils"; +import type { ZVecCollection, ZVecDocInput, ZVecMultiQuery, ZVecQuery } from "@zvec/zvec"; + +/** + * Module-level cache of the `@zvec/zvec` native addon. + * + * The module is an optional peer dependency — it carries a native addon that + * may be absent at runtime. We load it once via `createRequire` (synchronous in + * Bun) and cache the result so that all subsequent operations are sync. + */ +let zvecModule: typeof import("@zvec/zvec") | null | undefined; + +/** + * Synchronously load (and cache) the `@zvec/zvec` native module. + * Returns `null` when the package is not installed. + */ +function getZvec(): typeof import("@zvec/zvec") | null { + if (zvecModule !== undefined) return zvecModule; + try { + const require = createRequire(import.meta.url); + zvecModule = require("@zvec/zvec") as typeof import("@zvec/zvec"); + } catch (error) { + logger.debug("zvec: @zvec/zvec native module not available", { error: String(error) }); + zvecModule = null; + } + return zvecModule; +} + +/** A code chunk with its pre-computed embedding vector. */ +export interface CodeChunk { + /** Unique chunk ID (e.g. `${fileHash}-${lineStart}`). Must not contain `:`, `/`, or `|` — Zvec rejects those characters. */ + id: string; + /** Relative file path. */ + filePath: string; + /** Programming language. */ + language: string; + /** Start line (1-based, inclusive). */ + lineStart: number; + /** End line (1-based, inclusive). */ + lineEnd: number; + /** Code text. */ + content: string; + /** Hash of the content. */ + chunkHash: string; + /** Pre-computed embedding vector. */ + embedding: number[]; +} + +/** A search result from the vector store. */ +export interface ZVecSearchResult { + id: string; + filePath: string; + language: string; + lineStart: number; + lineEnd: number; + content: string; + chunkHash: string; + score: number; +} + +/** Names of scalar fields stored in the Zvec collection. */ +const FIELD_FILE_PATH = "file_path"; +const FIELD_LANGUAGE = "language"; +const FIELD_LINE_START = "line_start"; +const FIELD_LINE_END = "line_end"; +const FIELD_CONTENT = "content"; +const FIELD_CHUNK_HASH = "chunk_hash"; +const VECTOR_FIELD = "embedding"; + +/** All scalar field names, for output projection. */ +const ALL_FIELDS = [ + FIELD_FILE_PATH, + FIELD_LANGUAGE, + FIELD_LINE_START, + FIELD_LINE_END, + FIELD_CONTENT, + FIELD_CHUNK_HASH, +] as const; + +/** Maximum topk Zvec allows in a single query (enforced by the native engine). */ +const ZVEC_MAX_TOPK = 100_000; + +/** Convert a ZVecDoc (with `fields: Record`) into a typed result. */ +function docToResult(doc: { id: string; fields: Record; score: number }): ZVecSearchResult { + const f = doc.fields; + return { + id: doc.id, + filePath: String(f[FIELD_FILE_PATH] ?? ""), + language: String(f[FIELD_LANGUAGE] ?? ""), + lineStart: Number(f[FIELD_LINE_START] ?? 0), + lineEnd: Number(f[FIELD_LINE_END] ?? 0), + content: String(f[FIELD_CONTENT] ?? ""), + chunkHash: String(f[FIELD_CHUNK_HASH] ?? ""), + score: doc.score, + }; +} + +/** + * A vector store for code chunks backed by Zvec (`@zvec/zvec`). + * + * This is a pure vector-DB wrapper — it stores and searches pre-computed + * embeddings. Embedding computation is the caller's responsibility. + */ +export class ZVecCodeStore { + #collection: ZVecCollection; + #dimension: number; + + constructor(collection: ZVecCollection, dimension: number) { + this.#collection = collection; + this.#dimension = dimension; + } + + /** + * Create and open a new Zvec collection at the given path. + * @param path Directory for the collection files. + * @param dimension Embedding vector dimension (must match the embedding model). + */ + static create(path: string, dimension: number): ZVecCodeStore { + const zvec = getZvec(); + if (!zvec) { + throw new Error("@zvec/zvec is not installed. Install it with: bun add @zvec/zvec"); + } + + const schema = new zvec.ZVecCollectionSchema({ + name: "code_index", + fields: [ + { + name: FIELD_FILE_PATH, + dataType: zvec.ZVecDataType.STRING, + indexParams: { indexType: zvec.ZVecIndexType.INVERT }, + }, + { + name: FIELD_LANGUAGE, + dataType: zvec.ZVecDataType.STRING, + indexParams: { indexType: zvec.ZVecIndexType.INVERT }, + }, + { name: FIELD_LINE_START, dataType: zvec.ZVecDataType.INT32 }, + { name: FIELD_LINE_END, dataType: zvec.ZVecDataType.INT32 }, + { + name: FIELD_CHUNK_HASH, + dataType: zvec.ZVecDataType.STRING, + indexParams: { indexType: zvec.ZVecIndexType.INVERT }, + }, + { + name: FIELD_CONTENT, + dataType: zvec.ZVecDataType.STRING, + nullable: false, + indexParams: { + indexType: zvec.ZVecIndexType.FTS, + tokenizerName: "standard", + filters: ["lowercase"], + }, + }, + ], + vectors: [ + { + name: VECTOR_FIELD, + dataType: zvec.ZVecDataType.VECTOR_FP32, + dimension, + indexParams: { indexType: zvec.ZVecIndexType.HNSW, metricType: zvec.ZVecMetricType.COSINE }, + }, + ], + }); + + const collection = zvec.ZVecCreateAndOpen(path, schema); + return new ZVecCodeStore(collection, dimension); + } + + /** Open an existing Zvec collection at the given path. */ + static open(path: string): ZVecCodeStore { + const zvec = getZvec(); + if (!zvec) { + throw new Error("@zvec/zvec is not installed. Install it with: bun add @zvec/zvec"); + } + const collection = zvec.ZVecOpen(path); + // Read the dimension from the collection schema. + const vectorSchema = collection.schema.vectors().find(v => v.name === VECTOR_FIELD); + const dimension = vectorSchema?.dimension ?? 0; + return new ZVecCodeStore(collection, dimension); + } + + /** Check if `@zvec/zvec` is available (installed and loadable). */ + static available(): boolean { + return getZvec() !== null; + } + + /** Upsert code chunks (batch). Chunks without embeddings get a zero vector (FTS-only). */ + upsertChunks(chunks: readonly CodeChunk[]): void { + if (chunks.length === 0) return; + const docs: ZVecDocInput[] = chunks.map(chunk => { + const fields: Record = { + [FIELD_FILE_PATH]: chunk.filePath, + [FIELD_LANGUAGE]: chunk.language, + [FIELD_LINE_START]: chunk.lineStart, + [FIELD_LINE_END]: chunk.lineEnd, + [FIELD_CONTENT]: chunk.content, + [FIELD_CHUNK_HASH]: chunk.chunkHash, + }; + // The vector field is required by the Zvec schema. When no embedding + // is provided, use a zero vector so the doc can still be stored for FTS. + const embedding = chunk.embedding.length > 0 ? chunk.embedding : new Array(this.#dimension).fill(0); + return { id: chunk.id, fields, vectors: { [VECTOR_FIELD]: embedding } }; + }); + this.#collection.upsertSync(docs); + } + + /** Vector similarity search. Returns results sorted by descending score. */ + search(queryEmbedding: readonly number[], topK = 10): ZVecSearchResult[] { + if (queryEmbedding.length === 0) return []; + const query: ZVecQuery = { + fieldName: VECTOR_FIELD, + vector: [...queryEmbedding], + topk: topK, + outputFields: [...ALL_FIELDS], + }; + const docs = this.#collection.querySync(query); + return docs.map(docToResult); + } + + /** Full-text search (BM25) over the `content` field. */ + searchFts(query: string, topK = 10): ZVecSearchResult[] { + if (!query.trim()) return []; + const zvecQuery: ZVecQuery = { + fieldName: FIELD_CONTENT, + fts: { matchString: query }, + topk: topK, + outputFields: [...ALL_FIELDS], + }; + const docs = this.#collection.querySync(zvecQuery); + return docs.map(docToResult); + } + + /** + * Hybrid search combining vector similarity and FTS. + * Uses Zvec's native `multiQuerySync` with reciprocal rank fusion (RRF). + */ + searchHybrid(query: { vector?: readonly number[]; fts?: string }, topK = 10): ZVecSearchResult[] { + const subQueries: ZVecMultiQuery["queries"] = []; + + if (query.vector && query.vector.length > 0) { + subQueries.push({ + fieldName: VECTOR_FIELD, + vector: [...query.vector], + }); + } + + if (query.fts?.trim()) { + subQueries.push({ + fieldName: FIELD_CONTENT, + fts: { matchString: query.fts }, + }); + } + + if (subQueries.length === 0) return []; + if (subQueries.length === 1) { + // Single query mode — use plain querySync for direct scores. + const sq = subQueries[0]; + if (sq.fts) { + return this.searchFts(query.fts ?? "", topK); + } + return this.search(query.vector ?? [], topK); + } + + const multiQuery: ZVecMultiQuery = { + queries: subQueries, + topk: topK, + outputFields: [...ALL_FIELDS], + rerank: { type: "rrf", rankConstant: 60 }, + }; + + const docs = this.#collection.multiQuerySync(multiQuery); + return docs.map(docToResult); + } + + /** + * Scan all documents in the collection, projecting only the requested scalar + * fields. Uses the maximum topk allowed by Zvec and warns if the cap is hit, + * which would indicate silent truncation. + */ + #scanAll(outputFields: readonly string[]): { id: string; fields: Record }[] { + const docs = this.#collection.querySync({ + filter: `${FIELD_FILE_PATH} LIKE '%'`, + topk: ZVEC_MAX_TOPK, + outputFields: [...outputFields], + }); + if (docs.length >= ZVEC_MAX_TOPK) { + logger.warn("zvec: scanAll hit the topk cap — results may be truncated", { + returned: docs.length, + cap: ZVEC_MAX_TOPK, + }); + } + return docs; + } + + /** Remove all chunks for a given file path. */ + removeFile(filePath: string): void { + // For paths without single quotes, use deleteByFilterSync directly — + // it's precise and avoids scanning the collection. + if (!filePath.includes("'")) { + this.#collection.deleteByFilterSync(`${FIELD_FILE_PATH} = '${filePath}'`); + return; + } + // Zvec's SQL filter parser mishandles single quotes inside `=` + // comparisons, so for paths containing single quotes we scan and + // delete by document ID instead. + const docs = this.#scanAll([FIELD_FILE_PATH]); + const idsToDelete: string[] = []; + for (const doc of docs) { + if (String(doc.fields[FIELD_FILE_PATH] ?? "") === filePath) { + idsToDelete.push(doc.id); + } + } + if (idsToDelete.length > 0) { + this.#collection.deleteSync(idsToDelete); + } + } + + /** + * Get all distinct file_path → chunk_hash mappings for change detection. + * Returns a Map where each key is a file path and each value is the set of + * chunk hashes currently stored for that file. + */ + getFileHashes(): Map> { + const result = new Map>(); + const docs = this.#scanAll([FIELD_FILE_PATH, FIELD_CHUNK_HASH]); + for (const doc of docs) { + const fp = String(doc.fields[FIELD_FILE_PATH] ?? ""); + const hash = String(doc.fields[FIELD_CHUNK_HASH] ?? ""); + if (!fp) continue; + const hashes = result.get(fp) ?? new Set(); + hashes.add(hash); + result.set(fp, hashes); + } + return result; + } + + /** Optimize the collection (build vector index). */ + optimize(): void { + this.#collection.optimizeSync(); + } + + /** Get the current document count. */ + get docCount(): number { + return this.#collection.stats.docCount; + } + + /** Close the collection, releasing resources. */ + close(): void { + this.#collection.closeSync(); + } +} diff --git a/packages/mnemopi/src/index.ts b/packages/mnemopi/src/index.ts index 7a83864d1f3..0f89d85bf7f 100644 --- a/packages/mnemopi/src/index.ts +++ b/packages/mnemopi/src/index.ts @@ -31,3 +31,5 @@ export { storeMemory, update, } from "./core/memory"; +export * from "./core/zvec-config"; +export * from "./core/zvec-store"; diff --git a/packages/mnemopi/test/zvec-store.test.ts b/packages/mnemopi/test/zvec-store.test.ts new file mode 100644 index 00000000000..4486e565f58 --- /dev/null +++ b/packages/mnemopi/test/zvec-store.test.ts @@ -0,0 +1,379 @@ +import { afterEach, describe, expect, it } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { zvecChunkOverlap, zvecChunkSize, zvecCodeIndexPath, zvecEnabled, zvecTopK } from "../src/core/zvec-config"; +import { type CodeChunk, ZVecCodeStore } from "../src/core/zvec-store"; + +const ZVEC_AVAILABLE = ZVecCodeStore.available(); + +/** Create a temp directory path for a Zvec collection (does not create the dir — Zvec creates it). */ +async function makeTempDir(): Promise { + const parent = await fs.mkdtemp(path.join(os.tmpdir(), "zvec-test-")); + return path.join(parent, "collection"); +} + +/** Recursively remove a directory (and its parent), ignoring errors. */ +async function cleanupDir(dir: string): Promise { + // `dir` is /collection — remove the parent to clean everything. + const parent = path.dirname(dir); + await fs.rm(parent, { recursive: true, force: true }).catch(() => {}); +} + +function makeChunk(overrides: Partial & { id: string }): CodeChunk { + return { + filePath: `src/${overrides.id}.ts`, + language: "typescript", + lineStart: 1, + lineEnd: 10, + content: `// chunk ${overrides.id}`, + chunkHash: `hash-${overrides.id}`, + embedding: [0.1, 0.2, 0.3, 0.4], + ...overrides, + }; +} + +afterEach(() => { + // No global state to restore — ZVecCodeStore is instance-based. + // This hook exists to satisfy the full-suite-safe contract. +}); + +describe("ZVecCodeStore.available", () => { + it("returns a boolean", () => { + const result = ZVecCodeStore.available(); + expect(typeof result).toBe("boolean"); + }); +}); + +describe("zvec-config", () => { + it("returns default values", () => { + expect(zvecEnabled()).toBe(true); + expect(zvecChunkSize()).toBe(100); + expect(zvecChunkOverlap()).toBe(10); + expect(zvecTopK()).toBe(20); + expect(zvecCodeIndexPath()).toContain("zvec"); + expect(zvecCodeIndexPath()).toContain("code-index"); + }); + + it("respects environment overrides", () => { + const env = { ...process.env, OMP_ZVEC_CHUNK_SIZE: "200", OMP_ZVEC_TOP_K: "50", OMP_ZVEC_ENABLED: "false" }; + + expect(zvecChunkSize(env)).toBe(200); + expect(zvecTopK(env)).toBe(50); + expect(zvecEnabled(env)).toBe(false); + }); +}); + +describe.skipIf(!ZVEC_AVAILABLE)("ZVecCodeStore lifecycle", () => { + let store: ZVecCodeStore; + let dir: string; + + afterEach(async () => { + if (store) { + store.close(); + } + if (dir) { + await cleanupDir(dir); + } + }); + + it("creates a collection, upserts chunks, and searches by vector", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + + const chunks: CodeChunk[] = [ + makeChunk({ + id: "a_1", + filePath: "src/foo.ts", + content: "export function add(a, b) { return a + b }", + embedding: [1, 0, 0, 0], + lineStart: 1, + lineEnd: 1, + chunkHash: "h1", + }), + makeChunk({ + id: "b_1", + filePath: "src/bar.ts", + content: "export class Calculator { multiply(x, y) { return x * y } }", + embedding: [0, 1, 0, 0], + lineStart: 1, + lineEnd: 1, + chunkHash: "h2", + }), + ]; + + store.upsertChunks(chunks); + store.optimize(); + + expect(store.docCount).toBe(2); + + const results = store.search([1, 0, 0, 0], 2); + expect(results.length).toBe(2); + expect(results[0].id).toBe("a_1"); + expect(results[0].filePath).toBe("src/foo.ts"); + expect(results[0].language).toBe("typescript"); + expect(results[0].lineStart).toBe(1); + expect(results[0].lineEnd).toBe(1); + expect(results[0].content).toContain("add"); + expect(results[0].chunkHash).toBe("h1"); + // Zvec COSINE returns a distance (0 = identical vectors), not a similarity. + expect(results[0].score).toBe(0); + // The orthogonal vector should have a higher distance. + expect(results[1].score).toBeGreaterThan(results[0].score); + }); + + it("performs full-text search (BM25)", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + + store.upsertChunks([ + makeChunk({ + id: "fts1", + filePath: "src/search.ts", + content: "binary search tree implementation", + embedding: [0.1, 0.2, 0.3, 0.4], + chunkHash: "fh1", + }), + makeChunk({ + id: "fts2", + filePath: "src/graph.ts", + content: "depth first graph traversal", + embedding: [0.5, 0.6, 0.7, 0.8], + chunkHash: "fh2", + }), + ]); + store.optimize(); + + const results = store.searchFts("search", 10); + expect(results.length).toBeGreaterThanOrEqual(1); + expect(results[0].id).toBe("fts1"); + expect(results[0].content).toContain("search"); + expect(results[0].filePath).toBe("src/search.ts"); + expect(results[0].score).toBeGreaterThan(0); + }); + + it("performs hybrid search combining vector and FTS", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + + store.upsertChunks([ + makeChunk({ + id: "hyb1", + filePath: "src/vector.ts", + content: "vector dot product cosine similarity", + embedding: [1, 0, 0, 0], + chunkHash: "hh1", + }), + makeChunk({ + id: "hyb2", + filePath: "src/parser.ts", + content: "recursive descent parser tokenizer", + embedding: [0, 1, 0, 0], + chunkHash: "hh2", + }), + ]); + store.optimize(); + + const results = store.searchHybrid({ vector: [1, 0, 0, 0], fts: "vector" }, 10); + expect(results.length).toBeGreaterThanOrEqual(1); + // The doc matching both signals should be first. + expect(results[0].id).toBe("hyb1"); + + // Every result has a valid score. + for (const r of results) { + expect(Number.isFinite(r.score)).toBe(true); + } + }); + + it("removes all chunks for a given file path", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + + store.upsertChunks([ + makeChunk({ + id: "rm1", + filePath: "src/remove.ts", + content: "first chunk", + embedding: [1, 0, 0, 0], + chunkHash: "rh1", + }), + makeChunk({ + id: "rm2", + filePath: "src/remove.ts", + content: "second chunk", + embedding: [0, 1, 0, 0], + chunkHash: "rh2", + }), + makeChunk({ + id: "rm3", + filePath: "src/keep.ts", + content: "keep this", + embedding: [0, 0, 1, 0], + chunkHash: "rh3", + }), + ]); + store.optimize(); + + expect(store.docCount).toBe(3); + + store.removeFile("src/remove.ts"); + + expect(store.docCount).toBe(1); + + const results = store.search([0, 0, 1, 0], 10); + expect(results.length).toBe(1); + expect(results[0].id).toBe("rm3"); + expect(results[0].filePath).toBe("src/keep.ts"); + }); + + it("returns correct file_path → chunk_hash mappings from getFileHashes", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + + store.upsertChunks([ + makeChunk({ id: "fh-a1", filePath: "src/a.ts", chunkHash: "hash-a1", embedding: [1, 0, 0, 0] }), + makeChunk({ id: "fh-a2", filePath: "src/a.ts", chunkHash: "hash-a2", embedding: [0, 1, 0, 0] }), + makeChunk({ id: "fh-b1", filePath: "src/b.ts", chunkHash: "hash-b1", embedding: [0, 0, 1, 0] }), + ]); + store.optimize(); + + const hashes = store.getFileHashes(); + + expect(hashes.size).toBe(2); + expect(hashes.has("src/a.ts")).toBe(true); + expect(hashes.has("src/b.ts")).toBe(true); + + const aHashes = hashes.get("src/a.ts"); + expect(aHashes).not.toBeUndefined(); + expect(aHashes?.size).toBe(2); + expect(aHashes?.has("hash-a1")).toBe(true); + expect(aHashes?.has("hash-a2")).toBe(true); + + const bHashes = hashes.get("src/b.ts"); + expect(bHashes?.size).toBe(1); + expect(bHashes?.has("hash-b1")).toBe(true); + }); + + it("handles file paths with single quotes in removeFile", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + + const trickyPath = "src/it's a file.ts"; + store.upsertChunks([ + makeChunk({ + id: "sq1", + filePath: trickyPath, + content: "quoted path", + embedding: [1, 0, 0, 0], + chunkHash: "qh1", + }), + makeChunk({ + id: "sq2", + filePath: "src/normal.ts", + content: "normal path", + embedding: [0, 1, 0, 0], + chunkHash: "qh2", + }), + ]); + store.optimize(); + + expect(store.docCount).toBe(2); + + store.removeFile(trickyPath); + + expect(store.docCount).toBe(1); + const results = store.search([0, 1, 0, 0], 10); + expect(results.length).toBe(1); + expect(results[0].filePath).toBe("src/normal.ts"); + }); + + it("opens an existing collection with open()", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + store.upsertChunks([ + makeChunk({ + id: "open1", + filePath: "src/open.ts", + content: "persisted content", + embedding: [1, 0, 0, 0], + chunkHash: "oh1", + }), + ]); + store.optimize(); + store.close(); + + // Reopen + const reopened = ZVecCodeStore.open(dir); + expect(reopened.docCount).toBe(1); + + const results = reopened.search([1, 0, 0, 0], 1); + expect(results.length).toBe(1); + expect(results[0].id).toBe("open1"); + expect(results[0].content).toBe("persisted content"); + + reopened.close(); + }); + + it("returns empty array for search with empty query embedding", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + store.upsertChunks([ + makeChunk({ id: "empty1", filePath: "src/e.ts", content: "some content", embedding: [1, 0, 0, 0] }), + ]); + store.optimize(); + + const results = store.search([], 10); + expect(results).toEqual([]); + }); + + it("returns empty array for FTS with empty query string", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + store.upsertChunks([ + makeChunk({ id: "emptyfts1", filePath: "src/e.ts", content: "some content", embedding: [1, 0, 0, 0] }), + ]); + store.optimize(); + + const results = store.searchFts(" ", 10); + expect(results).toEqual([]); + }); + + it("returns empty array for hybrid with no vector and no fts", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + store.upsertChunks([ + makeChunk({ id: "hybempty", filePath: "src/e.ts", content: "content", embedding: [1, 0, 0, 0] }), + ]); + store.optimize(); + + expect(store.searchHybrid({}, 10)).toEqual([]); + expect(store.searchHybrid({ vector: [] }, 10)).toEqual([]); + expect(store.searchHybrid({ fts: " " }, 10)).toEqual([]); + }); + + it("upserts chunks without embeddings (FTS-only)", async () => { + dir = await makeTempDir(); + store = ZVecCodeStore.create(dir, 4); + + const chunkNoEmbed: CodeChunk = { + id: "noembed1", + filePath: "src/noembed.ts", + language: "typescript", + lineStart: 1, + lineEnd: 5, + content: "function without embedding", + chunkHash: "neh1", + embedding: [], + }; + + store.upsertChunks([chunkNoEmbed]); + store.optimize(); + + expect(store.docCount).toBe(1); + + // FTS still works + const ftsResults = store.searchFts("without", 10); + expect(ftsResults.length).toBe(1); + expect(ftsResults[0].id).toBe("noembed1"); + }); +});