Citation tagging over generated text

A legal assistant writes an answer; for each source document we have phrases that were quoted or closely used. Mark which parts of the answer came from which sources, so a lawyer can click through and check. Input:

• text: the generated answer.
• sources: a map from source id to a list of phrases from that source.

Output: the text with each supported region wrapped as <cite ids="S2,S1">…</cite>. Ask:

• Matching rule. Exact substring, case-insensitive? Or only at word boundaries, so the phrase "tax" does not light up inside "syntax"? Support both; word level is the safer default for legal text.
• Overlaps. When matches overlap or touch, merge them into one tagged block; otherwise the reader gets nested or fragmented tags.
• Order of ids in a block. By how many matches each source has in the block, most first; ties by id, so output is deterministic.
• Escaping. If the output is HTML, escape the text and the ids. Here: plain text in, tagged text out.

1. Find all matches. For each phrase, find every occurrence (including overlapping ones) in the lowercased text. With word-level matching, keep an occurrence only if the characters just before and after it are not letters or digits. Each match is (start, end, source_id).
2. Merge intervals. Sort matches by start. Sweep: a match that starts at or before the current block's end (touching counts) extends the block; otherwise close the block and start a new one. Each block keeps a counter of source ids.
3. Emit. Copy text between blocks unchanged and wrap each block in a tag whose ids are sorted by (-count, id).

Cost: the simple search is O(total phrase occurrences scanned), roughly O(n × P) for text length n and P phrases using str.find repeatedly. That is fine for an answer of a few thousand characters and a few hundred phrases. For thousands of phrases, build an Aho-Corasick automaton: one pass over the text finds every occurrence of every phrase in O(n + total phrase length + matches).

from collections import Counter
from typing import Dict, List, Tuple

Match = Tuple[int, int, str]


def _is_word_char(ch: str) -> bool:
    return ch.isalnum() or ch == "_"


def find_matches(text: str, sources: Dict[str, List[str]], word_level: bool = True) -> List[Match]:
    lower = text.lower()
    matches: List[Match] = []
    for sid, phrases in sources.items():
        for phrase in phrases:
            p = phrase.lower().strip()
            if not p:
                continue
            start = lower.find(p)
            while start != -1:
                end = start + len(p)
                ok = True
                if word_level:
                    before = lower[start - 1] if start > 0 else " "
                    after = lower[end] if end < len(lower) else " "
                    ok = not (_is_word_char(p[0]) and _is_word_char(before)) and \
                        not (_is_word_char(p[-1]) and _is_word_char(after))
                if ok:
                    matches.append((start, end, sid))
                start = lower.find(p, start + 1)     # overlapping occurrences too
    return matches


def merge_blocks(matches: List[Match]) -> List[Tuple[int, int, Counter]]:
    blocks: List[Tuple[int, int, Counter]] = []
    for start, end, sid in sorted(matches):
        if blocks and start <= blocks[-1][1]:            # overlaps or touches
            b_start, b_end, counts = blocks[-1]
            counts[sid] += 1
            blocks[-1] = (b_start, max(b_end, end), counts)
        else:
            blocks.append((start, end, Counter({sid: 1})))
    return blocks


def tag_citations(text: str, sources: Dict[str, List[str]], word_level: bool = True) -> str:
    out: List[str] = []
    pos = 0
    for start, end, counts in merge_blocks(find_matches(text, sources, word_level)):
        ids = ",".join(sid for sid, _ in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0])))
        out.append(text[pos:start])
        out.append(f'<cite ids="{ids}">{text[start:end]}</cite>')
        pos = end
    out.append(text[pos:])
    return "".join(out)


def coverage(text: str, sources: Dict[str, List[str]]) -> float:
    """Share of the answer's characters inside some citation: a groundedness signal."""
    covered = sum(e - s for s, e, _ in merge_blocks(find_matches(text, sources)))
    return covered / len(text) if text else 0.0

Matching is done on text.lower(), whose length equals the original's for ordinary text. A few Unicode characters change length when lowercased (for example "İ"); for those, match on a case-folded copy with an index map back to the original.

• Matching: O(P × n) in the worst case with repeated find (P phrases, text length n), plus the matches found.
• Merging: O(M log M) for M matches (the sort), then linear.
• Output: O(n + M).

1. Many phrases. Aho-Corasick: build a trie of all phrases with failure links, then scan the text once. Each position reports every phrase ending there. This is the right tool when every source contributes hundreds of sentences.
2. Fuzzy support. Models paraphrase. Exact matching misses "must pay rent each month" for "rent shall be paid monthly". Use n-gram overlap or embeddings at the sentence level to find the best-supporting source passage, with exact spans as a precision bonus. Show fuzzy citations differently from verbatim quotes.
3. Claims without support. Report the sentences with no citation at all. For legal work, an uncited factual claim is the thing a reviewer most needs to see; a low coverage score can block or flag an answer.
4. Markup safety. If the answer is rendered as HTML, escape the text before tagging and build tags with escaped attribute values, or emit structured spans ([{start, end, ids}]) and let the client render them. Never let model output inject markup.
5. Streaming. Tag as the answer streams: run matching on each completed sentence and hold back the last partial sentence until it ends.

Mid-level (L4)

• Finds all occurrences, merges overlapping intervals and wraps them correctly.
• Handles no matches, case differences and adjacent matches.
• Writes tests with overlapping phrases from two sources.

Senior (L5)

• Word-boundary matching, overlapping occurrences of one phrase, and deterministic id ordering by frequency.
• States complexity and knows when to switch to Aho-Corasick.
• Thinks about output safety (escaping or structured spans).

Staff (L6+)

• Moves from exact spans to support detection with paraphrase, and flags uncited claims as the core product signal.
• Designs for streaming answers and large source sets.
• Explains how citations fit the lawyer's review workflow and what precision and recall trade-off to pick.