lexquery/schema.py

from typing import Literal
from pydantic import BaseModel, ConfigDict, Field


class Strict(BaseModel):
    model_config = ConfigDict(extra="forbid")


class Page(Strict):
    number: int
    original: str
    english: str = ""
    language: str = "Unknown"
    method: str = "text"
    reviewed: bool = True
    ocr_confidence: float | None = None
    warnings: list[str] = Field(default_factory=list)


class Document(Strict):
    id: str
    name: str
    pages: list[Page]

    @property
    def ready(self):
        return bool(self.pages) and any(p.english.strip() for p in self.pages) and all(
            p.english.strip() or not p.original.strip() for p in self.pages)


class Evidence(Strict):
    id: str
    document: str
    page: int
    text: str
    original: str


class Statement(Strict):
    text: str = Field(min_length=1, max_length=350, description="One short factual English sentence answering the task.")
    source_id: str = Field(description="Copy the exact id of the source evidence item.")


class Claim(Statement):
    quote: str = Field(min_length=1)


class Draft(Strict):
    answerable: bool
    claims: list[Statement] = Field(max_length=4)


class Summary(Strict):
    claims: list[Statement] = Field(min_length=1, max_length=6)


class Fact(Strict):
    field: Literal["case_number", "court", "parties", "judges", "order_date", "outcome"]
    value: str = Field(max_length=250, description="Concise value; include party roles or order scope when needed.")
    source_id: str


class Fields(Strict):
    facts: list[Fact] = Field(max_length=12)


class Check(Strict):
    supported: list[bool] = Field(description="One true/false decision per claim, in the same order.")
    complete: bool
    reason: str


class Rewrite(Strict):
    question: str
    needs_clarification: bool


class Equivalent(Strict):
    equivalent: bool
    reason: str


class Result(Strict):
    status: Literal["checked", "conversational", "abstained", "needs_review"]
    claims: list[Claim] = Field(default_factory=list)
    sources: list[Evidence] = Field(default_factory=list)
    notes: list[str] = Field(default_factory=list)
    displayed: list[str] = Field(default_factory=list)
    language: str = "English"


class PDFInfo(Strict):
    page_count: int = Field(ge=1, description="Physical PDF page count, including blank pages.")


class ExtractedPage(Strict):
    number: int = Field(ge=1)
    original: str = Field(description="Complete faithful transcript in the original scripts. Empty only for a blank page.")
    english: str = Field(description="Faithful English translation, preserving numeric tokens and legal meaning. Empty for a blank page.")
    language: str = Field(description="Language name, Mixed, or Unknown.")
    warnings: list[str] = Field(default_factory=list)


class PDFBatch(Strict):
    pages: list[ExtractedPage] = Field(min_length=1)