diff --git a/backend/src/docmind/rag/service.py b/backend/src/docmind/rag/service.py index 9546297..50e8651 100644 --- a/backend/src/docmind/rag/service.py +++ b/backend/src/docmind/rag/service.py @@ -13,7 +13,7 @@ from docmind.text import best_sentence TitleLookup = Callable[[str, Sequence[str]], Awaitable[dict[str, str]]] -_CITE = re.compile(r"\[(\d{1,2})\]") +_CITE = re.compile(r"\[(\d{1,2}(?:\s*,\s*\d{1,2})*)\]") @dataclass(slots=True, frozen=True) @@ -40,7 +40,7 @@ class DoneEvent: def cited_indices(answer: str, available: int) -> list[int]: - found = {int(m) for m in _CITE.findall(answer)} + found = {int(num) for group in _CITE.findall(answer) for num in group.split(",")} return sorted(i for i in found if 1 <= i <= available) diff --git a/backend/tests/test_citations.py b/backend/tests/test_citations.py new file mode 100644 index 0000000..0dbf193 --- /dev/null +++ b/backend/tests/test_citations.py @@ -0,0 +1,29 @@ +"""Unit tests for parsing citation markers out of an answer.""" + +from __future__ import annotations + +from docmind.rag.service import cited_indices + + +def test_separate_markers() -> None: + assert cited_indices("See [1] and also [2].", 2) == [1, 2] + + +def test_repeated_bracket_markers() -> None: + assert cited_indices("Both [1][2] agree.", 2) == [1, 2] + + +def test_grouped_marker() -> None: + assert cited_indices("Both sources [1, 2] agree.", 2) == [1, 2] + + +def test_grouped_marker_without_spaces() -> None: + assert cited_indices("See [1,2,3].", 3) == [1, 2, 3] + + +def test_grouped_marker_drops_out_of_range() -> None: + assert cited_indices("See [1, 5].", 2) == [1] + + +def test_zero_and_out_of_range_ignored() -> None: + assert cited_indices("Nothing here: [0], [7].", 2) == [] diff --git a/frontend/src/components/AnswerText.test.tsx b/frontend/src/components/AnswerText.test.tsx index cdc5b10..615bb74 100644 --- a/frontend/src/components/AnswerText.test.tsx +++ b/frontend/src/components/AnswerText.test.tsx @@ -19,6 +19,24 @@ describe('AnswerText', () => { expect(screen.getByText(/See \[0\] and/)).toBeInTheDocument() }) + it('turns a grouped marker into one button per source', () => { + const onCite = vi.fn() + render() + expect(citeButtons().map((b) => b.textContent)).toEqual(['1', '2']) + }) + + it('keeps out-of-range numbers in a group as text', () => { + render() + expect(citeButtons().map((b) => b.textContent)).toEqual(['1']) + expect(screen.getByText(/5/)).toBeInTheDocument() + }) + + it('keeps a fully out-of-range group as literal text', () => { + render() + expect(citeButtons()).toHaveLength(0) + expect(screen.getByText(/See \[0, 9\]\./)).toBeInTheDocument() + }) + it('reports the clicked source number', async () => { const onCite = vi.fn() render() diff --git a/frontend/src/components/AnswerText.tsx b/frontend/src/components/AnswerText.tsx index 2c825b8..d6cc266 100644 --- a/frontend/src/components/AnswerText.tsx +++ b/frontend/src/components/AnswerText.tsx @@ -6,7 +6,7 @@ interface Props { available?: number } -const INLINE = /(\[\d{1,2}\])|(\*\*[^*]+\*\*)|(`[^`]+`)|(_[^_\n]+_)/g +const INLINE = /(\[\d{1,2}(?:\s*,\s*\d{1,2})*\])|(\*\*[^*]+\*\*)|(`[^`]+`)|(_[^_\n]+_)/g /** * Tiny, safe Markdown subset renderer: paragraphs, lists, bold/italic/code and [n] citations. @@ -53,22 +53,31 @@ function renderInline(text: string, onCite: Props['onCite'], available: number): if (start > last) out.push(text.slice(last, start)) const token = match[0] if (match[1]) { - const n = Number(token.slice(1, -1)) - out.push( - n >= 1 && n <= available ? ( - - ) : ( - token - ), - ) + // A citation can list several sources in one pair of brackets, e.g. [1, 2]. + const nums = token.slice(1, -1).split(',').map((part) => Number(part.trim())) + const inRange = (n: number) => n >= 1 && n <= available + if (!nums.some(inRange)) { + out.push(token) + } else { + nums.forEach((n, i) => { + if (i > 0) out.push(', ') + out.push( + inRange(n) ? ( + + ) : ( + String(n) + ), + ) + }) + } } else if (match[2]) { out.push({token.slice(2, -2)}) } else if (match[3]) {