@takara-ai/miru-code 1.7.8 → 1.8.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +22 -11
- package/docs/kiro-transform-research-synopsis.md +79 -0
- package/docs/miru-aws-transform-concept.md +74 -0
- package/docs/miru-aws-transform-kiro-power-spec.md +98 -0
- package/docs/self-hosted-sagemaker-identity-center.md +109 -0
- package/docs/self-hosted-sagemaker.md +107 -24
- package/package.json +7 -1
- package/src/benchmark/grep.ts +9 -0
- package/src/benchmark/locate-compare.ts +33 -12
- package/src/benchmark/rg-literal.ts +130 -24
- package/src/embeddings/sagemaker.ts +166 -3
- package/src/literal.ts +1 -1
- package/src/mcp/index-cache.ts +52 -7
- package/src/mcp/server.ts +77 -16
- package/src/miru-index.ts +4 -1
- package/src/ranking/boosting.ts +6 -9
- package/src/ranking/config.ts +35 -0
- package/src/ranking/location.ts +5 -4
- package/src/ranking/penalties.ts +4 -3
- package/src/ranking/weighting.ts +4 -1
- package/src/search.ts +13 -8
- package/src/setup.ts +1 -0
- package/src/snippet.ts +343 -6
- package/src/utils.ts +8 -2
- package/src/version.ts +4 -2
package/src/mcp/index-cache.ts
CHANGED
|
@@ -13,6 +13,7 @@ import {
|
|
|
13
13
|
} from "../utils.ts";
|
|
14
14
|
|
|
15
15
|
const CACHE_MAX_SIZE = 10;
|
|
16
|
+
const BUN_WATCH_RECONCILE_INTERVAL_MS = 500;
|
|
16
17
|
|
|
17
18
|
/** Directory names we skip for MCP fs.watch update triggers (aligned with file-walker). */
|
|
18
19
|
const WATCH_IGNORED_DIR_NAMES = new Set([
|
|
@@ -35,7 +36,7 @@ const WATCH_IGNORED_DIR_NAMES = new Set([
|
|
|
35
36
|
".eggs",
|
|
36
37
|
]);
|
|
37
38
|
|
|
38
|
-
type WatcherHandle =
|
|
39
|
+
type WatcherHandle = { close(): void };
|
|
39
40
|
|
|
40
41
|
type CacheEntry = {
|
|
41
42
|
source: string;
|
|
@@ -274,13 +275,19 @@ export class IndexCache {
|
|
|
274
275
|
|
|
275
276
|
try {
|
|
276
277
|
await index.applyFileChanges(paths);
|
|
277
|
-
if (!isGitUrl(source)) {
|
|
278
|
-
await index.saveToCache(resolve(source), { force: true });
|
|
279
|
-
}
|
|
280
278
|
} catch {
|
|
281
279
|
for (const p of paths) {
|
|
282
280
|
entry.pendingPaths.add(p);
|
|
283
281
|
}
|
|
282
|
+
return;
|
|
283
|
+
}
|
|
284
|
+
if (!isGitUrl(source)) {
|
|
285
|
+
try {
|
|
286
|
+
await index.saveToCache(resolve(source), { force: true });
|
|
287
|
+
} catch {
|
|
288
|
+
// The in-memory index is already current. A cache write failure must
|
|
289
|
+
// not requeue the change and repeatedly re-embed the same file.
|
|
290
|
+
}
|
|
284
291
|
}
|
|
285
292
|
};
|
|
286
293
|
|
|
@@ -351,10 +358,48 @@ export class IndexCache {
|
|
|
351
358
|
return;
|
|
352
359
|
}
|
|
353
360
|
|
|
354
|
-
|
|
355
|
-
|
|
361
|
+
let nativeWatcher: ReturnType<typeof watch> | null = null;
|
|
362
|
+
try {
|
|
363
|
+
nativeWatcher = watch(resolved, { recursive: true }, (_event, filename) => {
|
|
364
|
+
this.noteFileChange(path, filename);
|
|
365
|
+
});
|
|
366
|
+
} catch {
|
|
367
|
+
// Recursive fs.watch is unavailable on some platforms. The reconciliation
|
|
368
|
+
// fallback below keeps local indexes current there.
|
|
369
|
+
}
|
|
370
|
+
|
|
371
|
+
// Bun 1.3 on macOS can create a recursive watcher without delivering file
|
|
372
|
+
// events. Reconcile periodically in Bun so a silent watcher cannot leave an
|
|
373
|
+
// MCP index stale. Native watchers remain the fast path everywhere else.
|
|
374
|
+
let reconcileInFlight = false;
|
|
375
|
+
const reconcile = (): void => {
|
|
376
|
+
if (reconcileInFlight) {
|
|
377
|
+
return;
|
|
378
|
+
}
|
|
379
|
+
const cacheKey = computeSourceCacheKey(path);
|
|
380
|
+
const entry = this.entries.get(cacheKey);
|
|
381
|
+
if (!entry?.index) {
|
|
382
|
+
return;
|
|
383
|
+
}
|
|
384
|
+
reconcileInFlight = true;
|
|
385
|
+
void this.checkAndQueueStaleFiles(path, entry.index, cacheKey).finally(() => {
|
|
386
|
+
reconcileInFlight = false;
|
|
387
|
+
});
|
|
388
|
+
};
|
|
389
|
+
const interval =
|
|
390
|
+
process.versions.bun || !nativeWatcher
|
|
391
|
+
? setInterval(reconcile, BUN_WATCH_RECONCILE_INTERVAL_MS)
|
|
392
|
+
: null;
|
|
393
|
+
interval?.unref();
|
|
394
|
+
|
|
395
|
+
this.watchers.set(resolved, {
|
|
396
|
+
close: () => {
|
|
397
|
+
nativeWatcher?.close();
|
|
398
|
+
if (interval) {
|
|
399
|
+
clearInterval(interval);
|
|
400
|
+
}
|
|
401
|
+
},
|
|
356
402
|
});
|
|
357
|
-
this.watchers.set(resolved, watcher);
|
|
358
403
|
}
|
|
359
404
|
|
|
360
405
|
get watcher(): WatcherHandle | null {
|
package/src/mcp/server.ts
CHANGED
|
@@ -1,3 +1,4 @@
|
|
|
1
|
+
import { join } from "node:path";
|
|
1
2
|
import * as z from "zod";
|
|
2
3
|
import packageJson from "../../package.json";
|
|
3
4
|
import { benchmarkSearchComparison, toAgentBenchmarkSummary } from "../benchmark/compare.ts";
|
|
@@ -9,6 +10,7 @@ import {
|
|
|
9
10
|
recordFromBenchmark,
|
|
10
11
|
} from "../benchmark/history.ts";
|
|
11
12
|
import { benchmarkLocateComparison } from "../benchmark/locate-compare.ts";
|
|
13
|
+
import { selectComparableLiteralSearchTool } from "../benchmark/rg-literal.ts";
|
|
12
14
|
import { appendAgentBenchmark } from "../benchmark/summary.ts";
|
|
13
15
|
import {
|
|
14
16
|
MCP_BENCHMARK_SERVER_INSTRUCTIONS,
|
|
@@ -20,7 +22,7 @@ import {
|
|
|
20
22
|
MCP_SERVER_INSTRUCTIONS,
|
|
21
23
|
} from "../installer/search-policy.ts";
|
|
22
24
|
import { formatLiteralLocate } from "../literal.ts";
|
|
23
|
-
import type { ContentType } from "../types.ts";
|
|
25
|
+
import type { Chunk, ContentType, SearchResult } from "../types.ts";
|
|
24
26
|
import {
|
|
25
27
|
clampMcpTopK,
|
|
26
28
|
DEFAULT_EXPAND_AFTER,
|
|
@@ -50,6 +52,10 @@ const REPO_DESCRIPTION =
|
|
|
50
52
|
|
|
51
53
|
const BENCHMARK_SKIP_NOTES = {
|
|
52
54
|
local_repo_only: "Benchmark comparisons require a local repo path; git URL repos are skipped.",
|
|
55
|
+
limited_locate:
|
|
56
|
+
"Benchmark comparison skipped: locate.limit is global, while rg/grep limits are per file. Omit limit for a valid token and recall comparison.",
|
|
57
|
+
incompatible_literal_baseline:
|
|
58
|
+
"Benchmark comparison skipped: a comparable literal baseline requires rg or compatible grep. Windows findstr does not preserve locate scope, context, or count semantics.",
|
|
53
59
|
grep_timeout:
|
|
54
60
|
"Benchmark Grep baseline timed out; Miru results are still returned. " +
|
|
55
61
|
"Raise MIRU_BENCHMARK_SEARCH_TIMEOUT (seconds) if needed.",
|
|
@@ -64,6 +70,45 @@ function withBenchmarkSkippedNote(body: string, reason: BenchmarkSkipReason): st
|
|
|
64
70
|
})}`;
|
|
65
71
|
}
|
|
66
72
|
|
|
73
|
+
/**
|
|
74
|
+
* Use the source file for final snippet shaping when it is locally available.
|
|
75
|
+
* Index chunks can intentionally split a large declaration; source-backed
|
|
76
|
+
* snippets let the structural guard return a complete declaration instead.
|
|
77
|
+
*/
|
|
78
|
+
async function loadSnippetSources(
|
|
79
|
+
root: string | null,
|
|
80
|
+
results: SearchResult[],
|
|
81
|
+
): Promise<Map<string, Chunk> | undefined> {
|
|
82
|
+
if (!root) {
|
|
83
|
+
return undefined;
|
|
84
|
+
}
|
|
85
|
+
const sources = new Map<string, Chunk>();
|
|
86
|
+
await Promise.all(
|
|
87
|
+
[...new Set(results.map((result) => result.chunk.file_path))].map(async (filePath) => {
|
|
88
|
+
try {
|
|
89
|
+
const content = await Bun.file(join(root, filePath)).text();
|
|
90
|
+
const indexed = results.find((result) => result.chunk.file_path === filePath)?.chunk;
|
|
91
|
+
// Synthetic index entries (for example package entry-point metadata) do
|
|
92
|
+
// not occur verbatim in the source file and must retain their own chunk.
|
|
93
|
+
if (!indexed || !content.includes(indexed.content)) {
|
|
94
|
+
return;
|
|
95
|
+
}
|
|
96
|
+
const lines = content.split("\n");
|
|
97
|
+
sources.set(filePath, {
|
|
98
|
+
content,
|
|
99
|
+
file_path: filePath,
|
|
100
|
+
start_line: 1,
|
|
101
|
+
end_line: lines.length,
|
|
102
|
+
language: indexed.language,
|
|
103
|
+
});
|
|
104
|
+
} catch {
|
|
105
|
+
// A stale index entry is still safe to format from its indexed chunk.
|
|
106
|
+
}
|
|
107
|
+
}),
|
|
108
|
+
);
|
|
109
|
+
return sources;
|
|
110
|
+
}
|
|
111
|
+
|
|
67
112
|
async function persistBenchmarkQuery(
|
|
68
113
|
...args: Parameters<typeof appendBenchmarkQuery>
|
|
69
114
|
): Promise<void> {
|
|
@@ -139,7 +184,11 @@ export function createMcpServer(
|
|
|
139
184
|
}
|
|
140
185
|
await persistBenchmarkQuery(recordFromBenchmark(repoRoot, comparison.benchmark));
|
|
141
186
|
const body = formatResultsText(
|
|
142
|
-
formatResults(query, results, {
|
|
187
|
+
formatResults(query, results, {
|
|
188
|
+
repoRoot,
|
|
189
|
+
snippet: true,
|
|
190
|
+
snippetSourceChunks: await loadSnippetSources(index.root, results),
|
|
191
|
+
}),
|
|
143
192
|
);
|
|
144
193
|
return toolText(
|
|
145
194
|
appendAgentBenchmark(body, toAgentBenchmarkSummary(comparison.benchmark)),
|
|
@@ -157,7 +206,11 @@ export function createMcpServer(
|
|
|
157
206
|
if (results.length === 0) {
|
|
158
207
|
return toolText("No results found.");
|
|
159
208
|
}
|
|
160
|
-
const payload = formatResults(query, results, {
|
|
209
|
+
const payload = formatResults(query, results, {
|
|
210
|
+
repoRoot,
|
|
211
|
+
snippet: true,
|
|
212
|
+
snippetSourceChunks: await loadSnippetSources(index.root, results),
|
|
213
|
+
});
|
|
161
214
|
const body = formatResultsText(payload);
|
|
162
215
|
return toolText(skip ? withBenchmarkSkippedNote(body, skip) : body);
|
|
163
216
|
} catch (err) {
|
|
@@ -224,20 +277,27 @@ export function createMcpServer(
|
|
|
224
277
|
let skip: BenchmarkSkipReason | undefined;
|
|
225
278
|
const repoRoot = localRepoRoot(repo);
|
|
226
279
|
if (benchmark && repoRoot && typeof literal === "string") {
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
|
|
280
|
+
if (locateOpts.limit != null) {
|
|
281
|
+
skip = "limited_locate";
|
|
282
|
+
} else if (!selectComparableLiteralSearchTool()) {
|
|
283
|
+
skip = "incompatible_literal_baseline";
|
|
284
|
+
}
|
|
285
|
+
if (!skip) {
|
|
286
|
+
const comparison = await withGrepTimeoutFallback(() =>
|
|
287
|
+
benchmarkLocateComparison({
|
|
288
|
+
literal,
|
|
289
|
+
repoPath: repoRoot,
|
|
290
|
+
index,
|
|
291
|
+
locate: locateOpts,
|
|
292
|
+
}),
|
|
293
|
+
);
|
|
294
|
+
if (comparison) {
|
|
295
|
+
await persistBenchmarkQuery(recordFromAgentSummary(repoRoot, comparison.benchmark));
|
|
296
|
+
const body = formatLiteralLocateText(comparison.payload);
|
|
297
|
+
return toolText(appendAgentBenchmark(body, comparison.benchmark));
|
|
298
|
+
}
|
|
299
|
+
skip = "grep_timeout";
|
|
239
300
|
}
|
|
240
|
-
skip = "grep_timeout";
|
|
241
301
|
} else if (benchmark && !repoRoot) {
|
|
242
302
|
skip = "local_repo_only";
|
|
243
303
|
}
|
|
@@ -356,6 +416,7 @@ export function createMcpServer(
|
|
|
356
416
|
formatResults(`Chunks related to ${filePath}:${anchorLine}`, results, {
|
|
357
417
|
repoRoot,
|
|
358
418
|
snippet: true,
|
|
419
|
+
snippetSourceChunks: await loadSnippetSources(index.root, results),
|
|
359
420
|
}),
|
|
360
421
|
),
|
|
361
422
|
);
|
package/src/miru-index.ts
CHANGED
|
@@ -396,13 +396,15 @@ export class MiruIndex {
|
|
|
396
396
|
*/
|
|
397
397
|
async search(options: {
|
|
398
398
|
query: string;
|
|
399
|
+
/** Optional precomputed embedding for high-throughput callers. */
|
|
400
|
+
queryVector?: Float32Array;
|
|
399
401
|
topK?: number;
|
|
400
402
|
alpha?: number | null;
|
|
401
403
|
filterLanguages?: string[];
|
|
402
404
|
filterPaths?: string[];
|
|
403
405
|
rerank?: boolean;
|
|
404
406
|
}): Promise<SearchResult[]> {
|
|
405
|
-
const { query, topK = 10, alpha, filterLanguages, filterPaths, rerank } = options;
|
|
407
|
+
const { query, queryVector, topK = 10, alpha, filterLanguages, filterPaths, rerank } = options;
|
|
406
408
|
|
|
407
409
|
if (!this.chunks.length || !query.trim()) {
|
|
408
410
|
return [];
|
|
@@ -412,6 +414,7 @@ export class MiruIndex {
|
|
|
412
414
|
|
|
413
415
|
return hybridSearch({
|
|
414
416
|
query,
|
|
417
|
+
queryVector,
|
|
415
418
|
embeddings: this.embeddings,
|
|
416
419
|
semanticIndex: this.semanticIndex,
|
|
417
420
|
bm25Index: this.bm25Index,
|
package/src/ranking/boosting.ts
CHANGED
|
@@ -2,6 +2,7 @@ import { basename, dirname } from "node:path";
|
|
|
2
2
|
import { splitIdentifier } from "../tokens.ts";
|
|
3
3
|
import type { Chunk } from "../types.ts";
|
|
4
4
|
import { chunkKey } from "../types.ts";
|
|
5
|
+
import { rankingConfig } from "./config.ts";
|
|
5
6
|
import { searchImprovementsEnabled } from "./features.ts";
|
|
6
7
|
import {
|
|
7
8
|
boostExactStemMatches,
|
|
@@ -16,7 +17,6 @@ const EMBEDDED_SYMBOL_RE =
|
|
|
16
17
|
/\b(?:[A-Z][a-z][a-zA-Z0-9]*[A-Z][a-zA-Z0-9]*|[a-z][a-zA-Z0-9]*[A-Z][a-zA-Z0-9]+)\b/g;
|
|
17
18
|
|
|
18
19
|
const EMBEDDED_STEM_MIN_LEN = 4;
|
|
19
|
-
const EMBEDDED_SYMBOL_BOOST_SCALE = 0.5;
|
|
20
20
|
|
|
21
21
|
const DEFINITION_KEYWORDS = [
|
|
22
22
|
"class",
|
|
@@ -49,10 +49,6 @@ const SQL_DEFINITION_KEYWORDS = [
|
|
|
49
49
|
"CREATE FUNCTION",
|
|
50
50
|
];
|
|
51
51
|
|
|
52
|
-
const DEFINITION_BOOST_MULTIPLIER = 3.0;
|
|
53
|
-
const STEM_BOOST_MULTIPLIER = 1.0;
|
|
54
|
-
const FILE_COHERENCE_BOOST_FRAC = 0.2;
|
|
55
|
-
|
|
56
52
|
const STOPWORDS = new Set(
|
|
57
53
|
"a an and are as at be by do does for from has have how if in is it not of on or the to was what when where which who why with".split(
|
|
58
54
|
" ",
|
|
@@ -147,7 +143,7 @@ export function boostMultiChunkFiles(
|
|
|
147
143
|
}
|
|
148
144
|
|
|
149
145
|
const maxFileSum = Math.max(...fileSum.values());
|
|
150
|
-
const boostUnit = maxScore *
|
|
146
|
+
const boostUnit = maxScore * rankingConfig().fileCoherenceBoost;
|
|
151
147
|
|
|
152
148
|
for (const [fp, key] of bestChunk) {
|
|
153
149
|
const current = scores.get(key) ?? 0;
|
|
@@ -192,7 +188,7 @@ function boostSymbolDefinitions(
|
|
|
192
188
|
if (symbolName !== query.trim()) {
|
|
193
189
|
names.add(query.trim());
|
|
194
190
|
}
|
|
195
|
-
const boostUnit = maxScore *
|
|
191
|
+
const boostUnit = maxScore * rankingConfig().definitionBoost;
|
|
196
192
|
|
|
197
193
|
for (const key of [...boosted.keys()]) {
|
|
198
194
|
const chunk = chunksByKey.get(key);
|
|
@@ -235,7 +231,8 @@ function boostEmbeddedSymbols(
|
|
|
235
231
|
return;
|
|
236
232
|
}
|
|
237
233
|
|
|
238
|
-
const
|
|
234
|
+
const tuning = rankingConfig();
|
|
235
|
+
const boostUnit = maxScore * tuning.definitionBoost * tuning.embeddedSymbolScale;
|
|
239
236
|
const symbolsLower = new Set([...names].map((s) => s.toLowerCase()));
|
|
240
237
|
|
|
241
238
|
for (const key of [...boosted.keys()]) {
|
|
@@ -309,7 +306,7 @@ function boostStemMatches(
|
|
|
309
306
|
return;
|
|
310
307
|
}
|
|
311
308
|
|
|
312
|
-
const boost = maxScore *
|
|
309
|
+
const boost = maxScore * rankingConfig().exactStemBoost;
|
|
313
310
|
const pathCache = new Map<string, Set<string>>();
|
|
314
311
|
|
|
315
312
|
for (const key of [...boosted.keys()]) {
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
/** Optional ranking overrides, primarily for reproducible offline calibration. */
|
|
2
|
+
export interface RankingConfig {
|
|
3
|
+
rrfK: number;
|
|
4
|
+
candidateMultiplier: number;
|
|
5
|
+
locationCandidateMultiplier: number;
|
|
6
|
+
fileCoherenceBoost: number;
|
|
7
|
+
definitionBoost: number;
|
|
8
|
+
embeddedSymbolScale: number;
|
|
9
|
+
exactStemBoost: number;
|
|
10
|
+
locationBoost: number;
|
|
11
|
+
pathPenalties: boolean;
|
|
12
|
+
fileSaturationDecay: number;
|
|
13
|
+
}
|
|
14
|
+
|
|
15
|
+
function numberEnv(name: string, fallback: number, min: number, max: number): number {
|
|
16
|
+
const value = Number.parseFloat(process.env[name] ?? "");
|
|
17
|
+
return Number.isFinite(value) && value >= min && value <= max ? value : fallback;
|
|
18
|
+
}
|
|
19
|
+
|
|
20
|
+
export function rankingConfig(): RankingConfig {
|
|
21
|
+
return {
|
|
22
|
+
// Validated on the cached 3k CodeSearchNet corpus. Overrides preserve a
|
|
23
|
+
// straightforward rollback/A-B path for repository-specific evaluation.
|
|
24
|
+
rrfK: numberEnv("MIRU_RANK_RRF_K", 10, 1, 500),
|
|
25
|
+
candidateMultiplier: numberEnv("MIRU_RANK_CANDIDATE_MULTIPLIER", 5, 1, 50),
|
|
26
|
+
locationCandidateMultiplier: numberEnv("MIRU_RANK_LOCATION_CANDIDATE_MULTIPLIER", 10, 1, 100),
|
|
27
|
+
fileCoherenceBoost: numberEnv("MIRU_RANK_FILE_COHERENCE_BOOST", 0, 0, 2),
|
|
28
|
+
definitionBoost: numberEnv("MIRU_RANK_DEFINITION_BOOST", 3, 0, 10),
|
|
29
|
+
embeddedSymbolScale: numberEnv("MIRU_RANK_EMBEDDED_SYMBOL_SCALE", 0.5, 0, 2),
|
|
30
|
+
exactStemBoost: numberEnv("MIRU_RANK_EXACT_STEM_BOOST", 0.75, 0, 3),
|
|
31
|
+
locationBoost: numberEnv("MIRU_RANK_LOCATION_BOOST", 2.5, 0, 10),
|
|
32
|
+
pathPenalties: process.env.MIRU_RANK_PATH_PENALTIES !== "0",
|
|
33
|
+
fileSaturationDecay: numberEnv("MIRU_RANK_FILE_SATURATION_DECAY", 0.5, 0.05, 1),
|
|
34
|
+
};
|
|
35
|
+
}
|
package/src/ranking/location.ts
CHANGED
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
import { basename } from "node:path";
|
|
2
2
|
import type { Chunk } from "../types.ts";
|
|
3
3
|
import { chunkKey } from "../types.ts";
|
|
4
|
+
import { rankingConfig } from "./config.ts";
|
|
4
5
|
|
|
5
6
|
const LOCATION_QUERY_RE =
|
|
6
7
|
/\b(where(?:'s| is)?|entry\s*point|bootstrap|starts?|live[s]?|located|defined|wiring)\b/i;
|
|
@@ -16,7 +17,6 @@ const ENTRY_POINT_CONTENT_RE =
|
|
|
16
17
|
|
|
17
18
|
const PACKAGE_ENTRY_RE = /^\[package entry\]/;
|
|
18
19
|
|
|
19
|
-
const LOCATION_BOOST_MULTIPLIER = 2.5;
|
|
20
20
|
const INSTALLER_LOCATION_PENALTY = 0.35;
|
|
21
21
|
|
|
22
22
|
/** Too common as filenames to get exact-stem boosts from incidental query tokens. */
|
|
@@ -126,8 +126,9 @@ export function boostLocationSignals(
|
|
|
126
126
|
}
|
|
127
127
|
|
|
128
128
|
const entryPointQuery = isEntryPointQuery(query);
|
|
129
|
-
const
|
|
130
|
-
const
|
|
129
|
+
const tuning = rankingConfig();
|
|
130
|
+
const packageBoost = maxScore * (entryPointQuery ? tuning.locationBoost : 0.75);
|
|
131
|
+
const codeBoost = maxScore * tuning.locationBoost;
|
|
131
132
|
|
|
132
133
|
for (const chunk of allChunks) {
|
|
133
134
|
if (!chunkLooksLikeEntryPoint(chunk)) {
|
|
@@ -205,7 +206,7 @@ export function boostExactStemMatches(
|
|
|
205
206
|
return;
|
|
206
207
|
}
|
|
207
208
|
|
|
208
|
-
const boost = maxScore *
|
|
209
|
+
const boost = maxScore * rankingConfig().exactStemBoost;
|
|
209
210
|
const preferImplementation = hasSpecificImplementationTarget(query, chunksByKey);
|
|
210
211
|
|
|
211
212
|
for (const key of [...boosted.keys()]) {
|
package/src/ranking/penalties.ts
CHANGED
|
@@ -1,5 +1,6 @@
|
|
|
1
1
|
import { basename } from "node:path";
|
|
2
2
|
import type { Chunk } from "../types.ts";
|
|
3
|
+
import { rankingConfig } from "./config.ts";
|
|
3
4
|
|
|
4
5
|
const TEST_FILE_RE =
|
|
5
6
|
/(?:^|\/)(?:test_[^/]*\.py|[^/]*_test\.py|[^/]*_test\.go|[^/]*Tests?\.java|[^/]*Test\.php|[^/]*_spec\.rb|[^/]*_test\.rb|[^/]*\.test\.[jt]sx?|[^/]*\.spec\.[jt]sx?|[^/]*Tests?\.kt|[^/]*Spec\.kt|[^/]*Tests?\.swift|[^/]*Spec\.swift|[^/]*Tests?\.cs|test_[^/]*\.cpp|[^/]*_test\.cpp|test_[^/]*\.c|[^/]*_test\.c|[^/]*Spec\.scala|[^/]*Suite\.scala|[^/]*Test\.scala|[^/]*_test\.dart|test_[^/]*\.dart|[^/]*_spec\.lua|[^/]*_test\.lua|test_[^/]*\.lua|test_helpers?[^/]*\.\w+)$/;
|
|
@@ -15,7 +16,6 @@ const MILD_PENALTY = 0.7;
|
|
|
15
16
|
|
|
16
17
|
const REEXPORT_FILENAMES = new Set(["__init__.py", "package-info.java"]);
|
|
17
18
|
const FILE_SATURATION_THRESHOLD = 1;
|
|
18
|
-
const FILE_SATURATION_DECAY = 0.5;
|
|
19
19
|
|
|
20
20
|
function filePathPenalty(filePath: string): number {
|
|
21
21
|
const normalised = filePath.replace(/\\/g, "/");
|
|
@@ -49,6 +49,7 @@ export function rerankTopk(
|
|
|
49
49
|
}
|
|
50
50
|
|
|
51
51
|
const penaltyCache = new Map<string, number>();
|
|
52
|
+
const tuning = rankingConfig();
|
|
52
53
|
const penalised = new Map<string, number>();
|
|
53
54
|
|
|
54
55
|
for (const [key, score] of scores) {
|
|
@@ -57,7 +58,7 @@ export function rerankTopk(
|
|
|
57
58
|
continue;
|
|
58
59
|
}
|
|
59
60
|
let pen = score;
|
|
60
|
-
if (penalisePaths) {
|
|
61
|
+
if (penalisePaths && tuning.pathPenalties) {
|
|
61
62
|
let mult = penaltyCache.get(chunk.file_path);
|
|
62
63
|
if (mult === undefined) {
|
|
63
64
|
mult = filePathPenalty(chunk.file_path);
|
|
@@ -88,7 +89,7 @@ export function rerankTopk(
|
|
|
88
89
|
let effScore = penScore;
|
|
89
90
|
if (already >= FILE_SATURATION_THRESHOLD) {
|
|
90
91
|
const excess = already - FILE_SATURATION_THRESHOLD + 1;
|
|
91
|
-
effScore *=
|
|
92
|
+
effScore *= tuning.fileSaturationDecay ** excess;
|
|
92
93
|
}
|
|
93
94
|
|
|
94
95
|
selected.push({ score: effScore, key });
|
package/src/ranking/weighting.ts
CHANGED
|
@@ -1,7 +1,10 @@
|
|
|
1
1
|
import { isSymbolQuery } from "./boosting.ts";
|
|
2
2
|
|
|
3
3
|
const ALPHA_SYMBOL = 0.3;
|
|
4
|
-
|
|
4
|
+
// Natural-language code queries benefit from a modest lexical preference:
|
|
5
|
+
// CodeSearchNet factorial sweep (3k full validation) improved exact-match@1,
|
|
6
|
+
// hit@3, and MRR at 0.25 versus the former 0.50.
|
|
7
|
+
const ALPHA_NL = 0.25;
|
|
5
8
|
|
|
6
9
|
export function resolveAlpha(query: string, alpha: number | null | undefined): number {
|
|
7
10
|
if (alpha != null) {
|
package/src/search.ts
CHANGED
|
@@ -4,6 +4,7 @@ import type { SemanticIndex } from "./index/semantic-index.ts";
|
|
|
4
4
|
import { selectorToMask } from "./index/sparse.ts";
|
|
5
5
|
import { selectTopKScoreIndices } from "./index/topk.ts";
|
|
6
6
|
import { applyQueryBoost, boostMultiChunkFiles } from "./ranking/boosting.ts";
|
|
7
|
+
import { rankingConfig } from "./ranking/config.ts";
|
|
7
8
|
import { searchImprovementsEnabled } from "./ranking/features.ts";
|
|
8
9
|
import { isLocationQuery } from "./ranking/location.ts";
|
|
9
10
|
import { rerankTopk } from "./ranking/penalties.ts";
|
|
@@ -12,16 +13,14 @@ import { tokenize } from "./tokens.ts";
|
|
|
12
13
|
import type { Chunk, SearchResult } from "./types.ts";
|
|
13
14
|
import { chunkKey } from "./types.ts";
|
|
14
15
|
|
|
15
|
-
|
|
16
|
-
|
|
17
|
-
function rrfScores(scores: Map<string, number>): Map<string, number> {
|
|
16
|
+
function rrfScores(scores: Map<string, number>, rrfK: number): Map<string, number> {
|
|
18
17
|
if (scores.size === 0) {
|
|
19
18
|
return scores;
|
|
20
19
|
}
|
|
21
20
|
const ranked = [...scores.entries()].sort((a, b) => b[1] - a[1]);
|
|
22
21
|
const out = new Map<string, number>();
|
|
23
22
|
ranked.forEach(([key], i) => {
|
|
24
|
-
out.set(key, 1.0 / (
|
|
23
|
+
out.set(key, 1.0 / (rrfK + i + 1));
|
|
25
24
|
});
|
|
26
25
|
return out;
|
|
27
26
|
}
|
|
@@ -73,6 +72,8 @@ async function searchBm25(
|
|
|
73
72
|
|
|
74
73
|
export async function hybridSearch(options: {
|
|
75
74
|
query: string;
|
|
75
|
+
/** Reuse a caller-cached query vector when evaluating or batching searches. */
|
|
76
|
+
queryVector?: Float32Array;
|
|
76
77
|
embeddings: EmbeddingBackend;
|
|
77
78
|
semanticIndex: SemanticIndex;
|
|
78
79
|
bm25Index: BM25Index;
|
|
@@ -84,6 +85,7 @@ export async function hybridSearch(options: {
|
|
|
84
85
|
}): Promise<SearchResult[]> {
|
|
85
86
|
const {
|
|
86
87
|
query,
|
|
88
|
+
queryVector,
|
|
87
89
|
embeddings,
|
|
88
90
|
semanticIndex,
|
|
89
91
|
bm25Index,
|
|
@@ -95,11 +97,14 @@ export async function hybridSearch(options: {
|
|
|
95
97
|
} = options;
|
|
96
98
|
|
|
97
99
|
const alphaWeight = resolveAlpha(query, alpha);
|
|
100
|
+
const tuning = rankingConfig();
|
|
98
101
|
const candidateCount =
|
|
99
|
-
searchImprovementsEnabled() && isLocationQuery(query)
|
|
102
|
+
searchImprovementsEnabled() && isLocationQuery(query)
|
|
103
|
+
? topK * tuning.locationCandidateMultiplier
|
|
104
|
+
: topK * tuning.candidateMultiplier;
|
|
100
105
|
const chunksByKey = new Map(chunks.map((c) => [chunkKey(c), c]));
|
|
101
106
|
|
|
102
|
-
const queryVecPromise = embeddings.embedQuery(query);
|
|
107
|
+
const queryVecPromise = queryVector ? Promise.resolve(queryVector) : embeddings.embedQuery(query);
|
|
103
108
|
const bm25Promise = searchBm25(query, bm25Index, chunks, candidateCount, selector);
|
|
104
109
|
const queryVec = await queryVecPromise;
|
|
105
110
|
const bm25Hits = await bm25Promise;
|
|
@@ -123,8 +128,8 @@ export async function hybridSearch(options: {
|
|
|
123
128
|
}
|
|
124
129
|
}
|
|
125
130
|
|
|
126
|
-
const normalizedSemantic = rrfScores(semanticScores);
|
|
127
|
-
const normalizedBm25 = rrfScores(bm25Scores);
|
|
131
|
+
const normalizedSemantic = rrfScores(semanticScores, tuning.rrfK);
|
|
132
|
+
const normalizedBm25 = rrfScores(bm25Scores, tuning.rrfK);
|
|
128
133
|
|
|
129
134
|
const allKeys = new Set([...normalizedSemantic.keys(), ...normalizedBm25.keys()]);
|
|
130
135
|
const sortedKeys = [...allKeys].sort((a, b) => {
|
package/src/setup.ts
CHANGED
|
@@ -192,6 +192,7 @@ export async function runSageMakerSetup(options: RunSetupOptions = {}): Promise<
|
|
|
192
192
|
const config: SageMakerEmbeddingConfig = {
|
|
193
193
|
endpointName: parsed.endpointName,
|
|
194
194
|
region: parsed.region,
|
|
195
|
+
profile,
|
|
195
196
|
normalize: true,
|
|
196
197
|
truncate: true,
|
|
197
198
|
truncationDirection: "Right",
|