From 7fc116ee50a103f7f84c9965546053b77d183b59 Mon Sep 17 00:00:00 2001 From: AdarshVMore Date: Sun, 28 Jun 2026 14:48:05 +0530 Subject: [PATCH] implemented the AST for codebase indexing and replaced existing RAG with it --- Server/.npmrc | 1 + Server/app/api-service/Dockerfile | 8 +- Server/app/webhook-service/Dockerfile | 7 +- Server/app/worker-service/Dockerfile | 8 +- .../worker-service/services/rag.service.ts | 219 ++++++-- .../services/repoSetup.service.ts | 29 +- Server/package-lock.json | 241 +++++++-- Server/package.json | 7 + Server/package/ai/review.ts | 26 +- Server/package/ai/searchCode.tool.ts | 39 +- Server/package/ast/parser.ts | 473 ++++++++++++++++++ 11 files changed, 923 insertions(+), 135 deletions(-) create mode 100644 Server/.npmrc create mode 100644 Server/package/ast/parser.ts diff --git a/Server/.npmrc b/Server/.npmrc new file mode 100644 index 0000000..521a9f7 --- /dev/null +++ b/Server/.npmrc @@ -0,0 +1 @@ +legacy-peer-deps=true diff --git a/Server/app/api-service/Dockerfile b/Server/app/api-service/Dockerfile index d7dc685..be9e208 100644 --- a/Server/app/api-service/Dockerfile +++ b/Server/app/api-service/Dockerfile @@ -2,8 +2,14 @@ FROM node:20-alpine WORKDIR /app +RUN apk add --no-cache \ + python3 \ + make \ + g++ \ + libc6-compat + COPY package.json package-lock.json ./ -RUN npm ci +RUN npm install --force COPY . . diff --git a/Server/app/webhook-service/Dockerfile b/Server/app/webhook-service/Dockerfile index 14c54d1..b961935 100644 --- a/Server/app/webhook-service/Dockerfile +++ b/Server/app/webhook-service/Dockerfile @@ -2,8 +2,13 @@ FROM node:20-alpine WORKDIR /app +RUN apk add --no-cache \ + python3 \ + make \ + g++ \ + libc6-compat COPY package.json package-lock.json ./ -RUN npm ci +RUN npm install --force COPY . . diff --git a/Server/app/worker-service/Dockerfile b/Server/app/worker-service/Dockerfile index 610d264..40d85fc 100644 --- a/Server/app/worker-service/Dockerfile +++ b/Server/app/worker-service/Dockerfile @@ -2,8 +2,14 @@ FROM node:20-alpine WORKDIR /app +RUN apk add --no-cache \ + python3 \ + make \ + g++ \ + libc6-compat + COPY package.json package-lock.json ./ -RUN npm ci +RUN npm install --force COPY . . diff --git a/Server/app/worker-service/services/rag.service.ts b/Server/app/worker-service/services/rag.service.ts index 1b16fd9..ab3e41e 100644 --- a/Server/app/worker-service/services/rag.service.ts +++ b/Server/app/worker-service/services/rag.service.ts @@ -5,12 +5,15 @@ import { readCodeFiles, cleanup, } from "./repoSetup.service.js"; -import fs from "fs"; +import { + parseFileToASTChunks, + parseChangedFilesFromDiff, + parseChangedSymbolsFromDiff, + type ASTCodeChunk, +} from "../../../package/ast/parser.js"; const EMBEDDING_MODEL = "llama-text-embed-v2"; const EMBEDDING_DIMENSION = 1024; -const CHUNK_SIZE = 800; -const CHUNK_OVERLAP = 100; const EMBED_BATCH_SIZE = 96; type RAGPipelineInput = { @@ -19,12 +22,7 @@ type RAGPipelineInput = { repo: string; }; -type CodeChunk = { - id: string; - content: string; - filePath: string; - chunkIndex: number; -}; +export type CodeChunk = ASTCodeChunk; type EmbeddingVector = { id: string; @@ -33,31 +31,33 @@ type EmbeddingVector = { content: string; filePath: string; chunkIndex: number; + nodeType: string; + symbolName: string; + startLine: number; + endLine: number; + imports: string; }; }; +export type RelevantCodeMatch = { + id: string; + score: number; + content: string; + filePath: string; + nodeType?: string; + symbolName?: string; + startLine?: number; + endLine?: number; +}; + function getPinecone(): Pinecone { return new Pinecone({ apiKey: process.env.PINECONE_API_KEY! }); } -function chunkFile(filePath: string, content: string): CodeChunk[] { - const chunks: CodeChunk[] = []; - let i = 0; - let chunkIndex = 0; - - while (i < content.length) { - const chunk = content.slice(i, i + CHUNK_SIZE); - chunks.push({ - id: `${filePath}::${chunkIndex}`, - content: chunk, - filePath, - chunkIndex, - }); - i += CHUNK_SIZE - CHUNK_OVERLAP; - chunkIndex++; - } - - return chunks; +function chunkFilesWithAST( + files: { path: string; content: string }[], +): CodeChunk[] { + return files.flatMap((file) => parseFileToASTChunks(file.path, file.content)); } async function ensureIndexExists(pc: Pinecone, indexName: string) { @@ -106,11 +106,16 @@ export async function createEmbeddings( const embedding = result.data[j]; vectors.push({ id: chunk.id, - values: (embedding as any).values as number[], + values: (embedding as { values: number[] }).values, metadata: { content: chunk.content, filePath: chunk.filePath, chunkIndex: chunk.chunkIndex, + nodeType: chunk.nodeType, + symbolName: chunk.symbolName, + startLine: chunk.startLine, + endLine: chunk.endLine, + imports: JSON.stringify(chunk.imports), }, }); } @@ -145,40 +150,157 @@ export function toIndexName(owner: string, repo: string): string { .slice(0, 45); } +function toMatch(m: { + id?: string; + score?: number; + metadata?: Record; +}): RelevantCodeMatch { + return { + id: m.id ?? "", + score: m.score ?? 0, + content: (m.metadata?.content as string) ?? "", + filePath: (m.metadata?.filePath as string) ?? "", + nodeType: m.metadata?.nodeType as string | undefined, + symbolName: m.metadata?.symbolName as string | undefined, + startLine: m.metadata?.startLine as number | undefined, + endLine: m.metadata?.endLine as number | undefined, + }; +} -export async function searchRelevantEmbeddings(query: object) { - const { - text, - indexName, - topK = 5, - } = query as { text: string; indexName: string; topK?: number }; - +async function embedQuery(text: string): Promise { const pc = getPinecone(); - const result = await pc.inference.embed({ model: EMBEDDING_MODEL, inputs: [text], parameters: { inputType: "query", truncate: "END" }, }); + return (result.data[0] as { values: number[] }).values; +} + +export async function searchRelevantEmbeddings(query: object) { + const { + text, + indexName, + topK = 5, + filter, + } = query as { + text: string; + indexName: string; + topK?: number; + filter?: Record; + }; - const queryVector = (result.data[0] as any).values as number[]; + const queryVector = await embedQuery(text); + const pc = getPinecone(); const index = pc.index({ name: indexName }); const searchResult = await index.query({ vector: queryVector, topK, includeMetadata: true, + ...(filter ? { filter } : {}), }); - return searchResult.matches.map((m) => ({ - id: m.id, - score: m.score, - content: m.metadata?.content as string, - filePath: m.metadata?.filePath as string, + return searchResult.matches.map(toMatch); +} + +function dedupeMatches(matches: RelevantCodeMatch[]): RelevantCodeMatch[] { + const seen = new Set(); + const ranked = [...matches].sort((a, b) => b.score - a.score); + const unique: RelevantCodeMatch[] = []; + + for (const match of ranked) { + const key = `${match.filePath}::${match.symbolName ?? match.id}`; + if (seen.has(key)) continue; + seen.add(key); + unique.push(match); + } + + return unique; +} + +function scoreMatchForDiff( + match: RelevantCodeMatch, + changedFiles: string[], + changedSymbols: string[], +): number { + let boost = match.score; + + if (changedFiles.some((file) => match.filePath.endsWith(file) || file.endsWith(match.filePath))) { + boost += 0.15; + } + + if ( + match.symbolName && + changedSymbols.some( + (symbol) => + symbol === match.symbolName || + match.symbolName?.includes(symbol) || + symbol.includes(match.symbolName ?? ""), + ) + ) { + boost += 0.2; + } + + if (match.nodeType === "import" || match.nodeType === "export") { + boost += 0.05; + } + + return boost; +} + +export async function retrieveContextForDiff(options: { + diff: string; + semanticQuery: string; + indexName: string; + topK?: number; +}): Promise { + const { diff, semanticQuery, indexName, topK = 8 } = options; + const changedFiles = parseChangedFilesFromDiff(diff); + const changedSymbols = parseChangedSymbolsFromDiff(diff); + + const semanticMatches = await searchRelevantEmbeddings({ + text: semanticQuery, + indexName, + topK: topK * 2, + }); + + const fileMatches: RelevantCodeMatch[] = []; + for (const filePath of changedFiles.slice(0, 5)) { + const matches = await searchRelevantEmbeddings({ + text: `${semanticQuery} ${filePath}`, + indexName, + topK: 4, + filter: { filePath: { $eq: filePath } }, + }); + fileMatches.push(...matches); + } + + const symbolMatches: RelevantCodeMatch[] = []; + for (const symbol of changedSymbols.slice(0, 5)) { + const matches = await searchRelevantEmbeddings({ + text: `${semanticQuery} ${symbol}`, + indexName, + topK: 3, + }); + symbolMatches.push(...matches); + } + + const combined = dedupeMatches([ + ...semanticMatches, + ...fileMatches, + ...symbolMatches, + ]).map((match) => ({ + ...match, + score: scoreMatchForDiff(match, changedFiles, changedSymbols), })); + + return combined + .sort((a, b) => b.score - a.score) + .slice(0, topK); } -export async function vectorDBExists(DBName: string): Promise { +export async function vectorDBExists(DBName: string): Promise { if (!DBName) return false; const pc = getPinecone(); @@ -186,8 +308,6 @@ export async function vectorDBExists(DBName: string): Promise { return indexes?.some((idx) => idx.name === DBName) ?? false; } -// yo bro this is the main funciton - export async function runRAGPipeline(data: object) { const { installationId, owner, repo } = data as RAGPipelineInput; @@ -196,16 +316,13 @@ export async function runRAGPipeline(data: object) { const extractedDir = extractZIP(zipPath); const files = readCodeFiles(extractedDir); - const filesJSONStrgified = JSON.stringify(files) - - - const chunks = files.flatMap((f) => chunkFile(f.path, f.content)); + const chunks = chunkFilesWithAST(files); const vectors = await createEmbeddings({ chunks, indexName }); await saveToVectorDB({ vectors, indexName }); cleanup([zipPath, extractedDir]); console.log( - `RAG pipeline complete for ${owner}/${repo}: ${chunks.length} chunks indexed` + `RAG pipeline complete for ${owner}/${repo}: ${files.length} files → ${chunks.length} AST chunks indexed`, ); -} \ No newline at end of file +} diff --git a/Server/app/worker-service/services/repoSetup.service.ts b/Server/app/worker-service/services/repoSetup.service.ts index be3d60a..52085b5 100644 --- a/Server/app/worker-service/services/repoSetup.service.ts +++ b/Server/app/worker-service/services/repoSetup.service.ts @@ -51,8 +51,18 @@ export function extractZIP(zipPath: string): string { return destDir; } +function resolveRepoRoot(extractedDir: string): string { + const entries = fs.readdirSync(extractedDir); + if (entries.length === 1) { + const nested = path.join(extractedDir, entries[0]); + if (fs.statSync(nested).isDirectory()) return nested; + } + return extractedDir; +} + export function readCodeFiles(dir: string) { const files: { path: string; content: string }[] = []; + const repoRoot = resolveRepoRoot(dir); function walk(current: string) { const items = fs.readdirSync(current); @@ -61,20 +71,25 @@ export function readCodeFiles(dir: string) { const fullPath = path.join(current, item); if (fs.statSync(fullPath).isDirectory()) { - if (item === "node_modules" || item === ".git" || item === "dist") { + if ( + item === "node_modules" || + item === ".git" || + item === "dist" || + item === "build" || + item === "coverage" + ) { continue; } walk(fullPath); - } else { - if (ALLOWED_EXT.includes(path.extname(fullPath))) { - const content = fs.readFileSync(fullPath, "utf-8"); - files.push({ path: fullPath, content }); - } + } else if (ALLOWED_EXT.includes(path.extname(fullPath))) { + const content = fs.readFileSync(fullPath, "utf-8"); + const relativePath = path.relative(repoRoot, fullPath); + files.push({ path: relativePath, content }); } } } - walk(dir); + walk(repoRoot); return files; } diff --git a/Server/package-lock.json b/Server/package-lock.json index d93d6fe..5b74c17 100644 --- a/Server/package-lock.json +++ b/Server/package-lock.json @@ -27,6 +27,14 @@ "pg": "^8.20.0", "prisma": "^7.6.0", "redis": "^5.11.0", + "tree-sitter": "^0.25.0", + "tree-sitter-go": "^0.25.0", + "tree-sitter-java": "^0.23.5", + "tree-sitter-javascript": "^0.25.0", + "tree-sitter-python": "^0.25.0", + "tree-sitter-ruby": "^0.23.1", + "tree-sitter-rust": "^0.24.0", + "tree-sitter-typescript": "^0.23.2", "zod": "^4.3.6" }, "devDependencies": { @@ -2207,16 +2215,6 @@ "dev": true, "license": "MIT" }, - "node_modules/@types/react": { - "version": "19.2.14", - "resolved": "https://registry.npmjs.org/@types/react/-/react-19.2.14.tgz", - "integrity": "sha512-ilcTH/UniCkMdtexkoCN0bI7pMcJDvmQFPvuPvmEaYA/NSfFTAgdUSLAoVjaRJm7+6PvcM+q1zYOwS4wTYMF9w==", - "license": "MIT", - "peer": true, - "dependencies": { - "csstype": "^3.2.2" - } - }, "node_modules/@types/send": { "version": "1.2.1", "resolved": "https://registry.npmjs.org/@types/send/-/send-1.2.1.tgz", @@ -2770,13 +2768,6 @@ "deprecated": "This package is no longer supported. It's now a built-in Node module. If you've depended on crypto, you should switch to the one that's built-in.", "license": "ISC" }, - "node_modules/csstype": { - "version": "3.2.3", - "resolved": "https://registry.npmjs.org/csstype/-/csstype-3.2.3.tgz", - "integrity": "sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==", - "license": "MIT", - "peer": true - }, "node_modules/debug": { "version": "4.4.3", "resolved": "https://registry.npmjs.org/debug/-/debug-4.4.3.tgz", @@ -3992,6 +3983,15 @@ "node": ">= 0.6" } }, + "node_modules/node-addon-api": { + "version": "8.9.0", + "resolved": "https://registry.npmjs.org/node-addon-api/-/node-addon-api-8.9.0.tgz", + "integrity": "sha512-ekZMeaaIzSQTSpr7X2X3iJM7lTzgnx8ahAG9pJfT/7+14mlEM8ZYQ9cgCDvSSRbReFK0oHli3WrZdCiRsgAT9Q==", + "license": "MIT", + "engines": { + "node": "^18 || ^20 || >= 21" + } + }, "node_modules/node-fetch": { "version": "2.7.0", "resolved": "https://registry.npmjs.org/node-fetch/-/node-fetch-2.7.0.tgz", @@ -4018,6 +4018,17 @@ "integrity": "sha512-g9yhqoedzIUm0nTnTqAQvueMPVOuIY16bqgAJJC8XOOubYFNwz6IER9qs0Gq2Xd0+CecCKFjtdDTMA4u4xG06Q==", "license": "MIT" }, + "node_modules/node-gyp-build": { + "version": "4.8.4", + "resolved": "https://registry.npmjs.org/node-gyp-build/-/node-gyp-build-4.8.4.tgz", + "integrity": "sha512-LA4ZjwlnUblHVgq0oBF3Jl/6h/Nvs5fzBLwdEF4nuxnFdsfajde4WfxtJr3CaiH+F6ewcIB/q4jQ4UzPyid+CQ==", + "license": "MIT", + "bin": { + "node-gyp-build": "bin.js", + "node-gyp-build-optional": "optional.js", + "node-gyp-build-test": "build-test.js" + } + }, "node_modules/normalize-path": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/normalize-path/-/normalize-path-3.0.0.tgz", @@ -4473,29 +4484,6 @@ "destr": "^2.0.3" } }, - "node_modules/react": { - "version": "19.2.4", - "resolved": "https://registry.npmjs.org/react/-/react-19.2.4.tgz", - "integrity": "sha512-9nfp2hYpCwOjAN+8TZFGhtWEwgvWHXqESH8qT89AT/lWklpLON22Lc8pEtnpsZz7VmawabSU0gCjnj8aC0euHQ==", - "license": "MIT", - "peer": true, - "engines": { - "node": ">=0.10.0" - } - }, - "node_modules/react-dom": { - "version": "19.2.4", - "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.4.tgz", - "integrity": "sha512-AXJdLo8kgMbimY95O2aKQqsz2iWi9jMgKJhRBAxECE4IFxfcazB2LmzloIoibJI3C12IlY20+KFaLv+71bUJeQ==", - "license": "MIT", - "peer": true, - "dependencies": { - "scheduler": "^0.27.0" - }, - "peerDependencies": { - "react": "^19.2.4" - } - }, "node_modules/readdirp": { "version": "3.6.0", "resolved": "https://registry.npmjs.org/readdirp/-/readdirp-3.6.0.tgz", @@ -4639,13 +4627,6 @@ "integrity": "sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==", "license": "MIT" }, - "node_modules/scheduler": { - "version": "0.27.0", - "resolved": "https://registry.npmjs.org/scheduler/-/scheduler-0.27.0.tgz", - "integrity": "sha512-eNv+WrVbKu1f3vbYJT/xtiF5syA5HPIMtf9IgY/nKg0sWqzAUEvqY/xm7OcZc/qafLx/iO9FgOmeSAp4v5ti/Q==", - "license": "MIT", - "peer": true - }, "node_modules/semver": { "version": "7.7.4", "resolved": "https://registry.npmjs.org/semver/-/semver-7.7.4.tgz", @@ -4968,6 +4949,170 @@ "tree-kill": "cli.js" } }, + "node_modules/tree-sitter": { + "version": "0.25.0", + "resolved": "https://registry.npmjs.org/tree-sitter/-/tree-sitter-0.25.0.tgz", + "integrity": "sha512-PGZZzFW63eElZJDe/b/R/LbsjDDYJa5UEjLZJB59RQsMX+fo0j54fqBPn1MGKav/QNa0JR0zBiVaikYDWCj5KQ==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.3.0", + "node-gyp-build": "^4.8.4" + } + }, + "node_modules/tree-sitter-go": { + "version": "0.25.0", + "resolved": "https://registry.npmjs.org/tree-sitter-go/-/tree-sitter-go-0.25.0.tgz", + "integrity": "sha512-APBc/Dq3xz/e35Xpkhb1blu5UgW+2E3RyGWawZSCNcbGwa7jhSQPS8KsUupuzBla8PCo8+lz9W/JDJjmfRa2tw==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.3.1", + "node-gyp-build": "^4.8.4" + }, + "peerDependencies": { + "tree-sitter": "^0.25.0" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-java": { + "version": "0.23.5", + "resolved": "https://registry.npmjs.org/tree-sitter-java/-/tree-sitter-java-0.23.5.tgz", + "integrity": "sha512-Yju7oQ0Xx7GcUT01mUglPP+bYfvqjNCGdxqigTnew9nLGoII42PNVP3bHrYeMxswiCRM0yubWmN5qk+zsg0zMA==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.2.2", + "node-gyp-build": "^4.8.2" + }, + "peerDependencies": { + "tree-sitter": "^0.21.1" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-javascript": { + "version": "0.25.0", + "resolved": "https://registry.npmjs.org/tree-sitter-javascript/-/tree-sitter-javascript-0.25.0.tgz", + "integrity": "sha512-1fCbmzAskZkxcZzN41sFZ2br2iqTYP3tKls1b/HKGNPQUVOpsUxpmGxdN/wMqAk3jYZnYBR1dd/y/0avMeU7dw==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.3.1", + "node-gyp-build": "^4.8.4" + }, + "peerDependencies": { + "tree-sitter": "^0.25.0" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-python": { + "version": "0.25.0", + "resolved": "https://registry.npmjs.org/tree-sitter-python/-/tree-sitter-python-0.25.0.tgz", + "integrity": "sha512-eCmJx6zQa35GxaCtQD+wXHOhYqBxEL+bp71W/s3fcDMu06MrtzkVXR437dRrCrbrDbyLuUDJpAgycs7ncngLXw==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.5.0", + "node-gyp-build": "^4.8.4" + }, + "peerDependencies": { + "tree-sitter": "^0.25.0" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-ruby": { + "version": "0.23.1", + "resolved": "https://registry.npmjs.org/tree-sitter-ruby/-/tree-sitter-ruby-0.23.1.tgz", + "integrity": "sha512-d9/RXgWjR6HanN7wTYhS5bpBQLz1VkH048Vm3CodPGyJVnamXMGb8oEhDypVCBq4QnHui9sTXuJBBP3WtCw5RA==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.2.2", + "node-gyp-build": "^4.8.2" + }, + "peerDependencies": { + "tree-sitter": "^0.21.1" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-rust": { + "version": "0.24.0", + "resolved": "https://registry.npmjs.org/tree-sitter-rust/-/tree-sitter-rust-0.24.0.tgz", + "integrity": "sha512-NWemUDf629Tfc90Y0Z55zuwPCAHkLxWnMf2RznYu4iBkkrQl2o/CHGB7Cr52TyN5F1DAx8FmUnDtCy9iUkXZEQ==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.2.2", + "node-gyp-build": "^4.8.4" + }, + "peerDependencies": { + "tree-sitter": "^0.22.1" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-typescript": { + "version": "0.23.2", + "resolved": "https://registry.npmjs.org/tree-sitter-typescript/-/tree-sitter-typescript-0.23.2.tgz", + "integrity": "sha512-e04JUUKxTT53/x3Uq1zIL45DoYKVfHH4CZqwgZhPg5qYROl5nQjV+85ruFzFGZxu+QeFVbRTPDRnqL9UbU4VeA==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.2.2", + "node-gyp-build": "^4.8.2", + "tree-sitter-javascript": "^0.23.1" + }, + "peerDependencies": { + "tree-sitter": "^0.21.0" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-typescript/node_modules/tree-sitter-javascript": { + "version": "0.23.1", + "resolved": "https://registry.npmjs.org/tree-sitter-javascript/-/tree-sitter-javascript-0.23.1.tgz", + "integrity": "sha512-/bnhbrTD9frUYHQTiYnPcxyHORIw157ERBa6dqzaKxvR/x3PC4Yzd+D1pZIMS6zNg2v3a8BZ0oK7jHqsQo9fWA==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.2.2", + "node-gyp-build": "^4.8.2" + }, + "peerDependencies": { + "tree-sitter": "^0.21.1" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, "node_modules/ts-algebra": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/ts-algebra/-/ts-algebra-2.0.0.tgz", @@ -5104,7 +5249,7 @@ "version": "6.0.2", "resolved": "https://registry.npmjs.org/typescript/-/typescript-6.0.2.tgz", "integrity": "sha512-bGdAIrZ0wiGDo5l8c++HWtbaNCWTS4UTv7RaTH/ThVIgjkveJt83m74bBHMJkuCbslY8ixgLBVZJIOiQlQTjfQ==", - "devOptional": true, + "dev": true, "license": "Apache-2.0", "bin": { "tsc": "bin/tsc", diff --git a/Server/package.json b/Server/package.json index 7e433bf..edfdfb0 100644 --- a/Server/package.json +++ b/Server/package.json @@ -31,6 +31,13 @@ "pg": "^8.20.0", "prisma": "^7.6.0", "redis": "^5.11.0", + "tree-sitter": "^0.25.0", + "tree-sitter-go": "^0.25.0", + "tree-sitter-javascript": "^0.25.0", + "tree-sitter-python": "^0.25.0", + "tree-sitter-ruby": "^0.23.1", + "tree-sitter-rust": "^0.24.0", + "tree-sitter-typescript": "^0.23.2", "zod": "^4.3.6" }, "devDependencies": { diff --git a/Server/package/ai/review.ts b/Server/package/ai/review.ts index 156462e..4a7a8e4 100644 --- a/Server/package/ai/review.ts +++ b/Server/package/ai/review.ts @@ -174,17 +174,29 @@ Do NOT include markdown, explanations, or text outside JSON.`, } type RelevantCodeMatch = { - id: string; - score: number; - content: string; - filePath: string; -}; + id: string + score: number + content: string + filePath: string + nodeType?: string + symbolName?: string + startLine?: number + endLine?: number +} function formatRelevantCode(matches: RelevantCodeMatch[]): string { if (!matches.length) return ""; return matches - .map((m) => `// ${m.filePath}\n${m.content}`) - .join("\n\n---\n\n"); + .map((m) => { + const header = [ + m.symbolName ? `// ${m.filePath} — ${m.symbolName} (${m.nodeType ?? "symbol"})` : `// ${m.filePath}`, + m.startLine && m.endLine ? `// Lines ${m.startLine}-${m.endLine}` : null, + ] + .filter(Boolean) + .join("\n") + return `${header}\n${m.content}` + }) + .join("\n\n---\n\n") } export function reviewPrompt( diff --git a/Server/package/ai/searchCode.tool.ts b/Server/package/ai/searchCode.tool.ts index 431f570..4aeb584 100644 --- a/Server/package/ai/searchCode.tool.ts +++ b/Server/package/ai/searchCode.tool.ts @@ -1,22 +1,23 @@ -import { searchRelevantEmbeddings } from "../../app/worker-service/services/rag.service.js"; - -async function searchFile() { +import { + searchRelevantEmbeddings, + retrieveContextForDiff, + toIndexName, +} from "../../app/worker-service/services/rag.service.js"; +export async function searchCodeContext(options: { + owner: string; + repo: string; + diff: string; + semanticQuery: string; + topK?: number; +}) { + const indexName = toIndexName(options.owner, options.repo); + return retrieveContextForDiff({ + diff: options.diff, + semanticQuery: options.semanticQuery, + indexName, + topK: options.topK ?? 8, + }); } - -// Some Steps to Chunk => -// Get Repo zip => Unzip => get Folder structure some how -// create a Tree of that folder structure => take one branch => Read all Code Files in it => create Functions / Classes + Imports + Variables + etc as Nodes -// we can create such Tree by telling LLM to create and then we Store it in a DB [ will do this for now => Big Context Window with Many Tokens used] -// Based on that Tres and Branches and Nodes => find Functions / Classes / Files and Chunk those and embed and Store in Vectore DB -// -// -// 1. Identify PR Type => Feature / Bug Fix / Update Code -// 2. If Feature -// -// -// -// -// -// \ No newline at end of file +export { searchRelevantEmbeddings, retrieveContextForDiff, toIndexName }; diff --git a/Server/package/ast/parser.ts b/Server/package/ast/parser.ts new file mode 100644 index 0000000..0b96bb4 --- /dev/null +++ b/Server/package/ast/parser.ts @@ -0,0 +1,473 @@ +import Parser from "tree-sitter"; +import JavaScript from "tree-sitter-javascript"; +import TypeScript from "tree-sitter-typescript"; +import Python from "tree-sitter-python"; +import Go from "tree-sitter-go"; +// import Java from "tree-sitter-java"; +import Rust from "tree-sitter-rust"; +import Ruby from "tree-sitter-ruby"; + +export type ASTNodeType = + | "function" + | "class" + | "method" + | "interface" + | "type" + | "import" + | "export" + | "variable" + | "struct" + | "enum" + | "module" + | "file"; + +export type ASTCodeChunk = { + id: string; + content: string; + filePath: string; + chunkIndex: number; + nodeType: ASTNodeType; + symbolName: string; + startLine: number; + endLine: number; + imports: string[]; +}; + +type TreeSitterLanguage = ReturnType; + +type LanguageConfig = { + language: TreeSitterLanguage; + chunkNodeTypes: Set; + nameFields: Record; +}; + +const JS_TS_CHUNK_TYPES = new Set([ + "function_declaration", + "generator_function_declaration", + "class_declaration", + "method_definition", + "interface_declaration", + "type_alias_declaration", + "import_statement", + "export_statement", + "lexical_declaration", + "variable_declaration", +]); + +const LANGUAGE_CONFIG: Record = { + ".js": { + language: JavaScript as TreeSitterLanguage, + chunkNodeTypes: JS_TS_CHUNK_TYPES, + nameFields: { + function_declaration: "name", + generator_function_declaration: "name", + class_declaration: "name", + method_definition: "name", + interface_declaration: "name", + type_alias_declaration: "name", + lexical_declaration: "declarator", + variable_declaration: "declarator", + }, + }, + ".jsx": { + language: JavaScript as TreeSitterLanguage, + chunkNodeTypes: JS_TS_CHUNK_TYPES, + nameFields: { + function_declaration: "name", + generator_function_declaration: "name", + class_declaration: "name", + method_definition: "name", + lexical_declaration: "declarator", + variable_declaration: "declarator", + }, + }, + ".ts": { + language: TypeScript.typescript as TreeSitterLanguage, + chunkNodeTypes: JS_TS_CHUNK_TYPES, + nameFields: { + function_declaration: "name", + generator_function_declaration: "name", + class_declaration: "name", + method_definition: "name", + interface_declaration: "name", + type_alias_declaration: "name", + lexical_declaration: "declarator", + variable_declaration: "declarator", + }, + }, + ".tsx": { + language: TypeScript.tsx as TreeSitterLanguage, + chunkNodeTypes: JS_TS_CHUNK_TYPES, + nameFields: { + function_declaration: "name", + generator_function_declaration: "name", + class_declaration: "name", + method_definition: "name", + interface_declaration: "name", + type_alias_declaration: "name", + lexical_declaration: "declarator", + variable_declaration: "declarator", + }, + }, + ".py": { + language: Python as TreeSitterLanguage, + chunkNodeTypes: new Set([ + "function_definition", + "class_definition", + "import_statement", + "import_from_statement", + "decorated_definition", + ]), + nameFields: { + function_definition: "name", + class_definition: "name", + decorated_definition: "definition", + }, + }, + ".go": { + language: Go as TreeSitterLanguage, + chunkNodeTypes: new Set([ + "function_declaration", + "method_declaration", + "type_declaration", + "import_declaration", + ]), + nameFields: { + function_declaration: "name", + method_declaration: "name", + type_declaration: "type_spec", + }, + }, + // ".java": { + // language: Java as TreeSitterLanguage, + // chunkNodeTypes: new Set([ + // "method_declaration", + // "class_declaration", + // "interface_declaration", + // "import_declaration", + // ]), + // nameFields: { + // method_declaration: "name", + // class_declaration: "name", + // interface_declaration: "name", + // }, + // }, + ".rs": { + language: Rust as TreeSitterLanguage, + chunkNodeTypes: new Set([ + "function_item", + "impl_item", + "struct_item", + "enum_item", + "mod_item", + "use_declaration", + ]), + nameFields: { + function_item: "name", + struct_item: "name", + enum_item: "name", + mod_item: "name", + }, + }, + ".rb": { + language: Ruby as TreeSitterLanguage, + chunkNodeTypes: new Set([ + "method", + "class", + "module", + "singleton_method", + ]), + nameFields: { + method: "name", + class: "name", + module: "name", + singleton_method: "name", + }, + }, +}; + +const NODE_TYPE_MAP: Record = { + function_declaration: "function", + generator_function_declaration: "function", + function_definition: "function", + function_item: "function", + class_declaration: "class", + class_definition: "class", + class: "class", + method_definition: "method", + method_declaration: "method", + method: "method", + singleton_method: "method", + interface_declaration: "interface", + interface: "interface", + type_alias_declaration: "type", + type_declaration: "type", + struct_item: "struct", + enum_item: "enum", + import_statement: "import", + import_from_statement: "import", + import_declaration: "import", + use_declaration: "import", + export_statement: "export", + lexical_declaration: "variable", + variable_declaration: "variable", + impl_item: "class", + mod_item: "module", + module: "module", + decorated_definition: "function", +}; + +const parserCache = new Map(); + +function getParser(ext: string): Parser | null { + const config = LANGUAGE_CONFIG[ext]; + if (!config) return null; + + if (!parserCache.has(ext)) { + const parser = new Parser(); + parser.setLanguage(config.language); + parserCache.set(ext, parser); + } + return parserCache.get(ext)!; +} + +function lineAt(content: string, byteIndex: number): number { + return content.slice(0, byteIndex).split("\n").length; +} + +function extractSymbolName( + node: Parser.SyntaxNode, + config: LanguageConfig, +): string { + if ( + node.type === "import_statement" || + node.type === "import_from_statement" || + node.type === "import_declaration" || + node.type === "use_declaration" + ) { + const source = node.childForFieldName("source") ?? node.childForFieldName("path"); + if (source) return source.text.replace(/['"]/g, "").slice(0, 120); + } + + const nameField = config.nameFields[node.type]; + if (nameField) { + const nameNode = node.childForFieldName(nameField); + if (nameNode) { + if (nameNode.type === "variable_declarator") { + const id = nameNode.childForFieldName("name"); + if (id) return id.text.slice(0, 120); + } + return nameNode.text.slice(0, 120); + } + } + + for (const child of node.namedChildren) { + if ( + child.type === "identifier" || + child.type === "type_identifier" || + child.type === "property_identifier" + ) { + return child.text.slice(0, 120); + } + if (child.type === "variable_declarator") { + const id = child.childForFieldName("name"); + if (id) return id.text.slice(0, 120); + } + } + + return node.type; +} + +function extractImports(content: string, ext: string): string[] { + const parser = getParser(ext); + if (!parser) return []; + + const tree = parser.parse(content); + const imports: string[] = []; + + function walk(node: Parser.SyntaxNode) { + if ( + node.type === "import_statement" || + node.type === "import_from_statement" || + node.type === "import_declaration" || + node.type === "use_declaration" + ) { + imports.push(node.text.trim()); + } + for (const child of node.namedChildren) walk(child); + } + + walk(tree.rootNode); + return imports; +} + +function mapNodeType(tsType: string): ASTNodeType { + return NODE_TYPE_MAP[tsType] ?? "file"; +} + +function formatChunkForEmbedding( + chunk: Omit, +): string { + const header = [ + `// File: ${chunk.filePath}`, + `// Symbol: ${chunk.symbolName} (${chunk.nodeType})`, + `// Lines: ${chunk.startLine}-${chunk.endLine}`, + ]; + if (chunk.imports.length > 0) { + header.push(`// Imports: ${chunk.imports.slice(0, 8).join(", ")}`); + } + return `${header.join("\n")}\n${chunk.content}`; +} + +function walkAST( + node: Parser.SyntaxNode, + config: LanguageConfig, + filePath: string, + content: string, + fileImports: string[], + chunks: ASTCodeChunk[], + seen: Set, +) { + if (config.chunkNodeTypes.has(node.type)) { + const startLine = lineAt(content, node.startIndex); + const endLine = lineAt(content, node.endIndex); + const symbolName = extractSymbolName(node, config); + const nodeType = mapNodeType(node.type); + const dedupeKey = `${startLine}:${endLine}:${nodeType}:${symbolName}`; + + if (!seen.has(dedupeKey) && node.text.trim().length > 0) { + seen.add(dedupeKey); + chunks.push({ + id: "", + content: node.text, + filePath, + chunkIndex: chunks.length, + nodeType, + symbolName, + startLine, + endLine, + imports: fileImports, + }); + } + } + + for (const child of node.namedChildren) { + walkAST(child, config, filePath, content, fileImports, chunks, seen); + } +} + +export function parseFileToASTChunks( + filePath: string, + content: string, +): ASTCodeChunk[] { + const ext = filePath.includes(".") + ? `.${filePath.split(".").pop()}` + : ""; + const config = LANGUAGE_CONFIG[ext]; + const parser = getParser(ext); + + if (!config || !parser) { + return fallbackChunk(filePath, content); + } + + try { + const tree = parser.parse(content); + if (tree.rootNode.hasError) { + return fallbackChunk(filePath, content); + } + + const fileImports = extractImports(content, ext); + const chunks: ASTCodeChunk[] = []; + walkAST( + tree.rootNode, + config, + filePath, + content, + fileImports, + chunks, + new Set(), + ); + + if (chunks.length === 0) { + return fallbackChunk(filePath, content); + } + + return chunks.map((chunk, index) => { + const enriched = formatChunkForEmbedding(chunk); + return { + ...chunk, + chunkIndex: index, + content: enriched, + id: `${filePath}::${chunk.nodeType}::${chunk.symbolName}::${index}`, + }; + }); + } catch { + return fallbackChunk(filePath, content); + } +} + +const FALLBACK_CHUNK_SIZE = 800; +const FALLBACK_OVERLAP = 100; + +function fallbackChunk(filePath: string, content: string): ASTCodeChunk[] { + const chunks: ASTCodeChunk[] = []; + let i = 0; + let chunkIndex = 0; + + while (i < content.length) { + const slice = content.slice(i, i + FALLBACK_CHUNK_SIZE); + const startLine = content.slice(0, i).split("\n").length; + const endLine = startLine + slice.split("\n").length - 1; + + chunks.push({ + id: `${filePath}::file::chunk-${chunkIndex}::${chunkIndex}`, + content: `// File: ${filePath}\n// Symbol: chunk-${chunkIndex} (file)\n// Lines: ${startLine}-${endLine}\n${slice}`, + filePath, + chunkIndex, + nodeType: "file", + symbolName: `chunk-${chunkIndex}`, + startLine, + endLine, + imports: [], + }); + i += FALLBACK_CHUNK_SIZE - FALLBACK_OVERLAP; + chunkIndex++; + } + + return chunks; +} + +export function parseChangedFilesFromDiff(diff: string): string[] { + const files = new Set(); + + for (const line of diff.split("\n")) { + if (line.startsWith("+++ b/")) { + const filePath = line.slice(6).trim(); + if (filePath && filePath !== "/dev/null") files.add(filePath); + continue; + } + if (line.startsWith("--- a/")) { + const filePath = line.slice(6).trim(); + if (filePath && filePath !== "/dev/null") files.add(filePath); + } + } + + return [...files]; +} + +export function parseChangedSymbolsFromDiff(diff: string): string[] { + const symbols = new Set(); + const identifierPattern = + /(?:function|class|const|let|var|async\s+function|def|func|type|interface|struct|enum)\s+([A-Za-z_$][\w$]*)/g; + + for (const line of diff.split("\n")) { + if (!line.startsWith("+") || line.startsWith("+++")) continue; + const added = line.slice(1); + let match: RegExpExecArray | null; + while ((match = identifierPattern.exec(added)) !== null) { + symbols.add(match[1]); + } + } + + return [...symbols]; +}