roa2web-service-auto/backend/modules/data_entry/schemas/ocr.py

"""Pydantic schemas for OCR API."""

from datetime import date
from decimal import Decimal
from typing import Optional, List

from pydantic import BaseModel, Field


class TvaEntry(BaseModel):
    """Single TVA entry with code, percentage and amount."""
    code: Optional[str] = Field(default=None, description="TVA code: A, B, C, D")
    percent: int = Field(description="TVA percentage: 0, 5, 9, 19, 21")
    amount: Decimal = Field(description="TVA amount for this rate")


class PaymentMethod(BaseModel):
    """Payment method entry from OCR."""
    method: str = Field(description="CARD or NUMERAR")
    amount: Decimal = Field(description="Amount paid")


class ExtractionData(BaseModel):
    """Extracted receipt data from OCR."""

    receipt_type: str = Field(default='bon_fiscal', description="Receipt type: bon_fiscal or chitanta")
    receipt_number: Optional[str] = Field(default=None, description="Receipt number")
    receipt_series: Optional[str] = Field(default=None, description="Receipt series")
    receipt_date: Optional[date] = Field(default=None, description="Receipt date")
    amount: Optional[Decimal] = Field(default=None, description="Total amount")
    partner_name: Optional[str] = Field(default=None, description="Vendor/partner name")
    cui: Optional[str] = Field(default=None, description="CUI (fiscal identification code)")
    description: Optional[str] = Field(default=None, description="Optional description")

    # Additional extracted fields - Multiple TVA entries support
    tva_entries: List[TvaEntry] = Field(default=[], description="List of TVA entries by rate (A, B, C, D)")
    tva_total: Optional[Decimal] = Field(default=None, description="Total TVA amount")
    address: Optional[str] = Field(default=None, description="Vendor address")
    items_count: Optional[int] = Field(default=None, description="Number of items/articles")

    # Payment methods extracted from receipt
    payment_methods: List[PaymentMethod] = Field(default=[], description="Payment methods from receipt (CARD, NUMERAR)")
    suggested_payment_mode: Optional[str] = Field(default=None, description="Auto-suggested payment mode based on OCR (casa/banca)")

    # Client data (for B2B receipts - buyer information)
    client_name: Optional[str] = Field(default=None, description="Client/customer company name")
    client_cui: Optional[str] = Field(default=None, description="Client CUI/CIF fiscal code")
    client_address: Optional[str] = Field(default=None, description="Client address")

    confidence_amount: float = Field(default=0.0, ge=0, le=1, description="Amount extraction confidence")
    confidence_date: float = Field(default=0.0, ge=0, le=1, description="Date extraction confidence")
    confidence_vendor: float = Field(default=0.0, ge=0, le=1, description="Vendor extraction confidence")
    confidence_client: float = Field(default=0.0, ge=0, le=1, description="Client extraction confidence")
    overall_confidence: float = Field(default=0.0, ge=0, le=1, description="Overall confidence score")
    raw_text: str = Field(default="", description="Raw OCR text")
    ocr_engine: str = Field(default="", description="OCR engine used: paddleocr or tesseract")
    processing_time_ms: int = Field(default=0, ge=0, description="Processing time in milliseconds")

    class Config:
        """Pydantic config."""
        json_schema_extra = {
            "example": {
                "receipt_type": "bon_fiscal",
                "receipt_number": "1360760",
                "receipt_series": "0146",
                "receipt_date": "2025-10-11",
                "amount": 186.16,
                "partner_name": "FIVE-HOLDING S.A.",
                "cui": "10562600",
                "description": None,
                "tva_entries": [
                    {"code": "A", "percent": 19, "amount": 25.00},
                    {"code": "B", "percent": 9, "amount": 7.31}
                ],
                "tva_total": 32.31,
                "address": "JUD. CONSTANTA, MUN. CONSTANTA, STR. ION ROATA NR. 3",
                "items_count": 17,
                "confidence_amount": 0.98,
                "confidence_date": 0.98,
                "confidence_vendor": 0.95,
                "overall_confidence": 0.97,
                "raw_text": "FIVE-HOLDING S.A.\nCIF: RO10562600\n..."
            }
        }


class OCRResponse(BaseModel):
    """OCR API response."""

    success: bool = Field(description="Whether OCR processing was successful")
    message: str = Field(description="Status message")
    data: Optional[ExtractionData] = Field(default=None, description="Extracted data")

    class Config:
        """Pydantic config."""
        json_schema_extra = {
            "example": {
                "success": True,
                "message": "OCR processing successful. Found: amount, date, vendor",
                "data": {
                    "receipt_type": "bon_fiscal",
                    "receipt_number": "12345",
                    "receipt_date": "2024-01-15",
                    "amount": 125.50,
                    "partner_name": "MEGA IMAGE SRL",
                    "cui": "12345678",
                    "confidence_amount": 0.95,
                    "confidence_date": 0.90,
                    "confidence_vendor": 0.75,
                    "overall_confidence": 0.87,
                    "raw_text": "BON FISCAL\nMEGA IMAGE SRL\n..."
                }
            }
        }


class OCRStatusResponse(BaseModel):
    """OCR service status response."""

    available: bool = Field(description="Whether OCR service is available")
    engines: list[str] = Field(description="Available OCR engines")
    message: str = Field(description="Status message")