simbarag/blueprints/ynab/tools.py

"""LangChain agent tools that answer YNAB questions from the local mirror.

Every tool calls ``ensure_fresh`` first, which runs a cheap delta sync when
the mirror is more than a couple of minutes old.
"""

import functools
import json
import logging
from datetime import date, datetime, timedelta
from decimal import Decimal
from zoneinfo import ZoneInfo

from langchain.tools import tool
from tortoise import connections

from .sync import ensure_fresh

logger = logging.getLogger(__name__)

NOT_SYNCED = "YNAB data is unavailable: the budget has not been synced yet."

# Keep in step with TOOL_RESULT_MAX_CHARS in blueprints/conversation/logic.py;
# longer results get truncated there anyway and crowd the 32k context.
MAX_TRANSACTIONS = 40

# ynab_recent_changes accepts the names it prints as well as stored types
ENTITY_TYPE_ALIASES = {
    "assignment": "month_category",
    "transfer": "transaction",
    "split": "subtransaction",
    "group": "category_group",
}
ENTITY_TYPES = (
    "transaction",
    "subtransaction",
    "month_category",
    "category",
    "category_group",
    "payee",
    "account",
)


def ynab_tool_errors(fn):
    """Return errors to the model as text: LangGraph's ToolNode re-raises
    anything else, which would fail the whole chat turn."""

    @functools.wraps(fn)
    async def wrapper(*args, **kwargs):
        try:
            return await fn(*args, **kwargs)
        except Exception as e:
            logger.exception(f"{fn.__name__} failed")
            return f"YNAB error in {fn.__name__}: {type(e).__name__}: {e}"

    return wrapper


# Every transaction line, with split transactions expanded into their
# subtransactions so category/payee filters see the real allocation.
LINES_CTE = """
WITH lines AS (
    SELECT t.id, t.date, t.amount, t.memo, t.cleared, t.approved,
           t.account_id, t.account_name, t.payee_name, t.category_id,
           t.category_name, t.transfer_account_id
    FROM ynab_transactions t
    WHERE NOT t.deleted AND NOT EXISTS (
        SELECT 1 FROM ynab_subtransactions s
        WHERE s.transaction_id = t.id AND NOT s.deleted
    )
    UNION ALL
    SELECT s.id, t.date, s.amount, COALESCE(s.memo, t.memo), t.cleared, t.approved,
           t.account_id, t.account_name, COALESCE(sp.name, t.payee_name),
           s.category_id, sc.name, s.transfer_account_id
    FROM ynab_subtransactions s
    JOIN ynab_transactions t ON t.id = s.transaction_id
    LEFT JOIN ynab_payees sp ON sp.id = s.payee_id
    LEFT JOIN ynab_categories sc ON sc.id = s.category_id
    WHERE NOT s.deleted AND NOT t.deleted
)
"""

# Spending = outflows (net of refunds) on budget accounts, excluding transfers
# (credit card payments, moves to savings) and income.
SPENDING_FILTER = """
    a.on_budget AND l.transfer_account_id IS NULL
    AND NOT COALESCE(c.internal, false)
"""

GROUP_BY_EXPR = {
    "category": "COALESCE(l.category_name, 'Uncategorized')",
    "category_group": "COALESCE(g.name, 'Uncategorized')",
    "payee": "COALESCE(l.payee_name, 'Unknown payee')",
    "account": "l.account_name",
    "month": "to_char(l.date, 'YYYY-MM')",
}


def _today() -> date:
    return datetime.now(ZoneInfo("America/New_York")).date()


def _usd(value: Decimal | float | None) -> str:
    value = Decimal(value or 0)
    sign = "-" if value < 0 else ""
    return f"{sign}${abs(value):,.2f}"


def _parse_date(value: str, default: date | None) -> date | None:
    return date.fromisoformat(value) if value else default


def _month_start(month: str) -> date:
    if not month:
        return _today().replace(day=1)
    return date.fromisoformat(month[:7] + "-01")


async def _query(sql: str, *params) -> list[dict]:
    return await connections.get("default").execute_query_dict(sql, list(params))


@tool
@ynab_tool_errors
async def ynab_budget_summary(month: str = "") -> str:
    """Get YNAB budget overview for a month: Ready to Assign, income, assigned, spent, age of money, overspent categories, and net worth.

    Args:
        month: YYYY-MM (default: current month)
    """
    state = await ensure_fresh()
    if state is None:
        return NOT_SYNCED
    try:
        start = _month_start(month)
    except ValueError:
        return "Month must be YYYY-MM."

    months = await _query("SELECT * FROM ynab_months WHERE month = $1", start)
    if not months:
        return f"No YNAB budget data for {start:%Y-%m}."
    m = months[0]

    overspent = await _query(
        """SELECT c.name, mc.balance FROM ynab_month_categories mc
           JOIN ynab_categories c ON c.id = mc.category_id
           WHERE mc.month = $1 AND mc.balance < 0 AND NOT mc.deleted
             AND NOT c.deleted AND NOT c.internal
           ORDER BY mc.balance""",
        start,
    )
    # YNAB's month "activity" nets in credit card payment categories, which
    # go up whenever a card is used, so compute spending from transactions.
    next_month = (start.replace(day=28) + timedelta(days=4)).replace(day=1)
    spent = await _query(
        f"""{LINES_CTE}
            SELECT COALESCE(-SUM(l.amount), 0) AS spent FROM lines l
            JOIN ynab_accounts a ON a.id = l.account_id
            LEFT JOIN ynab_categories c ON c.id = l.category_id
            WHERE {SPENDING_FILTER} AND l.date >= $1 AND l.date < $2""",
        start,
        next_month,
    )
    accounts = await _query(
        """SELECT on_budget, SUM(balance) AS total FROM ynab_accounts
           WHERE NOT deleted AND NOT closed GROUP BY on_budget"""
    )
    by_budget = {r["on_budget"]: r["total"] for r in accounts}

    lines = [
        f"YNAB budget '{state.budget_name}' for {start:%Y-%m}:",
        f"Ready to Assign: {_usd(m['to_be_budgeted'])}",
        f"Income: {_usd(m['income'])}",
        f"Assigned: {_usd(m['budgeted'])}",
        f"Spent: {_usd(spent[0]['spent'])}",
    ]
    if m["age_of_money"] is not None:
        lines.append(f"Age of money: {m['age_of_money']} days")
    if overspent:
        lines.append(
            "Overspent: "
            + ", ".join(f"{r['name']} ({_usd(r['balance'])})" for r in overspent)
        )
    else:
        lines.append("No overspent categories.")
    lines.append(
        f"Net worth: {_usd(sum(by_budget.values(), Decimal(0)))} "
        f"(budget accounts {_usd(by_budget.get(True))}, "
        f"tracking accounts {_usd(by_budget.get(False))})"
    )
    return "\n".join(lines)


@tool
@ynab_tool_errors
async def ynab_account_balances(include_closed: bool = False) -> str:
    """List YNAB accounts (checking, savings, credit cards, investments) with current balances.

    Args:
        include_closed: Also list closed accounts
    """
    if await ensure_fresh() is None:
        return NOT_SYNCED
    rows = await _query(
        """SELECT name, type, on_budget, closed, balance, cleared_balance
           FROM ynab_accounts WHERE NOT deleted AND (NOT closed OR $1)
           ORDER BY on_budget DESC, type, name""",
        include_closed,
    )
    if not rows:
        return "No YNAB accounts found."
    lines = []
    for r in rows:
        tags = [r["type"]] + ([] if r["on_budget"] else ["tracking"])
        if r["closed"]:
            tags.append("closed")
        line = f"- {r['name']} [{', '.join(tags)}]: {_usd(r['balance'])}"
        if r["cleared_balance"] != r["balance"]:
            line += f" (cleared {_usd(r['cleared_balance'])})"
        lines.append(line)
    total = sum((r["balance"] for r in rows if not r["closed"]), Decimal(0))
    return "\n".join(lines) + f"\n\nNet worth (open accounts): {_usd(total)}"


@tool
@ynab_tool_errors
async def ynab_search_transactions(
    start_date: str = "",
    end_date: str = "",
    payee: str = "",
    category: str = "",
    account: str = "",
    memo: str = "",
    min_amount: float = 0,
    max_amount: float = 0,
    limit: int = 20,
) -> str:
    """Search YNAB transactions. Split transactions are expanded into their category lines.

    Args:
        start_date: YYYY-MM-DD (default: 30 days ago if no other filter, else all time)
        end_date: YYYY-MM-DD (default: today)
        payee: Partial match on payee/store name
        category: Partial match on category name
        account: Partial match on account name
        memo: Partial match on memo
        min_amount: Minimum absolute dollar amount
        max_amount: Maximum absolute dollar amount
        limit: Max transactions to list (default 20, max 40)
    """
    if await ensure_fresh() is None:
        return NOT_SYNCED

    has_filter = any([payee, category, account, memo, min_amount, max_amount])
    try:
        default_start = (
            None if has_filter else date.fromordinal(_today().toordinal() - 30)
        )
        start = _parse_date(start_date, default_start)
        end = _parse_date(end_date, _today())
    except ValueError:
        return "Dates must be YYYY-MM-DD."

    conditions, params = ["l.date <= $1"], [end]

    def add(condition: str, value) -> None:
        params.append(value)
        conditions.append(condition.replace("?", f"${len(params)}"))

    if start:
        add("l.date >= ?", start)
    if payee:
        add("l.payee_name ILIKE ?", f"%{payee}%")
    if category:
        add("l.category_name ILIKE ?", f"%{category}%")
    if account:
        add("l.account_name ILIKE ?", f"%{account}%")
    if memo:
        add("l.memo ILIKE ?", f"%{memo}%")
    if min_amount:
        add("abs(l.amount) >= ?", Decimal(str(min_amount)))
    if max_amount:
        add("abs(l.amount) <= ?", Decimal(str(max_amount)))

    where = " AND ".join(conditions)
    totals = await _query(
        f"""{LINES_CTE}
            SELECT count(*) AS n,
                   COALESCE(SUM(amount) FILTER (WHERE amount < 0), 0) AS outflow,
                   COALESCE(SUM(amount) FILTER (WHERE amount > 0), 0) AS inflow
            FROM lines l WHERE {where}""",
        *params,
    )
    count = totals[0]["n"]
    range_text = f"{start or 'the beginning'} to {end}"
    if not count:
        return f"No transactions found from {range_text}."

    limit = min(max(1, limit), MAX_TRANSACTIONS)
    rows = await _query(
        f"""{LINES_CTE}
            SELECT * FROM lines l WHERE {where}
            ORDER BY l.date DESC, abs(l.amount) DESC LIMIT {limit}""",
        *params,
    )
    lines = [
        f"Found {count} transactions from {range_text}. "
        f"Outflows {_usd(totals[0]['outflow'])}, inflows {_usd(totals[0]['inflow'])}."
    ]
    for r in rows:
        line = (
            f"- {r['date']}: {r['payee_name'] or 'Unknown payee'} {_usd(r['amount'])} "
            f"({r['category_name'] or 'Uncategorized'}; {r['account_name']})"
        )
        if r["memo"]:
            line += f" — {r['memo']}"
        if not r["approved"]:
            line += " [unapproved]"
        lines.append(line)
    if count > len(rows):
        lines.append(f"(Showing {len(rows)} most recent of {count})")
    return "\n".join(lines)


@tool
@ynab_tool_errors
async def ynab_spending_report(
    group_by: str = "category",
    start_date: str = "",
    end_date: str = "",
    category: str = "",
    payee: str = "",
    top: int = 15,
) -> str:
    """Total YNAB spending grouped by category, category_group, payee, account, or month. Use group_by='month' for trends over time.

    Args:
        group_by: category, category_group, payee, account, or month
        start_date: YYYY-MM-DD (default: first of current month)
        end_date: YYYY-MM-DD (default: today)
        category: Partial match on category name
        payee: Partial match on payee name
        top: Max groups to list (default 15)
    """
    if await ensure_fresh() is None:
        return NOT_SYNCED
    key = GROUP_BY_EXPR.get(group_by)
    if key is None:
        return f"group_by must be one of: {', '.join(GROUP_BY_EXPR)}."
    try:
        start = _parse_date(start_date, _today().replace(day=1))
        end = _parse_date(end_date, _today())
    except ValueError:
        return "Dates must be YYYY-MM-DD."

    conditions, params = [SPENDING_FILTER, "l.date >= $1", "l.date <= $2"], [start, end]
    if category:
        params.append(f"%{category}%")
        conditions.append(f"l.category_name ILIKE ${len(params)}")
    if payee:
        params.append(f"%{payee}%")
        conditions.append(f"l.payee_name ILIKE ${len(params)}")

    order = "key" if group_by == "month" else "spent DESC"
    rows = await _query(
        f"""{LINES_CTE}
            SELECT {key} AS key, -SUM(l.amount) AS spent, count(*) AS n
            FROM lines l
            JOIN ynab_accounts a ON a.id = l.account_id
            LEFT JOIN ynab_categories c ON c.id = l.category_id
            LEFT JOIN ynab_category_groups g ON g.id = c.category_group_id
            WHERE {' AND '.join(conditions)}
            GROUP BY 1 ORDER BY {order}""",
        *params,
    )
    if not rows:
        return f"No spending found from {start} to {end}."

    total = sum((r["spent"] for r in rows), Decimal(0))
    top = max(1, top)
    shown = rows if group_by == "month" else rows[:top]
    lines = [f"Spending by {group_by} from {start} to {end}: total {_usd(total)}"]
    lines += [f"- {r['key']}: {_usd(r['spent'])} ({r['n']} txns)" for r in shown]
    if len(rows) > len(shown):
        rest = sum((r["spent"] for r in rows[len(shown) :]), Decimal(0))
        lines.append(f"- {len(rows) - len(shown)} others: {_usd(rest)}")
    if group_by == "month" and len(rows) > 1:
        lines.append(f"Monthly average: {_usd(total / len(rows))}")
    return "\n".join(lines)


@tool
@ynab_tool_errors
async def ynab_category_budget(month: str = "", category_group: str = "") -> str:
    """Get YNAB categories for a month grouped by category group: assigned, activity (negative = spent), available, and goal target. Flags overspending.

    Args:
        month: YYYY-MM (default: current month)
        category_group: Partial match to show only one group
    """
    if await ensure_fresh() is None:
        return NOT_SYNCED
    try:
        start = _month_start(month)
    except ValueError:
        return "Month must be YYYY-MM."

    rows = await _query(
        """SELECT g.name AS group_name, c.name, c.goal_target,
                  mc.budgeted, mc.activity, mc.balance
           FROM ynab_month_categories mc
           JOIN ynab_categories c ON c.id = mc.category_id
           JOIN ynab_category_groups g ON g.id = c.category_group_id
           WHERE mc.month = $1 AND NOT mc.deleted
             AND NOT c.deleted AND NOT c.hidden AND NOT c.internal
             AND NOT g.deleted AND NOT g.hidden
             AND g.name ILIKE $2
           ORDER BY g.name, c.name""",
        start,
        f"%{category_group}%",
    )
    if not rows:
        return f"No YNAB category data for {start:%Y-%m}."

    lines = [f"Categories for {start:%Y-%m}:"]
    current = None
    for r in rows:
        if not (r["budgeted"] or r["activity"] or r["balance"]):
            continue
        if r["group_name"] != current:
            current = r["group_name"]
            lines.append(f"\n{current}")
        line = (
            f"  {r['name']}: assigned {_usd(r['budgeted'])}, "
            f"activity {_usd(r['activity'])}, available {_usd(r['balance'])}"
        )
        if r["goal_target"]:
            line += f", goal {_usd(r['goal_target'])}"
        if r["balance"] < 0:
            line += " [OVERSPENT]"
        lines.append(line)
    return "\n".join(lines)


@tool
@ynab_tool_errors
async def ynab_recent_changes(days: int = 7, entity_type: str = "") -> str:
    """Show what changed in YNAB recently: new/edited/deleted transactions, assignment changes, renamed categories or payees.

    Args:
        days: Look back this many days (default 7)
        entity_type: Optional filter: transaction, assignment, category, payee, account
    """
    if await ensure_fresh() is None:
        return NOT_SYNCED
    entity_type = entity_type.strip().lower().removesuffix("s")
    entity_type = ENTITY_TYPE_ALIASES.get(entity_type, entity_type)
    if entity_type and entity_type not in ENTITY_TYPES:
        return f"Unknown entity_type '{entity_type}'. Use one of: transaction, assignment, category, category_group, payee, account."
    rows = await _query(
        """SELECT synced_at, entity_type, action, label, changes FROM ynab_changes
           WHERE synced_at >= now() - make_interval(days => $1)
             AND ($2 = '' OR entity_type = $2)
           ORDER BY id DESC LIMIT 40""",
        max(1, days),
        entity_type,
    )
    if not rows:
        return f"No YNAB changes recorded in the last {days} days."

    lines = []
    for r in rows:
        when = r["synced_at"].astimezone(ZoneInfo("America/New_York"))
        changes = r["changes"]
        if isinstance(changes, str):
            changes = json.loads(changes)
        if r["action"] == "updated":
            detail = "; ".join(
                f"{field}: {old} → {new}" for field, (old, new) in changes.items()
            )
        else:
            detail = ""
        entity = r["entity_type"].replace("month_category", "assignment")
        line = f"- {when:%Y-%m-%d %H:%M} {r['action']} {entity} '{r['label']}'"
        lines.append(line + (f": {detail}" if detail else ""))
    note = " (newest 40)" if len(rows) == 40 else ""
    return f"YNAB changes in the last {days} days{note}:\n" + "\n".join(lines)


YNAB_TOOLS = [
    ynab_budget_summary,
    ynab_account_balances,
    ynab_search_transactions,
    ynab_spending_report,
    ynab_category_budget,
    ynab_recent_changes,
]