mirror of
https://github.com/abhigyanpatwari/GitNexus.git
synced 2026-09-11 22:53:04 +00:00
feat(embeddings): add chunking support for long code nodes
Split large code nodes (functions, methods, classes) into chunks for better embedding quality. Includes AST-aware statement-boundary chunking for functions/methods, character-based sliding window fallback, and enriched metadata headers (repo, server, export status, description). - Add character-chunk.ts as a pure module with O(log n) line estimation - Add chunker.ts with AST-based and character-based chunking strategies - Add server-mapping.ts for microservice context enrichment - Extend EmbeddingConfig with chunkSize, overlap, maxDescriptionLength - Update CodeEmbedding schema with chunkIndex, startLine, endLine columns - Deduplicate semantic search results by nodeId (keep best chunk) - Add unit tests for chunker, text-generator, and embedding pipeline - Fix progress bar not restoring after console.log interception
This commit is contained in:
wangjichao
parent
ba5de0bde4
commit
0cc2ce5cf6
15 changed files with 1061 additions and 262 deletions
7
.gitignore
vendored
7
.gitignore
vendored
|
|
@ -95,4 +95,9 @@ GitNexus.sln
|
|||
|
||||
.swarm/
|
||||
|
||||
local_docs/
|
||||
local_docs/
|
||||
|
||||
.cursor
|
||||
.github
|
||||
openspec
|
||||
CLAUDE.md
|
||||
|
|
|
|||
|
|
@ -137,9 +137,11 @@ export const analyzeCommand = async (inputPath?: string, options?: AnalyzeOption
|
|||
const origLog = console.log.bind(console);
|
||||
const origWarn = console.warn.bind(console);
|
||||
const origError = console.error.bind(console);
|
||||
let barCurrentValue = 0;
|
||||
const barLog = (...args: any[]) => {
|
||||
process.stdout.write('\x1b[2K\r');
|
||||
origLog(args.map((a) => (typeof a === 'string' ? a : String(a))).join(' '));
|
||||
bar.update(barCurrentValue);
|
||||
};
|
||||
console.log = barLog;
|
||||
console.warn = barLog;
|
||||
|
|
@ -150,6 +152,7 @@ export const analyzeCommand = async (inputPath?: string, options?: AnalyzeOption
|
|||
let phaseStart = Date.now();
|
||||
|
||||
const updateBar = (value: number, phaseLabel: string) => {
|
||||
barCurrentValue = value;
|
||||
if (phaseLabel !== lastPhaseLabel) {
|
||||
lastPhaseLabel = phaseLabel;
|
||||
phaseStart = Date.now();
|
||||
|
|
|
|||
70
gitnexus/src/core/embeddings/character-chunk.ts
Normal file
70
gitnexus/src/core/embeddings/character-chunk.ts
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
/**
|
||||
* Character-based sliding window chunking (pure, no tree-sitter dependency)
|
||||
*/
|
||||
|
||||
export interface Chunk {
|
||||
text: string;
|
||||
chunkIndex: number;
|
||||
startLine: number;
|
||||
endLine: number;
|
||||
}
|
||||
|
||||
export const characterChunk = (
|
||||
content: string,
|
||||
startLine: number,
|
||||
endLine: number,
|
||||
chunkSize: number = 1200,
|
||||
overlap: number = 120,
|
||||
): Chunk[] => {
|
||||
if (content.length <= chunkSize) {
|
||||
return [{ text: content, chunkIndex: 0, startLine, endLine }];
|
||||
}
|
||||
|
||||
const chunks: Chunk[] = [];
|
||||
let offset = 0;
|
||||
const lines = content.split('\n');
|
||||
|
||||
// Build cumulative offset lookup for O(1) line estimation
|
||||
const lineOffsets = new Int32Array(lines.length);
|
||||
let acc = 0;
|
||||
for (let i = 0; i < lines.length; i++) {
|
||||
lineOffsets[i] = acc;
|
||||
acc += lines[i].length + 1;
|
||||
}
|
||||
|
||||
while (offset < content.length) {
|
||||
const end = Math.min(offset + chunkSize, content.length);
|
||||
const chunkText = content.slice(offset, end);
|
||||
|
||||
chunks.push({
|
||||
text: chunkText,
|
||||
chunkIndex: chunks.length,
|
||||
startLine: estimateLineFromOffset(lineOffsets, offset, startLine),
|
||||
endLine: estimateLineFromOffset(lineOffsets, end, startLine),
|
||||
});
|
||||
|
||||
offset = end - overlap;
|
||||
if (offset >= content.length) break;
|
||||
if (end >= content.length) break;
|
||||
if (offset <= (chunks.length > 1 ? end - chunkSize : 0)) {
|
||||
offset = end;
|
||||
}
|
||||
}
|
||||
|
||||
return chunks;
|
||||
};
|
||||
|
||||
const estimateLineFromOffset = (
|
||||
lineOffsets: Int32Array,
|
||||
charOffset: number,
|
||||
startLine: number,
|
||||
): number => {
|
||||
let lo = 0;
|
||||
let hi = lineOffsets.length - 1;
|
||||
while (lo < hi) {
|
||||
const mid = (lo + hi + 1) >> 1;
|
||||
if (lineOffsets[mid] <= charOffset) lo = mid;
|
||||
else hi = mid - 1;
|
||||
}
|
||||
return startLine + lo + 1;
|
||||
};
|
||||
169
gitnexus/src/core/embeddings/chunker.ts
Normal file
169
gitnexus/src/core/embeddings/chunker.ts
Normal file
|
|
@ -0,0 +1,169 @@
|
|||
/**
|
||||
* Chunker Module
|
||||
*
|
||||
* Splits code nodes into chunks for embedding.
|
||||
* - Function/Method/Constructor: AST-aware chunking by statement boundaries
|
||||
* - Other types: character-based sliding window fallback
|
||||
* - Short content (≤ chunkSize): no chunking
|
||||
*/
|
||||
|
||||
export { type Chunk, characterChunk } from './character-chunk.js';
|
||||
|
||||
import { characterChunk } from './character-chunk.js';
|
||||
import type { Chunk } from './character-chunk.js';
|
||||
|
||||
/**
|
||||
* Main chunkNode function: dispatches by label
|
||||
*/
|
||||
export const chunkNode = async (
|
||||
label: string,
|
||||
content: string,
|
||||
filePath: string,
|
||||
startLine: number,
|
||||
endLine: number,
|
||||
chunkSize: number = 1200,
|
||||
overlap: number = 120,
|
||||
): Promise<Chunk[]> => {
|
||||
// Content fits in one chunk — no splitting needed
|
||||
if (content.length <= chunkSize) {
|
||||
return [{ text: content, chunkIndex: 0, startLine, endLine }];
|
||||
}
|
||||
|
||||
// Only Function/Method get AST chunking
|
||||
if (label === 'Function' || label === 'Method') {
|
||||
try {
|
||||
const astChunks = await astChunk(content, filePath, startLine, endLine, chunkSize, overlap);
|
||||
if (astChunks.length > 0) return astChunks;
|
||||
} catch {
|
||||
// AST parsing failed — fall through to character fallback
|
||||
}
|
||||
}
|
||||
|
||||
// Character-based fallback for everything else
|
||||
return characterChunk(content, startLine, endLine, chunkSize, overlap);
|
||||
};
|
||||
|
||||
/**
|
||||
* AST-based chunking for Function/Method
|
||||
* Parse file, locate node by startLine/endLine, split body by statements
|
||||
*/
|
||||
const astChunk = async (
|
||||
content: string,
|
||||
filePath: string,
|
||||
startLine: number,
|
||||
endLine: number,
|
||||
chunkSize: number,
|
||||
overlap: number,
|
||||
): Promise<Chunk[]> => {
|
||||
const { getLanguageFromFilename } = await import('gitnexus-shared');
|
||||
const language = getLanguageFromFilename(filePath);
|
||||
if (!language) return [];
|
||||
|
||||
const { loadParser, loadLanguage, isLanguageAvailable } =
|
||||
await import('../tree-sitter/parser-loader.js');
|
||||
if (!isLanguageAvailable(language)) return [];
|
||||
const parser = await loadParser();
|
||||
await loadLanguage(language, filePath);
|
||||
|
||||
const tree = parser.parse(content);
|
||||
const root = tree.rootNode;
|
||||
|
||||
// Find the node matching our startLine/endLine (0-based in tree-sitter)
|
||||
const targetNode = findNodeByRange(root, startLine, endLine);
|
||||
if (!targetNode) return [];
|
||||
|
||||
// Get the body (statements) via childForFieldName('body')
|
||||
const bodyNode = targetNode.childForFieldName('body');
|
||||
if (!bodyNode) return [];
|
||||
|
||||
// Extract individual statements
|
||||
const statements: Array<{ text: string; startRow: number; endRow: number }> = [];
|
||||
for (let i = 0; i < bodyNode.namedChildCount; i++) {
|
||||
const child = bodyNode.namedChild(i);
|
||||
if (!child) continue;
|
||||
statements.push({
|
||||
text: child.text,
|
||||
startRow: child.startPosition.row,
|
||||
endRow: child.endPosition.row,
|
||||
});
|
||||
}
|
||||
|
||||
if (statements.length === 0) return [];
|
||||
|
||||
// Extract signature (everything before the body)
|
||||
const bodyStart = bodyNode.startIndex;
|
||||
const signature = content.slice(0, bodyStart).trim();
|
||||
|
||||
// Greedy merge statements into chunks
|
||||
const chunks: Chunk[] = [];
|
||||
let currentText = '';
|
||||
let currentStartRow = startLine;
|
||||
let currentEndRow = startLine;
|
||||
let isFirst = true;
|
||||
|
||||
for (const stmt of statements) {
|
||||
const candidateText = currentText ? `${currentText}\n${stmt.text}` : stmt.text;
|
||||
|
||||
// For first chunk, include signature
|
||||
const fullCandidate = isFirst ? `${signature}\n${candidateText}` : candidateText;
|
||||
|
||||
if (fullCandidate.length > chunkSize && currentText.length > 0) {
|
||||
// Current chunk is full — emit it
|
||||
chunks.push({
|
||||
text: isFirst ? `${signature}\n${currentText}` : currentText,
|
||||
chunkIndex: chunks.length,
|
||||
startLine: currentStartRow + 1, // 1-based
|
||||
endLine: currentEndRow + 1,
|
||||
});
|
||||
// Start new chunk with overlap
|
||||
currentText = overlapText(currentText, overlap) + '\n' + stmt.text;
|
||||
currentStartRow = stmt.startRow;
|
||||
isFirst = false;
|
||||
} else {
|
||||
currentText = candidateText;
|
||||
}
|
||||
currentEndRow = stmt.endRow;
|
||||
}
|
||||
|
||||
// Emit remaining chunk
|
||||
if (currentText.length > 0) {
|
||||
chunks.push({
|
||||
text: isFirst ? `${signature}\n${currentText}` : currentText,
|
||||
chunkIndex: chunks.length,
|
||||
startLine: currentStartRow + 1,
|
||||
endLine: currentEndRow + 1,
|
||||
});
|
||||
}
|
||||
|
||||
// Handle single statement longer than chunkSize — character fallback
|
||||
if (chunks.length === 1 && chunks[0].text.length > chunkSize) {
|
||||
return characterChunk(content, startLine, endLine, chunkSize, overlap);
|
||||
}
|
||||
|
||||
return chunks;
|
||||
};
|
||||
|
||||
/**
|
||||
* Find a node in the AST that matches the given line range
|
||||
*/
|
||||
const findNodeByRange = (node: any, startLine: number, endLine: number): any | null => {
|
||||
if (node.startPosition.row === startLine && node.endPosition.row === endLine) {
|
||||
return node;
|
||||
}
|
||||
|
||||
if (node.startPosition.row <= startLine && node.endPosition.row >= endLine) {
|
||||
for (let i = 0; i < node.namedChildCount; i++) {
|
||||
const child = node.namedChild(i);
|
||||
if (!child) continue;
|
||||
const found = findNodeByRange(child, startLine, endLine);
|
||||
if (found) return found;
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
};
|
||||
|
||||
const overlapText = (text: string, overlapSize: number): string => {
|
||||
if (text.length <= overlapSize) return text;
|
||||
return text.slice(text.length - overlapSize);
|
||||
};
|
||||
|
|
@ -3,9 +3,9 @@
|
|||
*
|
||||
* Orchestrates the background embedding process:
|
||||
* 1. Query embeddable nodes from LadybugDB
|
||||
* 2. Generate text representations
|
||||
* 3. Batch embed using transformers.js
|
||||
* 4. Update LadybugDB with embeddings
|
||||
* 2. Generate text representations with enriched metadata
|
||||
* 3. Chunk long nodes, batch embed
|
||||
* 4. Update LadybugDB with chunk-aware embeddings
|
||||
* 5. Create vector index for semantic search
|
||||
*/
|
||||
|
||||
|
|
@ -16,15 +16,19 @@ import {
|
|||
embeddingToArray,
|
||||
isEmbedderReady,
|
||||
} from './embedder.js';
|
||||
import { generateBatchEmbeddingTexts } from './text-generator.js';
|
||||
import { generateEmbeddingText } from './text-generator.js';
|
||||
import { chunkNode, characterChunk } from './chunker.js';
|
||||
import {
|
||||
type EmbeddingProgress,
|
||||
type EmbeddingConfig,
|
||||
type EmbeddableNode,
|
||||
type SemanticSearchResult,
|
||||
type ModelProgress,
|
||||
type EmbeddingContext,
|
||||
DEFAULT_EMBEDDING_CONFIG,
|
||||
EMBEDDABLE_LABELS,
|
||||
isShortLabel,
|
||||
LABELS_WITH_EXPORTED,
|
||||
} from './types.js';
|
||||
|
||||
const isDev = process.env.NODE_ENV === 'development';
|
||||
|
|
@ -36,32 +40,44 @@ export type EmbeddingProgressCallback = (progress: EmbeddingProgress) => void;
|
|||
|
||||
/**
|
||||
* Query all embeddable nodes from LadybugDB
|
||||
* Uses table-specific queries (File has different schema than code elements)
|
||||
* Uses table-specific queries for different label types
|
||||
*/
|
||||
const queryEmbeddableNodes = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
): Promise<EmbeddableNode[]> => {
|
||||
const allNodes: EmbeddableNode[] = [];
|
||||
|
||||
// Query each embeddable table with table-specific columns
|
||||
for (const label of EMBEDDABLE_LABELS) {
|
||||
try {
|
||||
let query: string;
|
||||
|
||||
if (label === 'File') {
|
||||
// File nodes don't have startLine/endLine
|
||||
if (label === 'Method') {
|
||||
// Method has parameterCount and returnType
|
||||
query = `
|
||||
MATCH (n:File)
|
||||
RETURN n.id AS id, n.name AS name, 'File' AS label,
|
||||
n.filePath AS filePath, n.content AS content
|
||||
MATCH (n:Method)
|
||||
RETURN n.id AS id, n.name AS name, 'Method' AS label,
|
||||
n.filePath AS filePath, n.content AS content,
|
||||
n.startLine AS startLine, n.endLine AS endLine,
|
||||
n.isExported AS isExported, n.description AS description,
|
||||
n.parameterCount AS parameterCount, n.returnType AS returnType
|
||||
`;
|
||||
} else if (LABELS_WITH_EXPORTED.has(label)) {
|
||||
// Function, Class, Interface have isExported and description
|
||||
query = `
|
||||
MATCH (n:\`${label}\`)
|
||||
RETURN n.id AS id, n.name AS name, '${label}' AS label,
|
||||
n.filePath AS filePath, n.content AS content,
|
||||
n.startLine AS startLine, n.endLine AS endLine,
|
||||
n.isExported AS isExported, n.description AS description
|
||||
`;
|
||||
} else {
|
||||
// Code elements have startLine/endLine
|
||||
// Multi-language tables (Struct, Enum, etc.) — have description but no isExported
|
||||
query = `
|
||||
MATCH (n:${label})
|
||||
RETURN n.id AS id, n.name AS name, '${label}' AS label,
|
||||
MATCH (n:\`${label}\`)
|
||||
RETURN n.id AS id, n.name AS name, '${label}' AS label,
|
||||
n.filePath AS filePath, n.content AS content,
|
||||
n.startLine AS startLine, n.endLine AS endLine
|
||||
n.startLine AS startLine, n.endLine AS endLine,
|
||||
n.description AS description
|
||||
`;
|
||||
}
|
||||
|
||||
|
|
@ -75,10 +91,17 @@ const queryEmbeddableNodes = async (
|
|||
content: row.content ?? row[4] ?? '',
|
||||
startLine: row.startLine ?? row[5],
|
||||
endLine: row.endLine ?? row[6],
|
||||
isExported: row.isExported ?? row[7],
|
||||
description: row.description ?? (label === 'Method' ? row[8] : row[7]),
|
||||
...(label === 'Method'
|
||||
? {
|
||||
parameterCount: row.parameterCount ?? row[9],
|
||||
returnType: row.returnType ?? row[10],
|
||||
}
|
||||
: {}),
|
||||
});
|
||||
}
|
||||
} catch (error) {
|
||||
// Table might not exist or be empty, continue
|
||||
if (isDev) {
|
||||
console.warn(`Query for ${label} nodes failed:`, error);
|
||||
}
|
||||
|
|
@ -89,40 +112,47 @@ const queryEmbeddableNodes = async (
|
|||
};
|
||||
|
||||
/**
|
||||
* Batch INSERT embeddings into separate CodeEmbedding table
|
||||
* Using a separate lightweight table avoids copy-on-write overhead
|
||||
* that occurs when UPDATEing nodes with large content fields
|
||||
* Batch INSERT chunk-aware embeddings into CodeEmbedding table
|
||||
*/
|
||||
const batchInsertEmbeddings = async (
|
||||
export const batchInsertEmbeddings = async (
|
||||
executeWithReusedStatement: (
|
||||
cypher: string,
|
||||
paramsList: Array<Record<string, any>>,
|
||||
) => Promise<void>,
|
||||
updates: Array<{ id: string; embedding: number[] }>,
|
||||
updates: Array<{
|
||||
nodeId: string;
|
||||
chunkIndex: number;
|
||||
startLine: number;
|
||||
endLine: number;
|
||||
embedding: number[];
|
||||
}>,
|
||||
): Promise<void> => {
|
||||
// INSERT into separate embedding table - much more memory efficient!
|
||||
const cypher = `CREATE (e:CodeEmbedding {nodeId: $nodeId, embedding: $embedding})`;
|
||||
const paramsList = updates.map((u) => ({ nodeId: u.id, embedding: u.embedding }));
|
||||
const cypher = `CREATE (e:CodeEmbedding {id: $id, nodeId: $nodeId, chunkIndex: $chunkIndex, startLine: $startLine, endLine: $endLine, embedding: $embedding})`;
|
||||
const paramsList = updates.map((u) => ({
|
||||
id: `${u.nodeId}:${u.chunkIndex}`,
|
||||
nodeId: u.nodeId,
|
||||
chunkIndex: u.chunkIndex,
|
||||
startLine: u.startLine,
|
||||
endLine: u.endLine,
|
||||
embedding: u.embedding,
|
||||
}));
|
||||
await executeWithReusedStatement(cypher, paramsList);
|
||||
};
|
||||
|
||||
/**
|
||||
* Create the vector index for semantic search
|
||||
* Now indexes the separate CodeEmbedding table
|
||||
*/
|
||||
let vectorExtensionLoaded = false;
|
||||
|
||||
const createVectorIndex = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
): Promise<void> => {
|
||||
// LadybugDB v0.15+ requires explicit VECTOR extension loading (once per session)
|
||||
if (!vectorExtensionLoaded) {
|
||||
try {
|
||||
await executeQuery('INSTALL VECTOR');
|
||||
await executeQuery('LOAD EXTENSION VECTOR');
|
||||
vectorExtensionLoaded = true;
|
||||
} catch {
|
||||
// Extension may already be loaded — CREATE_VECTOR_INDEX will fail clearly if not
|
||||
vectorExtensionLoaded = true;
|
||||
}
|
||||
}
|
||||
|
|
@ -134,7 +164,6 @@ const createVectorIndex = async (
|
|||
try {
|
||||
await executeQuery(cypher);
|
||||
} catch (error) {
|
||||
// Index might already exist
|
||||
if (isDev) {
|
||||
console.warn('Vector index creation warning:', error);
|
||||
}
|
||||
|
|
@ -149,6 +178,7 @@ const createVectorIndex = async (
|
|||
* @param onProgress - Callback for progress updates
|
||||
* @param config - Optional configuration override
|
||||
* @param skipNodeIds - Optional set of node IDs that already have embeddings (incremental mode)
|
||||
* @param context - Optional repo/server context for metadata enrichment
|
||||
*/
|
||||
export const runEmbeddingPipeline = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
|
|
@ -159,6 +189,7 @@ export const runEmbeddingPipeline = async (
|
|||
onProgress: EmbeddingProgressCallback,
|
||||
config: Partial<EmbeddingConfig> = {},
|
||||
skipNodeIds?: Set<string>,
|
||||
context?: EmbeddingContext,
|
||||
): Promise<void> => {
|
||||
const finalConfig = { ...DEFAULT_EMBEDDING_CONFIG, ...config };
|
||||
|
||||
|
|
@ -194,6 +225,14 @@ export const runEmbeddingPipeline = async (
|
|||
// Phase 2: Query embeddable nodes
|
||||
let nodes = await queryEmbeddableNodes(executeQuery);
|
||||
|
||||
// Apply context metadata
|
||||
if (context?.repoName) {
|
||||
for (const node of nodes) {
|
||||
node.repoName = context.repoName;
|
||||
node.serverName = context.serverName;
|
||||
}
|
||||
}
|
||||
|
||||
// Incremental mode: filter out nodes that already have embeddings
|
||||
if (skipNodeIds && skipNodeIds.size > 0) {
|
||||
const beforeCount = nodes.length;
|
||||
|
|
@ -221,10 +260,12 @@ export const runEmbeddingPipeline = async (
|
|||
return;
|
||||
}
|
||||
|
||||
// Phase 3: Batch embed nodes
|
||||
// Phase 3: Chunk + embed nodes
|
||||
const batchSize = finalConfig.batchSize;
|
||||
const totalBatches = Math.ceil(totalNodes / batchSize);
|
||||
const chunkSize = finalConfig.chunkSize;
|
||||
const overlap = finalConfig.overlap;
|
||||
let processedNodes = 0;
|
||||
let totalChunks = 0;
|
||||
|
||||
onProgress({
|
||||
phase: 'embedding',
|
||||
|
|
@ -232,39 +273,100 @@ export const runEmbeddingPipeline = async (
|
|||
nodesProcessed: 0,
|
||||
totalNodes,
|
||||
currentBatch: 0,
|
||||
totalBatches,
|
||||
totalBatches: Math.ceil(totalNodes / batchSize),
|
||||
});
|
||||
|
||||
for (let batchIndex = 0; batchIndex < totalBatches; batchIndex++) {
|
||||
const start = batchIndex * batchSize;
|
||||
const end = Math.min(start + batchSize, totalNodes);
|
||||
const batch = nodes.slice(start, end);
|
||||
// Process in batches of nodes
|
||||
for (let batchIndex = 0; batchIndex < totalNodes; batchIndex += batchSize) {
|
||||
const batch = nodes.slice(batchIndex, batchIndex + batchSize);
|
||||
|
||||
// Generate texts for this batch
|
||||
const texts = generateBatchEmbeddingTexts(batch, finalConfig);
|
||||
// Chunk each node and generate text
|
||||
const allTexts: string[] = [];
|
||||
const allUpdates: Array<{
|
||||
nodeId: string;
|
||||
chunkIndex: number;
|
||||
startLine: number;
|
||||
endLine: number;
|
||||
}> = [];
|
||||
|
||||
// Embed the batch
|
||||
const embeddings = await embedBatch(texts);
|
||||
for (const node of batch) {
|
||||
const isShort = isShortLabel(node.label);
|
||||
const startLine = node.startLine ?? 0;
|
||||
const endLine = node.endLine ?? 0;
|
||||
|
||||
// Update LadybugDB with embeddings
|
||||
const updates = batch.map((node, i) => ({
|
||||
id: node.id,
|
||||
embedding: embeddingToArray(embeddings[i]),
|
||||
}));
|
||||
let chunks: Array<{ text: string; chunkIndex: number; startLine: number; endLine: number }>;
|
||||
if (isShort) {
|
||||
chunks = [{ text: node.content, chunkIndex: 0, startLine, endLine }];
|
||||
} else {
|
||||
try {
|
||||
chunks = await chunkNode(
|
||||
node.label,
|
||||
node.content,
|
||||
node.filePath,
|
||||
startLine,
|
||||
endLine,
|
||||
chunkSize,
|
||||
overlap,
|
||||
);
|
||||
} catch (chunkErr) {
|
||||
if (isDev) {
|
||||
console.warn(
|
||||
`⚠️ AST chunking failed for ${node.label} "${node.name}" (${node.filePath}), falling back to character-based chunking:`,
|
||||
chunkErr,
|
||||
);
|
||||
}
|
||||
chunks = characterChunk(node.content, startLine, endLine, chunkSize, overlap);
|
||||
}
|
||||
}
|
||||
|
||||
await batchInsertEmbeddings(executeWithReusedStatement, updates);
|
||||
for (const chunk of chunks) {
|
||||
const text = generateEmbeddingText(node, chunk.text, finalConfig);
|
||||
allTexts.push(text);
|
||||
allUpdates.push({
|
||||
nodeId: node.id,
|
||||
chunkIndex: chunk.chunkIndex,
|
||||
startLine: chunk.startLine,
|
||||
endLine: chunk.endLine,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Embed chunk texts in sub-batches to control memory
|
||||
const EMBED_SUB_BATCH = 8;
|
||||
for (let si = 0; si < allTexts.length; si += EMBED_SUB_BATCH) {
|
||||
const subTexts = allTexts.slice(si, si + EMBED_SUB_BATCH);
|
||||
const subUpdates = allUpdates.slice(si, si + EMBED_SUB_BATCH);
|
||||
|
||||
let embeddings: Float32Array[];
|
||||
try {
|
||||
embeddings = await embedBatch(subTexts);
|
||||
} catch (embedErr) {
|
||||
console.error(
|
||||
`❌ embedBatch failed for ${subTexts.length} texts (first: "${subTexts[0]?.substring(0, 80)}..."):`,
|
||||
embedErr,
|
||||
);
|
||||
throw embedErr;
|
||||
}
|
||||
|
||||
const dbUpdates = subUpdates.map((u, i) => ({
|
||||
...u,
|
||||
embedding: embeddingToArray(embeddings[i]),
|
||||
}));
|
||||
|
||||
await batchInsertEmbeddings(executeWithReusedStatement, dbUpdates);
|
||||
}
|
||||
|
||||
processedNodes += batch.length;
|
||||
totalChunks += allUpdates.length;
|
||||
|
||||
// Report progress (20-90% for embedding phase)
|
||||
const embeddingProgress = 20 + (processedNodes / totalNodes) * 70;
|
||||
onProgress({
|
||||
phase: 'embedding',
|
||||
percent: Math.round(embeddingProgress),
|
||||
nodesProcessed: processedNodes,
|
||||
totalNodes,
|
||||
currentBatch: batchIndex + 1,
|
||||
totalBatches,
|
||||
currentBatch: Math.floor(batchIndex / batchSize) + 1,
|
||||
totalBatches: Math.ceil(totalNodes / batchSize),
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -282,7 +384,6 @@ export const runEmbeddingPipeline = async (
|
|||
|
||||
await createVectorIndex(executeQuery);
|
||||
|
||||
// Complete
|
||||
onProgress({
|
||||
phase: 'ready',
|
||||
percent: 100,
|
||||
|
|
@ -291,7 +392,9 @@ export const runEmbeddingPipeline = async (
|
|||
});
|
||||
|
||||
if (isDev) {
|
||||
console.log('✅ Embedding pipeline complete!');
|
||||
console.log(
|
||||
`✅ Embedding pipeline complete! (${totalChunks} chunks from ${totalNodes} nodes)`,
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
|
||||
|
|
@ -311,15 +414,7 @@ export const runEmbeddingPipeline = async (
|
|||
};
|
||||
|
||||
/**
|
||||
* Perform semantic search using the vector index
|
||||
*
|
||||
* Uses CodeEmbedding table and queries each node table to get metadata
|
||||
*
|
||||
* @param executeQuery - Function to execute Cypher queries
|
||||
* @param query - Search query text
|
||||
* @param k - Number of results to return (default: 10)
|
||||
* @param maxDistance - Maximum distance threshold (default: 0.5)
|
||||
* @returns Array of search results ordered by relevance
|
||||
* Perform semantic search using the vector index with chunk deduplication
|
||||
*/
|
||||
export const semanticSearch = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
|
|
@ -331,19 +426,19 @@ export const semanticSearch = async (
|
|||
throw new Error('Embedding model not initialized. Run embedding pipeline first.');
|
||||
}
|
||||
|
||||
// Embed the query
|
||||
const queryEmbedding = await embedText(query);
|
||||
const queryVec = embeddingToArray(queryEmbedding);
|
||||
const queryVecStr = `[${queryVec.join(',')}]`;
|
||||
|
||||
// Query the vector index on CodeEmbedding to get nodeIds and distances
|
||||
// Query vector index — get nodeId, chunkIndex, distance
|
||||
const vectorQuery = `
|
||||
CALL QUERY_VECTOR_INDEX('CodeEmbedding', 'code_embedding_idx',
|
||||
CALL QUERY_VECTOR_INDEX('CodeEmbedding', 'code_embedding_idx',
|
||||
CAST(${queryVecStr} AS FLOAT[${queryVec.length}]), ${k})
|
||||
YIELD node AS emb, distance
|
||||
WITH emb, distance
|
||||
WHERE distance < ${maxDistance}
|
||||
RETURN emb.nodeId AS nodeId, distance
|
||||
RETURN emb.nodeId AS nodeId, emb.chunkIndex AS chunkIndex,
|
||||
emb.startLine AS startLine, emb.endLine AS endLine, distance
|
||||
ORDER BY distance
|
||||
`;
|
||||
|
||||
|
|
@ -353,15 +448,35 @@ export const semanticSearch = async (
|
|||
return [];
|
||||
}
|
||||
|
||||
// Deduplicate by nodeId — keep chunk with smallest distance
|
||||
const bestChunks = new Map<
|
||||
string,
|
||||
{ chunkIndex: number; startLine: number; endLine: number; distance: number }
|
||||
>();
|
||||
|
||||
for (const row of embResults) {
|
||||
const nodeId = row.nodeId ?? row[0];
|
||||
const chunkIndex = row.chunkIndex ?? row[1] ?? 0;
|
||||
const startLine = row.startLine ?? row[2] ?? 0;
|
||||
const endLine = row.endLine ?? row[3] ?? 0;
|
||||
const distance = row.distance ?? row[4];
|
||||
|
||||
const existing = bestChunks.get(nodeId);
|
||||
if (!existing || distance < existing.distance) {
|
||||
bestChunks.set(nodeId, { chunkIndex, startLine, endLine, distance });
|
||||
}
|
||||
}
|
||||
|
||||
// Group results by label for batched metadata queries
|
||||
const byLabel = new Map<string, Array<{ nodeId: string; distance: number }>>();
|
||||
for (const embRow of embResults) {
|
||||
const nodeId = embRow.nodeId ?? embRow[0];
|
||||
const distance = embRow.distance ?? embRow[1];
|
||||
const byLabel = new Map<
|
||||
string,
|
||||
Array<{ nodeId: string; distance: number } & Record<string, any>>
|
||||
>();
|
||||
for (const [nodeId, chunk] of bestChunks) {
|
||||
const labelEndIdx = nodeId.indexOf(':');
|
||||
const label = labelEndIdx > 0 ? nodeId.substring(0, labelEndIdx) : 'Unknown';
|
||||
if (!byLabel.has(label)) byLabel.set(label, []);
|
||||
byLabel.get(label)!.push({ nodeId, distance });
|
||||
byLabel.get(label)!.push({ nodeId, ...chunk });
|
||||
}
|
||||
|
||||
// Batch-fetch metadata per label
|
||||
|
|
@ -370,19 +485,11 @@ export const semanticSearch = async (
|
|||
for (const [label, items] of byLabel) {
|
||||
const idList = items.map((i) => `'${i.nodeId.replace(/'/g, "''")}'`).join(', ');
|
||||
try {
|
||||
let nodeQuery: string;
|
||||
if (label === 'File') {
|
||||
nodeQuery = `
|
||||
MATCH (n:File) WHERE n.id IN [${idList}]
|
||||
RETURN n.id AS id, n.name AS name, n.filePath AS filePath
|
||||
`;
|
||||
} else {
|
||||
nodeQuery = `
|
||||
MATCH (n:${label}) WHERE n.id IN [${idList}]
|
||||
RETURN n.id AS id, n.name AS name, n.filePath AS filePath,
|
||||
n.startLine AS startLine, n.endLine AS endLine
|
||||
`;
|
||||
}
|
||||
const nodeQuery = `
|
||||
MATCH (n:\`${label}\`) WHERE n.id IN [${idList}]
|
||||
RETURN n.id AS id, n.name AS name, n.filePath AS filePath,
|
||||
n.startLine AS startLine, n.endLine AS endLine
|
||||
`;
|
||||
const nodeRows = await executeQuery(nodeQuery);
|
||||
const rowMap = new Map<string, any>();
|
||||
for (const row of nodeRows) {
|
||||
|
|
@ -398,8 +505,8 @@ export const semanticSearch = async (
|
|||
label,
|
||||
filePath: nodeRow.filePath ?? nodeRow[2] ?? '',
|
||||
distance: item.distance,
|
||||
startLine: label !== 'File' ? (nodeRow.startLine ?? nodeRow[3]) : undefined,
|
||||
endLine: label !== 'File' ? (nodeRow.endLine ?? nodeRow[4]) : undefined,
|
||||
startLine: item.startLine,
|
||||
endLine: item.endLine,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -408,7 +515,6 @@ export const semanticSearch = async (
|
|||
}
|
||||
}
|
||||
|
||||
// Re-sort by distance since batch queries may have mixed order
|
||||
results.sort((a, b) => a.distance - b.distance);
|
||||
|
||||
return results;
|
||||
|
|
@ -416,16 +522,6 @@ export const semanticSearch = async (
|
|||
|
||||
/**
|
||||
* Semantic search with graph expansion (flattened results)
|
||||
*
|
||||
* Note: With multi-table schema, graph traversal is simplified.
|
||||
* Returns semantic matches with their metadata.
|
||||
* For full graph traversal, use execute_vector_cypher tool directly.
|
||||
*
|
||||
* @param executeQuery - Function to execute Cypher queries
|
||||
* @param query - Search query text
|
||||
* @param k - Number of initial semantic matches (default: 5)
|
||||
* @param _hops - Unused (kept for API compatibility).
|
||||
* @returns Semantic matches with metadata
|
||||
*/
|
||||
export const semanticSearchWithContext = async (
|
||||
executeQuery: (cypher: string) => Promise<any[]>,
|
||||
|
|
@ -433,8 +529,6 @@ export const semanticSearchWithContext = async (
|
|||
k: number = 5,
|
||||
_hops: number = 1,
|
||||
): Promise<any[]> => {
|
||||
// For multi-table schema, just return semantic search results
|
||||
// Graph traversal is complex with separate tables - use execute_vector_cypher instead
|
||||
const results = await semanticSearch(executeQuery, query, k, 0.5);
|
||||
|
||||
return results.map((r) => ({
|
||||
|
|
|
|||
37
gitnexus/src/core/embeddings/server-mapping.ts
Normal file
37
gitnexus/src/core/embeddings/server-mapping.ts
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
/**
|
||||
* Server Mapping Configuration
|
||||
*
|
||||
* Reads ~/.gitnexus/server-mapping.json to map repo names to service names.
|
||||
* Used in embedding text to enrich metadata with microservice context.
|
||||
*/
|
||||
|
||||
import fs from 'fs/promises';
|
||||
import path from 'path';
|
||||
import os from 'os';
|
||||
|
||||
const MAPPING_FILE = path.join(os.homedir(), '.gitnexus', 'server-mapping.json');
|
||||
|
||||
let cachedMapping: Record<string, string> | null = null;
|
||||
|
||||
/**
|
||||
* Read the server mapping file and return the serverName for a given repoName.
|
||||
* Returns undefined if no mapping exists.
|
||||
*/
|
||||
export const readServerMapping = async (repoName: string): Promise<string | undefined> => {
|
||||
try {
|
||||
if (!cachedMapping) {
|
||||
const raw = await fs.readFile(MAPPING_FILE, 'utf-8');
|
||||
cachedMapping = JSON.parse(raw);
|
||||
}
|
||||
return cachedMapping[repoName];
|
||||
} catch {
|
||||
return undefined;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* Clear the cached mapping (useful for testing or after file changes)
|
||||
*/
|
||||
export const clearServerMappingCache = (): void => {
|
||||
cachedMapping = null;
|
||||
};
|
||||
|
|
@ -1,206 +1,239 @@
|
|||
/**
|
||||
* Text Generator Module
|
||||
*
|
||||
* Pure functions to generate embedding text from code nodes.
|
||||
* Combines node metadata with code snippets for semantic matching.
|
||||
* Generates enriched embedding text from code nodes with metadata.
|
||||
* Supports chunkable labels (Function/Method with AST chunking),
|
||||
* Class-specific structural text, and short-node direct embed.
|
||||
*/
|
||||
|
||||
import type { EmbeddableNode, EmbeddingConfig } from './types.js';
|
||||
import { DEFAULT_EMBEDDING_CONFIG } from './types.js';
|
||||
import { DEFAULT_EMBEDDING_CONFIG, isShortLabel } from './types.js';
|
||||
|
||||
/**
|
||||
* Extract the filename from a file path
|
||||
* Truncate description to max length at sentence/word boundary
|
||||
*/
|
||||
const getFileName = (filePath: string): string => {
|
||||
const parts = filePath.split('/');
|
||||
return parts[parts.length - 1] || filePath;
|
||||
};
|
||||
const truncateDescription = (text: string, maxLength: number): string => {
|
||||
if (text.length <= maxLength) return text;
|
||||
|
||||
/**
|
||||
* Extract the directory path from a file path
|
||||
*/
|
||||
const getDirectory = (filePath: string): string => {
|
||||
const parts = filePath.split('/');
|
||||
parts.pop();
|
||||
return parts.join('/') || '';
|
||||
};
|
||||
const truncated = text.slice(0, maxLength);
|
||||
|
||||
/**
|
||||
* Truncate content to max length, preserving word boundaries
|
||||
*/
|
||||
const truncateContent = (content: string, maxLength: number): string => {
|
||||
if (content.length <= maxLength) {
|
||||
return content;
|
||||
// Try sentence boundary (. ! ?)
|
||||
const sentenceEnd = Math.max(
|
||||
truncated.lastIndexOf('. '),
|
||||
truncated.lastIndexOf('! '),
|
||||
truncated.lastIndexOf('? '),
|
||||
);
|
||||
if (sentenceEnd > maxLength * 0.5) {
|
||||
return truncated.slice(0, sentenceEnd + 1);
|
||||
}
|
||||
|
||||
// Find last space before maxLength to avoid cutting words
|
||||
const truncated = content.slice(0, maxLength);
|
||||
// Try word boundary
|
||||
const lastSpace = truncated.lastIndexOf(' ');
|
||||
|
||||
if (lastSpace > maxLength * 0.8) {
|
||||
return truncated.slice(0, lastSpace) + '...';
|
||||
if (lastSpace > maxLength * 0.5) {
|
||||
return truncated.slice(0, lastSpace);
|
||||
}
|
||||
|
||||
return truncated + '...';
|
||||
return truncated;
|
||||
};
|
||||
|
||||
/**
|
||||
* Clean code content for embedding
|
||||
* Removes excessive whitespace while preserving structure
|
||||
*/
|
||||
const cleanContent = (content: string): string => {
|
||||
return (
|
||||
content
|
||||
// Normalize line endings
|
||||
.replace(/\r\n/g, '\n')
|
||||
// Remove excessive blank lines (more than 2)
|
||||
.replace(/\n{3,}/g, '\n\n')
|
||||
// Trim each line
|
||||
.split('\n')
|
||||
.map((line) => line.trimEnd())
|
||||
.join('\n')
|
||||
.trim()
|
||||
);
|
||||
return content
|
||||
.replace(/\r\n/g, '\n')
|
||||
.replace(/\n{3,}/g, '\n\n')
|
||||
.split('\n')
|
||||
.map((line) => line.trimEnd())
|
||||
.join('\n')
|
||||
.trim();
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for a Function node
|
||||
* Build metadata header for a node
|
||||
*/
|
||||
const generateFunctionText = (node: EmbeddableNode, maxSnippetLength: number): string => {
|
||||
const parts: string[] = [`Function: ${node.name}`, `File: ${getFileName(node.filePath)}`];
|
||||
const buildMetadataHeader = (node: EmbeddableNode, config: Partial<EmbeddingConfig>): string => {
|
||||
const parts: string[] = [];
|
||||
|
||||
const dir = getDirectory(node.filePath);
|
||||
if (dir) {
|
||||
parts.push(`Directory: ${dir}`);
|
||||
// Label + name
|
||||
parts.push(`${node.label}: ${node.name}`);
|
||||
|
||||
// Repo name
|
||||
if (node.repoName) {
|
||||
parts.push(`Repo: ${node.repoName}`);
|
||||
}
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const snippet = truncateContent(cleanedContent, maxSnippetLength);
|
||||
parts.push('', snippet);
|
||||
// Server name (optional)
|
||||
if (node.serverName) {
|
||||
parts.push(`Server: ${node.serverName}`);
|
||||
}
|
||||
|
||||
// Full file path
|
||||
parts.push(`Path: ${node.filePath}`);
|
||||
|
||||
// Export status
|
||||
if (node.isExported !== undefined) {
|
||||
parts.push(`Export: ${node.isExported}`);
|
||||
}
|
||||
|
||||
// Description (truncated)
|
||||
if (node.description) {
|
||||
const maxLen = config.maxDescriptionLength ?? DEFAULT_EMBEDDING_CONFIG.maxDescriptionLength;
|
||||
const truncated = truncateDescription(node.description, maxLen);
|
||||
if (truncated) {
|
||||
parts.push(truncated);
|
||||
}
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for a Class node
|
||||
* Generate embedding text for Function/Method nodes
|
||||
* Includes metadata header + code body (chunk text passed separately)
|
||||
*/
|
||||
const generateClassText = (node: EmbeddableNode, maxSnippetLength: number): string => {
|
||||
const parts: string[] = [`Class: ${node.name}`, `File: ${getFileName(node.filePath)}`];
|
||||
const generateFunctionText = (
|
||||
node: EmbeddableNode,
|
||||
codeBody: string,
|
||||
config: Partial<EmbeddingConfig>,
|
||||
): string => {
|
||||
const header = buildMetadataHeader(node, config);
|
||||
const cleaned = cleanContent(codeBody);
|
||||
return `${header}\n\n${cleaned}`;
|
||||
};
|
||||
|
||||
const dir = getDirectory(node.filePath);
|
||||
if (dir) {
|
||||
parts.push(`Directory: ${dir}`);
|
||||
/**
|
||||
* Generate embedding text for Class nodes
|
||||
* Signature + properties + method name list only (no method bodies)
|
||||
*
|
||||
* NOTE: Method/property regex is currently tuned for JS/TS syntax.
|
||||
* Multi-language support (Python, Kotlin, Rust, etc.) is a TODO.
|
||||
*/
|
||||
const generateClassText = (
|
||||
node: EmbeddableNode,
|
||||
codeBody: string,
|
||||
config: Partial<EmbeddingConfig>,
|
||||
): string => {
|
||||
const header = buildMetadataHeader(node, config);
|
||||
const parts: string[] = [header];
|
||||
|
||||
// Extract method names and properties from content
|
||||
const cleaned = cleanContent(codeBody);
|
||||
const lines = cleaned.split('\n');
|
||||
|
||||
const methods: string[] = [];
|
||||
const properties: string[] = [];
|
||||
const classBodyLines: string[] = [];
|
||||
let inClass = false;
|
||||
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
// Detect class opening (JS/TS only — Python `class Foo:`, Kotlin `class Foo` etc. are not yet handled)
|
||||
if (
|
||||
trimmed.match(/^(?:export\s+)?(?:abstract\s+)?class\s/) ||
|
||||
trimmed.startsWith('data class ')
|
||||
) {
|
||||
inClass = true;
|
||||
classBodyLines.push(trimmed);
|
||||
continue;
|
||||
}
|
||||
if (!inClass) {
|
||||
classBodyLines.push(trimmed);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Extract method names
|
||||
const methodMatch = trimmed.match(
|
||||
/^(?:public|private|protected|static|async|abstract|\s)*\s*(\w+)\s*\(/,
|
||||
);
|
||||
if (methodMatch && !trimmed.startsWith('//') && !trimmed.startsWith('*')) {
|
||||
methods.push(methodMatch[1]);
|
||||
}
|
||||
|
||||
// Extract property declarations
|
||||
const propMatch = trimmed.match(
|
||||
/^(?:public|private|protected|static|readonly)\s+(\w+)\s*[=:(]/,
|
||||
);
|
||||
if (propMatch) {
|
||||
properties.push(propMatch[1]);
|
||||
}
|
||||
|
||||
// Keep class declaration + property lines (no method bodies)
|
||||
if (
|
||||
trimmed.match(/^(?:export\s+)?(?:abstract\s+)?class\s/) ||
|
||||
trimmed.startsWith('data class ') ||
|
||||
trimmed.startsWith('{') ||
|
||||
trimmed.startsWith('}') ||
|
||||
trimmed === '' ||
|
||||
propMatch ||
|
||||
trimmed.endsWith(';') ||
|
||||
!trimmed.includes('{')
|
||||
) {
|
||||
classBodyLines.push(trimmed);
|
||||
}
|
||||
}
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const snippet = truncateContent(cleanedContent, maxSnippetLength);
|
||||
parts.push('', snippet);
|
||||
if (methods.length > 0) parts.push(`Methods: ${methods.join(', ')}`);
|
||||
if (properties.length > 0) parts.push(`Properties: ${properties.join(', ')}`);
|
||||
|
||||
// Class declaration only (no method bodies)
|
||||
const declarationOnly = classBodyLines.join('\n').trim();
|
||||
if (declarationOnly) {
|
||||
parts.push('', declarationOnly);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for a Method node
|
||||
* Generate embedding text for short nodes (TypeAlias, Const, etc.)
|
||||
* No chunking, just metadata + full content
|
||||
*/
|
||||
const generateMethodText = (node: EmbeddableNode, maxSnippetLength: number): string => {
|
||||
const parts: string[] = [`Method: ${node.name}`, `File: ${getFileName(node.filePath)}`];
|
||||
|
||||
const dir = getDirectory(node.filePath);
|
||||
if (dir) {
|
||||
parts.push(`Directory: ${dir}`);
|
||||
}
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const snippet = truncateContent(cleanedContent, maxSnippetLength);
|
||||
parts.push('', snippet);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
const generateShortNodeText = (node: EmbeddableNode, config: Partial<EmbeddingConfig>): string => {
|
||||
const header = buildMetadataHeader(node, config);
|
||||
const cleaned = cleanContent(node.content);
|
||||
return `${header}\n\n${cleaned}`;
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for an Interface node
|
||||
* Generate embedding text for Interface/Struct/Enum/Trait/etc. (chunkable but non-function)
|
||||
*/
|
||||
const generateInterfaceText = (node: EmbeddableNode, maxSnippetLength: number): string => {
|
||||
const parts: string[] = [`Interface: ${node.name}`, `File: ${getFileName(node.filePath)}`];
|
||||
|
||||
const dir = getDirectory(node.filePath);
|
||||
if (dir) {
|
||||
parts.push(`Directory: ${dir}`);
|
||||
}
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const snippet = truncateContent(cleanedContent, maxSnippetLength);
|
||||
parts.push('', snippet);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for a File node
|
||||
* Uses file name and first N characters of content
|
||||
*/
|
||||
const generateFileText = (node: EmbeddableNode, maxSnippetLength: number): string => {
|
||||
const parts: string[] = [`File: ${node.name}`, `Path: ${node.filePath}`];
|
||||
|
||||
if (node.content) {
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
// For files, use a shorter snippet since they can be very long
|
||||
const snippet = truncateContent(cleanedContent, Math.min(maxSnippetLength, 300));
|
||||
parts.push('', snippet);
|
||||
}
|
||||
|
||||
return parts.join('\n');
|
||||
const generateChunkableNonFunctionText = (
|
||||
node: EmbeddableNode,
|
||||
codeBody: string,
|
||||
config: Partial<EmbeddingConfig>,
|
||||
): string => {
|
||||
const header = buildMetadataHeader(node, config);
|
||||
const cleaned = cleanContent(codeBody);
|
||||
return `${header}\n\n${cleaned}`;
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for any embeddable node
|
||||
* Dispatches to the appropriate generator based on node label
|
||||
*
|
||||
* @param node - The node to generate text for
|
||||
* @param config - Optional configuration for max snippet length
|
||||
* @returns Text suitable for embedding
|
||||
*/
|
||||
export const generateEmbeddingText = (
|
||||
node: EmbeddableNode,
|
||||
codeBody: string,
|
||||
config: Partial<EmbeddingConfig> = {},
|
||||
): string => {
|
||||
const maxSnippetLength = config.maxSnippetLength ?? DEFAULT_EMBEDDING_CONFIG.maxSnippetLength;
|
||||
|
||||
switch (node.label) {
|
||||
case 'Function':
|
||||
return generateFunctionText(node, maxSnippetLength);
|
||||
case 'Class':
|
||||
return generateClassText(node, maxSnippetLength);
|
||||
case 'Method':
|
||||
return generateMethodText(node, maxSnippetLength);
|
||||
case 'Interface':
|
||||
return generateInterfaceText(node, maxSnippetLength);
|
||||
case 'File':
|
||||
return generateFileText(node, maxSnippetLength);
|
||||
default:
|
||||
// Fallback for any other embeddable type
|
||||
return `${node.label}: ${node.name}\nPath: ${node.filePath}`;
|
||||
if (isShortLabel(node.label)) {
|
||||
return generateShortNodeText(node, config);
|
||||
}
|
||||
|
||||
if (node.label === 'Class') {
|
||||
return generateClassText(node, codeBody, config);
|
||||
}
|
||||
|
||||
if (node.label === 'Function' || node.label === 'Method') {
|
||||
return generateFunctionText(node, codeBody, config);
|
||||
}
|
||||
|
||||
// Other chunkable types (Interface, Struct, Enum, Trait, Impl, Macro, Namespace)
|
||||
return generateChunkableNonFunctionText(node, codeBody, config);
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding texts for a batch of nodes
|
||||
*
|
||||
* @param nodes - Array of nodes to generate text for
|
||||
* @param config - Optional configuration
|
||||
* @returns Array of texts in the same order as input nodes
|
||||
* Export truncation helper for testing
|
||||
*/
|
||||
export const generateBatchEmbeddingTexts = (
|
||||
nodes: EmbeddableNode[],
|
||||
config: Partial<EmbeddingConfig> = {},
|
||||
): string[] => {
|
||||
return nodes.map((node) => generateEmbeddingText(node, config));
|
||||
};
|
||||
export { truncateDescription };
|
||||
|
|
|
|||
|
|
@ -5,10 +5,39 @@
|
|||
*/
|
||||
|
||||
/**
|
||||
* Node labels that should be embedded for semantic search
|
||||
* These are code elements that benefit from semantic matching
|
||||
* Node labels that need chunking (have code body, potentially long)
|
||||
*/
|
||||
export const EMBEDDABLE_LABELS = ['Function', 'Class', 'Method', 'Interface', 'File'] as const;
|
||||
export const CHUNKABLE_LABELS = [
|
||||
'Function',
|
||||
'Method',
|
||||
'Class',
|
||||
'Interface',
|
||||
'Struct',
|
||||
'Enum',
|
||||
'Trait',
|
||||
'Impl',
|
||||
'Macro',
|
||||
'Namespace',
|
||||
] as const;
|
||||
|
||||
/**
|
||||
* Node labels that are short (no chunking needed, embed directly)
|
||||
*/
|
||||
export const SHORT_LABELS = [
|
||||
'TypeAlias',
|
||||
'Typedef',
|
||||
'Const',
|
||||
'Property',
|
||||
'Record',
|
||||
'Union',
|
||||
'Static',
|
||||
'Variable',
|
||||
] as const;
|
||||
|
||||
/**
|
||||
* All embeddable labels (union of CHUNKABLE + SHORT)
|
||||
*/
|
||||
export const EMBEDDABLE_LABELS = [...CHUNKABLE_LABELS, ...SHORT_LABELS] as const;
|
||||
|
||||
export type EmbeddableLabel = (typeof EMBEDDABLE_LABELS)[number];
|
||||
|
||||
|
|
@ -18,6 +47,29 @@ export type EmbeddableLabel = (typeof EMBEDDABLE_LABELS)[number];
|
|||
export const isEmbeddableLabel = (label: string): label is EmbeddableLabel =>
|
||||
EMBEDDABLE_LABELS.includes(label as EmbeddableLabel);
|
||||
|
||||
/**
|
||||
* Check if a label needs chunking
|
||||
*/
|
||||
export const isChunkableLabel = (label: string): boolean =>
|
||||
(CHUNKABLE_LABELS as readonly string[]).includes(label);
|
||||
|
||||
/**
|
||||
* Check if a label is a short type (no chunking)
|
||||
*/
|
||||
export const isShortLabel = (label: string): boolean =>
|
||||
(SHORT_LABELS as readonly string[]).includes(label);
|
||||
|
||||
/**
|
||||
* Node labels that have isExported column in their schema
|
||||
*/
|
||||
export const LABELS_WITH_EXPORTED = new Set([
|
||||
'Function',
|
||||
'Class',
|
||||
'Interface',
|
||||
'Method',
|
||||
'CodeElement',
|
||||
]) as ReadonlySet<string>;
|
||||
|
||||
/**
|
||||
* Embedding pipeline phases
|
||||
*/
|
||||
|
|
@ -57,6 +109,12 @@ export interface EmbeddingConfig {
|
|||
device: 'auto' | 'dml' | 'cuda' | 'cpu' | 'wasm';
|
||||
/** Maximum characters of code snippet to include */
|
||||
maxSnippetLength: number;
|
||||
/** Maximum code chunk size in characters (for chunking long code) */
|
||||
chunkSize: number;
|
||||
/** Overlap between chunks in characters */
|
||||
overlap: number;
|
||||
/** Maximum description length in characters */
|
||||
maxDescriptionLength: number;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -70,6 +128,9 @@ export const DEFAULT_EMBEDDING_CONFIG: EmbeddingConfig = {
|
|||
dimensions: 384,
|
||||
device: 'auto',
|
||||
maxSnippetLength: 500,
|
||||
chunkSize: 1200,
|
||||
overlap: 120,
|
||||
maxDescriptionLength: 150,
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -96,6 +157,31 @@ export interface EmbeddableNode {
|
|||
content: string;
|
||||
startLine?: number;
|
||||
endLine?: number;
|
||||
isExported?: boolean;
|
||||
description?: string;
|
||||
parameterCount?: number;
|
||||
returnType?: string;
|
||||
repoName?: string;
|
||||
serverName?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Cached embedding entry restored from LadybugDB before a graph rebuild
|
||||
*/
|
||||
export interface CachedEmbedding {
|
||||
nodeId: string;
|
||||
chunkIndex: number;
|
||||
startLine: number;
|
||||
endLine: number;
|
||||
embedding: number[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Context info for embedding pipeline (repo/server metadata enrichment)
|
||||
*/
|
||||
export interface EmbeddingContext {
|
||||
repoName?: string;
|
||||
serverName?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import {
|
|||
NodeTableName,
|
||||
} from './schema.js';
|
||||
import { streamAllCSVsToDisk } from './csv-generator.js';
|
||||
import type { CachedEmbedding } from '../embeddings/types.js';
|
||||
|
||||
let db: lbug.Database | null = null;
|
||||
let conn: lbug.Connection | null = null;
|
||||
|
|
@ -727,30 +728,49 @@ export const getLbugStats = async (): Promise<{ nodes: number; edges: number }>
|
|||
* Load cached embeddings from LadybugDB before a rebuild.
|
||||
* Returns all embedding vectors so they can be re-inserted after the graph is reloaded,
|
||||
* avoiding expensive re-embedding of unchanged nodes.
|
||||
*
|
||||
* Detects old schema (no chunkIndex column) and returns empty cache to trigger rebuild.
|
||||
*/
|
||||
export const loadCachedEmbeddings = async (): Promise<{
|
||||
embeddingNodeIds: Set<string>;
|
||||
embeddings: Array<{ nodeId: string; embedding: number[] }>;
|
||||
embeddings: CachedEmbedding[];
|
||||
}> => {
|
||||
if (!conn) {
|
||||
return { embeddingNodeIds: new Set(), embeddings: [] };
|
||||
}
|
||||
|
||||
const embeddingNodeIds = new Set<string>();
|
||||
const embeddings: Array<{ nodeId: string; embedding: number[] }> = [];
|
||||
const embeddings: CachedEmbedding[] = [];
|
||||
|
||||
try {
|
||||
// Schema migration detection: query with new columns to verify schema version.
|
||||
// Old schema only had (nodeId, embedding); new schema adds (id, chunkIndex, startLine, endLine).
|
||||
// If the query fails (column missing), we return empty cache to force a full rebuild.
|
||||
try {
|
||||
const check = await conn.query(
|
||||
`MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex LIMIT 1`,
|
||||
);
|
||||
const checkResult = Array.isArray(check) ? check[0] : check;
|
||||
await checkResult.getAll();
|
||||
} catch {
|
||||
return { embeddingNodeIds: new Set(), embeddings: [] };
|
||||
}
|
||||
|
||||
const rows = await conn.query(
|
||||
`MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.embedding AS embedding`,
|
||||
`MATCH (e:${EMBEDDING_TABLE_NAME}) RETURN e.nodeId AS nodeId, e.chunkIndex AS chunkIndex, e.startLine AS startLine, e.endLine AS endLine, e.embedding AS embedding`,
|
||||
);
|
||||
const result = Array.isArray(rows) ? rows[0] : rows;
|
||||
for (const row of await result.getAll()) {
|
||||
const nodeId = String(row.nodeId ?? row[0] ?? '');
|
||||
if (!nodeId) continue;
|
||||
embeddingNodeIds.add(nodeId);
|
||||
const embedding = row.embedding ?? row[1];
|
||||
const embedding = row.embedding ?? row[4];
|
||||
if (embedding) {
|
||||
embeddings.push({
|
||||
nodeId,
|
||||
chunkIndex: Number(row.chunkIndex ?? row[1] ?? 0),
|
||||
startLine: Number(row.startLine ?? row[2] ?? 0),
|
||||
endLine: Number(row.endLine ?? row[3] ?? 0),
|
||||
embedding: Array.isArray(embedding)
|
||||
? embedding.map(Number)
|
||||
: Array.from(embedding as any).map(Number),
|
||||
|
|
|
|||
|
|
@ -438,9 +438,13 @@ export const EMBEDDING_DIMS = _rawDims;
|
|||
|
||||
export const EMBEDDING_SCHEMA = `
|
||||
CREATE NODE TABLE ${EMBEDDING_TABLE_NAME} (
|
||||
id STRING,
|
||||
nodeId STRING,
|
||||
chunkIndex INT32,
|
||||
startLine INT64,
|
||||
endLine INT64,
|
||||
embedding FLOAT[${EMBEDDING_DIMS}],
|
||||
PRIMARY KEY (nodeId)
|
||||
PRIMARY KEY (id)
|
||||
)`;
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ import {
|
|||
cleanupOldKuzuFiles,
|
||||
} from '../storage/repo-manager.js';
|
||||
import { getCurrentCommit, hasGitDir } from '../storage/git.js';
|
||||
import type { CachedEmbedding } from './embeddings/types.js';
|
||||
import { generateAIContextFiles } from '../cli/ai-context.js';
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
|
@ -136,7 +137,7 @@ export async function runFullAnalysis(
|
|||
|
||||
// ── Cache embeddings from existing index before rebuild ────────────
|
||||
let cachedEmbeddingNodeIds = new Set<string>();
|
||||
let cachedEmbeddings: Array<{ nodeId: string; embedding: number[] }> = [];
|
||||
let cachedEmbeddings: CachedEmbedding[] = [];
|
||||
|
||||
if (options.embeddings && existingMeta && !options.force) {
|
||||
try {
|
||||
|
|
@ -214,15 +215,13 @@ export async function runFullAnalysis(
|
|||
cachedEmbeddingNodeIds = new Set();
|
||||
} else {
|
||||
progress('embeddings', 88, `Restoring ${cachedEmbeddings.length} cached embeddings...`);
|
||||
const { batchInsertEmbeddings: batchInsert } =
|
||||
await import('./embeddings/embedding-pipeline.js');
|
||||
const EMBED_BATCH = 200;
|
||||
for (let i = 0; i < cachedEmbeddings.length; i += EMBED_BATCH) {
|
||||
const batch = cachedEmbeddings.slice(i, i + EMBED_BATCH);
|
||||
const paramsList = batch.map((e) => ({ nodeId: e.nodeId, embedding: e.embedding }));
|
||||
try {
|
||||
await executeWithReusedStatement(
|
||||
`CREATE (e:CodeEmbedding {nodeId: $nodeId, embedding: $embedding})`,
|
||||
paramsList,
|
||||
);
|
||||
await batchInsert(executeWithReusedStatement, batch);
|
||||
} catch {
|
||||
/* some may fail if node was removed, that's fine */
|
||||
}
|
||||
|
|
@ -249,6 +248,9 @@ export async function runFullAnalysis(
|
|||
httpMode ? 'Connecting to embedding endpoint...' : 'Loading embedding model...',
|
||||
);
|
||||
const { runEmbeddingPipeline } = await import('./embeddings/embedding-pipeline.js');
|
||||
const { readServerMapping } = await import('./embeddings/server-mapping.js');
|
||||
const projectName = path.basename(repoPath);
|
||||
const serverName = await readServerMapping(projectName);
|
||||
await runEmbeddingPipeline(
|
||||
executeQuery,
|
||||
executeWithReusedStatement,
|
||||
|
|
@ -264,6 +266,7 @@ export async function runFullAnalysis(
|
|||
},
|
||||
{},
|
||||
cachedEmbeddingNodeIds.size > 0 ? cachedEmbeddingNodeIds : undefined,
|
||||
{ repoName: projectName, serverName },
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -807,12 +807,13 @@ export class LocalBackend {
|
|||
const queryVecStr = `[${queryVec.join(',')}]`;
|
||||
|
||||
const vectorQuery = `
|
||||
CALL QUERY_VECTOR_INDEX('CodeEmbedding', 'code_embedding_idx',
|
||||
CALL QUERY_VECTOR_INDEX('CodeEmbedding', 'code_embedding_idx',
|
||||
CAST(${queryVecStr} AS FLOAT[${dims}]), ${limit})
|
||||
YIELD node AS emb, distance
|
||||
WITH emb, distance
|
||||
WHERE distance < 0.6
|
||||
RETURN emb.nodeId AS nodeId, distance
|
||||
RETURN emb.nodeId AS nodeId, emb.chunkIndex AS chunkIndex,
|
||||
emb.startLine AS chunkStartLine, emb.endLine AS chunkEndLine, distance
|
||||
ORDER BY distance
|
||||
`;
|
||||
|
||||
|
|
@ -820,12 +821,26 @@ export class LocalBackend {
|
|||
|
||||
if (embResults.length === 0) return [];
|
||||
|
||||
// Deduplicate by nodeId — keep chunk with smallest distance
|
||||
const bestChunks = new Map<
|
||||
string,
|
||||
{ chunkIndex: number; startLine: number; endLine: number; distance: number }
|
||||
>();
|
||||
for (const row of embResults) {
|
||||
const nodeId = row.nodeId ?? row[0];
|
||||
const chunkIndex = row.chunkIndex ?? row[1] ?? 0;
|
||||
const startLine = row.chunkStartLine ?? row[2] ?? 0;
|
||||
const endLine = row.chunkEndLine ?? row[3] ?? 0;
|
||||
const distance = row.distance ?? row[4];
|
||||
const existing = bestChunks.get(nodeId);
|
||||
if (!existing || distance < existing.distance) {
|
||||
bestChunks.set(nodeId, { chunkIndex, startLine, endLine, distance });
|
||||
}
|
||||
}
|
||||
|
||||
const results: any[] = [];
|
||||
|
||||
for (const embRow of embResults) {
|
||||
const nodeId = embRow.nodeId ?? embRow[0];
|
||||
const distance = embRow.distance ?? embRow[1];
|
||||
|
||||
for (const [nodeId, chunk] of bestChunks) {
|
||||
const labelEndIdx = nodeId.indexOf(':');
|
||||
const label = labelEndIdx > 0 ? nodeId.substring(0, labelEndIdx) : 'Unknown';
|
||||
|
||||
|
|
@ -836,7 +851,7 @@ export class LocalBackend {
|
|||
const nodeQuery =
|
||||
label === 'File'
|
||||
? `MATCH (n:File {id: $nodeId}) RETURN n.name AS name, n.filePath AS filePath`
|
||||
: `MATCH (n:\`${label}\` {id: $nodeId}) RETURN n.name AS name, n.filePath AS filePath, n.startLine AS startLine, n.endLine AS endLine`;
|
||||
: `MATCH (n:\`${label}\` {id: $nodeId}) RETURN n.name AS name, n.filePath AS filePath`;
|
||||
|
||||
const nodeRows = await executeParameterized(repo.id, nodeQuery, { nodeId });
|
||||
if (nodeRows.length > 0) {
|
||||
|
|
@ -846,9 +861,9 @@ export class LocalBackend {
|
|||
name: nodeRow.name ?? nodeRow[0] ?? '',
|
||||
type: label,
|
||||
filePath: nodeRow.filePath ?? nodeRow[1] ?? '',
|
||||
distance,
|
||||
startLine: label !== 'File' ? (nodeRow.startLine ?? nodeRow[2]) : undefined,
|
||||
endLine: label !== 'File' ? (nodeRow.endLine ?? nodeRow[3]) : undefined,
|
||||
distance: chunk.distance,
|
||||
startLine: chunk.startLine,
|
||||
endLine: chunk.endLine,
|
||||
});
|
||||
}
|
||||
} catch {}
|
||||
|
|
|
|||
42
gitnexus/test/unit/chunker.test.ts
Normal file
42
gitnexus/test/unit/chunker.test.ts
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
/**
|
||||
* Unit tests for character chunking logic
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { characterChunk } from '../../src/core/embeddings/character-chunk.js';
|
||||
|
||||
describe('characterChunk', () => {
|
||||
it('returns single chunk when content fits', () => {
|
||||
const result = characterChunk('short content', 1, 5, 1200, 120);
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0].text).toBe('short content');
|
||||
expect(result[0].chunkIndex).toBe(0);
|
||||
expect(result[0].startLine).toBe(1);
|
||||
expect(result[0].endLine).toBe(5);
|
||||
});
|
||||
|
||||
it('splits long content into multiple chunks', () => {
|
||||
const longContent = 'a'.repeat(3000);
|
||||
const result = characterChunk(longContent, 1, 100, 1200, 120);
|
||||
expect(result.length).toBeGreaterThan(1);
|
||||
for (const chunk of result) {
|
||||
expect(chunk.text.length).toBeLessThanOrEqual(1200);
|
||||
}
|
||||
});
|
||||
|
||||
it('maintains sequential chunkIndex', () => {
|
||||
const longContent = 'x'.repeat(3000);
|
||||
const result = characterChunk(longContent, 1, 100, 1200, 120);
|
||||
for (let i = 0; i < result.length; i++) {
|
||||
expect(result[i].chunkIndex).toBe(i);
|
||||
}
|
||||
});
|
||||
|
||||
it('includes overlap between chunks', () => {
|
||||
const content = 'abcdefghij'.repeat(200);
|
||||
const result = characterChunk(content, 1, 50, 500, 50);
|
||||
if (result.length > 1) {
|
||||
const endOfFirst = result[0].text.slice(-50);
|
||||
expect(result[1].text.startsWith(endOfFirst)).toBe(true);
|
||||
}
|
||||
});
|
||||
});
|
||||
106
gitnexus/test/unit/embedding-chunking.test.ts
Normal file
106
gitnexus/test/unit/embedding-chunking.test.ts
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
/**
|
||||
* Integration test: Embedding chunking pipeline
|
||||
*
|
||||
* Tests the chunking + text generation pipeline together.
|
||||
*/
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { characterChunk } from '../../src/core/embeddings/character-chunk.js';
|
||||
import { generateEmbeddingText } from '../../src/core/embeddings/text-generator.js';
|
||||
import type { EmbeddableNode } from '../../src/core/embeddings/types.js';
|
||||
|
||||
describe('embedding-chunking integration', () => {
|
||||
const makeNode = (overrides: Partial<EmbeddableNode>): EmbeddableNode => ({
|
||||
id: 'Function:src/test.ts:test',
|
||||
name: 'test',
|
||||
label: 'Function',
|
||||
filePath: 'src/test.ts',
|
||||
content: '',
|
||||
startLine: 1,
|
||||
endLine: 10,
|
||||
...overrides,
|
||||
});
|
||||
|
||||
it('short function produces single chunk with metadata', () => {
|
||||
const node = makeNode({
|
||||
content: 'function hello() { return "world"; }',
|
||||
isExported: true,
|
||||
repoName: 'my-project',
|
||||
serverName: 'my-service',
|
||||
});
|
||||
|
||||
const chunks = characterChunk(node.content, 1, 3, 1200, 120);
|
||||
expect(chunks).toHaveLength(1);
|
||||
|
||||
const text = generateEmbeddingText(node, chunks[0].text);
|
||||
expect(text).toContain('Function: test');
|
||||
expect(text).toContain('Repo: my-project');
|
||||
expect(text).toContain('Server: my-service');
|
||||
expect(text).toContain('Export: true');
|
||||
expect(text).toContain('function hello()');
|
||||
});
|
||||
|
||||
it('long function produces multiple chunks', () => {
|
||||
const longContent = Array.from({ length: 100 }, (_, i) => ` const line${i} = ${i};`).join(
|
||||
'\n',
|
||||
);
|
||||
const node = makeNode({
|
||||
content: `function longFn() {\n${longContent}\n}`,
|
||||
startLine: 1,
|
||||
endLine: 102,
|
||||
});
|
||||
|
||||
const chunks = characterChunk(node.content, 1, 102, 1200, 120);
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
expect(chunks[0].chunkIndex).toBe(0);
|
||||
});
|
||||
|
||||
it('short labels (TypeAlias) skip chunking and embed directly', () => {
|
||||
const node = makeNode({
|
||||
label: 'TypeAlias',
|
||||
name: 'Result',
|
||||
content: 'type Result<T> = Success<T> | Error;',
|
||||
});
|
||||
|
||||
const chunks = characterChunk(node.content, 1, 1, 1200, 120);
|
||||
expect(chunks).toHaveLength(1);
|
||||
|
||||
const text = generateEmbeddingText(node, chunks[0].text);
|
||||
expect(text).toContain('TypeAlias: Result');
|
||||
expect(text).toContain('type Result<T> = Success<T> | Error;');
|
||||
});
|
||||
|
||||
it('long enum uses character fallback', () => {
|
||||
const enumContent = Array.from(
|
||||
{ length: 200 },
|
||||
(_, i) => ` Value${i} = "${'x'.repeat(20)}${i}",`,
|
||||
).join('\n');
|
||||
const node = makeNode({
|
||||
label: 'Enum',
|
||||
name: 'LargeEnum',
|
||||
content: `enum LargeEnum {\n${enumContent}\n}`,
|
||||
startLine: 1,
|
||||
endLine: 202,
|
||||
});
|
||||
|
||||
const chunks = characterChunk(node.content, 1, 202, 1200, 120);
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
});
|
||||
|
||||
it('metadata is present in every chunk', () => {
|
||||
const longContent = 'x'.repeat(3000);
|
||||
const node = makeNode({
|
||||
content: longContent,
|
||||
repoName: 'test-repo',
|
||||
});
|
||||
|
||||
const chunks = characterChunk(node.content, 1, 100, 1200, 120);
|
||||
expect(chunks.length).toBeGreaterThan(1);
|
||||
|
||||
for (const chunk of chunks) {
|
||||
const text = generateEmbeddingText(node, chunk.text);
|
||||
expect(text).toContain('Function: test');
|
||||
expect(text).toContain('Repo: test-repo');
|
||||
expect(text).toContain('Path: src/test.ts');
|
||||
}
|
||||
});
|
||||
});
|
||||
112
gitnexus/test/unit/text-generator.test.ts
Normal file
112
gitnexus/test/unit/text-generator.test.ts
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
import { describe, it, expect } from 'vitest';
|
||||
import {
|
||||
generateEmbeddingText,
|
||||
truncateDescription,
|
||||
} from '../../src/core/embeddings/text-generator.js';
|
||||
import type { EmbeddableNode } from '../../src/core/embeddings/types.js';
|
||||
|
||||
const baseNode: EmbeddableNode = {
|
||||
id: 'Function:src/utils.ts:parseJSON',
|
||||
name: 'parseJSON',
|
||||
label: 'Function',
|
||||
filePath: 'src/utils/parser.ts',
|
||||
content: 'function parseJSON(text: string): Result<any> {\n return JSON.parse(text);\n}',
|
||||
startLine: 10,
|
||||
endLine: 12,
|
||||
};
|
||||
|
||||
describe('text-generator', () => {
|
||||
describe('generateEmbeddingText', () => {
|
||||
it('includes metadata header for Function', () => {
|
||||
const node: EmbeddableNode = {
|
||||
...baseNode,
|
||||
isExported: true,
|
||||
repoName: 'backend-user-ms',
|
||||
};
|
||||
const text = generateEmbeddingText(node, node.content);
|
||||
expect(text).toContain('Function: parseJSON');
|
||||
expect(text).toContain('Repo: backend-user-ms');
|
||||
expect(text).toContain('Path: src/utils/parser.ts');
|
||||
expect(text).toContain('Export: true');
|
||||
expect(text).toContain('function parseJSON');
|
||||
});
|
||||
|
||||
it('includes Server line when serverName is set', () => {
|
||||
const node: EmbeddableNode = {
|
||||
...baseNode,
|
||||
repoName: 'backend-user-ms',
|
||||
serverName: 'user-service',
|
||||
};
|
||||
const text = generateEmbeddingText(node, node.content);
|
||||
expect(text).toContain('Server: user-service');
|
||||
});
|
||||
|
||||
it('omits Server line when serverName is undefined', () => {
|
||||
const text = generateEmbeddingText(baseNode, baseNode.content);
|
||||
expect(text).not.toContain('Server:');
|
||||
});
|
||||
|
||||
it('includes truncated description', () => {
|
||||
const node: EmbeddableNode = {
|
||||
...baseNode,
|
||||
description: 'This function parses JSON text and returns a typed result object.',
|
||||
};
|
||||
const text = generateEmbeddingText(node, node.content);
|
||||
expect(text).toContain('This function parses JSON text');
|
||||
});
|
||||
|
||||
it('generates short node text for TypeAlias without chunking', () => {
|
||||
const node: EmbeddableNode = {
|
||||
...baseNode,
|
||||
label: 'TypeAlias',
|
||||
name: 'Result',
|
||||
content: 'type Result<T> = Success<T> | Error;',
|
||||
};
|
||||
const text = generateEmbeddingText(node, node.content);
|
||||
expect(text).toContain('TypeAlias: Result');
|
||||
expect(text).toContain('type Result<T> = Success<T> | Error;');
|
||||
});
|
||||
|
||||
it('generates Class text with method names', () => {
|
||||
const node: EmbeddableNode = {
|
||||
...baseNode,
|
||||
label: 'Class',
|
||||
name: 'Parser',
|
||||
content: `class Parser {
|
||||
options: ParserOptions;
|
||||
private cache: Map<string, any>;
|
||||
parseJSON(text: string) { return JSON.parse(text); }
|
||||
validate() { return true; }
|
||||
}`,
|
||||
};
|
||||
const text = generateEmbeddingText(node, node.content);
|
||||
expect(text).toContain('Class: Parser');
|
||||
expect(text).toContain('Methods:');
|
||||
expect(text).toContain('parseJSON');
|
||||
expect(text).toContain('validate');
|
||||
expect(text).toContain('Properties:');
|
||||
expect(text).toContain('options');
|
||||
});
|
||||
});
|
||||
|
||||
describe('truncateDescription', () => {
|
||||
it('returns short text unchanged', () => {
|
||||
expect(truncateDescription('short text', 150)).toBe('short text');
|
||||
});
|
||||
|
||||
it('truncates at sentence boundary', () => {
|
||||
const text = 'First sentence. Second sentence. Third very long sentence that goes on and on.';
|
||||
const result = truncateDescription(text, 40);
|
||||
expect(result).toContain('First sentence');
|
||||
expect(result.length).toBeLessThan(text.length);
|
||||
});
|
||||
|
||||
it('truncates at word boundary when no sentence end', () => {
|
||||
const text =
|
||||
'this is a long description without any sentence ending punctuation marks at all';
|
||||
const result = truncateDescription(text, 30);
|
||||
expect(result.length).toBeLessThanOrEqual(30);
|
||||
expect(result.length).toBeLessThan(text.length);
|
||||
});
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Reference in a new issue