"""Affordability module — now with graceful UX and smart category detection."""
import datetime
from typing import Dict, Any
from utils.date_utils import parse_flexible_date
from utils.budget_core import (
calculate_remaining_budget,
get_budget_for_month,
budget_exists_for_month,
get_current_month
)
from utils.affordability_core import check_affordability_logic
from utils.category_detection import detect_category
from config import CATEGORIES
def register_affordability_tools(mcp):
@mcp.tool()
def check_affordability(
amount: float,
category: str = '',
item_name: str = '',
date: str = ''
) -> Dict[str, Any]:
"""
Conversational affordability check:
- Smart category detection
- Graceful inconsistency handling
- Soft fallbacks for date parsing
"""
# ---------------------------------------------------------------
# 1) Parse date (gracefully)
# ---------------------------------------------------------------
try:
expense_date = parse_flexible_date(date) if date else datetime.date.today().isoformat()
date_warning = None
except Exception:
expense_date = datetime.date.today().isoformat()
date_warning = "The date was unclear. I used today's date."
month = expense_date[:7]
# ---------------------------------------------------------------
# 2) If category missing OR invalid → detect automatically
# ---------------------------------------------------------------
detected_meta = None
if not category or category not in CATEGORIES:
detection = detect_category(item_name or category or "")
if detection["auto_assign"]:
category = detection["detected_category"]
detected_meta = {
"auto_detected": True,
"category": category,
"confidence": detection["confidence"],
}
else:
return {
"status": "needs_clarification",
"message": "I couldn't confidently determine the category for this purchase.",
"suggestions": [d["category"] for d in detection["alternatives"]],
"top_guess": detection["detected_category"],
"next_steps": [
"Pick a category from the suggestions.",
"Specify a category manually.",
"Say 'use total budget' to skip category mode."
]
}
# ---------------------------------------------------------------
# 3) Check if budget exists
# ---------------------------------------------------------------
if not budget_exists_for_month(month):
return {
"status": "no_budget",
"message": f"No budget found for {month}.",
"options": [
"Create a new budget",
"Check affordability using total income",
"Proceed without budget constraints"
]
}
# ---------------------------------------------------------------
# 4) Compute remaining budget
# ---------------------------------------------------------------
success, b = calculate_remaining_budget(month, category)
if not success:
return {
"status": "soft_error",
"message": "I couldn't compute the remaining budget.",
"suggestion": "Would you like to check your overall monthly spending instead?"
}
# ---------------------------------------------------------------
# 5) Normalize values for affordability engine
# ---------------------------------------------------------------
if b.get("budget_mode") == "total_only":
category_remaining = None
category_budget = None
total_remaining = b.get("remaining_total", 0)
total_budget = b.get("total_budget", 0)
mode = "total_only"
else:
breakdown = b.get("category_breakdown") or b.get("categories") or {}
cat_data = breakdown.get(category, {})
category_remaining = cat_data.get("remaining")
category_budget = cat_data.get("budget")
total_remaining = b.get("remaining_total")
total_budget = b.get("total_budget")
mode = "category_mode"
# ---------------------------------------------------------------
# 6) Run the affordability logic
# ---------------------------------------------------------------
aff = check_affordability_logic(
amount=amount,
category_remaining=category_remaining or 0,
total_remaining=total_remaining or 0,
category_budget=category_budget or 0,
total_budget=total_budget or 0,
category=category,
date=expense_date
)
# ---------------------------------------------------------------
# 7) Build response
# ---------------------------------------------------------------
response = {
"status": "ok",
"amount": amount,
"category": category,
"month": month,
"budget_mode": mode,
"affordability": aff,
"timestamp": datetime.datetime.now().isoformat()
}
if date_warning:
response["date_warning"] = date_warning
if detected_meta:
response["auto_category_detection"] = detected_meta
return response
