Codebase Rag: Build a Codebase-Aware AI Assistant That Does Not Hallucinate Your API

// 1. Extract semantic units from your codebase
import * as ts from 'typescript';
import * as fs from 'fs';
import * as path from 'path';

interface CodeChunk {
 id: string;
 type: 'function' | 'class' | 'interface' | 'example';
 name: string;
 signature: string;
 docstring: string;
 body: string; // First 20 lines
 filePath: string;
 embedding?: number[];
}

function extractChunks(filePath) {
 const source = fs.readFileSync(filePath, 'utf8');
 const sf = ts.createSourceFile(filePath, source, ts.ScriptTarget.Latest, true);
 const chunks = [];

 function visit(node) {
 if (ts.isFunctionDeclaration(node) && node.name) {
 chunks.push({
 id: filePath + '#' + node.name.text,
 type: 'function',
 name: node.name.text,
 signature: source.slice(node.pos, node.body?.pos || node.end).split('{')[0].trim(),
 docstring: getJSDoc(node, source),
 body: source.slice(node.pos, node.end).slice(0, 500),
 filePath
 });
 }
 ts.forEachChild(node, visit);
 }
 ts.forEachChild(sf, visit);
 return chunks;
}

// 2. Embed all chunks
async function buildIndex(codebaseRoot) {
 const files = getAllTsFiles(codebaseRoot);
 const allChunks = files.flatMap(f => extractChunks(f));

 // Batch embed for efficiency
 const embeddings = await Promise.all(
 chunk(allChunks, 100).map(batch =>
 embedBatch(batch.map(c => c.signature + '\n' + c.docstring + '\n' + c.body))
 )
 );

 return allChunks.map((c, i) => ({ ...c, embedding: embeddings[i] }));
}

// Retrieve relevant code chunks for a given query
async function retrieveContext(query, index, topK = 8) {
 const queryEmbedding = await embed(query);

 // Cosine similarity search
 const scored = index.map(chunk => ({
 ...chunk,
 score: cosineSimilarity(queryEmbedding, chunk.embedding)
 }));

 // Sort by relevance, return top K
 return scored
 .sort((a, b) => b.score - a.score)
 .slice(0, topK)
 .filter(c => c.score > 0.7); // Filter low-relevance results
}

// Format for LLM context injection
function formatContext(chunks) {
 return chunks.map(c =>
 `// File: ${c.filePath}\n// Type: ${c.type}\n${c.signature}\n${c.docstring}\n${c.body.slice(0, 200)}`
 ).join('\n\n---\n\n');
}

// Generate code with codebase context
async function generateWithRAG(query, index) {
 const relevantCode = await retrieveContext(query, index);
 const context = formatContext(relevantCode);

 return await client.messages.create({
 model: 'claude-opus-4-5',
 max_tokens: 2000,
 system: 'You are a code assistant for this specific codebase.'
 + '\n\nUse ONLY the following functions and patterns from the actual codebase:'
 + '\n\n' + context
 + '\n\nNever invent functions not shown above.',
 messages: [{ role: 'user', content: query }]
 });
}

-- PostgreSQL with pgvector extension (production setup)
CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE code_chunks (
 id TEXT PRIMARY KEY,
 type TEXT NOT NULL, -- function, class, interface, example
 name TEXT NOT NULL,
 signature TEXT NOT NULL,
 docstring TEXT,
 body TEXT NOT NULL,
 file_path TEXT NOT NULL,
 embedding vector(1536), -- text-embedding-3-small dimension
 updated_at TIMESTAMP DEFAULT NOW()
);

-- IVFFlat index for fast approximate nearest neighbor search
CREATE INDEX ON code_chunks USING ivfflat (embedding vector_cosine_ops)
 WITH (lists = 100);

-- Query: find top-5 most similar chunks to a query embedding
SELECT id, name, signature, docstring, body, file_path,
 1 - (embedding <=> $1) as similarity
FROM code_chunks
ORDER BY embedding <=> $1 -- cosine distance
LIMIT 5;

-- Incremental update: only re-embed changed files
UPDATE code_chunks SET embedding = $1, updated_at = NOW()
WHERE id = $2;

-- Check if file has changed before re-indexing:
SELECT MAX(updated_at) FROM code_chunks WHERE file_path = $1;