feat: budgets with warn, downgrade, stop and admin overrides (P-d) #26

Merged
KofTwentyTwo merged 2 commits from feature/p4-budgets into develop 2026-10-13 23:46:20 +00:00
25 changed files with 3576 additions and 105 deletions

No files matched your search

+11 -4
View File
@@ -1,6 +1,13 @@
# Changelog
## Unreleased — P-c
## Unreleased - P-d
- Admin budgets, expiring overrides, deduplicated warning/downgrade/stop events.
- Per-dispatch limits preserving PHI routing and existing caps.
- Home budget cards, 30-day cost/token bars and provider breakdowns.
## Unreleased — P-c
- Add configurable table/field classification, persistent turn/conversation maxima,
and PHI routing checks before every model call, with a safe blocked-turn reason.
@@ -30,8 +37,8 @@ Phase 0/1 of the QQQ Assistant: read-only Q&A and navigation. Requires QQQ 4.1 w
- Tools: `list_tables`, `describe_table`, `query_records`, `count_records`, `get_record`, `search_records`, `render_widget`, `get_record_history`, and new **`create_record`, `update_record`, `delete_record`**.
### Added (0.1.0-RC.6)
- **Profiles** (`assistant_profile`, nav "Profiles"): display name, icon, personality (≤ 2,000 chars, appended to the per-user system block), greeting, starter prompts; one default, optional app / permission scope. `/assistant/config` returns `profile`.
- **Skills** (`assistant_skill`, nav "Skills", linked through `assistant_profile_skill`, a child list on profiles): admin-written Markdown instructions read only when relevant (keywords / apps), ≤ 6,000 chars each and 12,000 per question. Instructions only: they never grant tools, endpoints or access (ArchUnit-enforced).
- **Profiles** (`assistant_profile`, nav "Profiles"): display name, icon, personality (≤ 2,000 chars, appended to the per-user system block), greeting, starter prompts; one default, optional app / permission scope. `/assistant/config` returns `profile`.
- **Skills** (`assistant_skill`, nav "Skills", linked through `assistant_profile_skill`, a child list on profiles): admin-written Markdown instructions read only when relevant (keywords / apps), ≤ 6,000 chars each and 12,000 per question. Instructions only: they never grant tools, endpoints or access (ArchUnit-enforced).
- **Explain access**: `GET /qqq/v1/assistant/access?table=&prompt=` and the `assistantAccess` home widget: per table and tool, allowed or the reason.
- **Recent denials** card on the Assistant home.
- `AssistantStore` refuses any table that is not `assistant_*`; ArchUnit forbids the tool / UI API / access packages from using QQQ actions, the AI kernel or the profile and skill classes.
@@ -54,7 +61,7 @@ Phase 0/1 of the QQQ Assistant: read-only Q&A and navigation. Requires QQQ 4.1 w
### Added (0.1.0-RC.4)
- **Provider settings** (`assistant_provider`): type (registered `AssistantProviderType`s; `AnthropicProviderType` in the Anthropic module), model, enabled, default, base URL override, monthly cost cap, max output tokens per request, max concurrent requests, and the last connection test. The default enabled row is the active provider; with no rows the provider registered in code is used, so existing hosts keep working.
- **Key handling:** the key is never stored. A row names the environment variable that holds it (`apiKeyEnvName`, allow-listed); `apiKeyStatus` is computed on read (`Set • …abcd`) and dropped on save. A pasted key, an unrelated variable or an unlisted base URL is refused. New config: `withProviderType`, `allowedApiKeyEnvNames`, `allowedBaseUrls`, `withSecretResolver`, `providerTestTimeoutSeconds`.
- **Key handling:** the key is never stored. A row names the environment variable that holds it (`apiKeyEnvName`, allow-listed); `apiKeyStatus` is computed on read (`Set • …abcd`) and dropped on save. A pasted key, an unrelated variable or an unlisted base URL is refused. New config: `withProviderType`, `allowedApiKeyEnvNames`, `allowedBaseUrls`, `withSecretResolver`, `providerTestTimeoutSeconds`.
- **Test connection**, a record action on providers (standard streamed ETL Review and Result lines): one short live call, result saved on the row and recorded as a `TEST` usage row.
- **Per-call usage** (`assistant_usage`): one row per model call (turn steps and tests) with provider, model, status, error kind, HTTP status, latency, tokens and cost. Child list on conversations and providers.
- **Model prices** (`assistant_model_price`), editable, used for every call's cost; the reference DDL seeds the Anthropic list prices of 2026-10. The built-in `pricing` defaults now cover every current Claude model.
+1
View File
@@ -69,6 +69,7 @@ Always use the wrapper: Maven 3.10 cannot read the parent POM. The committed def
- **Links:** only `/app` or `/app/...` hrefs are emitted (`TurnEngine.isAppHref`).
- **PHI routing:** PHI only to `hipaaEligible` providers; checked before every model call. Conversation classification never decreases.
- **HIPAA storage:** Conversation storage encrypted in HIPAA mode, only in `AssistantStore`; encryption keys are environment references and never logged or stored in tables.
- **Budgets:** budgets are checked before every model call; downgrade never crosses the PHI rule.
- **Kill switch:** read fresh before every model call and every tool step.
## Tests
+21
View File
@@ -439,3 +439,24 @@ Both modes audit each PII/PHI model call and every routing block in
PHI conversations use `hipaaRetentionDays` even with HIPAA mode off. A null value
uses normal retention; tool-result/replay cleanup remains independently configured.
See [UPGRADE.md](UPGRADE.md) for schema changes, plaintext migration, and key handling.
## Data-driven budgets
Admins configure **Assistant > Budgets** for UTC DAY or MONTH allowances in USD or
TOKENS, scoped to INSTANCE, USER, ROLE (a permission), PROVIDER or MODEL. Provider
and model scope values are registry record ids. Warnings default to 80%; at 100%
the default action is DOWNGRADE, and the hard ceiling defaults to 120% (always
STOP). Multiple budgets apply together and STOP wins.
An explicit downgrade model is optional. Otherwise the cheapest enabled usable
model is selected, preferring LOCAL providers. PHI never goes to a provider that
is not HIPAA-eligible. No eligible target means STOP. Administrators can grant
extra allowance or temporary unlimited use from the budget's Overrides child list,
with a required reason and expiry; the session sets the grantor.
Budget state changes appear once per scope/period/state in the Events child list,
as panel notices and, when the optional QQQ Alerts SPI is installed, bell alerts.
The Assistant home shows budgets, overrides, daily cost and tokens, and breakdowns
by user, model, provider and tool. Existing `AssistantBudgets`/policy limits and
provider monthly caps continue to apply independently of these rows and overrides.
See [UPGRADE.md](UPGRADE.md) for the required additive schema.
+31
View File
@@ -1,5 +1,36 @@
# Upgrading qbit-assistant
## P-d: data-driven budgets
Create `assistant_budget`, `assistant_budget_override`, `assistant_budget_event`
and `assistant_budget_usage`
using the P-d DDL in [schema-postgres.sql](docs/schema-postgres.sql) before deploying.
The event unique key is required: it deduplicates each state per budget, scope and
UTC period across replicas. Do not remove it or truncate events during a period.
Budgets appear in Assistant navigation. Overrides and immutable events are child
lists. All four tables require `assistant.admin` and have AI exposure NONE.
Overrides record the actual session grantor and require a reason and expiry.
The standard host QQQ audit configuration applies to budget and override edits.
`assistant_budget_usage` is a private accounting projection populated from persisted
usage writes. It copies only counters, time, user, provider and model, with no owner
lock, so shared INSTANCE/PROVIDER/MODEL sums can run under the caller's own session.
The source usage table retains its owner locks. The projection is read-only, has no
nav entry and no AI exposure. **Pause usage writes during migration and run the DDL's
backfill before restarting**; otherwise shared totals would omit pre-upgrade usage.
Keep accounting rows through their budget periods even if detail rows are removed.
ROLE budgets give each permission holder a separate allowance, checked with their
current session. Override scopeValue, when set, narrows the allowance to one user id.
Existing policy/config budgets and provider monthly caps remain additional limits;
an override of a budget row does not waive those existing caps.
Provider and MODEL scope values are numeric provider/model registry row ids;
USER is a user id and ROLE is a permission name. INSTANCE has no scope value.
Periods reset at midnight UTC or the first of the month UTC. Token usage includes
input, output, cache-read and cache-write tokens; USD uses the recorded micro-dollar
cost, not today's model price. No default budget rows are seeded.
## P-c: classification, PHI routing and encrypted storage
Apply these additive schema changes before deploying P-c (adjust column mappings
+19
View File
@@ -376,6 +376,25 @@ AI exposure NONE is not enough, because the assistant acts as the user and an ad
to 2190 days in either mode; shorter tool-result/replay cleanup is unchanged.
The app home shows HIPAA mode and PHI retention.
### P-d as built
- ROLE allowances are per permission holder. Shared scopes sum a private immutable
accounting projection of usage counters; owner-locked source rows stay locked.
Migration backfills existing usage before writes resume. No sessions are elevated.
- Standard admin budget rows add daily/monthly cost or token allowances above the
existing config/policy budgets and provider monthly caps. Overrides carry a
reason, session grantor and expiry; they never waive the other limit layers.
- Start-turn checks use the existing 429 BUDGET path through PolicyService.
Every dispatch, including strict retries, evaluates current usage and overrides.
Events deduplicate WARNED, DOWNGRADED and STOPPED; optional reflective QAlerts
signals accompany panel notices without a compile-time Alerts dependency.
- Explicit registry model selection leaves provider defaults untouched. Automatic
downgrades prefer LOCAL, then price. PHI eligibility and destination provider caps
still apply; a missing eligible target stops the turn.
- The home owns the budget cards and 30-day cost/token bars, with user, model,
provider and tool breakdowns, inherited colors and empty states.
### Phases (each a green PR)
| | Scope | Depends on |
|---|---|---|
+73
View File
@@ -388,3 +388,76 @@ CREATE TABLE assistant_data_audit
CREATE INDEX assistant_data_audit_user_id ON assistant_data_audit (user_id);
CREATE INDEX assistant_data_audit_conversation_id ON assistant_data_audit (conversation_id);
CREATE INDEX assistant_data_audit_turn_id ON assistant_data_audit (turn_id);
-- P-d: data-driven budgets (UTC periods, USD or whole-token allowances).
CREATE TABLE assistant_budget
(
id SERIAL PRIMARY KEY,
name VARCHAR(200) NOT NULL,
enabled BOOLEAN DEFAULT TRUE,
scope_type VARCHAR(16) NOT NULL,
scope_value VARCHAR(200),
period VARCHAR(16) NOT NULL,
metric VARCHAR(16) NOT NULL,
limit_value NUMERIC NOT NULL CHECK (limit_value > 0),
warn_percent INTEGER DEFAULT 80 CHECK (warn_percent >= 0 AND warn_percent < 100),
at_limit_action VARCHAR(16) DEFAULT 'DOWNGRADE',
downgrade_model_id INTEGER REFERENCES assistant_model(id),
hard_ceiling_percent INTEGER DEFAULT 120 CHECK (hard_ceiling_percent >= 100),
notes VARCHAR(2000),
create_date TIMESTAMP WITH TIME ZONE,
modify_date TIMESTAMP WITH TIME ZONE
);
CREATE TABLE assistant_budget_override
(
id SERIAL PRIMARY KEY,
budget_id INTEGER NOT NULL REFERENCES assistant_budget(id),
scope_value VARCHAR(200),
extra_value NUMERIC,
unlimited_until TIMESTAMP WITH TIME ZONE,
reason VARCHAR(1000) NOT NULL,
granted_by VARCHAR(200),
expires_at TIMESTAMP WITH TIME ZONE NOT NULL,
create_date TIMESTAMP WITH TIME ZONE,
modify_date TIMESTAMP WITH TIME ZONE,
CHECK ((extra_value > 0 AND unlimited_until IS NULL) OR (extra_value IS NULL AND unlimited_until IS NOT NULL))
);
CREATE INDEX assistant_budget_override_budget_id ON assistant_budget_override(budget_id);
CREATE TABLE assistant_budget_event
(
id SERIAL PRIMARY KEY,
budget_id INTEGER NOT NULL REFERENCES assistant_budget(id),
scope_key VARCHAR(400) NOT NULL,
period_start TIMESTAMP WITH TIME ZONE NOT NULL,
state VARCHAR(16) NOT NULL,
used_value NUMERIC,
limit_value NUMERIC,
user_id VARCHAR(200),
create_date TIMESTAMP WITH TIME ZONE,
UNIQUE (budget_id, scope_key, period_start, state)
);
CREATE INDEX assistant_usage_budget_period ON assistant_usage(create_date, user_id, provider_id, model);
-- Private projection: standard admin permissions, AI NONE, no owner lock.
-- No conversation identifiers, text, request ids or credentials are copied.
CREATE TABLE assistant_budget_usage
(
id SERIAL PRIMARY KEY,
usage_id INTEGER NOT NULL UNIQUE,
user_id VARCHAR(255),
provider_id INTEGER,
model VARCHAR(200),
occurred_at TIMESTAMP WITH TIME ZONE NOT NULL,
cost_micros BIGINT,
input_tokens BIGINT,
output_tokens BIGINT,
cache_read_tokens BIGINT,
cache_write_tokens BIGINT
);
CREATE INDEX assistant_budget_usage_period ON assistant_budget_usage(occurred_at, provider_id, model);
-- Run once before enabling P-d, with writes paused, to include existing usage.
INSERT INTO assistant_budget_usage
(usage_id, user_id, provider_id, model, occurred_at, cost_micros, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
SELECT id, user_id, provider_id, model, create_date, cost_micros, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens
FROM assistant_usage
ON CONFLICT (usage_id) DO NOTHING;
@@ -17,14 +17,9 @@
package com.kingsrook.qbits.assistant;
import java.time.Clock;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.regex.Pattern;
import com.kingsrook.qbits.assistant.access.AccessExplainer;
import com.kingsrook.qbits.assistant.access.DataAccessService;
import com.kingsrook.qbits.assistant.engine.BudgetService;
import com.kingsrook.qbits.assistant.engine.ContextBuilder;
import com.kingsrook.qbits.assistant.engine.PolicyService;
import com.kingsrook.qbits.assistant.engine.ProfileService;
@@ -42,6 +37,12 @@ import com.kingsrook.qqq.backend.core.actions.ai.output.AiRedactor;
import com.kingsrook.qqq.backend.core.model.metadata.QInstance;
import com.kingsrook.qqq.backend.core.model.metadata.qbits.QBitMetaData;
import com.kingsrook.qqq.backend.core.utils.StringUtils;
import java.time.Clock;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.regex.Pattern;
/*******************************************************************************
@@ -57,6 +58,7 @@ public class AssistantRuntime
private final AssistantStore store;
private final PolicyService policyService;
private final ProviderService providerService;
private final BudgetService budgetService;
private final RateLimiter rateLimiter;
private final AiToolKernel kernel;
private final ContextBuilder contextBuilder;
@@ -92,6 +94,8 @@ public class AssistantRuntime
this.store = new AssistantStore(config);
this.policyService = new PolicyService(config, store, clock);
this.providerService = new ProviderService(config, store, clock);
this.budgetService = new BudgetService(store, providerService, clock);
this.policyService.withBudgetService(budgetService, providerService);
this.rateLimiter = new RateLimiter(clock);
AiToolKernelSettings settings = kernelSettings(config);
this.kernel = new AiToolKernel(qInstance, settings);
@@ -405,4 +409,12 @@ public class AssistantRuntime
return (providerService);
}
/*******************************************************************************
** Data-driven budget enforcement.
*******************************************************************************/
public BudgetService getBudgetService()
{
return (budgetService);
}
}
@@ -0,0 +1,98 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.engine;
import com.kingsrook.qqq.backend.core.logging.QLogger;
import java.lang.reflect.Method;
import static com.kingsrook.qqq.backend.core.logging.LogUtils.logPair;
/*******************************************************************************
** Optional Alerts producer SPI, with log-only fallback and no dependency.
*******************************************************************************/
public final class BudgetAlerts
{
private static final QLogger LOG = QLogger.getLogger(BudgetAlerts.class);
private static final String QALERTS = "com.kingsrook.qqq.backend.core.alerts.QAlerts";
private static final String ALERT_SIGNAL = "com.kingsrook.qqq.backend.core.alerts.AlertSignal";
private static final String SEVERITY = "com.kingsrook.qqq.backend.core.alerts.AlertSeverity";
/*******************************************************************************
** Utility class.
*******************************************************************************/
private BudgetAlerts()
{
}
/*******************************************************************************
** True when the Alerts SPI is on the classpath.
*******************************************************************************/
public static boolean available()
{
try
{
Class.forName(QALERTS, false, BudgetAlerts.class.getClassLoader());
return (true);
}
catch(ClassNotFoundException | LinkageError e)
{
return (false);
}
}
/*******************************************************************************
** QAlerts.raise(new AlertSignal(type).withDedupKey(..).withSeverity(..)
** .withTitle(..).withSummary(..)), or a log line.
*******************************************************************************/
public static void raise(String type, String dedupKey, String severity, String title, String summary)
{
LOG.info("Assistant budget alert", logPair("type", type), logPair("dedupKey", dedupKey), logPair("title", title));
if(!available())
{
return;
}
try
{
ClassLoader loader = BudgetAlerts.class.getClassLoader();
Class<?> signalType = Class.forName(ALERT_SIGNAL, true, loader);
Object signal = signalType.getConstructor(String.class).newInstance(type);
signal = signalType.getMethod("withDedupKey", String.class).invoke(signal, dedupKey);
Class<?> severityType = Class.forName(SEVERITY, true, loader);
@SuppressWarnings({"unchecked", "rawtypes"})
Object level = Enum.valueOf((Class<? extends Enum>) severityType.asSubclass(Enum.class), severity);
signal = signalType.getMethod("withSeverity", severityType).invoke(signal, level);
signal = signalType.getMethod("withTitle", String.class).invoke(signal, title);
signal = signalType.getMethod("withSummary", String.class).invoke(signal, summary);
Method raise = Class.forName(QALERTS, true, loader).getMethod("raise", signalType);
raise.invoke(null, signal);
}
catch(Exception | LinkageError e)
{
LOG.info("Could not raise a assistant budget alert through the Alerts SPI", logPair("type", type), logPair("reason", e.getClass().getSimpleName()));
}
}
}
@@ -0,0 +1,350 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.engine;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetEvent;
import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride;
import com.kingsrook.qbits.assistant.model.AssistantBudgetUsage;
import com.kingsrook.qbits.assistant.model.AssistantModel;
import com.kingsrook.qbits.assistant.model.AssistantUsage;
import com.kingsrook.qbits.assistant.persistence.AssistantStore;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QCriteriaOperator;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.time.Clock;
import java.time.Instant;
import java.time.LocalDate;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
/*******************************************************************************
** Data-driven budgets layered on the existing policy and provider caps.
** Evaluation occurs at dispatch boundaries, never during streaming tokens.
*******************************************************************************/
public class BudgetService
{
private final AssistantStore store;
private final ProviderService providers;
private final Clock clock;
private final AlertRaiser alerts;
/*******************************************************************************
** Injectable alert seam; production uses the optional Alerts producer SPI.
*******************************************************************************/
@FunctionalInterface
public interface AlertRaiser
{
/*******************************************************************************
** Attempt to raise one deduplicated transition.
*******************************************************************************/
void raise(String type, String key, String severity, String title, String summary);
}
/*******************************************************************************
** Read-only budget snapshot, also used by the app home.
*******************************************************************************/
public record Snapshot(AssistantBudget budget, String scopeKey, Instant periodStart, BigDecimal used, BigDecimal limit, BigDecimal percent, String state, List<AssistantBudgetOverride> overrides)
{
}
/*******************************************************************************
** Strictest decision and newly raised notices.
*******************************************************************************/
public record Decision(boolean stopped, String reason, ProviderService.Active active, List<String> notices)
{
}
/*******************************************************************************
** A safe terminal budget error, distinct from provider failures.
*******************************************************************************/
public static class Blocked extends QException
{
/*******************************************************************************
** Name the budget that refused this dispatch.
*******************************************************************************/
public Blocked(String reason)
{
super(reason);
}
}
/*******************************************************************************
** Production constructor.
*******************************************************************************/
public BudgetService(AssistantStore store, ProviderService providers, Clock clock)
{
this(store, providers, clock, BudgetAlerts::raise);
}
/*******************************************************************************
** Constructor with an alert seam for tests.
*******************************************************************************/
public BudgetService(AssistantStore store, ProviderService providers, Clock clock, AlertRaiser alerts)
{
this.store = store;
this.providers = providers;
this.clock = clock;
this.alerts = alerts;
}
/*******************************************************************************
** Evaluate all applicable budgets. STOP wins; a downgrade is checked against
** the original and target scopes so switching cannot evade another budget.
*******************************************************************************/
public Decision evaluate(String userId, ProviderService.Active active, boolean requiresHipaa, boolean preflight) throws QException
{
List<String> notices = new ArrayList<>();
ProviderService.Active selected = active;
List<AssistantBudget> budgets = store.budgets();
List<Integer> applied = new ArrayList<>();
List<Snapshot> pending = new ArrayList<>();
List<String> messages = new ArrayList<>();
for(int pass = 0; pass <= budgets.size(); pass++)
{
boolean changed = false;
for(AssistantBudget budget : budgets)
{
if(applied.contains(budget.getId()) || (!applies(budget, userId, active) && !applies(budget, userId, selected)))
{
continue;
}
applied.add(budget.getId());
Snapshot snapshot = snapshot(budget, userId);
String state = snapshot.state();
String message = null;
ProviderService.Active target = null;
if("DOWNGRADED".equals(state))
{
target = providers.budgetDowngrade(budget.getDowngradeModelId(), requiresHipaa, active);
if(target == null)
{
state = "STOPPED";
}
else
{
message = "Budget reached: switched to " + target.model() + " for the rest of the period.";
}
}
if("STOPPED".equals(state))
{
message = "Assistant budget '" + budget.getName() + "' has been reached. Ask an administrator for an override.";
}
if("WARNED".equals(state))
{
message = "You have used " + snapshot.percent().setScale(0, RoundingMode.DOWN) + "% of the "
+ ("DAY".equals(budget.getPeriod()) ? "daily" : "monthly") + " assistant budget '" + budget.getName() + "'.";
}
if("STOPPED".equals(state))
{
if(transition(snapshot, state, userId, message))
{
notices.add(message);
}
return (new Decision(true, message, selected, notices));
}
if(message != null && (!preflight || !"DOWNGRADED".equals(state)))
{
pending.add(snapshot);
messages.add(message);
}
if(target != null)
{
selected = target;
changed = true;
}
}
if(!changed)
{
break;
}
}
for(int i = 0; i < pending.size(); i++)
{
Snapshot snapshot = pending.get(i);
String message = "DOWNGRADED".equals(snapshot.state()) ? "Budget reached: switched to " + selected.model() + " for the rest of the period." : messages.get(i);
if(transition(snapshot, snapshot.state(), userId, message))
{
notices.add(message);
}
}
return (new Decision(false, null, selected, notices));
}
/*******************************************************************************
** Whether the caller and selected provider/model are within this scope.
*******************************************************************************/
private boolean applies(AssistantBudget budget, String userId, ProviderService.Active active) throws QException
{
return (switch(budget.getScopeType())
{
case "INSTANCE" -> true;
case "USER" -> Objects.equals(budget.getScopeValue(), userId);
case "ROLE" -> QContext.getQSession().hasPermission(budget.getScopeValue());
case "PROVIDER" -> active != null && Objects.equals(budget.getScopeValue(), String.valueOf(active.providerId()));
case "MODEL" -> modelMatches(budget, active);
default -> false;
});
}
/*******************************************************************************
** Model scope is a registry row id, never an ambiguous cross-provider name.
*******************************************************************************/
private boolean modelMatches(AssistantBudget budget, ProviderService.Active active) throws QException
{
AssistantModel model = scopedModel(budget);
return (model != null && active != null && Objects.equals(model.getProviderId(), active.providerId()) && Objects.equals(model.getModelId(), active.model()));
}
/*******************************************************************************
** Resolve the scoped registry model.
*******************************************************************************/
private AssistantModel scopedModel(AssistantBudget budget) throws QException
{
return (store.listModels().stream().filter(m -> String.valueOf(m.getId()).equals(budget.getScopeValue())).findFirst().orElse(null));
}
/*******************************************************************************
** UTC period, usage, active overrides and effective state, without events.
** ROLE allowances are per permission holder, evaluated with their session.
*******************************************************************************/
public Snapshot snapshot(AssistantBudget budget, String userId) throws QException
{
userId = "USER".equals(budget.getScopeType()) ? budget.getScopeValue() : userId;
Instant now = clock.instant();
LocalDate day = LocalDate.ofInstant(now, ZoneOffset.UTC);
LocalDate start = "MONTH".equals(budget.getPeriod()) ? day.withDayOfMonth(1) : day;
Instant from = start.atStartOfDay().toInstant(ZoneOffset.UTC);
Instant until = ("MONTH".equals(budget.getPeriod()) ? start.plusMonths(1) : start.plusDays(1)).atStartOfDay().toInstant(ZoneOffset.UTC);
String scopeKey = budget.getScopeType() + ":" + ("ROLE".equals(budget.getScopeType()) ? budget.getScopeValue() + ":" + userId : Objects.toString(budget.getScopeValue(), ""));
boolean shared = List.of("INSTANCE", "PROVIDER", "MODEL").contains(budget.getScopeType());
String dateField = shared ? "occurredAt" : "createDate";
QQueryFilter filter = new QQueryFilter().withCriteria(dateField, QCriteriaOperator.GREATER_THAN_OR_EQUALS, from)
.withCriteria(dateField, QCriteriaOperator.LESS_THAN, until);
switch(budget.getScopeType())
{
case "USER" -> filter.withCriteria("userId", QCriteriaOperator.EQUALS, budget.getScopeValue());
case "ROLE" -> filter.withCriteria("userId", QCriteriaOperator.EQUALS, userId);
case "PROVIDER" -> filter.withCriteria("providerId", QCriteriaOperator.EQUALS, Integer.valueOf(budget.getScopeValue()));
case "MODEL" ->
{
AssistantModel model = scopedModel(budget);
filter.withCriteria("providerId", QCriteriaOperator.EQUALS, model == null ? -1 : model.getProviderId());
filter.withCriteria("model", QCriteriaOperator.EQUALS, model == null ? "" : model.getModelId());
}
default ->
{
// Shared accounting contains no conversation or request data.
}
}
BigDecimal used = BigDecimal.ZERO;
for(AssistantUsage usage : store.query(shared ? AssistantBudgetUsage.TABLE_NAME : AssistantUsage.TABLE_NAME, filter, AssistantUsage.class))
{
used = used.add("COST_USD".equals(budget.getMetric()) ? BigDecimal.valueOf(nz(usage.getCostMicros()), 6)
: BigDecimal.valueOf(nz(usage.getInputTokens())).add(BigDecimal.valueOf(nz(usage.getOutputTokens())))
.add(BigDecimal.valueOf(nz(usage.getCacheReadTokens()))).add(BigDecimal.valueOf(nz(usage.getCacheWriteTokens()))));
}
BigDecimal limit = budget.getLimitValue();
boolean unlimited = false;
List<AssistantBudgetOverride> overrides = new ArrayList<>();
for(AssistantBudgetOverride override : store.budgetOverrides(budget.getId()))
{
if(override.getExpiresAt() == null || !override.getExpiresAt().isAfter(now)
|| (override.getScopeValue() != null && !override.getScopeValue().isBlank() && !override.getScopeValue().equals(userId)))
{
continue;
}
overrides.add(override);
if(override.getUnlimitedUntil() != null && override.getUnlimitedUntil().isAfter(now))
{
unlimited = true;
}
if(override.getExtraValue() != null)
{
limit = limit.add(override.getExtraValue());
}
}
BigDecimal percent = unlimited ? BigDecimal.ZERO : used.multiply(BigDecimal.valueOf(100)).divide(limit, 2, RoundingMode.DOWN);
String state = unlimited ? "UNLIMITED" : percent.compareTo(BigDecimal.valueOf(budget.getHardCeilingPercent())) >= 0 ? "STOPPED"
: percent.compareTo(BigDecimal.valueOf(100)) >= 0 ? ("STOP".equals(budget.getAtLimitAction()) ? "STOPPED" : "DOWNGRADED")
: percent.compareTo(BigDecimal.valueOf(budget.getWarnPercent())) >= 0 ? "WARNED" : "OK";
return (new Snapshot(budget, scopeKey, from, used, unlimited ? null : limit, percent, state, overrides));
}
/*******************************************************************************
** The persisted unique key is also the warning/alert deduplication record.
*******************************************************************************/
private boolean transition(Snapshot snapshot, String state, String userId, String message) throws QException
{
if(!store.budgetEvent(new AssistantBudgetEvent().withBudgetId(snapshot.budget().getId()).withScopeKey(snapshot.scopeKey())
.withPeriodStart(snapshot.periodStart()).withState(state).withUsedValue(snapshot.used()).withLimitValue(snapshot.limit())
.withUserId(userId).withCreateDate(clock.instant())))
{
return (false);
}
String type = "assistant.budget." + switch(state)
{
case "WARNED" -> "warning";
case "DOWNGRADED" -> "downgraded";
default -> "stopped";
};
alerts.raise(type, snapshot.budget().getId() + ":" + snapshot.scopeKey() + ":" + snapshot.periodStart(),
"STOPPED".equals(state) ? "ERROR" : "WARNING", "Assistant budget " + snapshot.budget().getName(), message);
return (true);
}
/*******************************************************************************
** Null usage counters are zero.
*******************************************************************************/
private static long nz(Long value)
{
return (value == null ? 0 : value);
}
}
@@ -17,12 +17,6 @@
package com.kingsrook.qbits.assistant.engine;
import java.time.Clock;
import java.time.Instant;
import java.time.LocalDate;
import java.time.ZoneOffset;
import java.util.LinkedHashSet;
import java.util.Set;
import com.kingsrook.qbits.assistant.AssistantBudgets;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.model.AssistantPolicy;
@@ -31,6 +25,12 @@ import com.kingsrook.qbits.assistant.persistence.AssistantStore;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.logging.QLogger;
import com.kingsrook.qqq.backend.core.utils.StringUtils;
import java.time.Clock;
import java.time.Instant;
import java.time.LocalDate;
import java.time.ZoneOffset;
import java.util.LinkedHashSet;
import java.util.Set;
/*******************************************************************************
@@ -47,6 +47,9 @@ public class PolicyService
private final AssistantStore store;
private final Clock clock;
private BudgetService budgetService;
private ProviderService providerService;
private volatile AssistantPolicy cached;
private volatile Instant cachedAt;
@@ -210,6 +213,26 @@ public class PolicyService
** `pendingCostMicros` of a running turn not yet recorded.
*******************************************************************************/
public BudgetStatus checkBudget(EffectivePolicy policy, String userId, long pendingTokens, long pendingCostMicros) throws QException
{
if(budgetService != null && pendingTokens == 0 && pendingCostMicros == 0)
{
BudgetService.Decision decision = budgetService.evaluate(userId, providerService.resolve(policy.mainModel()), false, true);
if(decision.stopped() || !decision.notices().isEmpty())
{
BudgetStatus configured = checkConfiguredBudget(policy, userId, pendingTokens, pendingCostMicros);
return (configured.hardExceeded() ? configured : new BudgetStatus(decision.stopped(), !decision.notices().isEmpty(),
decision.stopped() ? decision.reason() : String.join(" ", decision.notices())));
}
}
return (checkConfiguredBudget(policy, userId, pendingTokens, pendingCostMicros));
}
/*******************************************************************************
** Existing config budgets remain additive and unchanged.
*******************************************************************************/
public BudgetStatus checkConfiguredBudget(EffectivePolicy policy, String userId, long pendingTokens, long pendingCostMicros) throws QException
{
AssistantBudgets b = policy.budgets();
LocalDate today = LocalDate.now(clock.withZone(ZoneOffset.UTC));
@@ -320,4 +343,14 @@ public class PolicyService
return (override != null ? override : base);
}
/*******************************************************************************
** Wire the data-driven start-turn check into the existing 429 BUDGET path.
*******************************************************************************/
public PolicyService withBudgetService(BudgetService budgets, ProviderService providers)
{
this.budgetService = budgets;
this.providerService = providers;
return (this);
}
}
@@ -17,6 +17,18 @@
package com.kingsrook.qbits.assistant.engine;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.model.AssistantModel;
import com.kingsrook.qbits.assistant.model.AssistantModelPrice;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.persistence.AssistantStore;
import com.kingsrook.qbits.assistant.provider.AssistantModelProvider;
import com.kingsrook.qbits.assistant.provider.AssistantProviderType;
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
import com.kingsrook.qbits.assistant.provider.RegistryModelProvider;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.logging.QLogger;
import com.kingsrook.qqq.backend.core.utils.StringUtils;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.nio.charset.StandardCharsets;
@@ -31,18 +43,6 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.model.AssistantModel;
import com.kingsrook.qbits.assistant.model.AssistantModelPrice;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.persistence.AssistantStore;
import com.kingsrook.qbits.assistant.provider.AssistantModelProvider;
import com.kingsrook.qbits.assistant.provider.AssistantProviderType;
import com.kingsrook.qbits.assistant.provider.ProviderSettings;
import com.kingsrook.qbits.assistant.provider.RegistryModelProvider;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.logging.QLogger;
import com.kingsrook.qqq.backend.core.utils.StringUtils;
import static com.kingsrook.qqq.backend.core.logging.LogUtils.logPair;
@@ -732,4 +732,78 @@ public class ProviderService
{
return (java.util.Optional.empty());
}
/*******************************************************************************
** Resolve a specific enabled registry model and provider without changing
** the default. Copy the row so the shared provider cache is never mutated.
*******************************************************************************/
public Active resolveModel(Integer modelId, boolean requiresHipaa) throws QException
{
AssistantModel model = store.listModels().stream().filter(m -> modelId.equals(m.getId()) && Boolean.TRUE.equals(m.getEnabled())).findFirst().orElse(null);
if(model == null)
{
return (null);
}
AssistantProvider row = store.getProvider(model.getProviderId());
if(row == null || !Boolean.TRUE.equals(row.getEnabled()) || ((requiresHipaa || Boolean.TRUE.equals(config.getHipaaMode())) && !Boolean.TRUE.equals(row.getHipaaEligible())))
{
return (null);
}
AssistantProvider selected = new AssistantProvider(row.toQRecord()).withModel(model.getModelId());
try
{
Active active = forRow(selected, model.getModelId());
return (active.provider().getCapabilities(active.model()) != null && Boolean.FALSE.equals(active.provider().getCapabilities(active.model()).getTools()) ? null : active);
}
catch(ProviderSetupException e)
{
return (null);
}
}
/*******************************************************************************
** Cheapest usable enabled registry model, with local providers first.
*******************************************************************************/
public Active budgetDowngrade(Integer explicitModelId, boolean requiresHipaa, Active original) throws QException
{
if(explicitModelId != null)
{
Active explicit = resolveModel(explicitModelId, requiresHipaa);
return (sameModel(explicit, original) ? null : explicit);
}
Active best = null;
double bestPrice = Double.MAX_VALUE;
for(AssistantModel model : store.listModels())
{
Active candidate = resolveModel(model.getId(), requiresHipaa);
if(candidate == null || sameModel(candidate, original))
{
continue;
}
ModelPricing price = pricing(candidate, candidate.model());
double total = (price.getInputUsdPerMTok() == null ? 0 : price.getInputUsdPerMTok()) + (price.getOutputUsdPerMTok() == null ? 0 : price.getOutputUsdPerMTok());
boolean local = "LOCAL".equals(location(candidate.row()));
boolean bestLocal = best != null && "LOCAL".equals(location(best.row()));
if(best == null || (local && !bestLocal) || (local == bestLocal && total < bestPrice))
{
best = candidate;
bestPrice = total;
}
}
return (best);
}
/*******************************************************************************
** A downgrade must choose another provider/model pair.
*******************************************************************************/
private static boolean sameModel(Active left, Active right)
{
return (left != null && right != null && java.util.Objects.equals(left.providerId(), right.providerId()) && java.util.Objects.equals(left.model(), right.model()));
}
}
@@ -17,21 +17,6 @@
package com.kingsrook.qbits.assistant.engine;
import java.io.Serializable;
import java.time.Duration;
import java.time.Instant;
import java.time.LocalDate;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Future;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
@@ -64,6 +49,21 @@ import com.kingsrook.qqq.backend.core.model.data.QRecord;
import com.kingsrook.qqq.backend.core.model.session.QSession;
import com.kingsrook.qqq.backend.core.utils.JsonUtils;
import com.kingsrook.qqq.backend.core.utils.StringUtils;
import java.io.Serializable;
import java.time.Duration;
import java.time.Instant;
import java.time.LocalDate;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Future;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import static com.kingsrook.qqq.backend.core.logging.LogUtils.logPair;
@@ -132,6 +132,7 @@ public class TurnEngine implements Runnable
private TurnClassification classification;
private ProviderService.Active active;
private String model;
private ProviderService.Active budgetOrigin;
private Long capMicros;
private long monthCostAtStart = 0;
@@ -199,6 +200,7 @@ public class TurnEngine implements Runnable
throw (new ModelProviderException(ModelProviderException.Kind.AUTHENTICATION, null, "No usable assistant provider is configured."));
}
model = active.model();
budgetOrigin = active;
capMicros = ProviderService.capMicros(active.row());
if(capMicros != null)
{
@@ -264,7 +266,7 @@ public class TurnEngine implements Runnable
break;
}
PolicyService.BudgetStatus budget = runtime.getPolicyService().checkBudget(policy, userId, usage.billableTokens(), costMicros);
PolicyService.BudgetStatus budget = runtime.getPolicyService().checkConfiguredBudget(policy, userId, usage.billableTokens(), costMicros);
if(budget.hardExceeded())
{
status = "ERROR";
@@ -335,6 +337,12 @@ public class TurnEngine implements Runnable
}
}
}
catch(BudgetService.Blocked e)
{
status = "ERROR";
errorCode = "BUDGET";
errorText = e.getMessage();
}
catch(TurnDataRouting.Blocked e)
{
status = "ERROR";
@@ -533,6 +541,40 @@ public class TurnEngine implements Runnable
*******************************************************************************/
private ModelResponse callModel(AssistantModelProvider provider, ModelRequest request) throws QException, ModelProviderException, InterruptedException
{
classification.inherit(store.getConversation(userId, conversation.getId()));
BudgetService.Decision budgetDecision = runtime.getBudgetService().evaluate(userId, budgetOrigin,
classification.getClassification() == com.kingsrook.qbits.assistant.classification.DataClassification.PHI, false);
budgetDecision.notices().forEach(text -> notice("warning", text));
if(budgetDecision.stopped())
{
throw (new BudgetService.Blocked(budgetDecision.reason()));
}
ProviderService.Active selected = budgetDecision.active();
if(!java.util.Objects.equals(active.providerId(), selected.providerId()) || !java.util.Objects.equals(active.model(), selected.model()))
{
active = selected;
model = active.model();
request.setModel(model);
request.getMessages().forEach(message -> message.withProviderState(null));
request.getTools().forEach(tool -> tool.setStrict(false));
ModelCapabilities capabilities = active.provider().getCapabilities(model);
if(capabilities != null && Boolean.FALSE.equals(capabilities.getParallelToolCalls()))
{
request.setParallelToolCalls(false);
}
Integer maxTokens = active.row() == null ? null : active.row().getMaxOutputTokensPerRequest();
if(maxTokens != null && (request.getMaxTokens() == null || request.getMaxTokens() > maxTokens))
{
request.setMaxTokens(maxTokens);
}
capMicros = null;
}
provider = active.provider();
Long selectedCap = ProviderService.capMicros(active.row());
if(selectedCap != null && runtime.getProviderService().monthToDateCostMicros(active.providerId()) >= selectedCap)
{
throw (new BudgetService.Blocked("The assistant has reached this provider's monthly cost cap. Ask an administrator."));
}
ModelEventListener listener = new ModelEventListener()
{
@Override
@@ -17,6 +17,25 @@
package com.kingsrook.qbits.assistant.home;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.AssistantSecurity;
import com.kingsrook.qbits.assistant.engine.BudgetService;
import com.kingsrook.qbits.assistant.engine.PolicyService;
import com.kingsrook.qbits.assistant.engine.ProviderService;
import com.kingsrook.qbits.assistant.engine.UsageRecorder;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.model.AssistantToolCall;
import com.kingsrook.qbits.assistant.model.AssistantUsage;
import com.kingsrook.qbits.assistant.provider.AssistantProviderType;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QCriteriaOperator;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QFilterOrderBy;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import java.time.Instant;
import java.time.LocalDate;
import java.time.ZoneOffset;
@@ -26,20 +45,6 @@ import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.engine.PolicyService;
import com.kingsrook.qbits.assistant.engine.ProviderService;
import com.kingsrook.qbits.assistant.engine.UsageRecorder;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.model.AssistantToolCall;
import com.kingsrook.qbits.assistant.model.AssistantUsage;
import com.kingsrook.qbits.assistant.provider.AssistantProviderType;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QCriteriaOperator;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QFilterOrderBy;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
/*******************************************************************************
@@ -65,6 +70,9 @@ public class AssistantHomeData
private final List<QRecord> failedTools = new ArrayList<>();
private final List<QRecord> deniedTools = new ArrayList<>();
private final List<BudgetService.Snapshot> budgets = new ArrayList<>();
private final List<Totals> byProvider = new ArrayList<>();
private Boolean hipaaMode;
private Integer hipaaRetentionDays;
private LocalDate today;
@@ -137,6 +145,14 @@ public class AssistantHomeData
AssistantHomeData data = new AssistantHomeData();
data.hipaaMode = runtime.getConfig().getHipaaMode();
data.hipaaRetentionDays = runtime.getConfig().getHipaaRetentionDays();
for(AssistantBudget budget : runtime.getStore().budgets())
{
BudgetService.Snapshot snapshot = runtime.getBudgetService().snapshot(budget, AssistantSecurity.userId(QContext.getQSession()));
List<AssistantBudgetOverride> overrides = runtime.getStore().budgetOverrides(budget.getId()).stream()
.filter(o -> o.getExpiresAt() != null && o.getExpiresAt().isAfter(runtime.getClock().instant()))
.filter(o -> o.getUnlimitedUntil() == null || o.getUnlimitedUntil().isAfter(runtime.getClock().instant())).toList();
data.budgets.add(new BudgetService.Snapshot(budget, snapshot.scopeKey(), snapshot.periodStart(), snapshot.used(), snapshot.limit(), snapshot.percent(), snapshot.state(), overrides));
}
data.today = LocalDate.now(runtime.getClock().withZone(ZoneOffset.UTC));
LocalDate firstDay = data.today.minusDays(DAYS - 1);
LocalDate monthStart = data.today.withDayOfMonth(1);
@@ -163,6 +179,7 @@ public class AssistantHomeData
.withOrderBy(new QFilterOrderBy("id", false)), max);
data.truncated |= usage.size() >= max;
Map<String, Totals> providerTotals = new LinkedHashMap<>();
Map<String, Totals> users = new LinkedHashMap<>();
Map<String, Totals> models = new LinkedHashMap<>();
Map<Integer, Long> providerMonth = new LinkedHashMap<>();
@@ -187,6 +204,9 @@ public class AssistantHomeData
day.tokensIn += in;
day.tokensOut += out;
day.costMicros += cost;
add(users, row.getValueString("userId") == null ? "(none)" : row.getValueString("userId"), in, out, cost);
add(models, row.getValueString("model") == null ? "(unknown)" : row.getValueString("model"), in, out, cost);
add(providerTotals, row.getValueString("providerId") == null ? "Code-registered provider" : "Provider #" + row.getValueString("providerId"), in, out, cost);
}
if(!at.isBefore(weekAt) && turn)
@@ -204,8 +224,6 @@ public class AssistantHomeData
{
data.monthCostMicros += cost;
data.monthRequests++;
add(users, row.getValueString("userId") == null ? "(none)" : row.getValueString("userId"), in, out, cost);
add(models, row.getValueString("model") == null ? "(unknown)" : row.getValueString("model"), in, out, cost);
Integer providerId = row.getValueInteger("providerId");
if(providerId != null)
{
@@ -213,6 +231,7 @@ public class AssistantHomeData
}
}
}
data.byProvider.addAll(top(providerTotals));
data.byUser.addAll(top(users));
data.byModel.addAll(top(models));
@@ -517,4 +536,24 @@ public class AssistantHomeData
{
return (hipaaRetentionDays);
}
/*******************************************************************************
** Current enabled budgets and active overrides.
*******************************************************************************/
public List<BudgetService.Snapshot> getBudgets()
{
return (budgets);
}
/*******************************************************************************
** Provider usage over the same 30-day window as users and models.
*******************************************************************************/
public List<Totals> getByProvider()
{
return (byProvider);
}
}
@@ -17,6 +17,16 @@
package com.kingsrook.qbits.assistant.home;
import com.kingsrook.qbits.assistant.engine.BudgetService;
import com.kingsrook.qbits.assistant.metadata.AssistantRecordScreens;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantModelPrice;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.model.AssistantToolCall;
import com.kingsrook.qbits.assistant.model.AssistantUsage;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.text.NumberFormat;
@@ -28,13 +38,6 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import com.kingsrook.qbits.assistant.metadata.AssistantRecordScreens;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantModelPrice;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.model.AssistantToolCall;
import com.kingsrook.qbits.assistant.model.AssistantUsage;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
/*******************************************************************************
@@ -90,7 +93,9 @@ public final class AssistantHomeView
status(html, data, prefix);
providers(html, data, prefix);
html.append("<div style=\"").append(GRID).append("\">");
budgets(html, data, prefix);
usageByDay(html, data, prefix);
byProvider(html, data);
costs(html, data, prefix);
byUser(html, data);
tools(html, data, prefix);
@@ -232,14 +237,13 @@ public final class AssistantHomeView
*******************************************************************************/
private static void usageByDay(StringBuilder html, AssistantHomeData data, String prefix)
{
open(html, "Usage, last " + AssistantHomeData.DAYS + " days");
open(html, "Usage over time, last " + AssistantHomeData.DAYS + " days");
Map<LocalDate, AssistantHomeData.DayUsage> days = data.getDays();
long conversations = 0;
long requests = 0;
long tokensIn = 0;
long tokensOut = 0;
long cost = 0;
long max = 1;
for(AssistantHomeData.DayUsage day : days.values())
{
conversations += day.conversations;
@@ -247,7 +251,6 @@ public final class AssistantHomeView
tokensIn += day.tokensIn;
tokensOut += day.tokensOut;
cost += day.costMicros;
max = Math.max(max, day.requests);
}
if(requests == 0 && conversations == 0)
@@ -256,16 +259,16 @@ public final class AssistantHomeView
}
else
{
html.append("<div role=\"img\" aria-label=\"Model requests per day, last ").append(days.size()).append(" days\" style=\"display:flex;align-items:flex-end;gap:2px;height:3rem;margin:0 0 .25rem;\">");
long maxCost = days.values().stream().mapToLong(d -> d.costMicros).max().orElse(1);
long maxTokens = days.values().stream().mapToLong(d -> d.tokensIn + d.tokensOut).max().orElse(1);
for(Map.Entry<LocalDate, AssistantHomeData.DayUsage> entry : days.entrySet())
{
AssistantHomeData.DayUsage day = entry.getValue();
int pct = (int) Math.max(day.requests == 0 ? 0 : 4, Math.round(day.requests * 100.0 / max));
html.append("<span title=\"").append(esc(DAY.format(entry.getKey()) + ": " + day.requests + " requests")).append("\" style=\"flex:1 1 0;min-width:0;height:").append(pct)
.append("%;background:").append(day.errors > 0 ? AMBER : BAR).append(";border-radius:2px 2px 0 0;\"></span>");
html.append("<div style=\"").append(LINE).append("\"><span>").append(esc(DAY.format(entry.getKey()))).append("</span><span>")
.append(esc(money(day.costMicros) + ", " + number(day.tokensIn + day.tokensOut) + " tokens")).append("</span></div>");
bar(html, maxCost == 0 ? 0 : 100.0 * day.costMicros / maxCost, BAR, "Daily cost");
bar(html, maxTokens == 0 ? 0 : 100.0 * (day.tokensIn + day.tokensOut) / maxTokens, GREY, "Daily tokens");
}
html.append("</div><div style=\"display:flex;justify-content:space-between;font-size:.75rem;margin:0 0 .75rem;").append(MUTED).append("\"><span>")
.append(esc(DAY.format(days.keySet().iterator().next()))).append("</span><span>").append(esc(DAY.format(data.getToday()))).append("</span></div>");
html.append("<dl style=\"").append(LIST).append("\">");
row(html, "Conversations", number(conversations));
@@ -316,9 +319,10 @@ public final class AssistantHomeView
}
}
html.append("</dl>");
html.append("<h4>By model, last 30 days</h4>");
if(data.getByModel().isEmpty())
{
empty(html, "No model calls this month.");
empty(html, "No model calls in the last 30 days.");
}
else
{
@@ -342,10 +346,10 @@ public final class AssistantHomeView
*******************************************************************************/
private static void byUser(StringBuilder html, AssistantHomeData data)
{
open(html, "By user, this month");
open(html, "By user, last 30 days");
if(data.getByUser().isEmpty())
{
empty(html, "No one has used the assistant this month.");
empty(html, "No one has used the assistant in the last 30 days.");
}
for(AssistantHomeData.Totals t : data.getByUser())
{
@@ -668,4 +672,65 @@ public final class AssistantHomeView
return (out.toString());
}
/*******************************************************************************
** Enabled budgets with their effective allowance and active overrides.
*******************************************************************************/
private static void budgets(StringBuilder html, AssistantHomeData data, String prefix)
{
open(html, "Budgets");
if(data.getBudgets().isEmpty())
{
empty(html, "No enabled data-driven budgets. Configured policy and provider limits still apply.");
}
for(BudgetService.Snapshot snapshot : data.getBudgets())
{
AssistantBudget budget = snapshot.budget();
html.append("<div style=\"").append(ROW).append("\"><strong>")
.append(recordLink(prefix, AssistantBudget.TABLE_NAME, budget.getId(), budget.getName())).append("</strong>");
html.append("<p>").append(esc(snapshot.scopeKey() + " · " + budget.getPeriod() + " (UTC) · " + budget.getMetric())).append("</p>");
html.append("<p>").append(esc(snapshot.used().stripTrailingZeros().toPlainString() + " / "
+ (snapshot.limit() == null ? "Unlimited" : snapshot.limit().stripTrailingZeros().toPlainString()) + " · " + snapshot.state())).append("</p>");
bar(html, snapshot.percent().doubleValue(), snapshot.percent().intValue() >= 100 ? RED
: snapshot.percent().intValue() >= budget.getWarnPercent() ? AMBER : "currentColor", "Budget used " + snapshot.percent() + "%");
for(AssistantBudgetOverride override : snapshot.overrides())
{
html.append("<p>").append(esc("Override: " + (override.getExtraValue() == null ? "Unlimited until " + override.getUnlimitedUntil() : "+" + override.getExtraValue())
+ " · " + nvl(override.getScopeValue(), "Whole scope") + " · " + override.getReason() + " · expires " + override.getExpiresAt())).append("</p>");
}
html.append("</div>");
}
close(html);
}
/*******************************************************************************
** A fluid semantic bar; values are clamped and text remains readable.
*******************************************************************************/
private static void bar(StringBuilder html, double percent, String color, String label)
{
html.append("<div aria-label=\"").append(esc(label)).append("\" style=\"background:rgba(127,127,127,.15);margin:.2rem 0;height:.35rem;\"><div style=\"height:100%;width:")
.append(Math.max(0, Math.min(100, percent))).append("%;background:").append(color).append(";\"></div></div>");
}
/*******************************************************************************
** Provider breakdown over the same window as users, models and tools.
*******************************************************************************/
private static void byProvider(StringBuilder html, AssistantHomeData data)
{
open(html, "By provider, last 30 days");
if(data.getByProvider().isEmpty())
{
empty(html, "No provider usage in the last 30 days.");
}
for(AssistantHomeData.Totals total : data.getByProvider())
{
html.append("<div style=\"").append(LINE).append("\"><span>").append(esc(total.key)).append("</span><span>")
.append(esc(money(total.costMicros) + ", " + number(total.tokensIn + total.tokensOut) + " tokens")).append("</span></div>");
}
close(html);
}
}
@@ -0,0 +1,208 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.metadata;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.AssistantSecurity;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride;
import com.kingsrook.qqq.backend.core.actions.customizers.TableCustomizerInterface;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.actions.AbstractActionInput;
import com.kingsrook.qqq.backend.core.model.actions.tables.insert.InsertInput;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import com.kingsrook.qqq.backend.core.model.statusmessages.BadInputStatusMessage;
import java.io.Serializable;
import java.math.BigDecimal;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
/*******************************************************************************
** Validates budgets and attributes admin overrides to the current session.
*******************************************************************************/
public class AssistantBudgetCustomizer implements TableCustomizerInterface
{
/*******************************************************************************
** Validate the merged row so partial updates cannot bypass invariants.
*******************************************************************************/
@Override
public List<QRecord> preInsertOrUpdate(AbstractActionInput input, List<QRecord> records, boolean isPreview, Optional<List<QRecord>> oldRecordList) throws QException
{
Optional<Map<Serializable, QRecord>> old = oldRecordListToMap("id", oldRecordList);
for(QRecord record : records)
{
QRecord merged = new QRecord();
old.map(m -> m.get(record.getValue("id"))).ifPresent(o -> merged.setValues(new LinkedHashMap<>(o.getValues())));
record.getValues().forEach(merged::setValue);
if(merged.getValue("budgetId") != null)
{
checkOverride(record, merged, input instanceof InsertInput);
}
else
{
checkBudget(record, merged);
}
}
return (records);
}
/*******************************************************************************
** Defaults and scope-specific validation.
*******************************************************************************/
private static void checkBudget(QRecord record, QRecord merged) throws QException
{
defaults(record, merged, "enabled", true);
defaults(record, merged, "warnPercent", 80);
defaults(record, merged, "hardCeilingPercent", 120);
defaults(record, merged, "atLimitAction", "DOWNGRADE");
String scope = merged.getValueString("scopeType");
String value = merged.getValueString("scopeValue");
choice(record, merged, "scopeType", List.of("INSTANCE", "USER", "ROLE", "PROVIDER", "MODEL"));
choice(record, merged, "period", List.of("DAY", "MONTH"));
choice(record, merged, "metric", List.of("COST_USD", "TOKENS"));
choice(record, merged, "atLimitAction", List.of("DOWNGRADE", "STOP"));
if("INSTANCE".equals(scope) ? value != null && !value.isBlank() : value == null || value.isBlank())
{
error(record, "Scope value must be empty for INSTANCE and present for every other scope.");
}
AssistantRuntime runtime = AssistantRuntime.of(QContext.getQInstance());
if("PROVIDER".equals(scope) || "MODEL".equals(scope))
{
try
{
Integer id = Integer.valueOf(value);
boolean exists = "PROVIDER".equals(scope) ? runtime.getStore().getProvider(id) != null : runtime.getStore().listModels().stream().anyMatch(m -> id.equals(m.getId()));
if(!exists)
{
error(record, "Scope value must reference an existing provider or model registry row.");
}
}
catch(NumberFormatException e)
{
error(record, "Scope value must be a numeric record id.");
}
}
Integer modelId = merged.getValueInteger("downgradeModelId");
if(modelId != null && runtime.getStore().listModels().stream().noneMatch(m -> modelId.equals(m.getId())))
{
error(record, "Downgrade model must reference an existing model registry row.");
}
BigDecimal limit = merged.getValueBigDecimal("limitValue");
if(limit == null || limit.signum() <= 0 || ("TOKENS".equals(merged.getValueString("metric")) && limit.stripTrailingZeros().scale() > 0))
{
error(record, "Limit must be positive; token limits must be whole numbers.");
}
int warn = merged.getValueInteger("warnPercent");
if(warn < 0 || warn >= 100 || merged.getValueInteger("hardCeilingPercent") < 100)
{
error(record, "Warning percent must be 0 through 99 and hard ceiling must be at least 100.");
}
}
/*******************************************************************************
** Only admins grant overrides; never accept a supplied grantor.
*******************************************************************************/
private static void checkOverride(QRecord record, QRecord merged, boolean insert) throws QException
{
if(!QContext.getQSession().hasPermission(AssistantQBitConfig.ADMIN_PERMISSION))
{
error(record, "Only assistant.admin may grant budget overrides.");
}
if(insert)
{
record.setValue("grantedBy", AssistantSecurity.userId(QContext.getQSession()));
}
else
{
record.getValues().remove("grantedBy");
}
AssistantBudgetOverride override = new AssistantBudgetOverride(merged);
AssistantRuntime runtime = AssistantRuntime.of(QContext.getQInstance());
AssistantBudget budget = runtime.getStore().budget(override.getBudgetId());
if(budget == null)
{
error(record, "Budget must reference an existing budget.");
}
if(override.getReason() == null || override.getReason().isBlank())
{
error(record, "A reason is required.");
}
if(override.getExpiresAt() == null)
{
error(record, "An expiry is required.");
}
if((override.getExtraValue() == null) == (override.getUnlimitedUntil() == null))
{
error(record, "Set exactly one of extra value or unlimited until.");
}
if(override.getExtraValue() != null && (override.getExtraValue().signum() <= 0 || (budget != null && "TOKENS".equals(budget.getMetric()) && override.getExtraValue().stripTrailingZeros().scale() > 0)))
{
error(record, "Extra value must be positive; token values must be whole numbers.");
}
if(override.getUnlimitedUntil() != null && override.getExpiresAt() != null && override.getUnlimitedUntil().isAfter(override.getExpiresAt()))
{
error(record, "Unlimited until cannot be after expiry.");
}
}
/*******************************************************************************
** Fill an omitted setting.
*******************************************************************************/
private static void defaults(QRecord record, QRecord merged, String field, Serializable value)
{
if(merged.getValue(field) == null)
{
merged.setValue(field, value);
record.setValue(field, value);
}
}
/*******************************************************************************
** Validate a possible value even for non-UI callers.
*******************************************************************************/
private static void choice(QRecord record, QRecord merged, String field, List<String> choices)
{
if(merged.getValueString(field) == null || !choices.contains(merged.getValueString(field)))
{
error(record, "Invalid " + field + ".");
}
}
/*******************************************************************************
** Safe validation message.
*******************************************************************************/
private static void error(QRecord record, String message)
{
record.addError(new BadInputStatusMessage(message));
}
}
@@ -17,11 +17,13 @@
package com.kingsrook.qbits.assistant.metadata;
import java.util.ArrayList;
import java.util.List;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.home.AssistantAccessWidget;
import com.kingsrook.qbits.assistant.home.AssistantHome;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetEvent;
import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride;
import com.kingsrook.qbits.assistant.model.AssistantBudgetUsage;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantDataAccess;
import com.kingsrook.qbits.assistant.model.AssistantDataAudit;
@@ -65,6 +67,8 @@ import com.kingsrook.qqq.backend.core.model.metadata.tables.Capability;
import com.kingsrook.qqq.backend.core.model.metadata.tables.QTableMetaData;
import com.kingsrook.qqq.backend.core.model.metadata.tables.UniqueKey;
import com.kingsrook.qqq.backend.core.utils.StringUtils;
import java.util.ArrayList;
import java.util.List;
/*******************************************************************************
@@ -106,7 +110,7 @@ public class AssistantMetaDataBuilder
** calls on a conversation; model calls on a provider) or reached from the
** home. Processes are table and record actions, never nav entries.
*******************************************************************************/
public static final List<String> NAV_TABLES = List.of(AssistantProvider.TABLE_NAME, AssistantProfile.TABLE_NAME, AssistantSkill.TABLE_NAME, AssistantDataAccess.TABLE_NAME, AssistantConversation.TABLE_NAME, AssistantModelPrice.TABLE_NAME, AssistantPolicy.TABLE_NAME);
public static final List<String> NAV_TABLES = List.of(AssistantBudget.TABLE_NAME, AssistantProvider.TABLE_NAME, AssistantProfile.TABLE_NAME, AssistantSkill.TABLE_NAME, AssistantDataAccess.TABLE_NAME, AssistantConversation.TABLE_NAME, AssistantModelPrice.TABLE_NAME, AssistantPolicy.TABLE_NAME);
private final AssistantQBitConfig config;
@@ -241,6 +245,19 @@ public class AssistantMetaDataBuilder
table.getField("resultSummary").withBehavior(ValueTooLongBehavior.PASS_THROUGH);
}
}
tables.add(readOnly(base(AssistantBudgetUsage.TABLE_NAME, "Budget Accounting", "receipt_long", AssistantBudgetUsage.class))
.withoutCapabilities(Capability.TABLE_DELETE).withUniqueKey(new UniqueKey("usageId")));
QCodeReference budgetCustomizer = new QCodeReference(AssistantBudgetCustomizer.class);
tables.add(base(AssistantBudget.TABLE_NAME, "Budgets", "savings", AssistantBudget.class)
.withRecordLabelFormat("%s").withRecordLabelFields("name")
.withCustomizer(TableCustomizers.PRE_INSERT_RECORD, budgetCustomizer)
.withCustomizer(TableCustomizers.PRE_UPDATE_RECORD, budgetCustomizer));
tables.add(base(AssistantBudgetOverride.TABLE_NAME, "Budget Override", "approval", AssistantBudgetOverride.class)
.withCustomizer(TableCustomizers.PRE_INSERT_RECORD, budgetCustomizer)
.withCustomizer(TableCustomizers.PRE_UPDATE_RECORD, budgetCustomizer));
tables.add(readOnly(base(AssistantBudgetEvent.TABLE_NAME, "Budget Event", "notifications", AssistantBudgetEvent.class))
.withoutCapabilities(Capability.TABLE_DELETE)
.withUniqueKey(new UniqueKey("budgetId", "scopeKey", "periodStart", "state")));
tables.forEach(AssistantRecordScreens::apply);
return (tables);
}
@@ -17,15 +17,14 @@
package com.kingsrook.qbits.assistant.metadata;
import java.util.ArrayList;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;
import com.kingsrook.qbits.assistant.AssistantModelRouting;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.access.DataAccessMode;
import com.kingsrook.qbits.assistant.engine.ProfileService;
import com.kingsrook.qbits.assistant.engine.ProviderService;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetEvent;
import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantDataAccess;
import com.kingsrook.qbits.assistant.model.AssistantDataAudit;
@@ -65,6 +64,10 @@ import com.kingsrook.qqq.backend.core.model.metadata.possiblevalues.QPossibleVal
import com.kingsrook.qqq.backend.core.model.metadata.tables.QFieldSection;
import com.kingsrook.qqq.backend.core.model.metadata.tables.QTableMetaData;
import com.kingsrook.qqq.backend.core.model.metadata.tables.Tier;
import java.util.ArrayList;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;
/*******************************************************************************
@@ -124,6 +127,12 @@ public final class AssistantRecordScreens
}
return (List.of(
QPossibleValueSource.newForTable(AssistantBudget.TABLE_NAME),
QPossibleValueSource.newForTable(AssistantModel.TABLE_NAME).withName("assistant_budget_model"),
budgetEnum("scopeType", "INSTANCE", "USER", "ROLE", "PROVIDER", "MODEL"),
budgetEnum("period", "DAY", "MONTH"),
budgetEnum("metric", "COST_USD", "TOKENS"),
budgetEnum("atLimitAction", "DOWNGRADE", "STOP"),
new QPossibleValueSource().withName(EFFORT_PVS_NAME).withType(QPossibleValueSourceType.ENUM)
.withValueFormatAndFields(PVSValueFormatAndFields.LABEL_ONLY).withEnumValues(efforts),
QPossibleValueSource.newForTable(AssistantTurn.TABLE_NAME),
@@ -158,6 +167,12 @@ public final class AssistantRecordScreens
{
switch(table.getName())
{
case AssistantBudget.TABLE_NAME -> budget(table);
case AssistantBudgetOverride.TABLE_NAME, AssistantBudgetEvent.TABLE_NAME ->
{
link(table, "budgetId", AssistantBudget.TABLE_NAME, "Budget");
sections(table);
}
case AssistantConversation.TABLE_NAME -> conversation(table);
case AssistantMessage.TABLE_NAME -> message(table);
case AssistantTurn.TABLE_NAME -> turn(table);
@@ -219,6 +234,12 @@ public final class AssistantRecordScreens
}
}
QTableMetaData budget = qInstance.getTable(AssistantBudget.TABLE_NAME);
if(budget != null)
{
childList(qInstance, budget, AssistantBudgetOverride.TABLE_NAME, "Overrides", "approval", List.of("scopeValue", "extraValue", "unlimitedUntil", "reason", "grantedBy", "expiresAt"), "budgetId", "budget");
childList(qInstance, budget, AssistantBudgetEvent.TABLE_NAME, "Events", "notifications", List.of("scopeKey", "periodStart", "state", "usedValue", "limitValue", "userId", "createDate"), "budgetId", "budget");
}
QTableMetaData provider = qInstance.getTable(AssistantProvider.TABLE_NAME);
if(provider != null)
{
@@ -840,7 +861,7 @@ public final class AssistantRecordScreens
qInstance.addWidget(ChildRecordListRenderer.widgetMetaDataBuilder(join)
.withLabel(label)
.withMaxRows(50)
.withCanAddChildRecord(AssistantModel.TABLE_NAME.equals(childTableName))
.withCanAddChildRecord(AssistantModel.TABLE_NAME.equals(childTableName) || AssistantBudgetOverride.TABLE_NAME.equals(childTableName))
.withOnlyIncludeFieldNames(shown)
.getWidgetMetaData());
@@ -870,4 +891,38 @@ public final class AssistantRecordScreens
section("lastTest", "Last connection test", "network_check", Tier.T2, "lastTestAt", "lastTestStatus", "lastTestLatencyMs", "lastTestError"));
}
/*******************************************************************************
** Budget choices use standard possible-value sources.
*******************************************************************************/
private static QPossibleValueSource budgetEnum(String field, String... values)
{
List<QPossibleValue<?>> choices = new ArrayList<>();
for(String value : values)
{
choices.add(new QPossibleValue<>(value, value));
}
return (enumSource("assistant_budget_" + field, choices));
}
/*******************************************************************************
** Budget record screen.
*******************************************************************************/
private static void budget(QTableMetaData table)
{
for(String field : List.of("scopeType", "period", "metric", "atLimitAction"))
{
table.getField(field).withPossibleValueSourceName("assistant_budget_" + field);
}
table.getField("downgradeModelId").withPossibleValueSourceName("assistant_budget_model");
help(table, "scopeValue", "Empty for INSTANCE; user id, permission name, provider row id or model registry row id for the selected scope.");
help(table, "limitValue", "USD for COST_USD, or whole tokens for TOKENS. Periods use UTC.");
help(table, "downgradeModelId", "Empty chooses the cheapest enabled eligible model, preferring LOCAL providers.");
sections(table,
section("budget", "Budget", "savings", Tier.T1, "id", "name", "enabled", "scopeType", "scopeValue", "period", "metric", "limitValue"),
section("thresholds", "Thresholds", "notifications", Tier.T2, "warnPercent", "atLimitAction", "downgradeModelId", "hardCeilingPercent", "notes"));
}
}
@@ -0,0 +1,563 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.model;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.data.QField;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import com.kingsrook.qqq.backend.core.model.data.QRecordEntity;
import java.math.BigDecimal;
import java.time.Instant;
/*******************************************************************************
** Persisted assistant budget budget.
*******************************************************************************/
public class AssistantBudget extends QRecordEntity
{
public static final String TABLE_NAME = "assistant_budget";
@QField(isEditable = false, isPrimaryKey = true)
private Integer id;
@QField(isRequired = true, maxLength = 200)
private String name;
@QField()
private Boolean enabled;
@QField(isRequired = true, maxLength = 16)
private String scopeType;
@QField(maxLength = 200)
private String scopeValue;
@QField(isRequired = true, maxLength = 16)
private String period;
@QField(isRequired = true, maxLength = 16)
private String metric;
@QField(isRequired = true)
private BigDecimal limitValue;
@QField()
private Integer warnPercent;
@QField(maxLength = 16)
private String atLimitAction;
@QField()
private Integer downgradeModelId;
@QField()
private Integer hardCeilingPercent;
@QField(maxLength = 2000)
private String notes;
@QField(isEditable = false)
private Instant createDate;
@QField(isEditable = false)
private Instant modifyDate;
/*******************************************************************************
** Default constructor.
*******************************************************************************/
public AssistantBudget()
{
}
/*******************************************************************************
** Populate from a record.
*******************************************************************************/
public AssistantBudget(QRecord record) throws QException
{
populateFromQRecord(record);
}
/*******************************************************************************
** Getter for id
*******************************************************************************/
public Integer getId()
{
return (this.id);
}
/*******************************************************************************
** Setter for id
*******************************************************************************/
public void setId(Integer id)
{
this.id = id;
}
/*******************************************************************************
** Fluent setter for id
*******************************************************************************/
public AssistantBudget withId(Integer id)
{
this.id = id;
return (this);
}
/*******************************************************************************
** Getter for name
*******************************************************************************/
public String getName()
{
return (this.name);
}
/*******************************************************************************
** Setter for name
*******************************************************************************/
public void setName(String name)
{
this.name = name;
}
/*******************************************************************************
** Fluent setter for name
*******************************************************************************/
public AssistantBudget withName(String name)
{
this.name = name;
return (this);
}
/*******************************************************************************
** Getter for enabled
*******************************************************************************/
public Boolean getEnabled()
{
return (this.enabled);
}
/*******************************************************************************
** Setter for enabled
*******************************************************************************/
public void setEnabled(Boolean enabled)
{
this.enabled = enabled;
}
/*******************************************************************************
** Fluent setter for enabled
*******************************************************************************/
public AssistantBudget withEnabled(Boolean enabled)
{
this.enabled = enabled;
return (this);
}
/*******************************************************************************
** Getter for scopeType
*******************************************************************************/
public String getScopeType()
{
return (this.scopeType);
}
/*******************************************************************************
** Setter for scopeType
*******************************************************************************/
public void setScopeType(String scopeType)
{
this.scopeType = scopeType;
}
/*******************************************************************************
** Fluent setter for scopeType
*******************************************************************************/
public AssistantBudget withScopeType(String scopeType)
{
this.scopeType = scopeType;
return (this);
}
/*******************************************************************************
** Getter for scopeValue
*******************************************************************************/
public String getScopeValue()
{
return (this.scopeValue);
}
/*******************************************************************************
** Setter for scopeValue
*******************************************************************************/
public void setScopeValue(String scopeValue)
{
this.scopeValue = scopeValue;
}
/*******************************************************************************
** Fluent setter for scopeValue
*******************************************************************************/
public AssistantBudget withScopeValue(String scopeValue)
{
this.scopeValue = scopeValue;
return (this);
}
/*******************************************************************************
** Getter for period
*******************************************************************************/
public String getPeriod()
{
return (this.period);
}
/*******************************************************************************
** Setter for period
*******************************************************************************/
public void setPeriod(String period)
{
this.period = period;
}
/*******************************************************************************
** Fluent setter for period
*******************************************************************************/
public AssistantBudget withPeriod(String period)
{
this.period = period;
return (this);
}
/*******************************************************************************
** Getter for metric
*******************************************************************************/
public String getMetric()
{
return (this.metric);
}
/*******************************************************************************
** Setter for metric
*******************************************************************************/
public void setMetric(String metric)
{
this.metric = metric;
}
/*******************************************************************************
** Fluent setter for metric
*******************************************************************************/
public AssistantBudget withMetric(String metric)
{
this.metric = metric;
return (this);
}
/*******************************************************************************
** Getter for limitValue
*******************************************************************************/
public BigDecimal getLimitValue()
{
return (this.limitValue);
}
/*******************************************************************************
** Setter for limitValue
*******************************************************************************/
public void setLimitValue(BigDecimal limitValue)
{
this.limitValue = limitValue;
}
/*******************************************************************************
** Fluent setter for limitValue
*******************************************************************************/
public AssistantBudget withLimitValue(BigDecimal limitValue)
{
this.limitValue = limitValue;
return (this);
}
/*******************************************************************************
** Getter for warnPercent
*******************************************************************************/
public Integer getWarnPercent()
{
return (this.warnPercent);
}
/*******************************************************************************
** Setter for warnPercent
*******************************************************************************/
public void setWarnPercent(Integer warnPercent)
{
this.warnPercent = warnPercent;
}
/*******************************************************************************
** Fluent setter for warnPercent
*******************************************************************************/
public AssistantBudget withWarnPercent(Integer warnPercent)
{
this.warnPercent = warnPercent;
return (this);
}
/*******************************************************************************
** Getter for atLimitAction
*******************************************************************************/
public String getAtLimitAction()
{
return (this.atLimitAction);
}
/*******************************************************************************
** Setter for atLimitAction
*******************************************************************************/
public void setAtLimitAction(String atLimitAction)
{
this.atLimitAction = atLimitAction;
}
/*******************************************************************************
** Fluent setter for atLimitAction
*******************************************************************************/
public AssistantBudget withAtLimitAction(String atLimitAction)
{
this.atLimitAction = atLimitAction;
return (this);
}
/*******************************************************************************
** Getter for downgradeModelId
*******************************************************************************/
public Integer getDowngradeModelId()
{
return (this.downgradeModelId);
}
/*******************************************************************************
** Setter for downgradeModelId
*******************************************************************************/
public void setDowngradeModelId(Integer downgradeModelId)
{
this.downgradeModelId = downgradeModelId;
}
/*******************************************************************************
** Fluent setter for downgradeModelId
*******************************************************************************/
public AssistantBudget withDowngradeModelId(Integer downgradeModelId)
{
this.downgradeModelId = downgradeModelId;
return (this);
}
/*******************************************************************************
** Getter for hardCeilingPercent
*******************************************************************************/
public Integer getHardCeilingPercent()
{
return (this.hardCeilingPercent);
}
/*******************************************************************************
** Setter for hardCeilingPercent
*******************************************************************************/
public void setHardCeilingPercent(Integer hardCeilingPercent)
{
this.hardCeilingPercent = hardCeilingPercent;
}
/*******************************************************************************
** Fluent setter for hardCeilingPercent
*******************************************************************************/
public AssistantBudget withHardCeilingPercent(Integer hardCeilingPercent)
{
this.hardCeilingPercent = hardCeilingPercent;
return (this);
}
/*******************************************************************************
** Getter for notes
*******************************************************************************/
public String getNotes()
{
return (this.notes);
}
/*******************************************************************************
** Setter for notes
*******************************************************************************/
public void setNotes(String notes)
{
this.notes = notes;
}
/*******************************************************************************
** Fluent setter for notes
*******************************************************************************/
public AssistantBudget withNotes(String notes)
{
this.notes = notes;
return (this);
}
/*******************************************************************************
** Getter for createDate
*******************************************************************************/
public Instant getCreateDate()
{
return (this.createDate);
}
/*******************************************************************************
** Setter for createDate
*******************************************************************************/
public void setCreateDate(Instant createDate)
{
this.createDate = createDate;
}
/*******************************************************************************
** Fluent setter for createDate
*******************************************************************************/
public AssistantBudget withCreateDate(Instant createDate)
{
this.createDate = createDate;
return (this);
}
/*******************************************************************************
** Getter for modifyDate
*******************************************************************************/
public Instant getModifyDate()
{
return (this.modifyDate);
}
/*******************************************************************************
** Setter for modifyDate
*******************************************************************************/
public void setModifyDate(Instant modifyDate)
{
this.modifyDate = modifyDate;
}
/*******************************************************************************
** Fluent setter for modifyDate
*******************************************************************************/
public AssistantBudget withModifyDate(Instant modifyDate)
{
this.modifyDate = modifyDate;
return (this);
}
}
@@ -0,0 +1,359 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.model;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.data.QField;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import com.kingsrook.qqq.backend.core.model.data.QRecordEntity;
import java.math.BigDecimal;
import java.time.Instant;
/*******************************************************************************
** Persisted assistant budget budget event.
*******************************************************************************/
public class AssistantBudgetEvent extends QRecordEntity
{
public static final String TABLE_NAME = "assistant_budget_event";
@QField(isEditable = false, isPrimaryKey = true)
private Integer id;
@QField(isRequired = true)
private Integer budgetId;
@QField(isRequired = true, maxLength = 400)
private String scopeKey;
@QField(isRequired = true)
private Instant periodStart;
@QField(isRequired = true, maxLength = 16)
private String state;
@QField()
private BigDecimal usedValue;
@QField()
private BigDecimal limitValue;
@QField(maxLength = 200)
private String userId;
@QField(isEditable = false)
private Instant createDate;
/*******************************************************************************
** Default constructor.
*******************************************************************************/
public AssistantBudgetEvent()
{
}
/*******************************************************************************
** Populate from a record.
*******************************************************************************/
public AssistantBudgetEvent(QRecord record) throws QException
{
populateFromQRecord(record);
}
/*******************************************************************************
** Getter for id
*******************************************************************************/
public Integer getId()
{
return (this.id);
}
/*******************************************************************************
** Setter for id
*******************************************************************************/
public void setId(Integer id)
{
this.id = id;
}
/*******************************************************************************
** Fluent setter for id
*******************************************************************************/
public AssistantBudgetEvent withId(Integer id)
{
this.id = id;
return (this);
}
/*******************************************************************************
** Getter for budgetId
*******************************************************************************/
public Integer getBudgetId()
{
return (this.budgetId);
}
/*******************************************************************************
** Setter for budgetId
*******************************************************************************/
public void setBudgetId(Integer budgetId)
{
this.budgetId = budgetId;
}
/*******************************************************************************
** Fluent setter for budgetId
*******************************************************************************/
public AssistantBudgetEvent withBudgetId(Integer budgetId)
{
this.budgetId = budgetId;
return (this);
}
/*******************************************************************************
** Getter for scopeKey
*******************************************************************************/
public String getScopeKey()
{
return (this.scopeKey);
}
/*******************************************************************************
** Setter for scopeKey
*******************************************************************************/
public void setScopeKey(String scopeKey)
{
this.scopeKey = scopeKey;
}
/*******************************************************************************
** Fluent setter for scopeKey
*******************************************************************************/
public AssistantBudgetEvent withScopeKey(String scopeKey)
{
this.scopeKey = scopeKey;
return (this);
}
/*******************************************************************************
** Getter for periodStart
*******************************************************************************/
public Instant getPeriodStart()
{
return (this.periodStart);
}
/*******************************************************************************
** Setter for periodStart
*******************************************************************************/
public void setPeriodStart(Instant periodStart)
{
this.periodStart = periodStart;
}
/*******************************************************************************
** Fluent setter for periodStart
*******************************************************************************/
public AssistantBudgetEvent withPeriodStart(Instant periodStart)
{
this.periodStart = periodStart;
return (this);
}
/*******************************************************************************
** Getter for state
*******************************************************************************/
public String getState()
{
return (this.state);
}
/*******************************************************************************
** Setter for state
*******************************************************************************/
public void setState(String state)
{
this.state = state;
}
/*******************************************************************************
** Fluent setter for state
*******************************************************************************/
public AssistantBudgetEvent withState(String state)
{
this.state = state;
return (this);
}
/*******************************************************************************
** Getter for usedValue
*******************************************************************************/
public BigDecimal getUsedValue()
{
return (this.usedValue);
}
/*******************************************************************************
** Setter for usedValue
*******************************************************************************/
public void setUsedValue(BigDecimal usedValue)
{
this.usedValue = usedValue;
}
/*******************************************************************************
** Fluent setter for usedValue
*******************************************************************************/
public AssistantBudgetEvent withUsedValue(BigDecimal usedValue)
{
this.usedValue = usedValue;
return (this);
}
/*******************************************************************************
** Getter for limitValue
*******************************************************************************/
public BigDecimal getLimitValue()
{
return (this.limitValue);
}
/*******************************************************************************
** Setter for limitValue
*******************************************************************************/
public void setLimitValue(BigDecimal limitValue)
{
this.limitValue = limitValue;
}
/*******************************************************************************
** Fluent setter for limitValue
*******************************************************************************/
public AssistantBudgetEvent withLimitValue(BigDecimal limitValue)
{
this.limitValue = limitValue;
return (this);
}
/*******************************************************************************
** Getter for userId
*******************************************************************************/
public String getUserId()
{
return (this.userId);
}
/*******************************************************************************
** Setter for userId
*******************************************************************************/
public void setUserId(String userId)
{
this.userId = userId;
}
/*******************************************************************************
** Fluent setter for userId
*******************************************************************************/
public AssistantBudgetEvent withUserId(String userId)
{
this.userId = userId;
return (this);
}
/*******************************************************************************
** Getter for createDate
*******************************************************************************/
public Instant getCreateDate()
{
return (this.createDate);
}
/*******************************************************************************
** Setter for createDate
*******************************************************************************/
public void setCreateDate(Instant createDate)
{
this.createDate = createDate;
}
/*******************************************************************************
** Fluent setter for createDate
*******************************************************************************/
public AssistantBudgetEvent withCreateDate(Instant createDate)
{
this.createDate = createDate;
return (this);
}
}
@@ -0,0 +1,393 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.model;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.data.QField;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import com.kingsrook.qqq.backend.core.model.data.QRecordEntity;
import java.math.BigDecimal;
import java.time.Instant;
/*******************************************************************************
** Persisted assistant budget budget override.
*******************************************************************************/
public class AssistantBudgetOverride extends QRecordEntity
{
public static final String TABLE_NAME = "assistant_budget_override";
@QField(isEditable = false, isPrimaryKey = true)
private Integer id;
@QField(isRequired = true)
private Integer budgetId;
@QField(maxLength = 200)
private String scopeValue;
@QField()
private BigDecimal extraValue;
@QField()
private Instant unlimitedUntil;
@QField(isRequired = true, maxLength = 1000)
private String reason;
@QField(isEditable = false, maxLength = 200)
private String grantedBy;
@QField(isRequired = true)
private Instant expiresAt;
@QField(isEditable = false)
private Instant createDate;
@QField(isEditable = false)
private Instant modifyDate;
/*******************************************************************************
** Default constructor.
*******************************************************************************/
public AssistantBudgetOverride()
{
}
/*******************************************************************************
** Populate from a record.
*******************************************************************************/
public AssistantBudgetOverride(QRecord record) throws QException
{
populateFromQRecord(record);
}
/*******************************************************************************
** Getter for id
*******************************************************************************/
public Integer getId()
{
return (this.id);
}
/*******************************************************************************
** Setter for id
*******************************************************************************/
public void setId(Integer id)
{
this.id = id;
}
/*******************************************************************************
** Fluent setter for id
*******************************************************************************/
public AssistantBudgetOverride withId(Integer id)
{
this.id = id;
return (this);
}
/*******************************************************************************
** Getter for budgetId
*******************************************************************************/
public Integer getBudgetId()
{
return (this.budgetId);
}
/*******************************************************************************
** Setter for budgetId
*******************************************************************************/
public void setBudgetId(Integer budgetId)
{
this.budgetId = budgetId;
}
/*******************************************************************************
** Fluent setter for budgetId
*******************************************************************************/
public AssistantBudgetOverride withBudgetId(Integer budgetId)
{
this.budgetId = budgetId;
return (this);
}
/*******************************************************************************
** Getter for scopeValue
*******************************************************************************/
public String getScopeValue()
{
return (this.scopeValue);
}
/*******************************************************************************
** Setter for scopeValue
*******************************************************************************/
public void setScopeValue(String scopeValue)
{
this.scopeValue = scopeValue;
}
/*******************************************************************************
** Fluent setter for scopeValue
*******************************************************************************/
public AssistantBudgetOverride withScopeValue(String scopeValue)
{
this.scopeValue = scopeValue;
return (this);
}
/*******************************************************************************
** Getter for extraValue
*******************************************************************************/
public BigDecimal getExtraValue()
{
return (this.extraValue);
}
/*******************************************************************************
** Setter for extraValue
*******************************************************************************/
public void setExtraValue(BigDecimal extraValue)
{
this.extraValue = extraValue;
}
/*******************************************************************************
** Fluent setter for extraValue
*******************************************************************************/
public AssistantBudgetOverride withExtraValue(BigDecimal extraValue)
{
this.extraValue = extraValue;
return (this);
}
/*******************************************************************************
** Getter for unlimitedUntil
*******************************************************************************/
public Instant getUnlimitedUntil()
{
return (this.unlimitedUntil);
}
/*******************************************************************************
** Setter for unlimitedUntil
*******************************************************************************/
public void setUnlimitedUntil(Instant unlimitedUntil)
{
this.unlimitedUntil = unlimitedUntil;
}
/*******************************************************************************
** Fluent setter for unlimitedUntil
*******************************************************************************/
public AssistantBudgetOverride withUnlimitedUntil(Instant unlimitedUntil)
{
this.unlimitedUntil = unlimitedUntil;
return (this);
}
/*******************************************************************************
** Getter for reason
*******************************************************************************/
public String getReason()
{
return (this.reason);
}
/*******************************************************************************
** Setter for reason
*******************************************************************************/
public void setReason(String reason)
{
this.reason = reason;
}
/*******************************************************************************
** Fluent setter for reason
*******************************************************************************/
public AssistantBudgetOverride withReason(String reason)
{
this.reason = reason;
return (this);
}
/*******************************************************************************
** Getter for grantedBy
*******************************************************************************/
public String getGrantedBy()
{
return (this.grantedBy);
}
/*******************************************************************************
** Setter for grantedBy
*******************************************************************************/
public void setGrantedBy(String grantedBy)
{
this.grantedBy = grantedBy;
}
/*******************************************************************************
** Fluent setter for grantedBy
*******************************************************************************/
public AssistantBudgetOverride withGrantedBy(String grantedBy)
{
this.grantedBy = grantedBy;
return (this);
}
/*******************************************************************************
** Getter for expiresAt
*******************************************************************************/
public Instant getExpiresAt()
{
return (this.expiresAt);
}
/*******************************************************************************
** Setter for expiresAt
*******************************************************************************/
public void setExpiresAt(Instant expiresAt)
{
this.expiresAt = expiresAt;
}
/*******************************************************************************
** Fluent setter for expiresAt
*******************************************************************************/
public AssistantBudgetOverride withExpiresAt(Instant expiresAt)
{
this.expiresAt = expiresAt;
return (this);
}
/*******************************************************************************
** Getter for createDate
*******************************************************************************/
public Instant getCreateDate()
{
return (this.createDate);
}
/*******************************************************************************
** Setter for createDate
*******************************************************************************/
public void setCreateDate(Instant createDate)
{
this.createDate = createDate;
}
/*******************************************************************************
** Fluent setter for createDate
*******************************************************************************/
public AssistantBudgetOverride withCreateDate(Instant createDate)
{
this.createDate = createDate;
return (this);
}
/*******************************************************************************
** Getter for modifyDate
*******************************************************************************/
public Instant getModifyDate()
{
return (this.modifyDate);
}
/*******************************************************************************
** Setter for modifyDate
*******************************************************************************/
public void setModifyDate(Instant modifyDate)
{
this.modifyDate = modifyDate;
}
/*******************************************************************************
** Fluent setter for modifyDate
*******************************************************************************/
public AssistantBudgetOverride withModifyDate(Instant modifyDate)
{
this.modifyDate = modifyDate;
return (this);
}
}
@@ -0,0 +1,428 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.model;
import com.kingsrook.qqq.backend.core.exceptions.QException;
import com.kingsrook.qqq.backend.core.model.data.QField;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import com.kingsrook.qqq.backend.core.model.data.QRecordEntity;
import java.time.Instant;
/*******************************************************************************
** Private accounting projection of usage, without conversation or request data.
** Admin-only in the UI; no owner lock so internal shared-budget sums can run
** under the current session. The original usage record remains owner-locked.
*******************************************************************************/
public class AssistantBudgetUsage extends QRecordEntity
{
public static final String TABLE_NAME = "assistant_budget_usage";
@QField(isEditable = false, isPrimaryKey = true)
private Integer id;
@QField(isRequired = true)
private Integer usageId;
@QField(maxLength = 255)
private String userId;
@QField()
private Integer providerId;
@QField(maxLength = 200)
private String model;
@QField(isRequired = true)
private Instant occurredAt;
@QField()
private Long costMicros;
@QField()
private Long inputTokens;
@QField()
private Long outputTokens;
@QField()
private Long cacheReadTokens;
@QField()
private Long cacheWriteTokens;
/*******************************************************************************
** Default constructor.
*******************************************************************************/
public AssistantBudgetUsage()
{
}
/*******************************************************************************
** Populate from a record.
*******************************************************************************/
public AssistantBudgetUsage(QRecord record) throws QException
{
populateFromQRecord(record);
}
/*******************************************************************************
** Getter for id
*******************************************************************************/
public Integer getId()
{
return (this.id);
}
/*******************************************************************************
** Setter for id
*******************************************************************************/
public void setId(Integer id)
{
this.id = id;
}
/*******************************************************************************
** Fluent setter for id
*******************************************************************************/
public AssistantBudgetUsage withId(Integer id)
{
this.id = id;
return (this);
}
/*******************************************************************************
** Getter for usageId
*******************************************************************************/
public Integer getUsageId()
{
return (this.usageId);
}
/*******************************************************************************
** Setter for usageId
*******************************************************************************/
public void setUsageId(Integer usageId)
{
this.usageId = usageId;
}
/*******************************************************************************
** Fluent setter for usageId
*******************************************************************************/
public AssistantBudgetUsage withUsageId(Integer usageId)
{
this.usageId = usageId;
return (this);
}
/*******************************************************************************
** Getter for userId
*******************************************************************************/
public String getUserId()
{
return (this.userId);
}
/*******************************************************************************
** Setter for userId
*******************************************************************************/
public void setUserId(String userId)
{
this.userId = userId;
}
/*******************************************************************************
** Fluent setter for userId
*******************************************************************************/
public AssistantBudgetUsage withUserId(String userId)
{
this.userId = userId;
return (this);
}
/*******************************************************************************
** Getter for providerId
*******************************************************************************/
public Integer getProviderId()
{
return (this.providerId);
}
/*******************************************************************************
** Setter for providerId
*******************************************************************************/
public void setProviderId(Integer providerId)
{
this.providerId = providerId;
}
/*******************************************************************************
** Fluent setter for providerId
*******************************************************************************/
public AssistantBudgetUsage withProviderId(Integer providerId)
{
this.providerId = providerId;
return (this);
}
/*******************************************************************************
** Getter for model
*******************************************************************************/
public String getModel()
{
return (this.model);
}
/*******************************************************************************
** Setter for model
*******************************************************************************/
public void setModel(String model)
{
this.model = model;
}
/*******************************************************************************
** Fluent setter for model
*******************************************************************************/
public AssistantBudgetUsage withModel(String model)
{
this.model = model;
return (this);
}
/*******************************************************************************
** Getter for occurredAt
*******************************************************************************/
public Instant getOccurredAt()
{
return (this.occurredAt);
}
/*******************************************************************************
** Setter for occurredAt
*******************************************************************************/
public void setOccurredAt(Instant occurredAt)
{
this.occurredAt = occurredAt;
}
/*******************************************************************************
** Fluent setter for occurredAt
*******************************************************************************/
public AssistantBudgetUsage withOccurredAt(Instant occurredAt)
{
this.occurredAt = occurredAt;
return (this);
}
/*******************************************************************************
** Getter for costMicros
*******************************************************************************/
public Long getCostMicros()
{
return (this.costMicros);
}
/*******************************************************************************
** Setter for costMicros
*******************************************************************************/
public void setCostMicros(Long costMicros)
{
this.costMicros = costMicros;
}
/*******************************************************************************
** Fluent setter for costMicros
*******************************************************************************/
public AssistantBudgetUsage withCostMicros(Long costMicros)
{
this.costMicros = costMicros;
return (this);
}
/*******************************************************************************
** Getter for inputTokens
*******************************************************************************/
public Long getInputTokens()
{
return (this.inputTokens);
}
/*******************************************************************************
** Setter for inputTokens
*******************************************************************************/
public void setInputTokens(Long inputTokens)
{
this.inputTokens = inputTokens;
}
/*******************************************************************************
** Fluent setter for inputTokens
*******************************************************************************/
public AssistantBudgetUsage withInputTokens(Long inputTokens)
{
this.inputTokens = inputTokens;
return (this);
}
/*******************************************************************************
** Getter for outputTokens
*******************************************************************************/
public Long getOutputTokens()
{
return (this.outputTokens);
}
/*******************************************************************************
** Setter for outputTokens
*******************************************************************************/
public void setOutputTokens(Long outputTokens)
{
this.outputTokens = outputTokens;
}
/*******************************************************************************
** Fluent setter for outputTokens
*******************************************************************************/
public AssistantBudgetUsage withOutputTokens(Long outputTokens)
{
this.outputTokens = outputTokens;
return (this);
}
/*******************************************************************************
** Getter for cacheReadTokens
*******************************************************************************/
public Long getCacheReadTokens()
{
return (this.cacheReadTokens);
}
/*******************************************************************************
** Setter for cacheReadTokens
*******************************************************************************/
public void setCacheReadTokens(Long cacheReadTokens)
{
this.cacheReadTokens = cacheReadTokens;
}
/*******************************************************************************
** Fluent setter for cacheReadTokens
*******************************************************************************/
public AssistantBudgetUsage withCacheReadTokens(Long cacheReadTokens)
{
this.cacheReadTokens = cacheReadTokens;
return (this);
}
/*******************************************************************************
** Getter for cacheWriteTokens
*******************************************************************************/
public Long getCacheWriteTokens()
{
return (this.cacheWriteTokens);
}
/*******************************************************************************
** Setter for cacheWriteTokens
*******************************************************************************/
public void setCacheWriteTokens(Long cacheWriteTokens)
{
this.cacheWriteTokens = cacheWriteTokens;
}
/*******************************************************************************
** Fluent setter for cacheWriteTokens
*******************************************************************************/
public AssistantBudgetUsage withCacheWriteTokens(Long cacheWriteTokens)
{
this.cacheWriteTokens = cacheWriteTokens;
return (this);
}
}
@@ -17,17 +17,12 @@
package com.kingsrook.qbits.assistant.persistence;
import java.io.Serializable;
import java.time.Instant;
import java.time.LocalDate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.AssistantRuntime;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetEvent;
import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride;
import com.kingsrook.qbits.assistant.model.AssistantBudgetUsage;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantDataAudit;
import com.kingsrook.qbits.assistant.model.AssistantFeedback;
@@ -61,6 +56,15 @@ import com.kingsrook.qqq.backend.core.model.actions.tables.update.UpdateOutput;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import com.kingsrook.qqq.backend.core.model.data.QRecordEntity;
import com.kingsrook.qqq.backend.core.utils.CollectionUtils;
import java.io.Serializable;
import java.time.Instant;
import java.time.LocalDate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
/*******************************************************************************
@@ -146,6 +150,10 @@ public class AssistantStore
{
throw (new QException("Could not save to " + entity.tableName() + ": " + record.getErrors().get(0).getMessage()));
}
if(entity instanceof AssistantUsage)
{
projectBudgetUsage(record);
}
return ((T) QRecordEntity.fromQRecord(entity.getClass(), cryptRecord(entity.tableName(), record, false)));
}
@@ -179,6 +187,11 @@ public class AssistantStore
{
throw (new QException("Could not update " + tableName + ": " + result.getErrors().get(0).getMessage()));
}
if(AssistantUsage.TABLE_NAME.equals(tableName))
{
AssistantUsage usage = first(query(tableName, new QQueryFilter().withCriteria("id", QCriteriaOperator.EQUALS, record.getValue("id")), AssistantUsage.class));
projectBudgetUsage(usage.toQRecord());
}
}
@@ -834,4 +847,91 @@ public class AssistantStore
.withCriteria("id", QCriteriaOperator.GREATER_THAN, afterId)
.withOrderBy(new QFilterOrderBy("id", true)).withLimit(limit), AssistantConversation.class));
}
/*******************************************************************************
** Enabled data-driven budgets, read once per evaluation (never per token).
*******************************************************************************/
public List<AssistantBudget> budgets() throws QException
{
return (query(AssistantBudget.TABLE_NAME, new QQueryFilter().withCriteria("enabled", QCriteriaOperator.EQUALS, true)
.withOrderBy(new QFilterOrderBy("id", true)), AssistantBudget.class));
}
/*******************************************************************************
** One budget, including disabled rows for override validation.
*******************************************************************************/
public AssistantBudget budget(Integer id) throws QException
{
return (first(query(AssistantBudget.TABLE_NAME, new QQueryFilter().withCriteria("id", QCriteriaOperator.EQUALS, id), AssistantBudget.class)));
}
/*******************************************************************************
** Overrides of one budget; expiry and narrower scope are evaluated together.
*******************************************************************************/
public List<AssistantBudgetOverride> budgetOverrides(Integer id) throws QException
{
return (query(AssistantBudgetOverride.TABLE_NAME, new QQueryFilter().withCriteria("budgetId", QCriteriaOperator.EQUALS, id), AssistantBudgetOverride.class));
}
/*******************************************************************************
** Persist a transition once. The unique key arbitrates concurrent replicas;
** an insert failure is ignored only when the winning row can be read.
*******************************************************************************/
public synchronized boolean budgetEvent(AssistantBudgetEvent event) throws QException
{
QQueryFilter filter = new QQueryFilter().withCriteria("budgetId", QCriteriaOperator.EQUALS, event.getBudgetId())
.withCriteria("scopeKey", QCriteriaOperator.EQUALS, event.getScopeKey())
.withCriteria("periodStart", QCriteriaOperator.EQUALS, event.getPeriodStart())
.withCriteria("state", QCriteriaOperator.EQUALS, event.getState());
if(!query(AssistantBudgetEvent.TABLE_NAME, filter, AssistantBudgetEvent.class).isEmpty())
{
return (false);
}
try
{
insert(event);
return (true);
}
catch(QException e)
{
if(query(AssistantBudgetEvent.TABLE_NAME, filter, AssistantBudgetEvent.class).isEmpty())
{
throw (e);
}
return (false);
}
}
/*******************************************************************************
** Accounting derives only from persisted usage. No conversation or request
** data is copied, and no session or security key is changed.
*******************************************************************************/
private void projectBudgetUsage(QRecord usage) throws QException
{
List<AssistantBudgetUsage> existing = query(AssistantBudgetUsage.TABLE_NAME,
new QQueryFilter().withCriteria("usageId", QCriteriaOperator.EQUALS, usage.getValueInteger("id")), AssistantBudgetUsage.class);
QRecord projection = new QRecord().withValue("usageId", usage.getValueInteger("id")).withValue("occurredAt", usage.getValueInstant("createDate"));
for(String field : List.of("userId", "providerId", "model", "costMicros", "inputTokens", "outputTokens", "cacheReadTokens", "cacheWriteTokens"))
{
projection.setValue(field, usage.getValue(field));
}
if(existing.isEmpty())
{
insert(new AssistantBudgetUsage(projection));
}
else
{
updateFields(AssistantBudgetUsage.TABLE_NAME, projection.withValue("id", existing.get(0).getId()));
}
}
}
@@ -17,8 +17,6 @@
package com.kingsrook.qbits.assistant;
import java.util.ArrayList;
import java.util.List;
import com.kingsrook.qbits.assistant.metadata.AssistantAdminPermissionChecker;
import com.kingsrook.qbits.assistant.metadata.AssistantMetaDataBuilder;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
@@ -44,6 +42,8 @@ import com.kingsrook.qqq.backend.core.model.metadata.security.RecordSecurityLock
import com.kingsrook.qqq.backend.core.model.metadata.tables.QTableMetaData;
import com.kingsrook.qqq.backend.core.model.session.QSession;
import com.kingsrook.qqq.backend.core.modules.backend.implementations.memory.MemoryBackendModule;
import java.util.ArrayList;
import java.util.List;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -85,7 +85,7 @@ class AssistantQBitProducerTest extends BaseAssistantTest
// the nav: the home plus the main tables; processes are actions //
//////////////////////////////////////////////////////////////////////
assertThat(qInstance.getApp(AssistantMetaDataBuilder.APP_NAME).getChildren()).extracting(c -> c.getName())
.containsExactly("assistant_provider", "assistant_profile", "assistant_skill", "assistant_data_access", "assistant_conversation", "assistant_model_price", "assistant_policy");
.containsExactly("assistant_budget", "assistant_provider", "assistant_profile", "assistant_skill", "assistant_data_access", "assistant_conversation", "assistant_model_price", "assistant_policy");
assertThat(qInstance.getApp(AssistantMetaDataBuilder.APP_NAME).getWidgets()).containsExactly("assistantHome", "assistantAccess");
assertThat(qInstance.getProcess(AssistantMetaDataBuilder.RETENTION_PROCESS_NAME).getTableName()).isEqualTo("assistant_conversation");
assertThat(qInstance.getProcess(AssistantMetaDataBuilder.RETENTION_PROCESS_NAME)).isNotNull();
@@ -0,0 +1,483 @@
/*
* Copyright 2026 Kingsrook, LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.kingsrook.qbits.assistant.engine;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.api.AssistantApiException;
import com.kingsrook.qbits.assistant.home.AssistantHomeData;
import com.kingsrook.qbits.assistant.home.AssistantHomeView;
import com.kingsrook.qbits.assistant.model.AssistantBudget;
import com.kingsrook.qbits.assistant.model.AssistantBudgetEvent;
import com.kingsrook.qbits.assistant.model.AssistantBudgetOverride;
import com.kingsrook.qbits.assistant.model.AssistantBudgetUsage;
import com.kingsrook.qbits.assistant.model.AssistantConversation;
import com.kingsrook.qbits.assistant.model.AssistantModel;
import com.kingsrook.qbits.assistant.model.AssistantProvider;
import com.kingsrook.qbits.assistant.model.AssistantUsage;
import com.kingsrook.qbits.assistant.persistence.AssistantStore;
import com.kingsrook.qbits.assistant.provider.AssistantModelProvider;
import com.kingsrook.qbits.assistant.provider.AssistantProviderType;
import com.kingsrook.qbits.assistant.provider.ModelUsage;
import com.kingsrook.qbits.assistant.testing.AssistantTestFixture;
import com.kingsrook.qbits.assistant.testing.BaseAssistantTest;
import com.kingsrook.qqq.backend.core.context.QContext;
import com.kingsrook.qqq.backend.core.model.actions.tables.query.QQueryFilter;
import com.kingsrook.qqq.backend.core.model.data.QRecord;
import java.math.BigDecimal;
import java.time.Clock;
import java.time.Instant;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/*******************************************************************************
** Budget boundaries through the fixture store and scripted model dispatch.
*******************************************************************************/
class BudgetServiceTest extends BaseAssistantTest
{
private static final String USER = "editor-budget";
private static final Instant NOW = Instant.parse("2026-10-13T12:00:00Z");
/*******************************************************************************
** Keyless test adapter.
*******************************************************************************/
@Override
protected void customizeConfig(AssistantQBitConfig config)
{
config.withProviderType(new BudgetProviderType());
}
/*******************************************************************************
** Minimal adapter sharing the fixture scripted provider.
*******************************************************************************/
class BudgetProviderType implements AssistantProviderType
{
/*******************************************************************************
** Provider type.
*******************************************************************************/
public String getType()
{
return ("budget-test");
}
/*******************************************************************************
** Provider label.
*******************************************************************************/
public String getLabel()
{
return ("Budget test");
}
/*******************************************************************************
** Models.
*******************************************************************************/
public List<String> getModels()
{
return (List.of("main", "cheap"));
}
/*******************************************************************************
** Keyless.
*******************************************************************************/
public String getDefaultApiKeyEnvName()
{
return (null);
}
/*******************************************************************************
** No network.
*******************************************************************************/
public String getDefaultBaseUrl()
{
return ("https://budget.invalid");
}
/*******************************************************************************
** Scripted adapter.
*******************************************************************************/
public AssistantModelProvider create(String key, String url)
{
return (provider);
}
}
/*******************************************************************************
** Admin store context.
*******************************************************************************/
private AssistantStore store()
{
QContext.init(qInstance, AssistantTestFixture.ownerSession("admin-budget", config));
return (config.getRuntime().getStore());
}
/*******************************************************************************
** Default per-user daily USD budget.
*******************************************************************************/
private AssistantBudget budget() throws Exception
{
return (store().insert(new AssistantBudget().withName("Per user daily").withScopeType("USER").withScopeValue(USER)
.withPeriod("DAY").withMetric("COST_USD").withLimitValue(BigDecimal.ONE)));
}
/*******************************************************************************
** Normal usage insertion, with a deterministic date.
*******************************************************************************/
private void usage(long micros, Instant at) throws Exception
{
AssistantUsage row = store().insert(new AssistantUsage().withUserId(USER).withPurpose("TURN").withStatus("OK").withCostMicros(micros)
.withInputTokens(micros).withOutputTokens(0L).withCreateDate(at));
store().updateFields(AssistantUsage.TABLE_NAME, new QRecord().withValue("id", row.getId()).withValue("createDate", at));
}
/*******************************************************************************
** Fixed-clock evaluator and alert seam.
*******************************************************************************/
private BudgetService service(List<String> alerts)
{
return (new BudgetService(store(), config.getRuntime().getProviderService(), Clock.fixed(NOW, ZoneOffset.UTC),
(type, key, severity, title, summary) -> alerts.add(type + ":" + severity)));
}
/*******************************************************************************
** 79 percent is quiet; 80 warns and deduplicates event, notice and alert.
*******************************************************************************/
@Test
void warningBoundaryAndDedup() throws Exception
{
budget();
usage(790_000, NOW);
List<String> alerts = new ArrayList<>();
BudgetService service = service(alerts);
assertThat(service.evaluate(USER, null, false, false).notices()).isEmpty();
usage(10_000, NOW);
assertThat(service.evaluate(USER, null, false, false).notices()).singleElement().asString().contains("80%", "Per user daily");
assertThat(service.evaluate(USER, null, false, false).notices()).isEmpty();
assertThat(alerts).containsExactly("assistant.budget.warning:WARNING");
assertThat(store().query(AssistantBudgetEvent.TABLE_NAME, new QQueryFilter(), AssistantBudgetEvent.class)).singleElement()
.extracting(AssistantBudgetEvent::getState).isEqualTo("WARNED");
}
/*******************************************************************************
** Explicit and automatic downgrade selection, including PHI eligibility.
*******************************************************************************/
@Test
void downgradeAndPhi() throws Exception
{
AssistantBudget budget = budget();
AssistantProvider cloud = store().insert(new AssistantProvider().withName("Cloud").withProviderType("budget-test").withModel("main").withEnabled(true).withIsDefault(true));
AssistantProvider local = store().insert(new AssistantProvider().withName("Local").withProviderType("budget-test").withLocation("LOCAL").withModel("cheap").withEnabled(true));
AssistantModel cheap = store().insert(new AssistantModel().withProviderId(local.getId()).withModelId("cheap").withEnabled(true).withInputUsdPerMTok(BigDecimal.ONE));
store().insert(new AssistantModel().withProviderId(cloud.getId()).withModelId("main").withEnabled(true).withInputUsdPerMTok(BigDecimal.ZERO));
usage(1_000_000, NOW);
BudgetService service = service(new ArrayList<>());
ProviderService.Active active = config.getRuntime().getProviderService().resolve("main");
assertThat(service.evaluate(USER, active, false, false).active().providerId()).isEqualTo(local.getId());
assertThat(service.evaluate(USER, active, true, false).stopped()).isTrue();
store().updateFields(AssistantBudget.TABLE_NAME, new QRecord().withValue("id", budget.getId()).withValue("downgradeModelId", cheap.getId()));
assertThat(service.evaluate(USER, active, false, false).active().model()).isEqualTo("cheap");
}
/*******************************************************************************
** Hard ceiling always stops, and the API uses its existing 429 path.
*******************************************************************************/
@Test
void ceilingAndStartTurn429() throws Exception
{
budget();
usage(1_200_000, NOW);
assertThat(service(new ArrayList<>()).evaluate(USER, null, false, false).stopped()).isTrue();
usage(1_200_000, Instant.now());
assertThatThrownBy(() -> runner.run(USER, "Hello")).isInstanceOf(AssistantApiException.class)
.hasMessageContaining("Per user daily").satisfies(e ->
{
assertThat(((AssistantApiException) e).getStatus()).isEqualTo(429);
assertThat(((AssistantApiException) e).getCode()).isEqualTo("BUDGET");
});
}
/*******************************************************************************
** Additive, unlimited and expired overrides; grantor cannot be supplied.
*******************************************************************************/
@Test
void overridesAndPermission() throws Exception
{
AssistantBudget budget = budget();
usage(1_100_000, NOW);
AssistantBudgetOverride extra = store().insert(new AssistantBudgetOverride().withBudgetId(budget.getId()).withExtraValue(BigDecimal.ONE)
.withReason("Project launch").withExpiresAt(NOW.plusSeconds(3600)).withGrantedBy("forged"));
assertThat(extra.getGrantedBy()).isEqualTo("admin-budget");
BudgetService service = service(new ArrayList<>());
assertThat(service.snapshot(budget, USER).limit()).isEqualByComparingTo("2");
store().insert(new AssistantBudgetOverride().withBudgetId(budget.getId()).withUnlimitedUntil(NOW.plusSeconds(60))
.withExpiresAt(NOW.plusSeconds(60)).withReason("Investigation"));
assertThat(service.snapshot(budget, USER).state()).isEqualTo("UNLIMITED");
BudgetService later = new BudgetService(store(), config.getRuntime().getProviderService(), Clock.fixed(NOW.plusSeconds(3601), ZoneOffset.UTC));
assertThat(later.snapshot(budget, USER).limit()).isEqualByComparingTo("1");
AssistantStore store = store();
QContext.init(qInstance, AssistantTestFixture.ownerSession(USER, config));
assertThatThrownBy(() -> store.insert(new AssistantBudgetOverride().withBudgetId(budget.getId()).withExtraValue(BigDecimal.ONE)
.withExpiresAt(NOW.plusSeconds(60)).withReason("No permission"))).hasMessageContaining("assistant.admin");
}
/*******************************************************************************
** Monthly UTC boundaries and token accounting include all token counters.
*******************************************************************************/
@Test
void monthAndTokens() throws Exception
{
AssistantBudget budget = store().insert(new AssistantBudget().withName("Month").withScopeType("INSTANCE").withPeriod("MONTH")
.withMetric("TOKENS").withLimitValue(BigDecimal.valueOf(100)).withAtLimitAction("STOP"));
usage(500, Instant.parse("2026-09-30T23:59:59Z"));
usage(50, Instant.parse("2026-10-01T00:00:00Z"));
usage(500, Instant.parse("2026-11-01T00:00:00Z"));
assertThat(service(new ArrayList<>()).snapshot(budget, USER).used()).isEqualByComparingTo("50");
}
/*******************************************************************************
** A stop dominates a warning or downgrade in another applicable budget.
*******************************************************************************/
@Test
void strictestWinsAndRoleMembership() throws Exception
{
budget();
AssistantBudget role = store().insert(new AssistantBudget().withName("Role").withScopeType("ROLE").withScopeValue("assistant.admin")
.withPeriod("DAY").withMetric("TOKENS").withLimitValue(BigDecimal.TEN).withAtLimitAction("STOP"));
usage(800_000, NOW);
BudgetService service = service(new ArrayList<>());
assertThat(service.evaluate(USER, null, false, false).stopped()).isTrue();
QContext.init(qInstance, AssistantTestFixture.ownerSession(USER, config));
assertThat(service.evaluate(USER, null, false, false).stopped()).isFalse();
store().updateFields(AssistantBudget.TABLE_NAME, new QRecord().withValue("id", role.getId()).withValue("scopeType", "INSTANCE").withValue("scopeValue", ""));
assertThat(service.evaluate(USER, null, false, false).stopped()).isTrue();
}
/*******************************************************************************
** App home has budget and time-series data and honest empty states.
*******************************************************************************/
@Test
void homeDataAndEmptyStates() throws Exception
{
store();
String empty = AssistantHomeView.render(AssistantHomeData.load(config.getRuntime()), "/app/assistantAdmin");
assertThat(empty).contains("No enabled data-driven budgets", "No assistant use", "No provider usage");
budget();
usage(800_000, Instant.now());
String html = AssistantHomeView.render(AssistantHomeData.load(config.getRuntime()), "/app/assistantAdmin");
assertThat(html).contains("Budgets", "Per user daily", "Usage over time", "Daily tokens", "Daily cost", "By provider");
assertThat(html).doesNotContain("<table", "overflow:auto");
}
/*******************************************************************************
** Next actual model dispatch uses the explicit downgrade model.
*******************************************************************************/
@Test
void turnDispatchUsesDowngradeModel() throws Exception
{
AssistantBudget budget = budget();
AssistantProvider row = store().insert(new AssistantProvider().withName("Local").withProviderType("budget-test").withModel("main").withLocation("LOCAL").withEnabled(true));
AssistantModel cheap = store().insert(new AssistantModel().withProviderId(row.getId()).withModelId("cheap").withEnabled(true));
store().updateFields(AssistantBudget.TABLE_NAME, new QRecord().withValue("id", budget.getId()).withValue("downgradeModelId", cheap.getId()));
usage(1_000_000, Instant.now());
provider.thenText("Done");
assertThat(runner.run(USER, "Hello").errorCode()).isNull();
assertThat(provider.getRequests()).singleElement().extracting(r -> r.getModel()).isEqualTo("cheap");
}
/*******************************************************************************
** Provider and model scopes isolate the other provider's same-named model.
*******************************************************************************/
@Test
void providerAndModelScopes() throws Exception
{
AssistantProvider row = store().insert(new AssistantProvider().withName("Scope provider").withProviderType("budget-test").withModel("main").withEnabled(true));
AssistantModel model = store().insert(new AssistantModel().withProviderId(row.getId()).withModelId("main").withEnabled(true));
AssistantBudget budget = store().insert(new AssistantBudget().withName("Provider").withScopeType("PROVIDER").withScopeValue(row.getId().toString())
.withPeriod("DAY").withMetric("COST_USD").withLimitValue(BigDecimal.ONE).withAtLimitAction("STOP"));
AssistantUsage usage = store().insert(new AssistantUsage().withUserId(USER).withProviderId(row.getId()).withModel("main").withCostMicros(1_000_000L));
store().updateFields(AssistantUsage.TABLE_NAME, new QRecord().withValue("id", usage.getId()).withValue("createDate", NOW));
BudgetService service = service(new ArrayList<>());
ProviderService.Active active = config.getRuntime().getProviderService().resolve("main");
assertThat(service.evaluate(USER, active, false, false).stopped()).isTrue();
assertThat(service.evaluate(USER, null, false, false).stopped()).isFalse();
store().updateFields(AssistantBudget.TABLE_NAME, new QRecord().withValue("id", budget.getId()).withValue("scopeType", "MODEL").withValue("scopeValue", model.getId().toString()));
assertThat(service.evaluate(USER, active, false, false).stopped()).isTrue();
assertThat(service.evaluate(USER, new ProviderService.Active(row, "budget-test", "another-model", provider), false, false).stopped()).isFalse();
}
/*******************************************************************************
** Usage from the first step stops the next dispatch with BUDGET.
*******************************************************************************/
@Test
void midTurnStop() throws Exception
{
store().insert(new AssistantBudget().withName("Token steps").withScopeType("USER").withScopeValue(USER)
.withPeriod("DAY").withMetric("TOKENS").withLimitValue(BigDecimal.valueOf(100)).withAtLimitAction("STOP"));
provider.withUsagePerCall(new ModelUsage().withInputTokens(100L).withOutputTokens(0L));
provider.thenToolUse("count_records", Map.of("table", "page")).thenText("Should not dispatch");
assertThat(runner.run(USER, "Count pages").errorCode()).isEqualTo("BUDGET");
assertThat(provider.getRequests()).hasSize(1);
}
/*******************************************************************************
** The start-turn warning reaches the panel once and is not replayed.
*******************************************************************************/
@Test
void panelWarningOnce() throws Exception
{
budget();
usage(800_000, Instant.now());
provider.thenText("Hello").thenText("Hello again");
assertThat(runner.run(USER, "Hello").events(TurnEvent.NOTICE)).hasSize(1);
assertThat(runner.run(USER, "Hello again").events(TurnEvent.NOTICE)).isEmpty();
}
/*******************************************************************************
** Another user's cost counts against shared budgets without unlocking detail.
*******************************************************************************/
@Test
void sharedBudgetKeepsUsageOwnerLocks() throws Exception
{
store().insert(new AssistantBudget().withName("Whole instance").withScopeType("INSTANCE").withPeriod("DAY")
.withMetric("COST_USD").withLimitValue(BigDecimal.ONE).withAtLimitAction("STOP"));
usage(1_000_000, NOW);
BudgetService service = service(new ArrayList<>());
AssistantStore store = store();
QContext.init(qInstance, AssistantTestFixture.ownerSession("editor-other", config));
assertThat(store.query(AssistantUsage.TABLE_NAME, new QQueryFilter(), AssistantUsage.class)).isEmpty();
assertThat(service.evaluate("editor-other", null, false, false).stopped()).isTrue();
assertThat(qInstance.getTable(AssistantBudgetUsage.TABLE_NAME).getRecordSecurityLocks()).isNullOrEmpty();
assertThat(qInstance.getTable(AssistantUsage.TABLE_NAME).getRecordSecurityLocks()).isNotEmpty();
}
/*******************************************************************************
** A provider STOP is not skipped after a user budget selects another provider.
*******************************************************************************/
@Test
void originalProviderStopWinsOverDowngrade() throws Exception
{
AssistantBudget user = budget();
AssistantProvider original = store().insert(new AssistantProvider().withName("Original").withProviderType("budget-test").withModel("main").withEnabled(true).withIsDefault(true));
AssistantProvider local = store().insert(new AssistantProvider().withName("Target").withProviderType("budget-test").withModel("cheap").withEnabled(true).withLocation("LOCAL"));
AssistantModel model = store().insert(new AssistantModel().withProviderId(local.getId()).withModelId("cheap").withEnabled(true));
store().updateFields(AssistantBudget.TABLE_NAME, new QRecord().withValue("id", user.getId()).withValue("downgradeModelId", model.getId()));
store().insert(new AssistantBudget().withName("Provider stop").withScopeType("PROVIDER").withScopeValue(original.getId().toString())
.withPeriod("DAY").withMetric("COST_USD").withLimitValue(BigDecimal.ONE).withAtLimitAction("STOP"));
AssistantUsage usage = store().insert(new AssistantUsage().withUserId(USER).withProviderId(original.getId()).withModel("main").withCostMicros(1_000_000L));
store().updateFields(AssistantUsage.TABLE_NAME, new QRecord().withValue("id", usage.getId()).withValue("createDate", NOW));
BudgetService service = service(new ArrayList<>());
assertThat(service.evaluate(USER, config.getRuntime().getProviderService().resolve("main"), false, false).reason()).contains("Provider stop");
assertThat(store().query(AssistantBudgetEvent.TABLE_NAME, new QQueryFilter(), AssistantBudgetEvent.class)).extracting(AssistantBudgetEvent::getState).containsExactly("STOPPED");
}
/*******************************************************************************
** A PHI conversation cannot downgrade back to its original model or to an
** ineligible candidate; no model receives the request.
*******************************************************************************/
@Test
void phiConversationWithoutEligibleDowngradeStops() throws Exception
{
budget();
AssistantProvider original = store().insert(new AssistantProvider().withName("Eligible original").withProviderType("budget-test")
.withModel("main").withEnabled(true).withIsDefault(true).withHipaaEligible(true));
AssistantProvider target = store().insert(new AssistantProvider().withName("Ineligible target").withProviderType("budget-test")
.withModel("cheap").withEnabled(true).withLocation("LOCAL").withHipaaEligible(false));
store().insert(new AssistantModel().withProviderId(original.getId()).withModelId("main").withEnabled(true));
store().insert(new AssistantModel().withProviderId(target.getId()).withModelId("cheap").withEnabled(true));
AssistantConversation conversation = store().insert(new AssistantConversation().withUserId(USER).withTitle("Classified")
.withMaxClassification("PHI").withStatus("ACTIVE"));
usage(1_000_000, Instant.now());
assertThat(runner.run(USER, conversation.getId(), "Continue", Map.of()).errorCode()).isEqualTo("BUDGET");
assertThat(provider.getRequests()).isEmpty();
}
/*******************************************************************************
** Partial edits validate the full row; narrow overrides affect only their user.
*******************************************************************************/
@Test
void validationAndNarrowOverrides() throws Exception
{
AssistantBudget budget = store().insert(new AssistantBudget().withName("Shared").withScopeType("INSTANCE").withPeriod("DAY")
.withMetric("COST_USD").withLimitValue(BigDecimal.ONE));
assertThatThrownBy(() -> store().updateFields(AssistantBudget.TABLE_NAME,
new QRecord().withValue("id", budget.getId()).withValue("hardCeilingPercent", 99))).hasMessageContaining("hard ceiling");
store().insert(new AssistantBudgetOverride().withBudgetId(budget.getId()).withScopeValue(USER).withExtraValue(BigDecimal.ONE)
.withExpiresAt(NOW.plusSeconds(3600)).withReason("One user"));
BudgetService service = service(new ArrayList<>());
assertThat(service.snapshot(budget, USER).limit()).isEqualByComparingTo("2");
assertThat(service.snapshot(budget, "editor-other").limit()).isEqualByComparingTo("1");
assertThatThrownBy(() -> store().insert(new AssistantBudget().withName("Invalid provider").withScopeType("PROVIDER")
.withScopeValue("missing").withPeriod("DAY").withMetric("TOKENS").withLimitValue(BigDecimal.TEN))).hasMessageContaining("numeric");
}
}
@@ -17,11 +17,6 @@
package com.kingsrook.qbits.assistant.settings;
import java.io.Serializable;
import java.math.BigDecimal;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import com.kingsrook.qbits.assistant.AssistantQBitConfig;
import com.kingsrook.qbits.assistant.api.AssistantApiException;
import com.kingsrook.qbits.assistant.engine.ConnectionTest;
@@ -71,6 +66,11 @@ import com.kingsrook.qqq.backend.core.model.data.QRecord;
import com.kingsrook.qqq.backend.core.model.metadata.tables.QTableMetaData;
import com.kingsrook.qqq.backend.core.processes.implementations.etl.streamedwithfrontend.StreamedETLWithFrontendProcess;
import com.kingsrook.qqq.backend.core.utils.JsonUtils;
import java.io.Serializable;
import java.math.BigDecimal;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -516,7 +516,7 @@ class AssistantSettingsTest extends BaseAssistantTest
asAdmin();
String html = render();
assertThat(html).contains("No providers yet").contains("No assistant use in the last 30 days").contains("No tool calls yet")
.contains("No failed model calls").contains("No one has used the assistant this month").doesNotContain("<table");
.contains("No failed model calls").contains("No one has used the assistant in the last 30 days").doesNotContain("<table");
}