diff --git a/README.md b/README.md index e6d86006..236ab4af 100644 --- a/README.md +++ b/README.md @@ -50,9 +50,59 @@ The install script downloads the correct tarball for your platform, extracts it SmallCode includes [BoneScript](https://github.com/Doorman11991/BoneScript) and [budget-aware-mcp](https://github.com/Doorman11991/budget-aware-mcp) as dependencies — everything installs in one go. +### Fresh GitHub checkout quick start + +If you just cloned/pulled this repository, run it directly from the checkout first: + +```bash +cd smallcode +npm install + +# Start your local model server first (LM Studio, llama.cpp, Ollama, etc.) +cat > .env <<'EOF' +SMALLCODE_MODEL=your-local-model-name +SMALLCODE_BASE_URL=http://localhost:1234/v1 +EOF + +node bin/smallcode.js +``` + +Optional: make the `smallcode` and `smallcode-rag-index` commands available globally from this checkout: + +```bash +npm link +smallcode --help +``` + +If the fullscreen UI has display issues in your terminal, start with `node bin/smallcode.js --classic`. + +### RAG harness quick run + +SmallCode runs as a terminal UI harness by default: + +```bash +smallcode # fullscreen TUI +smallcode --classic # readline UI fallback +node bin/smallcode.js # from a repo checkout +``` + +To build the local GitHub RAG database, run the Python scraper/indexer with the curated starter corpus, or use the broader preset for a larger multi-language corpus: + +```bash +npm run rag:index +npm run rag:index -- --preset broad +# or, after install: +smallcode-rag-index --preset broad +``` + +For custom repos, create `.smallcode/rag/repos.json` with `preset`, `repos`, and chunking limits. + +See [docs/rag-harness.md](docs/rag-harness.md) for the full LM Studio/llama.cpp setup, UI walkthrough, RAG config, indexing, and web-fallback flow. + ### Requirements - Node.js 18+ (LTS recommended — 20.x or 22.x have prebuilt binaries for SQLite) +- Python 3 + Git for the RAG scraper/indexer (`npm run rag:index`) - A local LLM server (LM Studio, Ollama, or any OpenAI-compatible endpoint) **Optional** (for code graph + FTS5 memory search): diff --git a/bin/rag-index.js b/bin/rag-index.js new file mode 100755 index 00000000..814b3883 --- /dev/null +++ b/bin/rag-index.js @@ -0,0 +1,110 @@ +#!/usr/bin/env node +'use strict'; + +const fs = require('fs'); +const os = require('os'); +const path = require('path'); +const cp = require('child_process'); +const readline = require('readline'); +const { RagIndexStore } = require('../src/rag/index_store'); + +function usage() { + console.log(`Usage: smallcode-rag-index [--config PATH] [--preset starter|broad|none] [--repo URL_OR_PATH] + +Builds .smallcode/rag/index.json from snippet-sized code chunks scraped by scripts/rag_scraper.py. + +Config example (.smallcode/rag/repos.json): +{ + "preset": "starter", + "repos": ["https://github.com/owner/repo.git"], + "maxFilesPerRepo": 1000, + "maxSnippetsPerRepo": 4000, + "chunkLines": 80 +} + +Presets: + starter curated multi-language framework/library set (default when no config exists) + broad larger curated corpus across Python, JS/TS, Go, Rust, Java, C#, Ruby, PHP, C/C++ + none only scrape repos listed in config or --repo +`); +} + +function parseArgs(argv) { + const args = { repos: [] }; + for (let i = 0; i < argv.length; i++) { + const a = argv[i]; + if (a === '--help' || a === '-h') args.help = true; + else if (a === '--config') args.config = argv[++i]; + else if (a === '--preset') args.preset = argv[++i]; + else if (a === '--repo') args.repos.push(argv[++i]); + else if (a === '--index-path') args.indexPath = argv[++i]; + else throw new Error(`Unknown argument: ${a}`); + } + return args; +} + +function readConfig(configPath) { + if (!fs.existsSync(configPath)) return {}; + return JSON.parse(fs.readFileSync(configPath, 'utf-8')); +} + +function findPython() { + for (const bin of ['python3', 'python']) { + const r = cp.spawnSync(bin, ['--version'], { encoding: 'utf-8' }); + if (r.status === 0) return bin; + } + throw new Error('Python 3 is required for the RAG scraper (tried python3 and python).'); +} + +async function loadJsonl(file, onRecord) { + const rl = readline.createInterface({ input: fs.createReadStream(file), crlfDelay: Infinity }); + let count = 0; + for await (const line of rl) { + if (!line.trim()) continue; + onRecord(JSON.parse(line)); + count++; + } + return count; +} + +async function main() { + const args = parseArgs(process.argv.slice(2)); + if (args.help) { usage(); return; } + + const configPath = args.config || process.env.SMALLCODE_RAG_REPOS || path.join(process.cwd(), '.smallcode', 'rag', 'repos.json'); + const cfg = readConfig(configPath); + const indexPath = args.indexPath || cfg.indexPath || path.join(process.cwd(), '.smallcode', 'rag', 'index.json'); + const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'smallcode-rag-')); + const jsonlPath = path.join(tmpDir, 'snippets.jsonl'); + const scraper = path.join(__dirname, '..', 'scripts', 'rag_scraper.py'); + + const pyArgs = [scraper, '--out', jsonlPath]; + if (fs.existsSync(configPath)) pyArgs.push('--config', configPath); + if (args.preset) pyArgs.push('--preset', args.preset); + for (const repo of args.repos) pyArgs.push('--repo', repo); + + const py = findPython(); + const scrape = cp.spawnSync(py, pyArgs, { cwd: process.cwd(), stdio: ['ignore', 'inherit', 'inherit'] }); + if (scrape.status !== 0) process.exit(scrape.status || 1); + + const store = new RagIndexStore({ path: indexPath }); + store.load(); + const batch = []; + let total = 0; + await loadJsonl(jsonlPath, (rec) => { + batch.push(rec); + if (batch.length >= 1000) { + store.upsertMany(batch.splice(0, batch.length)); + } + total++; + }); + if (batch.length) store.upsertMany(batch); + store.save(); + fs.rmSync(tmpDir, { recursive: true, force: true }); + process.stdout.write(`indexed ${total} snippets into ${indexPath}\n`); +} + +main().catch((err) => { + console.error(err.stack || err.message || String(err)); + process.exit(1); +}); diff --git a/bin/smallcode.js b/bin/smallcode.js index 91b688a0..3cdf2317 100755 --- a/bin/smallcode.js +++ b/bin/smallcode.js @@ -245,6 +245,7 @@ let _planTracker = null; let _bootstrapDetector = null; let _testRunnerDetector = null; let _knowledgeLoader = null; +let _ragRetriever = null; const improvementAttempts = {}; // filePath → attempt count async function runTUI(config) { @@ -684,6 +685,12 @@ async function runAgentLoop(userMessage, config) { const { KnowledgeLoader } = require('../src/knowledge/loader'); _knowledgeLoader = new KnowledgeLoader({ rootDir: process.cwd() }); } catch { _knowledgeLoader = null; } + // RAG retriever: local code examples indexed from GitHub repos or other source trees. + try { + const { RagRetriever } = require('../src/rag/retriever'); + _ragRetriever = new RagRetriever(); + _ragRetriever.load(); + } catch { _ragRetriever = null; } // Trust decay (Feature 13) resets per agent loop turn so TUI sessions // don't accumulate decay from unrelated prior requests. try { @@ -1942,7 +1949,7 @@ CRITICAL — large file rule: write_file calls are limited to 60 lines / ~8KB. l } // Legacy behavior: everything in the system prompt - prompt += getMemoryContext(messages) + getSkillContext(messages) + getPluginPrompts() + getKnowledgeContext(messages) + getActivePlanContext() + getTestRunnerContext(); + prompt += getMemoryContext(messages) + getSkillContext(messages) + getPluginPrompts() + getKnowledgeContext(messages) + getRagContext(messages) + getActivePlanContext() + getTestRunnerContext(); return prompt; } @@ -1957,6 +1964,7 @@ function buildDynamicContext(messages) { getMemoryContext(messages), getSkillContext(messages), getKnowledgeContext(messages), + getRagContext(messages), // Note: getPluginPrompts() stays in the system prompt — plugin instructions // are authoritative and should come from the system role. // Note: getActivePlanContext() stays in the system prompt — the model needs @@ -2006,6 +2014,20 @@ function getKnowledgeContext(messages) { } } +// Auto-load similar code snippets from the local RAG index. The index is +// created with `npm run rag:index` / `smallcode-rag-index` and stays fully local. +function getRagContext(messages) { + try { + if (!_ragRetriever) return ''; + const lastUser = [...messages].reverse().find(m => m.role === 'user'); + if (!lastUser || typeof lastUser.content !== 'string') return ''; + const maxChars = Math.min(6000, Math.floor(((config?.context?.detected_window || 32768) * 0.06) * 4)); + return _ragRetriever.formatForPrompt(lastUser.content, { limit: 6, maxChars }); + } catch { + return ''; + } +} + // Auto-load relevant memory for the current task (injected into system prompt) // Fix #15: Only inject memory objects with score >= 2 (at least 2 word matches) // to avoid burning tokens on low-relevance hits. diff --git a/docs/rag-harness.md b/docs/rag-harness.md new file mode 100644 index 00000000..132a766a --- /dev/null +++ b/docs/rag-harness.md @@ -0,0 +1,206 @@ +# RAG Coding Harness Guide + +This repo runs as a terminal UI coding harness. The default interface is the fullscreen TUI; use `--classic` if your terminal does not support alternate-screen rendering. + +## 0. From a fresh GitHub clone + +If you just downloaded or pulled the repo, you do **not** need to publish/install anything first. Run it from the checkout: + +```bash +cd smallcode +npm install +node bin/smallcode.js --help +``` + +Then start a local model server and create a `.env` file in the directory where you will run the harness. For testing against LM Studio's default local server this usually looks like: + +```bash +cat > .env <<'ENV' +SMALLCODE_MODEL=your-loaded-model-name +SMALLCODE_BASE_URL=http://localhost:1234/v1 +ENV +``` + +Start the UI from the checkout with: + +```bash +node bin/smallcode.js +``` + +If you want the normal `smallcode` command instead of typing `node bin/smallcode.js`, link the checkout once: + +```bash +npm link +smallcode +``` + +Use `node bin/smallcode.js --classic` or `smallcode --classic` if your terminal has trouble with the fullscreen interface. + +## 1. Start a local model server + +SmallCode talks to any OpenAI-compatible local endpoint. + +### LM Studio + +1. Open LM Studio. +2. Download/load a coder model such as Qwen Coder or another 8B-35B local coding model. +3. Start the **Local Server**. +4. Note the model name shown by LM Studio and the base URL, usually `http://localhost:1234/v1`. + +Create `.env` in the project you want to edit: + +```bash +SMALLCODE_MODEL=your-loaded-model-name +SMALLCODE_BASE_URL=http://localhost:1234/v1 +``` + +### llama.cpp server + +Start llama.cpp with its OpenAI-compatible server, then point SmallCode at it: + +```bash +SMALLCODE_MODEL=local-model +SMALLCODE_BASE_URL=http://localhost:8080/v1 +``` + +## 2. Run the UI harness + +From the project directory you want the agent to edit: + +```bash +smallcode +``` + +If you are developing from this repository checkout instead of a global install: + +```bash +node bin/smallcode.js +``` + +Useful launch modes: + +```bash +smallcode --classic # readline UI instead of fullscreen UI +smallcode -P "fix the parser bug" # one-shot prompt +smallcode --non-interactive "refactor" # stdin/script friendly mode +smallcode --resume # continue previous session +``` + +Inside the UI: + +- Type your task and press Enter. +- Use `/help` for commands. +- Use `/plan` to inspect the active plan. +- Use `/undo` to revert the last edit. +- Use `/quit` to exit. + +## 3. Create the local GitHub RAG database + +The scraper is a Python pipeline at `scripts/rag_scraper.py`. It shallow-clones or updates repositories, walks source files, and emits **snippet-sized chunks** around functions/classes/types plus sliding-window chunks for files without clear symbols. It does not index whole files as a single blob. + +### Fast starter corpus + +If you do not create a config file, the indexer uses the built-in `starter` preset: a curated multi-language set of popular frameworks/libraries across Python, JavaScript/TypeScript, Go, Rust, Java, C#, Ruby, PHP, and C/C++. + +```bash +npm run rag:index +``` + +### Larger broad corpus + +For a bigger language-modeling corpus, use the `broad` preset. This scrapes more well-known, high-signal codebases, but takes longer and uses more disk space. + +```bash +npm run rag:index -- --preset broad +``` + +The curated presets live in `src/rag/curated_repos.json`, so you can review or change the selected repositories. + +### Custom corpus + +Create `.smallcode/rag/repos.json` in the workspace where you run SmallCode: + +```json +{ + "preset": "starter", + "repos": [ + "https://github.com/owner/framework-example.git", + "https://github.com/owner/language-examples.git", + { "url": "/absolute/path/to/local/repo", "tags": ["local", "examples"] } + ], + "maxFilesPerRepo": 1000, + "maxSnippetsPerRepo": 4000, + "chunkLines": 80, + "overlap": 20 +} +``` + +Set `"preset": "none"` if you only want your own repos. + +Optional fields: + +```json +{ + "cacheDir": ".smallcode/rag/repos", + "indexPath": ".smallcode/rag/index.json", + "languages": ["python", "typescript", "go"], + "maxFileBytes": 250000, + "minChars": 120, + "repos": ["https://github.com/owner/repo.git"] +} +``` + +After package installation, the same command is also available as: + +```bash +smallcode-rag-index --preset broad +``` + +The indexer saves `.smallcode/rag/index.json` by default. + +## 4. How code search works + +SmallCode searches **code snippets**, not full files. The pipeline stores each snippet with repo, language, path, symbol name, start/end lines, tags, term frequencies, and a sparse local hashed vector. + +At query time the retriever runs a hybrid search: + +1. **BM25 lexical search** over identifiers, paths, symbols, tags, and snippet text. This is strong for exact APIs, framework names, error names, and language constructs. +2. **Local hashed-vector similarity** over the same snippet text. This is dependency-free and helps related naming patterns match even when exact words differ. +3. The final rank combines BM25 and vector scores, then injects only the top bounded snippets into model context. + +This approach is intentionally fast enough for local models and large local corpora without requiring a separate vector database or cloud embeddings. + +## 5. Use RAG in the harness + +Once `.smallcode/rag/index.json` exists, start the UI normally: + +```bash +smallcode +``` + +For each user turn, SmallCode now: + +1. plans/classifies the request, +2. retrieves similar snippets from the local RAG index, +3. injects the best snippets into the model context, +4. asks the local model to do one step at a time through the normal tool loop. + +No cloud embedding service is required. The default embedding path is dependency-free and optimized for fast local startup. + +## 6. Optional web fallback when RAG is weak + +Web search is disabled by default. Enable it only when you want the model to search externally after local RAG confidence is low: + +```bash +SMALLCODE_WEB_BROWSE=true smallcode +``` + +When enabled, low-confidence RAG context tells the model to use `web_search` with a GitHub/code-example query if it gets blocked. + +## 7. Speed tips for local models + +- Prefer the fullscreen UI (`smallcode`) for normal work; use `--classic` only for terminal compatibility issues. +- Keep `SMALLCODE_CACHE_SPLIT` at its default (`true`) so llama.cpp-style KV cache reuse is not invalidated by dynamic context. +- Keep RAG repos focused when you need fast indexing; use `--preset broad` when you intentionally want a large reference corpus. +- Use 8B-35B coder models; very small models often fail multi-step tool use. +- If the model struggles, ask for a smaller concrete task first, then continue with follow-up prompts. diff --git a/package-lock.json b/package-lock.json index 7012edc7..8c11189b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "smallcode", - "version": "0.9.6", + "version": "1.3.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "smallcode", - "version": "0.9.6", + "version": "1.3.0", "license": "MIT", "dependencies": { "bonescript-compiler": "0.14.0", @@ -18,7 +18,8 @@ }, "bin": { "smallcode": "bin/smallcode.js", - "smallcode-init": "bin/init.js" + "smallcode-init": "bin/init.js", + "smallcode-rag-index": "bin/rag-index.js" }, "devDependencies": { "@types/node": "^25.9.0", @@ -30,9 +31,19 @@ "node": ">=18.0.0" }, "optionalDependencies": { - "budget-aware-mcp": "0.6.1", - "playwright-extra": "4.3.0", - "puppeteer-extra-plugin-stealth": "2.11.0" + "budget-aware-mcp": "0.6.1" + }, + "peerDependencies": { + "playwright-extra": "^4.3.0", + "puppeteer-extra-plugin-stealth": "^2.11.0" + }, + "peerDependenciesMeta": { + "playwright-extra": { + "optional": true + }, + "puppeteer-extra-plugin-stealth": { + "optional": true + } } }, "node_modules/@babel/code-frame": { @@ -1226,6 +1237,7 @@ "integrity": "sha512-KSVgmQmzMwPlmtljOomayoR89W4FynCAi3E8PPs7vmDVPe84hT+vGPKkJfThkmXs0x0jAaa9U8uW8bbfyS2fWw==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "@types/ms": "*" } @@ -1262,7 +1274,8 @@ "resolved": "https://registry.npmjs.org/@types/ms/-/ms-2.1.0.tgz", "integrity": "sha512-GsCCIZDE/p3i96vtEqx+7dBUGXrc7zeSK3wwPHIaRThS+9OhWIXRqzs4d6k1SVU8g91DrNRWxWUGhp5KXQb2VA==", "license": "MIT", - "optional": true + "optional": true, + "peer": true }, "node_modules/@types/node": { "version": "25.9.0", @@ -1747,6 +1760,7 @@ "integrity": "sha512-sKpyeERZ02v1FeCZT8lrfJq5u6goHCtpTAzPwJYe7c8SPFOboNjNg1vz2L4VTn9T4PQxEx13TbXLmYUcS6Ug7Q==", "license": "MIT", "optional": true, + "peer": true, "engines": { "node": ">=0.10.0" } @@ -2257,6 +2271,7 @@ "integrity": "sha512-we+NuQo2DHhSl+DP6jlUiAhyAjBQrYnpOk15rN6c6JSPScjiCLh8IbSU+VTcph6YS3o7mASE8a0+gbZ7ChLpgg==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "for-own": "^0.1.3", "is-plain-object": "^2.0.1", @@ -2912,6 +2927,7 @@ "integrity": "sha512-7EwmXrOjyL+ChxMhmG5lnW9MPt1aIeZEwKhQzoBUdTV0N3zuwWDZYVJatDvZ2OyzPUvdIAZDsCetk3coyMfcnQ==", "license": "MIT", "optional": true, + "peer": true, "engines": { "node": ">=0.10.0" } @@ -2922,6 +2938,7 @@ "integrity": "sha512-SKmowqGTJoPzLO1T0BBJpkfp3EMacCMOuH40hOUbrbzElVktk4DioXVM99QkLCyKoiuOmyjgcWMpVz2xjE7LZw==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "for-in": "^1.0.1" }, @@ -2960,6 +2977,7 @@ "integrity": "sha512-oRXApq54ETRj4eMiFzGnHWGy+zo5raudjuxN0b8H7s/RU2oW0Wvsx9O0ACRN/kRq9E8Vu/ReskGB5o3ji+FzHQ==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "graceful-fs": "^4.2.0", "jsonfile": "^6.0.1", @@ -3316,7 +3334,8 @@ "resolved": "https://registry.npmjs.org/is-buffer/-/is-buffer-1.1.6.tgz", "integrity": "sha512-NcdALwpXkTm5Zvvbk7owOUSvVvBKDgKP5/ewfXEznmQFfs4ZRmanOeKBTjRVjka3QFoN6XJ+9F3USqfHqTaU5w==", "license": "MIT", - "optional": true + "optional": true, + "peer": true }, "node_modules/is-extendable": { "version": "0.1.1", @@ -3324,6 +3343,7 @@ "integrity": "sha512-5BMULNob1vgFX6EjQw5izWDxrecWK9AM72rugNr0TFldMOi0fj6Jk+zeKIt0xGj4cEfQIJth4w3OKWOJ4f+AFw==", "license": "MIT", "optional": true, + "peer": true, "engines": { "node": ">=0.10.0" } @@ -3353,6 +3373,7 @@ "integrity": "sha512-h5PpgXkWitc38BBMYawTYMWJHFZJVnBquFE57xFpjB8pJFiF6gZ+bU+WyI/yqXiFR5mdLsgYNaPe8uao6Uv9Og==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "isobject": "^3.0.1" }, @@ -3392,6 +3413,7 @@ "integrity": "sha512-WhB9zCku7EGTj/HQQRz5aUQEUeoQZH2bWcltRErOpymJ4boYE6wL9Tbr23krRPSZ+C5zqNSrSw+Cc7sZZ4b7vg==", "license": "MIT", "optional": true, + "peer": true, "engines": { "node": ">=0.10.0" } @@ -4201,6 +4223,7 @@ "integrity": "sha512-zwOTdL3rFQ/lRdBnntKVOX6k5cKJwEc1HdilT71BWEu7J41gXIB2MRp+vxduPSwZJPWBxEzv4yH1wYLJGUHX4Q==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "universalify": "^2.0.0" }, @@ -4214,6 +4237,7 @@ "integrity": "sha512-NOW9QQXMoZGg/oqnVNoNTTIFEIid1627WCffUBJEdMxYApq7mNE7CpzucIPc+ZQg25Phej7IJSmX3hO+oblOtQ==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "is-buffer": "^1.1.5" }, @@ -4227,6 +4251,7 @@ "integrity": "sha512-RE2g0b5VGZsOCFOCgP7omTRYFqydmZkBwl5oNnQ1lDYC57uyO9KqNnNVxT7COSHTxrRCWVcAVOcbjk+tvh/rgQ==", "license": "MIT", "optional": true, + "peer": true, "engines": { "node": ">=0.10.0" } @@ -4401,6 +4426,7 @@ "integrity": "sha512-qtmzAS6t6grwEkNrunqTBdn0qKwFgNWvlxUbAV8es9M7Ot1EbyApytCnvE0jALPa46ZpKDUo527kKiaWplmlFA==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "arr-union": "^3.1.0", "clone-deep": "^0.2.4", @@ -4519,6 +4545,7 @@ "integrity": "sha512-ALGF1Jt9ouehcaXaHhn6t1yGWRqGaHkPFndtFVHfZXOvkIZ/yoGaSi0AHVTafb3ZBGg4dr/bDwnaEKqCXzchMA==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "for-in": "^0.1.3", "is-extendable": "^0.1.1" @@ -4533,6 +4560,7 @@ "integrity": "sha512-F0to7vbBSHP8E3l6dCjxNOLuSFAACIxFy3UehTUlG7svlXi37HHsDkyVcHo0Pq8QwrE+pXvWSVX3ZT1T9wAZ9g==", "license": "MIT", "optional": true, + "peer": true, "engines": { "node": ">=0.10.0" } @@ -4796,7 +4824,8 @@ "resolved": "https://registry.npmjs.org/package-json-from-dist/-/package-json-from-dist-1.0.1.tgz", "integrity": "sha512-UEZIS3/by4OC8vL3P2dTXRETpebLI2NiI5vIrjaD/5UtrkFX/tNbwjTSRAGC/+7CAo2pIcBaRgWmcBBHcsaCIw==", "license": "BlueOak-1.0.0", - "optional": true + "optional": true, + "peer": true }, "node_modules/parse-json": { "version": "5.2.0", @@ -4963,6 +4992,7 @@ "integrity": "sha512-/Hec3BmYMY/GznBo0ZPsSqm6IHil7jInxLz+9/UnBBC5Ozh2abVNnv+vYNJ+JKKRVtiKyHCrJpKRibJ9uribZw==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "debug": "^4.3.4" }, @@ -5057,6 +5087,7 @@ "integrity": "sha512-6RNy0e6pH8vaS3akPIKGg28xcryKscczt4wIl0ePciZENGE2yoaQJNd17UiEbdmh5/6WW6dPcfRWT9lxBwCi2Q==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "@types/debug": "^4.1.0", "debug": "^4.1.1", @@ -5084,6 +5115,7 @@ "integrity": "sha512-BqckPV95MHP25quZgzBnZJD8S38ZYP4B3HJ3Kr/vibqxJxhK6L1VQ6jnu/JcFKV0wzCIQPrCiiavZnwE5u1C2A==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "debug": "^4.1.1", "puppeteer-extra-plugin": "^3.2.2", @@ -5111,6 +5143,7 @@ "integrity": "sha512-kH1GnCcqEDoBXO7epAse4TBPJh9tEpVEK/vkedKfjOVOhZAvLkHGc9swMs5ChrJbRnf8Hdpug6TJlEuimXNQ+g==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "debug": "^4.1.1", "fs-extra": "^10.0.0", @@ -5139,6 +5172,7 @@ "integrity": "sha512-i1oAZxRbc1bk8MZufKCruCEC3CCafO9RKMkkodZltI4OqibLFXF3tj6HZ4LZ9C5vCXZjYcDWazgtY69mnmrQ9A==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "debug": "^4.1.1", "deepmerge": "^4.2.2", @@ -5322,6 +5356,7 @@ "integrity": "sha512-LKg+Cr2ZF61fkcaK1UdkH2yEBBKnYjTyWzTJT6KNPcSPaiT7HSdhtMXQuN5wkTX0Xu72KQ1l8S42rlmexS2hSA==", "license": "BlueOak-1.0.0", "optional": true, + "peer": true, "dependencies": { "glob": "^13.0.3", "package-json-from-dist": "^1.0.1" @@ -5446,6 +5481,7 @@ "integrity": "sha512-J1zdXCky5GmNnuauESROVu31MQSnLoYvlyEn6j2Ztk6Q5EHFIhxkMhYcv6vuDzl2XEzoRr856QwzMgWM/TmZgw==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "is-extendable": "^0.1.1", "kind-of": "^2.0.1", @@ -5462,6 +5498,7 @@ "integrity": "sha512-0u8i1NZ/mg0b+W3MGGw5I7+6Eib2nx72S/QvXa0hYjEkjTknYmEYQJwGu3mLC0BrhtJjtQafTkyRUQ75Kx0LVg==", "license": "MIT", "optional": true, + "peer": true, "dependencies": { "is-buffer": "^1.0.2" }, @@ -5475,6 +5512,7 @@ "integrity": "sha512-gkX52wvU/R8DVMMt78ATVPFMJqfW8FPz1GZ1sVHBVQHmu/WvhIWE4cE1GBzhJNFicDeYhnwp6Rl35BcAIM3YOQ==", "license": "MIT", "optional": true, + "peer": true, "engines": { "node": ">=0.10.0" } @@ -6104,6 +6142,7 @@ "integrity": "sha512-gptHNQghINnc/vTGIk0SOFGFNXw7JVrlRUtConJRlvaw6DuX0wO5Jeko9sWrMBhh+PsYAZ7oXAiOnf/UKogyiw==", "license": "MIT", "optional": true, + "peer": true, "engines": { "node": ">= 10.0.0" } diff --git a/package.json b/package.json index d9ca38d9..a31023de 100644 --- a/package.json +++ b/package.json @@ -5,7 +5,8 @@ "main": "src/api/index.js", "bin": { "smallcode": "bin/smallcode.js", - "smallcode-init": "bin/init.js" + "smallcode-init": "bin/init.js", + "smallcode-rag-index": "bin/rag-index.js" }, "scripts": { "build": "node build.js", @@ -21,7 +22,8 @@ "bench:diff": "node bench/diff.js", "test": "node --test test/*.test.js", "test:e2e": "node test/e2e_smoke.js", - "test:e2e:offline": "node test/e2e_offline.js" + "test:e2e:offline": "node test/e2e_offline.js", + "rag:index": "node bin/rag-index.js" }, "dependencies": { "bonescript-compiler": "0.14.0", @@ -39,8 +41,12 @@ "puppeteer-extra-plugin-stealth": "^2.11.0" }, "peerDependenciesMeta": { - "playwright-extra": { "optional": true }, - "puppeteer-extra-plugin-stealth": { "optional": true } + "playwright-extra": { + "optional": true + }, + "puppeteer-extra-plugin-stealth": { + "optional": true + } }, "keywords": [ "ai", diff --git a/scripts/rag_scraper.py b/scripts/rag_scraper.py new file mode 100755 index 00000000..3c6f0b6c --- /dev/null +++ b/scripts/rag_scraper.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +"""SmallCode RAG scraper. + +Clones/pulls curated or user-provided repositories, extracts symbol-sized code +snippets (not whole files), and emits JSONL records consumed by bin/rag-index.js. +The implementation is intentionally dependency-free so a fresh checkout can run +it with only Python 3 and git. +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shutil +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, Iterable, Iterator, List, Optional, Sequence, Tuple + +ROOT = Path(__file__).resolve().parents[1] +CURATED_PATH = ROOT / "src" / "rag" / "curated_repos.json" + +SKIP_DIRS = { + ".git", ".hg", ".svn", "node_modules", "vendor", "dist", "build", "target", + "coverage", ".next", ".nuxt", ".venv", "venv", "__pycache__", ".pytest_cache", + "bin", "obj", ".gradle", ".idea", ".vscode", "Pods", ".dart_tool", +} + +LANG_BY_EXT = { + ".py": "python", ".js": "javascript", ".jsx": "javascript", ".mjs": "javascript", + ".cjs": "javascript", ".ts": "typescript", ".tsx": "typescript", ".go": "go", + ".rs": "rust", ".java": "java", ".kt": "kotlin", ".kts": "kotlin", + ".cs": "csharp", ".php": "php", ".rb": "ruby", ".c": "c", ".h": "c", + ".cc": "cpp", ".cpp": "cpp", ".cxx": "cpp", ".hpp": "cpp", ".hh": "cpp", + ".swift": "swift", ".dart": "dart", ".scala": "scala", ".clj": "clojure", + ".ex": "elixir", ".exs": "elixir", ".erl": "erlang", ".hrl": "erlang", + ".lua": "lua", ".sh": "shell", ".bash": "shell", ".zsh": "shell", +} + +SYMBOL_PATTERNS: Dict[str, re.Pattern[str]] = { + "python": re.compile(r"^\s*(?:async\s+def|def|class)\s+([A-Za-z_][\w]*)"), + "javascript": re.compile(r"^\s*(?:export\s+)?(?:async\s+)?(?:function\s+([A-Za-z_$][\w$]*)|class\s+([A-Za-z_$][\w$]*)|const\s+([A-Za-z_$][\w$]*)\s*=\s*(?:async\s*)?(?:\(|function))"), + "typescript": re.compile(r"^\s*(?:export\s+)?(?:async\s+)?(?:function\s+([A-Za-z_$][\w$]*)|class\s+([A-Za-z_$][\w$]*)|interface\s+([A-Za-z_$][\w$]*)|type\s+([A-Za-z_$][\w$]*)|const\s+([A-Za-z_$][\w$]*)\s*=\s*(?:async\s*)?(?:\(|function))"), + "go": re.compile(r"^\s*(?:func\s+(?:\([^)]*\)\s*)?([A-Za-z_][\w]*)|type\s+([A-Za-z_][\w]*)\s+(?:struct|interface))"), + "rust": re.compile(r"^\s*(?:pub\s+)?(?:async\s+)?(?:fn\s+([A-Za-z_][\w]*)|struct\s+([A-Za-z_][\w]*)|enum\s+([A-Za-z_][\w]*)|trait\s+([A-Za-z_][\w]*)|impl\b)"), + "java": re.compile(r"^\s*(?:public|private|protected|static|final|abstract|\s)*\s*(?:class|interface|enum|record)\s+([A-Za-z_][\w]*)|^\s*(?:public|private|protected|static|final|synchronized|abstract|native|\s)+[\w<>\[\], ?]+\s+([A-Za-z_][\w]*)\s*\("), + "csharp": re.compile(r"^\s*(?:public|private|protected|internal|static|sealed|abstract|partial|async|\s)*\s*(?:class|interface|enum|record|struct)\s+([A-Za-z_][\w]*)|^\s*(?:public|private|protected|internal|static|async|override|virtual|\s)+[\w<>\[\], ?]+\s+([A-Za-z_][\w]*)\s*\("), + "php": re.compile(r"^\s*(?:final\s+|abstract\s+)?(?:class|interface|trait|enum)\s+([A-Za-z_][\w]*)|^\s*(?:public|private|protected|static|final|abstract|\s)*function\s+([A-Za-z_][\w]*)"), + "ruby": re.compile(r"^\s*(?:class|module|def)\s+([A-Za-z_][\w:!?=]*)"), + "cpp": re.compile(r"^\s*(?:template\s*<[^>]+>\s*)?(?:class|struct|enum)\s+([A-Za-z_][\w]*)|^\s*(?:inline\s+|static\s+|constexpr\s+|virtual\s+)*[\w:<>&*\s]+\s+([A-Za-z_][\w]*)\s*\([^;]*\)\s*(?:const\s*)?(?:\{|$)"), + "c": re.compile(r"^\s*(?:static\s+|inline\s+)*[\w\s\*]+\s+([A-Za-z_][\w]*)\s*\([^;]*\)\s*(?:\{|$)"), +} + +DEFAULT_SYMBOL_PATTERN = re.compile(r"^\s*(?:function|class|def|func|fn)\s+([A-Za-z_][\w]*)") + +@dataclass +class RepoSpec: + url: str + tags: List[str] + name: Optional[str] = None + max_files: Optional[int] = None + max_snippets: Optional[int] = None + + +def run_git(args: Sequence[str], cwd: Optional[Path] = None) -> None: + subprocess.run(["git", *args], cwd=str(cwd) if cwd else None, check=True, stdout=subprocess.DEVNULL) + + +def safe_repo_name(url: str) -> str: + if os.path.exists(url): + return Path(url).resolve().name + clean = url.rstrip("/").removesuffix(".git") + parts = clean.split("/")[-2:] + return "__".join(parts) if len(parts) == 2 else hashlib.sha1(url.encode()).hexdigest()[:12] + + +def ensure_repo(spec: RepoSpec, cache_dir: Path) -> Path: + if os.path.isdir(spec.url) and not spec.url.startswith(("http://", "https://", "git@")): + return Path(spec.url).resolve() + cache_dir.mkdir(parents=True, exist_ok=True) + dest = cache_dir / safe_repo_name(spec.url) + if (dest / ".git").exists(): + run_git(["pull", "--ff-only"], dest) + else: + run_git(["clone", "--depth=1", "--filter=blob:none", spec.url, str(dest)], cache_dir) + return dest + + +def iter_code_files(root: Path, max_bytes: int, languages: Optional[set[str]]) -> Iterator[Tuple[Path, str]]: + for dirpath, dirnames, filenames in os.walk(root): + dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS and not d.startswith(".")] + for filename in filenames: + path = Path(dirpath) / filename + lang = LANG_BY_EXT.get(path.suffix.lower()) + if not lang or (languages and lang not in languages): + continue + try: + if path.stat().st_size > max_bytes: + continue + except OSError: + continue + yield path, lang + + +def symbol_name(match: re.Match[str]) -> str: + for group in match.groups(): + if group: + return group + return "symbol" + + +def find_symbols(lines: List[str], lang: str) -> List[Tuple[int, str]]: + pat = SYMBOL_PATTERNS.get(lang, DEFAULT_SYMBOL_PATTERN) + out = [] + for idx, line in enumerate(lines): + m = pat.search(line) + if m: + out.append((idx, symbol_name(m))) + return out + + +def trim_chunk(lines: List[str], max_lines: int) -> List[str]: + if len(lines) <= max_lines: + return lines + return lines[:max_lines] + + +def sliding_chunks(lines: List[str], chunk_lines: int, overlap: int) -> Iterator[Tuple[int, List[str]]]: + step = max(1, chunk_lines - overlap) + for start in range(0, len(lines), step): + chunk = lines[start:start + chunk_lines] + if len("\n".join(chunk).strip()) >= 80: + yield start, chunk + + +def chunk_file(text: str, lang: str, chunk_lines: int, overlap: int, min_chars: int) -> Iterator[Dict[str, object]]: + lines = text.splitlines() + symbols = find_symbols(lines, lang) + emitted = 0 + used_ranges: List[Tuple[int, int]] = [] + + for pos, (start, name) in enumerate(symbols): + next_start = symbols[pos + 1][0] if pos + 1 < len(symbols) else len(lines) + end = min(next_start, start + chunk_lines) + chunk = trim_chunk(lines[start:end], chunk_lines) + code = "\n".join(chunk).strip() + if len(code) < min_chars: + continue + used_ranges.append((start, end)) + emitted += 1 + yield { + "kind": "symbol", + "symbol": name, + "startLine": start + 1, + "endLine": start + len(chunk), + "code": code, + } + + # Fallback/window chunks keep files without obvious symbols useful, while + # still indexing snippets rather than complete files. + if emitted == 0: + for start, chunk in sliding_chunks(lines, chunk_lines, overlap): + code = "\n".join(chunk).strip() + if len(code) >= min_chars: + yield {"kind": "window", "symbol": None, "startLine": start + 1, "endLine": start + len(chunk), "code": code} + + +def snippet_id(repo: str, rel: str, start: int, code: str) -> str: + return hashlib.sha1(f"{repo}:{rel}:{start}:{code}".encode("utf-8", "ignore")).hexdigest() + + +def scrape_repo(spec: RepoSpec, repo_root: Path, options: argparse.Namespace) -> Iterator[Dict[str, object]]: + max_files = spec.max_files or options.max_files_per_repo + max_snippets = spec.max_snippets or options.max_snippets_per_repo + languages = set(options.languages.split(",")) if options.languages else None + file_count = 0 + snippet_count = 0 + for file_path, lang in iter_code_files(repo_root, options.max_file_bytes, languages): + if file_count >= max_files or snippet_count >= max_snippets: + break + file_count += 1 + try: + text = file_path.read_text(encoding="utf-8", errors="ignore") + except OSError: + continue + rel = str(file_path.relative_to(repo_root)).replace(os.sep, "/") + for chunk in chunk_file(text, lang, options.chunk_lines, options.overlap, options.min_chars): + if snippet_count >= max_snippets: + break + code = str(chunk["code"]) + record = { + "id": snippet_id(spec.url, rel, int(chunk["startLine"]), code), + "repo": spec.url, + "repoName": spec.name or safe_repo_name(spec.url), + "tags": spec.tags, + "path": rel, + "lang": lang, + **chunk, + } + snippet_count += 1 + yield record + + +def load_curated(preset: str) -> List[RepoSpec]: + data = json.loads(CURATED_PATH.read_text(encoding="utf-8")) + items = data.get(preset) + if not isinstance(items, list): + raise SystemExit(f"Unknown RAG preset '{preset}'. Available: {', '.join(sorted(data))}") + return [RepoSpec(url=i["url"], tags=list(i.get("tags", [])), name=i.get("name")) for i in items] + + +def load_config(path: Optional[Path], preset_override: Optional[str]) -> Tuple[Dict[str, object], List[RepoSpec]]: + cfg: Dict[str, object] = {} + if path and path.exists(): + cfg = json.loads(path.read_text(encoding="utf-8")) + preset = preset_override if preset_override is not None else str(cfg.get("preset", "none" if cfg.get("repos") else "starter")) + specs: List[RepoSpec] = [] + if preset and preset != "none": + specs.extend(load_curated(preset)) + for item in cfg.get("repos", []) or []: + if isinstance(item, str): + specs.append(RepoSpec(url=item, tags=[])) + elif isinstance(item, dict) and item.get("url"): + specs.append(RepoSpec( + url=str(item["url"]), + tags=list(item.get("tags", [])), + name=item.get("name"), + max_files=item.get("maxFiles"), + max_snippets=item.get("maxSnippets"), + )) + # Stable de-dupe by URL. + seen = set() + unique = [] + for spec in specs: + if spec.url in seen: + continue + seen.add(spec.url) + unique.append(spec) + return cfg, unique + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + ap = argparse.ArgumentParser(description="Scrape curated GitHub/local repos into SmallCode RAG JSONL snippets.") + ap.add_argument("--config", type=Path, default=None, help="Path to .smallcode/rag/repos.json") + ap.add_argument("--cache-dir", type=Path, default=Path(".smallcode/rag/repos")) + ap.add_argument("--out", type=Path, required=True, help="Output JSONL file") + ap.add_argument("--preset", default=None, help="Curated preset: starter, broad, or none") + ap.add_argument("--repo", action="append", default=[], help="Extra Git URL or local path to scrape") + ap.add_argument("--languages", default="", help="Optional comma-separated language allowlist") + ap.add_argument("--max-files-per-repo", type=int, default=1000) + ap.add_argument("--max-snippets-per-repo", type=int, default=4000) + ap.add_argument("--max-file-bytes", type=int, default=250_000) + ap.add_argument("--chunk-lines", type=int, default=80) + ap.add_argument("--overlap", type=int, default=20) + ap.add_argument("--min-chars", type=int, default=120) + return ap.parse_args(argv) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + cfg, specs = load_config(args.config, args.preset) + if cfg: + args.cache_dir = Path(str(cfg.get("cacheDir", args.cache_dir))) + args.max_files_per_repo = int(cfg.get("maxFilesPerRepo", args.max_files_per_repo)) + args.max_snippets_per_repo = int(cfg.get("maxSnippetsPerRepo", args.max_snippets_per_repo)) + args.max_file_bytes = int(cfg.get("maxFileBytes", args.max_file_bytes)) + args.chunk_lines = int(cfg.get("chunkLines", args.chunk_lines)) + args.overlap = int(cfg.get("overlap", args.overlap)) + args.min_chars = int(cfg.get("minChars", args.min_chars)) + if cfg.get("languages") and not args.languages: + args.languages = ",".join(cfg.get("languages")) if isinstance(cfg.get("languages"), list) else str(cfg.get("languages")) + for repo in args.repo: + specs.append(RepoSpec(url=repo, tags=[])) + if not specs: + print("No repositories configured. Use --preset starter/broad or add repos to repos.json.", file=sys.stderr) + return 2 + + args.out.parent.mkdir(parents=True, exist_ok=True) + total = 0 + with args.out.open("w", encoding="utf-8") as fh: + for spec in specs: + print(f"scraping {spec.url}", file=sys.stderr) + try: + repo_root = ensure_repo(spec, args.cache_dir) + count = 0 + for rec in scrape_repo(spec, repo_root, args): + fh.write(json.dumps(rec, ensure_ascii=False) + "\n") + count += 1 + total += count + print(f" snippets: {count}", file=sys.stderr) + except (subprocess.CalledProcessError, OSError) as exc: + print(f" skipped: {exc}", file=sys.stderr) + print(json.dumps({"snippets": total, "out": str(args.out)})) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/rag/curated_repos.json b/src/rag/curated_repos.json new file mode 100644 index 00000000..dd84069d --- /dev/null +++ b/src/rag/curated_repos.json @@ -0,0 +1,47 @@ +{ + "starter": [ + { "url": "https://github.com/pallets/flask.git", "tags": ["python", "web"] }, + { "url": "https://github.com/fastapi/fastapi.git", "tags": ["python", "api"] }, + { "url": "https://github.com/expressjs/express.git", "tags": ["javascript", "web"] }, + { "url": "https://github.com/vuejs/core.git", "tags": ["typescript", "frontend"] }, + { "url": "https://github.com/gin-gonic/gin.git", "tags": ["go", "web"] }, + { "url": "https://github.com/tokio-rs/tokio.git", "tags": ["rust", "async"] }, + { "url": "https://github.com/spring-projects/spring-petclinic.git", "tags": ["java", "spring"] }, + { "url": "https://github.com/dotnet/aspnetcore.git", "tags": ["csharp", "web"] }, + { "url": "https://github.com/rails/rails.git", "tags": ["ruby", "web"] }, + { "url": "https://github.com/laravel/framework.git", "tags": ["php", "web"] }, + { "url": "https://github.com/fmtlib/fmt.git", "tags": ["cpp", "library"] } + ], + "broad": [ + { "url": "https://github.com/psf/requests.git", "tags": ["python", "http"] }, + { "url": "https://github.com/django/django.git", "tags": ["python", "web"] }, + { "url": "https://github.com/pallets/flask.git", "tags": ["python", "web"] }, + { "url": "https://github.com/fastapi/fastapi.git", "tags": ["python", "api"] }, + { "url": "https://github.com/encode/httpx.git", "tags": ["python", "http"] }, + { "url": "https://github.com/expressjs/express.git", "tags": ["javascript", "web"] }, + { "url": "https://github.com/facebook/react.git", "tags": ["javascript", "frontend"] }, + { "url": "https://github.com/vercel/next.js.git", "tags": ["typescript", "react"] }, + { "url": "https://github.com/vuejs/core.git", "tags": ["typescript", "frontend"] }, + { "url": "https://github.com/remix-run/remix.git", "tags": ["typescript", "web"] }, + { "url": "https://github.com/gin-gonic/gin.git", "tags": ["go", "web"] }, + { "url": "https://github.com/prometheus/prometheus.git", "tags": ["go", "monitoring"] }, + { "url": "https://github.com/go-gorm/gorm.git", "tags": ["go", "database"] }, + { "url": "https://github.com/spf13/cobra.git", "tags": ["go", "cli"] }, + { "url": "https://github.com/tokio-rs/tokio.git", "tags": ["rust", "async"] }, + { "url": "https://github.com/rust-lang/cargo.git", "tags": ["rust", "cli"] }, + { "url": "https://github.com/serde-rs/serde.git", "tags": ["rust", "serialization"] }, + { "url": "https://github.com/tauri-apps/tauri.git", "tags": ["rust", "desktop"] }, + { "url": "https://github.com/spring-projects/spring-boot.git", "tags": ["java", "spring"] }, + { "url": "https://github.com/spring-projects/spring-petclinic.git", "tags": ["java", "spring"] }, + { "url": "https://github.com/apache/dubbo.git", "tags": ["java", "rpc"] }, + { "url": "https://github.com/dotnet/runtime.git", "tags": ["csharp", "runtime"] }, + { "url": "https://github.com/dotnet/aspnetcore.git", "tags": ["csharp", "web"] }, + { "url": "https://github.com/rails/rails.git", "tags": ["ruby", "web"] }, + { "url": "https://github.com/rubocop/rubocop.git", "tags": ["ruby", "cli"] }, + { "url": "https://github.com/laravel/framework.git", "tags": ["php", "web"] }, + { "url": "https://github.com/symfony/symfony.git", "tags": ["php", "web"] }, + { "url": "https://github.com/fmtlib/fmt.git", "tags": ["cpp", "library"] }, + { "url": "https://github.com/nlohmann/json.git", "tags": ["cpp", "json"] }, + { "url": "https://github.com/catchorg/Catch2.git", "tags": ["cpp", "testing"] } + ] +} diff --git a/src/rag/github_scraper.js b/src/rag/github_scraper.js new file mode 100644 index 00000000..2bccafd2 --- /dev/null +++ b/src/rag/github_scraper.js @@ -0,0 +1,57 @@ +'use strict'; + +const fs = require('fs'); +const path = require('path'); +const crypto = require('crypto'); +const cp = require('child_process'); + +const CODE_EXTS = new Set(['.js', '.ts', '.tsx', '.jsx', '.py', '.go', '.rs', '.java', '.cpp', '.c', '.cs', '.php', '.rb']); + +function runGit(args, cwd) { + cp.execFileSync('git', args, { cwd, stdio: 'pipe' }); +} + +function ensureRepo(url, targetDir) { + const name = url.replace(/\.git$/, '').split('/').slice(-2).join('__'); + const out = path.join(targetDir, name); + if (fs.existsSync(path.join(out, '.git'))) runGit(['pull', '--ff-only'], out); + else runGit(['clone', '--depth=1', url, out], targetDir); + return out; +} + +function* walk(dir) { + const list = fs.readdirSync(dir, { withFileTypes: true }); + for (const ent of list) { + if (ent.name === '.git' || ent.name === 'node_modules' || ent.name === 'dist' || ent.name === 'build') continue; + const full = path.join(dir, ent.name); + if (ent.isDirectory()) yield* walk(full); + else if (ent.isFile()) yield full; + } +} + +function chunkCode(text, chunkLines = 60, overlap = 15) { + const lines = text.split('\n'); + const out = []; + for (let i = 0; i < lines.length; i += (chunkLines - overlap)) { + const piece = lines.slice(i, i + chunkLines).join('\n').trim(); + if (piece.length >= 80) out.push({ startLine: i + 1, code: piece }); + } + return out; +} + +function collectSnippets(repoRoot, sourceUrl) { + const snippets = []; + for (const file of walk(repoRoot)) { + const ext = path.extname(file).toLowerCase(); + if (!CODE_EXTS.has(ext)) continue; + const rel = path.relative(repoRoot, file); + const body = fs.readFileSync(file, 'utf-8'); + for (const c of chunkCode(body)) { + const id = crypto.createHash('sha1').update(`${sourceUrl}:${rel}:${c.startLine}:${c.code}`).digest('hex'); + snippets.push({ id, repo: sourceUrl, path: rel, startLine: c.startLine, code: c.code, lang: ext.slice(1) }); + } + } + return snippets; +} + +module.exports = { ensureRepo, collectSnippets, chunkCode }; diff --git a/src/rag/index_store.js b/src/rag/index_store.js new file mode 100644 index 00000000..055a6724 --- /dev/null +++ b/src/rag/index_store.js @@ -0,0 +1,170 @@ +'use strict'; + +const fs = require('fs'); +const path = require('path'); + +const DIMS = parseInt(process.env.SMALLCODE_RAG_DIMS || '1024', 10); +const BM25_K1 = 1.4; +const BM25_B = 0.72; + +function splitIdentifier(token) { + return String(token || '') + .replace(/([a-z0-9])([A-Z])/g, '$1 $2') + .replace(/[_-]+/g, ' ') + .toLowerCase() + .split(/\s+/) + .filter(Boolean); +} + +function tokenize(text) { + const raw = String(text || '').match(/[A-Za-z_][A-Za-z0-9_]*|\d+|#[A-Za-z0-9_-]+/g) || []; + const out = []; + for (const tok of raw) { + const parts = splitIdentifier(tok); + out.push(...parts); + if (parts.length > 1) out.push(parts.join('_')); + } + return out.filter(t => t.length >= 2 && !STOP_WORDS.has(t)); +} + +function hashToken(token, dims = DIMS) { + let h = 2166136261; + for (let i = 0; i < token.length; i++) { + h ^= token.charCodeAt(i); + h = Math.imul(h, 16777619); + } + return Math.abs(h) % dims; +} + +function termFrequency(tokens) { + const tf = Object.create(null); + for (const t of tokens) tf[t] = (tf[t] || 0) + 1; + return tf; +} + +function embed(text, dims = DIMS) { + const vec = Object.create(null); + const toks = tokenize(text); + for (const t of toks) { + const key = String(hashToken(t, dims)); + vec[key] = (vec[key] || 0) + 1; + } + let norm = 0; + for (const v of Object.values(vec)) norm += v * v; + norm = Math.sqrt(norm) || 1; + for (const k of Object.keys(vec)) vec[k] = Number((vec[k] / norm).toFixed(6)); + return vec; +} + +function cosine(a, b) { + if (Array.isArray(a) && Array.isArray(b)) { + const n = Math.min(a.length, b.length); + let s = 0; + for (let i = 0; i < n; i++) s += a[i] * b[i]; + return s; + } + const left = a || {}; + const right = b || {}; + let s = 0; + const small = Object.keys(left).length <= Object.keys(right).length ? left : right; + const large = small === left ? right : left; + for (const [k, v] of Object.entries(small)) s += v * (large[k] || 0); + return s; +} + +function bm25Score(queryTerms, doc, stats) { + if (!queryTerms.length || !doc.termFreq) return 0; + const dl = doc.docLength || 1; + const avgdl = stats.avgDocLength || 1; + let score = 0; + for (const term of queryTerms) { + const tf = doc.termFreq[term] || 0; + if (!tf) continue; + const df = stats.df.get(term) || 0; + const idf = Math.log(1 + (stats.totalDocs - df + 0.5) / (df + 0.5)); + score += idf * ((tf * (BM25_K1 + 1)) / (tf + BM25_K1 * (1 - BM25_B + BM25_B * (dl / avgdl)))); + } + return score; +} + +class RagIndexStore { + constructor(options = {}) { + this.path = options.path || path.join(process.cwd(), '.smallcode', 'rag', 'index.json'); + this.docs = []; + this._stats = null; + } + + load() { + try { + const raw = JSON.parse(fs.readFileSync(this.path, 'utf-8')); + this.docs = Array.isArray(raw.docs) ? raw.docs : []; + this._stats = null; + } catch { + this.docs = []; + this._stats = null; + } + return this.docs.length; + } + + save() { + fs.mkdirSync(path.dirname(this.path), { recursive: true }); + fs.writeFileSync(this.path, JSON.stringify({ version: 2, dims: DIMS, scoring: 'hybrid-bm25-hash-vector', docs: this.docs }, null, 2)); + } + + _prepare(snippet) { + const code = snippet.code || snippet.content || ''; + const searchable = [snippet.repoName, snippet.repo, snippet.path, snippet.lang, snippet.symbol, ...(snippet.tags || []), code].filter(Boolean).join('\n'); + const tokens = tokenize(searchable); + return { + ...snippet, + code, + termFreq: termFrequency(tokens), + docLength: tokens.length, + embedding: embed(searchable), + }; + } + + upsertMany(snippets) { + const byId = new Map(this.docs.map(d => [d.id, d])); + for (const s of snippets) byId.set(s.id, this._prepare(s)); + this.docs = [...byId.values()]; + this._stats = null; + return this.docs.length; + } + + _buildStats(queryTerms = []) { + const df = new Map(queryTerms.map(t => [t, 0])); + let totalLen = 0; + for (const d of this.docs) { + totalLen += d.docLength || 0; + if (!d.termFreq) continue; + for (const t of queryTerms) if (d.termFreq[t]) df.set(t, (df.get(t) || 0) + 1); + } + return { df, totalDocs: this.docs.length || 1, avgDocLength: totalLen / (this.docs.length || 1) || 1 }; + } + + search(query, limit = 8, options = {}) { + if (!this.docs.length) return []; + const queryTerms = [...new Set(tokenize(query))]; + const queryEmbedding = embed(query); + const stats = this._buildStats(queryTerms); + const vectorWeight = options.vectorWeight ?? 0.75; + return this.docs + .map(d => { + const bm25 = bm25Score(queryTerms, d, stats); + const vector = cosine(queryEmbedding, d.embedding || {}); + return { ...d, bm25Score: bm25, vectorScore: vector, score: bm25 + vectorWeight * vector }; + }) + .filter(d => d.score > 0) + .sort((a, b) => b.score - a.score) + .slice(0, limit); + } +} + +const STOP_WORDS = new Set([ + 'the', 'and', 'for', 'that', 'this', 'with', 'from', 'into', 'your', 'you', 'are', 'was', 'were', + 'will', 'would', 'could', 'should', 'have', 'has', 'had', 'not', 'but', 'what', 'when', 'where', + 'why', 'how', 'can', 'need', 'make', 'create', 'add', 'fix', 'code', 'file', 'class', 'function', +]); + +module.exports = { RagIndexStore, tokenize, embed, cosine, bm25Score }; diff --git a/src/rag/retriever.js b/src/rag/retriever.js new file mode 100644 index 00000000..855461c0 --- /dev/null +++ b/src/rag/retriever.js @@ -0,0 +1,84 @@ +'use strict'; + +const { RagIndexStore } = require('./index_store'); + +function planQuery(query) { + const q = String(query || '').trim(); + const intent = /bug|error|fix|failing|stack/.test(q.toLowerCase()) ? 'debug' : 'implement'; + const focus = q.split(/\s+/).filter(w => w.length > 3).slice(0, 8); + return { intent, focus }; +} + +function googleFallbackUrl(query) { + return `https://www.google.com/search?q=${encodeURIComponent(query + ' github code example')}`; +} + +class RagRetriever { + constructor(options = {}) { + this.index = options.index || new RagIndexStore(options.store || {}); + this.maxLoops = options.maxLoops || 3; + this.loaded = false; + } + + load() { + const count = this.index.load(); + this.loaded = true; + return count; + } + + ensureLoaded() { + if (!this.loaded) return this.load(); + return this.index.docs.length; + } + + retrieve(query, opts = {}) { + this.ensureLoaded(); + const plan = planQuery(query); + const loops = []; + let hits = []; + for (let i = 0; i < (opts.maxLoops || this.maxLoops); i++) { + const subquery = i === 0 ? query : `${query} ${plan.focus.slice(0, i + 2).join(' ')}`; + hits = this.index.search(subquery, opts.limit || 8); + loops.push({ loop: i + 1, subquery, hitCount: hits.length, topScore: hits[0]?.score || 0 }); + if ((hits[0]?.score || 0) >= 0.3) break; + } + return { + plan, + loops, + hits, + stuck: (hits[0]?.score || 0) < 0.2, + googleFallback: (hits[0]?.score || 0) < 0.2 ? googleFallbackUrl(query) : null, + }; + } + + formatForPrompt(query, opts = {}) { + if (process.env.SMALLCODE_RAG_DISABLE === 'true') return ''; + const result = this.retrieve(query, opts); + if (!result.hits.length) return ''; + + const maxChars = opts.maxChars || 6000; + let used = 0; + const parts = []; + for (const hit of result.hits.slice(0, opts.limit || 6)) { + const code = String(hit.code || '').slice(0, opts.snippetChars || 1200); + const lang = String(hit.lang || '').replace(/[^a-z0-9+#-]/gi, ''); + const scoreBits = [`score=${hit.score.toFixed(3)}`]; + if (typeof hit.bm25Score === 'number') scoreBits.push(`bm25=${hit.bm25Score.toFixed(3)}`); + if (typeof hit.vectorScore === 'number') scoreBits.push(`vector=${hit.vectorScore.toFixed(3)}`); + const symbol = hit.symbol ? ` ${hit.symbol}` : ''; + const header = `### ${hit.repo || 'local'}:${hit.path}:${hit.startLine || 1}${symbol} ${scoreBits.join(' ')}`; + const block = `${header}\n` + '```' + `${lang}\n${code}\n` + '```'; + if (used + block.length > maxChars) break; + used += block.length; + parts.push(block); + } + if (!parts.length) return ''; + + const stuckHint = result.stuck && process.env.SMALLCODE_WEB_BROWSE === 'true' + ? `\nRAG confidence is low. If blocked, use web_search with: ${query} github code example` + : ''; + return `\n[RAG_CODE_CONTEXT] Retrieved similar code snippets using hybrid BM25 + local hashed-vector search. Use these as examples, not as authoritative project files.${stuckHint}\n${parts.join('\n\n')}\n[/RAG_CODE_CONTEXT]\n`; + } +} + +module.exports = { RagRetriever, planQuery, googleFallbackUrl }; diff --git a/test/rag.test.js b/test/rag.test.js new file mode 100644 index 00000000..b5632bc2 --- /dev/null +++ b/test/rag.test.js @@ -0,0 +1,83 @@ +'use strict'; +const test = require('node:test'); +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const os = require('node:os'); +const path = require('node:path'); +const cp = require('node:child_process'); +const { RagIndexStore } = require('../src/rag/index_store'); +const { RagRetriever, planQuery } = require('../src/rag/retriever'); + +test('hybrid RAG index ranks related snippets first', () => { + const store = new RagIndexStore({ path: '/tmp/nonexistent.json' }); + store.upsertMany([ + { id: '1', path: 'a.py', lang: 'python', symbol: 'binary_search', code: 'def binary_search(arr, x):\n pass' }, + { id: '2', path: 'b.js', lang: 'javascript', symbol: 'renderButton', code: 'function renderButton() { return "ok" }' }, + ]); + const hits = store.search('how to implement binary search in python', 1); + assert.equal(hits[0].id, '1'); + assert.ok(hits[0].bm25Score > 0); +}); + +test('BM25 exact API names beat vague vector-only matches', () => { + const store = new RagIndexStore({ path: '/tmp/nonexistent.json' }); + store.upsertMany([ + { id: 'api', path: 'auth.ts', lang: 'typescript', symbol: 'createAccessToken', code: 'export function createAccessToken(userId: string) { return jwt.sign({ userId }, secret) }' }, + { id: 'vague', path: 'misc.ts', lang: 'typescript', symbol: 'createThing', code: 'export function createThing(value: string) { return value.toLowerCase() }' }, + ]); + const hits = store.search('createAccessToken jwt auth token', 2); + assert.equal(hits[0].id, 'api'); +}); + +test('query planning detects debug intent', () => { + const p = planQuery('fix stack error in parser'); + assert.equal(p.intent, 'debug'); +}); + +test('retriever formats snippets for prompt injection', () => { + const store = new RagIndexStore({ path: '/tmp/nonexistent.json' }); + store.upsertMany([ + { id: '1', repo: 'local', path: 'router.ts', startLine: 7, lang: 'ts', symbol: 'routeRequest', code: 'export function routeRequest() { return true }' }, + ]); + const retriever = new RagRetriever({ index: store }); + retriever.loaded = true; + const prompt = retriever.formatForPrompt('route request in typescript', { limit: 1 }); + assert.match(prompt, /RAG_CODE_CONTEXT/); + assert.match(prompt, /router\.ts:7 routeRequest/); + assert.match(prompt, /bm25=/); +}); + +test('python scraper emits snippet-sized symbol chunks from local repos', () => { + const tmp = fs.mkdtempSync(path.join(os.tmpdir(), 'smallcode-rag-scraper-')); + const repo = path.join(tmp, 'repo'); + fs.mkdirSync(repo, { recursive: true }); + const source = [ + 'def alpha_search(items, needle):', + ' for item in items:', + ' if item == needle:', + ' return item', + ' return None', + '', + 'class BetaCache:', + ' def __init__(self):', + ' self.items = {}', + ' def get(self, key):', + ' return self.items.get(key)', + ].join('\n'); + fs.writeFileSync(path.join(repo, 'sample.py'), source); + const out = path.join(tmp, 'snippets.jsonl'); + const result = cp.spawnSync('python3', [ + path.join(__dirname, '..', 'scripts', 'rag_scraper.py'), + '--preset', 'none', + '--repo', repo, + '--out', out, + '--min-chars', '20', + '--chunk-lines', '6', + ], { encoding: 'utf-8' }); + assert.equal(result.status, 0, result.stderr); + const lines = fs.readFileSync(out, 'utf-8').trim().split('\n').map(JSON.parse); + assert.ok(lines.length >= 2); + assert.equal(lines[0].kind, 'symbol'); + assert.equal(lines[0].path, 'sample.py'); + assert.ok(lines.every(s => s.code.split('\n').length <= 6)); +});