Models

"""Typed data structures shared by indexing, retrieval, generation, and evaluation."""

from dataclasses import dataclass, field
from typing import Any, Dict, Mapping, Sequence, Tuple


@dataclass(frozen=True)
class Document:
    """A source document before chunking.

    ``source`` should be a resolvable URL or stable local identifier.  Metadata
    is intentionally open-ended; the reference graph retriever reads an
    optional ``entities`` sequence from it.
    """

    id: str
    text: str
    title: str = ""
    source: str = ""
    date: str = ""
    metadata: Mapping[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.id.strip():
            raise ValueError("document id must be non-empty")
        if not self.text.strip():
            raise ValueError("document text must be non-empty")


@dataclass(frozen=True)
class Chunk:
    """A token-window view of a document."""

    id: str
    document_id: str
    text: str
    start_token: int
    end_token: int
    title: str = ""
    source: str = ""
    metadata: Mapping[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class SearchResult:
    """A ranked chunk and transparent component scores."""

    chunk: Chunk
    score: float
    rank: int = 0
    retriever: str = ""
    component_scores: Mapping[str, float] = field(default_factory=dict)


@dataclass(frozen=True)
class Citation:
    """Evidence actually used in an answer."""

    chunk_id: str
    document_id: str
    title: str
    source: str
    quote: str


@dataclass(frozen=True)
class TraceEvent:
    """One observable pipeline decision."""

    stage: str
    detail: str
    values: Mapping[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class Answer:
    """A generated answer, its evidence, and its execution trace."""

    text: str
    citations: Tuple[Citation, ...]
    contexts: Tuple[SearchResult, ...]
    trace: Tuple[TraceEvent, ...]
    confidence: float
    abstained: bool = False


@dataclass(frozen=True)
class QuestionExample:
    """A retrieval/generation evaluation case."""

    id: str
    question: str
    relevant_document_ids: Tuple[str, ...]
    reference_answer: str
    tags: Tuple[str, ...] = ()


MetricRow = Dict[str, float]
MetricTable = Sequence[Mapping[str, Any]]