ai-tax-agent/libs/schemas/coverage/core.py

"""Core coverage policy models."""

from collections.abc import Callable
from datetime import datetime
from typing import Any

from pydantic import BaseModel, Field

from ..enums import Role


class Validity(BaseModel):
    """Validity constraints for evidence"""

    within_tax_year: bool = False
    available_by: str | None = None
    date_tolerance_days: int = 30


class StatusClassifier(BaseModel):
    """Rules for classifying evidence status"""

    min_ocr: float = 0.82
    min_extract: float = 0.85
    date_in_year: bool = True
    date_in_year_or_tolerance: bool = True
    conflict_rules: list[str] = Field(default_factory=list)


class StatusClassifierConfig(BaseModel):
    """Complete status classifier configuration"""

    present_verified: StatusClassifier
    present_unverified: StatusClassifier
    conflicting: StatusClassifier
    missing: StatusClassifier = Field(default_factory=lambda: StatusClassifier())


class EvidenceItem(BaseModel):
    """Evidence requirement definition"""

    id: str
    role: Role
    condition: str | None = None
    boxes: list[str] = Field(default_factory=list)
    acceptable_alternatives: list[str] = Field(default_factory=list)
    validity: Validity = Field(default_factory=Validity)
    reasons: dict[str, str] = Field(default_factory=dict)


class CrossCheck(BaseModel):
    """Cross-validation rule"""

    name: str
    logic: str


class SchedulePolicy(BaseModel):
    """Policy for a specific tax schedule"""

    guidance_hint: str | None = None
    evidence: list[EvidenceItem] = Field(default_factory=list)
    cross_checks: list[CrossCheck] = Field(default_factory=list)
    selection_rule: dict[str, str] = Field(default_factory=dict)
    notes: dict[str, Any] = Field(default_factory=dict)


class Trigger(BaseModel):
    """Schedule trigger condition"""

    any_of: list[str] = Field(default_factory=list)
    all_of: list[str] = Field(default_factory=list)


class GuidanceRef(BaseModel):
    """Reference to guidance document"""

    doc_id: str
    kind: str


class QuestionTemplates(BaseModel):
    """Templates for generating clarifying questions"""

    default: dict[str, str] = Field(default_factory=dict)
    reasons: dict[str, str] = Field(default_factory=dict)


class ConflictRules(BaseModel):
    """Rules for handling conflicting evidence"""

    precedence: list[str] = Field(default_factory=list)
    escalation: dict[str, Any] = Field(default_factory=dict)


class TaxYearBoundary(BaseModel):
    """Tax year date boundaries"""

    start: str
    end: str


class Defaults(BaseModel):
    """Default configuration values"""

    confidence_thresholds: dict[str, float] = Field(default_factory=dict)
    date_tolerance_days: int = 30
    require_lineage_bbox: bool = True
    allow_bank_substantiation: bool = True


class Privacy(BaseModel):
    """Privacy and PII handling configuration"""

    vector_pii_free: bool = True
    redact_patterns: list[str] = Field(default_factory=list)


class CoveragePolicy(BaseModel):
    """Complete coverage policy definition"""

    version: str
    jurisdiction: str
    tax_year: str
    tax_year_boundary: TaxYearBoundary
    defaults: Defaults
    document_kinds: list[str] = Field(default_factory=list)
    guidance_refs: dict[str, GuidanceRef] = Field(default_factory=dict)
    triggers: dict[str, Trigger] = Field(default_factory=dict)
    schedules: dict[str, SchedulePolicy] = Field(default_factory=dict)
    status_classifier: StatusClassifierConfig
    conflict_resolution: ConflictRules
    question_templates: QuestionTemplates
    privacy: Privacy


class CompiledCoveragePolicy(BaseModel):
    """Coverage policy with compiled predicates"""

    policy: CoveragePolicy
    compiled_predicates: dict[str, Callable[[str, str], bool]] = Field(
        default_factory=dict
    )
    compiled_at: datetime
    hash: str
    source_files: list[str] = Field(default_factory=list)