from pydantic import BaseModel, Field, validator from typing import Optional, List from datetime import datetime from enum import Enum class ExpenseType(str, Enum): INCOME = "income" EXPENSE = "expense" class Expense(BaseModel): user_id: str amount: float = Field(..., gt=0, description="Expense amount") category: str = Field(..., description="Expense category (e.g., Groceries, Transport)") description: Optional[str] = None date: datetime type: ExpenseType = ExpenseType.EXPENSE class Budget(BaseModel): user_id: str category: str amount: float = Field(..., gt=0, description="Budget amount") period: str = Field(..., description="Budget period: daily, weekly, monthly, yearly") start_date: datetime end_date: Optional[datetime] = None class BudgetRecommendation(BaseModel): category: str = Field(..., description="Category name (e.g., Groceries, Transport)") category_id: Optional[str] = Field(None, description="Category ID from database") average_expense: float recommended_budget: float reason: str confidence: float = Field(..., ge=0, le=1, description="Confidence score (0-1)") action: Optional[str] = Field( None, description="AI suggestion: increase, decrease, or keep" ) class CategoryExpense(BaseModel): category: str average_monthly_expense: float total_expenses: int months_analyzed: int class RecommendationRequest(BaseModel): user_id: str = Field(..., description="User identifier") category_id: str = Field(..., description="Category ID to check for previous data") budget_amount: Optional[float] = Field(None, gt=0, description="Current budget amount for this category (optional)") @validator('budget_amount') def validate_budget_amount(cls, v): if v is not None: if v <= 0: raise ValueError('budget_amount must be greater than 0') # Auto-cap unreasonably large numbers instead of rejecting # This handles corrupted data gracefully - cap at 1 trillion if v > 1e12: # More than 1 trillion is likely corrupted data print(f"⚠️ Auto-capping corrupted budget_amount: {v:,.2e} -> 1,000,000,000,000") return 1e12 # Cap at 1 trillion instead of rejecting return v class RecommendationResponse(BaseModel): has_previous_data: bool message: Optional[str] = None recommendations: Optional[List[BudgetRecommendation]] = None class RecommendationByNameRequest(BaseModel): user_id: str = Field(..., description="User identifier") category_name: str = Field(..., description="Head Category name (e.g., 'Food & Drinks', 'Shopping', 'Housing')") budget_amount: Optional[float] = Field(None, gt=0, description="Current budget amount for this category (optional)") @validator('budget_amount') def validate_budget_amount(cls, v): if v is not None: if v <= 0: raise ValueError('budget_amount must be greater than 0') if v > 1e12: print(f"⚠️ Auto-capping corrupted budget_amount: {v:,.2e} -> 1,000,000,000,000") return 1e12 return v