app/rag/models.py

"""Typed contracts for structured academic evidence and adaptive context.

These models are intentionally storage- and agent-neutral. SQLite is the
canonical evidence store, Chroma indexes paper evidence and project memory,
and Cognee supplies cross-project student memory. Consumers receive the three
channels separately through :class:`EvidenceBundle`.
"""

from __future__ import annotations

import hashlib
import json
import re
import unicodedata
from datetime import datetime, timezone
from enum import StrEnum
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator


class EvidenceType(StrEnum):
    TITLE = "title"
    HEADING = "heading"
    PARAGRAPH = "paragraph"
    LIST = "list"
    TABLE = "table"
    FIGURE = "figure"
    PLOT = "plot"
    DIAGRAM = "diagram"
    FORMULA = "formula"
    CAPTION = "caption"
    FOOTNOTE = "footnote"
    REFERENCE = "reference"
    ANNOTATION = "annotation"
    SECTION_CARD = "section_card"
    LOCAL_WINDOW = "local_window"


class SourceKind(StrEnum):
    PAPER = "paper"
    STUDENT_ANNOTATION = "student_annotation"
    DERIVED_SUMMARY = "derived_summary"


class RelationType(StrEnum):
    PARENT = "parent"
    NEXT = "next"
    PREVIOUS = "previous"
    CAPTION_OF = "caption_of"
    DESCRIBES = "describes"
    CITES = "cites"
    ANNOTATES = "annotates"
    SAME_TABLE = "same_table"
    CONTINUATION = "continuation"
    SUPPORTS = "supports"


class ConsumerType(StrEnum):
    CHAT = "chat"
    WIKI = "wiki"
    PAIR_BUDDY = "pair_buddy"
    TUTOR = "tutor"
    QUIZ = "quiz"
    FLASHCARDS = "flashcards"
    VISUALIZATION = "visualization"
    DRAFT = "draft"
    REPORT = "report"
    PAPER_GRAPH = "paper_graph"
    PROJECT_GRAPH = "project_graph"
    CITATION_GRAPH = "citation_graph"
    LIBRARY = "library"


class BoundingBox(BaseModel):
    """Normalized PDF-space rectangle in the closed interval ``[0, 1]``."""

    model_config = ConfigDict(frozen=True)

    x: float = Field(ge=0.0, le=1.0)
    y: float = Field(ge=0.0, le=1.0)
    w: float = Field(gt=0.0, le=1.0)
    h: float = Field(gt=0.0, le=1.0)

    @model_validator(mode="after")
    def validate_extent(self) -> "BoundingBox":
        tolerance = 1e-6
        if self.x + self.w > 1.0 + tolerance or self.y + self.h > 1.0 + tolerance:
            raise ValueError("normalized bounding box extends beyond the page")
        return self

    def rounded(self, digits: int = 6) -> dict[str, float]:
        return {name: round(float(getattr(self, name)), digits) for name in ("x", "y", "w", "h")}

    def intersection_ratio(self, other: "BoundingBox") -> float:
        left = max(self.x, other.x)
        top = max(self.y, other.y)
        right = min(self.x + self.w, other.x + other.w)
        bottom = min(self.y + self.h, other.y + other.h)
        if right <= left or bottom <= top:
            return 0.0
        intersection = (right - left) * (bottom - top)
        smaller_area = min(self.w * self.h, other.w * other.h)
        return intersection / smaller_area if smaller_area else 0.0


class SelectionAnchor(BaseModel):
    document_id: str
    page_number: int = Field(ge=1)
    boxes: list[BoundingBox] = Field(default_factory=list)
    evidence_id: str | None = None
    region_id: str | None = None
    text: str = ""


class AcademicDocument(BaseModel):
    project_id: str
    document_id: str
    filename: str
    title: str = ""
    authors: list[str] = Field(default_factory=list)
    abstract: str = ""
    page_count: int = Field(ge=0)
    parse_quality: Literal["good", "degraded", "empty"] = "good"
    quality_flags: list[str] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))


class EvidenceUnit(BaseModel):
    evidence_id: str
    project_id: str
    document_id: str
    element_type: EvidenceType
    page_start: int = Field(ge=1)
    page_end: int = Field(ge=1)
    parent_id: str | None = None
    section_path: list[str] = Field(default_factory=list)
    ordinal: int = Field(default=0, ge=0)
    bbox_norm: BoundingBox | None = None
    raw_text: str = ""
    retrieval_text: str = ""
    caption: str = ""
    table_markdown: str = ""
    visual_description: str = ""
    quality_flags: list[str] = Field(default_factory=list)
    annotation_ids: list[str] = Field(default_factory=list)
    source_kind: SourceKind = SourceKind.PAPER
    metadata: dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="after")
    def validate_pages_and_content(self) -> "EvidenceUnit":
        if self.page_end < self.page_start:
            raise ValueError("page_end must be greater than or equal to page_start")
        if not any((self.raw_text, self.caption, self.table_markdown, self.visual_description)):
            self.quality_flags.append("content_empty")
        return self

    @property
    def index_text(self) -> str:
        return self.retrieval_text or self.raw_text or self.table_markdown or self.visual_description or self.caption


class EvidenceRelation(BaseModel):
    source_evidence_id: str
    target_evidence_id: str
    relation_type: RelationType
    confidence: float = Field(default=1.0, ge=0.0, le=1.0)
    derivation: Literal["parser", "deterministic_linker", "vision_model", "user"]


class ParsedAcademicDocument(BaseModel):
    document: AcademicDocument
    units: list[EvidenceUnit] = Field(default_factory=list)
    relations: list[EvidenceRelation] = Field(default_factory=list)

    def indexable_units(self) -> list[EvidenceUnit]:
        return [unit for unit in self.units if unit.index_text.strip()]


ProjectMemoryPolicy = Literal["none", "relevant"]
StudentMemoryPolicy = Literal["none", "cached", "live"]


class EvidenceRequest(BaseModel):
    query: str
    consumer: ConsumerType
    project_id: str
    document_ids: list[str] = Field(default_factory=list)
    anchor_evidence_ids: list[str] = Field(default_factory=list)
    selection_anchors: list[SelectionAnchor] = Field(default_factory=list)
    graph_node_id: str | None = None
    modalities: set[EvidenceType] = Field(default_factory=set)
    token_budget: int = Field(default=6000, ge=256, le=100_000)
    project_memory_policy: ProjectMemoryPolicy = "relevant"
    student_memory_policy: StudentMemoryPolicy = "cached"


class RetrievedEvidence(BaseModel):
    evidence: EvidenceUnit
    fused_score: float = 0.0
    dense_score: float | None = None
    lexical_score: float | None = None
    structural_score: float = 0.0
    rerank_score: float | None = None
    retrieval_reasons: list[str] = Field(default_factory=list)


class MemoryContextItem(BaseModel):
    memory_id: str
    source: Literal["project_memory", "student_memory"]
    statement: str
    project_id: str | None = None
    kind: str
    score: float | None = None
    evidence_ids: list[str] = Field(default_factory=list)
    observed_at: datetime | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class EvidenceCitation(BaseModel):
    evidence_id: str
    document_id: str
    filename: str = ""
    page_start: int = Field(ge=1)
    page_end: int = Field(ge=1)
    bbox_norm: BoundingBox | None = None
    section_path: list[str] = Field(default_factory=list)


class RetrievalDiagnostics(BaseModel):
    terminal_state: Literal["success", "success_empty", "degraded", "error_fallback"] = "success"
    anchor_count: int = 0
    lexical_candidate_count: int = 0
    dense_candidate_count: int = 0
    fused_candidate_count: int = 0
    duplicates_removed: int = 0
    documents_represented: int = 0
    sections_represented: int = 0
    rerank_used: bool = False
    rerank_reason: str = ""
    context_truncated: bool = False
    source_context_chars: int = 0
    project_memory_chars: int = 0
    student_memory_chars: int = 0
    stage_ms: dict[str, float] = Field(default_factory=dict)
    warnings: list[str] = Field(default_factory=list)


class EvidenceBundle(BaseModel):
    source_evidence: list[RetrievedEvidence] = Field(default_factory=list)
    project_memory: list[MemoryContextItem] = Field(default_factory=list)
    student_memory: list[MemoryContextItem] = Field(default_factory=list)
    citations: list[EvidenceCitation] = Field(default_factory=list)
    diagnostics: RetrievalDiagnostics = Field(default_factory=RetrievalDiagnostics)


_WHITESPACE = re.compile(r"\s+")


def normalize_evidence_text(text: str) -> str:
    normalized = unicodedata.normalize("NFKC", text or "")
    return _WHITESPACE.sub(" ", normalized).strip()


def make_evidence_id(
    project_id: str,
    document_id: str,
    page_number: int,
    bbox: BoundingBox | None,
    element_type: EvidenceType | str,
    content: str,
) -> str:
    """Return a deterministic ID that changes when source provenance changes."""

    payload = {
        "project_id": project_id,
        "document_id": document_id,
        "page_number": int(page_number),
        "bbox": bbox.rounded() if bbox else None,
        "element_type": str(element_type),
        "content": normalize_evidence_text(content),
    }
    digest = hashlib.sha256(
        json.dumps(payload, sort_keys=True, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    ).hexdigest()
    return f"ev_{digest}"