mirror of
https://github.com/abhigyanpatwari/GitNexus.git
synced 2026-08-28 05:25:25 +00:00
feat(embeddings): structural chunking with data-driven CHUNKING_RULES dispatch (#987)
* feat(embeddings): structural chunking with data-driven CHUNKING_RULES dispatch Replace hardcoded label comparisons with a CHUNKING_RULES lookup table that drives chunking strategy and text generation. Key changes: - Data-driven dispatch: CHUNKING_RULES table maps labels to chunking mode (ast-function / ast-declaration), prefix/suffix, field grouping, and structural text mode - Struct support: add Struct to AST declaration chunking with field grouping (same as Class) - Multi-chunk context: preceding chunk tail (prevTail) injected into embedding text for cross-chunk coherence - Version-gated hashes: EMBEDDING_TEXT_VERSION prefix in content hashes invalidates stale vectors when text template changes - Compact container context: first declaration line preserved in every structural chunk for identity * fix(embeddings): address PR review findings for CHUNKING_RULES refactor - Remove LABEL_ENUM from STRUCTURAL_LABELS to avoid wasted AST parses - Add maintenance note about extractStructuralNames and EMBEDDING_TEXT_VERSION - Clarify CHUNK_MODE_CHARACTER is a no-op in CHUNKING_RULES - Strengthen EMBEDDING_TEXT_VERSION test assertion to exact value --------- Co-authored-by: wangjichao <wangjichao@inke.cn>
This commit is contained in:
parent
d976038dc8
commit
2b7cff5fd2
8 changed files with 520 additions and 93 deletions
|
|
@ -13,6 +13,12 @@ import { characterChunk } from './character-chunk.js';
|
|||
import type { Chunk } from './character-chunk.js';
|
||||
import { ensureAndParse, findDeclarationNode, findFunctionNode } from './ast-utils.js';
|
||||
import { buildLineIndex, resolveChunkLines } from './line-index.js';
|
||||
import {
|
||||
CHUNKING_RULES,
|
||||
CHUNK_MODE_AST_DECLARATION,
|
||||
CHUNK_MODE_AST_FUNCTION,
|
||||
type ChunkingRule,
|
||||
} from './types.js';
|
||||
|
||||
/**
|
||||
* Main chunkNode function: dispatches by label
|
||||
|
|
@ -40,31 +46,39 @@ export const chunkNode = async (
|
|||
];
|
||||
}
|
||||
|
||||
// Only function-like labels get AST chunking
|
||||
if (label === 'Function' || label === 'Method' || label === 'Constructor') {
|
||||
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
|
||||
}
|
||||
const rule = CHUNKING_RULES[label];
|
||||
if (!rule) {
|
||||
return characterChunk(content, startLine, endLine, chunkSize, overlap);
|
||||
}
|
||||
|
||||
if (label === 'Class' || label === 'Interface') {
|
||||
try {
|
||||
const declarationChunks = await declarationChunk(
|
||||
label,
|
||||
try {
|
||||
if (rule.mode === CHUNK_MODE_AST_FUNCTION) {
|
||||
const astChunks = await astChunk(
|
||||
content,
|
||||
filePath,
|
||||
startLine,
|
||||
endLine,
|
||||
chunkSize,
|
||||
overlap,
|
||||
rule,
|
||||
);
|
||||
if (astChunks.length > 0) return astChunks;
|
||||
}
|
||||
|
||||
if (rule.mode === CHUNK_MODE_AST_DECLARATION) {
|
||||
const declarationChunks = await declarationChunk(
|
||||
content,
|
||||
filePath,
|
||||
startLine,
|
||||
endLine,
|
||||
chunkSize,
|
||||
overlap,
|
||||
rule,
|
||||
);
|
||||
if (declarationChunks.length > 0) return declarationChunks;
|
||||
} catch {
|
||||
// AST parsing failed — fall through to character fallback
|
||||
}
|
||||
} catch {
|
||||
// AST parsing failed — fall through to character fallback
|
||||
}
|
||||
|
||||
// Character-based fallback for everything else
|
||||
|
|
@ -83,6 +97,7 @@ const astChunk = async (
|
|||
endLine: number,
|
||||
chunkSize: number,
|
||||
overlap: number,
|
||||
rule: ChunkingRule,
|
||||
): Promise<Chunk[]> => {
|
||||
const tree = await ensureAndParse(content, filePath);
|
||||
if (!tree) return [];
|
||||
|
|
@ -121,8 +136,8 @@ const astChunk = async (
|
|||
statements,
|
||||
targetNode.startIndex,
|
||||
targetNode.endIndex,
|
||||
true,
|
||||
true,
|
||||
rule.includePrefix,
|
||||
rule.includeSuffix,
|
||||
);
|
||||
};
|
||||
|
||||
|
|
@ -145,13 +160,13 @@ const FIELD_LIKE_MEMBER_TYPES = new Set([
|
|||
]);
|
||||
|
||||
const declarationChunk = async (
|
||||
label: 'Class' | 'Interface',
|
||||
content: string,
|
||||
filePath: string,
|
||||
startLine: number,
|
||||
endLine: number,
|
||||
chunkSize: number,
|
||||
overlap: number,
|
||||
rule: ChunkingRule,
|
||||
): Promise<Chunk[]> => {
|
||||
const tree = await ensureAndParse(content, filePath);
|
||||
if (!tree) return [];
|
||||
|
|
@ -162,7 +177,7 @@ const declarationChunk = async (
|
|||
const bodyNode = getDeclarationBodyNode(targetNode);
|
||||
if (!bodyNode) return [];
|
||||
|
||||
const members = collectDeclarationUnits(bodyNode, label);
|
||||
const members = collectDeclarationUnits(bodyNode, rule.groupFields);
|
||||
if (members.length === 0) return [];
|
||||
|
||||
return chunkByUnits(
|
||||
|
|
@ -174,8 +189,8 @@ const declarationChunk = async (
|
|||
members,
|
||||
targetNode.startIndex,
|
||||
targetNode.endIndex,
|
||||
false,
|
||||
false,
|
||||
rule.includePrefix,
|
||||
rule.includeSuffix,
|
||||
);
|
||||
};
|
||||
|
||||
|
|
@ -237,14 +252,22 @@ const chunkByUnits = (
|
|||
|
||||
if (candidateEndOffset - chunkStartOffset > chunkSize) {
|
||||
const oversizedUnit = units[chunkStartUnitIdx];
|
||||
const oversizedStartOffset =
|
||||
chunkStartUnitIdx === 0 && includeContainerPrefixOnFirstChunk
|
||||
? containerStartOffset
|
||||
: oversizedUnit.startIndex;
|
||||
const oversizedEndOffset =
|
||||
chunkStartUnitIdx === units.length - 1 && includeContainerSuffixOnLastChunk
|
||||
? containerEndOffset
|
||||
: oversizedUnit.endIndex;
|
||||
const oversizedLineRange = resolveChunkLines(
|
||||
lineOffsets,
|
||||
oversizedUnit.startIndex,
|
||||
oversizedUnit.endIndex,
|
||||
oversizedStartOffset,
|
||||
oversizedEndOffset,
|
||||
baseStartLine,
|
||||
);
|
||||
const oversizedChunks = characterChunk(
|
||||
content.slice(oversizedUnit.startIndex, oversizedUnit.endIndex),
|
||||
content.slice(oversizedStartOffset, oversizedEndOffset),
|
||||
oversizedLineRange.startLine,
|
||||
oversizedLineRange.endLine,
|
||||
chunkSize,
|
||||
|
|
@ -252,8 +275,8 @@ const chunkByUnits = (
|
|||
).map((chunk, offsetIdx) => ({
|
||||
...chunk,
|
||||
chunkIndex: chunks.length + offsetIdx,
|
||||
startOffset: chunk.startOffset + oversizedUnit.startIndex,
|
||||
endOffset: chunk.endOffset + oversizedUnit.startIndex,
|
||||
startOffset: chunk.startOffset + oversizedStartOffset,
|
||||
endOffset: chunk.endOffset + oversizedStartOffset,
|
||||
}));
|
||||
chunks.push(...oversizedChunks);
|
||||
chunkStartUnitIdx += 1;
|
||||
|
|
@ -325,7 +348,7 @@ const getDeclarationBodyNode = (node: any): any | null => {
|
|||
|
||||
const collectDeclarationUnits = (
|
||||
bodyNode: any,
|
||||
label: 'Class' | 'Interface',
|
||||
groupFields: boolean,
|
||||
): Array<{ startIndex: number; endIndex: number }> => {
|
||||
const members: Array<{ startIndex: number; endIndex: number; groupable: boolean }> = [];
|
||||
|
||||
|
|
@ -335,7 +358,7 @@ const collectDeclarationUnits = (
|
|||
members.push({
|
||||
startIndex: child.startIndex,
|
||||
endIndex: child.endIndex,
|
||||
groupable: label === 'Class' && FIELD_LIKE_MEMBER_TYPES.has(child.type),
|
||||
groupable: groupFields && FIELD_LIKE_MEMBER_TYPES.has(child.type),
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ import {
|
|||
DEFAULT_EMBEDDING_CONFIG,
|
||||
EMBEDDABLE_LABELS,
|
||||
isShortLabel,
|
||||
LABEL_METHOD,
|
||||
LABELS_WITH_EXPORTED,
|
||||
STRUCTURAL_LABELS,
|
||||
collectBestChunks,
|
||||
|
|
@ -43,6 +44,12 @@ import {
|
|||
import { loadVectorExtension } from '../lbug/lbug-adapter.js';
|
||||
|
||||
const isDev = process.env.NODE_ENV === 'development';
|
||||
/**
|
||||
* Bump this when the embedding text template changes in a way that should
|
||||
* invalidate existing vectors, such as metadata/header shape changes,
|
||||
* structural container context changes, or preceding-context formatting rules.
|
||||
*/
|
||||
export const EMBEDDING_TEXT_VERSION = 'v2';
|
||||
|
||||
/**
|
||||
* Compute a stable content fingerprint for an embeddable node.
|
||||
|
|
@ -57,12 +64,13 @@ export const contentHashForNode = (
|
|||
// Hash must be deterministic across runs, so exclude methodNames/fieldNames
|
||||
// which are populated during the batch loop via AST extraction.
|
||||
// Using only node.content ensures the hash stays stable.
|
||||
// NOTE: A change to extractStructuralNames behavior requires bumping EMBEDDING_TEXT_VERSION.
|
||||
const text = generateEmbeddingText(
|
||||
{ ...node, methodNames: undefined, fieldNames: undefined },
|
||||
node.content,
|
||||
config,
|
||||
);
|
||||
return createHash('sha1').update(text).digest('hex');
|
||||
return createHash('sha1').update(EMBEDDING_TEXT_VERSION).update('\n').update(text).digest('hex');
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -83,7 +91,7 @@ const queryEmbeddableNodes = async (
|
|||
try {
|
||||
let query: string;
|
||||
|
||||
if (label === 'Method') {
|
||||
if (label === LABEL_METHOD) {
|
||||
// Method has parameterCount and returnType
|
||||
query = `
|
||||
MATCH (n:Method)
|
||||
|
|
@ -115,7 +123,7 @@ const queryEmbeddableNodes = async (
|
|||
|
||||
const rows = await executeQuery(query);
|
||||
for (const row of rows) {
|
||||
const hasExportedColumn = label === 'Method' || LABELS_WITH_EXPORTED.has(label);
|
||||
const hasExportedColumn = label === LABEL_METHOD || LABELS_WITH_EXPORTED.has(label);
|
||||
allNodes.push({
|
||||
id: row.id ?? row[0],
|
||||
name: row.name ?? row[1],
|
||||
|
|
@ -126,7 +134,7 @@ const queryEmbeddableNodes = async (
|
|||
endLine: row.endLine ?? row[6],
|
||||
isExported: hasExportedColumn ? (row.isExported ?? row[7]) : undefined,
|
||||
description: row.description ?? (hasExportedColumn ? row[8] : row[7]),
|
||||
...(label === 'Method'
|
||||
...(label === LABEL_METHOD
|
||||
? {
|
||||
parameterCount: row.parameterCount ?? row[9],
|
||||
returnType: row.returnType ?? row[10],
|
||||
|
|
@ -415,8 +423,15 @@ export const runEmbeddingPipeline = async (
|
|||
}
|
||||
}
|
||||
|
||||
let prevTail = '';
|
||||
for (const chunk of chunks) {
|
||||
const text = generateEmbeddingText(node, chunk.text, finalConfig);
|
||||
const text = generateEmbeddingText(
|
||||
node,
|
||||
chunk.text,
|
||||
finalConfig,
|
||||
chunk.chunkIndex,
|
||||
prevTail,
|
||||
);
|
||||
allTexts.push(text);
|
||||
allUpdates.push({
|
||||
nodeId: node.id,
|
||||
|
|
@ -425,6 +440,7 @@ export const runEmbeddingPipeline = async (
|
|||
endLine: chunk.endLine,
|
||||
contentHash: hash,
|
||||
});
|
||||
prevTail = overlap > 0 ? chunk.text.slice(-overlap) : '';
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,12 @@
|
|||
*/
|
||||
|
||||
import type { EmbeddableNode, EmbeddingConfig } from './types.js';
|
||||
import { DEFAULT_EMBEDDING_CONFIG, isShortLabel } from './types.js';
|
||||
import {
|
||||
CHUNKING_RULES,
|
||||
DEFAULT_EMBEDDING_CONFIG,
|
||||
STRUCTURAL_TEXT_MODE_DECLARATION,
|
||||
isShortLabel,
|
||||
} from './types.js';
|
||||
|
||||
/**
|
||||
* Truncate description to max length at sentence/word boundary
|
||||
|
|
@ -95,47 +100,62 @@ const generateCodeBodyText = (
|
|||
node: EmbeddableNode,
|
||||
codeBody: string,
|
||||
config: Partial<EmbeddingConfig>,
|
||||
prevTail?: string,
|
||||
): string => {
|
||||
const header = buildMetadataHeader(node, config);
|
||||
const cleaned = cleanContent(codeBody);
|
||||
return `${header}\n\n${cleaned}`;
|
||||
const parts = [header];
|
||||
if (prevTail) {
|
||||
parts.push(`[preceding context]: ...${cleanContent(prevTail)}`);
|
||||
}
|
||||
parts.push('', cleanContent(codeBody));
|
||||
return parts.join('\n');
|
||||
};
|
||||
|
||||
/**
|
||||
* Generate embedding text for Class nodes
|
||||
* Signature + properties + method name list only (no method bodies)
|
||||
* Method/field names come from AST extractors via node.methodNames/node.fieldNames.
|
||||
*/
|
||||
const generateClassText = (
|
||||
node: EmbeddableNode,
|
||||
codeBody: string,
|
||||
config: Partial<EmbeddingConfig>,
|
||||
): string => {
|
||||
return generateStructuralTypeText(node, codeBody, config);
|
||||
const getCompactContainerContext = (
|
||||
cleanedContent: string,
|
||||
declarationOnly: string,
|
||||
): string | undefined => {
|
||||
const source = declarationOnly || cleanedContent;
|
||||
const nlIdx = source.indexOf('\n');
|
||||
const firstLine = (nlIdx === -1 ? source : source.substring(0, nlIdx)).trim();
|
||||
return firstLine ? `Container: ${firstLine}` : undefined;
|
||||
};
|
||||
|
||||
const generateStructuralTypeText = (
|
||||
node: EmbeddableNode,
|
||||
codeBody: string,
|
||||
config: Partial<EmbeddingConfig>,
|
||||
chunkIndex?: number,
|
||||
prevTail?: string,
|
||||
): string => {
|
||||
const header = buildMetadataHeader(node, config);
|
||||
const parts: string[] = [header];
|
||||
const isFirstChunk = chunkIndex === undefined || chunkIndex === 0;
|
||||
const cleanedContent = cleanContent(node.content);
|
||||
const declarationOnly = extractDeclarationOnly(cleanedContent);
|
||||
const compactContainerContext = getCompactContainerContext(cleanedContent, declarationOnly);
|
||||
|
||||
if (node.methodNames?.length) {
|
||||
if (compactContainerContext) {
|
||||
parts.push(compactContainerContext);
|
||||
}
|
||||
|
||||
if (prevTail) {
|
||||
parts.push(`[preceding context]: ...${cleanContent(prevTail)}`);
|
||||
}
|
||||
|
||||
if (isFirstChunk && node.methodNames?.length) {
|
||||
parts.push(`Methods: ${node.methodNames.join(', ')}`);
|
||||
}
|
||||
if (node.fieldNames?.length) {
|
||||
if (isFirstChunk && node.fieldNames?.length) {
|
||||
parts.push(`Properties: ${node.fieldNames.join(', ')}`);
|
||||
}
|
||||
|
||||
const declarationOnly = extractDeclarationOnly(cleanContent(node.content));
|
||||
if (declarationOnly) {
|
||||
if (isFirstChunk && declarationOnly) {
|
||||
parts.push('', declarationOnly);
|
||||
}
|
||||
|
||||
const cleanedChunk = cleanContent(codeBody);
|
||||
if (cleanedChunk && cleanedChunk !== cleanContent(node.content)) {
|
||||
if (cleanedChunk && cleanedChunk !== cleanedContent) {
|
||||
parts.push('', cleanedChunk);
|
||||
}
|
||||
|
||||
|
|
@ -229,6 +249,8 @@ export const generateEmbeddingText = (
|
|||
node: EmbeddableNode,
|
||||
codeBody: string,
|
||||
config: Partial<EmbeddingConfig> = {},
|
||||
chunkIndex?: number,
|
||||
prevTail?: string,
|
||||
): string => {
|
||||
if (isShortLabel(node.label)) {
|
||||
const header = buildMetadataHeader(node, config);
|
||||
|
|
@ -236,15 +258,12 @@ export const generateEmbeddingText = (
|
|||
return `${header}\n\n${cleaned}`;
|
||||
}
|
||||
|
||||
if (node.label === 'Class') {
|
||||
return generateClassText(node, codeBody, config);
|
||||
const chunkingRule = CHUNKING_RULES[node.label];
|
||||
if (chunkingRule?.structuralTextMode === STRUCTURAL_TEXT_MODE_DECLARATION) {
|
||||
return generateStructuralTypeText(node, codeBody, config, chunkIndex, prevTail);
|
||||
}
|
||||
|
||||
if (node.label === 'Interface') {
|
||||
return generateStructuralTypeText(node, codeBody, config);
|
||||
}
|
||||
|
||||
return generateCodeBodyText(node, codeBody, config);
|
||||
return generateCodeBodyText(node, codeBody, config, prevTail);
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -4,35 +4,76 @@
|
|||
* Type definitions for the embedding generation and semantic search system.
|
||||
*/
|
||||
|
||||
export const LABEL_FUNCTION = 'Function' as const;
|
||||
export const LABEL_METHOD = 'Method' as const;
|
||||
export const LABEL_CONSTRUCTOR = 'Constructor' as const;
|
||||
export const LABEL_CLASS = 'Class' as const;
|
||||
export const LABEL_INTERFACE = 'Interface' as const;
|
||||
export const LABEL_STRUCT = 'Struct' as const;
|
||||
export const LABEL_ENUM = 'Enum' as const;
|
||||
export const LABEL_TRAIT = 'Trait' as const;
|
||||
export const LABEL_IMPL = 'Impl' as const;
|
||||
export const LABEL_MACRO = 'Macro' as const;
|
||||
export const LABEL_NAMESPACE = 'Namespace' as const;
|
||||
export const LABEL_TYPE_ALIAS = 'TypeAlias' as const;
|
||||
export const LABEL_TYPEDEF = 'Typedef' as const;
|
||||
export const LABEL_CONST = 'Const' as const;
|
||||
export const LABEL_PROPERTY = 'Property' as const;
|
||||
export const LABEL_RECORD = 'Record' as const;
|
||||
export const LABEL_UNION = 'Union' as const;
|
||||
export const LABEL_STATIC = 'Static' as const;
|
||||
export const LABEL_VARIABLE = 'Variable' as const;
|
||||
export const LABEL_CODE_ELEMENT = 'CodeElement' as const;
|
||||
|
||||
export const CHUNK_MODE_AST_FUNCTION = 'ast-function' as const;
|
||||
export const CHUNK_MODE_AST_DECLARATION = 'ast-declaration' as const;
|
||||
// CHUNK_MODE_CHARACTER exists for type completeness but is a no-op in CHUNKING_RULES —
|
||||
// omit the entry entirely to get character fallback via chunker.ts dispatch.
|
||||
export const CHUNK_MODE_CHARACTER = 'character' as const;
|
||||
|
||||
export const STRUCTURAL_TEXT_MODE_NONE = 'none' as const;
|
||||
export const STRUCTURAL_TEXT_MODE_DECLARATION = 'declaration' as const;
|
||||
|
||||
export interface ChunkingRule {
|
||||
mode:
|
||||
| typeof CHUNK_MODE_AST_FUNCTION
|
||||
| typeof CHUNK_MODE_AST_DECLARATION
|
||||
| typeof CHUNK_MODE_CHARACTER;
|
||||
includePrefix: boolean;
|
||||
includeSuffix: boolean;
|
||||
groupFields: boolean;
|
||||
structuralTextMode: typeof STRUCTURAL_TEXT_MODE_NONE | typeof STRUCTURAL_TEXT_MODE_DECLARATION;
|
||||
}
|
||||
|
||||
/**
|
||||
* Node labels that need chunking (have code body, potentially long)
|
||||
*/
|
||||
export const CHUNKABLE_LABELS = [
|
||||
'Function',
|
||||
'Method',
|
||||
'Constructor',
|
||||
'Class',
|
||||
'Interface',
|
||||
'Struct',
|
||||
'Enum',
|
||||
'Trait',
|
||||
'Impl',
|
||||
'Macro',
|
||||
'Namespace',
|
||||
LABEL_FUNCTION,
|
||||
LABEL_METHOD,
|
||||
LABEL_CONSTRUCTOR,
|
||||
LABEL_CLASS,
|
||||
LABEL_INTERFACE,
|
||||
LABEL_STRUCT,
|
||||
LABEL_ENUM,
|
||||
LABEL_TRAIT,
|
||||
LABEL_IMPL,
|
||||
LABEL_MACRO,
|
||||
LABEL_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',
|
||||
LABEL_TYPE_ALIAS,
|
||||
LABEL_TYPEDEF,
|
||||
LABEL_CONST,
|
||||
LABEL_PROPERTY,
|
||||
LABEL_RECORD,
|
||||
LABEL_UNION,
|
||||
LABEL_STATIC,
|
||||
LABEL_VARIABLE,
|
||||
] as const;
|
||||
|
||||
/**
|
||||
|
|
@ -61,26 +102,78 @@ export const isShortLabel = (label: string): boolean =>
|
|||
(SHORT_LABELS as readonly string[]).includes(label);
|
||||
|
||||
/**
|
||||
* Node labels that have structural names (methods/fields) extractable via AST
|
||||
* Node labels that have structural names (methods/fields) extractable via AST.
|
||||
* Only labels that consume methodNames/fieldNames in their embedding text should
|
||||
* be listed here — extra entries trigger wasted AST parses with no effect on output.
|
||||
*/
|
||||
export const STRUCTURAL_LABELS: ReadonlySet<string> = new Set([
|
||||
'Class',
|
||||
'Struct',
|
||||
'Interface',
|
||||
'Enum',
|
||||
LABEL_CLASS,
|
||||
LABEL_STRUCT,
|
||||
LABEL_INTERFACE,
|
||||
]);
|
||||
|
||||
/**
|
||||
* Node labels that have isExported column in their schema
|
||||
*/
|
||||
export const LABELS_WITH_EXPORTED = new Set([
|
||||
'Function',
|
||||
'Class',
|
||||
'Interface',
|
||||
'Method',
|
||||
'CodeElement',
|
||||
LABEL_FUNCTION,
|
||||
LABEL_CLASS,
|
||||
LABEL_INTERFACE,
|
||||
LABEL_METHOD,
|
||||
LABEL_CODE_ELEMENT,
|
||||
]) as ReadonlySet<string>;
|
||||
|
||||
/**
|
||||
* Labels that need special chunking and/or structural text semantics.
|
||||
* Any chunkable label omitted here intentionally falls back to characterChunk
|
||||
* plus generateCodeBodyText (for example Enum/Trait/Impl/Macro/Namespace).
|
||||
*/
|
||||
type ChunkableLabel = (typeof CHUNKABLE_LABELS)[number];
|
||||
export const CHUNKING_RULES: Readonly<Partial<Record<ChunkableLabel, ChunkingRule>>> = {
|
||||
[LABEL_FUNCTION]: {
|
||||
mode: CHUNK_MODE_AST_FUNCTION,
|
||||
includePrefix: true,
|
||||
includeSuffix: true,
|
||||
groupFields: false,
|
||||
structuralTextMode: STRUCTURAL_TEXT_MODE_NONE,
|
||||
},
|
||||
[LABEL_METHOD]: {
|
||||
mode: CHUNK_MODE_AST_FUNCTION,
|
||||
includePrefix: true,
|
||||
includeSuffix: true,
|
||||
groupFields: false,
|
||||
structuralTextMode: STRUCTURAL_TEXT_MODE_NONE,
|
||||
},
|
||||
[LABEL_CONSTRUCTOR]: {
|
||||
mode: CHUNK_MODE_AST_FUNCTION,
|
||||
includePrefix: true,
|
||||
includeSuffix: true,
|
||||
groupFields: false,
|
||||
structuralTextMode: STRUCTURAL_TEXT_MODE_NONE,
|
||||
},
|
||||
[LABEL_CLASS]: {
|
||||
mode: CHUNK_MODE_AST_DECLARATION,
|
||||
includePrefix: true,
|
||||
includeSuffix: false,
|
||||
groupFields: true,
|
||||
structuralTextMode: STRUCTURAL_TEXT_MODE_DECLARATION,
|
||||
},
|
||||
[LABEL_INTERFACE]: {
|
||||
mode: CHUNK_MODE_AST_DECLARATION,
|
||||
includePrefix: true,
|
||||
includeSuffix: false,
|
||||
groupFields: false,
|
||||
structuralTextMode: STRUCTURAL_TEXT_MODE_DECLARATION,
|
||||
},
|
||||
[LABEL_STRUCT]: {
|
||||
mode: CHUNK_MODE_AST_DECLARATION,
|
||||
includePrefix: true,
|
||||
includeSuffix: false,
|
||||
groupFields: true,
|
||||
structuralTextMode: STRUCTURAL_TEXT_MODE_DECLARATION,
|
||||
},
|
||||
};
|
||||
|
||||
/**
|
||||
* Embedding pipeline phases
|
||||
*/
|
||||
|
|
|
|||
|
|
@ -209,11 +209,12 @@ describe('chunkNode', () => {
|
|||
const result = await chunkNode('Class', content, 'test.ts', 1, 6, 90, 0);
|
||||
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0].text).toContain('class Parser {');
|
||||
expect(result[0].text).toContain('options: ParserOptions;');
|
||||
expect(result[0].text).toContain('cache: Map<string, any>;');
|
||||
expect(result[1].text).toContain('parseJSON()');
|
||||
expect(result[1].text).toContain('validate()');
|
||||
expect(result[0].startLine).toBe(2);
|
||||
expect(result[0].startLine).toBe(1);
|
||||
expect(result[1].startLine).toBe(4);
|
||||
});
|
||||
|
||||
|
|
@ -237,11 +238,44 @@ describe('chunkNode', () => {
|
|||
const result = await chunkNode('Interface', content, 'test.ts', 10, 14, 500, 0);
|
||||
|
||||
expect(result).toHaveLength(1);
|
||||
expect(result[0].text).toContain('interface Handler {');
|
||||
expect(result[0].text).toContain('handle(event: Event): void;');
|
||||
expect(result[0].text).toContain('validate(input: string): boolean;');
|
||||
expect(result[0].text).toContain('readonly name: string;');
|
||||
});
|
||||
|
||||
it('uses declaration-aware chunking for Struct labels', async () => {
|
||||
const content = [
|
||||
'struct User {',
|
||||
' name: String,',
|
||||
' email: String,',
|
||||
' age: u32,',
|
||||
' address: String,',
|
||||
'}',
|
||||
].join('\n');
|
||||
const tree = makeDeclarationTree('struct_item', 'declaration_list', content, [
|
||||
'name: String,',
|
||||
'email: String,',
|
||||
'age: u32,',
|
||||
'address: String,',
|
||||
]);
|
||||
createParserForLanguage.mockResolvedValue({
|
||||
parse: vi.fn().mockReturnValue(tree),
|
||||
});
|
||||
|
||||
const result = await chunkNode('Struct', content, 'test.rs', 40, 45, 45, 0);
|
||||
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0].text).toContain('struct User {');
|
||||
expect(result[0].text).toContain('name: String,');
|
||||
expect(result[0].text).toContain('email: String');
|
||||
const combinedText = result.map((chunk) => chunk.text).join('\n');
|
||||
expect(combinedText).toContain('email: String');
|
||||
expect(combinedText).toContain('age: u32');
|
||||
expect(combinedText).toContain('address: String');
|
||||
expect(result[0].startLine).toBe(40);
|
||||
});
|
||||
|
||||
it('splits a function into multiple AST-aware chunks using snippet offsets', async () => {
|
||||
const content = [
|
||||
'function example() {',
|
||||
|
|
|
|||
|
|
@ -18,12 +18,19 @@ vi.mock('../../src/core/tree-sitter/parser-loader.js', () => ({
|
|||
resolveLanguageKey: vi.fn((language: string) => language),
|
||||
}));
|
||||
|
||||
vi.mock('gitnexus-shared', () => ({
|
||||
const { getLanguageFromFilename } = vi.hoisted(() => ({
|
||||
getLanguageFromFilename: vi.fn().mockReturnValue('typescript'),
|
||||
}));
|
||||
|
||||
vi.mock('gitnexus-shared', () => ({
|
||||
getLanguageFromFilename,
|
||||
}));
|
||||
|
||||
import { chunkNode } from '../../src/core/embeddings/chunker.js';
|
||||
|
||||
const CLASS_PREV_TAIL_SAMPLE = 30;
|
||||
const STRUCT_PREV_TAIL_SAMPLE = 20;
|
||||
|
||||
type FakeNode = {
|
||||
type: string;
|
||||
startIndex: number;
|
||||
|
|
@ -183,10 +190,18 @@ describe('embedding-chunking integration', () => {
|
|||
const chunks = await chunkNode(node.label, node.content, node.filePath, 20, 25, 90, 0);
|
||||
expect(chunks).toHaveLength(2);
|
||||
|
||||
const secondText = generateEmbeddingText(node, chunks[1].text);
|
||||
const secondText = generateEmbeddingText(
|
||||
node,
|
||||
chunks[1].text,
|
||||
{},
|
||||
chunks[1].chunkIndex,
|
||||
chunks[0].text.slice(-CLASS_PREV_TAIL_SAMPLE),
|
||||
);
|
||||
expect(secondText).toContain('Class: Parser');
|
||||
expect(secondText).toContain('Methods: parseJSON, validate');
|
||||
expect(secondText).toContain('Properties: options, cache');
|
||||
expect(secondText).toContain('Container: class Parser {');
|
||||
expect(secondText).toContain('[preceding context]: ...');
|
||||
expect(secondText).not.toContain('Methods: parseJSON, validate');
|
||||
expect(secondText).not.toContain('Properties: options, cache');
|
||||
expect(secondText).toContain('parseJSON(text: string)');
|
||||
});
|
||||
|
||||
|
|
@ -220,9 +235,53 @@ describe('embedding-chunking integration', () => {
|
|||
const text = generateEmbeddingText(node, chunks[0].text);
|
||||
expect(text).toContain('Interface: Handler');
|
||||
expect(text).toContain('Methods: handle, validate');
|
||||
expect(text).toContain('Container: interface Handler {');
|
||||
expect(text).toContain('readonly name: string;');
|
||||
});
|
||||
|
||||
it('struct chunks retain structural container context', async () => {
|
||||
getLanguageFromFilename.mockReturnValue('rust');
|
||||
const node = makeNode({
|
||||
label: 'Struct',
|
||||
name: 'User',
|
||||
fieldNames: ['name', 'email', 'age', 'address'],
|
||||
content: `struct User {
|
||||
name: String,
|
||||
email: String,
|
||||
age: u32,
|
||||
address: String,
|
||||
}`,
|
||||
startLine: 40,
|
||||
endLine: 45,
|
||||
filePath: 'src/user.rs',
|
||||
});
|
||||
createParserForLanguage.mockResolvedValue({
|
||||
parse: vi.fn().mockReturnValue(
|
||||
makeDeclarationTree('struct_item', 'declaration_list', node.content, [
|
||||
{ text: 'name: String,', type: 'field_definition' },
|
||||
{ text: 'email: String,', type: 'field_definition' },
|
||||
{ text: 'age: u32,', type: 'field_definition' },
|
||||
{ text: 'address: String,', type: 'field_definition' },
|
||||
]),
|
||||
),
|
||||
});
|
||||
|
||||
const chunks = await chunkNode(node.label, node.content, node.filePath, 40, 45, 45, 0);
|
||||
expect(chunks).toHaveLength(2);
|
||||
|
||||
const secondText = generateEmbeddingText(
|
||||
node,
|
||||
chunks[1].text,
|
||||
{},
|
||||
chunks[1].chunkIndex,
|
||||
chunks[0].text.slice(-STRUCT_PREV_TAIL_SAMPLE),
|
||||
);
|
||||
expect(secondText).toContain('Struct: User');
|
||||
expect(secondText).toContain('Container: struct User {');
|
||||
expect(secondText).not.toContain('Properties: name, email, age, address');
|
||||
expect(secondText).toContain('age: u32,');
|
||||
});
|
||||
|
||||
it('metadata is present in every chunk', () => {
|
||||
const longContent = 'x'.repeat(3000);
|
||||
const node = makeNode({
|
||||
|
|
|
|||
|
|
@ -1,11 +1,17 @@
|
|||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { createHash } from 'crypto';
|
||||
import { contentHashForNode } from '../../src/core/embeddings/embedding-pipeline.js';
|
||||
import {
|
||||
contentHashForNode,
|
||||
EMBEDDING_TEXT_VERSION,
|
||||
} from '../../src/core/embeddings/embedding-pipeline.js';
|
||||
import { generateEmbeddingText } from '../../src/core/embeddings/text-generator.js';
|
||||
import type { EmbeddableNode, EmbeddingProgress } from '../../src/core/embeddings/types.js';
|
||||
import { DEFAULT_EMBEDDING_CONFIG, EMBEDDABLE_LABELS } from '../../src/core/embeddings/types.js';
|
||||
import { STALE_HASH_SENTINEL } from '../../src/core/lbug/schema.js';
|
||||
|
||||
const CLASS_CHUNK_SIZE = 90;
|
||||
const CLASS_OVERLAP = 10;
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// contentHashForNode
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
|
|
@ -32,6 +38,8 @@ describe('contentHashForNode', () => {
|
|||
it('matches sha1(generateEmbeddingText(node, node.content))', () => {
|
||||
const node = makeNode();
|
||||
const expected = createHash('sha1')
|
||||
.update(EMBEDDING_TEXT_VERSION)
|
||||
.update('\n')
|
||||
.update(generateEmbeddingText(node, node.content))
|
||||
.digest('hex');
|
||||
expect(contentHashForNode(node)).toBe(expected);
|
||||
|
|
@ -56,6 +64,10 @@ describe('contentHashForNode', () => {
|
|||
const hashWithFullDefaults = contentHashForNode(node, DEFAULT_EMBEDDING_CONFIG);
|
||||
expect(hashWithEmptyConfig).toBe(hashWithFullDefaults);
|
||||
});
|
||||
|
||||
it('exports a text template version marker', () => {
|
||||
expect(EMBEDDING_TEXT_VERSION).toBe('v2');
|
||||
});
|
||||
});
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
|
|
@ -439,6 +451,118 @@ describe('runEmbeddingPipeline incremental filter', () => {
|
|||
expect(vectorIndexCalls.length).toBeGreaterThanOrEqual(1);
|
||||
});
|
||||
|
||||
it('does not inject preceding context when overlap is disabled', async () => {
|
||||
const embedBatchSpy = vi
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => new Float32Array(384))),
|
||||
);
|
||||
vi.doMock('../../src/core/embeddings/embedder.js', () => ({
|
||||
initEmbedder: vi.fn().mockResolvedValue(undefined),
|
||||
embedBatch: embedBatchSpy,
|
||||
embedText: vi.fn().mockResolvedValue(new Float32Array(384)),
|
||||
embeddingToArray: vi.fn().mockImplementation((emb: Float32Array) => Array.from(emb)),
|
||||
isEmbedderReady: vi.fn().mockReturnValue(true),
|
||||
}));
|
||||
vi.doMock('../../src/core/lbug/lbug-adapter.js', () => ({
|
||||
loadVectorExtension: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
const node = makeNode({
|
||||
label: 'Class',
|
||||
name: 'Parser',
|
||||
content: `class Parser {
|
||||
options: ParserOptions;
|
||||
cache: Map<string, any>;
|
||||
parseJSON() { return JSON.parse("{}"); }
|
||||
validate() { return true; }
|
||||
}`,
|
||||
startLine: 1,
|
||||
endLine: 6,
|
||||
});
|
||||
|
||||
const executeQuery = mockExecuteQuery([node]);
|
||||
const executeWithReusedStatement = mockExecuteWithReusedStatement();
|
||||
|
||||
const { runEmbeddingPipeline } =
|
||||
await import('../../src/core/embeddings/embedding-pipeline.js');
|
||||
|
||||
await runEmbeddingPipeline(
|
||||
executeQuery,
|
||||
executeWithReusedStatement,
|
||||
onProgress,
|
||||
{ chunkSize: 90, overlap: 0 },
|
||||
undefined,
|
||||
undefined,
|
||||
new Map(),
|
||||
);
|
||||
|
||||
const embeddedTexts = embedBatchSpy.mock.calls.flatMap((call) => call[0] as string[]);
|
||||
const laterChunks = embeddedTexts.slice(1);
|
||||
expect(laterChunks.length).toBeGreaterThan(0);
|
||||
for (const text of laterChunks) {
|
||||
expect(text).not.toContain('[preceding context]:');
|
||||
}
|
||||
});
|
||||
|
||||
it('truncates preceding context to the configured overlap size', async () => {
|
||||
const embedBatchSpy = vi
|
||||
.fn()
|
||||
.mockImplementation((texts: string[]) =>
|
||||
Promise.resolve(texts.map(() => new Float32Array(384))),
|
||||
);
|
||||
vi.doMock('../../src/core/embeddings/embedder.js', () => ({
|
||||
initEmbedder: vi.fn().mockResolvedValue(undefined),
|
||||
embedBatch: embedBatchSpy,
|
||||
embedText: vi.fn().mockResolvedValue(new Float32Array(384)),
|
||||
embeddingToArray: vi.fn().mockImplementation((emb: Float32Array) => Array.from(emb)),
|
||||
isEmbedderReady: vi.fn().mockReturnValue(true),
|
||||
}));
|
||||
vi.doMock('../../src/core/lbug/lbug-adapter.js', () => ({
|
||||
loadVectorExtension: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
const node = makeNode({
|
||||
label: 'Class',
|
||||
name: 'Parser',
|
||||
content: `class Parser {
|
||||
options: ParserOptions;
|
||||
cache: Map<string, any>;
|
||||
parseJSON() { return JSON.parse("{}"); }
|
||||
validate() { return true; }
|
||||
}`,
|
||||
startLine: 1,
|
||||
endLine: 6,
|
||||
});
|
||||
|
||||
const executeQuery = mockExecuteQuery([node]);
|
||||
const executeWithReusedStatement = mockExecuteWithReusedStatement();
|
||||
|
||||
const { runEmbeddingPipeline } =
|
||||
await import('../../src/core/embeddings/embedding-pipeline.js');
|
||||
|
||||
await runEmbeddingPipeline(
|
||||
executeQuery,
|
||||
executeWithReusedStatement,
|
||||
onProgress,
|
||||
{ chunkSize: CLASS_CHUNK_SIZE, overlap: CLASS_OVERLAP },
|
||||
undefined,
|
||||
undefined,
|
||||
new Map(),
|
||||
);
|
||||
|
||||
const embeddedTexts = embedBatchSpy.mock.calls.flatMap((call) => call[0] as string[]);
|
||||
const laterChunk = embeddedTexts.find((text) => text.includes('[preceding context]:'));
|
||||
expect(laterChunk).toBeDefined();
|
||||
expect(laterChunk).toContain('[preceding context]: ...');
|
||||
const precedingContextLine = laterChunk
|
||||
?.split('\n')
|
||||
.find((line) => line.startsWith('[preceding context]: ...'));
|
||||
expect(precedingContextLine).toBeDefined();
|
||||
expect(precedingContextLine).toContain('ring, any>');
|
||||
expect(precedingContextLine).not.toContain('parseJSON() {');
|
||||
});
|
||||
|
||||
it('throws when DELETE for stale nodes fails with non-trivial error', async () => {
|
||||
mockEmbedderSetup();
|
||||
|
||||
|
|
|
|||
|
|
@ -148,6 +148,65 @@ describe('text-generator', () => {
|
|||
expect(text).toContain('class Parser {');
|
||||
expect(text).toContain('parseJSON(text: string) { return JSON.parse(text); }');
|
||||
});
|
||||
|
||||
it('generates Struct text with structural metadata', () => {
|
||||
const node: EmbeddableNode = {
|
||||
...baseNode,
|
||||
label: 'Struct',
|
||||
name: 'User',
|
||||
fieldNames: ['name', 'age'],
|
||||
content: `struct User {
|
||||
name: String,
|
||||
age: u32,
|
||||
}`,
|
||||
};
|
||||
const text = generateEmbeddingText(node, node.content);
|
||||
expect(text).toContain('Struct: User');
|
||||
expect(text).toContain('Properties: name, age');
|
||||
expect(text).toContain('Container: struct User {');
|
||||
expect(text).toContain('struct User {');
|
||||
});
|
||||
|
||||
it('keeps compact container context on later structural chunks', () => {
|
||||
const node: EmbeddableNode = {
|
||||
...baseNode,
|
||||
label: 'Class',
|
||||
name: 'Parser',
|
||||
methodNames: ['parseJSON', 'validate'],
|
||||
fieldNames: ['options', 'cache'],
|
||||
content: `class Parser {
|
||||
options: ParserOptions;
|
||||
cache: Map<string, any>;
|
||||
parseJSON(text: string) { return JSON.parse(text); }
|
||||
validate() { return true; }
|
||||
}`,
|
||||
};
|
||||
const text = generateEmbeddingText(
|
||||
node,
|
||||
'validate() { return true; }',
|
||||
{},
|
||||
1,
|
||||
'parseJSON(text: string) { return JSON.parse(text); }',
|
||||
);
|
||||
expect(text).toContain('Class: Parser');
|
||||
expect(text).toContain('Container: class Parser {');
|
||||
expect(text).toContain('[preceding context]: ...parseJSON(text: string)');
|
||||
expect(text).not.toContain('Methods: parseJSON, validate');
|
||||
expect(text).not.toContain('Properties: options, cache');
|
||||
});
|
||||
|
||||
it('adds preceding context to non-structural chunk text', () => {
|
||||
const text = generateEmbeddingText(
|
||||
baseNode,
|
||||
'return JSON.parse(text);',
|
||||
{},
|
||||
1,
|
||||
'function parseJSON(text: string): Result<any> {',
|
||||
);
|
||||
expect(text).toContain('Function: parseJSON');
|
||||
expect(text).toContain('[preceding context]: ...function parseJSON');
|
||||
expect(text).toContain('return JSON.parse(text);');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Constructor label', () => {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue