Skip to content

Commit 4032501

Browse files
neo-gptneo-gpt
authored
fix(ai): guard local embedding inputs (#13928) (#13929)
1 parent e0a5922 commit 4032501

2 files changed

Lines changed: 298 additions & 51 deletions

File tree

‎ai/services/knowledge-base/VectorService.mjs‎

Lines changed: 146 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,11 @@
1-
import aiConfig from '../../mcp/server/knowledge-base/config.mjs';
2-
import TextEmbeddingService from '../memory-core/TextEmbeddingService.mjs';
3-
import mcConfig from '../../mcp/server/memory-core/config.mjs';
4-
import Base from '../../../src/core/Base.mjs';
1+
import aiConfig from '../../mcp/server/knowledge-base/config.mjs';
2+
import TextEmbeddingService from '../memory-core/TextEmbeddingService.mjs';
3+
import mcConfig from '../../mcp/server/memory-core/config.mjs';
4+
import Base from '../../../src/core/Base.mjs';
5+
import {
6+
bytesToTokens,
7+
emitConsumerFriction
8+
} from '../memory-core/helpers/consumerFrictionHelper.mjs';
59
import ChromaManager from './ChromaManager.mjs';
610
import crypto from 'crypto';
711
import fs from 'fs-extra';
@@ -147,10 +151,10 @@ class VectorService extends Base {
147151
*/
148152
resolveTenantStamp(tenantContext = {}) {
149153
const config = this.getTenantIsolationConfig();
150-
const stamp = {
151-
tenantId : tenantContext.tenantId ?? config.defaultTenantId ?? 'neo-shared',
152-
repoSlug : tenantContext.repoSlug ?? config.defaultRepoSlug ?? 'neo',
153-
visibility : tenantContext.visibility ?? config.defaultVisibility ?? 'team',
154+
const stamp = {
155+
tenantId : tenantContext.tenantId ?? config.defaultTenantId,
156+
repoSlug : tenantContext.repoSlug ?? config.defaultRepoSlug,
157+
visibility : tenantContext.visibility ?? config.defaultVisibility,
154158
tenantConfigVersion: tenantContext.configVersion ?? 0,
155159
ingestedAt : Date.now()
156160
};
@@ -211,9 +215,9 @@ class VectorService extends Base {
211215
field,
212216
chunkId,
213217
clientValue,
214-
serverValue: serverValue ?? null,
215-
tenantId: stamp.tenantId,
216-
repoSlug: stamp.repoSlug,
218+
serverValue : serverValue ?? null,
219+
tenantId : stamp.tenantId,
220+
repoSlug : stamp.repoSlug,
217221
originAgentIdentity: stamp.originAgentIdentity ?? null
218222
};
219223

@@ -275,31 +279,134 @@ class VectorService extends Base {
275279
return deleteStale ? 'delete-upfront' : STALE_STRATEGY_SKIP;
276280
}
277281

282+
/**
283+
* Resolves the local embedding-model guardrail used before provider invocation.
284+
*
285+
* @returns {{enabled: Boolean, contextLimitTokens: Number, safeProcessingLimitTokens: Number, model: String}}
286+
*/
287+
resolveEmbeddingGuardrail() {
288+
const localEmbeddingProviders = new Set(['openAiCompatible', 'ollama']);
289+
const embeddingProvider = mcConfig.embeddingProvider;
290+
const contextLimitTokens = Number(aiConfig.localModels.embedding.contextLimitTokens);
291+
const safeProcessingLimitTokens = Number(aiConfig.localModels.embedding.safeProcessingLimitTokens);
292+
const model = embeddingProvider === 'ollama'
293+
? aiConfig.ollama.embeddingModel
294+
: embeddingProvider === 'openAiCompatible'
295+
? aiConfig.openAiCompatible.embeddingModel
296+
: embeddingProvider;
297+
298+
return {
299+
enabled : localEmbeddingProviders.has(embeddingProvider),
300+
contextLimitTokens,
301+
safeProcessingLimitTokens,
302+
model
303+
};
304+
}
305+
306+
/**
307+
* Emits a bounded, non-secret friction record for an over-budget embedding input.
308+
*
309+
* @param {Object} options
310+
* @param {Object} options.chunk Tenant-stamped chunk.
311+
* @param {String} options.text Provider input text.
312+
* @param {Object} options.guardrail Resolved embedding guardrail.
313+
* @returns {{skip: Boolean, inputBytes: Number, inputTokensEstimate: Number}}
314+
*/
315+
evaluateEmbeddingInput({chunk, text, guardrail}) {
316+
if (!guardrail.enabled) {
317+
return {skip: false, inputBytes: 0, inputTokensEstimate: 0};
318+
}
319+
320+
const inputBytes = Buffer.byteLength(text || '', 'utf8');
321+
const inputTokensEstimate = bytesToTokens(inputBytes);
322+
323+
if (inputTokensEstimate <= guardrail.safeProcessingLimitTokens) {
324+
return {skip: false, inputBytes, inputTokensEstimate};
325+
}
326+
327+
emitConsumerFriction({
328+
assetRef : chunk.id || chunk.source || chunk.name || 'kb-chunk',
329+
consumer : 'VectorService.embedChunks',
330+
model : guardrail.model,
331+
symptom : 'size-precheck-skip',
332+
emissionPoint : 'pre-invocation',
333+
suggestionKind : 'split-document',
334+
inputBytes,
335+
inputTokensEstimate,
336+
contextLimitTokens : guardrail.contextLimitTokens,
337+
safeProcessingLimitTokens: guardrail.safeProcessingLimitTokens,
338+
serviceDomain : 'other',
339+
note : 'KB embedding input exceeds safe processing band; split or reduce the source chunk before embedding.'
340+
});
341+
342+
logger.warn('[VectorService] Skipping over-budget embedding chunk before provider invocation.', {
343+
chunkId : chunk.id,
344+
tenantId : chunk.tenantId,
345+
repoSlug : chunk.repoSlug,
346+
source : chunk.source,
347+
inputBytes,
348+
inputTokensEstimate,
349+
safeProcessingLimitTokens: guardrail.safeProcessingLimitTokens,
350+
contextLimitTokens : guardrail.contextLimitTokens
351+
});
352+
353+
return {skip: true, inputBytes, inputTokensEstimate};
354+
}
355+
278356
/**
279357
* Embeds a set of chunks into the provided Chroma collection.
280358
*
281359
* @param {Object} options
282360
* @param {Object} options.collection Chroma collection target.
283361
* @param {Object[]} options.chunksToProcess Tenant-stamped chunks to embed.
284-
* @returns {Promise<void>}
362+
* @returns {Promise<{embedded: Number, skipped: Number}>}
285363
*/
286364
async embedChunks({collection, chunksToProcess}) {
287365
if (chunksToProcess.length === 0) {
288-
return;
366+
return {embedded: 0, skipped: 0};
289367
}
290368

291369
logger.log(`Using TextEmbeddingService with provider: ${mcConfig.embeddingProvider}.`);
292370
logger.log('Embedding chunks...');
293371

294372
const {batchSize, batchDelay, maxRetries} = aiConfig;
373+
const guardrail = this.resolveEmbeddingGuardrail();
374+
let embeddedCount = 0;
375+
let skippedCount = 0;
295376

296377
for (let i = 0; i < chunksToProcess.length; i += batchSize) {
297378
if (i > 0 && batchDelay) {
298379
await this.timeout(batchDelay);
299380
}
300381

301-
const batch = chunksToProcess.slice(i, i + batchSize);
302-
const textsToEmbed = batch.map(chunk => `${chunk.type}: ${chunk.name} in ${chunk.className || ''}\n${chunk.description || chunk.content || ''}`);
382+
const batch = chunksToProcess.slice(i, i + batchSize);
383+
const batchInputs = batch.map(chunk => ({
384+
chunk,
385+
text: `${chunk.type}: ${chunk.name} in ${chunk.className || ''}\n${chunk.description || chunk.content || ''}`
386+
}));
387+
const embeddable = [];
388+
389+
for (const input of batchInputs) {
390+
const result = this.evaluateEmbeddingInput({
391+
chunk: input.chunk,
392+
text : input.text,
393+
guardrail
394+
});
395+
396+
if (result.skip) {
397+
skippedCount++;
398+
} else {
399+
embeddable.push(input);
400+
}
401+
}
402+
403+
if (embeddable.length === 0) {
404+
logger.warn(`[VectorService] Skipped embedding batch ${i / batchSize + 1}; all ${batch.length} chunk(s) exceeded the embedding safe-processing band.`);
405+
continue;
406+
}
407+
408+
const batchToEmbed = embeddable.map(input => input.chunk);
409+
const textsToEmbed = embeddable.map(input => input.text);
303410

304411
let retries = 0;
305412
let success = false;
@@ -308,7 +415,7 @@ class VectorService extends Base {
308415
try {
309416
const embeddings = await TextEmbeddingService.embedTexts(textsToEmbed, mcConfig.embeddingProvider);
310417

311-
const metadatas = batch.map(chunk => {
418+
const metadatas = batchToEmbed.map(chunk => {
312419
const metadata = {};
313420
for (const [key, value] of Object.entries(chunk)) {
314421
metadata[key] = (value === null) ? 'null' : (typeof value === 'object') ? JSON.stringify(value) : value;
@@ -317,12 +424,13 @@ class VectorService extends Base {
317424
});
318425

319426
await collection.upsert({
320-
ids: batch.map(chunk => chunk.id),
427+
ids: batchToEmbed.map(chunk => chunk.id),
321428
embeddings,
322429
metadatas
323430
});
324431

325-
logger.log(`Processed and embedded batch ${i / batchSize + 1} of ${Math.ceil(chunksToProcess.length / batchSize)}`);
432+
embeddedCount += batchToEmbed.length;
433+
logger.log(`Processed and embedded batch ${i / batchSize + 1} of ${Math.ceil(chunksToProcess.length / batchSize)} (${batchToEmbed.length} embedded, ${batch.length - batchToEmbed.length} skipped).`);
326434
success = true;
327435
} catch (err) {
328436
retries++;
@@ -335,6 +443,8 @@ class VectorService extends Base {
335443
}
336444
}
337445
}
446+
447+
return {embedded: embeddedCount, skipped: skippedCount};
338448
}
339449

340450
/**
@@ -377,7 +487,11 @@ class VectorService extends Base {
377487
let shadowPromoted = false;
378488

379489
try {
380-
await this.embedChunks({collection: shadowCollection, chunksToProcess: knowledgeBase});
490+
const embedResult = await this.embedChunks({collection: shadowCollection, chunksToProcess: knowledgeBase});
491+
492+
if (embedResult.skipped > 0) {
493+
throw new Error(`KB_EMBEDDING_INPUT_SIZE_EXCEEDED: shadow-swap refused to promote an incomplete corpus after skipping ${embedResult.skipped} over-budget embedding chunk(s).`);
494+
}
381495

382496
logger.log(`Promoting shadow collection '${shadowName}' to '${aiConfig.collectionName}'.`);
383497
await liveCollection.modify({name: parkingName});
@@ -394,7 +508,7 @@ class VectorService extends Base {
394508

395509
return {
396510
message,
397-
embedded : knowledgeBase.length,
511+
embedded : embedResult.embedded,
398512
deleted : idsToDeleteCount,
399513
staleStrategy : 'shadow-swap',
400514
shadowCollection : shadowName,
@@ -498,16 +612,16 @@ class VectorService extends Base {
498612
knowledgeBase.forEach(chunk => {
499613
if (chunk.kind === 'module-context' && chunk.className) {
500614
classNameToDataMap[chunk.className] = {
501-
source : chunk.source,
502-
parent : chunk.extends || null
615+
source: chunk.source,
616+
parent: chunk.extends || null
503617
};
504618
}
505619
});
506620

507621
knowledgeBase.forEach(chunk => {
508-
let currentClass = chunk.className; // Metadata is now on every chunk
622+
let currentClass = chunk.className; // Metadata is now on every chunk
509623
const inheritanceChain = [];
510-
const visited = new Set();
624+
const visited = new Set();
511625

512626
// If no className metadata (e.g. non-class files), skip
513627
if (!currentClass) return;
@@ -529,8 +643,8 @@ class VectorService extends Base {
529643

530644
logger.log('Fetching existing documents from ChromaDB...');
531645
const existingIds = new Set();
532-
let offset = 0;
533-
const limit = 2000;
646+
let offset = 0;
647+
const limit = 2000;
534648
let batch;
535649

536650
// ChromaDB has a default limit (usually 10) if not specified.
@@ -539,7 +653,7 @@ class VectorService extends Base {
539653
batch = await collection.get({
540654
include: [],
541655
limit,
542-
offset: offset
656+
offset : offset
543657
});
544658

545659
batch.ids.forEach(id => existingIds.add(id));
@@ -593,7 +707,7 @@ class VectorService extends Base {
593707
if (viaMcp && workVolume > mcpThreshold) {
594708
// `logPath` is a Provider-owned leaf; read it directly so malformed config
595709
// shape fails loud instead of silently re-deriving a local default.
596-
const logDir = aiConfig.logPath;
710+
const logDir = aiConfig.logPath;
597711
const errorPayload = {
598712
error : `KB sync work volume exceeds MCP-callable threshold`,
599713
message: `${workVolume} chunks need re-embedding (threshold: ${mcpThreshold}). ` +
@@ -610,7 +724,7 @@ class VectorService extends Base {
610724

611725
if (shouldShadowSwap) {
612726
return await this.embedViaShadowSwap({
613-
liveCollection: collection,
727+
liveCollection : collection,
614728
knowledgeBase,
615729
idsToDeleteCount: idsToDelete.length
616730
});
@@ -621,12 +735,12 @@ class VectorService extends Base {
621735
logger.log(`Deleted ${idsToDelete.length} stale chunks.`);
622736
}
623737

624-
await this.embedChunks({collection, chunksToProcess});
738+
const embedResult = await this.embedChunks({collection, chunksToProcess});
625739

626-
const count = await collection.count();
740+
const count = await collection.count();
627741
const message = `Embedding complete. Collection now contains ${count} items.`;
628742
logger.log(message);
629-
return {message, embedded: chunksToProcess.length, deleted: idsToDelete.length};
743+
return {message, embedded: embedResult.embedded, deleted: idsToDelete.length};
630744
}
631745

632746
/**

0 commit comments

Comments
 (0)