pipeline/attribution.py

"""
DocMind — Attribution Parser

Extracts [SOURCE: chunk_id] tags from LLM responses and strips
sentences that cite invalid or missing chunks.
"""

import logging
import re
from dataclasses import dataclass
from typing import List, Optional, Set

logger = logging.getLogger(__name__)


@dataclass
class AttributedSentence:
    """A single sentence with its cited source chunk ID."""
    text: str                        # The clean sentence text (without the tag)
    chunk_id: Optional[str] = None   # Extracted chunk_id, or None if no valid citation
    raw_text: str = ""               # Original text including the [SOURCE: ...] tag


# Pattern to match [SOURCE: some_chunk_id]
_SOURCE_PATTERN = re.compile(r"\[SOURCE:\s*([^\]]+)\]", re.IGNORECASE)


def parse_attributed_response(response_text: str) -> List[AttributedSentence]:
    """
    Parse an LLM response into individual sentences with their source attributions.

    Each sentence is expected to end with [SOURCE: chunk_id].
    Sentences without attributions are still returned but with chunk_id=None.

    Args:
        response_text: Raw LLM response text.

    Returns:
        List of AttributedSentence objects.
    """
    if not response_text or response_text.strip() == "INSUFFICIENT_CONTEXT":
        return []

    sentences: List[AttributedSentence] = []

    # Split on sentence-ending punctuation followed by optional whitespace
    # We need to be careful: the [SOURCE: ...] tag comes AFTER the period
    # Strategy: split on [SOURCE: ...] boundaries
    parts = _SOURCE_PATTERN.split(response_text)

    # parts alternates between text segments and captured chunk_ids:
    # [text_before_first_tag, chunk_id_1, text_between, chunk_id_2, ...]
    i = 0
    while i < len(parts):
        text_part = parts[i].strip()
        chunk_id = parts[i + 1].strip() if i + 1 < len(parts) else None

        if text_part:
            # Further split if multiple sentences exist in one text segment
            # (shouldn't happen with well-formatted LLM output, but be safe)
            sub_sentences = _split_into_sentences(text_part)
            for j, sent in enumerate(sub_sentences):
                sent = sent.strip()
                if not sent:
                    continue
                # Only the last sub-sentence gets the chunk_id
                cid = chunk_id if j == len(sub_sentences) - 1 else None
                raw = f"{sent} [SOURCE: {cid}]" if cid else sent
                sentences.append(AttributedSentence(
                    text=sent,
                    chunk_id=cid,
                    raw_text=raw,
                ))

        i += 2  # Skip text + chunk_id pair

    logger.info(
        "Parsed %d attributed sentences (%d with valid sources)",
        len(sentences),
        sum(1 for s in sentences if s.chunk_id),
    )
    return sentences


def _split_into_sentences(text: str) -> List[str]:
    """Basic sentence splitter on . ! ? boundaries."""
    # Use a regex that splits on period/exclamation/question followed by space or EOL
    raw_parts = re.split(r"(?<=[.!?])\s+", text)
    return [p for p in raw_parts if p.strip()]


def strip_unattributed(
    sentences: List[AttributedSentence],
    valid_chunk_ids: Set[str],
) -> List[AttributedSentence]:
    """
    Remove sentences whose cited chunk_id is not in the valid set.

    Args:
        sentences: Parsed attributed sentences.
        valid_chunk_ids: Set of chunk IDs that actually exist in the current index.

    Returns:
        Filtered list with only validly-attributed sentences.
    """
    kept = []
    removed_count = 0

    for sent in sentences:
        if sent.chunk_id and sent.chunk_id in valid_chunk_ids:
            kept.append(sent)
        elif sent.chunk_id is None:
            # Sentence without any citation — remove it
            removed_count += 1
        else:
            # Sentence cites a chunk_id that doesn't exist
            removed_count += 1

    if removed_count > 0:
        logger.info("Stripped %d unattributed/invalid sentences", removed_count)

    return kept