From 7685650e6333f44ff6e57b8dda6497dc948fea62 Mon Sep 17 00:00:00 2001 From: Wes <43003576+wolstad@users.noreply.github.com> Date: Fri, 29 May 2026 15:51:16 -0700 Subject: [PATCH] feat(assistant): add get_budget function for budget tracking (#1966) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(assistant): add get_budget function for budget tracking Exposes the existing Budget / BudgetCategory pacing data to the AI assistant as a `get_budget` function. Supports a target month and an optional `prior_months` window for trend comparison, with the response shape matching the budget UI (totals, income, per-category status, suggested daily spend on the current month). Honors custom month_start_day by matching `Budget.param_to_date` semantics for explicit slug input, so `month` round-trips with the response's `month` field. Co-Authored-By: Claude Opus 4.7 (1M context) * test(assistant): use fixture reference for Food & Drink lookup Replace fragile string match on `bc.category.name == "Food & Drink"` with the `categories(:food_and_drink)` fixture so the test setup isn't sensitive to category-name translations. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(assistant): enforce strict month format in get_budget `Date.strptime` is lenient about trailing characters, so inputs like `"2026-05-01"` or `"may-2026foo"` were parsing successfully and being silently truncated to May 2026. Pre-validate the raw string with anchored regex patterns for the documented YYYY-MM and MMM-YYYY shapes so malformed tool arguments raise Assistant::Error instead. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(budgets): suggested_daily_spending handles custom-month periods The helper compared `budget.start_date.month/year` against `Date.current.month/year` and returned nil whenever the current period straddled two calendar months — common for families with `month_start_day != 1` (e.g., May 15–Jun 14 viewed on Jun 1). Replace the calendar-month check with `budget.current?` and compute remaining days from `budget.end_date` so the helper works for both standard and custom periods. This also restores the daily pacing row in the budget UI for custom-month families. Co-Authored-By: Claude Opus 4.7 (1M context) * fix(assistant): make get_budget read-only for prior months `prior_months: N` was calling `Budget.find_or_bootstrap` for every month, which created empty `Budget` rows (and synced `BudgetCategory` children) as a side effect of an AI query. Only the explicit target month now bootstraps; prior months use `Budget.find_by` and are dropped from the response if they don't exist. The response now includes `months_unavailable: N` so the LLM can phrase a sensible answer when fewer months come back than requested. Extract `Budget.period_for(date, family:)` to share the date-bracket math between `find_or_bootstrap`, `budget_date_valid?`, and the new read-only path in `get_budget`. Adds two tests covering the no-bootstrap behavior for prior months and the `prior_months` clamp at `MAX_PRIOR_MONTHS`. Updates the existing N+1 sorted-months test to seed prior budgets explicitly. Co-Authored-By: Claude Opus 4.7 (1M context) --------- Co-authored-by: wolstad Co-authored-by: Claude Opus 4.7 (1M context) --- app/models/assistant.rb | 1 + app/models/assistant/function/get_budget.rb | 200 ++++++++++++++++++ app/models/budget.rb | 23 +- app/models/budget_category.rb | 6 +- .../assistant/function/get_budget_test.rb | 171 +++++++++++++++ test/models/budget_category_test.rb | 16 ++ 6 files changed, 400 insertions(+), 17 deletions(-) create mode 100644 app/models/assistant/function/get_budget.rb create mode 100644 test/models/assistant/function/get_budget_test.rb diff --git a/app/models/assistant.rb b/app/models/assistant.rb index b07009396..78fcbb556 100644 --- a/app/models/assistant.rb +++ b/app/models/assistant.rb @@ -27,6 +27,7 @@ module Assistant Function::GetHoldings, Function::GetBalanceSheet, Function::GetIncomeStatement, + Function::GetBudget, Function::ImportBankStatement, Function::SearchFamilyFiles ] diff --git a/app/models/assistant/function/get_budget.rb b/app/models/assistant/function/get_budget.rb new file mode 100644 index 000000000..3cf95115b --- /dev/null +++ b/app/models/assistant/function/get_budget.rb @@ -0,0 +1,200 @@ +class Assistant::Function::GetBudget < Assistant::Function + include ActiveSupport::NumberHelper + + MAX_PRIOR_MONTHS = 11 + + class << self + def name + "get_budget" + end + + def description + <<~INSTRUCTIONS + Use this to see how the user is tracking against their monthly budget — total + budgeted vs spent and a parent/subcategory breakdown matching the budget UI. + + This is great for answering questions like: + - How am I tracking against my budget this month? + - Which categories am I over budget on? + - How does this month's spending compare to the last few months? + + Parameters: + - `month` (optional): "YYYY-MM" or "MMM-YYYY". Defaults to the current month. + - `prior_months` (optional): integer 0..#{MAX_PRIOR_MONTHS}. Number of months + preceding the target month to include for trend comparison. Default 0. + + Example (current month only): + + ``` + get_budget({}) + ``` + + Example (current month plus last 2 months): + + ``` + get_budget({ month: "#{Date.current.strftime('%Y-%m')}", prior_months: 2 }) + ``` + INSTRUCTIONS + end + end + + def strict_mode? + false + end + + def params_schema + build_schema( + properties: { + month: { + type: "string", + description: "Target month in YYYY-MM or MMM-YYYY format. Defaults to the current month." + }, + prior_months: { + type: "integer", + description: "Number of months before the target month to also return for trend comparison.", + minimum: 0, + maximum: MAX_PRIOR_MONTHS + } + } + ) + end + + def call(params = {}) + target_start = resolve_month_start(params["month"]) + prior = [ params["prior_months"].to_i, 0 ].max + prior = [ prior, MAX_PRIOR_MONTHS ].min + + month_starts = (0..prior).map { |offset| shift_months(target_start, -offset) }.reverse + requested = month_starts.count { |start_date| Budget.budget_date_valid?(start_date, family: family) } + + months = month_starts.filter_map do |start_date| + next unless Budget.budget_date_valid?(start_date, family: family) + build_month_payload(start_date, bootstrap: start_date == target_start) + end + + result = { + currency: family.currency, + months: months + } + unavailable = requested - months.length + result[:months_unavailable] = unavailable if unavailable > 0 + result + end + + private + def build_month_payload(start_date, bootstrap:) + budget = if bootstrap + Budget.find_or_bootstrap(family, start_date: start_date, user: user) + else + budget_start, budget_end = Budget.period_for(start_date, family: family) + family.budgets.find_by(start_date: budget_start, end_date: budget_end) + end + return nil unless budget + + groups = BudgetCategory::Group.for(budget.budget_categories) + + { + month: budget.to_param, + period: { + start_date: budget.start_date, + end_date: budget.end_date + }, + is_current: budget.current?, + initialized: budget.initialized?, + totals: { + budgeted_spending: format_money(budget.budgeted_spending), + allocated_spending: format_money(budget.allocated_spending), + available_to_allocate: format_money(budget.available_to_allocate), + actual_spending: format_money(budget.actual_spending), + available_to_spend: format_money(budget.available_to_spend), + percent_of_budget_spent: format_percent(budget.initialized? ? budget.percent_of_budget_spent : 0), + overage_percent: format_percent(budget.overage_percent) + }, + income: { + expected_income: format_money(budget.expected_income), + actual_income: format_money(budget.actual_income), + remaining_expected_income: format_money((budget.expected_income || 0) - budget.actual_income) + }, + categories: groups.map { |group| serialize_group(group, include_daily_suggestion: budget.current?) } + } + end + + def serialize_group(group, include_daily_suggestion:) + parent = group.budget_category + serialize_category(parent, include_daily_suggestion: include_daily_suggestion).merge( + color: parent.category.color, + subcategories: group.budget_subcategories.map do |sub| + serialize_category(sub, include_daily_suggestion: include_daily_suggestion).merge( + inherits_parent_budget: sub.inherits_parent_budget? + ) + end + ) + end + + def serialize_category(bc, include_daily_suggestion:) + payload = { + name: bc.name, + budgeted: format_money(bc.display_budgeted_spending), + actual: format_money(bc.actual_spending), + available: format_money(bc.available_to_spend), + percent_spent: format_percent(bc.percent_of_budget_spent || 0), + status: category_status(bc) + } + + if include_daily_suggestion + suggestion = bc.suggested_daily_spending + payload[:suggested_daily_spending] = suggestion[:amount].format if suggestion + end + + payload + end + + def category_status(bc) + return "over_budget" if bc.over_budget_with_budget? + return "unbudgeted" if bc.unbudgeted_with_spending? + return "near_limit" if bc.budgeted? && bc.near_limit? + return "on_track" if bc.on_track? + "no_activity" + end + + def resolve_month_start(raw) + base = parse_month(raw) + return (base || Date.current).beginning_of_month unless family.uses_custom_month_start? + + # Match Budget.param_to_date for explicit slugs so the input round-trips with the response. + base ? Date.new(base.year, base.month, family.month_start_day) : family.custom_month_start_for(Date.current) + end + + def parse_month(raw) + return nil if raw.blank? + + # Date.strptime ignores trailing characters, so guard with strict anchors first. + fmt = case raw + when /\A\d{4}-\d{2}\z/ then "%Y-%m" + when /\A[A-Za-z]{3}-\d{4}\z/ then "%b-%Y" + end + + raise Assistant::Error, "Invalid month: #{raw}. Use YYYY-MM or MMM-YYYY." if fmt.nil? + + Date.strptime(raw, fmt) + rescue ArgumentError + raise Assistant::Error, "Invalid month: #{raw}. Use YYYY-MM or MMM-YYYY." + end + + def shift_months(date, n) + shifted = date >> n + if family.uses_custom_month_start? + family.custom_month_start_for(shifted) + else + shifted.beginning_of_month + end + end + + def format_money(value) + Money.new(value || 0, family.currency).format + end + + def format_percent(value) + number_to_percentage(value || 0, precision: 1) + end +end diff --git a/app/models/budget.rb b/app/models/budget.rb index f08f1e0d0..7e673b3fa 100644 --- a/app/models/budget.rb +++ b/app/models/budget.rb @@ -31,27 +31,24 @@ class Budget < ApplicationRecord end def budget_date_valid?(date, family:) - budget_start = if family.uses_custom_month_start? - family.custom_month_start_for(date) - else - date.beginning_of_month - end - + budget_start, _ = period_for(date, family: family) budget_start >= oldest_valid_budget_date(family) && budget_start <= latest_valid_budget_start_date(family) end + def period_for(date, family:) + if family.uses_custom_month_start? + [ family.custom_month_start_for(date), family.custom_month_end_for(date) ] + else + [ date.beginning_of_month, date.end_of_month ] + end + end + def find_or_bootstrap(family, start_date:, user: nil) return nil unless budget_date_valid?(start_date, family: family) Budget.transaction do - if family.uses_custom_month_start? - budget_start = family.custom_month_start_for(start_date) - budget_end = family.custom_month_end_for(start_date) - else - budget_start = start_date.beginning_of_month - budget_end = start_date.end_of_month - end + budget_start, budget_end = period_for(start_date, family: family) budget = Budget.find_or_create_by!( family: family, diff --git a/app/models/budget_category.rb b/app/models/budget_category.rb index 312db666a..ec4930623 100644 --- a/app/models/budget_category.rb +++ b/app/models/budget_category.rb @@ -198,11 +198,9 @@ class BudgetCategory < ApplicationRecord # Returns hash with suggested daily spending info or nil if not applicable def suggested_daily_spending return nil unless available_to_spend > 0 + return nil unless budget.current? - budget_date = budget.start_date - return nil unless budget_date.month == Date.current.month && budget_date.year == Date.current.year - - days_remaining = (budget_date.end_of_month - Date.current).to_i + 1 + days_remaining = (budget.end_date - Date.current).to_i + 1 return nil unless days_remaining > 0 { diff --git a/test/models/assistant/function/get_budget_test.rb b/test/models/assistant/function/get_budget_test.rb new file mode 100644 index 000000000..84fb10adf --- /dev/null +++ b/test/models/assistant/function/get_budget_test.rb @@ -0,0 +1,171 @@ +require "test_helper" + +class Assistant::Function::GetBudgetTest < ActiveSupport::TestCase + setup do + @user = users(:family_admin) + @family = @user.family + @function = Assistant::Function::GetBudget.new(@user) + end + + test "has correct name" do + assert_equal "get_budget", @function.name + end + + test "has a description" do + assert_not_empty @function.description + end + + test "is not in strict mode" do + refute @function.strict_mode? + end + + test "params_schema declares optional month and prior_months" do + schema = @function.params_schema + assert schema[:properties].key?(:month) + assert schema[:properties].key?(:prior_months) + assert_empty schema[:required] + end + + test "returns current month when no month given" do + result = @function.call({}) + + assert_equal @family.currency, result[:currency] + assert_equal 1, result[:months].length + + month = result[:months].first + assert month[:is_current] + assert_equal Date.current.beginning_of_month, month[:period][:start_date] + assert_equal Date.current.end_of_month, month[:period][:end_date] + end + + test "returns N+1 months sorted oldest first when prior_months is set" do + current_start = Date.current.beginning_of_month + 2.times do |i| + prior_start = current_start << (i + 1) + Budget.create!( + family: @family, + start_date: prior_start, + end_date: prior_start.end_of_month, + currency: @family.currency + ) + end + + result = @function.call("prior_months" => 2) + + assert_equal 3, result[:months].length + starts = result[:months].map { |m| m[:period][:start_date] } + assert_equal starts.sort, starts + assert_equal current_start, starts.last + end + + test "does not bootstrap budgets for prior_months that do not exist" do + initial_count = Budget.count + + result = @function.call("prior_months" => 3) + + assert_equal initial_count, Budget.count, "no prior budgets should be created as a side effect" + assert_equal 1, result[:months].length + assert_equal 3, result[:months_unavailable] + end + + test "clamps prior_months above MAX_PRIOR_MONTHS" do + result = @function.call("prior_months" => 99) + considered = result[:months].length + (result[:months_unavailable] || 0) + assert_operator considered, :<=, Assistant::Function::GetBudget::MAX_PRIOR_MONTHS + 1 + end + + test "accepts YYYY-MM month format" do + target = Date.current.beginning_of_month << 1 + result = @function.call("month" => target.strftime("%Y-%m")) + + assert_equal 1, result[:months].length + assert_equal target, result[:months].first[:period][:start_date] + end + + test "accepts MMM-YYYY month format" do + target = Date.current.beginning_of_month << 1 + result = @function.call("month" => target.strftime("%b-%Y").downcase) + + assert_equal target, result[:months].first[:period][:start_date] + end + + test "respects custom month_start_day so slug input roundtrips" do + @family.update!(month_start_day: 15) + target = Date.new(Date.current.year, Date.current.month, 15) - 2.months + slug = target.strftime("%b-%Y").downcase + + result = @function.call("month" => slug) + + assert_equal 1, result[:months].length + month = result[:months].first + assert_equal slug, month[:month] + assert_equal target, month[:period][:start_date] + end + + test "raises on invalid month format" do + assert_raises(Assistant::Error) do + @function.call("month" => "not-a-month") + end + end + + test "rejects month strings with trailing characters" do + [ "2026-05-01", "2026-05foo", "may-2026foo" ].each do |raw| + assert_raises(Assistant::Error, "Expected #{raw.inspect} to be rejected") do + @function.call("month" => raw) + end + end + end + + test "nests subcategories under their parent" do + result = @function.call({}) + categories = result[:months].first[:categories] + + food = categories.find { |c| c[:name] == "Food & Drink" } + assert food, "Food & Drink parent should be present" + sub_names = food[:subcategories].map { |s| s[:name] } + assert_includes sub_names, "Restaurants" + end + + test "category status reflects over_budget helper" do + budget = Budget.find_or_bootstrap(@family, start_date: Date.current.beginning_of_month, user: @user) + food_bc = budget.budget_categories.find { |bc| bc.category == categories(:food_and_drink) } + food_bc.update!(budgeted_spending: 100) + + BudgetCategory.any_instance.stubs(:actual_spending).returns(150) + + result = @function.call({}) + food = result[:months].first[:categories].find { |c| c[:name] == "Food & Drink" } + assert_equal "over_budget", food[:status] + end + + test "suggested_daily_spending omitted on non-current months" do + target = Date.current.beginning_of_month << 1 + result = @function.call("month" => target.strftime("%Y-%m")) + past = result[:months].first + + refute past[:is_current] + past[:categories].each do |cat| + refute cat.key?(:suggested_daily_spending), "Past months should not include suggested_daily_spending" + cat[:subcategories].each do |sub| + refute sub.key?(:suggested_daily_spending), "Past month subcategories should not include suggested_daily_spending" + end + end + end + + test "includes color on parent categories" do + result = @function.call({}) + result[:months].first[:categories].each do |cat| + assert cat.key?(:color), "Each parent category should expose a color" + end + end + + test "totals expose budget pacing fields" do + result = @function.call({}) + totals = result[:months].first[:totals] + + %i[budgeted_spending allocated_spending available_to_allocate actual_spending + available_to_spend percent_of_budget_spent overage_percent].each do |key| + assert totals.key?(key), "totals should include #{key}" + end + end +end diff --git a/test/models/budget_category_test.rb b/test/models/budget_category_test.rb index d6408c530..2486f9a75 100644 --- a/test/models/budget_category_test.rb +++ b/test/models/budget_category_test.rb @@ -255,4 +255,20 @@ class BudgetCategoryTest < ActiveSupport::TestCase @subcategory_inheriting_bc.stubs(:actual_spending).returns(10) assert @subcategory_inheriting_bc.visible_on_track? end + + test "suggested_daily_spending uses budget.end_date for custom month periods" do + @family.update!(month_start_day: 15) + + # Today (Jun 1) is in the calendar month after the budget period start (May 15). + # The pre-fix helper compared start_date.month to Date.current.month and returned nil here. + travel_to Date.new(2026, 6, 1) do + @budget.update!(start_date: Date.new(2026, 5, 15), end_date: Date.new(2026, 6, 14)) + @parent_budget_category.stubs(:actual_spending).returns(0) + + suggestion = @parent_budget_category.suggested_daily_spending + + assert suggestion, "expected suggested_daily_spending when current period spans calendar months" + assert_equal 14, suggestion[:days_remaining] + end + end end