Skip to content

Commit c5dcd59

Browse files
github-actions[bot]github-actions[bot]
committed
feat(embeddings): graph-enriched strategy + context window overflow detection
Addresses #55 and #56 with a combined solution: - Add contextWindow to MODELS registry for all 7 embedding models - New 'structured' embedding strategy (default): uses callers/callees from the dependency graph, extracted parameters, and leading comments instead of raw source code (~100 tokens vs ~360 avg) - Detect and truncate symbols exceeding model context window, with user-facing warning and truncated_count in embedding_meta - Add --strategy flag to CLI embed command (structured|source) - Store strategy in embedding_meta for auditability - Export estimateTokens() and EMBEDDING_STRATEGIES for programmatic use Closes #55, closes #56 Impact: 5 functions changed, 3 affected
1 parent ce25da7 commit c5dcd59

4 files changed

Lines changed: 489 additions & 18 deletions

File tree

‎src/cli.js‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { buildGraph } from './builder.js';
88
import { loadConfig } from './config.js';
99
import { findCycles, formatCycles } from './cycles.js';
1010
import { findDbPath } from './db.js';
11-
import { buildEmbeddings, MODELS, search } from './embedder.js';
11+
import { buildEmbeddings, EMBEDDING_STRATEGIES, MODELS, search } from './embedder.js';
1212
import { exportDOT, exportJSON, exportMermaid } from './export.js';
1313
import { setVerbose } from './logger.js';
1414
import {
@@ -418,9 +418,12 @@ program
418418
console.log('\nAvailable embedding models:\n');
419419
for (const [key, config] of Object.entries(MODELS)) {
420420
const def = key === 'minilm' ? ' (default)' : '';
421-
console.log(` ${key.padEnd(12)} ${String(config.dim).padStart(4)}d ${config.desc}${def}`);
421+
const ctx = config.contextWindow ? `${config.contextWindow} ctx` : '';
422+
console.log(
423+
` ${key.padEnd(12)} ${String(config.dim).padStart(4)}d ${ctx.padEnd(9)} ${config.desc}${def}`,
424+
);
422425
}
423-
console.log('\nUsage: codegraph embed --model <name>');
426+
console.log('\nUsage: codegraph embed --model <name> --strategy <structured|source>');
424427
console.log(' codegraph search "query" --model <name>\n');
425428
});
426429

@@ -434,9 +437,20 @@ program
434437
'Embedding model: minilm (default), jina-small, jina-base, jina-code, nomic, nomic-v1.5, bge-large. Run `codegraph models` for details',
435438
'minilm',
436439
)
440+
.option(
441+
'-s, --strategy <name>',
442+
`Embedding strategy: ${EMBEDDING_STRATEGIES.join(', ')}. "structured" uses graph context (callers/callees), "source" embeds raw code`,
443+
'structured',
444+
)
437445
.action(async (dir, opts) => {
446+
if (!EMBEDDING_STRATEGIES.includes(opts.strategy)) {
447+
console.error(
448+
`Unknown strategy: ${opts.strategy}. Available: ${EMBEDDING_STRATEGIES.join(', ')}`,
449+
);
450+
process.exit(1);
451+
}
438452
const root = path.resolve(dir || '.');
439-
await buildEmbeddings(root, opts.model);
453+
await buildEmbeddings(root, opts.model, undefined, { strategy: opts.strategy });
440454
});
441455

442456
program

‎src/embedder.js‎

Lines changed: 163 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -26,47 +26,56 @@ export const MODELS = {
2626
minilm: {
2727
name: 'Xenova/all-MiniLM-L6-v2',
2828
dim: 384,
29+
contextWindow: 256,
2930
desc: 'Smallest, fastest (~23MB). General text.',
3031
quantized: true,
3132
},
3233
'jina-small': {
3334
name: 'Xenova/jina-embeddings-v2-small-en',
3435
dim: 512,
36+
contextWindow: 8192,
3537
desc: 'Small, good quality (~33MB). General text.',
3638
quantized: false,
3739
},
3840
'jina-base': {
3941
name: 'Xenova/jina-embeddings-v2-base-en',
4042
dim: 768,
43+
contextWindow: 8192,
4144
desc: 'Good quality (~137MB). General text, 8192 token context.',
4245
quantized: false,
4346
},
4447
'jina-code': {
4548
name: 'Xenova/jina-embeddings-v2-base-code',
4649
dim: 768,
50+
contextWindow: 8192,
4751
desc: 'Code-aware (~137MB). Trained on code+text, best for code search.',
4852
quantized: false,
4953
},
5054
nomic: {
5155
name: 'Xenova/nomic-embed-text-v1',
5256
dim: 768,
57+
contextWindow: 8192,
5358
desc: 'Good local quality (~137MB). 8192 context.',
5459
quantized: false,
5560
},
5661
'nomic-v1.5': {
5762
name: 'nomic-ai/nomic-embed-text-v1.5',
5863
dim: 768,
64+
contextWindow: 8192,
5965
desc: 'Improved nomic (~137MB). Matryoshka dimensions, 8192 context.',
6066
quantized: false,
6167
},
6268
'bge-large': {
6369
name: 'Xenova/bge-large-en-v1.5',
6470
dim: 1024,
71+
contextWindow: 512,
6572
desc: 'Best general retrieval (~335MB). Top MTEB scores.',
6673
quantized: false,
6774
},
6875
};
6976

77+
export const EMBEDDING_STRATEGIES = ['structured', 'source'];
78+
7079
export const DEFAULT_MODEL = 'minilm';
7180
const BATCH_SIZE_MAP = {
7281
minilm: 32,
@@ -89,6 +98,108 @@ function getModelConfig(modelKey) {
8998
return config;
9099
}
91100

101+
/**
102+
* Rough token estimate (~4 chars per token for code/English).
103+
* Conservative — avoids adding a tokenizer dependency.
104+
*/
105+
export function estimateTokens(text) {
106+
return Math.ceil(text.length / 4);
107+
}
108+
109+
/**
110+
* Extract leading comment text (JSDoc, //, #, etc.) above a function line.
111+
* Returns the cleaned comment text or null if none found.
112+
*/
113+
function extractLeadingComment(lines, fnLineIndex) {
114+
const raw = [];
115+
for (let i = fnLineIndex - 1; i >= Math.max(0, fnLineIndex - 15); i--) {
116+
const trimmed = lines[i].trim();
117+
if (/^(\/\/|\/\*|\*\/|\*|#|\/\/\/)/.test(trimmed)) {
118+
raw.unshift(trimmed);
119+
} else if (trimmed === '') {
120+
if (raw.length > 0) break;
121+
} else {
122+
break;
123+
}
124+
}
125+
if (raw.length === 0) return null;
126+
return raw
127+
.map((line) =>
128+
line
129+
.replace(/^\/\*\*?\s?|\*\/$/g, '') // opening /** or /* and closing */
130+
.replace(/^\*\s?/, '') // middle * lines
131+
.replace(/^\/\/\/?\s?/, '') // // or ///
132+
.replace(/^#\s?/, '') // # (Python/Ruby)
133+
.trim(),
134+
)
135+
.filter((l) => l.length > 0)
136+
.join(' ');
137+
}
138+
139+
/**
140+
* Build graph-enriched text for a symbol using dependency context.
141+
* Produces compact, semantic text (~100 tokens) instead of full source code.
142+
*/
143+
function buildStructuredText(node, file, lines, calleesStmt, callersStmt) {
144+
const readable = splitIdentifier(node.name);
145+
const parts = [`${node.kind} ${node.name} (${readable}) in ${file}`];
146+
const startLine = Math.max(0, node.line - 1);
147+
148+
// Extract parameters from signature (best-effort, single-line)
149+
const sigLine = lines[startLine] || '';
150+
const paramMatch = sigLine.match(/\(([^)]*)\)/);
151+
if (paramMatch?.[1]?.trim()) {
152+
parts.push(`Parameters: ${paramMatch[1].trim()}`);
153+
}
154+
155+
// Graph context: callees (capped at 10)
156+
const callees = calleesStmt.all(node.id);
157+
if (callees.length > 0) {
158+
parts.push(
159+
`Calls: ${callees
160+
.slice(0, 10)
161+
.map((c) => c.name)
162+
.join(', ')}`,
163+
);
164+
}
165+
166+
// Graph context: callers (capped at 10)
167+
const callers = callersStmt.all(node.id);
168+
if (callers.length > 0) {
169+
parts.push(
170+
`Called by: ${callers
171+
.slice(0, 10)
172+
.map((c) => c.name)
173+
.join(', ')}`,
174+
);
175+
}
176+
177+
// Leading comment (high semantic value) or first few lines of code
178+
const comment = extractLeadingComment(lines, startLine);
179+
if (comment) {
180+
parts.push(comment);
181+
} else {
182+
const endLine = Math.min(lines.length, startLine + 4);
183+
const snippet = lines.slice(startLine, endLine).join('\n').trim();
184+
if (snippet) parts.push(snippet);
185+
}
186+
187+
return parts.join('\n');
188+
}
189+
190+
/**
191+
* Build raw source-code text for a symbol (original strategy).
192+
*/
193+
function buildSourceText(node, file, lines) {
194+
const startLine = Math.max(0, node.line - 1);
195+
const endLine = node.end_line
196+
? Math.min(lines.length, node.end_line)
197+
: Math.min(lines.length, startLine + 15);
198+
const context = lines.slice(startLine, endLine).join('\n');
199+
const readable = splitIdentifier(node.name);
200+
return `${node.kind} ${node.name} (${readable}) in ${file}\n${context}`;
201+
}
202+
92203
/**
93204
* Lazy-load @huggingface/transformers.
94205
* This is an optional dependency — gives a clear error if not installed.
@@ -203,10 +314,14 @@ function initEmbeddingsSchema(db) {
203314

204315
/**
205316
* Build embeddings for all functions/methods/classes in the graph.
317+
* @param {string} rootDir - Project root directory
318+
* @param {string} modelKey - Model identifier from MODELS registry
319+
* @param {string} [customDbPath] - Override path to graph.db
320+
* @param {object} [options] - Embedding options
321+
* @param {string} [options.strategy='structured'] - 'structured' (graph-enriched) or 'source' (raw code)
206322
*/
207-
export async function buildEmbeddings(rootDir, modelKey, customDbPath) {
208-
// path already imported at top
209-
// fs already imported at top
323+
export async function buildEmbeddings(rootDir, modelKey, customDbPath, options = {}) {
324+
const strategy = options.strategy || 'structured';
210325
const dbPath = customDbPath || findDbPath(null);
211326

212327
const db = new Database(dbPath);
@@ -221,7 +336,24 @@ export async function buildEmbeddings(rootDir, modelKey, customDbPath) {
221336
)
222337
.all();
223338

224-
console.log(`Building embeddings for ${nodes.length} symbols...`);
339+
console.log(`Building embeddings for ${nodes.length} symbols (strategy: ${strategy})...`);
340+
341+
// Prepare graph-context queries for structured strategy
342+
let calleesStmt, callersStmt;
343+
if (strategy === 'structured') {
344+
calleesStmt = db.prepare(`
345+
SELECT DISTINCT n.name FROM edges e
346+
JOIN nodes n ON e.target_id = n.id
347+
WHERE e.source_id = ? AND e.kind = 'calls'
348+
ORDER BY n.name
349+
`);
350+
callersStmt = db.prepare(`
351+
SELECT DISTINCT n.name FROM edges e
352+
JOIN nodes n ON e.source_id = n.id
353+
WHERE e.target_id = ? AND e.kind = 'calls'
354+
ORDER BY n.name
355+
`);
356+
}
225357

226358
const byFile = new Map();
227359
for (const node of nodes) {
@@ -232,6 +364,9 @@ export async function buildEmbeddings(rootDir, modelKey, customDbPath) {
232364
const texts = [];
233365
const nodeIds = [];
234366
const previews = [];
367+
const config = getModelConfig(modelKey);
368+
const contextWindow = config.contextWindow;
369+
let overflowCount = 0;
235370

236371
for (const [file, fileNodes] of byFile) {
237372
const fullPath = path.join(rootDir, file);
@@ -244,20 +379,31 @@ export async function buildEmbeddings(rootDir, modelKey, customDbPath) {
244379
}
245380

246381
for (const node of fileNodes) {
247-
const startLine = Math.max(0, node.line - 1);
248-
const endLine = node.end_line
249-
? Math.min(lines.length, node.end_line)
250-
: Math.min(lines.length, startLine + 15);
251-
const context = lines.slice(startLine, endLine).join('\n');
252-
253-
const readable = splitIdentifier(node.name);
254-
const text = `${node.kind} ${node.name} (${readable}) in ${file}\n${context}`;
382+
let text =
383+
strategy === 'structured'
384+
? buildStructuredText(node, file, lines, calleesStmt, callersStmt)
385+
: buildSourceText(node, file, lines);
386+
387+
// Detect and handle context window overflow
388+
const tokens = estimateTokens(text);
389+
if (tokens > contextWindow) {
390+
overflowCount++;
391+
const maxChars = contextWindow * 4;
392+
text = text.slice(0, maxChars);
393+
}
394+
255395
texts.push(text);
256396
nodeIds.push(node.id);
257397
previews.push(`${node.name} (${node.kind}) -- ${file}:${node.line}`);
258398
}
259399
}
260400

401+
if (overflowCount > 0) {
402+
warn(
403+
`${overflowCount} symbol(s) exceeded model context window (${contextWindow} tokens) and were truncated`,
404+
);
405+
}
406+
261407
console.log(`Embedding ${texts.length} symbols...`);
262408
const { vectors, dim } = await embed(texts, modelKey);
263409

@@ -269,16 +415,19 @@ export async function buildEmbeddings(rootDir, modelKey, customDbPath) {
269415
for (let i = 0; i < vectors.length; i++) {
270416
insert.run(nodeIds[i], Buffer.from(vectors[i].buffer), previews[i]);
271417
}
272-
const config = getModelConfig(modelKey);
273418
insertMeta.run('model', config.name);
274419
insertMeta.run('dim', String(dim));
275420
insertMeta.run('count', String(vectors.length));
421+
insertMeta.run('strategy', strategy);
276422
insertMeta.run('built_at', new Date().toISOString());
423+
if (overflowCount > 0) {
424+
insertMeta.run('truncated_count', String(overflowCount));
425+
}
277426
});
278427
insertAll();
279428

280429
console.log(
281-
`\nStored ${vectors.length} embeddings (${dim}d, ${getModelConfig(modelKey).name}) in graph.db`,
430+
`\nStored ${vectors.length} embeddings (${dim}d, ${config.name}, strategy: ${strategy}) in graph.db`,
282431
);
283432
db.close();
284433
}

‎src/index.js‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ export {
2121
buildEmbeddings,
2222
cosineSim,
2323
DEFAULT_MODEL,
24+
EMBEDDING_STRATEGIES,
2425
embed,
26+
estimateTokens,
2527
MODELS,
2628
multiSearchData,
2729
search,

0 commit comments

Comments
 (0)