"""Budget module — now with graceful handling, helpful suggestions."""
from typing import Dict, Any
from config import CATEGORIES
from utils.budget_core import (
parse_flexible_month,
load_budget,
budget_exists_for_month,
set_budget,
get_budget_for_month,
list_all_budgets,
calculate_remaining_budget,
get_current_month
)
from utils.category_detection import detect_category
def register_budget_tools(mcp):
# =====================================================================
# SET MONTHLY BUDGET
# =====================================================================
@mcp.tool()
def set_monthly_budget(
total_amount: float,
month: str = "",
category_budgets: Dict[str, float] | None = None
) -> Dict[str, Any]:
"""User-friendly wrapper for setting budgets."""
# Default if omitted
category_budgets = category_budgets or {}
# Resolve month
try:
target_month = parse_flexible_month(month) if month else get_current_month()
except Exception:
return {
"status": "needs_clarification",
"message": f"I couldn't understand '{month}'.",
"suggestion": "Try formats like 'Jan 2025', 'next month', or '2024-12'."
}
# Validate categories
valid_cat_budgets = {}
invalid_entries = []
for cat, amt in category_budgets.items():
if cat not in CATEGORIES:
det = detect_category(cat)
invalid_entries.append({
"entered": cat,
"guess": det["detected_category"],
"alternatives": det["alternatives"]
})
else:
valid_cat_budgets[cat] = amt
if invalid_entries:
return {
"status": "invalid_categories",
"message": "Some categories were not recognized.",
"invalid_entries": invalid_entries,
"next_steps": [
"Correct the invalid categories",
"Use suggested category replacements",
"Remove uncertain categories"
]
}
# Check matching totals
if valid_cat_budgets:
sum_cats = sum(valid_cat_budgets.values())
if abs(sum_cats - total_amount) >= 1:
return {
"status": "needs_clarification",
"message": (
f"Category totals (₹{sum_cats:.2f}) do not match total budget "
f"(₹{total_amount:.2f})."
),
"options": [
"Normalize categories automatically",
"Adjust category values",
"Ignore category mode and use total-only"
]
}
try:
new_b = set_budget(total_amount, target_month, valid_cat_budgets)
except Exception as e:
return {
"status": "soft_error",
"message": f"Could not set budget: {e}",
"suggestion": "Try again or use a simpler budget structure."
}
mode = "category_mode" if valid_cat_budgets else "total_only"
return {
"status": "ok",
"message": f"Budget for {target_month} saved ({mode}).",
"budget": new_b
}
# =====================================================================
# GET REMAINING BUDGET
# =====================================================================
@mcp.tool()
def get_remaining_budget(month: str = '', category: str = '') -> Dict[str, Any]:
try:
target = parse_flexible_month(month) if month else get_current_month()
except Exception:
return {
"status": "needs_clarification",
"message": f"I couldn't understand '{month}'.",
"suggestion": "Try 'this month', 'last month', or '2025-01'."
}
success, data = calculate_remaining_budget(target, category)
if not success:
return {
"status": "soft_error",
"message": data.get("error", "Could not compute remaining budget."),
"suggestion": "Try checking total budget instead."
}
return {
"status": "ok",
**data
}
# =====================================================================
# LIST BUDGETS
# =====================================================================
@mcp.tool()
def list_budgets() -> Dict[str, Any]:
all_b = list_all_budgets()
if not all_b:
return {
"status": "info",
"message": "No budgets found.",
"suggestion": "Create one using set_monthly_budget()."
}
formatted = []
for month, b in sorted(all_b.items(), reverse=True):
formatted.append({
"month": month,
"total": b["total"],
"mode": "category_mode" if b["categories"] else "total_only",
"categories": list(b["categories"].keys()) if b["categories"] else []
})
return {"status": "ok", "budgets": formatted}
# =====================================================================
# CONVERT TO CATEGORY BUDGET (with suggestions)
# =====================================================================
@mcp.tool()
def convert_to_category_budget(month: str = '', **allocations) -> Dict[str, Any]:
try:
target = parse_flexible_month(month) if month else get_current_month()
except Exception:
return {
"status": "needs_clarification",
"message": f"I couldn't understand '{month}'.",
}
b = get_budget_for_month(target)
if not b:
return {"status": "no_budget", "message": f"No budget found for {target}."}
if b["categories"]:
return {
"status": "already_category_mode",
"message": f"Budget for {target} is already category-based.",
"categories": b["categories"]
}
if not allocations:
return {
"status": "needs_clarification",
"message": "You must provide allocation percentages (e.g., food=0.3)."
}
# Check allocations sum
total_allocation = sum(allocations.values())
if abs(total_allocation - 1.0) >= 0.02:
return {
"status": "needs_clarification",
"message": f"Allocations must sum to 1.0; you provided {total_allocation:.2f}.",
"options": [
"Normalize automatically",
"Retry with corrected values"
]
}
# Validate categories
invalid = [cat for cat in allocations if cat not in CATEGORIES]
if invalid:
suggestions = {cat: detect_category(cat)["detected_category"] for cat in invalid}
return {
"status": "invalid_categories",
"message": "Some categories were not recognized.",
"invalid_entries": suggestions,
"next_steps": ["Use suggestions", "Adjust allocations manually"]
}
new_cats = {cat: round(b["total"] * pct, 2) for cat, pct in allocations.items()}
try:
updated = set_budget(b["total"], target, new_cats)
except Exception as e:
return {
"status": "soft_error",
"message": f"Could not convert budget: {e}"
}
return {
"status": "ok",
"message": f"Converted to category-mode for {target}.",
"budget": updated
}
