src/schemas/fraud.py

"""Pydantic schemas for fraud detection."""

from typing import Dict, Optional

from pydantic import BaseModel, Field


class TransactionData(BaseModel):
    """Transaction data schema."""

    trans_date_trans_time: Optional[str] = None
    merchant: Optional[str] = None
    category: Optional[str] = None
    amt: Optional[float] = None
    first: Optional[str] = None
    last: Optional[str] = None
    gender: Optional[str] = None
    city: Optional[str] = None
    state: Optional[str] = None
    job: Optional[str] = None
    city_pop: Optional[int] = None


class TransactionSummary(BaseModel):
    """Transaction summary schema."""

    total_transactions: int
    fraud_count: int
    fraud_percentage: float
    total_amount: float
    average_amount: float
    categories: Dict[str, int]


class FraudAnalysisRequest(BaseModel):
    """Request schema for fraud analysis."""

    transaction_id: Optional[int] = Field(None, description="Transaction ID from dataset")
    transaction_data: Optional[TransactionData] = Field(None, description="Direct transaction data")
    use_rag: bool = Field(True, description="Whether to use RAG for context")

    class Config:
        """Pydantic config."""

        json_schema_extra = {
            "example": {
                "transaction_id": 0,
                "use_rag": True,
            }
        }


class FraudAnalysisResponse(BaseModel):
    """Response schema for fraud analysis."""

    transaction: Dict
    analysis: str
    formatted_transaction: str
    success: bool = True
    error: Optional[str] = None