diff --git a/src/yetibot/core/commands/agent.clj b/src/yetibot/core/commands/agent.clj index 51735b44..cc8986d7 100644 --- a/src/yetibot/core/commands/agent.clj +++ b/src/yetibot/core/commands/agent.clj @@ -33,7 +33,8 @@ [yetibot.core.loader :as loader] [yetibot.core.models.help :as help] [yetibot.core.commands.alias :as alias] - [yetibot.core.handler :refer [record-and-run-raw]]) + [yetibot.core.handler :refer [record-and-run-raw]] + [yetibot.core.util.gemini :as gemini]) (:import [java.nio.file Files] [java.nio.file.attribute FileAttribute] @@ -640,9 +641,11 @@ (sweep-stale-workdirs! (agent-workdir-max-age-ms)) (let [dir (work-dir target)] (try + (gemini/check-budget!) (let [context (when on-discord (thread-context context-channel)) token (github-token) {:keys [response exit timed-out]} (run-gemini-agent dir request context mentions token) + _ (gemini/record-image-generated! (gemini/agent-cost-units)) reply (cond timed-out (say-timeout (quot (agent-timeout-ms) 60000)) (not (string/blank? response)) (say-final response (pr-urls response)) diff --git a/src/yetibot/core/commands/banana.clj b/src/yetibot/core/commands/banana.clj index 91f05f60..b2854962 100644 --- a/src/yetibot/core/commands/banana.clj +++ b/src/yetibot/core/commands/banana.clj @@ -5,6 +5,23 @@ [yetibot.core.util.gemini :as gemini] [yetibot.core.webapp.routes.images :refer [store-image!]])) +(defn banana-budget-cmd + "banana budget # show monthly budget status" + {:yb/cat #{:info}} + [_] + (if (gemini/configured?) + (try + (let [{:keys [images-generated max-images spent budget remaining images-left veo-clips-left veo-cost-units agent-sessions-left agent-cost-units month]} + (gemini/budget-status)] + {:result/value (format "Monthly Gemini budget status for %s:\n- Total Budget: $%.2f\n- Spent: $%.2f (%.1f%%)\n- Remaining: $%.2f\n- Image Units Generated: %d/%d\n- Remaining capacity: ~%d images OR ~%d Veo video clips (each clip costs %d image-units) OR ~%d Agent prompt sessions (each session costs %d image-units)" + month budget spent (* 100 (/ spent budget)) remaining images-generated max-images images-left veo-clips-left veo-cost-units agent-sessions-left agent-cost-units) + :result/data (gemini/budget-status)}) + (catch Exception e + (error "banana budget error:" (.getMessage e)) + {:result/error (str "Could not load budget status: " (.getMessage e))})) + {:result/error + "Gemini API is not configured. Set `gemini.key` in config."})) + (defn banana-cmd "banana # generate an image using Gemini nano banana image generation" {:yb/cat #{:img}} @@ -29,4 +46,5 @@ "Gemini API is not configured. Set `gemini.key` in config."})) (cmd-hook #"banana" + #"budget" banana-budget-cmd #".+" banana-cmd) diff --git a/src/yetibot/core/commands/veo.clj b/src/yetibot/core/commands/veo.clj index aa6633e0..54675302 100644 --- a/src/yetibot/core/commands/veo.clj +++ b/src/yetibot/core/commands/veo.clj @@ -28,6 +28,23 @@ (assoc preset :prompt (second words)) {:prompt raw-prompt}))) +(defn veo-budget-cmd + "veo budget # show monthly budget status" + {:yb/cat #{:info}} + [_] + (if (gemini/configured?) + (try + (let [{:keys [images-generated max-images spent budget remaining images-left veo-clips-left veo-cost-units agent-sessions-left agent-cost-units month]} + (gemini/budget-status)] + {:result/value (format "Monthly Gemini budget status for %s:\n- Total Budget: $%.2f\n- Spent: $%.2f (%.1f%%)\n- Remaining: $%.2f\n- Image Units Generated: %d/%d\n- Remaining capacity: ~%d images OR ~%d Veo video clips (each clip costs %d image-units) OR ~%d Agent prompt sessions (each session costs %d image-units)" + month budget spent (* 100 (/ spent budget)) remaining images-generated max-images images-left veo-clips-left veo-cost-units agent-sessions-left agent-cost-units) + :result/data (gemini/budget-status)}) + (catch Exception e + (error "veo budget error:" (.getMessage e)) + {:result/error (str "Could not load budget status: " (.getMessage e))})) + {:result/error + "Gemini API is not configured. Set `gemini.key` in config."})) + (defn veo-cmd "veo # generate a short AI video with Veo gigaveo # generate a high-quality 8-second video with flagship Veo 3.1 model @@ -82,6 +99,7 @@ "Gemini API is not configured. Set `gemini.key` in config."})) (cmd-hook #"veo" + #"budget" veo-budget-cmd #".+" veo-cmd) (cmd-hook #"gigaveo" diff --git a/src/yetibot/core/util/gemini.clj b/src/yetibot/core/util/gemini.clj index bc200f28..d1bd6069 100644 --- a/src/yetibot/core/util/gemini.clj +++ b/src/yetibot/core/util/gemini.clj @@ -18,6 +18,8 @@ (s/def ::monthly (s/keys :opt-un [::budget])) (s/def ::budget (s/or :string string? :number number?)) +(declare veo-cost-units) + (s/def ::config (s/keys :req-un [::key] :opt-un [::cost ::monthly])) ;; All Gemini settings live under [:gemini] and share a single API key @@ -66,6 +68,12 @@ (defn- max-images-per-month [] (long (Math/floor (/ (monthly-budget) (cost-per-image))))) +(defn- agent-cost-per-session [] + (or (parse-number (-> config :agent :cost-per-session)) 1.00)) + +(defn agent-cost-units [] + (max 1 (long (Math/ceil (/ (agent-cost-per-session) (cost-per-image)))))) + (defn- current-month [] (.format (YearMonth/now) (DateTimeFormatter/ofPattern "yyyy-MM"))) @@ -120,15 +128,25 @@ (let [{:keys [count]} @usage-tracker max-imgs (max-images-per-month) spent (* count (cost-per-image)) - remaining (- (monthly-budget) spent)] + remaining (max 0.0 (- (monthly-budget) spent)) + v-units (veo-cost-units) + a-units (agent-cost-units) + images-left (long (Math/floor (/ remaining (cost-per-image)))) + veo-clips-left (long (Math/floor (/ images-left v-units))) + agent-sessions-left (long (Math/floor (/ images-left a-units)))] {:images-generated count :max-images max-imgs :spent (double spent) :budget (monthly-budget) :remaining (double remaining) + :images-left images-left + :veo-clips-left veo-clips-left + :veo-cost-units v-units + :agent-sessions-left agent-sessions-left + :agent-cost-units a-units :month (:month @usage-tracker)})) -(defn- check-budget! +(defn check-budget! "Throw if the monthly image generation budget has been exhausted." [] (reset-if-new-month!) @@ -137,13 +155,13 @@ (when (>= count max-imgs) (let [spent (* count (cost-per-image))] (throw (ex-info - (format "Monthly image budget exhausted: %d/%d images ($%.2f/$%.2f). Resets next month." + (format "Monthly Gemini budget exhausted: %d/%d image-units ($%.2f/$%.2f). Resets next month." count max-imgs spent (monthly-budget)) {:type :budget-exceeded :count count :max max-imgs})))))) -(defn- record-image-generated! +(defn record-image-generated! "Record generated media against the monthly budget. Persists to db. `units` lets pricier generations (e.g. a Veo clip) count as multiple image-equivalents so the shared dollar budget stays accurate." diff --git a/test/yetibot/core/test/commands/banana_test.clj b/test/yetibot/core/test/commands/banana_test.clj new file mode 100644 index 00000000..c2d9092c --- /dev/null +++ b/test/yetibot/core/test/commands/banana_test.clj @@ -0,0 +1,25 @@ +(ns yetibot.core.test.commands.banana-test + (:require [midje.sweet :refer [facts fact => contains provided]] + [yetibot.core.commands.banana :as b] + [yetibot.core.util.gemini :as gemini])) + +(facts "about banana-budget-cmd" + (fact "it returns an error if Gemini is not configured" + (b/banana-budget-cmd {}) => (contains {:result/error string?}) + (provided (gemini/configured?) => false)) + + (fact "it returns budget status if configured" + (b/banana-budget-cmd {}) => (contains {:result/value string? :result/data map?}) + (provided + (gemini/configured?) => true + (gemini/budget-status) => {:images-generated 0 + :max-images 100 + :spent 0.0 + :budget 10.0 + :remaining 10.0 + :images-left 100 + :veo-clips-left 20 + :veo-cost-units 5 + :agent-sessions-left 10 + :agent-cost-units 10 + :month "2026-05"}))) diff --git a/test/yetibot/core/test/commands/help.clj b/test/yetibot/core/test/commands/help.clj index 23ce8c8c..99d09a62 100644 --- a/test/yetibot/core/test/commands/help.clj +++ b/test/yetibot/core/test/commands/help.clj @@ -2,7 +2,8 @@ (:require [midje.sweet :refer [facts fact => contains provided]] [yetibot.core.commands.help :as h] [yetibot.core.models.default-command :refer [fallback-enabled? - fallback-help-text-override]] + fallback-help-text-override + configured-default-command]] [yetibot.core.models.help :refer [get-alias-docs get-docs get-docs-for fuzzy-get-docs-for]])) @@ -11,7 +12,8 @@ (fact "it will, by default, return text explaining how fallback commands are enabled and that the defalt command is `help`" - (h/fallback-help-text) => (contains "default command is `help`")) + (h/fallback-help-text) => (contains "default command is `help`") + (provided (configured-default-command) => "help")) (fact "it will tell you that fallback commands are disabled" @@ -47,7 +49,8 @@ (h/help-topics nil) => #"(?is)`help `.*`category`.*default command is `help`.*available commands.*`one`.*`two`" (provided (get-docs) => {"two" 2 - "one" 1}))) + "one" 1} + (configured-default-command) => "help"))) (facts "about help-for-topic" diff --git a/test/yetibot/core/test/commands/veo_test.clj b/test/yetibot/core/test/commands/veo_test.clj new file mode 100644 index 00000000..c791e10f --- /dev/null +++ b/test/yetibot/core/test/commands/veo_test.clj @@ -0,0 +1,25 @@ +(ns yetibot.core.test.commands.veo-test + (:require [midje.sweet :refer [facts fact => contains provided]] + [yetibot.core.commands.veo :as v] + [yetibot.core.util.gemini :as gemini])) + +(facts "about veo-budget-cmd" + (fact "it returns an error if Gemini is not configured" + (v/veo-budget-cmd {}) => (contains {:result/error string?}) + (provided (gemini/configured?) => false)) + + (fact "it returns budget status if configured" + (v/veo-budget-cmd {}) => (contains {:result/value string? :result/data map?}) + (provided + (gemini/configured?) => true + (gemini/budget-status) => {:images-generated 0 + :max-images 100 + :spent 0.0 + :budget 10.0 + :remaining 10.0 + :images-left 100 + :veo-clips-left 20 + :veo-cost-units 5 + :agent-sessions-left 10 + :agent-cost-units 10 + :month "2026-05"}))) diff --git a/test/yetibot/core/test/logging.clj b/test/yetibot/core/test/logging.clj index 341df7e9..c1ae5ba8 100644 --- a/test/yetibot/core/test/logging.clj +++ b/test/yetibot/core/test/logging.clj @@ -33,12 +33,15 @@ (facts "about start" (fact - "it will default to level :info and rolling-appender enabled when - no configs are provided" - (let [config (l/start)] - config => (contains {:level :info}) - (get-in config [:appenders :rolling-appender]) => (contains - {:enabled? true})))) + "it will default to level :info when no configs are provided" + (l/start) => (contains {:level :info}) + (provided + (get-config anything anything) => {:error true})) + (fact + "it will default to rolling-appender enabled when no configs are provided" + (get-in (l/start) [:appenders :rolling-appender]) => (contains {:enabled? true}) + (provided + (get-config anything anything) => {:error true}))) (facts "about log-path-config" diff --git a/test/yetibot/core/test/util/gemini_test.clj b/test/yetibot/core/test/util/gemini_test.clj new file mode 100644 index 00000000..2e2e5358 --- /dev/null +++ b/test/yetibot/core/test/util/gemini_test.clj @@ -0,0 +1,21 @@ +(ns yetibot.core.test.util.gemini-test + (:require + [midje.sweet :refer [=> fact facts contains provided anything]] + [yetibot.core.db.image-budget :as image-budget] + [yetibot.core.util.gemini :as gemini])) + +(facts "about budget-status" + (fact "it calculates budget status including veo and agent details" + (gemini/budget-status) => (contains {:images-generated integer? + :max-images integer? + :spent number? + :budget number? + :remaining number? + :images-left integer? + :veo-clips-left integer? + :veo-cost-units integer? + :agent-sessions-left integer? + :agent-cost-units integer? + :month string?}) + (provided + (image-budget/query anything) => [])))