"""BudgetBuddy — gradio.Server backend (v2). Custom dark frontend (frontend/) over JSON API endpoints that wrap the on-device small models: MiniCPM-V-4.6 vision on the Space's ZeroGPU, MiniCPM4.1-8B text on Modal (core.inference routes each). Flexible transaction model, an understanding agent, per-item analytics, budgets, and a tool-using chat agent. """ from __future__ import annotations import base64 import io import json from datetime import date from pathlib import Path import gradio as gr from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles from core.extract import extract_receipt, extract_payment, reconcile from core.categorize import CATEGORIES, understand from core import storage, analytics, auth, agent FRONTEND = Path(__file__).parent / "frontend" app = gr.Server() _PERSIST_KEYS = ["vendor", "date", "currency", "line_items", "charges", "total", "category", "understanding", "note", "source"] # --------------------------------------------------------------------------- # # Helpers # --------------------------------------------------------------------------- # def _open_image(image_b64: str): from PIL import Image if not image_b64: return None if "," in image_b64 and image_b64.strip().lower().startswith("data:"): image_b64 = image_b64.split(",", 1)[1] return Image.open(io.BytesIO(base64.b64decode(image_b64))).convert("RGB") def _loads(s, default): if isinstance(s, (dict, list)): return s try: return json.loads(s) if s else default except Exception: return default def _clean(record: dict) -> dict: out = {k: record.get(k) for k in _PERSIST_KEYS if k in record} out.setdefault("line_items", record.get("line_items") or []) out.setdefault("charges", record.get("charges") or []) return out # --------------------------------------------------------------------------- # # Auth # --------------------------------------------------------------------------- # @app.api(name="signup") def signup(username: str, pin: str) -> dict: ok, msg = auth.signup(username, pin) return {"ok": ok, "message": msg, "token": auth.issue_token(username) if ok else "", "user": auth.normalize_username(username) if ok else ""} @app.api(name="login") def login(username: str, pin: str) -> dict: ok, msg = auth.login(username, pin) return {"ok": ok, "message": msg, "token": auth.issue_token(username) if ok else "", "user": auth.normalize_username(username) if ok else ""} # --------------------------------------------------------------------------- # # Capture # --------------------------------------------------------------------------- # @app.api(name="extract_receipt") def api_extract_receipt(image_b64: str) -> dict: rec = extract_receipt(_open_image(image_b64)) rec["reconcile"] = reconcile(rec) return rec @app.api(name="extract_payment") def api_extract_payment(image_b64: str) -> dict: return extract_payment(_open_image(image_b64)) @app.api(name="understand") def api_understand(record_json: str) -> dict: """The 8B reasons over the bill: per-item + overall categories + summary.""" record = _loads(record_json, {}) record.setdefault("source", "receipt") try: record = understand(record) except Exception as e: # pragma: no cover print(f"[app] understand failed: {e}") record.setdefault("category", "Other") return {"record": record, "reconcile": reconcile(record), "uncertain": record.get("_uncertain") or [], "understanding": record.get("understanding", "")} # --------------------------------------------------------------------------- # # Data (token-scoped) # --------------------------------------------------------------------------- # @app.api(name="save") def api_save(token: str, record_json: str) -> dict: user = auth.verify_token(token) if not user: return {"ok": False, "message": "Please sign in again."} record = _clean(_loads(record_json, {})) if analytics._num(record.get("total", 0)) <= 0 and not record.get("line_items"): return {"ok": False, "message": "Nothing to save — add an amount first."} try: storage.save(user, record) return {"ok": True, "message": "Saved.", "count": len(storage.load(user, force=True))} except Exception as e: return {"ok": False, "message": f"Save failed: {e}"} def _txn_dict(r: dict) -> dict: return {"date": r.get("date", ""), "vendor": r.get("vendor", ""), "total": analytics._num(r.get("total")), "category": analytics._category(r), "source": r.get("source", "receipt"), "understanding": r.get("understanding", ""), "currency": r.get("currency", ""), "items": [{"name": str(it.get("name", "")), "amount": analytics._num(it.get("amount")), "qty": analytics._num(it.get("qty", 1)) or 1, "category": str(it.get("category") or analytics._category(r))} for it in (r.get("line_items") or [])], "charges": [{"label": str(ch.get("label", "")), "amount": analytics._num(ch.get("amount"))} for ch in (r.get("charges") or [])]} def _sort_recent(records): return sorted(records, key=lambda r: (analytics.parse_date(r.get("date")) or date.min, str(r.get("saved_at", ""))), reverse=True) @app.api(name="dashboard") def api_dashboard(token: str, start: str = "", end: str = "", category: str = "All", granularity: str = "Monthly", cal_month: str = "") -> dict: user = auth.verify_token(token) if not user: return {"ok": False, "error": "auth"} try: records = storage.load(user, force=True) summ = analytics.summary(records) cat = None if category in ("", "All", None) else category filt = analytics.filter_records( records, analytics.parse_date(start), analytics.parse_date(end), cat) today = date.today() try: cy, cm = (int(x) for x in cal_month.split("-")) if cal_month else (today.year, today.month) except Exception: cy, cm = today.year, today.month return { "ok": True, "summary": summ, "budget": storage.get_budget(user), "by_category": analytics.spend_by_category(filt), "over_time": analytics.spend_over_time(filt, granularity or "Monthly"), "comparison": analytics.category_comparison(records), "calendar": analytics.calendar_data(records, cy, cm), "calendar_month": f"{cy}-{cm:02d}", "transactions": [_txn_dict(r) for r in _sort_recent(filt)], "recent": [_txn_dict(r) for r in _sort_recent(records)[:8]], "has_any": len(records) > 0, } except Exception as e: import traceback print(f"[app] dashboard failed: {e}\n{traceback.format_exc()[-800:]}") return {"ok": False, "error": str(e)} @app.api(name="set_budget") def api_set_budget(token: str, amount: float) -> dict: user = auth.verify_token(token) if not user: return {"ok": False} try: return {"ok": True, "budget": storage.set_budget(user, amount)} except Exception as e: return {"ok": False, "message": str(e)} @app.api(name="chat") def api_chat(token: str, message: str, history_json: str = "") -> dict: user = auth.verify_token(token) if not user: return {"reply": "Please sign in again.", "trace": []} records = storage.load(user, force=True) try: budget = storage.get_budget(user) except Exception: budget = 0.0 result = agent.run(message, records, budget=budget) return {"reply": result["reply"], "trace": result.get("trace", [])} @app.api(name="categories") def api_categories() -> dict: return {"categories": CATEGORIES} # --------------------------------------------------------------------------- # # Static custom frontend # --------------------------------------------------------------------------- # if (FRONTEND / "assets").exists(): app.mount("/assets", StaticFiles(directory=str(FRONTEND / "assets")), name="assets") @app.get("/", response_class=HTMLResponse) async def home(): html = (FRONTEND / "index.html").read_text(encoding="utf-8") try: v = int((FRONTEND / "assets" / "app.js").stat().st_mtime) except Exception: v = 1 return html.replace("/assets/app.js", f"/assets/app.js?v={v}") if __name__ == "__main__": app.launch(show_error=True)