diff --git a/README.md b/README.md index bd538cb..c6b2e42 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,8 @@ pnpm exec rgr doctor ``` Setup writes local helper files for the selected agents. The MCP surface is intentionally bounded -and read-focused. +and read-focused. Agents can request compact evidence first, then expand one returned citation +without opening a second index or reading arbitrary files. ### Audit a knowledge base diff --git a/docs/agent-integration.md b/docs/agent-integration.md index e559214..6f81d14 100644 --- a/docs/agent-integration.md +++ b/docs/agent-integration.md @@ -28,9 +28,17 @@ Use `--scope user` only when you intentionally want a user-wide installation. Pr ## MCP tools -The server exposes `ragmir_status`, `ragmir_route_prompt`, `ragmir_search`, `ragmir_ask`, `ragmir_research`, `ragmir_audit`, `ragmir_evaluate`, `ragmir_usage_report`, and `ragmir_security_audit`. - -Use compact search output when context is limited. `ragmir_ask` returns cited evidence, not a model generated answer. A cloud agent can receive returned passages, so choose that handoff only when it matches the corpus’s confidentiality requirements. +The server exposes `ragmir_status`, `ragmir_route_prompt`, `ragmir_search`, `ragmir_ask`, +`ragmir_research`, `ragmir_expand`, `ragmir_audit`, `ragmir_evaluate`, `ragmir_usage_report`, and +`ragmir_security_audit`. + +Use compact retrieval first, then pass a returned citation to `ragmir_expand` when the agent needs +the exact chunk or a bounded neighbor window. Search, ask, research, and expansion accept `maxBytes`; +the server also enforces the configured `mcpMaxOutputBytes` ceiling. Their single JSON text result +stays parseable, while `_meta["ragmir/output"]` reports retrieved bytes, returned bytes, compacting, +and truncation. `ragmir_ask` returns cited evidence, not a model generated answer. A cloud agent can +receive returned passages, so choose that handoff only when it matches the corpus’s confidentiality +requirements. ## Verify diff --git a/docs/api-reference.md b/docs/api-reference.md index 03a3a01..1724b4f 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -29,11 +29,16 @@ All paths resolve from `cwd` or the current working directory. Retrieval results | `search(query, options?)` | Return ranked cited passages. | | `ask(query, options?)` | Return cited retrieval context without calling an LLM. | | `research(query, options?)` | Run audit-backed multi-query retrieval. | +| `expandCitation(citation, options?)` | Read one exact indexed chunk and a bounded neighbor window. | | `compactSearchResults(results)` | Produce compact search output for limited contexts. | | `compactResearchReport(report)` | Remove verbose evidence text from a research report. | | `evaluateGoldenQueries(options)` | Score retrieval against a local golden-query file. | `SearchOptions` accepts `cwd`, `topK`, `contextRadius`, `includePaths`, and `excludePaths`. +`ExpandCitationOptions` accepts `cwd` and a `contextRadius` clamped to three chunks. +Search results expose a `contextPath` derived from Markdown headings or JSON structure. This field +improves retrieval, while `text`, offsets, and citations continue to reference the exact source +passage. ## Operations and privacy @@ -70,3 +75,8 @@ All paths resolve from `cwd` or the current working directory. Retrieval results `kbCommand` and `ragmirCommand` remain compatibility helpers. New integration code should use `rgrCommand` and the `rgr` CLI name. + +The MCP retrieval tools accept an optional `maxBytes` below the configured `mcpMaxOutputBytes` +ceiling. `ragmir_search` and `ragmir_research` also accept `compact`; `ragmir_ask` supports the same +flag. Oversized full output is compacted and then reduced as valid JSON. Output metrics are available +under `_meta["ragmir/output"]` and in the metadata-only usage report. diff --git a/docs/configuration.md b/docs/configuration.md index 539302b..ffb78ff 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -21,11 +21,14 @@ edit JSON only for a real need. | `retrievalProfile` | `balanced` | Use `fast`, `quality`, or `custom` for different search budgets. | | `embeddingProvider` | `local-hash` | Set `transformers` only after an explicit preload. | | `topK` | `8` | Change the default number of returned passages. | +| `mcpMaxOutputBytes` | `32768` | Cap each retrieval tool's serialized MCP text output. | | `chunkSize` / `chunkOverlap` | `1200` / `200` | Tune chunking, then rebuild the index. | | `maxFileBytes` | `50000000` | Raise only when the target corpus justifies it. | | `includeExtensions` | `[]` | Add safe custom text extensions. | Changing an embedding provider, model, or chunking field requires `rgr ingest --rebuild`. +Ragmir also preserves Markdown heading paths and JSON or JSONL structure as retrieval-only context. +Rebuild indexes created by an older Ragmir version to populate that structural context. ## Privacy profiles @@ -61,7 +64,8 @@ they run without a shell and must print text to stdout. Use `RAGMIR_*` variables for local experiments, for example: ```bash -RAGMIR_TOP_K=5 RAGMIR_CHUNK_SIZE=1000 rgr search "migration" +RAGMIR_TOP_K=5 rgr search "migration" +RAGMIR_MCP_MAX_OUTPUT_BYTES=16384 rgr mcp ``` Environment overrides cover selected runtime settings such as models, retrieval limits, access logs, diff --git a/packages/ragmir-core/README.md b/packages/ragmir-core/README.md index 70a7494..d95ce24 100644 --- a/packages/ragmir-core/README.md +++ b/packages/ragmir-core/README.md @@ -89,7 +89,7 @@ Frequently used exports: | --- | --- | | `setupProject`, `addSourceEntries` | Initialize project state and select files | | `ingest`, `audit` | Build the index and compare it with files on disk | -| `search`, `ask`, `research` | Retrieve cited passages at different levels of breadth | +| `search`, `ask`, `research`, `expandCitation` | Retrieve or expand cited passages | | `doctor`, `securityAudit` | Inspect readiness and local privacy posture | | `serveMcp` | Start the read-focused local MCP server | | `configurePdfOcr`, `inspectPdfOcr` | Configure and inspect local PDF OCR | @@ -105,8 +105,9 @@ npx rgr doctor ``` Ragmir writes helper files for the selected clients and points them at the current project. MCP -exposes status, search, ask, research, audit, evaluation, usage, and security tools. It does not -expose index deletion. +exposes status, search, ask, research, exact citation expansion, audit, evaluation, usage, and +security tools. Retrieval responses have a global byte ceiling and expose metadata-only output +metrics. The server does not expose index deletion. ## Retrieval modes @@ -147,7 +148,8 @@ local executable, never a shell string or cloud OCR service. - Generated indexes, models, reports, audio, and access logs belong under ignored `.ragmir/` state. - Redaction runs before indexing when configured, but it is not a compliance certification. - External extractors and model downloads are opt-in system boundaries. -- MCP retrieval is bounded and access logs contain metadata, not query text or retrieved passages. +- MCP retrieval is bounded by `mcpMaxOutputBytes`; access logs contain byte metrics, not query text + or retrieved passages. Read [configuration and privacy](https://github.com/jcode-works/jcode-ragmir/blob/main/docs/configuration.md) and [security hardening](https://github.com/jcode-works/jcode-ragmir/blob/main/SECURITY-HARDENING.md) diff --git a/packages/ragmir-core/scripts/mcp-smoke.mjs b/packages/ragmir-core/scripts/mcp-smoke.mjs index 6280e24..128fffb 100644 --- a/packages/ragmir-core/scripts/mcp-smoke.mjs +++ b/packages/ragmir-core/scripts/mcp-smoke.mjs @@ -17,6 +17,7 @@ const requiredTools = [ "ragmir_search", "ragmir_ask", "ragmir_research", + "ragmir_expand", "ragmir_audit", "ragmir_evaluate", "ragmir_usage_report", @@ -61,6 +62,9 @@ try { if (status.chunksIndexed < 1) { throw new Error("MCP status reported an empty index.") } + if (status.mcpMaxOutputBytes !== 32_768) { + throw new Error(`MCP status reported an unexpected output budget: ${JSON.stringify(status)}`) + } if ( status.ingestionLimits?.maxFileBytes !== 50_000_000 || status.ingestionLimits?.maxFiles !== null || @@ -90,6 +94,47 @@ try { throw new Error("MCP compact search did not return snippet-only results.") } + const boundedSearchResult = await client.callTool( + { + name: "ragmir_search", + arguments: { + query: "offline retrieval approval", + topK: 2, + compact: true, + maxBytes: 1_024, + }, + }, + undefined, + { timeout: 10_000 }, + ) + const boundedSearch = parseJsonToolResult(boundedSearchResult, "ragmir_search") + const outputMeta = boundedSearchResult._meta?.["ragmir/output"] + if ( + !Array.isArray(boundedSearch) || + typeof outputMeta?.returnedBytes !== "number" || + outputMeta.returnedBytes > 1_024 || + outputMeta.budgetBytes !== 1_024 + ) { + throw new Error( + `MCP search did not enforce its byte budget: ${JSON.stringify(boundedSearchResult)}`, + ) + } + + const expanded = await callJsonTool(client, "ragmir_expand", { + citation: searchResults[0].citation, + contextRadius: 1, + maxBytes: 1_024, + }) + if ( + expanded.found !== true || + !Array.isArray(expanded.passages) || + !expanded.passages.some((passage) => passage.citation === searchResults[0].citation) + ) { + throw new Error( + `MCP citation expansion returned an unexpected result: ${JSON.stringify(expanded)}`, + ) + } + const filteredSearchResults = await callJsonTool(client, "ragmir_search", { query: "offline retrieval approval", topK: 5, @@ -108,10 +153,14 @@ try { const answer = await callJsonTool(client, "ragmir_ask", { query: "What evidence supports offline operation?", topK: 2, + compact: true, }) if (!Array.isArray(answer.sources) || answer.sources.length < 1) { throw new Error("MCP ask returned no cited sources.") } + if (!("snippet" in answer.sources[0]) || "text" in answer.sources[0]) { + throw new Error("MCP compact ask did not return snippet-only sources.") + } const research = await callJsonTool(client, "ragmir_research", { query: "offline retrieval approval", @@ -140,6 +189,9 @@ try { if (typeof usage.averageResultCountByAction?.search !== "number") { throw new Error(`MCP usage report omitted per-action averages: ${JSON.stringify(usage)}`) } + if (usage.mcpOutput?.responses < 1 || usage.mcpOutput.returnedBytes < 1) { + throw new Error(`MCP usage report omitted output metrics: ${JSON.stringify(usage)}`) + } if (JSON.stringify(usage).includes(demoRoot)) { throw new Error("MCP usage report should not expose local project paths.") } @@ -152,6 +204,7 @@ try { tools: requiredTools, chunksIndexed: status.chunksIndexed, searchResults: searchResults.length, + expandedPassages: expanded.passages.length, askSources: answer.sources.length, researchEvidence: research.evidence.length, evaluationRecall: evaluation.recall, @@ -182,6 +235,10 @@ function runCliJson(args) { async function callJsonTool(client, name, args) { const result = await client.callTool({ name, arguments: args }, undefined, { timeout: 10_000 }) + return parseJsonToolResult(result, name) +} + +function parseJsonToolResult(result, name) { if (result.isError) { throw new Error(`${name} returned an MCP error.`) } diff --git a/packages/ragmir-core/skills/ragmir/SKILL.md b/packages/ragmir-core/skills/ragmir/SKILL.md index 77c04ba..02282ba 100644 --- a/packages/ragmir-core/skills/ragmir/SKILL.md +++ b/packages/ragmir-core/skills/ragmir/SKILL.md @@ -48,6 +48,7 @@ acceptable. Normal confidential indexing keeps remote model loading disabled. | Check unindexed or unsupported files | `rgr audit --unsupported` | | Check privacy posture | `rgr security-audit` | | Retrieve exact passages | `rgr search "query" --compact` or `ragmir_search` | +| Expand one returned citation | `ragmir_expand` | | Gather broad cited evidence | `rgr research "topic" --compact` or `ragmir_research` | | Return deterministic context | `rgr ask "question"` or `ragmir_ask` | | Measure retrieval recall | `rgr evaluate --golden ` or `ragmir_evaluate` | @@ -99,6 +100,11 @@ Cline. A generic server configuration is: When a client cannot set `cwd`, set `RAGMIR_PROJECT_ROOT` for the server process. Prefer MCP tools when available. Use CLI commands when they are not. +Prefer compact search, ask, or research output first. Call `ragmir_expand` with a returned citation +only when the exact chunk or neighboring context is needed. Retrieval tools accept `maxBytes`, but +the configured `mcpMaxOutputBytes` remains the hard ceiling. Inspect `_meta["ragmir/output"]` to see +whether the response was compacted or truncated. + MCP is read-focused. Only remove an index with the explicit shell command: ```sh diff --git a/packages/ragmir-core/src/access-log.test.ts b/packages/ragmir-core/src/access-log.test.ts index bf8d823..643a085 100644 --- a/packages/ragmir-core/src/access-log.test.ts +++ b/packages/ragmir-core/src/access-log.test.ts @@ -2,7 +2,7 @@ import { appendFile, mkdtemp, readFile, rm, stat, writeFile } from "node:fs/prom import os from "node:os" import path from "node:path" import { afterEach, describe, expect, it } from "vitest" -import { accessLogUsageReport, recordAccess } from "./access-log.js" +import { accessLogUsageReport, recordAccess, recordMcpOutput } from "./access-log.js" import { loadConfig } from "./config.js" import { initProject } from "./init.js" @@ -77,6 +77,71 @@ describe("accessLogUsageReport", () => { expect(report.averageResultCountByAction.ask).toBe(2) expect(report.averageResultCountByAction.evaluate).toBeNull() }) + + it("should aggregate metadata-only MCP output metrics without storing corpus details", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-mcp-output-log-")) + tempDirs.push(root) + await initProject(root) + const config = await loadConfig(root) + + await recordMcpOutput(config, { + tool: "ragmir_search", + retrievedBytes: 8_000, + returnedBytes: 2_000, + compacted: true, + truncated: false, + }) + await recordMcpOutput(config, { + tool: "ragmir_expand", + retrievedBytes: 3_000, + returnedBytes: 1_000, + compacted: false, + truncated: true, + }) + + const report = await accessLogUsageReport({ cwd: root, days: 7 }) + const rawLog = await readFile(config.accessLogPath, "utf8") + + expect(report.mcpOutput).toEqual({ + responses: 2, + retrievedBytes: 11_000, + returnedBytes: 3_000, + savedBytes: 8_000, + reductionRatio: 8 / 11, + compactedResponses: 1, + truncatedResponses: 1, + }) + expect(rawLog).not.toContain("private query") + expect(rawLog).not.toContain("raw/policy.md") + expect(rawLog).not.toContain(root) + }) + + it("should reject malformed MCP metric lines while reading older access entries", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-mcp-output-invalid-")) + tempDirs.push(root) + await initProject(root) + const config = await loadConfig(root) + await recordAccess(config, { action: "search", query: "known query", resultCount: 1 }) + await appendFile( + config.accessLogPath, + `${JSON.stringify({ + timestamp: new Date().toISOString(), + kind: "mcp-output", + tool: "ragmir_search", + retrievedBytes: -1, + returnedBytes: 1, + compacted: false, + truncated: false, + })}\n`, + "utf8", + ) + + const report = await accessLogUsageReport({ cwd: root, days: 7 }) + + expect(report.eventsByAction.search).toBe(1) + expect(report.mcpOutput.responses).toBe(0) + expect(report.invalidLines).toBe(1) + }) }) describe("recordAccess retention", () => { diff --git a/packages/ragmir-core/src/access-log.ts b/packages/ragmir-core/src/access-log.ts index d5f1237..6206d38 100644 --- a/packages/ragmir-core/src/access-log.ts +++ b/packages/ragmir-core/src/access-log.ts @@ -9,6 +9,7 @@ import type { AccessLogUsageOptions, AccessLogUsageReport, Config, + McpOutputTool, } from "./types.js" export interface AccessLogEvent { @@ -19,6 +20,14 @@ export interface AccessLogEvent { redactions?: number } +export interface McpOutputLogEvent { + tool: McpOutputTool + retrievedBytes: number + returnedBytes: number + compacted: boolean + truncated: boolean +} + interface AccessLogLine { timestamp: string action: AccessLogAction @@ -26,6 +35,18 @@ interface AccessLogLine { resultCount?: number } +interface McpOutputLogLine { + timestamp: string + kind: "mcp-output" + tool: McpOutputTool + retrievedBytes: number + returnedBytes: number + compacted: boolean + truncated: boolean +} + +type ParsedAccessLogLine = AccessLogLine | McpOutputLogLine + interface ResultCountStats { total: number events: number @@ -40,6 +61,13 @@ const ACCESS_LOG_ACTIONS: AccessLogAction[] = [ "destroy-index", ] const ACCESS_LOG_ACTION_SET = new Set(ACCESS_LOG_ACTIONS) +const MCP_OUTPUT_TOOLS: McpOutputTool[] = [ + "ragmir_search", + "ragmir_ask", + "ragmir_research", + "ragmir_expand", +] +const MCP_OUTPUT_TOOL_SET = new Set(MCP_OUTPUT_TOOLS) const DEFAULT_USAGE_REPORT_DAYS = 7 const MILLISECONDS_PER_DAY = 24 * 60 * 60 * 1000 /** @@ -70,6 +98,33 @@ export async function recordAccess(config: Config, event: AccessLogEvent): Promi } } +export async function recordMcpOutput(config: Config, event: McpOutputLogEvent): Promise { + if (!config.accessLog) { + return + } + + try { + await ensurePrivateDirectory(path.dirname(config.accessLogPath)) + await trimAccessLogIfNeeded(config.accessLogPath) + const line: McpOutputLogLine = { + timestamp: new Date().toISOString(), + kind: "mcp-output", + tool: event.tool, + retrievedBytes: event.retrievedBytes, + returnedBytes: event.returnedBytes, + compacted: event.compacted, + truncated: event.truncated, + } + await appendFile(config.accessLogPath, `${JSON.stringify(line)}\n`, { + encoding: "utf8", + mode: 0o600, + }) + await hardenPrivateFile(config.accessLogPath) + } catch { + // Output metrics are best-effort and must never block local retrieval. + } +} + async function trimAccessLogIfNeeded(accessLogPath: string): Promise { let size = 0 try { @@ -109,6 +164,11 @@ export async function accessLogUsageReport( let invalidLines = 0 let resultCountTotal = 0 let resultCountEvents = 0 + let mcpOutputResponses = 0 + let mcpRetrievedBytes = 0 + let mcpReturnedBytes = 0 + let compactedResponses = 0 + let truncatedResponses = 0 let lastEventAt: string | null = null if (existsSync(config.accessLogPath)) { @@ -128,6 +188,17 @@ export async function accessLogUsageReport( continue } totalEvents += 1 + if (isMcpOutputLogLine(event)) { + mcpOutputResponses += 1 + mcpRetrievedBytes += event.retrievedBytes + mcpReturnedBytes += event.returnedBytes + compactedResponses += event.compacted ? 1 : 0 + truncatedResponses += event.truncated ? 1 : 0 + if (lastEventAt === null || event.timestamp > lastEventAt) { + lastEventAt = event.timestamp + } + continue + } eventsByAction[event.action] += 1 if (event.queryHash) { queryHashes.add(event.queryHash) @@ -154,6 +225,15 @@ export async function accessLogUsageReport( uniqueQueryHashes: queryHashes.size, averageResultCount: resultCountEvents === 0 ? null : resultCountTotal / resultCountEvents, averageResultCountByAction: averageResultCounts(resultCountsByAction), + mcpOutput: { + responses: mcpOutputResponses, + retrievedBytes: mcpRetrievedBytes, + returnedBytes: mcpReturnedBytes, + savedBytes: mcpRetrievedBytes - mcpReturnedBytes, + reductionRatio: mcpRetrievedBytes === 0 ? null : 1 - mcpReturnedBytes / mcpRetrievedBytes, + compactedResponses, + truncatedResponses, + }, lastEventAt, } } @@ -246,18 +326,35 @@ function normalizeUsageReportDays(days: number | undefined): number { return days } -function parseAccessLogLine(line: string): AccessLogLine | null { +function parseAccessLogLine(line: string): ParsedAccessLogLine | null { try { const parsed: unknown = JSON.parse(line) - if (!isAccessLogLine(parsed)) { - return null + if (isAccessLogLine(parsed) || isMcpOutputLogLine(parsed)) { + return parsed } - return parsed + return null } catch { return null } } +function isMcpOutputLogLine(value: unknown): value is McpOutputLogLine { + return ( + typeof value === "object" && + value !== null && + hasTimestamp(value) && + "kind" in value && + value.kind === "mcp-output" && + "tool" in value && + typeof value.tool === "string" && + MCP_OUTPUT_TOOL_SET.has(value.tool) && + hasNonNegativeNumber(value, "retrievedBytes") && + hasNonNegativeNumber(value, "returnedBytes") && + hasBoolean(value, "compacted") && + hasBoolean(value, "truncated") + ) +} + function isAccessLogLine(value: unknown): value is AccessLogLine { return ( typeof value === "object" && @@ -286,3 +383,12 @@ function hasOptionalQueryHash(value: object): value is { queryHash?: string } { function hasOptionalResultCount(value: object): value is { resultCount?: number } { return !("resultCount" in value) || typeof value.resultCount === "number" } + +function hasNonNegativeNumber(value: object, key: string): boolean { + const field = Reflect.get(value, key) + return typeof field === "number" && Number.isFinite(field) && field >= 0 +} + +function hasBoolean(value: object, key: string): boolean { + return typeof Reflect.get(value, key) === "boolean" +} diff --git a/packages/ragmir-core/src/chunking.test.ts b/packages/ragmir-core/src/chunking.test.ts index a5a3fe6..fa4e19e 100644 --- a/packages/ragmir-core/src/chunking.test.ts +++ b/packages/ragmir-core/src/chunking.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest" -import { chunkDocument } from "./chunking.js" +import { chunkDocument, chunkSearchText } from "./chunking.js" import type { ParsedDocument } from "./types.js" describe("chunkDocument", () => { @@ -119,4 +119,88 @@ describe("chunkDocument", () => { expect(chunks[0]?.text).toBe("const alpha = 1") expect(chunks[1]?.text).toBe("const beta = 2") }) + + it("should keep Markdown context separate from cited text when sections are chunked", () => { + const doc: ParsedDocument = { + file: { + absolutePath: "/tmp/guide.md", + relativePath: ".ragmir/raw/guide.md", + source: "guide.md", + extension: ".md", + bytes: 200, + mtimeMs: 1, + checksum: "guide", + }, + text: "# Install\nOverview text.\n\n## MCP\nBounded output details.", + } + + const chunks = chunkDocument(doc, 40, 8) + const mcpChunk = chunks.find((chunk) => chunk.contextPath === "Install > MCP") + + expect(mcpChunk?.text).toContain("## MCP") + expect(mcpChunk?.text).not.toContain("Install > MCP") + expect(mcpChunk && chunkSearchText(mcpChunk)).toContain("Install > MCP\n## MCP") + }) + + it("should avoid splitting a fenced block when the complete block fits", () => { + const doc: ParsedDocument = { + file: { + absolutePath: "/tmp/code.md", + relativePath: ".ragmir/raw/code.md", + source: "code.md", + extension: ".md", + bytes: 200, + mtimeMs: 1, + checksum: "code", + }, + text: [ + "# API", + "Intro text before code.", + "", + "```ts", + "const alpha = 1", + "const beta = 2", + "```", + "After.", + ].join("\n"), + } + + const chunks = chunkDocument(doc, 52, 8) + const codeChunk = chunks.find((chunk) => chunk.text.includes("const alpha")) + + expect(codeChunk?.text).toContain("const beta = 2") + expect(codeChunk?.text).toContain("```ts") + expect(codeChunk?.text).toContain("```") + }) + + it("should retain exact JSON offsets while adding JSONPath context", () => { + const text = JSON.stringify( + { + primary: { owner: "alice", description: "A".repeat(50) }, + secondary: { owner: "bob", description: "B".repeat(50) }, + }, + null, + 2, + ) + const doc: ParsedDocument = { + file: { + absolutePath: "/tmp/projects.json", + relativePath: ".ragmir/raw/projects.json", + source: "projects.json", + extension: ".json", + bytes: Buffer.byteLength(text), + mtimeMs: 1, + checksum: "json", + }, + text, + } + + const chunks = chunkDocument(doc, 60, 0) + + expect(chunks.some((chunk) => chunk.contextPath.startsWith("$.primary"))).toBe(true) + expect(chunks.some((chunk) => chunk.contextPath.startsWith("$.secondary"))).toBe(true) + expect(chunks.every((chunk) => text.slice(chunk.charStart, chunk.charEnd) === chunk.text)).toBe( + true, + ) + }) }) diff --git a/packages/ragmir-core/src/chunking.ts b/packages/ragmir-core/src/chunking.ts index 8e4a334..617a107 100644 --- a/packages/ragmir-core/src/chunking.ts +++ b/packages/ragmir-core/src/chunking.ts @@ -1,4 +1,5 @@ import { createHash } from "node:crypto" +import { markdownFenceSpans, structuralSpans } from "./document-structure.js" import type { ParsedDocument, TextChunk } from "./types.js" const PARAGRAPH_BREAK_MIN_RATIO = 0.45 @@ -18,44 +19,68 @@ export function chunkDocument( const chunks: TextChunk[] = [] const lineStarts = lineStartOffsets(document.text) - let cursor = 0 + const structured = structuralSpans(document.text, document.file.extension, chunkSize) + const regions = + structured.length > 0 + ? structured + : [ + { + charStart: 0, + charEnd: document.text.length, + contextPath: "", + kind: "markdown-section" as const, + }, + ] + const fences = [".md", ".mdx", ".markdown"].includes(document.file.extension) + ? markdownFenceSpans(document.text) + : [] let chunkIndex = 0 - while (cursor < document.text.length) { - const end = chooseChunkEnd(document.text, cursor, chunkSize) - const span = trimmedSpan(document.text, cursor, end) - - if (span.text) { - const id = createHash("sha256") - .update(`${document.file.relativePath}:${chunkIndex}:${span.text}`) - .digest("hex") - chunks.push({ - id, - source: document.file.source, - relativePath: document.file.relativePath, - chunkIndex, - text: span.text, - charStart: span.start, - charEnd: span.end, - lineStart: lineNumberForOffset(lineStarts, span.start), - lineEnd: lineNumberForOffset(lineStarts, Math.max(span.start, span.end - 1)), - ...pageRangeForSpan(document, span.start, span.end), - checksum: document.file.checksum, - bytes: document.file.bytes, - mtimeMs: document.file.mtimeMs, - }) - chunkIndex += 1 + for (const region of regions) { + let cursor = region.charStart + while (cursor < region.charEnd) { + const end = chooseChunkEnd(document.text, cursor, chunkSize, region.charEnd, fences) + const span = trimmedSpan(document.text, cursor, end) + + if (span.text) { + const id = createHash("sha256") + .update(`${document.file.relativePath}:${chunkIndex}:${region.contextPath}:${span.text}`) + .digest("hex") + chunks.push({ + id, + source: document.file.source, + relativePath: document.file.relativePath, + chunkIndex, + contextPath: region.contextPath, + text: span.text, + charStart: span.start, + charEnd: span.end, + lineStart: lineNumberForOffset(lineStarts, span.start), + lineEnd: lineNumberForOffset(lineStarts, Math.max(span.start, span.end - 1)), + ...pageRangeForSpan(document, span.start, span.end), + checksum: document.file.checksum, + bytes: document.file.bytes, + mtimeMs: document.file.mtimeMs, + }) + chunkIndex += 1 + } + + if (end >= region.charEnd) { + break + } + cursor = fences.some((fence) => fence.start === end) + ? end + : Math.max(end - chunkOverlap, cursor + 1, region.charStart) } - - if (end >= document.text.length) { - break - } - cursor = Math.max(end - chunkOverlap, cursor + 1) } return chunks } +export function chunkSearchText(chunk: Pick): string { + return chunk.contextPath ? `${chunk.contextPath}\n${chunk.text}` : chunk.text +} + function pageRangeForSpan( document: ParsedDocument, start: number, @@ -70,12 +95,25 @@ function pageRangeForSpan( return first && last ? { pageStart: first.pageNumber, pageEnd: last.pageNumber } : {} } -function chooseChunkEnd(text: string, cursor: number, chunkSize: number): number { - const hardEnd = Math.min(cursor + chunkSize, text.length) - if (hardEnd === text.length) { +function chooseChunkEnd( + text: string, + cursor: number, + chunkSize: number, + regionEnd: number, + fences: Array<{ start: number; end: number }>, +): number { + const hardEnd = Math.min(cursor + chunkSize, regionEnd) + if (hardEnd === regionEnd) { return hardEnd } + const intersectedFence = fences.find( + (fence) => fence.start < hardEnd && fence.end > hardEnd && fence.end <= regionEnd, + ) + if (intersectedFence && intersectedFence.end - intersectedFence.start <= chunkSize) { + return intersectedFence.start > cursor ? intersectedFence.start : intersectedFence.end + } + const window = text.slice(cursor, hardEnd) const paragraphBreak = window.lastIndexOf("\n\n") if (paragraphBreak > chunkSize * PARAGRAPH_BREAK_MIN_RATIO) { diff --git a/packages/ragmir-core/src/cli.ts b/packages/ragmir-core/src/cli.ts index a01866e..0ea2add 100644 --- a/packages/ragmir-core/src/cli.ts +++ b/packages/ragmir-core/src/cli.ts @@ -838,6 +838,7 @@ program redactionEnabled: config.redaction.enabled, accessLog: config.accessLog, mcpMaxTopK: config.mcpMaxTopK, + mcpMaxOutputBytes: config.mcpMaxOutputBytes, topK: config.topK, chunkSize: config.chunkSize, chunkOverlap: config.chunkOverlap, @@ -870,6 +871,7 @@ program console.log(`redactionEnabled=${config.redaction.enabled}`) console.log(`accessLog=${config.accessLog}`) console.log(`mcpMaxTopK=${config.mcpMaxTopK}`) + console.log(`mcpMaxOutputBytes=${config.mcpMaxOutputBytes}`) console.log(`topK=${config.topK}`) console.log(`chunkSize=${config.chunkSize}`) console.log(`chunkOverlap=${config.chunkOverlap}`) @@ -909,6 +911,7 @@ program console.log(`accessLogStoresRawQueries=${report.accessLog.storesRawQueries}`) console.log(`storageGitIgnored=${report.storage.gitIgnored}`) console.log(`mcpMaxTopK=${report.mcp.maxTopK}`) + console.log(`mcpMaxOutputBytes=${report.mcp.maxOutputBytes}`) console.log(`mcpDestructiveToolsExposed=${report.mcp.destructiveToolsExposed}`) for (const warning of report.warnings) { console.log(pc.yellow(`warning: ${warning}`)) diff --git a/packages/ragmir-core/src/config.test.ts b/packages/ragmir-core/src/config.test.ts index ef1b938..324a1fe 100644 --- a/packages/ragmir-core/src/config.test.ts +++ b/packages/ragmir-core/src/config.test.ts @@ -43,6 +43,7 @@ describe("loadConfig", () => { expect(config.redaction.enabled).toBe(true) expect(config.accessLog).toBe(true) expect(config.mcpMaxTopK).toBe(10) + expect(config.mcpMaxOutputBytes).toBe(32_768) expect(config.sources).toEqual([]) expect(config.includeExtensions).toEqual([]) expect(config.pdfOcrCommand).toEqual([]) @@ -308,6 +309,31 @@ describe("loadConfig", () => { } }) + it("should override the MCP output budget from env and reject undersized config values", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-mcp-output-budget-")) + tempDirs.push(root) + await mkdir(path.join(root, ".ragmir"), { recursive: true }) + await writeFile(path.join(root, ".ragmir/config.json"), "{}\n", "utf8") + + const original = process.env.RAGMIR_MCP_MAX_OUTPUT_BYTES + process.env.RAGMIR_MCP_MAX_OUTPUT_BYTES = "8192" + try { + expect((await loadConfig(root)).mcpMaxOutputBytes).toBe(8_192) + process.env.RAGMIR_MCP_MAX_OUTPUT_BYTES = "invalid" + expect((await loadConfig(root)).mcpMaxOutputBytes).toBe(32_768) + process.env.RAGMIR_MCP_MAX_OUTPUT_BYTES = "512" + expect((await loadConfig(root)).mcpMaxOutputBytes).toBe(32_768) + } finally { + restoreEnv("RAGMIR_MCP_MAX_OUTPUT_BYTES", original) + } + + await writeFile( + path.join(root, ".ragmir/config.json"), + JSON.stringify({ mcpMaxOutputBytes: 512 }), + ) + await expect(loadConfig(root)).rejects.toThrow() + }) + it("defaults hybridTextScanLimit and overrides it from env", async () => { const root = await mkdtemp(path.join(os.tmpdir(), "jcode-kb-scan-limit-")) tempDirs.push(root) @@ -349,6 +375,7 @@ describe("loadConfig", () => { transformersAllowRemoteModels: true, redaction: { enabled: false, builtIn: false }, mcpMaxTopK: 50, + mcpMaxOutputBytes: 100_000, pdfOcrCommand: ["ocr-wrapper"], }), ) @@ -362,6 +389,7 @@ describe("loadConfig", () => { expect(config.redaction.enabled).toBe(true) expect(config.redaction.builtIn).toBe(true) expect(config.mcpMaxTopK).toBe(5) + expect(config.mcpMaxOutputBytes).toBe(16_384) expect(config.pdfOcrCommand).toEqual([]) } finally { restoreEnv("RAGMIR_TRANSFORMERS_ALLOW_REMOTE_MODELS", originalRemote) diff --git a/packages/ragmir-core/src/config.ts b/packages/ragmir-core/src/config.ts index c3cddab..a442a86 100644 --- a/packages/ragmir-core/src/config.ts +++ b/packages/ragmir-core/src/config.ts @@ -52,6 +52,7 @@ const rawConfigSchema = z .default(DEFAULT_CONFIG.redaction), accessLog: z.boolean().default(DEFAULT_CONFIG.accessLog), mcpMaxTopK: z.number().int().positive().default(DEFAULT_CONFIG.mcpMaxTopK), + mcpMaxOutputBytes: z.number().int().min(1_024).default(DEFAULT_CONFIG.mcpMaxOutputBytes), topK: z.number().int().positive().default(DEFAULT_CONFIG.topK), chunkSize: z.number().int().positive().default(DEFAULT_CONFIG.chunkSize), chunkOverlap: z.number().int().nonnegative().default(DEFAULT_CONFIG.chunkOverlap), @@ -153,6 +154,7 @@ export async function loadConfig(start = process.cwd()): Promise { redaction: effective.redaction, accessLog: effective.accessLog, mcpMaxTopK: effective.mcpMaxTopK, + mcpMaxOutputBytes: effective.mcpMaxOutputBytes, topK: effective.topK, chunkSize: effective.chunkSize, chunkOverlap: effective.chunkOverlap, @@ -199,6 +201,7 @@ function applyPrivacyFloor(config: RawConfig): RawConfig { transformersAllowRemoteModels: false, redaction: { ...config.redaction, enabled: true, builtIn: true }, mcpMaxTopK: Math.min(config.mcpMaxTopK, 5), + mcpMaxOutputBytes: Math.min(config.mcpMaxOutputBytes, 16_384), pdfOcrCommand: [], imageOcrCommand: [], legacyWordCommand: [], @@ -260,6 +263,12 @@ function applyEnv(config: RawConfig): RawConfig { }, accessLog: readBooleanEnv("RAGMIR_ACCESS_LOG", "KB_ACCESS_LOG", config.accessLog), mcpMaxTopK: readPositiveIntEnv("RAGMIR_MCP_MAX_TOP_K", "KB_MCP_MAX_TOP_K", config.mcpMaxTopK), + mcpMaxOutputBytes: readIntegerAtLeastEnv( + "RAGMIR_MCP_MAX_OUTPUT_BYTES", + "KB_MCP_MAX_OUTPUT_BYTES", + 1_024, + config.mcpMaxOutputBytes, + ), topK: readPositiveIntEnv("RAGMIR_TOP_K", "KB_TOP_K", config.topK), chunkSize: readPositiveIntEnv("RAGMIR_CHUNK_SIZE", "KB_CHUNK_SIZE", config.chunkSize), chunkOverlap: readNonNegativeIntEnv( @@ -381,6 +390,24 @@ function readPositiveIntEnv(name: string, legacyName: string, fallback: number): return value } +function readIntegerAtLeastEnv( + name: string, + legacyName: string, + minimum: number, + fallback: number, +): number { + const raw = process.env[name] ?? process.env[legacyName] + if (!raw) { + return fallback + } + const value = Number.parseInt(raw, 10) + if (!(Number.isInteger(value) && value >= minimum)) { + warnInvalidEnv(name, raw, `an integer greater than or equal to ${minimum}`) + return fallback + } + return value +} + function warnInvalidEnv(name: string, raw: string, expected: string): void { // Only the CLI writes to stdout; config warnings go to stderr so the operator // notices an env override that was silently ignored (e.g. RAGMIR_TOP_K=abc). diff --git a/packages/ragmir-core/src/defaults.ts b/packages/ragmir-core/src/defaults.ts index b6d54b8..0c9cf3e 100644 --- a/packages/ragmir-core/src/defaults.ts +++ b/packages/ragmir-core/src/defaults.ts @@ -39,6 +39,7 @@ export const DEFAULT_CONFIG: Omit = { }, accessLog: true, mcpMaxTopK: 10, + mcpMaxOutputBytes: 32_768, topK: 8, chunkSize: 1200, chunkOverlap: 200, diff --git a/packages/ragmir-core/src/document-structure.test.ts b/packages/ragmir-core/src/document-structure.test.ts new file mode 100644 index 0000000..be46751 --- /dev/null +++ b/packages/ragmir-core/src/document-structure.test.ts @@ -0,0 +1,108 @@ +import { describe, expect, it } from "vitest" +import { markdownFenceSpans, structuralSpans } from "./document-structure.js" + +describe("structuralSpans", () => { + it("should preserve nested heading paths when Markdown sections are detected", () => { + const text = [ + "Preamble.", + "", + "# Install", + "Install overview.", + "", + "## MCP", + "MCP details.", + "", + "### Codex", + "Codex details.", + ].join("\n") + + const spans = structuralSpans(text, ".md", 1_200) + + expect(spans.map((span) => span.contextPath)).toEqual([ + "", + "Install", + "Install > MCP", + "Install > MCP > Codex", + ]) + expect(spans.map((span) => text.slice(span.charStart, span.charEnd))).toEqual([ + "Preamble.\n\n", + "# Install\nInstall overview.\n\n", + "## MCP\nMCP details.\n\n", + "### Codex\nCodex details.", + ]) + }) + + it("should recognize Setext headings and ignore heading markers inside fences", () => { + const text = [ + "Guide", + "=====", + "", + "```md", + "# Not a heading", + "```", + "", + "Usage", + "-----", + "Use it.", + ].join("\n") + + const spans = structuralSpans(text, ".md", 1_200) + + expect(spans.map((span) => span.contextPath)).toEqual(["Guide", "Guide > Usage"]) + expect(markdownFenceSpans(text)).toHaveLength(1) + }) + + it("should expose exact JSON ranges with JSONPath context", () => { + const text = JSON.stringify( + { + projects: [ + { owner: "alice", status: "ready" }, + { owner: "bob", status: "blocked" }, + ], + policy: { retentionDays: 30 }, + }, + null, + 2, + ) + + const spans = structuralSpans(text, ".json", 70) + + expect(spans.length).toBeGreaterThan(1) + expect(spans.some((span) => span.contextPath.includes("$.projects"))).toBe(true) + expect(spans.some((span) => span.contextPath.includes("$.policy"))).toBe(true) + expect(spans.every((span) => text.slice(span.charStart, span.charEnd).length > 0)).toBe(true) + expect(spans.every((span) => span.charStart >= 0 && span.charEnd <= text.length)).toBe(true) + }) + + it("should keep repeated JSON keys distinct through their parent paths", () => { + const text = JSON.stringify( + { + primary: { owner: "alice", description: "A".repeat(50) }, + secondary: { owner: "bob", description: "B".repeat(50) }, + }, + null, + 2, + ) + + const spans = structuralSpans(text, ".json", 55) + const paths = spans.map((span) => span.contextPath) + + expect(paths.some((value) => value.startsWith("$.primary"))).toBe(true) + expect(paths.some((value) => value.startsWith("$.secondary"))).toBe(true) + }) + + it("should group valid JSONL entries without changing their source ranges", () => { + const text = ['{"id":1,"state":"ready"}', '{"id":2,"state":"blocked"}'].join("\n") + + const spans = structuralSpans(text, ".jsonl", 35) + + expect(spans).toHaveLength(2) + expect(spans[0]?.contextPath).toBe("$[1]") + expect(text.slice(spans[1]?.charStart, spans[1]?.charEnd)).toBe('{"id":2,"state":"blocked"}') + }) + + it("should fall back when JSON content is invalid", () => { + expect(structuralSpans('{"broken":', ".json", 100)).toEqual([]) + expect(structuralSpans('{"valid":true}\nnot-json', ".jsonl", 100)).toEqual([]) + }) +}) diff --git a/packages/ragmir-core/src/document-structure.ts b/packages/ragmir-core/src/document-structure.ts new file mode 100644 index 0000000..d3315f2 --- /dev/null +++ b/packages/ragmir-core/src/document-structure.ts @@ -0,0 +1,396 @@ +export type StructuralSpanKind = "markdown-section" | "json-node" | "jsonl-entry" + +export interface StructuralSpan { + charStart: number + charEnd: number + contextPath: string + kind: StructuralSpanKind +} + +interface TextLine { + start: number + end: number + contentEnd: number + text: string +} + +interface MarkdownHeading { + start: number + level: number + title: string +} + +interface JsonNode { + start: number + end: number + path: string + parentPath: string + children: JsonNode[] +} + +const MARKDOWN_EXTENSIONS = new Set([".md", ".mdx", ".markdown"]) +const ATX_HEADING_PATTERN = /^ {0,3}(#{1,6})\s+(.+?)(?:\s+#+)?\s*$/u +const SETEXT_HEADING_PATTERN = /^ {0,3}(=+|-+)\s*$/u +const FENCE_PATTERN = /^\s*(`{3,}|~{3,})/u + +export function structuralSpans( + text: string, + extension: string, + chunkSize: number, +): StructuralSpan[] { + if (MARKDOWN_EXTENSIONS.has(extension)) { + return markdownSectionSpans(text) + } + if (extension === ".json") { + return jsonNodeSpans(text, chunkSize) + } + if (extension === ".jsonl" || extension === ".ndjson") { + return jsonlEntrySpans(text, chunkSize) + } + return [] +} + +export function markdownFenceSpans(text: string): Array<{ start: number; end: number }> { + const lines = textLines(text) + const spans: Array<{ start: number; end: number }> = [] + let opening: { marker: string; start: number } | null = null + + for (const line of lines) { + const match = FENCE_PATTERN.exec(line.text) + if (!match?.[1]) { + continue + } + const marker = match[1] + if (!opening) { + opening = { marker: marker[0] ?? "`", start: line.start } + continue + } + if ((marker[0] ?? "") === opening.marker) { + spans.push({ start: opening.start, end: line.end }) + opening = null + } + } + + if (opening) { + spans.push({ start: opening.start, end: text.length }) + } + return spans +} + +function markdownSectionSpans(text: string): StructuralSpan[] { + const lines = textLines(text) + const headings: MarkdownHeading[] = [] + let fenceMarker: string | null = null + + for (let index = 0; index < lines.length; index += 1) { + const line = lines[index] + if (!line) { + continue + } + const fence = FENCE_PATTERN.exec(line.text)?.[1] + if (fence) { + const marker = fence[0] ?? "" + fenceMarker = fenceMarker === null ? marker : fenceMarker === marker ? null : fenceMarker + continue + } + if (fenceMarker !== null) { + continue + } + + const atx = ATX_HEADING_PATTERN.exec(line.text) + if (atx?.[1] && atx[2]) { + headings.push({ + start: line.start, + level: atx[1].length, + title: normalizeHeading(atx[2]), + }) + continue + } + + const setext = SETEXT_HEADING_PATTERN.exec(line.text) + const previous = lines[index - 1] + if (setext?.[1] && previous?.text.trim()) { + headings.push({ + start: previous.start, + level: setext[1].startsWith("=") ? 1 : 2, + title: normalizeHeading(previous.text), + }) + } + } + + if (headings.length === 0) { + return [] + } + + const spans: StructuralSpan[] = [] + const firstHeading = headings[0] + if (firstHeading && firstHeading.start > 0) { + spans.push({ + charStart: 0, + charEnd: firstHeading.start, + contextPath: "", + kind: "markdown-section", + }) + } + + const hierarchy: string[] = [] + for (let index = 0; index < headings.length; index += 1) { + const heading = headings[index] + if (!heading) { + continue + } + hierarchy.length = Math.max(0, heading.level - 1) + hierarchy[heading.level - 1] = heading.title + const next = headings[index + 1] + spans.push({ + charStart: heading.start, + charEnd: next?.start ?? text.length, + contextPath: hierarchy.filter(Boolean).join(" > "), + kind: "markdown-section", + }) + } + return spans.filter((span) => span.charEnd > span.charStart) +} + +function jsonNodeSpans(text: string, chunkSize: number): StructuralSpan[] { + try { + const parser = new JsonRangeParser(text) + const root = parser.parse() + const selected = selectJsonNodes(root, chunkSize) + return groupJsonNodes(selected, chunkSize).map((node) => ({ + charStart: node.start, + charEnd: node.end, + contextPath: node.path, + kind: "json-node" as const, + })) + } catch { + return [] + } +} + +function selectJsonNodes(node: JsonNode, chunkSize: number): JsonNode[] { + if (node.end - node.start <= chunkSize || node.children.length === 0) { + return [node] + } + return node.children.flatMap((child) => selectJsonNodes(child, chunkSize)) +} + +function groupJsonNodes(nodes: JsonNode[], chunkSize: number): JsonNode[] { + const grouped: JsonNode[] = [] + for (const node of nodes) { + const previous = grouped.at(-1) + if ( + previous && + previous.parentPath === node.parentPath && + node.end - previous.start <= chunkSize + ) { + previous.end = node.end + previous.path = node.parentPath + continue + } + grouped.push({ ...node, children: [] }) + } + return grouped +} + +function jsonlEntrySpans(text: string, chunkSize: number): StructuralSpan[] { + const spans: StructuralSpan[] = [] + let groupStart = -1 + let groupEnd = -1 + let firstLine = 0 + let lastLine = 0 + + for (const [index, line] of textLines(text).entries()) { + if (!line.text.trim()) { + continue + } + try { + JSON.parse(line.text) + } catch { + return [] + } + if (groupStart >= 0 && line.contentEnd - groupStart > chunkSize) { + spans.push(jsonlSpan(groupStart, groupEnd, firstLine, lastLine)) + groupStart = -1 + } + if (groupStart < 0) { + groupStart = line.start + firstLine = index + 1 + } + groupEnd = line.contentEnd + lastLine = index + 1 + } + + if (groupStart >= 0) { + spans.push(jsonlSpan(groupStart, groupEnd, firstLine, lastLine)) + } + return spans +} + +function jsonlSpan( + start: number, + end: number, + firstLine: number, + lastLine: number, +): StructuralSpan { + return { + charStart: start, + charEnd: end, + contextPath: firstLine === lastLine ? `$[${firstLine}]` : `$[${firstLine}..${lastLine}]`, + kind: "jsonl-entry", + } +} + +function normalizeHeading(value: string): string { + return value.replace(/\s+/gu, " ").trim() +} + +function textLines(text: string): TextLine[] { + const lines: TextLine[] = [] + let start = 0 + for (let index = 0; index <= text.length; index += 1) { + if (index !== text.length && text[index] !== "\n") { + continue + } + const hasNewline = index < text.length + const contentEnd = index > start && text[index - 1] === "\r" ? index - 1 : index + lines.push({ + start, + end: hasNewline ? index + 1 : index, + contentEnd, + text: text.slice(start, contentEnd), + }) + start = index + 1 + } + return lines +} + +class JsonRangeParser { + private index = 0 + + constructor(private readonly text: string) {} + + parse(): JsonNode { + this.skipWhitespace() + const node = this.parseValue("$", "") + this.skipWhitespace() + if (this.index !== this.text.length) { + throw new Error("Unexpected JSON suffix.") + } + return node + } + + private parseValue(path: string, parentPath: string): JsonNode { + this.skipWhitespace() + const start = this.index + const current = this.text[this.index] + if (current === "{") { + return this.parseObject(path, parentPath, start) + } + if (current === "[") { + return this.parseArray(path, parentPath, start) + } + if (current === '"') { + this.parseString() + } else { + this.parsePrimitive() + } + return { start, end: this.index, path, parentPath, children: [] } + } + + private parseObject(path: string, parentPath: string, start: number): JsonNode { + this.index += 1 + const children: JsonNode[] = [] + this.skipWhitespace() + while (this.text[this.index] !== "}") { + const memberStart = this.index + const key = this.parseString() + this.skipWhitespace() + this.expect(":") + const childPath = appendJsonPath(path, key) + const value = this.parseValue(childPath, path) + children.push({ ...value, start: memberStart }) + this.skipWhitespace() + if (this.text[this.index] !== ",") { + break + } + this.index += 1 + this.skipWhitespace() + } + this.expect("}") + return { start, end: this.index, path, parentPath, children } + } + + private parseArray(path: string, parentPath: string, start: number): JsonNode { + this.index += 1 + const children: JsonNode[] = [] + let itemIndex = 0 + this.skipWhitespace() + while (this.text[this.index] !== "]") { + const childPath = `${path}[${itemIndex}]` + children.push(this.parseValue(childPath, path)) + itemIndex += 1 + this.skipWhitespace() + if (this.text[this.index] !== ",") { + break + } + this.index += 1 + this.skipWhitespace() + } + this.expect("]") + return { start, end: this.index, path, parentPath, children } + } + + private parseString(): string { + const start = this.index + this.expect('"') + let escaped = false + while (this.index < this.text.length) { + const character = this.text[this.index] + this.index += 1 + if (escaped) { + escaped = false + continue + } + if (character === "\\") { + escaped = true + continue + } + if (character === '"') { + const parsed: unknown = JSON.parse(this.text.slice(start, this.index)) + if (typeof parsed !== "string") { + throw new Error("Expected a JSON string.") + } + return parsed + } + } + throw new Error("Unterminated JSON string.") + } + + private parsePrimitive(): void { + const start = this.index + while (this.index < this.text.length && !/[\s,}\]]/u.test(this.text[this.index] ?? "")) { + this.index += 1 + } + JSON.parse(this.text.slice(start, this.index)) + } + + private skipWhitespace(): void { + while (/\s/u.test(this.text[this.index] ?? "")) { + this.index += 1 + } + } + + private expect(character: string): void { + if (this.text[this.index] !== character) { + throw new Error(`Expected ${character}.`) + } + this.index += 1 + } +} + +function appendJsonPath(path: string, key: string): string { + return /^[A-Za-z_$][A-Za-z0-9_$]*$/u.test(key) + ? `${path}.${key}` + : `${path}[${JSON.stringify(key)}]` +} diff --git a/packages/ragmir-core/src/index-diagnostics.ts b/packages/ragmir-core/src/index-diagnostics.ts index bc60ade..d365ac8 100644 --- a/packages/ragmir-core/src/index-diagnostics.ts +++ b/packages/ragmir-core/src/index-diagnostics.ts @@ -9,7 +9,7 @@ import type { Config } from "./types.js" * required metadata). A stored manifest with a lower schemaVersion means the * index predates the current code and should be rebuilt. */ -export const INDEX_SCHEMA_VERSION = 6 +export const INDEX_SCHEMA_VERSION = 7 /** * Detect a stale or incompatible index without re-scanning every source file. diff --git a/packages/ragmir-core/src/index-policy.ts b/packages/ragmir-core/src/index-policy.ts index 73bc2cf..d3c75ee 100644 --- a/packages/ragmir-core/src/index-policy.ts +++ b/packages/ragmir-core/src/index-policy.ts @@ -2,7 +2,7 @@ import { createHash } from "node:crypto" import type { Config } from "./types.js" const INDEX_CONTENT_POLICY_VERSION = 1 -const CHUNKING_ADAPTER_VERSION = 2 +const CHUNKING_ADAPTER_VERSION = 3 export function indexPolicyFingerprint(config: Config): string { const policy = { diff --git a/packages/ragmir-core/src/index.ts b/packages/ragmir-core/src/index.ts index e479359..9c71c71 100644 --- a/packages/ragmir-core/src/index.ts +++ b/packages/ragmir-core/src/index.ts @@ -27,7 +27,7 @@ export type { PackageManager, RagmirCommand } from "./package-manager.js" export { detectPackageManager, kbCommand, ragmirCommand, rgrCommand } from "./package-manager.js" export type { PromptRouteDecision, PromptRouteTool } from "./prompt-routing.js" export { routePrompt } from "./prompt-routing.js" -export { ask, search } from "./query.js" +export { ask, expandCitation, search } from "./query.js" export { redactText } from "./redaction.js" export { compactResearchReport, compactSearchResults, research } from "./research.js" export { securityAudit } from "./security.js" @@ -72,11 +72,15 @@ export type { EvaluationCaseResult, EvaluationOptions, EvaluationResult, + ExpandCitationOptions, + ExpandedCitation, GoldenQuery, IndexManifest, IndexManifestFile, IngestionLimitsReport, IngestResult, + McpOutputTool, + McpOutputUsageReport, ParsedPage, PrivacyProfile, ResearchEvidence, diff --git a/packages/ragmir-core/src/ingest.test.ts b/packages/ragmir-core/src/ingest.test.ts index 4bb6192..284ebe6 100644 --- a/packages/ragmir-core/src/ingest.test.ts +++ b/packages/ragmir-core/src/ingest.test.ts @@ -88,7 +88,7 @@ describe("ingest", () => { const manifest = await readIndexManifest(await loadConfig(root)) expect(manifest).not.toBeNull() expect(manifest?.embeddingProvider).toBe("local-hash") - expect(manifest?.schemaVersion).toBe(6) + expect(manifest?.schemaVersion).toBe(7) expect(manifest?.indexPolicyFingerprint).toBe(indexPolicyFingerprint(await loadConfig(root))) expect(manifest?.vectorDimension).toBeGreaterThan(0) expect(manifest?.vectorDistanceMetric).toBe("l2") diff --git a/packages/ragmir-core/src/ingest.ts b/packages/ragmir-core/src/ingest.ts index 04c6733..0a4b3ae 100644 --- a/packages/ragmir-core/src/ingest.ts +++ b/packages/ragmir-core/src/ingest.ts @@ -1,5 +1,5 @@ import { recordAccess } from "./access-log.js" -import { chunkDocument } from "./chunking.js" +import { chunkDocument, chunkSearchText } from "./chunking.js" import { loadConfig } from "./config.js" import { VECTOR_DISTANCE_METRIC } from "./defaults.js" import { embedTexts } from "./embeddings.js" @@ -134,10 +134,7 @@ export async function ingest(options: IngestOptions = {}): Promise const rows: VectorRow[] = [] for (let i = 0; i < allChunks.length; i += config.embeddingBatchSize) { const batch = allChunks.slice(i, i + config.embeddingBatchSize) - const embeddings = await embedTexts( - batch.map((chunk) => chunk.text), - config, - ) + const embeddings = await embedTexts(batch.map(chunkSearchText), config) for (const [index, chunk] of batch.entries()) { const vector = embeddings[index] if (!vector) { @@ -145,6 +142,7 @@ export async function ingest(options: IngestOptions = {}): Promise } rows.push({ ...chunk, + searchText: chunkSearchText(chunk), vector, embeddingProvider: config.embeddingProvider, embeddingModel: config.embeddingModel, diff --git a/packages/ragmir-core/src/mcp-output.test.ts b/packages/ragmir-core/src/mcp-output.test.ts new file mode 100644 index 0000000..3016cd1 --- /dev/null +++ b/packages/ragmir-core/src/mcp-output.test.ts @@ -0,0 +1,212 @@ +import { describe, expect, it } from "vitest" +import type { McpAskPayload, McpResearchPayload, McpSearchPayload } from "./mcp-output.js" +import { + budgetMcpJson, + fitAskPayload, + fitExpandedCitation, + fitResearchPayload, + fitSearchPayload, + resolveMcpOutputBudget, +} from "./mcp-output.js" +import type { ExpandedCitation, SearchResult } from "./types.js" + +describe("MCP output budgeting", () => { + it("should preserve the JSON root when the response exceeds the UTF-8 byte budget", () => { + const full = [searchResult("évidence ".repeat(600))] + const compact: McpSearchPayload = [ + { + source: "policy.md", + relativePath: "raw/policy.md", + chunkIndex: 0, + contextPath: "Policy > Controls", + citation: "raw/policy.md:L1-L2#0", + snippet: "évidence concise", + distance: 0.1, + lineStart: 1, + lineEnd: 2, + pageStart: null, + pageEnd: null, + }, + ] + + const bounded = budgetMcpJson({ + tool: "ragmir_search", + maxBytes: 1_024, + fullValue: full, + preferredValue: full, + compactValue: compact, + compacted: false, + reduce: fitSearchPayload, + }) + + expect(Array.isArray(JSON.parse(bounded.result.content[0].text))).toBe(true) + expect(Buffer.byteLength(bounded.result.content[0].text, "utf8")).toBeLessThanOrEqual(1_024) + expect(bounded.metadata.compacted).toBe(true) + expect(bounded.metadata.retrievedBytes).toBeGreaterThan(bounded.metadata.returnedBytes) + }) + + it("should keep the payload unchanged when it already fits", () => { + const value = [searchResult("short evidence")] + + const bounded = budgetMcpJson({ + tool: "ragmir_search", + maxBytes: 32_768, + fullValue: value, + preferredValue: value, + compacted: false, + reduce: fitSearchPayload, + }) + + expect(bounded.result.content).toEqual([{ type: "text", text: JSON.stringify(value) }]) + expect(bounded.metadata.truncated).toBe(false) + expect(bounded.metadata.returnedBytes).toBe(bounded.metadata.retrievedBytes) + }) + + it("should clamp requested budgets to the configured maximum", () => { + expect(resolveMcpOutputBudget(8_192, 20_000)).toBe(8_192) + expect(resolveMcpOutputBudget(8_192, 1_024)).toBe(1_024) + }) + + it("should retain an ask object when sources must be omitted", () => { + const value: McpAskPayload = { + answer: "answer ".repeat(400), + sources: [searchResult("source ".repeat(400))], + staleWarning: null, + } + + const fitted = fitAskPayload(value, 1_024) + + expect(Array.isArray(fitted.value.sources)).toBe(true) + expect(Buffer.byteLength(JSON.stringify(fitted.value), "utf8")).toBeLessThanOrEqual(1_024) + expect(fitted.truncated).toBe(true) + }) + + it("should trim research detail while preserving its report shape", () => { + const value = researchPayload("diagnostic ".repeat(300)) + value.query = "é".repeat(3_000) + + const fitted = fitResearchPayload(value, 1_024) + + expect(fitted.value.audit.totalChunks).toBe(4) + expect(fitted.value.sourceDiagnostics).toEqual({ + duplicateCandidates: [], + archiveCandidates: [], + mirrorCandidates: [], + }) + expect(Buffer.byteLength(JSON.stringify(fitted.value), "utf8")).toBeLessThanOrEqual(1_024) + }) + + it("should prioritize the requested passage when expansion must be truncated", () => { + const value: ExpandedCitation = { + requestedCitation: "raw/policy.md:L2-L3#1", + found: true, + relativePath: "raw/policy.md", + chunkIndex: 1, + contextRadius: 1, + passages: [ + contextPassage(0, "short before"), + contextPassage(1, "target évidence ".repeat(300)), + contextPassage(2, "short after"), + ], + } + + const fitted = fitExpandedCitation(value, 1_024) + + expect(fitted.value.passages[0]?.chunkIndex).toBe(1) + expect(Buffer.byteLength(JSON.stringify(fitted.value), "utf8")).toBeLessThanOrEqual(1_024) + expect(fitted.truncated).toBe(true) + }) + + it("should fit long Unicode citation metadata when the minimum budget is active", () => { + const relativePath = `raw/${"資料".repeat(900)}.md` + const value: ExpandedCitation = { + requestedCitation: `${relativePath}:L1-L2#0`, + found: false, + relativePath, + chunkIndex: 0, + contextRadius: 0, + passages: [], + } + + const bounded = budgetMcpJson({ + tool: "ragmir_expand", + maxBytes: 1_024, + fullValue: value, + preferredValue: value, + compacted: false, + reduce: fitExpandedCitation, + }) + + expect(Buffer.byteLength(bounded.result.content[0].text, "utf8")).toBeLessThanOrEqual(1_024) + expect(JSON.parse(bounded.result.content[0].text)).toMatchObject({ + found: false, + passages: [], + }) + expect(bounded.metadata.truncated).toBe(true) + }) +}) + +function searchResult(text: string): SearchResult { + return { + source: "policy.md", + relativePath: "raw/policy.md", + chunkIndex: 0, + contextPath: "Policy > Controls", + citation: "raw/policy.md:L1-L2#0", + text, + distance: 0.1, + charStart: 0, + charEnd: text.length, + lineStart: 1, + lineEnd: 2, + pageStart: null, + pageEnd: null, + context: [], + } +} + +function researchPayload(detail: string): McpResearchPayload { + return { + query: "release evidence", + generatedQueries: [detail, detail], + ready: true, + audit: { + supportedFiles: 1, + supportedBytes: 100, + largestFileBytes: 100, + skippedFiles: 0, + unsupportedFiles: 0, + oversizedFiles: 0, + indexedFiles: 1, + totalChunks: 4, + missingFromIndex: 0, + staleInIndex: 0, + emptyTextFiles: 0, + }, + securityWarnings: [detail], + sourceDiagnostics: { + duplicateCandidates: [{ key: "duplicate", files: [detail] }], + archiveCandidates: [{ relativePath: detail, reason: "archive" }], + mirrorCandidates: [{ relativePath: detail, reason: "mirror" }], + }, + evidence: [], + codeEvidence: [], + gaps: [detail], + nextSteps: [detail], + } +} + +function contextPassage(chunkIndex: number, text: string) { + return { + chunkIndex, + contextPath: "Policy > Controls", + text, + charStart: chunkIndex * 10, + charEnd: chunkIndex * 10 + text.length, + lineStart: chunkIndex + 1, + lineEnd: chunkIndex + 1, + pageStart: null, + pageEnd: null, + citation: `raw/policy.md:L${chunkIndex + 1}-L${chunkIndex + 1}#${chunkIndex}`, + } +} diff --git a/packages/ragmir-core/src/mcp-output.ts b/packages/ragmir-core/src/mcp-output.ts new file mode 100644 index 0000000..8fe844b --- /dev/null +++ b/packages/ragmir-core/src/mcp-output.ts @@ -0,0 +1,383 @@ +import type { + CompactSearchResult, + ExpandedCitation, + McpOutputTool, + ResearchEvidence, + ResearchReport, + SearchResult, +} from "./types.js" + +export const MIN_MCP_OUTPUT_BYTES = 1_024 + +export interface McpOutputMetadata { + tool: McpOutputTool + budgetBytes: number + retrievedBytes: number + returnedBytes: number + compacted: boolean + truncated: boolean + omittedItems: number + expandTool: "ragmir_expand" +} + +export interface McpTextResult { + [key: string]: unknown + content: [{ type: "text"; text: string }] + _meta: { + "ragmir/output": McpOutputMetadata + } +} + +export interface BudgetedMcpResult { + result: McpTextResult + metadata: McpOutputMetadata +} + +export type McpSearchPayload = Array + +export interface McpAskPayload { + answer: string + sources: McpSearchPayload + staleWarning: string | null +} + +interface CompactResearchEvidence extends Omit { + snippet: string +} + +export interface McpResearchPayload extends Omit { + evidence: Array +} + +interface ReducedPayload { + value: T + omittedItems: number + truncated: boolean +} + +interface BudgetMcpJsonOptions { + tool: McpOutputTool + maxBytes: number + fullValue: unknown + preferredValue: T + compactValue?: T + compacted: boolean + reduce: (value: T, maxBytes: number) => ReducedPayload +} + +export function resolveMcpOutputBudget(configured: number, requested?: number): number { + const configuredBudget = Math.max(MIN_MCP_OUTPUT_BYTES, Math.floor(configured)) + if (requested === undefined) { + return configuredBudget + } + return Math.max(MIN_MCP_OUTPUT_BYTES, Math.min(configuredBudget, Math.floor(requested))) +} + +export function budgetMcpJson(options: BudgetMcpJsonOptions): BudgetedMcpResult { + const budgetBytes = Math.max(MIN_MCP_OUTPUT_BYTES, Math.floor(options.maxBytes)) + const retrievedBytes = jsonBytes(options.fullValue) + let value = options.preferredValue + let compacted = options.compacted + + if ( + jsonBytes(value) > budgetBytes && + options.compactValue !== undefined && + jsonBytes(options.compactValue) < jsonBytes(value) + ) { + value = options.compactValue + compacted = true + } + + const reduced = options.reduce(value, budgetBytes) + const text = serializeJson(reduced.value) + const returnedBytes = Buffer.byteLength(text, "utf8") + if (returnedBytes > budgetBytes) { + throw new Error( + `MCP output reducer returned ${returnedBytes} bytes for a ${budgetBytes}-byte budget.`, + ) + } + + const metadata: McpOutputMetadata = { + tool: options.tool, + budgetBytes, + retrievedBytes, + returnedBytes, + compacted, + truncated: reduced.truncated, + omittedItems: reduced.omittedItems, + expandTool: "ragmir_expand", + } + return { + result: { + content: [{ type: "text", text }], + _meta: { "ragmir/output": metadata }, + }, + metadata, + } +} + +export function fitSearchPayload( + value: McpSearchPayload, + maxBytes: number, +): ReducedPayload { + if (jsonBytes(value) <= maxBytes) { + return { value, omittedItems: 0, truncated: false } + } + for (let length = value.length - 1; length >= 0; length -= 1) { + const candidate = value.slice(0, length) + if (jsonBytes(candidate) <= maxBytes) { + return { + value: candidate, + omittedItems: value.length - candidate.length, + truncated: true, + } + } + } + return { value: [], omittedItems: value.length, truncated: value.length > 0 } +} + +export function fitAskPayload( + value: McpAskPayload, + maxBytes: number, +): ReducedPayload { + if (jsonBytes(value) <= maxBytes) { + return { value, omittedItems: 0, truncated: false } + } + + for (let length = value.sources.length - 1; length >= 0; length -= 1) { + const candidate = { ...value, sources: value.sources.slice(0, length) } + if (jsonBytes(candidate) <= maxBytes) { + return { + value: candidate, + omittedItems: value.sources.length - candidate.sources.length, + truncated: true, + } + } + } + + const minimal: McpAskPayload = { + answer: "Retrieval output exceeded the active MCP byte budget.", + sources: [], + staleWarning: null, + } + return { + value: minimal, + omittedItems: value.sources.length, + truncated: true, + } +} + +export function fitResearchPayload( + value: McpResearchPayload, + maxBytes: number, +): ReducedPayload { + if (jsonBytes(value) <= maxBytes) { + return { value, omittedItems: 0, truncated: false } + } + + const candidate = cloneResearchPayload(value) + let omittedItems = 0 + while (jsonBytes(candidate) > maxBytes && removeResearchDetail(candidate)) { + omittedItems += 1 + } + if (jsonBytes(candidate) <= maxBytes) { + return { value: candidate, omittedItems, truncated: true } + } + + const minimal: McpResearchPayload = { + query: compactString(candidate.query, 80), + generatedQueries: [], + ready: candidate.ready, + audit: candidate.audit, + securityWarnings: [], + sourceDiagnostics: { + duplicateCandidates: [], + archiveCandidates: [], + mirrorCandidates: [], + }, + evidence: [], + codeEvidence: [], + gaps: [], + nextSteps: [], + } + return { value: minimal, omittedItems: omittedItems + 1, truncated: true } +} + +export function fitExpandedCitation( + value: ExpandedCitation, + maxBytes: number, +): ReducedPayload { + if (jsonBytes(value) <= maxBytes) { + return { value, omittedItems: 0, truncated: false } + } + + const prioritized = [...value.passages].sort( + (left, right) => + Math.abs(left.chunkIndex - value.chunkIndex) - + Math.abs(right.chunkIndex - value.chunkIndex) || left.chunkIndex - right.chunkIndex, + ) + const baseValue = fitExpandedMetadata(value, maxBytes) + const metadataTruncated = + baseValue.requestedCitation !== value.requestedCitation || + baseValue.relativePath !== value.relativePath + const target = prioritized.find((passage) => passage.chunkIndex === value.chunkIndex) + let passages: ExpandedCitation["passages"] = [] + if (target) { + const targetOnly: ExpandedCitation = { ...baseValue, passages: [target] } + if (jsonBytes(targetOnly) <= maxBytes) { + passages = [target] + } else { + const fittedTarget = fitExpandedTarget(baseValue, target, maxBytes) + if (!fittedTarget) { + return { + value: baseValue, + omittedItems: value.passages.length, + truncated: true, + } + } + passages = fittedTarget.passages + } + } + + for (const passage of prioritized.filter((candidate) => candidate !== target)) { + const candidate = { + ...baseValue, + passages: [...passages, passage].sort((left, right) => left.chunkIndex - right.chunkIndex), + } + if (jsonBytes(candidate) <= maxBytes) { + passages.push(passage) + } + } + if (passages.length > 0) { + return { + value: { + ...baseValue, + passages: passages.sort((left, right) => left.chunkIndex - right.chunkIndex), + }, + omittedItems: value.passages.length - passages.length, + truncated: true, + } + } + + return { + value: baseValue, + omittedItems: value.passages.length, + truncated: value.passages.length > 0 || metadataTruncated, + } +} + +function cloneResearchPayload(value: McpResearchPayload): McpResearchPayload { + return { + ...value, + generatedQueries: [...value.generatedQueries], + securityWarnings: [...value.securityWarnings], + sourceDiagnostics: { + duplicateCandidates: [...value.sourceDiagnostics.duplicateCandidates], + archiveCandidates: [...value.sourceDiagnostics.archiveCandidates], + mirrorCandidates: [...value.sourceDiagnostics.mirrorCandidates], + }, + evidence: [...value.evidence], + codeEvidence: [...value.codeEvidence], + gaps: [...value.gaps], + nextSteps: [...value.nextSteps], + } +} + +function removeResearchDetail(value: McpResearchPayload): boolean { + const arrays = [ + value.sourceDiagnostics.duplicateCandidates, + value.sourceDiagnostics.archiveCandidates, + value.sourceDiagnostics.mirrorCandidates, + value.codeEvidence, + value.evidence, + value.generatedQueries, + value.securityWarnings, + value.gaps, + value.nextSteps, + ] + const largest = arrays.reduce( + (current, items) => (items.length > (current?.length ?? 0) ? items : current), + null, + ) + if (!largest || largest.length === 0) { + return false + } + largest.pop() + return true +} + +function fitExpandedTarget( + value: ExpandedCitation, + target: ExpandedCitation["passages"][number], + maxBytes: number, +): ExpandedCitation | null { + const characters = [...target.text] + let low = 0 + let high = characters.length + let best: ExpandedCitation | null = null + while (low <= high) { + const middle = Math.floor((low + high) / 2) + const text = + middle < characters.length ? `${characters.slice(0, middle).join("")}...` : target.text + const candidate: ExpandedCitation = { + ...value, + passages: [{ ...target, text }], + } + if (jsonBytes(candidate) <= maxBytes) { + best = candidate + low = middle + 1 + } else { + high = middle - 1 + } + } + return best +} + +function fitExpandedMetadata(value: ExpandedCitation, maxBytes: number): ExpandedCitation { + const withoutPassages: ExpandedCitation = { ...value, passages: [] } + if (jsonBytes(withoutPassages) <= maxBytes) { + return withoutPassages + } + + let low = 0 + let high = Math.max([...value.requestedCitation].length, [...value.relativePath].length) + let best: ExpandedCitation = { + ...withoutPassages, + requestedCitation: "", + relativePath: "", + } + while (low <= high) { + const middle = Math.floor((low + high) / 2) + const candidate: ExpandedCitation = { + ...withoutPassages, + requestedCitation: compactString(value.requestedCitation, middle), + relativePath: compactString(value.relativePath, middle), + } + if (jsonBytes(candidate) <= maxBytes) { + best = candidate + low = middle + 1 + } else { + high = middle - 1 + } + } + return best +} + +function compactString(value: string, maxLength: number): string { + const characters = [...value] + if (characters.length <= maxLength) { + return value + } + if (maxLength <= 3) { + return ".".repeat(Math.max(0, maxLength)) + } + return `${characters.slice(0, maxLength - 3).join("")}...` +} + +function jsonBytes(value: unknown): number { + return Buffer.byteLength(serializeJson(value), "utf8") +} + +function serializeJson(value: unknown): string { + return JSON.stringify(value) +} diff --git a/packages/ragmir-core/src/mcp.ts b/packages/ragmir-core/src/mcp.ts index 42e3865..83b74e5 100644 --- a/packages/ragmir-core/src/mcp.ts +++ b/packages/ragmir-core/src/mcp.ts @@ -3,32 +3,54 @@ import path from "node:path" import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js" import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js" import { z } from "zod" -import { accessLogUsageReport } from "./access-log.js" +import { accessLogUsageReport, recordMcpOutput } from "./access-log.js" import { findProjectConfig, loadConfig } from "./config.js" import { RAGMIR_PROJECT_ROOT_ENV } from "./defaults.js" import { evaluateGoldenQueries } from "./evaluate.js" import { audit } from "./ingest.js" import { ingestionLimits } from "./limits.js" +import type { + BudgetedMcpResult, + McpAskPayload, + McpResearchPayload, + McpSearchPayload, +} from "./mcp-output.js" +import { + budgetMcpJson, + fitAskPayload, + fitExpandedCitation, + fitResearchPayload, + fitSearchPayload, + MIN_MCP_OUTPUT_BYTES, + resolveMcpOutputBudget, +} from "./mcp-output.js" import { routePrompt } from "./prompt-routing.js" -import { ask, search } from "./query.js" +import { ask, expandCitation, search } from "./query.js" import { compactResearchReport, compactSearchResults, research } from "./research.js" import { securityAudit } from "./security.js" import { countRows } from "./store.js" +import type { AskResult } from "./types.js" import { VERSION } from "./version.js" const queryToolInputSchema = z.object({ query: z.string().min(1), topK: z.number().int().positive().optional(), contextRadius: z.number().int().min(0).optional(), + maxBytes: z.number().int().min(MIN_MCP_OUTPUT_BYTES).optional(), includePaths: z.array(z.string().min(1).max(500)).max(20).optional(), excludePaths: z.array(z.string().min(1).max(500)).max(20).optional(), }) +const askToolInputSchema = queryToolInputSchema.extend({ + compact: z.boolean().optional(), +}) + const researchToolInputSchema = z.object({ query: z.string().min(1), topK: z.number().int().positive().optional(), includeCode: z.boolean().optional(), compact: z.boolean().optional(), + maxBytes: z.number().int().min(MIN_MCP_OUTPUT_BYTES).optional(), includePaths: z.array(z.string().min(1).max(500)).max(20).optional(), excludePaths: z.array(z.string().min(1).max(500)).max(20).optional(), }) @@ -51,6 +73,12 @@ const promptRouteInputSchema = z.object({ prompt: z.string().min(1), }) +const expandToolInputSchema = z.object({ + citation: z.string().min(1).max(2_000), + contextRadius: z.number().int().min(0).max(3).optional(), + maxBytes: z.number().int().min(MIN_MCP_OUTPUT_BYTES).optional(), +}) + export async function serveMcp(cwd = resolveMcpProjectRoot()): Promise { const server = new McpServer({ name: "ragmir", @@ -89,6 +117,7 @@ export async function serveMcp(cwd = resolveMcpProjectRoot()): Promise { llmGeneration: false, redactionEnabled: config.redaction.enabled, mcpMaxTopK: config.mcpMaxTopK, + mcpMaxOutputBytes: config.mcpMaxOutputBytes, maxFileBytes: config.maxFileBytes, ingestConcurrency: config.ingestConcurrency, embeddingBatchSize: config.embeddingBatchSize, @@ -125,14 +154,26 @@ export async function serveMcp(cwd = resolveMcpProjectRoot()): Promise { description: "Retrieve relevant passages from the local Ragmir knowledge base.", inputSchema: searchToolInputSchema, }, - async ({ query, topK, contextRadius, compact, includePaths, excludePaths }) => { + async ({ query, topK, contextRadius, compact, maxBytes, includePaths, excludePaths }) => { const config = await loadConfig(cwd) const results = await search( query, await searchOptions(cwd, topK, contextRadius, includePaths, excludePaths), ) - const compactOutput = compact ?? config.privacyProfile === "strict" - return textResult(compactOutput ? compactSearchResults(results) : results) + const compactResults = compactSearchResults(results) + const compactOutput = config.privacyProfile === "strict" || compact === true + const preferred: McpSearchPayload = compactOutput ? compactResults : results + const bounded = budgetMcpJson({ + tool: "ragmir_search", + maxBytes: resolveMcpOutputBudget(config.mcpMaxOutputBytes, maxBytes), + fullValue: results, + preferredValue: preferred, + compactValue: compactResults, + compacted: compactOutput, + reduce: fitSearchPayload, + }) + await recordBudgetedOutput(config, bounded) + return bounded.result }, ) @@ -141,19 +182,39 @@ export async function serveMcp(cwd = resolveMcpProjectRoot()): Promise { { title: "Ragmir Ask", description: "Return cited retrieval context for a question without calling an LLM.", - inputSchema: queryToolInputSchema, + inputSchema: askToolInputSchema, }, - async ({ query, topK, contextRadius, includePaths, excludePaths }) => { + async ({ query, topK, contextRadius, compact, maxBytes, includePaths, excludePaths }) => { const config = await loadConfig(cwd) const options = await searchOptions(cwd, topK, contextRadius, includePaths, excludePaths) + let fullPayload: AskResult if (config.privacyProfile === "strict") { const results = await search(query, options) - return textResult({ + fullPayload = { answer: "Strict privacy profile returns compact cited retrieval only.", - sources: compactSearchResults(results), - }) + sources: results, + staleWarning: null, + } + } else { + fullPayload = await ask(query, options) } - return textResult(await ask(query, options)) + const compactPayload: McpAskPayload = { + answer: "Ragmir returns compact cited retrieval only. Expand a citation when needed.", + sources: compactSearchResults(fullPayload.sources), + staleWarning: fullPayload.staleWarning, + } + const compactOutput = config.privacyProfile === "strict" || compact === true + const bounded = budgetMcpJson({ + tool: "ragmir_ask", + maxBytes: resolveMcpOutputBudget(config.mcpMaxOutputBytes, maxBytes), + fullValue: fullPayload, + preferredValue: compactOutput ? compactPayload : fullPayload, + compactValue: compactPayload, + compacted: compactOutput, + reduce: fitAskPayload, + }) + await recordBudgetedOutput(config, bounded) + return bounded.result }, ) @@ -165,7 +226,7 @@ export async function serveMcp(cwd = resolveMcpProjectRoot()): Promise { "Run an audit-backed multi-query research pass with cited evidence and optional code matches.", inputSchema: researchToolInputSchema, }, - async ({ query, topK, includeCode, compact, includePaths, excludePaths }) => { + async ({ query, topK, includeCode, compact, maxBytes, includePaths, excludePaths }) => { const config = await loadConfig(cwd) const options = await searchOptions(cwd, topK, undefined, includePaths, excludePaths) const researchOptions: Parameters[1] = { cwd } @@ -174,8 +235,46 @@ export async function serveMcp(cwd = resolveMcpProjectRoot()): Promise { addOption(researchOptions, "includePaths", options.includePaths) addOption(researchOptions, "excludePaths", options.excludePaths) const result = await research(query, researchOptions) - const compactOutput = compact ?? config.privacyProfile === "strict" - return textResult(compactOutput ? compactResearchReport(result) : result) + const compactResult = compactResearchReport(result) + const compactOutput = config.privacyProfile === "strict" || compact === true + const preferred: McpResearchPayload = compactOutput ? compactResult : result + const bounded = budgetMcpJson({ + tool: "ragmir_research", + maxBytes: resolveMcpOutputBudget(config.mcpMaxOutputBytes, maxBytes), + fullValue: result, + preferredValue: preferred, + compactValue: compactResult, + compacted: compactOutput, + reduce: fitResearchPayload, + }) + await recordBudgetedOutput(config, bounded) + return bounded.result + }, + ) + + server.registerTool( + "ragmir_expand", + { + title: "Ragmir Expand", + description: "Expand one Ragmir citation into a bounded exact passage window.", + inputSchema: expandToolInputSchema, + }, + async ({ citation, contextRadius, maxBytes }) => { + const config = await loadConfig(cwd) + const expanded = await expandCitation(citation, { + cwd, + ...(contextRadius === undefined ? {} : { contextRadius }), + }) + const bounded = budgetMcpJson({ + tool: "ragmir_expand", + maxBytes: resolveMcpOutputBudget(config.mcpMaxOutputBytes, maxBytes), + fullValue: expanded, + preferredValue: expanded, + compacted: false, + reduce: fitExpandedCitation, + }) + await recordBudgetedOutput(config, bounded) + return bounded.result }, ) @@ -292,6 +391,19 @@ function textResult(value: unknown): { content: Array<{ type: "text"; text: stri } } +async function recordBudgetedOutput( + config: Awaited>, + bounded: BudgetedMcpResult, +): Promise { + await recordMcpOutput(config, { + tool: bounded.metadata.tool, + retrievedBytes: bounded.metadata.retrievedBytes, + returnedBytes: bounded.metadata.returnedBytes, + compacted: bounded.metadata.compacted, + truncated: bounded.metadata.truncated, + }) +} + export async function searchOptions( cwd: string, topK: number | undefined, diff --git a/packages/ragmir-core/src/query.test.ts b/packages/ragmir-core/src/query.test.ts index 09cdad9..b5fe42e 100644 --- a/packages/ragmir-core/src/query.test.ts +++ b/packages/ragmir-core/src/query.test.ts @@ -5,7 +5,7 @@ import { fileURLToPath } from "node:url" import { afterEach, describe, expect, it } from "vitest" import { ingest } from "./ingest.js" import { initProject } from "./init.js" -import { ask, search, vectorCandidateLimit } from "./query.js" +import { ask, expandCitation, search, vectorCandidateLimit } from "./query.js" const tempDirs: string[] = [] const packageRoot = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "..") @@ -110,6 +110,81 @@ describe("search", () => { expect(results[0]?.relativePath).toBe(".ragmir/raw/zeta.md") }) + it("should retrieve later Markdown chunks through their heading context", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-query-markdown-context-")) + tempDirs.push(root) + await initProject(root) + await mkdir(path.join(root, ".ragmir", "raw"), { recursive: true }) + await writeFile( + path.join(root, ".ragmir", "config.json"), + JSON.stringify({ chunkSize: 90, chunkOverlap: 10, topK: 4 }), + "utf8", + ) + await writeFile( + path.join(root, ".ragmir", "raw", "guide.md"), + [ + "# Operations", + "", + "## Zephyr authorization", + "The first paragraph describes the ordinary workflow with neutral wording repeated here.", + "", + "The second paragraph contains consequences but deliberately omits the section title.", + ].join("\n"), + "utf8", + ) + + await ingest({ cwd: root }) + const results = await search("zephyr authorization", { cwd: root, topK: 4 }) + + expect( + results.some((result) => result.contextPath === "Operations > Zephyr authorization"), + ).toBe(true) + expect( + results.some( + (result) => + !result.text.toLowerCase().includes("zephyr") && + result.contextPath.includes("Zephyr authorization"), + ), + ).toBe(true) + }) + + it("should retrieve JSON values through their JSONPath without exposing it as text", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-query-json-context-")) + tempDirs.push(root) + await initProject(root) + await mkdir(path.join(root, ".ragmir", "raw"), { recursive: true }) + await writeFile( + path.join(root, ".ragmir", "config.json"), + JSON.stringify({ chunkSize: 55, chunkOverlap: 0, topK: 3 }), + "utf8", + ) + await writeFile( + path.join(root, ".ragmir", "raw", "controls.json"), + JSON.stringify( + { + compliance: { + owner: "alice", + description: "A neutral control description that is long enough to split the object.", + }, + operations: { owner: "bob" }, + }, + null, + 2, + ), + "utf8", + ) + + await ingest({ cwd: root }) + const results = await search("compliance", { cwd: root, topK: 3 }) + const nested = results.find( + (result) => + result.contextPath.startsWith("$.compliance") && !result.text.includes("compliance"), + ) + + expect(nested?.relativePath).toBe(".ragmir/raw/controls.json") + expect(nested?.contextPath).toContain("$.compliance") + }) + it("matches accent-folded lexical queries", async () => { const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-query-accent-")) tempDirs.push(root) @@ -127,6 +202,25 @@ describe("search", () => { expect(results[0]?.relativePath).toBe(".ragmir/raw/policy.md") }) + it("should recover a single-token transposition without matching a nearby wrong word", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-query-typo-")) + tempDirs.push(root) + await initProject(root) + await mkdir(path.join(root, ".ragmir", "raw"), { recursive: true }) + await writeFile( + path.join(root, ".ragmir", "raw", "security.md"), + "Security controls require signed evidence.\n", + "utf8", + ) + await ingest({ cwd: root }) + + const typo = await search("sceurity", { cwd: root, topK: 1 }) + const wrongWord = await search("secretary", { cwd: root, topK: 1 }) + + expect(typo[0]?.relativePath).toBe(".ragmir/raw/security.md") + expect(wrongWord).toEqual([]) + }) + it("blocks search when the active index policy differs", async () => { const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-query-policy-")) tempDirs.push(root) @@ -140,6 +234,9 @@ describe("search", () => { ) await expect(search("policy", { cwd: root })).rejects.toThrow("Rebuild") + await expect(expandCitation(".ragmir/raw/policy.md:L1-L1#0", { cwd: root })).rejects.toThrow( + "Rebuild", + ) }) it("abstains when local-hash retrieval has no lexical evidence", async () => { @@ -260,6 +357,67 @@ describe("search", () => { expect(results[0]?.pageStart).toBe(1) expect(results[0]?.pageEnd).toBe(1) expect(results[0]?.citation).toContain("brief.pdf:p1:L") + + const expanded = await expandCitation(results[0]?.citation ?? "", { cwd: root }) + expect(expanded.found).toBe(true) + expect(expanded.passages[0]?.pageStart).toBe(1) + }) +}) + +describe("expandCitation", () => { + it("should return the exact indexed passage and bounded neighbors when a citation is valid", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-expand-")) + tempDirs.push(root) + await initProject(root) + await writeFile( + path.join(root, ".ragmir", "config.json"), + JSON.stringify({ chunkSize: 36, chunkOverlap: 0, topK: 1 }), + "utf8", + ) + await writeFile( + path.join(root, ".ragmir", "raw", "policy's.md"), + [ + "Opening operational note.", + "Rare expansion target evidence.", + "Closing consequence line.", + ].join("\n\n"), + "utf8", + ) + await ingest({ cwd: root }) + const [result] = await search("expansion target evidence", { cwd: root, topK: 1 }) + + const expanded = await expandCitation(result?.citation ?? "", { cwd: root, contextRadius: 20 }) + + expect(expanded.found).toBe(true) + expect(expanded.relativePath).toBe(".ragmir/raw/policy's.md") + expect(expanded.contextRadius).toBe(3) + expect(expanded.passages.some((passage) => passage.chunkIndex === result?.chunkIndex)).toBe( + true, + ) + + const forgedCitation = (result?.citation ?? "").replace(/:L\d+-L\d+#/u, ":L999-L1000#") + await expect(expandCitation(forgedCitation, { cwd: root })).rejects.toThrow( + "do not match the indexed passage", + ) + }) + + it("should return no passages when the cited chunk does not exist", async () => { + const root = await mkdtemp(path.join(os.tmpdir(), "ragmir-expand-missing-")) + tempDirs.push(root) + await initProject(root) + + await expect(expandCitation(".ragmir/raw/missing.md:L1-L1#4", { cwd: root })).resolves.toEqual({ + requestedCitation: ".ragmir/raw/missing.md:L1-L1#4", + found: false, + relativePath: ".ragmir/raw/missing.md", + chunkIndex: 4, + contextRadius: 0, + passages: [], + }) + }) + + it("should reject malformed citation input", async () => { + await expect(expandCitation("raw/policy.md:L1-L2")).rejects.toThrow("chunk suffix") }) }) diff --git a/packages/ragmir-core/src/query.ts b/packages/ragmir-core/src/query.ts index bf815b3..5d1063d 100644 --- a/packages/ragmir-core/src/query.ts +++ b/packages/ragmir-core/src/query.ts @@ -8,6 +8,8 @@ import { openRowsTable, readIndexManifest } from "./store.js" import { tokenize } from "./text.js" import type { AskResult, + ExpandCitationOptions, + ExpandedCitation, RetrievalProfile, SearchContextChunk, SearchOptions, @@ -20,6 +22,8 @@ interface SearchRow { source: string relativePath: string chunkIndex: number + contextPath: string + searchText: string text: string charStart?: number charEnd?: number @@ -63,10 +67,14 @@ const RRF_LEXICAL_WEIGHT = 1 const BM25_K1 = 1.2 const BM25_B = 0.75 const MAX_CONTEXT_RADIUS = 3 +const MIN_FUZZY_TOKEN_LENGTH = 7 +const MIN_TRIGRAM_DICE_SIMILARITY = 0.5 const SEARCH_COLUMNS = [ "source", "relativePath", "chunkIndex", + "contextPath", + "searchText", "text", "charStart", "charEnd", @@ -84,6 +92,7 @@ export async function search(query: string, options: SearchOptions = {}): Promis if (!table) { return [] } + await assertIndexFreshness(config) const sanitized = sanitizeRetrievalQuery(query) const queryTokens = tokenize(sanitized.query) @@ -105,7 +114,7 @@ export async function search(query: string, options: SearchOptions = {}): Promis const rankedRows = rankHybridRows(sanitized.query, vectorRows, textRows) const relevantRows = config.embeddingProvider === "local-hash" - ? rankedRows.filter((ranked) => hasLexicalOverlap(queryTokens, ranked.row.text)) + ? rankedRows.filter((ranked) => hasLexicalOverlap(queryTokens, ranked.row.searchText)) : rankedRows const rows = diversifyRows(relevantRows, topK, config.retrievalProfile) const defaultContextRadius = config.retrievalProfile === "quality" ? 1 : 0 @@ -119,6 +128,7 @@ export async function search(query: string, options: SearchOptions = {}): Promis source: row.row.source, relativePath: row.row.relativePath, chunkIndex: row.row.chunkIndex, + contextPath: row.row.contextPath, citation: citationForRow(row.row), text: row.row.text, distance: typeof row.row._distance === "number" ? row.row._distance : null, @@ -144,7 +154,7 @@ function diversifyRows(rows: RankedRow[], topK: number, profile: RetrievalProfil const textIndexes = new Map() for (const row of rows) { - const textKey = row.row.text.replace(/\s+/gu, " ").trim().toLowerCase() + const textKey = `${row.row.contextPath}\0${row.row.text.replace(/\s+/gu, " ").trim().toLowerCase()}` const existingIndex = textIndexes.get(textKey) if (existingIndex === undefined) { textIndexes.set(textKey, uniqueRows.length) @@ -206,13 +216,48 @@ function preferCanonicalPath(candidate: string, current: string): boolean { function hasLexicalOverlap(queryTokens: string[], text: string): boolean { const textTokens = tokenize(text) return queryTokens.some((queryToken) => - textTokens.some( - (textToken) => - queryToken === textToken || - (queryToken.length >= 4 && - textToken.length >= 4 && - sharedPrefixLength(queryToken, textToken) >= 4), - ), + textTokens.some((textToken) => tokensAreLexicallyRelated(queryToken, textToken)), + ) +} + +function tokensAreLexicallyRelated(queryToken: string, textToken: string): boolean { + if (queryToken === textToken) { + return true + } + if ( + queryToken.length >= 4 && + textToken.length >= 4 && + sharedPrefixLength(queryToken, textToken) >= 4 + ) { + return true + } + if ( + queryToken.length < MIN_FUZZY_TOKEN_LENGTH || + textToken.length < MIN_FUZZY_TOKEN_LENGTH || + Math.abs(queryToken.length - textToken.length) > 1 || + !/^[a-z0-9_-]+$/u.test(queryToken) || + !/^[a-z0-9_-]+$/u.test(textToken) + ) { + return false + } + return trigramDiceSimilarity(queryToken, textToken) >= MIN_TRIGRAM_DICE_SIMILARITY +} + +function trigramDiceSimilarity(left: string, right: string): number { + const leftTrigrams = tokenTrigrams(left) + const rightTrigrams = tokenTrigrams(right) + let shared = 0 + for (const trigram of leftTrigrams) { + if (rightTrigrams.has(trigram)) { + shared += 1 + } + } + return (2 * shared) / (leftTrigrams.size + rightTrigrams.size) +} + +function tokenTrigrams(token: string): Set { + return new Set( + Array.from({ length: token.length - 2 }, (_value, index) => token.slice(index, index + 3)), ) } @@ -257,6 +302,71 @@ export async function ask(query: string, options: SearchOptions = {}): Promise { + const requestedCitation = citation.trim() + const target = parseCitationTarget(requestedCitation) + const contextRadius = normalizeContextRadius(options.contextRadius) + const config = await loadConfig(String(options.cwd ?? process.cwd())) + const table = await openRowsTable(config) + if (!table) { + return { + requestedCitation, + found: false, + relativePath: target.relativePath, + chunkIndex: target.chunkIndex, + contextRadius, + passages: [], + } + } + + const manifest = await readIndexManifest(config) + if (!manifest) { + throw new Error( + "Index manifest is missing. Rebuild with `rgr ingest --rebuild` before expanding citations.", + ) + } + const freshnessWarning = await getIndexFreshnessWarning(config) + if (freshnessWarning) { + throw new Error(freshnessWarning) + } + + const minimumChunkIndex = Math.max(0, target.chunkIndex - contextRadius) + const maximumChunkIndex = target.chunkIndex + contextRadius + const rows = (await table + .query() + .select(SEARCH_COLUMNS) + .where( + `relativePath = ${sqlString(target.relativePath)} AND chunkIndex >= ${minimumChunkIndex} AND chunkIndex <= ${maximumChunkIndex}`, + ) + .toArray()) as SearchRow[] + const targetRow = rows.find((row) => row.chunkIndex === target.chunkIndex) + if (!targetRow) { + return { + requestedCitation, + found: false, + relativePath: target.relativePath, + chunkIndex: target.chunkIndex, + contextRadius, + passages: [], + } + } + if (citationForRow(targetRow) !== requestedCitation) { + throw new Error("Citation coordinates do not match the indexed passage.") + } + + return { + requestedCitation, + found: true, + relativePath: target.relativePath, + chunkIndex: target.chunkIndex, + contextRadius, + passages: rows.sort(compareChunkRows).map(contextChunkForRow), + } +} + function retrievalOnlyAnswer(sources: SearchResult[]): string { const snippets = sources .map((source, index) => { @@ -281,7 +391,7 @@ async function lexicalCandidateRows( const ftsQuery = lexicalQuery(query) if (ftsQuery) { try { - const searchQuery = table.search(ftsQuery, "fts", "text").select(FTS_SEARCH_COLUMNS) + const searchQuery = table.search(ftsQuery, "fts", "searchText").select(FTS_SEARCH_COLUMNS) return (await (pathPredicate ? searchQuery.where(pathPredicate) : searchQuery) .limit(limit) .toArray()) as SearchRow[] @@ -402,6 +512,7 @@ function answerText(source: SearchResult): string { function contextChunkForRow(row: SearchRow): SearchContextChunk { return { chunkIndex: row.chunkIndex, + contextPath: row.contextPath, text: row.text, charStart: nullableNumber(row.charStart), charEnd: nullableNumber(row.charEnd), @@ -518,7 +629,7 @@ function bm25Scores(queryTokens: string[], rows: SearchRow[]): Map { - const tokens = tokenize(row.text) + const tokens = tokenize(row.searchText) const frequencies = new Map() for (const token of tokens) { frequencies.set(token, (frequencies.get(token) ?? 0) + 1) @@ -565,10 +676,50 @@ function rowDistance(row: SearchRow): number { : Number.POSITIVE_INFINITY } +async function assertIndexFreshness(config: Awaited>): Promise { + const manifest = await readIndexManifest(config) + if (!manifest) { + throw new Error( + "Index manifest is missing. Rebuild with `rgr ingest --rebuild` before searching.", + ) + } + const freshnessWarning = await getIndexFreshnessWarning(config) + if (freshnessWarning) { + throw new Error(freshnessWarning) + } +} + function rowKey(row: SearchRow): string { return `${row.relativePath}\0${row.chunkIndex}` } +function parseCitationTarget(citation: string): { relativePath: string; chunkIndex: number } { + const match = /#(\d+)$/u.exec(citation) + const chunkIndexText = match?.[1] + if (!match || chunkIndexText === undefined) { + throw new Error("Citation must end with a Ragmir chunk suffix such as `#3`.") + } + + const chunkIndex = Number.parseInt(chunkIndexText, 10) + let relativePath = citation.slice(0, match.index) + relativePath = relativePath.replace(/:L\d+-L\d+$/u, "") + relativePath = relativePath.replace(/:p\d+(?:-p\d+)?$/u, "") + if (!relativePath) { + throw new Error("Citation must include a source path before its chunk suffix.") + } + return { relativePath, chunkIndex } +} + +function normalizeContextRadius(contextRadius: number | undefined): number { + if (contextRadius === undefined) { + return 0 + } + if (!Number.isInteger(contextRadius) || contextRadius < 0) { + throw new Error("contextRadius must be a non-negative integer.") + } + return Math.min(contextRadius, MAX_CONTEXT_RADIUS) +} + function citationForRow(row: SearchRow): string { const lineStart = nullableNumber(row.lineStart) const lineEnd = nullableNumber(row.lineEnd) diff --git a/packages/ragmir-core/src/research.ts b/packages/ragmir-core/src/research.ts index a724248..4162395 100644 --- a/packages/ragmir-core/src/research.ts +++ b/packages/ragmir-core/src/research.ts @@ -169,6 +169,7 @@ export function compactSearchResults( source: result.source, relativePath: result.relativePath, chunkIndex: result.chunkIndex, + contextPath: result.contextPath, citation: result.citation, snippet: compactText(result.text, maxLength), distance: result.distance, @@ -188,6 +189,7 @@ export function compactResearchReport(report: ResearchReport): Omit { expect(report.providers.embedding).toBe("transformers") expect(report.providers.transformersAllowRemoteModels).toBe(true) + expect(report.mcp.maxOutputBytes).toBe(32_768) expect(report.warnings).toContain( "Transformers remote model loading is enabled; model files can be downloaded from Hugging Face.", ) diff --git a/packages/ragmir-core/src/security.ts b/packages/ragmir-core/src/security.ts index bef8170..4cbf066 100644 --- a/packages/ragmir-core/src/security.ts +++ b/packages/ragmir-core/src/security.ts @@ -122,6 +122,7 @@ export async function securityAudit(cwd = process.cwd()): Promise { expect(rows[0]?.vector[1]).toBeCloseTo(0.2) expect(rows[0]?.vector[2]).toBeCloseTo(0.3) expect(rows[0]?.embeddingProvider).toBe("local-hash") + expect(rows[0]?.contextPath).toBe("Evidence") + expect(rows[0]?.searchText).toBe("Evidence\ncontent 0") }) it("drops the table when writing zero rows", async () => { diff --git a/packages/ragmir-core/src/store.ts b/packages/ragmir-core/src/store.ts index 01f972d..964bf2b 100644 --- a/packages/ragmir-core/src/store.ts +++ b/packages/ragmir-core/src/store.ts @@ -75,13 +75,13 @@ interface EnsureVectorIndexResult { async function ensureLexicalIndex(table: lancedb.Table): Promise { const existing = await table.listIndices() - const hasTextIndex = existing.some((index) => index.name === "text_idx") + const hasTextIndex = existing.some((index) => index.name === "searchText_idx") if (hasTextIndex) { return { created: false, warning: null } } try { - await table.createIndex("text", { + await table.createIndex("searchText", { config: lancedb.Index.fts({ asciiFolding: true, lowercase: true }), }) return { created: true, warning: null } diff --git a/packages/ragmir-core/src/test-support/config.ts b/packages/ragmir-core/src/test-support/config.ts index f63f949..cc5fef9 100644 --- a/packages/ragmir-core/src/test-support/config.ts +++ b/packages/ragmir-core/src/test-support/config.ts @@ -35,6 +35,7 @@ export function testConfig( }, accessLog: DEFAULT_CONFIG.accessLog, mcpMaxTopK: DEFAULT_CONFIG.mcpMaxTopK, + mcpMaxOutputBytes: DEFAULT_CONFIG.mcpMaxOutputBytes, topK: DEFAULT_CONFIG.topK, chunkSize: DEFAULT_CONFIG.chunkSize, chunkOverlap: DEFAULT_CONFIG.chunkOverlap, diff --git a/packages/ragmir-core/src/types.ts b/packages/ragmir-core/src/types.ts index 9b2e7b2..0be7680 100644 --- a/packages/ragmir-core/src/types.ts +++ b/packages/ragmir-core/src/types.ts @@ -20,6 +20,7 @@ export interface Config { redaction: RedactionConfig accessLog: boolean mcpMaxTopK: number + mcpMaxOutputBytes: number topK: number chunkSize: number chunkOverlap: number @@ -62,9 +63,22 @@ export interface AccessLogUsageReport { uniqueQueryHashes: number averageResultCount: number | null averageResultCountByAction: Record + mcpOutput: McpOutputUsageReport lastEventAt: string | null } +export type McpOutputTool = "ragmir_search" | "ragmir_ask" | "ragmir_research" | "ragmir_expand" + +export interface McpOutputUsageReport { + responses: number + retrievedBytes: number + returnedBytes: number + savedBytes: number + reductionRatio: number | null + compactedResponses: number + truncatedResponses: number +} + export interface IngestionLimitsReport { maxFileBytes: number maxFiles: null @@ -174,6 +188,7 @@ export interface TextChunk { source: string relativePath: string chunkIndex: number + contextPath: string text: string charStart: number charEnd: number @@ -187,6 +202,7 @@ export interface TextChunk { } export interface VectorRow extends TextChunk { + searchText: string vector: number[] embeddingProvider: EmbeddingProvider embeddingModel: string @@ -229,6 +245,7 @@ export interface SearchOptions { export interface SearchContextChunk { chunkIndex: number + contextPath: string text: string charStart: number | null charEnd: number | null @@ -243,6 +260,7 @@ export interface SearchResult { source: string relativePath: string chunkIndex: number + contextPath: string citation: string text: string distance: number | null @@ -255,10 +273,25 @@ export interface SearchResult { context: SearchContextChunk[] } +export interface ExpandCitationOptions { + cwd?: PathLike + contextRadius?: number +} + +export interface ExpandedCitation { + requestedCitation: string + found: boolean + relativePath: string + chunkIndex: number + contextRadius: number + passages: SearchContextChunk[] +} + export interface CompactSearchResult { source: string relativePath: string chunkIndex: number + contextPath: string citation: string snippet: string distance: number | null @@ -296,6 +329,7 @@ export interface ResearchEvidence { source: string relativePath: string chunkIndex: number + contextPath: string citation: string text: string distance: number | null @@ -502,6 +536,7 @@ export interface SecurityAuditReport { } mcp: { maxTopK: number + maxOutputBytes: number destructiveToolsExposed: false } gitignore: { diff --git a/scripts/smoke.mjs b/scripts/smoke.mjs index 708b934..c2178c6 100644 --- a/scripts/smoke.mjs +++ b/scripts/smoke.mjs @@ -202,6 +202,11 @@ try { if (typeof statusJson.chunksIndexed !== "number" || statusJson.chunksIndexed <= 0) { throw new Error(`status --json should expose chunksIndexed, got ${statusJson.chunksIndexed}`) } + if (statusJson.mcpMaxOutputBytes !== 32_768) { + throw new Error( + `status --json should expose mcpMaxOutputBytes, got ${statusJson.mcpMaxOutputBytes}`, + ) + } const limitsJson = parseJson((await runKb(["limits", "--json"], tempRoot)).stdout, "limits JSON") if ( @@ -821,6 +826,7 @@ async function smokeMcp(cwd) { assertIncludes(JSON.stringify(tools), "ragmir_search", "MCP should expose ragmir_search") assertIncludes(JSON.stringify(tools), "ragmir_ask", "MCP should expose ragmir_ask") assertIncludes(JSON.stringify(tools), "ragmir_research", "MCP should expose ragmir_research") + assertIncludes(JSON.stringify(tools), "ragmir_expand", "MCP should expose ragmir_expand") assertIncludes(JSON.stringify(tools), "ragmir_audit", "MCP should expose ragmir_audit") assertIncludes(JSON.stringify(tools), "ragmir_evaluate", "MCP should expose ragmir_evaluate") assertIncludes( @@ -864,6 +870,39 @@ async function smokeMcp(cwd) { throw new Error(`MCP search should retrieve line-aware context chunks: ${mcpText(search)}`) } + const boundedSearch = await client.request("tools/call", { + name: "ragmir_search", + arguments: { + query: "French tax residency", + topK: 3, + compact: true, + maxBytes: 1_024, + }, + }) + const boundedSearchJson = parseJson(mcpText(boundedSearch), "bounded MCP search JSON") + const outputMeta = boundedSearch.result?._meta?.["ragmir/output"] + if ( + !Array.isArray(boundedSearchJson) || + outputMeta?.budgetBytes !== 1_024 || + typeof outputMeta.returnedBytes !== "number" || + outputMeta.returnedBytes > 1_024 + ) { + throw new Error(`MCP search should enforce its byte budget: ${JSON.stringify(boundedSearch)}`) + } + + const expanded = await client.request("tools/call", { + name: "ragmir_expand", + arguments: { citation: searchJson[0].citation, contextRadius: 1, maxBytes: 1_024 }, + }) + const expandedJson = parseJson(mcpText(expanded), "MCP citation expansion JSON") + if ( + expandedJson.found !== true || + !Array.isArray(expandedJson.passages) || + !expandedJson.passages.some((passage) => passage.citation === searchJson[0].citation) + ) { + throw new Error(`MCP should expand an exact returned citation: ${mcpText(expanded)}`) + } + const ask = await client.request("tools/call", { name: "ragmir_ask", arguments: { query: "What proves the French tax residency risk?", topK: 1, contextRadius: 1 }, @@ -933,6 +972,9 @@ async function smokeMcp(cwd) { if (usageJson.totalEvents < 1) { throw new Error(`MCP usage report should summarize local usage: ${mcpText(usage)}`) } + if (usageJson.mcpOutput?.responses < 1 || usageJson.mcpOutput.returnedBytes < 1) { + throw new Error(`MCP usage report should include output metrics: ${mcpText(usage)}`) + } const security = await client.request("tools/call", { name: "ragmir_security_audit",