Skip to content

Commit 9cec82a

Browse files
committed
fix(search): query-aware bestExample, golden file threshold, ONNX thread limits
- Lower golden file threshold from >=3 to >=2 for React/generic codebase coverage - Persist patterns field in IntelligenceGoldenFile for downstream category filtering - get_team_patterns now returns filtered goldenFiles for di/state/testing categories - search_codebase getBestExample uses directory overlap against top-5 results (no keyword maps) - discovery-harness falls back to parsed.bestExample when preflight.bestExample absent - ONNX Runtime limited to half cores by default (CODEBASE_CONTEXT_ONNX_THREADS override) - setImmediate yields added in embedBatch inner loop and indexer outer batch loop
1 parent 999faf7 commit 9cec82a

File tree

7 files changed

+119
-13
lines changed

7 files changed

+119
-13
lines changed

src/core/indexer.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -694,7 +694,7 @@ export class CodebaseIndexer {
694694
: [];
695695
const uniqueCategories = new Set(detectedPatterns.map((p) => p.category));
696696
const patternScore = uniqueCategories.size;
697-
if (patternScore >= 3) {
697+
if (patternScore >= 2) {
698698
const patternFlags: Record<string, boolean> = {};
699699
for (const p of detectedPatterns) {
700700
patternFlags[`${p.category}:${p.name}`] = true;
@@ -797,6 +797,10 @@ export class CodebaseIndexer {
797797
});
798798
}
799799

800+
// Yield to event loop between outer embedding batches — keeps system responsive
801+
// and allows SIGINT/SIGTERM signal handlers to fire during long indexing runs.
802+
await new Promise<void>((resolve) => setImmediate(resolve));
803+
800804
// Update progress
801805
const embeddingProgress = 50 + Math.round((i / chunksToEmbed.length) * 25);
802806
this.updateProgress('embedding', embeddingProgress);

src/core/reranker.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
*/
1010

1111
import type { SearchResult } from '../types/index.js';
12+
import os from 'os';
1213

1314
const DEFAULT_RERANKER_MODEL = 'Xenova/ms-marco-MiniLM-L-6-v2';
1415

@@ -47,7 +48,12 @@ async function ensureModelLoaded(): Promise<void> {
4748

4849
cachedTokenizer = await AutoTokenizer.from_pretrained(DEFAULT_RERANKER_MODEL);
4950
cachedModel = await AutoModelForSequenceClassification.from_pretrained(DEFAULT_RERANKER_MODEL, {
50-
dtype: 'q8'
51+
dtype: 'q8',
52+
// Limit ONNX Runtime to half cores by default — prevents system freeze during indexing.
53+
session_options: {
54+
intraOpNumThreads: Math.max(1, Math.floor(os.cpus().length / 2)),
55+
interOpNumThreads: 1
56+
}
5157
});
5258

5359
console.error('[reranker] Cross-encoder loaded successfully');

src/embeddings/transformers.ts

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,20 @@
11
import { EmbeddingProvider, DEFAULT_MODEL } from './types.js';
22
import type { FeatureExtractionPipelineType } from '@huggingface/transformers';
3+
import os from 'os';
4+
5+
/**
6+
* Returns the number of ONNX intra-op threads to use.
7+
* Defaults to half of available CPU cores to prevent system freeze during indexing.
8+
* Override via CODEBASE_CONTEXT_ONNX_THREADS env var.
9+
*/
10+
function getOnnxThreadCount(): number {
11+
const envVal = process.env.CODEBASE_CONTEXT_ONNX_THREADS;
12+
if (envVal) {
13+
const parsed = parseInt(envVal, 10);
14+
if (!isNaN(parsed) && parsed > 0) return parsed;
15+
}
16+
return Math.max(1, Math.floor(os.cpus().length / 2));
17+
}
318

419
interface ModelConfig {
520
dimensions: number;
@@ -62,7 +77,14 @@ export class TransformersEmbeddingProvider implements EmbeddingProvider {
6277
opts: Record<string, unknown>
6378
) => Promise<FeatureExtractionPipelineType>;
6479
this.pipeline = await (pipeline as PipelineFn)('feature-extraction', this.modelName, {
65-
dtype: 'q8'
80+
dtype: 'q8',
81+
// Limit ONNX Runtime to half cores by default — prevents system freeze during indexing.
82+
// interOpNumThreads: 1 — no benefit for single-model pipelines.
83+
// Override via CODEBASE_CONTEXT_ONNX_THREADS env var.
84+
session_options: {
85+
intraOpNumThreads: getOnnxThreadCount(),
86+
interOpNumThreads: 1
87+
}
6688
});
6789

6890
this.ready = true;
@@ -114,6 +136,10 @@ export class TransformersEmbeddingProvider implements EmbeddingProvider {
114136
});
115137
embeddings.push(...(output.tolist() as number[][]));
116138

139+
// Yield to event loop — allows signal handlers (SIGINT/SIGTERM) to fire during
140+
// tight embedding loops, keeping the system responsive during indexing.
141+
await new Promise<void>((resolve) => setImmediate(resolve));
142+
117143
if (texts.length > 100 && (i + batchSize) % 100 === 0) {
118144
console.error(`Embedded ${Math.min(i + batchSize, texts.length)}/${texts.length} chunks`);
119145
}

src/eval/discovery-harness.ts

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,9 @@ async function runSearchCodebase(
6363
const bestExample =
6464
preflight && typeof preflight === 'object' && preflight !== null && 'bestExample' in preflight
6565
? ((preflight.bestExample as string | undefined) ?? null)
66-
: null;
66+
: typeof parsed.bestExample === 'string'
67+
? parsed.bestExample
68+
: null;
6769

6870
return {
6971
payload,
@@ -145,7 +147,12 @@ export function matchSignals(
145147
payload: string,
146148
expectedSignals: string[],
147149
forbiddenSignals: string[] | undefined
148-
): { matchedSignals: string[]; missingSignals: string[]; forbiddenHits: string[]; usefulnessScore: number } {
150+
): {
151+
matchedSignals: string[];
152+
missingSignals: string[];
153+
forbiddenHits: string[];
154+
usefulnessScore: number;
155+
} {
149156
const normalizedPayload = normalizeText(payload);
150157
const matchedSignals = expectedSignals.filter((signal) =>
151158
normalizedPayload.includes(normalizeText(signal))

src/tools/get-team-patterns.ts

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,24 @@ import {
66
isComplementaryPatternConflict,
77
shouldSkipLegacyTestingFrameworkCategory
88
} from '../patterns/semantics.js';
9-
import type { IntelligenceData, PatternsData } from '../types/index.js';
9+
import type { IntelligenceData, IntelligenceGoldenFile, PatternsData } from '../types/index.js';
10+
11+
/**
12+
* Filter golden files by category prefix keys in their patterns map.
13+
* Returns files whose patterns include at least one key matching a given prefix.
14+
* Returns empty array when no golden files match — never fabricates.
15+
*/
16+
function filterGoldenFilesByCategory(
17+
goldenFiles: IntelligenceGoldenFile[] | undefined,
18+
categoryPrefixes: string[]
19+
): IntelligenceGoldenFile[] {
20+
if (!goldenFiles?.length) return [];
21+
return goldenFiles.filter(
22+
(gf) =>
23+
gf.patterns &&
24+
categoryPrefixes.some((prefix) => Object.keys(gf.patterns!).some((k) => k.startsWith(prefix)))
25+
);
26+
}
1027

1128
export const definition: Tool = {
1229
name: 'get_team_patterns',
@@ -59,11 +76,13 @@ export async function handle(
5976
if (intel.patterns?.dependencyInjection)
6077
(result.patterns as Record<string, unknown>).dependencyInjection =
6178
intel.patterns.dependencyInjection;
79+
result.goldenFiles = filterGoldenFilesByCategory(intel.goldenFiles, ['dependencyInjection:']);
6280
} else if (category === 'state') {
6381
result.patterns = {};
6482
if (intel.patterns?.stateManagement)
6583
(result.patterns as Record<string, unknown>).stateManagement =
6684
intel.patterns.stateManagement;
85+
result.goldenFiles = filterGoldenFilesByCategory(intel.goldenFiles, ['stateManagement:']);
6786
} else if (category === 'testing') {
6887
result.patterns = {};
6988
for (const k of [
@@ -75,6 +94,11 @@ export async function handle(
7594
if (intel.patterns?.[k])
7695
(result.patterns as Record<string, unknown>)[k] = intel.patterns[k];
7796
}
97+
result.goldenFiles = filterGoldenFilesByCategory(intel.goldenFiles, [
98+
'unitTestFramework:',
99+
'e2eFramework:',
100+
'testingFramework:'
101+
]);
78102
} else if (category === 'libraries') {
79103
result.topUsed = intel.importGraph?.topUsed || [];
80104
if (intel.tsconfigPaths) {

src/tools/search-codebase.ts

Lines changed: 44 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -587,9 +587,45 @@ export async function handle(
587587
return entries.length > 0 ? entries.join(', ') : undefined;
588588
}
589589

590-
// Return the top golden file path from intelligence.json (compact mode bestExample)
591-
function getBestExample(): string | undefined {
592-
return intelligence?.goldenFiles?.[0]?.file;
590+
// Return the golden file path most relevant to the given search results, using directory
591+
// overlap as a query-aware proximity signal. Falls back to goldenFiles[0] when no overlap.
592+
// Track B: use cross-encoder reranker to score golden files against the query for stronger
593+
// relevance — deferred to avoid reranker latency on compact mode in v2.
594+
function getBestExample(topResults: Array<{ filePath: string }>): string | undefined {
595+
const goldenFiles = intelligence?.goldenFiles;
596+
if (!goldenFiles?.length) return undefined;
597+
if (!topResults?.length) return goldenFiles[0].file;
598+
599+
// Extract directory paths from top-5 results
600+
const resultDirs = topResults.slice(0, 5).map((r) => path.dirname(r.filePath));
601+
602+
let bestFile = goldenFiles[0];
603+
let bestScore = -1;
604+
605+
for (const gf of goldenFiles) {
606+
const gfParts = path.dirname(gf.file).split(path.sep).filter(Boolean);
607+
let maxShared = 0;
608+
609+
for (const resDir of resultDirs) {
610+
const resParts = resDir.split(path.sep).filter(Boolean);
611+
let shared = 0;
612+
const minLen = Math.min(gfParts.length, resParts.length);
613+
for (let i = 0; i < minLen; i++) {
614+
if (gfParts[i] === resParts[i]) shared++;
615+
else break;
616+
}
617+
if (shared > maxShared) maxShared = shared;
618+
}
619+
620+
// Use golden file score as tie-breaker
621+
const score = maxShared * 1000 + (gf.score ?? 0);
622+
if (score > bestScore) {
623+
bestScore = score;
624+
bestFile = gf;
625+
}
626+
}
627+
628+
return bestFile.file;
593629
}
594630

595631
// Build up to 2 dynamic next-hop suggestions (compact mode)
@@ -887,9 +923,10 @@ export async function handle(
887923
};
888924
}
889925

890-
// Add bestExample (top 1 golden file)
891-
if (goldenFiles.length > 0) {
892-
decisionCard.bestExample = `${goldenFiles[0].file}`;
926+
// Add bestExample — query-aware via directory overlap against top search results
927+
const preflightBestExample = getBestExample(results.slice(0, 5));
928+
if (preflightBestExample) {
929+
decisionCard.bestExample = preflightBestExample;
893930
}
894931

895932
// Add impact (coverage + top 3 files)
@@ -1011,7 +1048,7 @@ export async function handle(
10111048
const compactResults = results.slice(0, 6);
10121049
const strongMemories = filterStrongMemories(relatedMemories, queryStr);
10131050
const patternSummary = buildPatternSummary();
1014-
const bestExample = getBestExample();
1051+
const bestExample = getBestExample(compactResults);
10151052
const nextHops = buildNextHops(compactResults, searchQuality);
10161053

10171054
return {

src/types/index.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -639,6 +639,8 @@ export type PatternsData = Record<string, PatternEntry>;
639639
export interface IntelligenceGoldenFile {
640640
file: string;
641641
score: number;
642+
/** Pattern flags persisted from runtime GoldenFile — e.g. "dependencyInjection:inject() function": true */
643+
patterns?: Record<string, boolean>;
642644
}
643645

644646
// ============================================================================

0 commit comments

Comments
 (0)