core/chat.py

"""Grounded spending chat agent with MiniCPM5-1B (Phase 6).

Answers questions about *your own* transactions ("How much did I spend on
groceries last month?", "What's my biggest category?"). To keep a 1B model
honest, we don't let it free-associate: we precompute the real numbers from
core.analytics, stuff them into the prompt as the ONLY source of truth, and
instruct the model to answer strictly from them.

Reuses the MiniCPM5-1B instance already loaded by core.categorize (no second
copy) and the ZeroGPU decorator from core.extract.
"""

from __future__ import annotations

from typing import Any

from core import analytics, inference

MAX_NEW_TOKENS = 256
RECENT_LIMIT = 12
MONTHS_LIMIT = 6

SYSTEM_PROMPT = (
    "You are BudgetBuddy, a friendly personal-spending assistant. The user's "
    "message contains their REAL spending data followed by a question. Always "
    "answer using that data — quote the exact numbers from it. Never say you "
    "lack access; the data is right there in the message. Be concise and warm, "
    "use the user's currency, and don't output code or JSON."
)


def _fmt_money(amount: float, currency: str) -> str:
    prefix = f"{currency} " if currency else ""
    return f"{prefix}{amount:,.2f}"


def build_context(records: list[dict[str, Any]]) -> str:
    """Compact, factual summary of the user's data — the model's only source."""
    if not records:
        return "The user has no saved transactions yet."

    summ = analytics.summary(records)
    cur = summ["currency"]
    lines: list[str] = []
    lines.append(f"Currency: {cur or 'unknown'}")
    lines.append(f"Total transactions saved: {summ['total_receipts']}")
    lines.append(
        f"This month spend: {_fmt_money(summ['this_month_total'], cur)} "
        f"across {summ['receipts_this_month']} transaction(s)"
    )
    lines.append(f"Last month spend: {_fmt_money(summ['prev_month_total'], cur)}")
    if summ["pct_change"] is not None:
        lines.append(f"Change vs last month: {summ['pct_change']:+.0f}%")
    if summ["top_category"]:
        lines.append(f"Top category this month: {summ['top_category']}")

    by_cat = analytics.spend_by_category(records)
    if by_cat:
        lines.append("Spend by category (all time):")
        for row in by_cat:
            lines.append(f"  - {row['category']}: {_fmt_money(row['amount'], cur)}")

    by_month = analytics.spend_over_time(records, "Monthly")
    if by_month:
        lines.append("Spend by month:")
        for row in by_month[-MONTHS_LIMIT:]:
            lines.append(f"  - {row['period']}: {_fmt_money(row['amount'], cur)}")

    recent = analytics.transactions_table(records)[:RECENT_LIMIT]
    if recent:
        lines.append(f"Most recent transactions (up to {RECENT_LIMIT}):")
        for date, vendor, total, category in recent:
            lines.append(
                f"  - {date or '?'} | {vendor or '(no name)'} | "
                f"{_fmt_money(float(total), cur)} | {category}"
            )
    return "\n".join(lines)


def _generate(messages: list[dict[str, str]]) -> str:
    return inference.text_generate(messages, max_new_tokens=MAX_NEW_TOKENS)


def answer(
    question: str,
    records: list[dict[str, Any]],
    history: list[tuple[str, str]] | None = None,
) -> str:
    """Answer a spending question grounded in the user's records.

    `history` is an optional list of prior (user, assistant) turns. Never raises.
    """
    question = (question or "").strip()
    if not question:
        return "Ask me about your spending — e.g. “How much did I spend on Dining this month?”"

    context = build_context(records)
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user_turn, bot_turn in (history or [])[-3:]:
        if user_turn:
            messages.append({"role": "user", "content": str(user_turn)})
        if bot_turn:
            messages.append({"role": "assistant", "content": str(bot_turn)})
    # Data sits in the user turn, right next to the question — small models
    # attend to it far more reliably there than in the system prompt.
    messages.append({
        "role": "user",
        "content": f"Here is my spending data:\n{context}\n\nQuestion: {question}",
    })

    for attempt in range(2):
        try:
            out = _generate(messages)
            if out:
                return out
        except Exception as e:  # pragma: no cover - model/runtime dependent
            print(f"[chat] generation failed (attempt {attempt + 1}): {e}")
    return "Sorry — I couldn't answer that right now. Please try again."