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:
evolution 2026-04-20 15:25:31 +08:00 committed by GitHub
parent d976038dc8
commit 2b7cff5fd2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 520 additions and 93 deletions

View file

@ -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),
});
}

View file

@ -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) : '';
}
}

View file

@ -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);
};
/**

View file

@ -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
*/

View file

@ -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() {',

View file

@ -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({

View file

@ -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();

View file

@ -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', () => {