diff --git a/src/__tests__/citations.test.ts b/src/__tests__/citations.test.ts new file mode 100644 index 0000000..58e5f42 --- /dev/null +++ b/src/__tests__/citations.test.ts @@ -0,0 +1,235 @@ +import { describe, expect, it } from "vitest"; +import { + citedIndices, + collectSources, + createSourceList, + linkCitations, + renderSources, +} from "../lib/citations"; +import type { ToolCallRecord } from "../lib/types"; + +/** One record, with only the fields the source helpers read. */ +function call(sources?: ToolCallRecord["sources"]): ToolCallRecord { + return { name: "web_search", arguments: {}, result: "", isError: false, sources }; +} + +describe("createSourceList", () => { + it("numbers sources from one in the order they were found", () => { + const list = createSourceList(); + expect(list.add("First", "https://first.example").index).toBe(1); + expect(list.add("Second", "https://second.example").index).toBe(2); + }); + + it("gives a URL the number it already has, so a re-fetch cites the same page", () => { + const list = createSourceList(); + list.add("First", "https://first.example"); + const again = list.add("First again", "https://first.example"); + + expect(again.index).toBe(1); + // Nothing new was added, which is how the caller knows not to store it twice. + expect(again.source).toBeUndefined(); + }); + + it("returns the entry only the first time, so a turn stores it once", () => { + const list = createSourceList(); + expect(list.add("First", "https://first.example").source).toEqual({ + index: 1, + title: "First", + url: "https://first.example", + }); + }); + + it("ignores surrounding whitespace when matching a URL", () => { + const list = createSourceList(); + list.add("First", "https://first.example"); + expect(list.add("First", " https://first.example ").index).toBe(1); + }); + + it("falls back to the URL as the title when there is none", () => { + const list = createSourceList(); + expect(list.add(" ", "https://first.example").source?.title).toBe("https://first.example"); + }); + + it("gives a result with no usable URL no number at all", () => { + // Zero rather than the next number: there is nothing to link to, so the tool + // prints no marker and the model cannot cite a page that has no address. + const list = createSourceList(); + expect(list.add("Title", " ").index).toBe(0); + expect(list.add("Next", "https://next.example").index).toBe(1); + }); +}); + +describe("collectSources", () => { + it("flattens a turn's records into one list in number order", () => { + const collected = collectSources([ + call([{ index: 1, title: "First", url: "https://first.example" }]), + call([{ index: 2, title: "Second", url: "https://second.example" }]), + ]); + + expect(collected.map((source) => source.index)).toEqual([1, 2]); + }); + + it("keeps the first record's version of a number", () => { + // A number is assigned once, so two records carrying it describe the same + // page, and the first one seen is the one that named it. + const collected = collectSources([ + call([{ index: 1, title: "First", url: "https://first.example" }]), + call([{ index: 1, title: "Duplicate", url: "https://first.example" }]), + ]); + + expect(collected).toEqual([{ index: 1, title: "First", url: "https://first.example" }]); + }); + + it("sorts by number even when the records arrive out of order", () => { + const collected = collectSources([ + call([{ index: 3, title: "Third", url: "https://third.example" }]), + call([{ index: 1, title: "First", url: "https://first.example" }]), + ]); + + expect(collected.map((source) => source.index)).toEqual([1, 3]); + }); + + it("is empty for a turn that consulted nothing", () => { + expect(collectSources([])).toEqual([]); + expect(collectSources([call()])).toEqual([]); + }); +}); + +describe("citedIndices", () => { + const sources = [ + { index: 1, title: "First", url: "https://first.example" }, + { index: 2, title: "Second", url: "https://second.example" }, + ]; + + it("finds the numbers an answer cites, in either marker shape", () => { + expect([...citedIndices("See [1] and [^2].", sources)].sort()).toEqual([1, 2]); + }); + + it("ignores a number that is not a source", () => { + expect([...citedIndices("See [9].", sources)]).toEqual([]); + }); +}); + +describe("linkCitations", () => { + const sources = [ + { index: 1, title: "First", url: "https://first.example" }, + { index: 3, title: "Third", url: "https://third.example" }, + ]; + + it("turns a marker into a link whose label is the marker", () => { + const result = linkCitations("Rain is likely [1].", sources); + + expect(result.text).toBe("Rain is likely [\\[1\\]](https://first.example)."); + expect([...result.cited]).toEqual([1]); + }); + + it("reads the footnote shape as the same thing", () => { + expect(linkCitations("Rain [^1].", sources).text).toBe("Rain [\\[1\\]](https://first.example)."); + }); + + it("links a list of markers and keeps the comma", () => { + const result = linkCitations("Both [1, 3] agree.", sources); + + expect(result.text).toBe( + "Both [\\[1\\]](https://first.example), [\\[3\\]](https://third.example) agree.", + ); + expect([...result.cited].sort()).toEqual([1, 3]); + }); + + it("leaves a marker alone when any number in it is not a source", () => { + // Half a link reads as a citation that works, which is worse than the marker + // the model wrote: the reader cannot tell which half is missing. + const result = linkCitations("See [1, 9].", sources); + + expect(result.text).toBe("See [1, 9]."); + expect([...result.cited]).toEqual([]); + }); + + it("leaves a marker alone when the destination is not one this window should open", () => { + const unsafe = [{ index: 1, title: "Nasty", url: "javascript:alert(1)" }]; + const result = linkCitations("See [1].", unsafe); + + expect(result.text).toBe("See [1]."); + expect([...result.cited]).toEqual([]); + }); + + it("leaves a destination containing whitespace as text", () => { + const odd = [{ index: 1, title: "Odd", url: "https://first.example/a b" }]; + expect(linkCitations("See [1].", odd).text).toBe("See [1]."); + }); + + it("does not touch a link the model wrote", () => { + const result = linkCitations("See [1](https://elsewhere.example).", sources); + expect(result.text).toBe("See [1](https://elsewhere.example)."); + }); + + it("does not touch a reference definition line", () => { + expect(linkCitations("[1]: https://elsewhere.example", sources).text).toBe( + "[1]: https://elsewhere.example", + ); + }); + + it("drops a footnote definition for a source it already lists", () => { + const answer = "Rain [^1].\n\n[^1]: First, https://first.example"; + expect(linkCitations(answer, sources).text).toBe("Rain [\\[1\\]](https://first.example).\n"); + }); + + it("keeps a markdown reference definition, which backs a link the model wrote", () => { + // `[1]:` with no caret is reference syntax, not a footnote. Only the footnote + // shape is unambiguously a duplicate of an entry we render ourselves. + const answer = "Rain [^1].\n\n[1]: https://elsewhere.example"; + expect(linkCitations(answer, sources).text).toBe( + "Rain [\\[1\\]](https://first.example).\n\n[1]: https://elsewhere.example", + ); + }); + + it("keeps a definition for a number it knows nothing about", () => { + const answer = "Rain [^1].\n\n[^7]: Something else"; + expect(linkCitations(answer, sources).text).toBe( + "Rain [\\[1\\]](https://first.example).\n\n[^7]: Something else", + ); + }); + + it("leaves a marker inside a code span as the text it is", () => { + expect(linkCitations("Type `curl [1]` now.", sources).text).toBe("Type `curl [1]` now."); + }); + + it("changes nothing when there are no sources", () => { + expect(linkCitations("See [1].", []).text).toBe("See [1]."); + }); +}); + +describe("renderSources", () => { + const sources = [ + { index: 1, title: "First", url: "https://first.example" }, + { index: 3, title: "Third", url: "https://third.example" }, + ]; + + it("lists the sources the answer cited, numbered as the model numbered them", () => { + const list = renderSources(sources, new Set([3]), false); + + expect(list).toBe("_Sources_\n\n- [\\[3\\]](https://third.example) Third"); + }); + + it("lists everything when the answer cited nothing, once the turn has ended", () => { + // Nothing to filter on, so the only thing left to show is what was consulted. + expect(renderSources(sources, new Set(), false)).toBe( + "_Sources_\n\n- [\\[1\\]](https://first.example) First\n- [\\[3\\]](https://third.example) Third", + ); + }); + + it("stays away while an answer is still arriving", () => { + // Otherwise the list appears in full at the first token and then shrinks as + // citations arrive, which reads as a glitch rather than as settling. + expect(renderSources(sources, new Set(), true)).toBe(""); + }); + + it("lists nothing when there was nothing to list", () => { + expect(renderSources([], new Set(), false)).toBe(""); + }); + + it("keeps an entry whose destination cannot be linked", () => { + const unsafe = [{ index: 1, title: "Nasty", url: "javascript:alert(1)" }]; + expect(renderSources(unsafe, new Set(), false)).toBe("_Sources_\n\n- [1] Nasty"); + }); +}); diff --git a/src/lib/citations.ts b/src/lib/citations.ts new file mode 100644 index 0000000..01e5b38 --- /dev/null +++ b/src/lib/citations.ts @@ -0,0 +1,229 @@ +import { mapOutsideCodeSpans, markdownLink } from "./code"; +import type { Source, ToolCallRecord } from "./types"; + +/** + * Numbered sources for one turn, and the citation markers that point at them. + * + * Vicinae has no footnote support to lean on: it parses markdown with + * `CMARK_OPT_DEFAULT` and attaches only the table and strikethrough extensions + * (`src/server/src/ui/quick/markdown-model.cpp` upstream), so a `[^1]` written by + * a model arrives in the pane as literal text. Everything a footnote would do + * therefore has to happen here, on the way to the renderer. + * + * Imports no Vicinae module, so it unit tests with no runtime. + */ + +/** What a turn uses to hand out numbers, one per page it consults. */ +export interface SourceList { + /** + * The number for a URL, registering it the first time it is seen. + * + * `index` is 0 for a result with no usable address, which is the caller's + * signal to print no marker at all. `source` comes back only when this call + * registered something new, so a caller records each page exactly once. + */ + add(title: string, url: string): { index: number; source?: Source }; +} + +/** + * A fresh list, numbering from one. + * + * Scoped to a turn rather than a conversation: a new question starts again at + * one, which is what lets a marker mean one thing for the whole answer beside + * it, and lets the transcript list sources per answer rather than per session. + */ +export function createSourceList(): SourceList { + const byUrl = new Map(); + let next = 1; + + return { + add(title, url) { + const address = url.trim(); + if (address.length === 0) return { index: 0 }; + + const known = byUrl.get(address); + if (known) return { index: known.index }; + + const source: Source = { index: next, title: title.trim() || address, url: address }; + next += 1; + byUrl.set(address, source); + + return { index: source.index, source }; + }, + }; +} + +/** + * Every source a turn consulted, in number order. + * + * A record holds only what its own call added, so the turn's list is the + * flattening of its records. That keeps the running total from being stored once + * per later call, and it means the list can be rebuilt from a stored turn without + * re-reading a tool result string and hoping its format held. + */ +export function collectSources(calls: ToolCallRecord[]): Source[] { + const byIndex = new Map(); + + for (const call of calls) { + for (const source of call.sources ?? []) { + // First wins. A number is assigned once, so a second record carrying it is + // describing the same page, and the first one is what named it. + if (!byIndex.has(source.index)) byIndex.set(source.index, source); + } + } + + return [...byIndex.values()].sort((left, right) => left.index - right.index); +} + +/** + * The numbers an answer cites that are actually sources. + * + * A plain scan, without the fence awareness the linking pass has, because its one + * caller is the clipboard export: what a copied document needs is the pages the + * answer leaned on, and a `[1]` inside a snippet there costs a reader nothing. + */ +export function citedIndices(answer: string, sources: Source[]): Set { + const known = new Set(sources.map((source) => source.index)); + const cited = new Set(); + + for (const match of answer.matchAll(/\[\^?(\d+(?:\s*,\s*\d+)*)\]/g)) { + for (const part of match[1].split(",")) { + const index = Number(part.trim()); + if (known.has(index)) cited.add(index); + } + } + + return cited; +} + +/** What linking did to one run of prose. */ +export interface CitationResult { + text: string; + /** The numbers that became links, which is what the list filters on. */ + cited: Set; +} + +/** + * A citation marker the model wrote, in the two shapes it reaches for. + * + * The trailing guards are what keep this off markdown the model wrote on purpose: + * `](` starts a link, and `]:` is a reference definition. The digit list allows + * `[1, 3]`, which is how a sentence resting on two sources is usually written. + */ +const CITATION_MARKER = /\[\^?(\d+(?:\s*,\s*\d+)*)\](?![(:])/g; + +/** + * A footnote definition line, which is `[^1]: title, url` on a line of its own. + * + * The caret is required rather than optional, and the two shapes mean different + * things. `[1]: something` is markdown's own reference definition syntax, which + * backs a link the model wrote, so dropping one because it happens to carry a + * source number would break that link. `[^1]:` is a footnote definition, which is + * a duplicate of an entry we render ourselves. + */ +const FOOTNOTE_DEFINITION = /^\s*\[\^(\d+)\]:\s*\S/; + +/** + * A destination this window can safely be asked to open, or nothing. + * + * A search result is untrusted input, so the scheme is checked here for the same + * reason `SAFE_URL` in `html.ts` checks it: a clickable `javascript:` link would + * be a real vulnerability rather than a rendering quirk. The rest of the test is + * characters that would break the link or smuggle a second one out of it. + */ +function linkableUrl(url: string): string | undefined { + const address = url.trim(); + if (!/^https?:\/\//i.test(address)) return undefined; + if (/[\s<>\\]/.test(address)) return undefined; + + // Control characters by code point. A character class would say the same thing + // and trip `no-control-regex`, and a NUL or a DEL in a destination is never + // anything but damage, so it is refused either way. + for (const character of address) { + const code = character.codePointAt(0) ?? 0; + if (code < 0x20 || code === 0x7f) return undefined; + } + + return address; +} + +/** + * Drops a definition line for a source the turn already lists. + * + * The entry is drawn from what the tool returned, so the model's copy of the same + * thing is a duplicate line rather than information. One for a number we know + * nothing about is left alone, because then it is the only place that text exists. + */ +function stripDefinitions(text: string, known: Set): string { + return text + .split("\n") + .filter((line) => { + const definition = FOOTNOTE_DEFINITION.exec(line); + return !definition || !known.has(Number(definition[1])); + }) + .join("\n"); +} + +/** + * `[1]`, `[^1]` and `[1, 3]` as links, plus the numbers that became links. + * + * One run of prose only: fences are the caller's business, because the caller + * already walks the document in segments for the maths pass and there is no + * reason for two ideas of where code starts. + */ +export function linkCitations(text: string, sources: Source[]): CitationResult { + const cited = new Set(); + if (sources.length === 0) return { text, cited }; + + const byIndex = new Map(sources.map((source) => [source.index, source])); + const prose = stripDefinitions(text, new Set(byIndex.keys())); + + const linked = mapOutsideCodeSpans(prose, (chunk) => + chunk.replace(CITATION_MARKER, (whole: string, digits: string) => { + const wanted = digits.split(",").map((part) => Number(part.trim())); + const links: string[] = []; + + for (const index of wanted) { + const source = byIndex.get(index); + const href = source ? linkableUrl(source.url) : undefined; + // All or nothing, and nothing at all if a destination is unusable: a + // marker that is partly a link reads as a citation that works. + if (!href) return whole; + + links.push(markdownLink(`[${index}]`, href)); + } + + for (const index of wanted) cited.add(index); + return links.join(", "); + }), + ); + + return { text: linked, cited }; +} + +/** + * The sources under an answer, or nothing when there is nothing to say. + * + * The marker is the link, matching the inline marker character for character, so + * one token means one thing in both places. Numbers are the model's rather than + * the reader's: `[3]` is the third source the turn found whether or not the answer + * cited the two before it, which keeps the transcript agreeing with the tool + * output and with anything the model said about its own numbering. + */ +export function renderSources(sources: Source[], cited: Set, streaming: boolean): string { + // While an answer is arriving the cited set is still filling up, so the list + // waits rather than appearing in full and then dropping most of itself. An + // answer that cites nothing at all is the one case left to fall back on + // everything consulted, because there is then nothing to filter by. + const shown = + cited.size > 0 ? sources.filter((source) => cited.has(source.index)) : streaming ? [] : sources; + if (shown.length === 0) return ""; + + const entries = shown.map((source) => { + const href = linkableUrl(source.url); + const marker = href ? markdownLink(`[${source.index}]`, href) : `[${source.index}]`; + return `- ${marker} ${source.title}`; + }); + + return ["_Sources_", "", ...entries].join("\n"); +} diff --git a/src/lib/types.ts b/src/lib/types.ts index d5d8b0a..799a341 100644 --- a/src/lib/types.ts +++ b/src/lib/types.ts @@ -84,12 +84,39 @@ export interface PendingAttachment extends Attachment { /** Base64 of the image, with no `data:` URL prefix. Ollama wants it bare. */ base64: string; } +/** + * One numbered source, as a web tool found it. + * + * The number is the whole point: it is assigned once per turn, printed to the + * model in the tool result, and rendered as the marker the reader clicks. The + * model and the reader therefore agree on what `[3]` means without either of them + * being asked to remember a URL. + */ +export interface Source { + /** The number the model cites. One-based, and never reused for a different page. */ + index: number; + /** Page title, or the URL when the page had none. */ + title: string; + url: string; +} + /** What a tool actually returned during a turn, kept for display and debugging. */ export interface ToolCallRecord { name: string; arguments: Record; result: string; isError: boolean; + /** + * The sources this call added to the turn, if it consulted any. + * + * Only what this call found, never the turn's running total: the whole list is + * the flattening of a turn's records, so storing the total here would write + * the same page once per later call. + * + * Absent on turns stored before sources were numbered, which is honest rather + * than empty: those answers never cited a number, so there is nothing to draw. + */ + sources?: Source[]; /** * Set when the call was refused because the turn had already made it twice. *