markindex 0.0.0-stage → 0.1.1
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/CHANGELOG.md +65 -0
- package/LICENSE +21 -0
- package/README.md +288 -2
- package/README.zh-CN.md +164 -0
- package/THIRD_PARTY_NOTICES.md +36 -0
- package/dist/adapters/ai-sdk.d.ts +8 -0
- package/dist/adapters/ai-sdk.js +20 -0
- package/dist/adapters/mcp.d.ts +6 -0
- package/dist/adapters/mcp.js +28 -0
- package/dist/cli.d.ts +3 -0
- package/dist/cli.js +297 -0
- package/dist/config.d.ts +18 -0
- package/dist/config.js +18 -0
- package/dist/eval/questions.d.ts +34 -0
- package/dist/eval/questions.js +80 -0
- package/dist/eval/run.d.ts +62 -0
- package/dist/eval/run.js +163 -0
- package/dist/eval/score.d.ts +74 -0
- package/dist/eval/score.js +104 -0
- package/dist/index.d.ts +23 -0
- package/dist/index.js +23 -0
- package/dist/library/indexer.d.ts +89 -0
- package/dist/library/indexer.js +199 -0
- package/dist/library/scan.d.ts +29 -0
- package/dist/library/scan.js +106 -0
- package/dist/library.d.ts +30 -0
- package/dist/library.js +42 -0
- package/dist/llm/provider.d.ts +10 -0
- package/dist/llm/provider.js +32 -0
- package/dist/llm/rate-limit.d.ts +33 -0
- package/dist/llm/rate-limit.js +95 -0
- package/dist/llm/summarize.d.ts +67 -0
- package/dist/llm/summarize.js +193 -0
- package/dist/parse/markdown.d.ts +38 -0
- package/dist/parse/markdown.js +58 -0
- package/dist/retrieve.d.ts +65 -0
- package/dist/retrieve.js +174 -0
- package/dist/store/db.d.ts +50 -0
- package/dist/store/db.js +399 -0
- package/dist/store/search.d.ts +27 -0
- package/dist/store/search.js +116 -0
- package/dist/text/segment.d.ts +4 -0
- package/dist/text/segment.js +37 -0
- package/dist/text/tokens.d.ts +3 -0
- package/dist/text/tokens.js +13 -0
- package/dist/tools/contract.d.ts +51 -0
- package/dist/tools/contract.js +159 -0
- package/dist/tools/envelope.d.ts +31 -0
- package/dist/tools/envelope.js +62 -0
- package/dist/tools/impl.d.ts +7 -0
- package/dist/tools/impl.js +441 -0
- package/dist/tools/similar.d.ts +8 -0
- package/dist/tools/similar.js +40 -0
- package/dist/tree/build.d.ts +27 -0
- package/dist/tree/build.js +139 -0
- package/dist/tree/optimize.d.ts +39 -0
- package/dist/tree/optimize.js +185 -0
- package/dist/types.d.ts +27 -0
- package/dist/types.js +1 -0
- package/package.json +88 -4
|
@@ -0,0 +1,193 @@
|
|
|
1
|
+
import { createHash } from 'node:crypto';
|
|
2
|
+
import { generateText } from 'ai';
|
|
3
|
+
import { SUMMARY_TOKEN_THRESHOLD } from '../config.js';
|
|
4
|
+
import { approxTokens } from '../text/tokens.js';
|
|
5
|
+
import { rateLimiterFor } from './rate-limit.js';
|
|
6
|
+
const DEFAULT_CONCURRENCY = 8;
|
|
7
|
+
const DEFAULT_MAX_RETRIES = 3;
|
|
8
|
+
/** PageIndex caps summaries at about this many words; stated in the prompt. */
|
|
9
|
+
const SUMMARY_WORD_LIMIT = 150;
|
|
10
|
+
/** Outline given to the description prompt is cut to this many chars so it cannot overflow the context. */
|
|
11
|
+
export const OUTLINE_MAX_CHARS = 12_000;
|
|
12
|
+
const OUTLINE_SUMMARY_CHARS = 160;
|
|
13
|
+
const sha256 = (s) => createHash('sha256').update(s).digest('hex');
|
|
14
|
+
/** Port of PageIndex `generate_node_summary`, plus language and length constraints. */
|
|
15
|
+
export function nodeSummaryPrompt(text) {
|
|
16
|
+
return `You are given a part of a document, your task is to generate a description of the partial document about what are main points covered in the partial document.
|
|
17
|
+
|
|
18
|
+
Write the description in the same language as the partial document text. Keep it concise (at most ${SUMMARY_WORD_LIMIT} words).
|
|
19
|
+
|
|
20
|
+
Partial Document Text: ${text}
|
|
21
|
+
|
|
22
|
+
Directly return the description, do not include any other text.`;
|
|
23
|
+
}
|
|
24
|
+
/** Port of PageIndex `generate_doc_description`, plus a language constraint. */
|
|
25
|
+
export function docDescriptionPrompt(outline) {
|
|
26
|
+
return `You are an expert in generating descriptions for a document.
|
|
27
|
+
You are given a structure of a document. Your task is to generate a one-sentence description for the document, which makes it easy to distinguish the document from other documents.
|
|
28
|
+
|
|
29
|
+
Write the description in the same language as the document structure's titles.
|
|
30
|
+
|
|
31
|
+
Document Structure:
|
|
32
|
+
${outline}
|
|
33
|
+
|
|
34
|
+
Directly return the description, do not include any other text.`;
|
|
35
|
+
}
|
|
36
|
+
function flatten(nodes, out = []) {
|
|
37
|
+
for (const n of nodes) {
|
|
38
|
+
out.push(n);
|
|
39
|
+
flatten(n.children, out);
|
|
40
|
+
}
|
|
41
|
+
return out;
|
|
42
|
+
}
|
|
43
|
+
/** Nodes of a document that get a summary: blank nodes are skipped; short ones carry their own text. */
|
|
44
|
+
function targets(store, docPath, tree, threshold) {
|
|
45
|
+
const out = [];
|
|
46
|
+
for (const node of flatten(tree)) {
|
|
47
|
+
const text = store.getLines(docPath, node.startLine, node.ownEndLine).join('\n');
|
|
48
|
+
if (text.trim() === '')
|
|
49
|
+
continue;
|
|
50
|
+
const field = node.children.length === 0 ? 'summary' : 'prefixSummary';
|
|
51
|
+
out.push(node.ownTokens < threshold ? { node, text, field } : { node, text, field, hash: sha256(text) });
|
|
52
|
+
}
|
|
53
|
+
return out;
|
|
54
|
+
}
|
|
55
|
+
/** Titles + truncated summaries, indented by depth, cut to OUTLINE_MAX_CHARS. */
|
|
56
|
+
export function buildOutline(tree, maxChars = OUTLINE_MAX_CHARS) {
|
|
57
|
+
const lines = [];
|
|
58
|
+
let size = 0;
|
|
59
|
+
const walk = (nodes, depth) => {
|
|
60
|
+
for (const n of nodes) {
|
|
61
|
+
const s = (n.summary ?? n.prefixSummary ?? '').replace(/\s+/g, ' ').trim();
|
|
62
|
+
const line = `${' '.repeat(depth)}- ${n.title}${s ? `: ${Array.from(s).slice(0, OUTLINE_SUMMARY_CHARS).join('')}` : ''}`;
|
|
63
|
+
if (size + line.length + 1 > maxChars)
|
|
64
|
+
return false;
|
|
65
|
+
lines.push(line);
|
|
66
|
+
size += line.length + 1;
|
|
67
|
+
if (!walk(n.children, depth + 1))
|
|
68
|
+
return false;
|
|
69
|
+
}
|
|
70
|
+
return true;
|
|
71
|
+
};
|
|
72
|
+
const complete = walk(tree, 0);
|
|
73
|
+
if (!complete)
|
|
74
|
+
lines.push('- ... (outline truncated)');
|
|
75
|
+
return lines.join('\n');
|
|
76
|
+
}
|
|
77
|
+
async function runPool(items, concurrency, fn) {
|
|
78
|
+
let next = 0;
|
|
79
|
+
const worker = async () => {
|
|
80
|
+
while (next < items.length)
|
|
81
|
+
await fn(items[next++]);
|
|
82
|
+
};
|
|
83
|
+
await Promise.all(Array.from({ length: Math.max(1, Math.min(concurrency, items.length)) }, worker));
|
|
84
|
+
}
|
|
85
|
+
const errorMessage = (e) => (e instanceof Error ? e.message : String(e));
|
|
86
|
+
/**
|
|
87
|
+
* Fills node summaries (leaf -> `summary`, non-leaf -> `prefixSummary` of the node's own text only) and, unless
|
|
88
|
+
* `describe` is false and the doc has none yet, the document description. Nodes under the threshold use their own
|
|
89
|
+
* text; others are looked up in the `summaries` cache by sha256 of their own text, then sent to the LLM. A failed
|
|
90
|
+
* call never throws and never touches the store: it is returned in `failures`, everything that succeeded is saved.
|
|
91
|
+
*/
|
|
92
|
+
export async function summarizeDocument(store, docPath, opts) {
|
|
93
|
+
const threshold = opts.threshold ?? SUMMARY_TOKEN_THRESHOLD;
|
|
94
|
+
const concurrency = opts.concurrency ?? DEFAULT_CONCURRENCY;
|
|
95
|
+
const maxRetries = opts.maxRetries ?? DEFAULT_MAX_RETRIES;
|
|
96
|
+
const result = { llmCalls: 0, cached: 0, usage: { inputTokens: 0, outputTokens: 0 }, failures: [] };
|
|
97
|
+
const doc = store.getDoc(docPath);
|
|
98
|
+
const tree = store.getTree(docPath);
|
|
99
|
+
if (!doc || !tree)
|
|
100
|
+
throw new Error(`document not found: ${docPath}`);
|
|
101
|
+
const limiter = opts.rateLimiter ?? rateLimiterFor(opts.rateLimit);
|
|
102
|
+
const ask = async (prompt) => {
|
|
103
|
+
result.llmCalls++;
|
|
104
|
+
await limiter?.acquire(approxTokens(prompt));
|
|
105
|
+
const res = await generateText({ model: opts.model, prompt, maxRetries });
|
|
106
|
+
result.usage.inputTokens += res.usage?.inputTokens ?? 0;
|
|
107
|
+
result.usage.outputTokens += res.usage?.outputTokens ?? 0;
|
|
108
|
+
const text = res.text.trim();
|
|
109
|
+
if (text === '')
|
|
110
|
+
throw new Error('model returned an empty response');
|
|
111
|
+
return text;
|
|
112
|
+
};
|
|
113
|
+
const summaryByHash = new Map();
|
|
114
|
+
const pending = new Map();
|
|
115
|
+
const updates = [];
|
|
116
|
+
for (const t of targets(store, docPath, tree, threshold)) {
|
|
117
|
+
if (t.hash === undefined) {
|
|
118
|
+
updates.push({ nodeId: t.node.nodeId, [t.field]: t.text });
|
|
119
|
+
continue;
|
|
120
|
+
}
|
|
121
|
+
const hit = summaryByHash.get(t.hash) ?? store.getCachedSummary(t.hash);
|
|
122
|
+
if (hit !== undefined) {
|
|
123
|
+
summaryByHash.set(t.hash, hit);
|
|
124
|
+
result.cached++;
|
|
125
|
+
updates.push({ nodeId: t.node.nodeId, [t.field]: hit });
|
|
126
|
+
}
|
|
127
|
+
else {
|
|
128
|
+
const list = pending.get(t.hash);
|
|
129
|
+
if (list)
|
|
130
|
+
list.push(t);
|
|
131
|
+
else
|
|
132
|
+
pending.set(t.hash, [t]);
|
|
133
|
+
}
|
|
134
|
+
}
|
|
135
|
+
await runPool([...pending.entries()], concurrency, async ([hash, group]) => {
|
|
136
|
+
try {
|
|
137
|
+
const summary = await ask(nodeSummaryPrompt(group[0].text));
|
|
138
|
+
store.putCachedSummary(hash, summary);
|
|
139
|
+
for (const t of group)
|
|
140
|
+
updates.push({ nodeId: t.node.nodeId, [t.field]: summary });
|
|
141
|
+
}
|
|
142
|
+
catch (e) {
|
|
143
|
+
for (const t of group)
|
|
144
|
+
result.failures.push({ nodeId: t.node.nodeId, message: errorMessage(e) });
|
|
145
|
+
}
|
|
146
|
+
});
|
|
147
|
+
if (updates.length > 0) {
|
|
148
|
+
try {
|
|
149
|
+
store.updateNodeSummaries(docPath, updates);
|
|
150
|
+
}
|
|
151
|
+
catch (e) {
|
|
152
|
+
result.failures.push({ nodeId: null, message: `could not save summaries: ${errorMessage(e)}` });
|
|
153
|
+
}
|
|
154
|
+
}
|
|
155
|
+
// No description from an incomplete outline: it would never be regenerated once stored.
|
|
156
|
+
if (opts.describe !== false && doc.description === undefined && result.failures.length === 0) {
|
|
157
|
+
try {
|
|
158
|
+
const description = await ask(docDescriptionPrompt(buildOutline(store.getTree(docPath))));
|
|
159
|
+
store.setDocDescription(docPath, description);
|
|
160
|
+
}
|
|
161
|
+
catch (e) {
|
|
162
|
+
result.failures.push({ nodeId: null, message: errorMessage(e) });
|
|
163
|
+
}
|
|
164
|
+
}
|
|
165
|
+
return result;
|
|
166
|
+
}
|
|
167
|
+
/**
|
|
168
|
+
* Dry-run over every indexed document: how many nodes are over the threshold, how many LLM calls a full
|
|
169
|
+
* `summarizeDocument` pass would make (unique uncached node texts + one description per undescribed document unless
|
|
170
|
+
* `describe` is false) and the node-text input tokens (description outlines are not counted).
|
|
171
|
+
*/
|
|
172
|
+
export function estimateSummaryCost(store, opts = {}) {
|
|
173
|
+
const threshold = opts.threshold ?? SUMMARY_TOKEN_THRESHOLD;
|
|
174
|
+
let nodes = 0;
|
|
175
|
+
let inputTokens = 0;
|
|
176
|
+
let descriptions = 0;
|
|
177
|
+
const seen = new Set();
|
|
178
|
+
for (const path of store.allDocPaths()) {
|
|
179
|
+
const tree = store.getTree(path);
|
|
180
|
+
if (opts.describe !== false && store.getDoc(path).description === undefined)
|
|
181
|
+
descriptions++;
|
|
182
|
+
for (const t of targets(store, path, tree, threshold)) {
|
|
183
|
+
if (t.hash === undefined)
|
|
184
|
+
continue;
|
|
185
|
+
nodes++;
|
|
186
|
+
if (seen.has(t.hash) || store.getCachedSummary(t.hash) !== undefined)
|
|
187
|
+
continue;
|
|
188
|
+
seen.add(t.hash);
|
|
189
|
+
inputTokens += t.node.ownTokens;
|
|
190
|
+
}
|
|
191
|
+
}
|
|
192
|
+
return { nodes, llmCalls: seen.size + descriptions, inputTokens };
|
|
193
|
+
}
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
export interface Heading {
|
|
2
|
+
/** plain heading text, capped at MAX_TITLE_CHARS code points */
|
|
3
|
+
title: string;
|
|
4
|
+
level: number;
|
|
5
|
+
/** 1-based, original file line */
|
|
6
|
+
line: number;
|
|
7
|
+
/** last line of the heading (setext headings span two lines) */
|
|
8
|
+
endLine: number;
|
|
9
|
+
}
|
|
10
|
+
export interface Block {
|
|
11
|
+
type: string;
|
|
12
|
+
startLine: number;
|
|
13
|
+
endLine: number;
|
|
14
|
+
}
|
|
15
|
+
export interface ParsedDoc {
|
|
16
|
+
/** BOM stripped, split on \r\n, \r or \n (the same line endings mdast counts) */
|
|
17
|
+
lines: string[];
|
|
18
|
+
frontmatter?: {
|
|
19
|
+
startLine: number;
|
|
20
|
+
endLine: number;
|
|
21
|
+
data: Record<string, unknown>;
|
|
22
|
+
};
|
|
23
|
+
/** direct children of the mdast root only */
|
|
24
|
+
headings: Heading[];
|
|
25
|
+
/** top-level non-heading, non-frontmatter nodes */
|
|
26
|
+
blocks: Block[];
|
|
27
|
+
/** first line after frontmatter (1 when there is none) */
|
|
28
|
+
firstContentLine: number;
|
|
29
|
+
}
|
|
30
|
+
export interface ParseOptions {
|
|
31
|
+
/**
|
|
32
|
+
* Recognize setext headings (`Title` underlined with `===` / `---`), as CommonMark does. Default true.
|
|
33
|
+
* Set false for converter output (PDF/HTML → Markdown) that marks page breaks with a bare `---` line right after
|
|
34
|
+
* a paragraph: `---` then parses as a thematic break and `===` lines as paragraph text.
|
|
35
|
+
*/
|
|
36
|
+
setext?: boolean;
|
|
37
|
+
}
|
|
38
|
+
export declare function parseMarkdown(src: string, opts?: ParseOptions): ParsedDoc;
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
import { fromMarkdown } from 'mdast-util-from-markdown';
|
|
2
|
+
import { gfm } from 'micromark-extension-gfm';
|
|
3
|
+
import { gfmFromMarkdown } from 'mdast-util-gfm';
|
|
4
|
+
import { frontmatter } from 'micromark-extension-frontmatter';
|
|
5
|
+
import { frontmatterFromMarkdown } from 'mdast-util-frontmatter';
|
|
6
|
+
import { toString } from 'mdast-util-to-string';
|
|
7
|
+
import { parse as parseYaml } from 'yaml';
|
|
8
|
+
import { MAX_TITLE_CHARS } from '../config.js';
|
|
9
|
+
const FM = ['yaml', 'toml'];
|
|
10
|
+
const NO_SETEXT = { disable: { null: ['setextUnderline'] } };
|
|
11
|
+
export function parseMarkdown(src, opts = {}) {
|
|
12
|
+
const text = src.charCodeAt(0) === 0xfeff ? src.slice(1) : src;
|
|
13
|
+
const lines = text.split(/\r\n|\r|\n/);
|
|
14
|
+
const root = fromMarkdown(text, {
|
|
15
|
+
extensions: opts.setext === false ? [gfm(), frontmatter(FM), NO_SETEXT] : [gfm(), frontmatter(FM)],
|
|
16
|
+
mdastExtensions: [gfmFromMarkdown(), frontmatterFromMarkdown(FM)],
|
|
17
|
+
});
|
|
18
|
+
const doc = { lines, headings: [], blocks: [], firstContentLine: 1 };
|
|
19
|
+
for (const node of root.children) {
|
|
20
|
+
const pos = node.position;
|
|
21
|
+
if (!pos)
|
|
22
|
+
continue;
|
|
23
|
+
const startLine = pos.start.line;
|
|
24
|
+
const endLine = pos.end.line;
|
|
25
|
+
const type = node.type; // 'toml' is absent from mdast's static types
|
|
26
|
+
if (type === 'yaml' || type === 'toml') {
|
|
27
|
+
if (!doc.frontmatter) {
|
|
28
|
+
const value = node.value;
|
|
29
|
+
doc.frontmatter = { startLine, endLine, data: parseFrontmatter(type, value) };
|
|
30
|
+
doc.firstContentLine = endLine + 1;
|
|
31
|
+
}
|
|
32
|
+
}
|
|
33
|
+
else if (node.type === 'heading') {
|
|
34
|
+
doc.headings.push({ title: capTitle(toString(node).trim()), level: node.depth, line: startLine, endLine });
|
|
35
|
+
}
|
|
36
|
+
else {
|
|
37
|
+
doc.blocks.push({ type: node.type, startLine, endLine });
|
|
38
|
+
}
|
|
39
|
+
}
|
|
40
|
+
return doc;
|
|
41
|
+
}
|
|
42
|
+
function capTitle(title) {
|
|
43
|
+
const cps = Array.from(title);
|
|
44
|
+
return cps.length <= MAX_TITLE_CHARS ? title : cps.slice(0, MAX_TITLE_CHARS - 1).join('') + '…';
|
|
45
|
+
}
|
|
46
|
+
function parseFrontmatter(kind, value) {
|
|
47
|
+
if (kind === 'toml')
|
|
48
|
+
return { raw: value };
|
|
49
|
+
try {
|
|
50
|
+
const data = parseYaml(value);
|
|
51
|
+
return data !== null && typeof data === 'object' && !Array.isArray(data)
|
|
52
|
+
? data
|
|
53
|
+
: {};
|
|
54
|
+
}
|
|
55
|
+
catch {
|
|
56
|
+
return {};
|
|
57
|
+
}
|
|
58
|
+
}
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
import { type LanguageModel } from 'ai';
|
|
2
|
+
import { type RateLimiter, type RateLimitOptions } from './llm/rate-limit.js';
|
|
3
|
+
import type { Library } from './library.js';
|
|
4
|
+
export interface Evidence {
|
|
5
|
+
docPath: string;
|
|
6
|
+
nodeId: string;
|
|
7
|
+
titlePath: string[];
|
|
8
|
+
startLine: number;
|
|
9
|
+
endLine: number;
|
|
10
|
+
/** Re-read from the store by line range; never the model's own text. */
|
|
11
|
+
text: string;
|
|
12
|
+
}
|
|
13
|
+
export interface RetrieveOptions {
|
|
14
|
+
model: LanguageModel;
|
|
15
|
+
/** Max model steps (default 12). */
|
|
16
|
+
maxSteps?: number;
|
|
17
|
+
/** Stop once cumulative input tokens over completed steps reach this (default 200_000). */
|
|
18
|
+
maxInputTokens?: number;
|
|
19
|
+
/** Restrict the search to this folder (prompt-level instruction; the tools take a `folder` argument). */
|
|
20
|
+
folder?: string;
|
|
21
|
+
/**
|
|
22
|
+
* Client-side rate limit, applied before every model step (sliding 60 s window). Input tokens per step are
|
|
23
|
+
* estimated from the instructions and messages sent (tool definitions are not counted). The estimate is
|
|
24
|
+
* conservative: it counts the JSON-serialized messages, whose syntax overhead exceeds the text actually tokenized. The same options object
|
|
25
|
+
* shares one budget across calls, e.g. across the questions of an eval run.
|
|
26
|
+
*/
|
|
27
|
+
rateLimit?: RateLimitOptions;
|
|
28
|
+
/** Advanced: use this limiter instead of building one from `rateLimit` (e.g. to share a budget or inject a clock). */
|
|
29
|
+
rateLimiter?: RateLimiter;
|
|
30
|
+
}
|
|
31
|
+
export interface RetrieveResult {
|
|
32
|
+
evidence: Evidence[];
|
|
33
|
+
usage: {
|
|
34
|
+
inputTokens: number;
|
|
35
|
+
outputTokens: number;
|
|
36
|
+
};
|
|
37
|
+
/** Number of model steps (generateText round trips) taken. */
|
|
38
|
+
steps: number;
|
|
39
|
+
/** 'done': the model finished (selected evidence, or answered without selecting). */
|
|
40
|
+
stopReason: 'done' | 'budget' | 'max_steps';
|
|
41
|
+
/** Number of selected items dropped because the doc / node / range was invalid. */
|
|
42
|
+
dropped: number;
|
|
43
|
+
/** 'model': from the model's select_evidence call; 'fallback': nodes it read via get_node_content (budget / max_steps). */
|
|
44
|
+
selectedBy: 'model' | 'fallback';
|
|
45
|
+
/** True when the model called select_evidence with input that failed validation (no valid selection followed). */
|
|
46
|
+
invalidSelection: boolean;
|
|
47
|
+
}
|
|
48
|
+
/** Default `RetrieveOptions.maxSteps`. */
|
|
49
|
+
export declare const RETRIEVE_MAX_STEPS = 12;
|
|
50
|
+
/** Default `RetrieveOptions.maxInputTokens`. */
|
|
51
|
+
export declare const RETRIEVE_MAX_INPUT_TOKENS = 200000;
|
|
52
|
+
/**
|
|
53
|
+
* Built-in evidence retrieval loop. Returns evidence + usage + stopReason, never an answer.
|
|
54
|
+
*
|
|
55
|
+
* Termination: `select_evidence` is registered without `execute`; per the AI SDK, a step that calls a tool without
|
|
56
|
+
* `execute` ends the loop on its own (no extra `hasToolCall` condition needed). Evidence is taken from the last
|
|
57
|
+
* `select_evidence` call. If it never calls it and the stop is 'budget' or 'max_steps', evidence falls back to the
|
|
58
|
+
* nodes it successfully read via `get_node_content` (full ranges, most recent first, max 10; selectedBy 'fallback').
|
|
59
|
+
* Plain-text finish without selecting: evidence []. A select_evidence call failing schema validation is
|
|
60
|
+
* rejected by the SDK (error fed back, loop continues); if no valid selection follows, `invalidSelection` is true
|
|
61
|
+
* and the normal budget / max_steps fallback applies.
|
|
62
|
+
* Invalid refs (unknown doc, unknown node, empty clamped range) are dropped and counted in `dropped`.
|
|
63
|
+
* `folder` is a prompt-level scoping instruction only.
|
|
64
|
+
*/
|
|
65
|
+
export declare function retrieve(lib: Library, question: string, opts: RetrieveOptions): Promise<RetrieveResult>;
|
package/dist/retrieve.js
ADDED
|
@@ -0,0 +1,174 @@
|
|
|
1
|
+
import { generateText, stepCountIs, tool } from 'ai';
|
|
2
|
+
import { z } from 'zod';
|
|
3
|
+
import { createAiSdkTools } from './adapters/ai-sdk.js';
|
|
4
|
+
import { estimateInputTokens, rateLimiterFor } from './llm/rate-limit.js';
|
|
5
|
+
import { AGENT_INSTRUCTIONS } from './tools/contract.js';
|
|
6
|
+
import { walk } from './tree/build.js';
|
|
7
|
+
/** Default `RetrieveOptions.maxSteps`. */
|
|
8
|
+
export const RETRIEVE_MAX_STEPS = 12;
|
|
9
|
+
/** Default `RetrieveOptions.maxInputTokens`. */
|
|
10
|
+
export const RETRIEVE_MAX_INPUT_TOKENS = 200_000;
|
|
11
|
+
const FALLBACK_CAP = 10;
|
|
12
|
+
const SELECT_EVIDENCE = 'select_evidence';
|
|
13
|
+
const selectEvidenceInput = z.object({
|
|
14
|
+
items: z.array(z.object({
|
|
15
|
+
doc_path: z.string(),
|
|
16
|
+
node_id: z.string(),
|
|
17
|
+
start_line: z.number().int().optional(),
|
|
18
|
+
end_line: z.number().int().optional(),
|
|
19
|
+
})),
|
|
20
|
+
});
|
|
21
|
+
const EVIDENCE_INSTRUCTIONS = [
|
|
22
|
+
'You are gathering evidence, not answering. Use the discovery tools to find the relevant sections, read them to confirm,',
|
|
23
|
+
`then call ${SELECT_EVIDENCE} exactly once with every section (doc_path + node_id, optionally a narrower start_line/end_line) that`,
|
|
24
|
+
'supports the question. Do not write a final answer; the evidence text is read back from the library by line range.',
|
|
25
|
+
].join(' ');
|
|
26
|
+
/**
|
|
27
|
+
* Built-in evidence retrieval loop. Returns evidence + usage + stopReason, never an answer.
|
|
28
|
+
*
|
|
29
|
+
* Termination: `select_evidence` is registered without `execute`; per the AI SDK, a step that calls a tool without
|
|
30
|
+
* `execute` ends the loop on its own (no extra `hasToolCall` condition needed). Evidence is taken from the last
|
|
31
|
+
* `select_evidence` call. If it never calls it and the stop is 'budget' or 'max_steps', evidence falls back to the
|
|
32
|
+
* nodes it successfully read via `get_node_content` (full ranges, most recent first, max 10; selectedBy 'fallback').
|
|
33
|
+
* Plain-text finish without selecting: evidence []. A select_evidence call failing schema validation is
|
|
34
|
+
* rejected by the SDK (error fed back, loop continues); if no valid selection follows, `invalidSelection` is true
|
|
35
|
+
* and the normal budget / max_steps fallback applies.
|
|
36
|
+
* Invalid refs (unknown doc, unknown node, empty clamped range) are dropped and counted in `dropped`.
|
|
37
|
+
* `folder` is a prompt-level scoping instruction only.
|
|
38
|
+
*/
|
|
39
|
+
export async function retrieve(lib, question, opts) {
|
|
40
|
+
const maxSteps = opts.maxSteps ?? RETRIEVE_MAX_STEPS;
|
|
41
|
+
const maxInputTokens = opts.maxInputTokens ?? RETRIEVE_MAX_INPUT_TOKENS;
|
|
42
|
+
const tools = {
|
|
43
|
+
...createAiSdkTools(lib),
|
|
44
|
+
[SELECT_EVIDENCE]: tool({
|
|
45
|
+
description: 'Submit the final evidence sections. Ends the retrieval.',
|
|
46
|
+
inputSchema: selectEvidenceInput,
|
|
47
|
+
}),
|
|
48
|
+
};
|
|
49
|
+
const overBudget = ({ steps }) => steps.reduce((sum, s) => sum + (s.usage.inputTokens ?? 0), 0) >= maxInputTokens;
|
|
50
|
+
const scope = opts.folder
|
|
51
|
+
? `\n\nRestrict the search to the folder "${opts.folder}": pass folder="${opts.folder}" to browse/search tools.`
|
|
52
|
+
: '';
|
|
53
|
+
const limiter = opts.rateLimiter ?? rateLimiterFor(opts.rateLimit);
|
|
54
|
+
const res = await generateText({
|
|
55
|
+
model: opts.model,
|
|
56
|
+
system: `${AGENT_INSTRUCTIONS}\n\n${EVIDENCE_INSTRUCTIONS}${scope}`,
|
|
57
|
+
prompt: question,
|
|
58
|
+
tools,
|
|
59
|
+
stopWhen: [stepCountIs(maxSteps), overBudget],
|
|
60
|
+
// prepareStep runs before each model call (AI SDK v7); awaiting here delays the step without any model wrapper.
|
|
61
|
+
...(limiter
|
|
62
|
+
? {
|
|
63
|
+
prepareStep: async ({ instructions, messages }) => {
|
|
64
|
+
await limiter.acquire(estimateInputTokens([instructions, messages]));
|
|
65
|
+
return {};
|
|
66
|
+
},
|
|
67
|
+
}
|
|
68
|
+
: {}),
|
|
69
|
+
});
|
|
70
|
+
let selected;
|
|
71
|
+
let invalidSelection = false;
|
|
72
|
+
const read = []; // successfully read nodes, first-read order
|
|
73
|
+
const readKeys = new Set();
|
|
74
|
+
for (const step of res.steps) {
|
|
75
|
+
for (const call of step.toolCalls) {
|
|
76
|
+
if (call.toolName === SELECT_EVIDENCE) {
|
|
77
|
+
const parsed = selectEvidenceInput.safeParse(call.input);
|
|
78
|
+
if (parsed.success) {
|
|
79
|
+
selected = parsed.data.items;
|
|
80
|
+
invalidSelection = false;
|
|
81
|
+
}
|
|
82
|
+
else
|
|
83
|
+
invalidSelection = true;
|
|
84
|
+
}
|
|
85
|
+
else if (call.toolName === 'get_node_content') {
|
|
86
|
+
const result = step.toolResults.find((r) => r.toolCallId === call.toolCallId);
|
|
87
|
+
if (!readSucceeded(result?.output))
|
|
88
|
+
continue;
|
|
89
|
+
const input = call.input;
|
|
90
|
+
if (typeof input.doc_path !== 'string' || !Array.isArray(input.node_ids))
|
|
91
|
+
continue;
|
|
92
|
+
for (const id of input.node_ids) {
|
|
93
|
+
if (typeof id !== 'string')
|
|
94
|
+
continue;
|
|
95
|
+
const key = `${input.doc_path}\0${id}`;
|
|
96
|
+
if (readKeys.has(key))
|
|
97
|
+
continue;
|
|
98
|
+
readKeys.add(key);
|
|
99
|
+
read.push({ doc_path: input.doc_path, node_id: id });
|
|
100
|
+
}
|
|
101
|
+
}
|
|
102
|
+
}
|
|
103
|
+
}
|
|
104
|
+
let stopReason = 'done';
|
|
105
|
+
if (selected === undefined) {
|
|
106
|
+
const input = res.steps.reduce((sum, s) => sum + (s.usage.inputTokens ?? 0), 0);
|
|
107
|
+
if (input >= maxInputTokens)
|
|
108
|
+
stopReason = 'budget';
|
|
109
|
+
else if (res.steps.length >= maxSteps && res.steps.at(-1)?.finishReason === 'tool-calls')
|
|
110
|
+
stopReason = 'max_steps';
|
|
111
|
+
}
|
|
112
|
+
// Budget / max_steps without a selection: fall back to the nodes the model read (most recent first, capped).
|
|
113
|
+
const useFallback = selected === undefined && stopReason !== 'done';
|
|
114
|
+
const items = useFallback ? read.slice(-FALLBACK_CAP).reverse() : (selected ?? []);
|
|
115
|
+
const { evidence, dropped } = resolveEvidence(lib, items);
|
|
116
|
+
return {
|
|
117
|
+
evidence,
|
|
118
|
+
usage: {
|
|
119
|
+
inputTokens: res.totalUsage.inputTokens ?? 0,
|
|
120
|
+
outputTokens: res.totalUsage.outputTokens ?? 0,
|
|
121
|
+
},
|
|
122
|
+
steps: res.steps.length,
|
|
123
|
+
stopReason,
|
|
124
|
+
dropped,
|
|
125
|
+
selectedBy: useFallback ? 'fallback' : 'model',
|
|
126
|
+
invalidSelection,
|
|
127
|
+
};
|
|
128
|
+
}
|
|
129
|
+
function readSucceeded(output) {
|
|
130
|
+
try {
|
|
131
|
+
return JSON.parse(String(output)).success === true;
|
|
132
|
+
}
|
|
133
|
+
catch {
|
|
134
|
+
return false;
|
|
135
|
+
}
|
|
136
|
+
}
|
|
137
|
+
function resolveEvidence(lib, items) {
|
|
138
|
+
const evidence = [];
|
|
139
|
+
const seen = new Set();
|
|
140
|
+
let dropped = 0;
|
|
141
|
+
const nodeIndex = new Map();
|
|
142
|
+
for (const item of items) {
|
|
143
|
+
let index = nodeIndex.get(item.doc_path);
|
|
144
|
+
if (!index) {
|
|
145
|
+
const roots = lib.store.getDoc(item.doc_path) ? lib.store.getTree(item.doc_path) : undefined;
|
|
146
|
+
index = new Map();
|
|
147
|
+
for (const e of walk(roots ?? []))
|
|
148
|
+
index.set(e.node.nodeId, e);
|
|
149
|
+
nodeIndex.set(item.doc_path, index);
|
|
150
|
+
}
|
|
151
|
+
const hit = index.get(item.node_id);
|
|
152
|
+
if (!hit) {
|
|
153
|
+
dropped++;
|
|
154
|
+
continue;
|
|
155
|
+
}
|
|
156
|
+
const startLine = Math.max(hit.node.startLine, item.start_line ?? hit.node.startLine);
|
|
157
|
+
const endLine = Math.min(hit.node.endLine, item.end_line ?? hit.node.endLine);
|
|
158
|
+
const key = `${item.doc_path}\0${item.node_id}\0${startLine}\0${endLine}`;
|
|
159
|
+
if (endLine < startLine || seen.has(key)) {
|
|
160
|
+
dropped++;
|
|
161
|
+
continue;
|
|
162
|
+
}
|
|
163
|
+
seen.add(key);
|
|
164
|
+
evidence.push({
|
|
165
|
+
docPath: item.doc_path,
|
|
166
|
+
nodeId: item.node_id,
|
|
167
|
+
titlePath: hit.titlePath,
|
|
168
|
+
startLine,
|
|
169
|
+
endLine,
|
|
170
|
+
text: lib.store.getLines(item.doc_path, startLine, endLine).join('\n'),
|
|
171
|
+
});
|
|
172
|
+
}
|
|
173
|
+
return { evidence, dropped };
|
|
174
|
+
}
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
import Database from 'better-sqlite3';
|
|
2
|
+
import type { DocumentRecord, TreeNode } from '../types.js';
|
|
3
|
+
type DB = Database.Database;
|
|
4
|
+
export interface Store {
|
|
5
|
+
readonly db: DB;
|
|
6
|
+
getDoc(path: string): DocumentRecord | undefined;
|
|
7
|
+
/** Direct subfolders (full paths) and a page of the docs directly inside `folder` ('' = library root). */
|
|
8
|
+
listDocs(folder: string, offset: number, limit: number): {
|
|
9
|
+
folders: string[];
|
|
10
|
+
docs: DocumentRecord[];
|
|
11
|
+
total: number;
|
|
12
|
+
};
|
|
13
|
+
allDocPaths(): string[];
|
|
14
|
+
getTree(path: string): TreeNode[] | undefined;
|
|
15
|
+
/** 1-based inclusive range, clamped to the document. */
|
|
16
|
+
getLines(path: string, start: number, end: number): string[];
|
|
17
|
+
/** Replaces every row of `doc.path` (docs, nodes, lines, FTS) in a single transaction. */
|
|
18
|
+
replaceDoc(doc: DocumentRecord, tree: TreeNode[], lines: string[]): void;
|
|
19
|
+
deleteDoc(path: string): void;
|
|
20
|
+
/**
|
|
21
|
+
* Writes summaries into the stored node JSON and the `summary` FTS column (segmented; fts_tri is untouched),
|
|
22
|
+
* in one transaction. Throws (and changes nothing) if a node id does not exist in `docPath`.
|
|
23
|
+
*/
|
|
24
|
+
updateNodeSummaries(docPath: string, updates: {
|
|
25
|
+
nodeId: string;
|
|
26
|
+
summary?: string;
|
|
27
|
+
prefixSummary?: string;
|
|
28
|
+
}[]): void;
|
|
29
|
+
setDocDescription(docPath: string, description: string): void;
|
|
30
|
+
/** Summary cache keyed by sha256 of a node's own text. */
|
|
31
|
+
getCachedSummary(contentHash: string): string | undefined;
|
|
32
|
+
putCachedSummary(contentHash: string, summary: string): void;
|
|
33
|
+
/** Index-wide key/value settings (e.g. `index_format`). Undefined when unset or the store predates the table. */
|
|
34
|
+
getMeta(key: string): string | undefined;
|
|
35
|
+
setMeta(key: string, value: string): void;
|
|
36
|
+
close(): void;
|
|
37
|
+
}
|
|
38
|
+
export declare class StoreLockedError extends Error {
|
|
39
|
+
readonly lockFile: string;
|
|
40
|
+
readonly pid: number | undefined;
|
|
41
|
+
constructor(lockFile: string, pid: number | undefined);
|
|
42
|
+
}
|
|
43
|
+
/**
|
|
44
|
+
* Opens (creating if needed) a store file. Write mode takes an exclusive `${file}.lock` holding the pid, created
|
|
45
|
+
* atomically (temp file + link). A lock whose pid is dead, or an empty / unparsable lock older than 10 s, is taken over. Read-only opens take no lock and never block on a writer (WAL).
|
|
46
|
+
*/
|
|
47
|
+
export declare function openStore(file: string, opts?: {
|
|
48
|
+
readonly?: boolean;
|
|
49
|
+
}): Store;
|
|
50
|
+
export {};
|