diff --git a/README.md b/README.md index 637c1cc..32fc742 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ On the SWE-PolyBench Verified agent benchmark (n=24): **R@1 19/24 (79%) at 226K **Indexing** (one-time per repo, then incremental): -1. **Chunk** — parse via [oxc-parser](https://github.com/oxc-project/oxc). Each function, class, exported variable, and tRPC-router-style sub-procedure becomes its own chunk. Drizzle table definitions are chunked. Boundaries follow real code structure. +1. **Chunk** — parse TS/JS via [oxc-parser](https://github.com/oxc-project/oxc) and Python via [tree-sitter-python](https://github.com/tree-sitter/tree-sitter-python). Each function, class, exported variable, and tRPC-router-style sub-procedure becomes its own chunk. Drizzle table definitions and decorated Python functions (`@app.route`, `@dataclass`) keep their decorators attached so the embedder sees the routing intent. Python respects `__all__` for module-level assignments; absent that, the leading-underscore convention is used. Boundaries follow real code structure. 2. **Describe** — each chunk is passed to a local LLM (default: `gemma4:26b`) which writes 2–3 sentences explaining what the code does, the domain, and key constraints. Bridges the vocabulary gap between query intent and code identifiers. 3. **Embed twice**: - The description (prefixed with `filePath [symbol]:`) is embedded → stored as `embedding` BLOB @@ -118,6 +118,8 @@ ollama pull gemma4:26b ollama pull embeddinggemma ``` +The Python chunker uses the native `tree-sitter` Node binding. Prebuilds ship for darwin-arm64, darwin-x64, and linux-x64; on Windows or musl-libc Linux `npm install` will compile from source via `node-gyp` (requires Python 3 + a C/C++ toolchain). + ### Install globally ```bash @@ -203,7 +205,7 @@ tail -f .search-code/last-index.log ## Roadmap -- **Multi-language chunkers** — describer + embedder are language-agnostic. Adding Python, Go, Rust chunkers (via oxc-parser equivalents or tree-sitter) extends coverage. TS/JS only today. +- **More language chunkers** — TS/JS (via oxc-parser) and Python (via tree-sitter) ship today. Go and Rust chunkers via the same tree-sitter mechanism would extend coverage further; the describer + embedder + BM25 layers are language-agnostic. - **Wider SWE-PolyBench coverage** — 24 instances across 3 repos today. Indexing more repos would tighten the variance bounds. - **End-to-end agent benchmark** — measure not just file-finding but task completion + total tokens for full fix cycles. Closest signal to production value. - **Smaller models** — test gemma3:1b describer + smaller embedder for teams with less GPU headroom. diff --git a/bun.lock b/bun.lock index 7fd87d3..52448bb 100644 --- a/bun.lock +++ b/bun.lock @@ -1,6 +1,6 @@ { "lockfileVersion": 1, - "configVersion": 0, + "configVersion": 1, "workspaces": { "": { "name": "semantic-search", @@ -9,6 +9,8 @@ "@huggingface/transformers": "^4.2.0", "@modelcontextprotocol/sdk": "^1.0.0", "oxc-parser": "0.127.0", + "tree-sitter": "^0.25.0", + "tree-sitter-python": "^0.25.0", }, "devDependencies": { "@types/node": "^22.0.0", @@ -17,6 +19,10 @@ }, }, }, + "trustedDependencies": [ + "tree-sitter-python", + "tree-sitter", + ], "packages": { "@anthropic-ai/sdk": ["@anthropic-ai/sdk@0.39.0", "", { "dependencies": { "@types/node": "^18.11.18", "@types/node-fetch": "^2.6.4", "abort-controller": "^3.0.0", "agentkeepalive": "^4.2.1", "form-data-encoder": "1.7.2", "formdata-node": "^4.3.2", "node-fetch": "^2.6.7" } }, "sha512-eMyDIPRZbt1CCLErRCi3exlAvNkBtRe+kW5vvJyef93PmNr/clstYgHhtvmkxN82nlKgzyGPCyGxrm0JQ1ZIdg=="], @@ -382,10 +388,14 @@ "negotiator": ["negotiator@1.0.0", "", {}, "sha512-8Ofs/AUQh8MaEcrlq5xOX0CQ9ypTF5dl78mjlMNfOK08fzpgTHQRQPBxcPlEtIw0yRpws+Zo/3r+5WRby7u3Gg=="], + "node-addon-api": ["node-addon-api@8.7.0", "", {}, "sha512-9MdFxmkKaOYVTV+XVRG8ArDwwQ77XIgIPyKASB1k3JPq3M8fGQQQE3YpMOrKm6g//Ktx8ivZr8xo1Qmtqub+GA=="], + "node-domexception": ["node-domexception@1.0.0", "", {}, "sha512-/jKZoMpw0F8GRwl4/eLROPA3cfcXtLApP0QzLmUT/HuPCZWyB7IY9ZrMeKw2O/nFIqPQB3PVM9aYm0F312AXDQ=="], "node-fetch": ["node-fetch@2.7.0", "", { "dependencies": { "whatwg-url": "^5.0.0" }, "peerDependencies": { "encoding": "^0.1.0" }, "optionalPeers": ["encoding"] }, "sha512-c4FRfUm/dbcWZ7U+1Wq0AwCyFL+3nt2bEw05wfxSz+DWpWsitgmSgYmy2dQdWyKC1694ELPqMs/YzUSNozLt8A=="], + "node-gyp-build": ["node-gyp-build@4.8.4", "", { "bin": { "node-gyp-build": "bin.js", "node-gyp-build-optional": "optional.js", "node-gyp-build-test": "build-test.js" } }, "sha512-LA4ZjwlnUblHVgq0oBF3Jl/6h/Nvs5fzBLwdEF4nuxnFdsfajde4WfxtJr3CaiH+F6ewcIB/q4jQ4UzPyid+CQ=="], + "object-assign": ["object-assign@4.1.1", "", {}, "sha512-rJgTQnkUnH1sFw8yT6VSU3zD3sWmu6sZhIseY8VX+GRu3P6F7Fu+JNDoXfklElbLJSnc3FUQHVe4cU5hj+BcUg=="], "object-inspect": ["object-inspect@1.13.4", "", {}, "sha512-W67iLl4J2EXEGTbfeHCffrjDfitvLANg0UlX3wFUUSTx92KXRFegMHUVgSqE+wvhAbi4WqjGg9czysTV2Epbew=="], @@ -468,6 +478,10 @@ "tr46": ["tr46@0.0.3", "", {}, "sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw=="], + "tree-sitter": ["tree-sitter@0.25.0", "", { "dependencies": { "node-addon-api": "^8.3.0", "node-gyp-build": "^4.8.4" } }, "sha512-PGZZzFW63eElZJDe/b/R/LbsjDDYJa5UEjLZJB59RQsMX+fo0j54fqBPn1MGKav/QNa0JR0zBiVaikYDWCj5KQ=="], + + "tree-sitter-python": ["tree-sitter-python@0.25.0", "", { "dependencies": { "node-addon-api": "^8.5.0", "node-gyp-build": "^4.8.4" }, "peerDependencies": { "tree-sitter": "^0.25.0" }, "optionalPeers": ["tree-sitter"] }, "sha512-eCmJx6zQa35GxaCtQD+wXHOhYqBxEL+bp71W/s3fcDMu06MrtzkVXR437dRrCrbrDbyLuUDJpAgycs7ncngLXw=="], + "tslib": ["tslib@2.8.1", "", {}, "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w=="], "tsx": ["tsx@4.21.0", "", { "dependencies": { "esbuild": "~0.27.0", "get-tsconfig": "^4.7.5" }, "optionalDependencies": { "fsevents": "~2.3.3" }, "bin": "dist/cli.mjs" }, "sha512-5C1sg4USs1lfG0GFb2RLXsdpXqBSEhAaA/0kPL01wxzpMqLILNxIxIOKiILz+cdg/pLnOUxFYOR5yhHU666wbw=="], diff --git a/config.json b/config.json index e036daa..396188d 100644 --- a/config.json +++ b/config.json @@ -15,7 +15,7 @@ }, "indexing": { "sourceRoots": ["src", "socket-server/src"], - "excludePatterns": ["**/*.test.ts", "**/*.test.tsx", "**/*.d.ts", "**/node_modules/**", "**/cloned-projects/**", "**/*icons*.tsx", "**/*icon*.tsx", "**/canvas-icons.tsx"], + "excludePatterns": ["**/*.test.ts", "**/*.test.tsx", "**/*.d.ts", "**/node_modules/**", "**/cloned-projects/**", "**/*icons*.tsx", "**/*icon*.tsx", "**/canvas-icons.tsx", "**/__pycache__/**", "**/*.pyc", "**/.venv/**", "**/venv/**", "**/.pytest_cache/**", "**/.mypy_cache/**", "**/dist/**"], "minChunkLines": 5, "maxChunkLines": 300, "concurrency": 1 diff --git a/package.json b/package.json index 8b1510e..2bd40db 100644 --- a/package.json +++ b/package.json @@ -15,11 +15,17 @@ "@anthropic-ai/sdk": "^0.39.0", "@huggingface/transformers": "^4.2.0", "@modelcontextprotocol/sdk": "^1.0.0", - "oxc-parser": "0.127.0" + "oxc-parser": "0.127.0", + "tree-sitter": "^0.25.0", + "tree-sitter-python": "^0.25.0" }, "devDependencies": { "@types/node": "^22.0.0", "tsx": "^4.21.0", "typescript": "^5.4.0" - } + }, + "trustedDependencies": [ + "tree-sitter", + "tree-sitter-python" + ] } diff --git a/src/chunker-python.ts b/src/chunker-python.ts new file mode 100644 index 0000000..ef8e1f6 --- /dev/null +++ b/src/chunker-python.ts @@ -0,0 +1,270 @@ +import Parser from 'tree-sitter'; +import Python from 'tree-sitter-python'; +import type { ExtractResult, RawChunk } from './chunker-ts.js'; + +// tree-sitter-python's exported Language shape lags the tree-sitter Language +// type (missing `name`); the runtime value works fine. Cast at the boundary. +// eslint-disable-next-line @typescript-eslint/no-explicit-any +const PythonLanguage = Python as any; + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +type TSNode = any; + +// ─── Parser singleton ───────────────────────────────────────────────────────── +// +// tree-sitter's Node binding is sync after construction; we keep a single +// Parser instance to avoid per-file setLanguage cost. initPython() exists +// to match the async init() contract and to make a future swap to +// web-tree-sitter (WASM, async init) a one-line change. + +let parser: InstanceType | null = null; + +export async function initPython(): Promise { + if (parser) return; + parser = new Parser(); + parser.setLanguage(PythonLanguage); +} + +function getParser(): InstanceType { + if (!parser) { + parser = new Parser(); + parser.setLanguage(PythonLanguage); + } + return parser; +} + +// ─── Helpers ────────────────────────────────────────────────────────────────── + +const PROPERTY_CHUNK_MIN_BYTES = 80; + +function nodeName(node: TSNode): string | null { + const id = node.childForFieldName?.('name'); + if (id && typeof id.text === 'string') return id.text; + return null; +} + +// `decorated_definition` wraps a `function_definition` or `class_definition`. +// Unwrap to the inner def for name lookup; keep the outer span so decorators +// (e.g. `@app.route("/...")`) stay attached to the chunk's source. +function innerDef(node: TSNode): TSNode { + if (node.type !== 'decorated_definition') return node; + for (const child of node.namedChildren) { + if (child.type === 'function_definition' || child.type === 'class_definition') return child; + } + return node; +} + +// ─── __all__ parsing ────────────────────────────────────────────────────────── +// +// If a module declares `__all__ = ["foo", "bar"]`, treat that as the export +// whitelist for module-level assignments. Functions and classes are always +// emitted regardless. Absent or unparseable → fall back to the leading- +// underscore convention. + +function parseAllList(rootChildren: TSNode[]): Set | null { + for (const node of rootChildren) { + if (node.type !== 'expression_statement') continue; + const assign = node.namedChild(0); + if (!assign || assign.type !== 'assignment') continue; + const left = assign.childForFieldName('left'); + const right = assign.childForFieldName('right'); + if (!left || !right) continue; + if (left.type !== 'identifier' || left.text !== '__all__') continue; + if (right.type !== 'list') continue; + const names = new Set(); + for (const item of right.namedChildren) { + if (item.type !== 'string') continue; + // string_content is the unquoted text + let content = ''; + for (const part of item.namedChildren) { + if (part.type === 'string_content') content += part.text; + } + if (content) names.add(content); + } + return names; + } + return null; +} + +// ─── Module-level assignment heuristic ──────────────────────────────────────── +// +// Python has no `export` keyword. Closest analogue to TS's "exported +// VariableDeclaration with non-function init" rule: top-level assignments +// whose RHS is a call expression (e.g. `app = FastAPI()`, `engine = +// create_engine(...)`, `Base = declarative_base()`). These are structural +// plumbing whose descriptions carry domain vocabulary worth indexing. +// +// Rules: +// - LHS must be a single identifier +// - LHS must not start with `_` (private convention) and must not be a dunder +// - RHS must be a `call` (and span at least PROPERTY_CHUNK_MIN_BYTES) +// - If __all__ is declared, LHS must be in __all__ +// +// Lambda RHS is also chunked (mirrors TS's arrow-init case). + +function tryExtractAssignment( + exprStmt: TSNode, + allowlist: Set | null, + out: RawChunk[], +): void { + const assign = exprStmt.namedChild(0); + if (!assign || assign.type !== 'assignment') return; + + const left = assign.childForFieldName('left'); + const right = assign.childForFieldName('right'); + if (!left || !right) return; + if (left.type !== 'identifier') return; + + const name: string = left.text; + if (name.startsWith('_')) return; // covers `_private` and dunders like `__all__` + if (allowlist && !allowlist.has(name)) return; + + if (right.type === 'lambda') { + out.push({ startOffset: exprStmt.startIndex, endOffset: exprStmt.endIndex, symbolName: name }); + return; + } + if (right.type !== 'call') return; + if (right.endIndex - right.startIndex < PROPERTY_CHUNK_MIN_BYTES) return; + + out.push({ startOffset: exprStmt.startIndex, endOffset: exprStmt.endIndex, symbolName: name }); +} + +// ─── Class body walk ────────────────────────────────────────────────────────── + +function extractMethods(classDef: TSNode, out: RawChunk[]): void { + const body = classDef.childForFieldName('body'); + if (!body) return; + for (const member of body.namedChildren) { + if (member.type === 'function_definition') { + const name = nodeName(member); + out.push({ + startOffset: member.startIndex, + endOffset: member.endIndex, + symbolName: name, + }); + } else if (member.type === 'decorated_definition') { + const inner = innerDef(member); + if (inner.type === 'function_definition') { + out.push({ + startOffset: member.startIndex, + endOffset: member.endIndex, + symbolName: nodeName(inner), + }); + } + } + } +} + +// ─── Top-level walk ─────────────────────────────────────────────────────────── + +function collectFromModule(rootChildren: TSNode[]): RawChunk[] { + const result: RawChunk[] = []; + const allowlist = parseAllList(rootChildren); + + for (const node of rootChildren) { + if (node.type === 'function_definition') { + result.push({ + startOffset: node.startIndex, + endOffset: node.endIndex, + symbolName: nodeName(node), + }); + continue; + } + + if (node.type === 'class_definition') { + result.push({ + startOffset: node.startIndex, + endOffset: node.endIndex, + symbolName: nodeName(node), + }); + extractMethods(node, result); + continue; + } + + if (node.type === 'decorated_definition') { + const inner = innerDef(node); + // Emit one chunk for the decorated outer span; methods walked inside if it's a class. + result.push({ + startOffset: node.startIndex, + endOffset: node.endIndex, + symbolName: nodeName(inner), + }); + if (inner.type === 'class_definition') { + extractMethods(inner, result); + } + continue; + } + + if (node.type === 'expression_statement') { + tryExtractAssignment(node, allowlist, result); + continue; + } + } + + return result; +} + +// ─── Import + docstring prefixes ────────────────────────────────────────────── + +function buildImportPrefix(rootChildren: TSNode[]): string { + const parts: string[] = []; + for (const node of rootChildren) { + if (node.type === 'import_statement' || node.type === 'import_from_statement') { + parts.push(node.text); + } + } + return parts.length > 0 ? parts.join('\n') + '\n\n' : ''; +} + +// Python module docstring: the first `expression_statement` whose only child +// is a string, occurring before any def/class/import. Mirrors TS's top-of-file +// JSDoc prefix. +function buildModuleDocstring(rootChildren: TSNode[]): string { + for (const node of rootChildren) { + if (node.type === 'expression_statement') { + const inner = node.namedChild(0); + if (inner && inner.type === 'string') { + return node.text + '\n\n'; + } + return ''; // first non-docstring expression statement → no module docstring + } + // Imports and dunder assignments may legitimately precede a module docstring + // in some style guides; skip past them rather than aborting. + if ( + node.type === 'import_statement' || + node.type === 'import_from_statement' || + node.type === 'comment' + ) { + continue; + } + // Any other top-level node (def, class, real assignment) → no docstring. + return ''; + } + return ''; +} + +// ─── Entry point ────────────────────────────────────────────────────────────── + +export function extractPython( + code: string, + _lines: string[], + _offsetToLine: (offset: number) => number, +): ExtractResult { + const tree = getParser().parse(code); + const root = tree.rootNode; + + // tree-sitter never throws on syntax errors; it produces an ERROR node and + // continues. Partial chunks are usually better than no chunks, so we accept + // hasError trees but surface a warning. + const parseWarnings: string[] = []; + if (root.hasError) { + parseWarnings.push('python parse tree contains ERROR nodes; chunks may be partial'); + } + + const children = root.namedChildren; + const rawChunks = collectFromModule(children); + const importPrefix = buildImportPrefix(children); + const fileDocPrefix = buildModuleDocstring(children); + + return { rawChunks, importPrefix, fileDocPrefix, parseWarnings }; +} diff --git a/src/chunker-ts.ts b/src/chunker-ts.ts new file mode 100644 index 0000000..fe5069e --- /dev/null +++ b/src/chunker-ts.ts @@ -0,0 +1,273 @@ +import { parseSync } from 'oxc-parser'; + +// ─── Types ──────────────────────────────────────────────────────────────────── + +export interface RawChunk { + startOffset: number; + endOffset: number; // exclusive (ESTree convention) + symbolName: string | null; +} + +export interface ExtractResult { + rawChunks: RawChunk[]; + importPrefix: string; + fileDocPrefix: string; + parseError?: string; + parseWarnings: string[]; +} + +// eslint-disable-next-line @typescript-eslint/no-explicit-any +type ASTNode = any; + +// ─── AST extraction ─────────────────────────────────────────────────────────── + +function extractFromDecl(node: ASTNode, exportWrapper: ASTNode | null, result: RawChunk[]): void { + const container = exportWrapper ?? node; + + switch (node.type) { + case 'FunctionDeclaration': { + result.push({ + startOffset: container.start, + endOffset: container.end, + symbolName: node.id?.name ?? null, + }); + break; + } + + case 'ClassDeclaration': { + result.push({ + startOffset: container.start, + endOffset: container.end, + symbolName: node.id?.name ?? null, + }); + for (const member of node.body?.body ?? []) { + if (member.type === 'MethodDefinition' && member.key?.type === 'Identifier') { + result.push({ + startOffset: member.start, + endOffset: member.end, + symbolName: member.key.name, + }); + } + } + break; + } + + case 'VariableDeclaration': { + for (const decl of node.declarations ?? []) { + const init = decl.init; + if (!init) continue; + const name = decl.id?.type === 'Identifier' ? decl.id.name : null; + + if (init.type === 'ArrowFunctionExpression' || init.type === 'FunctionExpression') { + result.push({ startOffset: container.start, endOffset: container.end, symbolName: name }); + continue; + } + + // Top-level EXPORTED declarations: always chunk regardless of init kind. + // Captures Drizzle table defs, zod schemas, exported config objects, + // route registries — plumbing whose descriptions carry domain vocabulary. + if (exportWrapper) { + result.push({ startOffset: container.start, endOffset: container.end, symbolName: name }); + + // tRPC-router pattern: per-procedure sub-chunks for properties whose + // value spans more than ~3 lines. Without these, handlers past the + // maxChunkLines cap are never indexed. + extractObjectPropertyChunks(init, result); + continue; + } + + // Non-exported CallExpression with a function first-arg — module-scope HOC. + if (init.type === 'CallExpression') { + const firstArg = init.arguments?.[0]; + if ( + firstArg && + (firstArg.type === 'ArrowFunctionExpression' || firstArg.type === 'FunctionExpression') + ) { + result.push({ + startOffset: container.start, + endOffset: container.end, + symbolName: name, + }); + } + } + } + break; + } + } +} + +const PROPERTY_CHUNK_MIN_BYTES = 80; // ≈3 lines; skip 1-line Drizzle column defs + +function extractObjectPropertyChunks(initNode: ASTNode, out: RawChunk[]): void { + if (!initNode || typeof initNode !== 'object') return; + if (initNode.type !== 'CallExpression') return; + for (const arg of initNode.arguments ?? []) { + if (!arg || arg.type !== 'ObjectExpression') continue; + for (const prop of arg.properties ?? []) { + if (!prop || prop.type !== 'Property') continue; + if (prop.computed) continue; + const keyName = + prop.key?.type === 'Identifier' + ? prop.key.name + : prop.key?.type === 'Literal' && typeof prop.key.value === 'string' + ? prop.key.value + : null; + if (!keyName) continue; + const value = prop.value; + if (!value || value.end - value.start < PROPERTY_CHUNK_MIN_BYTES) continue; + out.push({ startOffset: prop.start, endOffset: prop.end, symbolName: keyName }); + } + } +} + +function collectFromBody(body: ASTNode[]): RawChunk[] { + const result: RawChunk[] = []; + for (const node of body) { + if ( + (node.type === 'ExportNamedDeclaration' || node.type === 'ExportDefaultDeclaration') && + node.declaration + ) { + extractFromDecl(node.declaration, node, result); + } else { + extractFromDecl(node, null, result); + } + } + return result; +} + +// ─── Prefix extraction ──────────────────────────────────────────────────────── + +function buildImportPrefix(body: ASTNode[], lines: string[], offsetToLine: (n: number) => number): string { + const importLines: string[] = []; + for (const node of body) { + if (node.type === 'ImportDeclaration') { + const startL = offsetToLine(node.start); + const endL = offsetToLine(Math.max(node.start, node.end - 1)); + importLines.push(...lines.slice(startL - 1, endL)); + } + } + return importLines.length > 0 ? importLines.join('\n') + '\n\n' : ''; +} + +// Top-of-file JSDoc blocks before the first non-import declaration. These +// carry module-level domain context that individual function bodies omit. +function buildFileDocPrefix(lines: string[]): string { + const fileJsDocLines: string[] = []; + let idx = 0; + while (idx < lines.length) { + const trimmed = lines[idx].trim(); + if (trimmed === '') { + idx++; + continue; + } + if (trimmed.startsWith('//')) { + idx++; + continue; + } + if (trimmed.startsWith('import ')) { + while (idx < lines.length && !lines[idx].includes(';')) idx++; + idx++; + continue; + } + if (trimmed.startsWith('/**') || trimmed === '*') { + const blockStart = idx; + while (idx < lines.length && !lines[idx].includes('*/')) idx++; + idx++; // include closing `*/` + fileJsDocLines.push(...lines.slice(blockStart, idx)); + continue; + } + break; + } + return fileJsDocLines.length > 0 ? fileJsDocLines.join('\n') + '\n\n' : ''; +} + +// ─── Entry point ────────────────────────────────────────────────────────────── + +export function extractTs( + filePath: string, + code: string, + lines: string[], + offsetToLine: (n: number) => number, +): ExtractResult { + let parsed: ReturnType; + try { + parsed = parseSync(filePath, code); + } catch (err) { + return { + rawChunks: [], + importPrefix: '', + fileDocPrefix: '', + parseError: (err as Error).message, + parseWarnings: [], + }; + } + + const parseWarnings = (parsed.errors ?? []).map((e) => e.message); + const rawChunks = collectFromBody(parsed.program.body); + const importPrefix = buildImportPrefix(parsed.program.body, lines, offsetToLine); + const fileDocPrefix = buildFileDocPrefix(lines); + + return { rawChunks, importPrefix, fileDocPrefix, parseWarnings }; +} + +// ─── Manifest extraction ────────────────────────────────────────────────────── +// +// Collects identifier-shaped tokens from a chunk. Walks ESTree Identifier +// and identifier-shaped string Literal nodes, skipping any TS*-prefixed +// node so type-position primitives (`string`, `Promise`) don't leak in. + +const MANIFEST_TOKEN_LIMIT = 50; +const IDENT_LITERAL_RE = /^[a-zA-Z_][a-zA-Z0-9_]{1,49}$/; + +export function extractManifest(rawCode: string): string { + let parsed: ReturnType; + try { + parsed = parseSync('manifest.ts', rawCode); + } catch { + return ''; + } + + const tokens = new Set(); + + function visit(node: ASTNode, inTypePosition: boolean): void { + if (!node || typeof node !== 'object') return; + if (Array.isArray(node)) { + for (const child of node) visit(child, inTypePosition); + return; + } + + const type = node.type; + const entersTypePosition = typeof type === 'string' && type.startsWith('TS'); + const childrenInType = inTypePosition || entersTypePosition; + + if (!childrenInType && !entersTypePosition) { + if (type === 'Identifier' && typeof node.name === 'string') { + tokens.add(node.name); + } else if ( + type === 'Literal' && + typeof node.value === 'string' && + IDENT_LITERAL_RE.test(node.value) + ) { + tokens.add(node.value); + } + } + + for (const key in node) { + if (key === 'type' || key === 'start' || key === 'end' || key === 'loc' || key === 'range') + continue; + const child = node[key]; + if (child && typeof child === 'object') visit(child, childrenInType); + } + } + + visit(parsed.program as ASTNode, false); + + const filtered: string[] = []; + for (const t of tokens) { + if (t.length < 2) continue; + filtered.push(t); + if (filtered.length >= MANIFEST_TOKEN_LIMIT) break; + } + + return filtered.join(' '); +} diff --git a/src/chunker.ts b/src/chunker.ts index f587ac9..cdac784 100644 --- a/src/chunker.ts +++ b/src/chunker.ts @@ -1,10 +1,12 @@ import crypto from 'crypto'; import fs from 'fs'; -// oxc-parser ships native arm64 binaries; no WASM, no async initialization required -import { parseSync } from 'oxc-parser'; import path from 'path'; +import { extractTs, type ExtractResult, type RawChunk } from './chunker-ts.js'; +import { extractPython, initPython } from './chunker-python.js'; import { loadConfig, PROJECT_ROOT } from './project.js'; +export { extractManifest } from './chunker-ts.js'; + // ─── Types ──────────────────────────────────────────────────────────────────── export interface Chunk { @@ -12,19 +14,20 @@ export interface Chunk { startLine: number; // 1-indexed endLine: number; // 1-indexed inclusive symbolName: string | null; - language: string; // 'typescript' | 'javascript' + language: string; // 'typescript' | 'javascript' | 'python' rawCode: string; - codeHash: string; // sha256 hex of rawCode + codeHash: string; // sha256 hex of function body (not the prefixes) fileMtime: number; // file modification time as Unix ms } // ─── Language mapping ───────────────────────────────────────────────────────── -const EXT_TO_LANG: Record = { +const EXT_TO_LANG: Record = { '.ts': 'typescript', '.tsx': 'typescript', '.js': 'javascript', '.jsx': 'javascript', + '.py': 'python', }; // ─── Glob matching ──────────────────────────────────────────────────────────── @@ -59,234 +62,29 @@ function buildLineOffsets(code: string): number[] { return offsets; } -// Returns 1-indexed line number for a given character offset (binary search). -function offsetToLine(offsets: number[], offset: number): number { - let lo = 0, - hi = offsets.length - 1; - while (lo < hi) { - const mid = (lo + hi + 1) >> 1; - if (offsets[mid] <= offset) lo = mid; - else hi = mid - 1; - } - return lo + 1; -} - -// ─── AST extraction ─────────────────────────────────────────────────────────── - -interface RawChunk { - startOffset: number; - endOffset: number; // exclusive (ESTree convention) - symbolName: string | null; -} - -// eslint-disable-next-line @typescript-eslint/no-explicit-any -type ASTNode = any; - -function extractFromDecl(node: ASTNode, exportWrapper: ASTNode | null, result: RawChunk[]): void { - const container = exportWrapper ?? node; - - switch (node.type) { - case 'FunctionDeclaration': { - result.push({ - startOffset: container.start, - endOffset: container.end, - symbolName: node.id?.name ?? null, - }); - break; - } - - case 'ClassDeclaration': { - result.push({ - startOffset: container.start, - endOffset: container.end, - symbolName: node.id?.name ?? null, - }); - // Extract each method as its own chunk - for (const member of node.body?.body ?? []) { - if (member.type === 'MethodDefinition' && member.key?.type === 'Identifier') { - result.push({ - startOffset: member.start, - endOffset: member.end, - symbolName: member.key.name, - }); - } - } - break; +function offsetToLineFn(offsets: number[]): (offset: number) => number { + return (offset: number) => { + let lo = 0, + hi = offsets.length - 1; + while (lo < hi) { + const mid = (lo + hi + 1) >> 1; + if (offsets[mid] <= offset) lo = mid; + else hi = mid - 1; } - - case 'VariableDeclaration': { - for (const decl of node.declarations ?? []) { - const init = decl.init; - if (!init) continue; - const name = decl.id?.type === 'Identifier' ? decl.id.name : null; - - // Function-shaped initializers — always chunk (named function exports, hooks) - if (init.type === 'ArrowFunctionExpression' || init.type === 'FunctionExpression') { - result.push({ startOffset: container.start, endOffset: container.end, symbolName: name }); - continue; - } - - // Top-level EXPORTED declarations: always chunk regardless of init kind. - // Captures Drizzle table defs (`sqliteTable(...)`), zod schemas (`z.object(...)`), - // exported config objects, route registries — all structural plumbing whose - // descriptions and embed manifests carry domain vocabulary worth indexing. - if (exportWrapper) { - result.push({ startOffset: container.start, endOffset: container.end, symbolName: name }); - - // tRPC-router pattern: `export const fooRouter = router({ procA: ..., procB: ... })`. - // Without per-procedure chunks, only the first ~maxChunkLines lines of the router - // get chunked (Mode 1 failure). Walk into CallExpression arguments and emit a - // sub-chunk for any Property whose value spans more than a few lines — - // captures procedure handlers buried inside the router object literal. - extractObjectPropertyChunks(init, result); - continue; - } - - // Non-exported CallExpression with a function first-arg — HOC pattern at - // module scope (e.g. `const X = memo(() => {})` without re-export). - if (init.type === 'CallExpression') { - const firstArg = init.arguments?.[0]; - if ( - firstArg && - (firstArg.type === 'ArrowFunctionExpression' || firstArg.type === 'FunctionExpression') - ) { - result.push({ - startOffset: container.start, - endOffset: container.end, - symbolName: name, - }); - } - } - } - break; - } - } + return lo + 1; + }; } -// Mode 1 fix: tRPC routers, command palettes, route registries, and similar -// dispatch-tables are CallExpression arguments shaped like `{ procA, procB, … }` -// where each property's value is a long expression chain. The outer -// VariableDeclaration emits one chunk capped at maxChunkLines, so handlers -// past the cap are never indexed. Walk into CallExpression args' ObjectExpression -// properties and emit a sub-chunk per non-trivial Property. -const PROPERTY_CHUNK_MIN_BYTES = 80; // ≈3 lines of code; skip 1-line Drizzle column defs - -function extractObjectPropertyChunks(initNode: ASTNode, out: RawChunk[]): void { - if (!initNode || typeof initNode !== 'object') return; - if (initNode.type !== 'CallExpression') return; - for (const arg of initNode.arguments ?? []) { - if (!arg || arg.type !== 'ObjectExpression') continue; - for (const prop of arg.properties ?? []) { - if (!prop || prop.type !== 'Property') continue; - if (prop.computed) continue; - const keyName = - prop.key?.type === 'Identifier' - ? prop.key.name - : prop.key?.type === 'Literal' && typeof prop.key.value === 'string' - ? prop.key.value - : null; - if (!keyName) continue; - const value = prop.value; - if (!value || value.end - value.start < PROPERTY_CHUNK_MIN_BYTES) continue; - out.push({ startOffset: prop.start, endOffset: prop.end, symbolName: keyName }); - } - } -} - -function collectFromBody(body: ASTNode[]): RawChunk[] { - const result: RawChunk[] = []; - - for (const node of body) { - if ( - (node.type === 'ExportNamedDeclaration' || node.type === 'ExportDefaultDeclaration') && - node.declaration - ) { - extractFromDecl(node.declaration, node, result); - } else { - extractFromDecl(node, null, result); - } - } - - return result; -} - -// ─── Manifest extraction ────────────────────────────────────────────────────── -// -// Collects identifier-shaped tokens from the chunk AST. Prepended to the embed -// text so plumbing files (Drizzle schemas, repo readers) surface their domain -// vocabulary in dense space — column names like `sandbox_id` and call-site -// identifiers like `userIntegrations` that don't appear in the LLM-generated -// description. +// ─── Initialization ─────────────────────────────────────────────────────────── // -// Walks ESTree `Identifier` and identifier-shaped string `Literal` nodes, -// skipping subtrees under any TS*-prefixed node so type-position primitives -// (`string`, `number`, `Promise`, `Array`) don't leak into the manifest. -// No hand-coded keyword list — ESTree never emits language keywords as -// Identifier nodes. - -const MANIFEST_TOKEN_LIMIT = 50; -const IDENT_LITERAL_RE = /^[a-zA-Z_][a-zA-Z0-9_]{1,49}$/; - -export function extractManifest(rawCode: string): string { - let parsed: ReturnType; - try { - parsed = parseSync('manifest.ts', rawCode); - } catch { - return ''; - } - - const tokens = new Set(); +// oxc-parser is sync. tree-sitter native binding loads sync too, but we expose +// async init for forward compatibility with web-tree-sitter (WASM) if we ever +// need to swap. - function visit(node: ASTNode, inTypePosition: boolean): void { - if (!node || typeof node !== 'object') return; - if (Array.isArray(node)) { - for (const child of node) visit(child, inTypePosition); - return; - } - - const type = node.type; - // Any TypeScript type-AST subtree (TSTypeAnnotation, TSTypeReference, - // TSStringKeyword, TSTypeParameterInstantiation, …) — recurse without - // collecting so type primitives never enter the manifest. - const entersTypePosition = typeof type === 'string' && type.startsWith('TS'); - const childrenInType = inTypePosition || entersTypePosition; - - if (!childrenInType && !entersTypePosition) { - if (type === 'Identifier' && typeof node.name === 'string') { - tokens.add(node.name); - } else if ( - type === 'Literal' && - typeof node.value === 'string' && - IDENT_LITERAL_RE.test(node.value) - ) { - tokens.add(node.value); - } - } - - for (const key in node) { - if (key === 'type' || key === 'start' || key === 'end' || key === 'loc' || key === 'range') - continue; - const child = node[key]; - if (child && typeof child === 'object') visit(child, childrenInType); - } - } - - visit(parsed.program as ASTNode, false); - - const filtered: string[] = []; - for (const t of tokens) { - if (t.length < 2) continue; - filtered.push(t); - if (filtered.length >= MANIFEST_TOKEN_LIMIT) break; - } - - return filtered.join(' '); +export async function init(): Promise { + await initPython(); } -// ─── Initialization (no-op — oxc-parser requires no async setup) ────────────── - -export async function init(): Promise {} - // ─── File chunker ───────────────────────────────────────────────────────────── export async function chunkFile(filePath: string): Promise { @@ -310,90 +108,68 @@ export async function chunkFile(filePath: string): Promise { const code = await fs.promises.readFile(filePath, 'utf-8'); const relPath = path.relative(PROJECT_ROOT, filePath); - let parsed: ReturnType; - try { - parsed = parseSync(filePath, code); - } catch (err) { - console.warn(`[chunker] Parse error in ${relPath}: ${(err as Error).message}`); - return []; - } - - if (parsed.errors?.length > 0) { - for (const e of parsed.errors) { - console.warn(`[chunker] Parse warning in ${relPath}: ${e.message}`); - } - } - - const rawChunks = collectFromBody(parsed.program.body); - if (rawChunks.length === 0) return []; - const lineOffsets = buildLineOffsets(code); + const offsetToLine = offsetToLineFn(lineOffsets); const lines = code.split('\n'); - // Collect top-level import statements to prepend as context for the describer. - // This lets the LLM see `import { useState } from 'react'` etc. when describing - // a hook, dramatically improving descriptions for cross-cutting concerns. - const importLines: string[] = []; - for (const node of parsed.program.body) { - if (node.type === 'ImportDeclaration') { - const startL = offsetToLine(lineOffsets, node.start); - const endL = offsetToLine(lineOffsets, Math.max(node.start, node.end - 1)); - importLines.push(...lines.slice(startL - 1, endL)); - } + let result: ExtractResult; + if (lang === 'python') { + result = extractPython(code, lines, offsetToLine); + } else { + result = extractTs(filePath, code, lines, offsetToLine); } - const importPrefix = importLines.length > 0 ? importLines.join('\n') + '\n\n' : ''; - - // Collect top-of-file JSDoc blocks that appear before the first non-import declaration. - // These carry module-level domain context (e.g. "Why: agent executions allocate large - // per-stream heap") that the describer needs but individual function bodies omit. - const fileJsDocLines: string[] = []; - { - let idx = 0; - while (idx < lines.length) { - const trimmed = lines[idx].trim(); - if (trimmed === '') { - idx++; - continue; - } - if (trimmed.startsWith('//')) { - idx++; - continue; - } - if (trimmed.startsWith('import ')) { - // skip multi-line import statements - while (idx < lines.length && !lines[idx].includes(';')) idx++; - idx++; - continue; - } - if (trimmed.startsWith('/**') || trimmed === '*') { - const blockStart = idx; - while (idx < lines.length && !lines[idx].includes('*/')) idx++; - idx++; // include the closing `*/` line - fileJsDocLines.push(...lines.slice(blockStart, idx)); - continue; - } - break; // first non-import, non-comment declaration — stop - } + + if (result.parseError) { + console.warn(`[chunker] Parse error in ${relPath}: ${result.parseError}`); + return []; + } + for (const w of result.parseWarnings) { + console.warn(`[chunker] Parse warning in ${relPath}: ${w}`); } - const fileJsDocPrefix = fileJsDocLines.length > 0 ? fileJsDocLines.join('\n') + '\n\n' : ''; + if (result.rawChunks.length === 0) return []; + + return assembleChunks( + result.rawChunks, + result.importPrefix, + result.fileDocPrefix, + lines, + offsetToLine, + relPath, + lang, + fileMtime, + minChunkLines, + maxChunkLines, + ); +} + +function assembleChunks( + rawChunks: RawChunk[], + importPrefix: string, + fileDocPrefix: string, + lines: string[], + offsetToLine: (offset: number) => number, + relPath: string, + lang: string, + fileMtime: number, + minChunkLines: number, + maxChunkLines: number, +): Chunk[] { const chunks: Chunk[] = []; for (const raw of rawChunks) { - const startLine = offsetToLine(lineOffsets, raw.startOffset); - // endOffset is exclusive; last character is at endOffset - 1 - const rawEndLine = offsetToLine(lineOffsets, Math.max(raw.startOffset, raw.endOffset - 1)); + const startLine = offsetToLine(raw.startOffset); + const rawEndLine = offsetToLine(Math.max(raw.startOffset, raw.endOffset - 1)); const endLine = Math.min(rawEndLine, startLine + maxChunkLines - 1); const lineCount = endLine - startLine + 1; - // Named functions are complete semantic units — always include them. - // Only apply minChunkLines to unnamed chunks (anonymous exports) to avoid noise. + // Named symbols are complete semantic units — always keep. + // minChunkLines only applies to unnamed chunks (anonymous exports). if (raw.symbolName === null && lineCount < minChunkLines) continue; if (lineCount < 2) continue; // absolute floor: skip single-line no-ops const functionCode = lines.slice(startLine - 1, endLine).join('\n'); - // codeHash is based on the function body only — stable across prefix changes - const rawCode = fileJsDocPrefix + importPrefix + functionCode; + const rawCode = fileDocPrefix + importPrefix + functionCode; chunks.push({ filePath: relPath, @@ -420,10 +196,9 @@ export interface WalkStats { // `knownMtimes` (filePath → stored file_mtime) gates the parser: files whose // disk mtime is ≤ the stored mtime are skipped entirely. Pass an empty Map to -// force a full re-walk (initial index, or after schema/chunker rule change). -// `seenFilePaths`, if provided, is populated with every matched file path -// (parsed OR mtime-skipped) so callers can compute orphans correctly — files -// the walker skipped on disk still exist and must NOT be deleted from the DB. +// force a full re-walk. `seenFilePaths`, if provided, is populated with every +// matched file path (parsed OR mtime-skipped) so callers can compute orphans +// correctly — skipped files still exist and must NOT be deleted from the DB. export async function* walkAndChunk( projectRoot: string, knownMtimes: Map = new Map(), @@ -456,7 +231,7 @@ export async function* walkAndChunk( if (stat.isDirectory()) continue; const ext = path.extname(absFilePath).toLowerCase(); - if (!EXT_TO_LANG[ext]) continue; // skip non-TS/JS silently + if (!EXT_TO_LANG[ext]) continue; const relPath = path.relative(projectRoot, absFilePath); if (matchesAnyExcludePattern(relPath, excludePatterns)) continue; @@ -464,11 +239,6 @@ export async function* walkAndChunk( stats.filesSeen++; seenFilePaths?.add(relPath); - // mtime gate: skip parse + chunk if disk mtime hasn't advanced beyond - // the stored mtime for this path. Saves the heavy oxc-parser pass on - // unchanged files. New/changed files (and files absent from knownMtimes) - // still get parsed. seenFilePaths is populated above so orphan cleanup - // does NOT delete the existing chunks of an unchanged file. const known = knownMtimes.get(relPath); if (known !== undefined && stat.mtimeMs <= known) { stats.filesSkippedMtime++; diff --git a/src/project.ts b/src/project.ts index 4129e1d..125c345 100644 --- a/src/project.ts +++ b/src/project.ts @@ -30,6 +30,13 @@ export const DEFAULT_CONFIG: SemanticSearchConfig = { '**/*.d.ts', '**/node_modules/**', '**/.semantic-search/**', + '**/__pycache__/**', + '**/*.pyc', + '**/.venv/**', + '**/venv/**', + '**/.pytest_cache/**', + '**/.mypy_cache/**', + '**/dist/**', ], minChunkLines: 5, maxChunkLines: 300, diff --git a/src/test/chunker-python.test.ts b/src/test/chunker-python.test.ts new file mode 100644 index 0000000..6587ec3 --- /dev/null +++ b/src/test/chunker-python.test.ts @@ -0,0 +1,205 @@ +import { describe, it, before } from 'node:test'; +import assert from 'node:assert/strict'; +import { extractPython, initPython } from '../chunker-python.js'; + +before(async () => { + await initPython(); +}); + +function extract(code: string) { + const lines = code.split('\n'); + // The offsetToLine arg isn't consumed by the python extractor — startIndex/ + // endIndex byte offsets are returned directly and the caller in chunker.ts + // does the line mapping. Pass a stub so the signature matches. + return extractPython(code, lines, (offset: number) => { + let line = 1; + for (let i = 0; i < offset && i < code.length; i++) { + if (code[i] === '\n') line++; + } + return line; + }); +} + +function namesOf(result: ReturnType): (string | null)[] { + return result.rawChunks.map((c) => c.symbolName); +} + +describe('chunker-python', () => { + it('chunks a decorated function with the decorator span attached', () => { + const code = `import os + +@app.route("/hello") +def greet(name): + return f"hello {name}" +`; + const result = extract(code); + assert.equal(result.rawChunks.length, 1); + const chunk = result.rawChunks[0]; + assert.equal(chunk.symbolName, 'greet'); + const span = code.slice(chunk.startOffset, chunk.endOffset); + assert.ok( + span.startsWith('@app.route'), + `expected span to begin with decorator, got: ${span.slice(0, 40)}`, + ); + assert.ok(span.includes('def greet')); + }); + + it('emits one chunk for a class plus one per method', () => { + const code = `class AppConfig: + """Config class.""" + def __init__(self, debug=False): + self.debug = debug + + def reload(self): + return True +`; + const result = extract(code); + const names = namesOf(result); + assert.deepEqual(names, ['AppConfig', '__init__', 'reload']); + }); + + it('preserves decorators on methods inside a class', () => { + const code = `class Api: + @property + def status(self): + return self._status + + @staticmethod + def helper(x): + return x * 2 +`; + const result = extract(code); + const status = result.rawChunks.find((c) => c.symbolName === 'status'); + const helper = result.rawChunks.find((c) => c.symbolName === 'helper'); + assert.ok(status && helper); + assert.ok(code.slice(status.startOffset, status.endOffset).includes('@property')); + assert.ok(code.slice(helper.startOffset, helper.endOffset).includes('@staticmethod')); + }); + + it('respects __all__ for module-level assignments', () => { + const code = `__all__ = ["public_thing"] + +public_thing = make_thing( + name="public", + description="exported via __all__, long enough to pass byte threshold", + extra="padding to ensure the RHS exceeds the chunk minimum byte length", +) + +private_thing = make_thing( + name="private", + description="not in __all__, also padded to pass the byte threshold here", + extra="more padding to ensure the RHS exceeds the chunk minimum byte length", +) +`; + const result = extract(code); + const names = namesOf(result); + assert.ok(names.includes('public_thing')); + assert.ok(!names.includes('private_thing')); + }); + + it('falls back to leading-underscore convention when no __all__', () => { + const code = `app = FastAPI( + title="my long application that needs documenting here", + version="1.0.0", + description="A long-enough description for the chunker.", +) + +_private = make_private_thing( + arg="something", + other="value", +) +`; + const result = extract(code); + const names = namesOf(result); + assert.ok(names.includes('app')); + assert.ok(!names.includes('_private')); + }); + + it('emits a module docstring as fileDocPrefix', () => { + const code = `"""Module docstring. + +Spans multiple lines. +""" +import os + +def foo(): + return 1 +`; + const result = extract(code); + assert.ok(result.fileDocPrefix.startsWith('"""Module docstring.')); + assert.ok(result.fileDocPrefix.endsWith('\n\n')); + }); + + it('collects imports into importPrefix', () => { + const code = `import os +from typing import Optional + +def foo(): + return 1 +`; + const result = extract(code); + assert.ok(result.importPrefix.includes('import os')); + assert.ok(result.importPrefix.includes('from typing import Optional')); + assert.ok(result.importPrefix.endsWith('\n\n')); + }); + + it('survives parse errors and emits partial chunks with a warning', () => { + const code = `def foo(): + return 1 + +def broken( + # missing closing paren — syntax error + +def bar(): + return 2 +`; + const result = extract(code); + assert.ok(result.parseWarnings.length > 0, 'expected parse warning'); + const names = namesOf(result); + assert.ok(names.includes('foo')); + assert.ok(names.includes('bar')); + }); + + it('does not chunk dunders even without __all__', () => { + const code = `__version__ = "1.0" +__author__ = "someone" + +real_thing = make_real_thing( + arg="value-padded-out-to-clear-the-eighty-byte-minimum", + other="another value also padded for the byte threshold", +) +`; + const result = extract(code); + const names = namesOf(result); + assert.ok(!names.includes('__version__')); + assert.ok(!names.includes('__author__')); + assert.ok(names.includes('real_thing')); + }); + + it('skips short module-level call assignments', () => { + const code = `x = make(1) +`; + const result = extract(code); + assert.equal(result.rawChunks.length, 0); + }); + + it('does not chunk integer/string-literal assignments (not calls)', () => { + const code = `MAX_RETRIES = 5 +DEFAULT_NAME = "alice" +`; + const result = extract(code); + assert.equal(result.rawChunks.length, 0); + }); + + it('chunks an async function definition', () => { + const code = `async def fetch_data(url): + async with httpx.AsyncClient() as client: + return await client.get(url) +`; + const result = extract(code); + // tree-sitter-python represents async funcs as function_definition with an `async` modifier, + // so this should produce one chunk named `fetch_data`. + assert.equal(result.rawChunks.length, 1); + assert.equal(result.rawChunks[0].symbolName, 'fetch_data'); + }); +});