src/knowledge_extraction/extract/schemas.py

"""What the MODEL is asked to return.

Deliberately separate from `models.py`. The model never supplies `doc_id` (we
know it), never sets `extraction_status`, and never sets the conflict or diff
fields — validation owns those. **Asking a model for fields it cannot know is an
invitation to fabricate**, so the request schema is narrower than the stored one.
"""

from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, Field

from ..models import SubdomainEnum

Language = Literal["id", "en", "mixed"]


class ProvenanceDraft(BaseModel):
    section_no: str | None = None
    page: int
    span: str


class GlossaryDraft(BaseModel):
    term: str
    full_name: str | None = None
    # The wording exactly as the document writes it, even when "wrong" — the
    # standard heads its section "Physical of Availability (PA)". Surfacing the
    # discrepancy is the point; normalising it hides a decision the expert owns.
    source_wording: str | None = None
    definition: str | None = None
    formula_latex: str | None = None
    interpretation: str | None = None
    subdomain_tags: list[SubdomainEnum] = Field(default_factory=list)
    domain: str | None = None
    company: str | None = None
    language: Language | None = None
    provenance: ProvenanceDraft


class RuleDraft(BaseModel):
    rule_id: str
    statement: str | None = None
    condition: str | None = None
    consequence: str | None = None
    applies_to: str | None = None
    subdomain_tags: list[SubdomainEnum] = Field(default_factory=list)
    language: Language | None = None
    provenance: ProvenanceDraft


class VariableDraft(BaseModel):
    symbol: str
    meaning: str | None = None


class FormulaDraft(BaseModel):
    name: str | None = None
    formula_latex: str | None = None
    variables: list[VariableDraft] = Field(default_factory=list)
    unit: str | None = None
    provenance: ProvenanceDraft


class SummaryDraft(BaseModel):
    title: str | None = None
    purpose: str | None = None
    scope: str | None = None
    key_parameters: list[str] = Field(default_factory=list)
    summary_md: str | None = None
    provenance: ProvenanceDraft


DRAFTS = {
    "glossary": GlossaryDraft,
    "rule": RuleDraft,
    "formula": FormulaDraft,
    "summary": SummaryDraft,
}


def schema_for(branch: str) -> dict:
    return DRAFTS[branch].model_json_schema()