diff --git a/CLAUDE.md b/CLAUDE.md index e4e227d..b6b1011 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -311,6 +311,82 @@ embedding model and read through another raises no error — retrieval just degr (pgvector's `text`-column fallback has no dimension constraint and cosine similarity still returns a number). +## LLM Usage & Cost Accounting + +Every model call writes one `llm_usage` row (`V119__create_llm_usage.sql`), surfaced at +**Settings → AI Usage & Cost** (`LlmUsageTab.jsx`) via `GET /admin/llm-usage/summary`. +Admin-only and not connection-scoped, so it carries a class-level `@PreAuthorize` rather +than an `assertCan*` call — the case `Endpoint Authorization Rules` describes as +"an endpoint with no connection scope at all is admin-only". + +- **Recording happens at the two provider funnels, never in feature code.** + `RefreshableChatModel.call/stream` and `EmbeddingService.createEmbedding(s)` are the only + call sites, so a thirteenth service that reaches a model is accounted for without + touching it. `LlmUsageRecorder` swallows its own failures by design: accounting must + never be able to break the call it measures. +- **Rows are written in their own transaction** (`LlmUsageWriter`, `REQUIRES_NEW`). It is a + separate bean on purpose — self-invoking a `@Transactional` method bypasses the Spring + proxy, and the row would then roll back with a failed chat turn, which is exactly the + case it exists to record. Same trap `McpTokenRepository.deleteByUserId` documents. +- **`estimated` distinguishes metered from derived counts.** Chat providers return real + token counts; `LlmEmbeddingProvider` returns vectors only, so embedding tokens are + derived from input length at 3 chars/token (schema text is denser than prose — see + `EmbeddingService.truncate`). Both land in the same columns so one spend total is + possible, but a vendor-invoice reconciliation can tell the halves apart. +- **An unpriced model stores `NULL` cost, never `0`.** `LlmPricingService` reads rates from + `system_config` (`llm.pricing..{input,output,cached-input}-per-1m`) and ships **no + default prices** — a stale bundled price list produces confident wrong totals nobody + thinks to check. The UI reports unpriced calls rather than silently understating spend. +- **Rates are edited in the UI**, in the Model pricing panel of the same tab + (`LlmPricingPanel.jsx` → `GET|PUT /admin/llm-usage/pricing`). The list is every model the + ledger has seen plus every model with a rate configured, unpriced first. Writes take + effect on the next call with no restart, since `LlmPricingService` reads `system_config` + per call. Three details are load-bearing: + - **A cleared field writes `""`, not a deleted row.** `rate()` already treats blank as + absent, and `SystemConfigService` has no delete — adding one for a single caller would + widen a shared service. Sending the whole set means an emptied box genuinely clears + that rate rather than leaving the old value behind. + - **A model name can contain dots** (`gpt-5.4`), so `configuredModels()` strips the known + suffix from the *end* of the key. Splitting on the first `.` after the prefix reports + `gpt-5` and loses the row. The controller mapping is `{model:.+}` for the same reason. + - **A name containing a slash cannot go in the path at all**, even percent-encoded: + Spring Security's default `StrictHttpFirewall` rejects `%2F` with a bare 400 before any + controller runs — verified in QA, and it applies to every endpoint, not just this one. + Self-hosted ids look like `meta-llama/Llama-3-8b`, so `PUT /pricing` (no path segment) + takes the name in the body, and `client.js` switches to it when the name contains `/`. + Relaxing the firewall would be the wrong trade for a naming convenience. + - **A failed save must still say so.** Spring's default 500 body carries no `message`, so + the client had nothing to display and a save against a broken config store showed the + user *nothing at all* — no error, no toast, silent. The handler now returns a `message` + on any non-`IllegalArgumentException` failure, and the panel catches the `mutateAsync` + rejection rather than letting it escape the click handler (an uncaught rejection there + is what stopped the `isError` banner from rendering). Both found by taking + `system_config` away mid-save. + - **A negative rate is rejected at write time**, not just ignored on read: a value that + silently does nothing after the UI said "Saved" is worse than an error at entry. +- **Editing a rate does not re-cost recorded calls.** `estimated_cost_usd` is a snapshot + written at record time, which is what an audit trail wants — but it means a mid-window + price change shows as a step in the daily chart, not a uniform restatement. The panel + says so. +- **Attribution cannot ride a ThreadLocal alone.** `LlmUsageAttributionFilter` labels each + request from its URI, but chat returns a `Flux` and does its model work later on a + `CompletableFuture` — the servlet, and any thread-local set during it, is gone by then. + `ChatService` re-establishes `LlmUsageContext.with(...)` inside `runAsync` exactly where + it already re-establishes `QueryActorContextHolder.withActor`; `DashboardAlertService` + declares its own scope since scheduled work has no request at all. Verified live: every + row read `feature = unknown` until this was added, while a single-threaded unit test of + the filter passed. +- **`ResponsesApiChatModel.buildMetadata` used to discard `usage` entirely**, which cost + nothing while nothing read it and became a silent zero the moment accounting summed it — + real billed calls recorded 0 tokens and $0.00. It now reads both dialects + (`prompt_tokens`/`completion_tokens` and `input_tokens`/`output_tokens`) plus cached-token + details, and the streaming path sends `stream_options.include_usage` and emits the late + usage event as a text-free chunk. `RefreshableChatModel.meteredStream` records **one row + per stream**, taking the last reported usage rather than summing chunks: providers that + report running totals would otherwise have every partial added to the final figure. +- Usage belongs to `QueryActorContextHolder` first and the security principal second, so + under **View as** the spend is attributed to the target user, not the admin. + ## Key Rules & Patterns ### Backend Rules @@ -853,6 +929,10 @@ DEEPSQL_EMBEDDING_API_KEY= DEEPSQL_EMBEDDING_MODEL=text-embedding-3-large # Optional chat tuning: DEEPSQL_CHAT_TEMPERATURE, DEEPSQL_CHAT_API_VERSION, # DEEPSQL_CHAT_USE_RESPONSES_API (true|false|auto). +# Model prices are NOT environment variables. They live in system_config and are edited in +# Settings -> AI Usage & Cost -> Model pricing; there are deliberately no defaults: +# llm.pricing..input-per-1m / .output-per-1m / .cached-input-per-1m +# An unpriced model still records tokens; its cost is NULL and the UI flags it. # Only if using Azure AI Search instead of pgvector for the vector store. azure.search.api-key= diff --git a/backend/src/main/java/com/dbaagent/config/LlmConfig.java b/backend/src/main/java/com/dbaagent/config/LlmConfig.java index 5167807..bc6341b 100644 --- a/backend/src/main/java/com/dbaagent/config/LlmConfig.java +++ b/backend/src/main/java/com/dbaagent/config/LlmConfig.java @@ -4,6 +4,7 @@ import com.dbaagent.llm.LlmProviderRegistry; import com.dbaagent.llm.spring.ProviderBackedEmbeddingModel; import com.dbaagent.service.EmbeddingService; +import com.dbaagent.service.llm.LlmUsageRecorder; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.context.annotation.Bean; @@ -24,8 +25,9 @@ public class LlmConfig { @Bean @Primary - public ChatModel chatModel(LlmConfigResolver resolver, LlmProviderRegistry registry) { - return new RefreshableChatModel(resolver, registry); + public ChatModel chatModel(LlmConfigResolver resolver, LlmProviderRegistry registry, + LlmUsageRecorder usageRecorder) { + return new RefreshableChatModel(resolver, registry, usageRecorder); } /** diff --git a/backend/src/main/java/com/dbaagent/config/RefreshableChatModel.java b/backend/src/main/java/com/dbaagent/config/RefreshableChatModel.java index 52b3978..1cf206f 100644 --- a/backend/src/main/java/com/dbaagent/config/RefreshableChatModel.java +++ b/backend/src/main/java/com/dbaagent/config/RefreshableChatModel.java @@ -7,15 +7,20 @@ import com.dbaagent.llm.api.LlmErrorCategory; import com.dbaagent.llm.api.LlmNotConfiguredException; import com.dbaagent.llm.api.UnsupportedLlmProviderException; +import com.dbaagent.model.LlmUsageRole; +import com.dbaagent.service.llm.LlmUsageRecorder; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.chat.metadata.Usage; import reactor.core.publisher.Flux; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Supplier; /** * A {@link ChatModel} that resolves its credentials on every call through @@ -52,28 +57,62 @@ private record CachedDelegate(LlmCredentials key, ChatModel model) {} private final LlmConfigResolver resolver; private final LlmProviderRegistry registry; + /** + * Nullable so the accounting dependency stays optional. This class is constructed + * directly in tests and in {@code LlmConfig}; a hard dependency would make every + * existing construction site a compile error for a concern none of them care about. + */ + private final LlmUsageRecorder usageRecorder; + /** Serialises rebuilds only; the read path is a single lock-free volatile read. */ private final ReentrantLock buildLock = new ReentrantLock(); private volatile CachedDelegate cached; public RefreshableChatModel(LlmConfigResolver resolver, LlmProviderRegistry registry) { + this(resolver, registry, null); + } + + public RefreshableChatModel(LlmConfigResolver resolver, LlmProviderRegistry registry, + LlmUsageRecorder usageRecorder) { this.resolver = resolver; this.registry = registry; + this.usageRecorder = usageRecorder; } @Override public ChatResponse call(Prompt prompt) { CachedDelegate active = resolveDelegate(); try { - return active.model().call(prompt); + return metered(active, () -> active.model().call(prompt)); } catch (RuntimeException e) { if (shouldRetryWithEnvFallback(e, active.key())) { - return resolveDelegate().model().call(prompt); + CachedDelegate retry = resolveDelegate(); + return metered(retry, () -> retry.model().call(prompt)); } throw e; } } + /** + * Runs a chat call and records what it cost. + * + *

A failed call is recorded too. Providers bill for prompt tokens on responses that + * error partway, and an operator investigating a spend spike caused by a retry loop + * needs to see the failures — a ledger holding only successes hides exactly the + * pathology it would be consulted about. + */ + private ChatResponse metered(CachedDelegate active, Supplier call) { + long startedAt = System.nanoTime(); + try { + ChatResponse response = call.get(); + recordUsage(active, response, startedAt, null); + return response; + } catch (RuntimeException e) { + recordUsage(active, null, startedAt, e); + throw e; + } + } + /** * Streaming delegates report failures asynchronously through the sink rather than by * throwing, so the fallback hangs off the error signal as well as off a synchronous @@ -87,15 +126,65 @@ public Flux stream(Prompt prompt) { CachedDelegate active = resolveDelegate(); AtomicBoolean emitted = new AtomicBoolean(false); try { - return active.model().stream(prompt) + return meteredStream(active, active.model().stream(prompt) .doOnNext(chunk -> emitted.set(true)) - .onErrorResume(e -> resumeStream(prompt, e, active.key(), emitted.get())); + .onErrorResume(e -> resumeStream(prompt, e, active.key(), emitted.get()))); } catch (RuntimeException e) { - return resumeStream(prompt, e, active.key(), false); + return meteredStream(active, resumeStream(prompt, e, active.key(), false)); } }); } + /** + * Records one row for a whole stream, not one per chunk. + * + *

Usage on a stream arrives on a single late chunk — typically the last, after the + * provider has finished counting — while every earlier chunk carries either no + * metadata or a zero-filled {@link Usage}. Recording per chunk would write hundreds of + * rows for one call and inflate the call count enormously; summing across chunks would + * be worse, because providers that report cumulative running totals would have their + * final figure added on top of every partial. So the last non-zero usage seen wins, and + * exactly one row is written when the stream terminates. + * + *

{@code doFinally} rather than {@code doOnComplete}: a stream that errors or is + * cancelled partway still consumed prompt tokens, and a cancelled dashboard build is + * precisely the kind of silent spend an operator wants on the ledger. + */ + private Flux meteredStream(CachedDelegate active, Flux source) { + if (usageRecorder == null) { + return source; + } + long startedAt = System.nanoTime(); + AtomicReference lastWithUsage = new AtomicReference<>(); + AtomicReference failure = new AtomicReference<>(); + AtomicBoolean recorded = new AtomicBoolean(false); + + return source + .doOnNext(chunk -> { + if (hasUsage(chunk)) { + lastWithUsage.set(chunk); + } + }) + .doOnError(failure::set) + .doFinally(signal -> { + // doFinally can fire once per subscription; a Flux that is retried or + // resubscribed must not double-bill the same logical call. + if (recorded.compareAndSet(false, true)) { + Throwable error = failure.get(); + recordUsage(active, lastWithUsage.get(), startedAt, + error instanceof RuntimeException re ? re : null); + } + }); + } + + private static boolean hasUsage(ChatResponse chunk) { + if (chunk == null || chunk.getMetadata() == null) { + return false; + } + Usage usage = chunk.getMetadata().getUsage(); + return usage != null && zeroIfNull(usage.getTotalTokens()) > 0; + } + /** * Reports the active delegate's options, or neutral ones when nothing is configured. * @@ -128,6 +217,75 @@ public ChatOptions getDefaultOptions() { } } + /** + * Records one chat call, taking token counts from the response the provider returned. + * + *

The model name comes from the response metadata when the provider reports it and + * falls back to the configured model. Those can legitimately differ — an alias that + * resolves to a dated snapshot, for instance — and the served model is the one that + * was actually billed, so it wins. + * + *

Never throws. It runs inside the call path of every chat turn in the product; + * a defect here must not become a failed conversation. + */ + private void recordUsage(CachedDelegate active, ChatResponse response, + long startedAt, RuntimeException failure) { + if (usageRecorder == null) { + return; + } + try { + long latencyMs = (System.nanoTime() - startedAt) / 1_000_000L; + LlmCredentials key = active.key(); + + long prompt = 0; + long completion = 0; + long total = 0; + long cached = 0; + String model = key.getOrDefault("model", "unknown"); + + if (response != null && response.getMetadata() != null) { + Usage usage = response.getMetadata().getUsage(); + if (usage != null) { + prompt = zeroIfNull(usage.getPromptTokens()); + completion = zeroIfNull(usage.getCompletionTokens()); + total = zeroIfNull(usage.getTotalTokens()); + cached = zeroIfNullLong(usage.getCacheReadInputTokens()); + } + String served = response.getMetadata().getModel(); + if (served != null && !served.isBlank()) { + model = served; + } + } + + String errorCategory = failure == null ? null + : String.valueOf(registry.chatProvider(key.providerId()).classify(failure)); + + usageRecorder.record(new LlmUsageRecorder.Call( + LlmUsageRole.CHAT, + key.providerId(), + model, + prompt, + completion, + total, + cached, + false, + latencyMs, + failure == null, + errorCategory)); + } catch (RuntimeException e) { + log.warn("RefreshableChatModel: could not record usage; the call was unaffected", e); + } + } + + private static long zeroIfNull(Integer value) { + return value == null ? 0L : value.longValue(); + } + + /** Cache token accessors are {@code Long} in Spring AI 2.0, unlike prompt/completion. */ + private static long zeroIfNullLong(Long value) { + return value == null ? 0L : value; + } + // ── Internal ────────────────────────────────────────────────────────────── private CachedDelegate resolveDelegate() { diff --git a/backend/src/main/java/com/dbaagent/controller/LlmUsageController.java b/backend/src/main/java/com/dbaagent/controller/LlmUsageController.java new file mode 100644 index 0000000..4a52a2f --- /dev/null +++ b/backend/src/main/java/com/dbaagent/controller/LlmUsageController.java @@ -0,0 +1,118 @@ +package com.dbaagent.controller; + +import com.dbaagent.model.LlmUsage; +import com.dbaagent.service.llm.LlmPricingService; +import com.dbaagent.service.llm.LlmUsageQueryService; +import lombok.extern.slf4j.Slf4j; +import org.springframework.data.domain.Page; +import org.springframework.http.HttpStatus; +import org.springframework.http.ResponseEntity; +import org.springframework.security.access.prepost.PreAuthorize; +import org.springframework.web.bind.annotation.*; + +import java.math.BigDecimal; +import java.util.List; +import java.util.Map; + +/** + * LLM spend reporting. + * + *

Admin-only, and not scoped to a connection: usage spans every connection and every + * user, so there is no connection ACL that could authorize it. That makes this one of the + * cases {@code CLAUDE.md} calls out — an endpoint with no connection scope at all is + * admin-only, enforced with {@code @PreAuthorize} rather than an access-control assert. + */ +@Slf4j +@RestController +@RequestMapping("/admin/llm-usage") +@PreAuthorize("hasRole('ADMIN')") +public class LlmUsageController { + + private final LlmUsageQueryService usageQueryService; + private final LlmPricingService pricingService; + + public LlmUsageController(LlmUsageQueryService usageQueryService, + LlmPricingService pricingService) { + this.usageQueryService = usageQueryService; + this.pricingService = pricingService; + } + + @GetMapping("/summary") + public ResponseEntity summary( + @RequestParam(defaultValue = "30") int days) { + return ResponseEntity.ok(usageQueryService.summary(days)); + } + + @GetMapping("/recent") + public ResponseEntity> recent( + @RequestParam(defaultValue = "30") int days, + @RequestParam(defaultValue = "0") int page, + @RequestParam(defaultValue = "50") int size) { + return ResponseEntity.ok(usageQueryService.recent(days, page, size)); + } + + @GetMapping("/pricing") + public ResponseEntity> pricing() { + return ResponseEntity.ok(usageQueryService.pricing()); + } + + /** + * Rates for one model, in USD per 1M tokens. A null rate field clears that rate. + * + *

{@code model} is optional here and only used when the path cannot carry the name + * — see {@link #updatePricing}. + */ + public record PricingUpdate( + String model, + BigDecimal inputPer1m, + BigDecimal outputPer1m, + BigDecimal cachedInputPer1m) {} + + /** + * The model is normally a path variable, so the URL identifies what is being edited, + * and {@code :.+} keeps a dotted name like {@code gpt-5.4} intact. + * + *

A name containing a slash cannot travel in the path at all, even percent-encoded: + * Spring Security's default {@code StrictHttpFirewall} rejects {@code %2F} with a bare + * 400 before any controller runs — verified against this deployment, and it rejects + * such a URL on every endpoint, not just this one. Self-hosted model ids are routinely + * of the form {@code meta-llama/Llama-3-8b}, so the body may carry {@code model} + * instead and the caller PUTs to {@link #PRICING_BODY_MODEL}. Relaxing the firewall + * would be the wrong trade: it is a platform-wide security control, and this is a + * naming convenience. + */ + static final String PRICING_BODY_MODEL = "_"; + + @PutMapping({"/pricing/{model:.+}", "/pricing"}) + public ResponseEntity updatePricing( + @PathVariable(required = false) String model, @RequestBody PricingUpdate update) { + try { + String target = PRICING_BODY_MODEL.equals(model) || model == null || model.isBlank() + ? update.model() + : model; + LlmPricingService.ModelRates saved = pricingService.updateRates( + target, update.inputPer1m(), update.outputPer1m(), update.cachedInputPer1m()); + log.info("Updated LLM pricing for model {}", saved.model()); + return ResponseEntity.ok(saved); + } catch (IllegalArgumentException e) { + // A bad rate is the operator's typo, not a server fault, and the message names + // what to correct. + return ResponseEntity.badRequest().body(Map.of("message", e.getMessage())); + } catch (RuntimeException e) { + // Anything else — the config store being unreachable, say — must still come + // back with a `message`. Spring's default 500 body has none, and the client + // then has nothing to display: verified by taking system_config away, where + // the save failed and the UI showed the user nothing at all. + log.error("Could not save LLM pricing for model {}", model, e); + return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR) + .body(Map.of("message", "Could not save the rate: " + e.getMessage())); + } + } + + @DeleteMapping("/purge") + public ResponseEntity> purge(@RequestParam int olderThanDays) { + int deleted = usageQueryService.purgeOlderThan(olderThanDays); + log.info("Purged {} LLM usage rows older than {} days", deleted, olderThanDays); + return ResponseEntity.ok(Map.of("deleted", deleted)); + } +} diff --git a/backend/src/main/java/com/dbaagent/dto/LlmUsageDailyPoint.java b/backend/src/main/java/com/dbaagent/dto/LlmUsageDailyPoint.java new file mode 100644 index 0000000..881862a --- /dev/null +++ b/backend/src/main/java/com/dbaagent/dto/LlmUsageDailyPoint.java @@ -0,0 +1,12 @@ +package com.dbaagent.dto; + +import java.math.BigDecimal; +import java.time.LocalDate; + +/** One day of the spend trend. */ +public record LlmUsageDailyPoint( + LocalDate day, + long calls, + long totalTokens, + BigDecimal costUsd) { +} diff --git a/backend/src/main/java/com/dbaagent/dto/LlmUsageGroup.java b/backend/src/main/java/com/dbaagent/dto/LlmUsageGroup.java new file mode 100644 index 0000000..70e5a0e --- /dev/null +++ b/backend/src/main/java/com/dbaagent/dto/LlmUsageGroup.java @@ -0,0 +1,11 @@ +package com.dbaagent.dto; + +import java.math.BigDecimal; + +/** One row of a "by feature" / "by user" / "by model" breakdown. */ +public record LlmUsageGroup( + String key, + long calls, + long totalTokens, + BigDecimal costUsd) { +} diff --git a/backend/src/main/java/com/dbaagent/dto/LlmUsageTotals.java b/backend/src/main/java/com/dbaagent/dto/LlmUsageTotals.java new file mode 100644 index 0000000..772b1be --- /dev/null +++ b/backend/src/main/java/com/dbaagent/dto/LlmUsageTotals.java @@ -0,0 +1,22 @@ +package com.dbaagent.dto; + +import java.math.BigDecimal; + +/** + * Window totals. {@code unpricedCalls} is reported alongside cost so the UI can say + * "$12.40 across 900 calls, 40 of them unpriced" rather than presenting a partial sum as + * if it were complete. + */ +public record LlmUsageTotals( + long calls, + long promptTokens, + long completionTokens, + long totalTokens, + BigDecimal costUsd, + long unpricedCalls, + long failedCalls) { + + public static LlmUsageTotals empty() { + return new LlmUsageTotals(0, 0, 0, 0, BigDecimal.ZERO, 0, 0); + } +} diff --git a/backend/src/main/java/com/dbaagent/llm/openai/ResponsesApiChatModel.java b/backend/src/main/java/com/dbaagent/llm/openai/ResponsesApiChatModel.java index 7395493..a87a20d 100644 --- a/backend/src/main/java/com/dbaagent/llm/openai/ResponsesApiChatModel.java +++ b/backend/src/main/java/com/dbaagent/llm/openai/ResponsesApiChatModel.java @@ -9,6 +9,7 @@ import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.MessageType; import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.metadata.DefaultUsage; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.Generation; @@ -252,7 +253,7 @@ public Flux stream(Prompt prompt) { // -- Request building -- - private String buildRequestBody(Prompt prompt, boolean stream) { + String buildRequestBody(Prompt prompt, boolean stream) { ObjectNode body = objectMapper.createObjectNode(); body.put("model", model); body.put("stream", stream); @@ -264,6 +265,15 @@ private String buildRequestBody(Prompt prompt, boolean stream) { } body.put("temperature", effectiveTemperature); + // Chat Completions omits usage from a stream unless asked; the Responses API + // reports it on its completion event either way. Without this, a streamed turn is + // unbillable — the tokens were spent and nothing says how many. + if (stream && !useResponsesApi) { + ObjectNode streamOptions = objectMapper.createObjectNode(); + streamOptions.put("include_usage", true); + body.set("stream_options", streamOptions); + } + if (useResponsesApi) { buildResponsesApiBody(prompt, body); } else { @@ -374,8 +384,60 @@ private String parseChatCompletionsOutput(JsonNode root) { return ""; } - private ChatResponseMetadata buildMetadata(JsonNode usage) { - return ChatResponseMetadata.builder().build(); + /** + * Carries the provider's token counts onto the response. + * + *

This used to return empty metadata and drop {@code usage} on the floor, which + * cost nothing while nobody read it and became a silent zero the moment usage + * accounting started summing it — every chat row recorded 0 tokens and $0.00 against + * a real, billed call. + * + *

The two API shapes name the same numbers differently, and the difference is not + * cosmetic: Responses uses {@code input_tokens}/{@code output_tokens}, Chat + * Completions uses {@code prompt_tokens}/{@code completion_tokens}. Both are read + * rather than branching on {@code useResponsesApi}, so a mixed or proxied deployment + * cannot report zeros just because it answered in the other dialect. + * + *

Package-private as a test seam, like {@link #buildHttpRequest}: it lets both + * vendor payload shapes be asserted without a live call. + */ + ChatResponseMetadata buildMetadata(JsonNode usage) { + if (usage == null || usage.isMissingNode() || !usage.isObject()) { + return ChatResponseMetadata.builder().build(); + } + + Integer prompt = firstInt(usage, "prompt_tokens", "input_tokens"); + Integer completion = firstInt(usage, "completion_tokens", "output_tokens"); + Integer total = firstInt(usage, "total_tokens"); + if (total == null) { + total = (prompt == null ? 0 : prompt) + (completion == null ? 0 : completion); + } + + // Cached input is billed below fresh input, so it is read where reported. Both + // dialects nest it, under different parents. + Long cachedRead = null; + JsonNode details = usage.path("prompt_tokens_details"); + if (details.isMissingNode() || !details.isObject()) { + details = usage.path("input_tokens_details"); + } + if (details.isObject() && details.hasNonNull("cached_tokens")) { + cachedRead = details.get("cached_tokens").asLong(); + } + + return ChatResponseMetadata.builder() + .usage(new DefaultUsage(prompt, completion, total, null, cachedRead, null)) + .build(); + } + + /** First present, integral field among {@code names}, or null when none is reported. */ + private static Integer firstInt(JsonNode node, String... names) { + for (String name : names) { + JsonNode field = node.path(name); + if (!field.isMissingNode() && field.isNumber()) { + return field.asInt(); + } + } + return null; } // -- Streaming -- @@ -424,6 +486,15 @@ private void streamResponse(String body, FluxSink sink) { sink.next(new ChatResponse( List.of(new Generation(new AssistantMessage(delta))))); } + + // Usage arrives on its own late event, carrying no text. Emitted as + // a text-free chunk so accounting can read it; consumers concatenate + // deltas, so an empty generation adds nothing to the answer. + ChatResponseMetadata usage = extractStreamUsage(data); + if (usage != null) { + sink.next(new ChatResponse( + List.of(new Generation(new AssistantMessage(""))), usage)); + } } } sink.complete(); @@ -445,6 +516,41 @@ private void streamResponse(String body, FluxSink sink) { } } + /** + * Token counts from a streaming event, or null when this event carries none. + * + *

Both dialects report usage exactly once, late, on an event whose other fields are + * empty: Chat Completions sends a final chunk whose {@code choices} array is empty and + * whose {@code usage} is populated (only when {@code stream_options.include_usage} was + * requested), while the Responses API nests it under + * {@code response.usage} on its {@code response.completed} event. + * + *

A parse failure returns null rather than throwing. This runs per streamed event on + * the live answer path; an unexpected shape must cost the usage row, never the reply. + */ + ChatResponseMetadata extractStreamUsage(String data) { + try { + JsonNode root = objectMapper.readTree(data); + JsonNode usage = root.path("usage"); + if (usage.isMissingNode() || !usage.isObject()) { + usage = root.path("response").path("usage"); + } + if (usage.isMissingNode() || !usage.isObject()) { + return null; + } + ChatResponseMetadata metadata = buildMetadata(usage); + // buildMetadata returns empty metadata for a shape it did not recognise; + // emitting that would add a chunk that says nothing. + return metadata.getUsage() == null + || metadata.getUsage().getTotalTokens() == null + || metadata.getUsage().getTotalTokens() == 0 + ? null : metadata; + } catch (Exception e) { + log.debug("Could not read usage from a streamed event", e); + return null; + } + } + private String extractResponsesApiDelta(String data) { try { JsonNode event = objectMapper.readTree(data); diff --git a/backend/src/main/java/com/dbaagent/model/LlmUsage.java b/backend/src/main/java/com/dbaagent/model/LlmUsage.java new file mode 100644 index 0000000..c1b9098 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/model/LlmUsage.java @@ -0,0 +1,107 @@ +package com.dbaagent.model; + +import jakarta.persistence.*; +import lombok.Data; +import org.hibernate.annotations.CreationTimestamp; + +import java.math.BigDecimal; +import java.time.LocalDateTime; + +/** + * One metered LLM call. + * + *

Rows are written from the two provider funnels ({@code RefreshableChatModel} and + * {@code EmbeddingService}), never from feature code, so a new caller is accounted for + * without touching this table. + * + *

{@code estimated} is load-bearing rather than decorative. Chat providers return real + * token counts in the response; the embedding provider API returns vectors only, so those + * counts are derived from input length. Both land in the same columns because operators + * want one spend total, but a row that says 4,000 tokens has to be distinguishable from a + * row that guesses 4,000 — otherwise a reconciliation against the vendor invoice has no + * way to know which half of the ledger to trust. + */ +@Entity +@Table(name = "llm_usage") +@Data +public class LlmUsage { + + @Id + @GeneratedValue(strategy = GenerationType.IDENTITY) + private Long id; + + @Column(nullable = false, length = 16) + private String role; + + @Column(name = "provider_id", nullable = false) + private String providerId; + + @Column(nullable = false) + private String model; + + /** The feature that made the call, e.g. {@code chat}, {@code dashboard-generate}. */ + @Column(nullable = false, length = 64) + private String feature; + + /** Null for background work with no human actor (scheduled jobs, boot-time indexing). */ + @Column + private String username; + + @Column(name = "connection_id") + private String connectionId; + + @Column(name = "prompt_tokens", nullable = false) + private long promptTokens; + + @Column(name = "completion_tokens", nullable = false) + private long completionTokens; + + @Column(name = "total_tokens", nullable = false) + private long totalTokens; + + /** + * Cached prompt tokens, when the provider reports them. Billed well below fresh input + * by every provider that offers it, so cost is computed against + * {@code promptTokens - cachedPromptTokens} rather than the raw prompt count. + */ + @Column(name = "cached_prompt_tokens", nullable = false) + private long cachedPromptTokens; + + /** + * Cost in USD at the rates configured when the row was written, or null when the + * model has no configured rate. + * + *

Null is deliberately not zero. An unpriced model is an operator gap to surface, + * and writing 0.00 would silently understate spend in every total that sums this + * column. + */ + @Column(name = "estimated_cost_usd", precision = 12, scale = 6) + private BigDecimal estimatedCostUsd; + + /** True when token counts are derived locally rather than reported by the provider. */ + @Column(nullable = false) + private boolean estimated; + + @Column(name = "latency_ms") + private Long latencyMs; + + /** False for a call that threw; such rows still cost tokens upstream sometimes. */ + @Column(nullable = false) + private boolean succeeded; + + @Column(name = "error_category", length = 64) + private String errorCategory; + + @CreationTimestamp + @Column(name = "created_at", nullable = false, updatable = false) + private LocalDateTime createdAt; + + @Transient + public LlmUsageRole getRoleEnum() { + return role == null ? null : LlmUsageRole.valueOf(role); + } + + public void setRoleEnum(LlmUsageRole value) { + this.role = value != null ? value.name() : null; + } +} diff --git a/backend/src/main/java/com/dbaagent/model/LlmUsageRole.java b/backend/src/main/java/com/dbaagent/model/LlmUsageRole.java new file mode 100644 index 0000000..c86cfe7 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/model/LlmUsageRole.java @@ -0,0 +1,6 @@ +package com.dbaagent.model; + +public enum LlmUsageRole { + CHAT, + EMBEDDING +} diff --git a/backend/src/main/java/com/dbaagent/repository/LlmUsageRepository.java b/backend/src/main/java/com/dbaagent/repository/LlmUsageRepository.java new file mode 100644 index 0000000..6e0a222 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/repository/LlmUsageRepository.java @@ -0,0 +1,111 @@ +package com.dbaagent.repository; + +import com.dbaagent.model.LlmUsage; +import org.springframework.data.domain.Page; +import org.springframework.data.domain.Pageable; +import org.springframework.data.jpa.repository.JpaRepository; +import org.springframework.data.jpa.repository.Modifying; +import org.springframework.data.jpa.repository.Query; +import org.springframework.data.repository.query.Param; + +import java.time.LocalDateTime; +import java.util.List; + +public interface LlmUsageRepository extends JpaRepository { + + /** + * Totals for a window. Aggregated in SQL rather than by summing entities in Java — + * this table grows by one row per model call, so loading a month of them to add up a + * cost column would be the most expensive query in the product. + */ + @Query(""" + SELECT new com.dbaagent.dto.LlmUsageTotals( + COUNT(u), + COALESCE(SUM(u.promptTokens), 0L), + COALESCE(SUM(u.completionTokens), 0L), + COALESCE(SUM(u.totalTokens), 0L), + COALESCE(SUM(u.estimatedCostUsd), 0), + SUM(CASE WHEN u.estimatedCostUsd IS NULL THEN 1L ELSE 0L END), + SUM(CASE WHEN u.succeeded = false THEN 1L ELSE 0L END)) + FROM LlmUsage u + WHERE u.createdAt >= :since + """) + com.dbaagent.dto.LlmUsageTotals totalsSince(@Param("since") LocalDateTime since); + + @Query(""" + SELECT new com.dbaagent.dto.LlmUsageGroup( + COALESCE(u.feature, 'unknown'), + COUNT(u), + COALESCE(SUM(u.totalTokens), 0L), + COALESCE(SUM(u.estimatedCostUsd), 0)) + FROM LlmUsage u + WHERE u.createdAt >= :since + GROUP BY u.feature + ORDER BY COALESCE(SUM(u.estimatedCostUsd), 0) DESC, COUNT(u) DESC + """) + List byFeatureSince(@Param("since") LocalDateTime since); + + @Query(""" + SELECT new com.dbaagent.dto.LlmUsageGroup( + COALESCE(u.username, 'background'), + COUNT(u), + COALESCE(SUM(u.totalTokens), 0L), + COALESCE(SUM(u.estimatedCostUsd), 0)) + FROM LlmUsage u + WHERE u.createdAt >= :since + GROUP BY u.username + ORDER BY COALESCE(SUM(u.estimatedCostUsd), 0) DESC, COUNT(u) DESC + """) + List byUserSince(@Param("since") LocalDateTime since); + + @Query(""" + SELECT new com.dbaagent.dto.LlmUsageGroup( + COALESCE(u.model, 'unknown'), + COUNT(u), + COALESCE(SUM(u.totalTokens), 0L), + COALESCE(SUM(u.estimatedCostUsd), 0)) + FROM LlmUsage u + WHERE u.createdAt >= :since + GROUP BY u.model + ORDER BY COALESCE(SUM(u.estimatedCostUsd), 0) DESC, COUNT(u) DESC + """) + List byModelSince(@Param("since") LocalDateTime since); + + /** Daily buckets for the spend chart. */ + @Query(""" + SELECT new com.dbaagent.dto.LlmUsageDailyPoint( + CAST(u.createdAt AS java.time.LocalDate), + COUNT(u), + COALESCE(SUM(u.totalTokens), 0L), + COALESCE(SUM(u.estimatedCostUsd), 0)) + FROM LlmUsage u + WHERE u.createdAt >= :since + GROUP BY CAST(u.createdAt AS java.time.LocalDate) + ORDER BY CAST(u.createdAt AS java.time.LocalDate) ASC + """) + List dailySince(@Param("since") LocalDateTime since); + + Page findByCreatedAtGreaterThanEqualOrderByCreatedAtDesc( + LocalDateTime since, Pageable pageable); + + /** Models seen in the window that have no configured price, for the operator nudge. */ + @Query(""" + SELECT DISTINCT u.model FROM LlmUsage u + WHERE u.createdAt >= :since AND u.estimatedCostUsd IS NULL + """) + List unpricedModelsSince(@Param("since") LocalDateTime since); + + /** + * Every model the ledger has ever seen, busiest first. + * + *

Deliberately unwindowed, unlike the reporting queries: the pricing editor is + * about configuration, not about a reporting period, and a model that went quiet last + * month still needs its rate visible and editable. + */ + @Query("SELECT u.model FROM LlmUsage u GROUP BY u.model ORDER BY COUNT(u) DESC") + List distinctModels(); + + @Modifying + @Query("DELETE FROM LlmUsage u WHERE u.createdAt < :before") + int deleteByCreatedAtBefore(@Param("before") LocalDateTime before); +} diff --git a/backend/src/main/java/com/dbaagent/repository/SystemConfigRepository.java b/backend/src/main/java/com/dbaagent/repository/SystemConfigRepository.java index 7962958..c184ba6 100644 --- a/backend/src/main/java/com/dbaagent/repository/SystemConfigRepository.java +++ b/backend/src/main/java/com/dbaagent/repository/SystemConfigRepository.java @@ -2,8 +2,16 @@ import com.dbaagent.model.SystemConfig; import org.springframework.data.jpa.repository.JpaRepository; +import org.springframework.data.jpa.repository.Query; +import org.springframework.data.repository.query.Param; import org.springframework.stereotype.Repository; +import java.util.List; + @Repository public interface SystemConfigRepository extends JpaRepository { + + /** Keys beginning with {@code prefix}. Used to enumerate namespaced config families. */ + @Query("SELECT c.key FROM SystemConfig c WHERE c.key LIKE CONCAT(:prefix, '%')") + List findKeysByPrefix(@Param("prefix") String prefix); } diff --git a/backend/src/main/java/com/dbaagent/service/ChatService.java b/backend/src/main/java/com/dbaagent/service/ChatService.java index 440d735..8d1d5f5 100644 --- a/backend/src/main/java/com/dbaagent/service/ChatService.java +++ b/backend/src/main/java/com/dbaagent/service/ChatService.java @@ -69,6 +69,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.dbaagent.util.QueryNormalizer; import com.dbaagent.service.security.AccessControlService; +import com.dbaagent.service.llm.LlmUsageContext; import org.springframework.ai.chat.client.ChatClient; import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor; import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor; @@ -3435,7 +3436,8 @@ public StreamResult streamProcessMessage(String connectionId, String message, St String actualUserQuestion = extractActualUserQuestion(message); String actorUsername = resolveExecutionActorUsername(userId); boolean actorIsAdmin = isAdminUser(actorUsername); - return QueryActorContextHolder.withActor(actorUsername, () -> { + return QueryActorContextHolder.withActor(actorUsername, () -> + LlmUsageContext.with("chat", connectionId, () -> { PreparedConversationTurn prepared = prepareConversationTurn(connectionId, chatId, actualUserQuestion); ChatResponse scopeGuardResponse = maybeBuildScopeGuardrailResponse( connectionId, @@ -3449,7 +3451,7 @@ public StreamResult streamProcessMessage(String connectionId, String message, St return singleMessageStream(scopeGuardResponse); } return buildAgenticStreamResult(connectionId, actualUserQuestion, prepared, null, actorUsername, actorIsAdmin); - }); + })); } catch (Exception e) { log.error("Error initializing streaming chat", e); return new StreamResult(Flux.empty(), Flux.empty(), Flux.error(e), Flux.error(e)); @@ -3486,7 +3488,13 @@ private StreamResult buildAgenticStreamResult( } CompletableFuture.runAsync(() -> { - QueryActorContextHolder.withActor(actorUsername, () -> { + // Both contexts are re-established here rather than inherited. The controller + // returns its Flux immediately, so the servlet request — and any thread-local + // set during it — is long gone by the time this runs. The actor was already + // re-set for exactly this reason; usage attribution needs the same treatment, + // or every streamed turn bills to "unknown". + QueryActorContextHolder.withActor(actorUsername, () -> + LlmUsageContext.with("chat", connectionId, () -> { AgentProgressListener progressListener = event -> { String progressJson = buildAgentProgressJson(event); if (progressJson != null) { @@ -3530,7 +3538,7 @@ private StreamResult buildAgenticStreamResult( resultSink.tryEmitComplete(); tokenSink.tryEmitComplete(); } - }); + })); }); return new StreamResult( diff --git a/backend/src/main/java/com/dbaagent/service/DashboardAlertService.java b/backend/src/main/java/com/dbaagent/service/DashboardAlertService.java index fc5a4ec..67abe73 100644 --- a/backend/src/main/java/com/dbaagent/service/DashboardAlertService.java +++ b/backend/src/main/java/com/dbaagent/service/DashboardAlertService.java @@ -4,6 +4,7 @@ import com.dbaagent.model.SavedDashboard; import com.dbaagent.repository.DashboardAlertRepository; import com.dbaagent.repository.SavedDashboardRepository; +import com.dbaagent.service.llm.LlmUsageContext; import jakarta.mail.MessagingException; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -107,7 +108,16 @@ public void evaluate(UUID alertId) { DashboardAlert alert = requireAlert(alertId); alert.setLastCheckedAt(LocalDateTime.now()); try { - Verdict verdict = runCheck(alert); + // Scheduled work has no request, so the usage filter cannot label it. Declared + // here instead, or every alert evaluation would bill to "unknown" — and alerts + // are the one feature that spends money with nobody watching. + // + // The actor is set for the same reason the check itself runs as the creator: + // the spend belongs to whoever owns the alert, not to nobody. + Verdict verdict = QueryActorContextHolder.withActor( + alert.getCreatedByUsername(), + () -> LlmUsageContext.with( + "dashboard-alert", alert.getConnectionId(), () -> runCheck(alert))); alert.setLastVerdict(verdict.fired() ? "FIRED" : "OK"); alert.setLastReason(verdict.reason()); alert.setLastError(null); diff --git a/backend/src/main/java/com/dbaagent/service/EmbeddingService.java b/backend/src/main/java/com/dbaagent/service/EmbeddingService.java index a2d1742..de8ad80 100644 --- a/backend/src/main/java/com/dbaagent/service/EmbeddingService.java +++ b/backend/src/main/java/com/dbaagent/service/EmbeddingService.java @@ -6,6 +6,8 @@ import com.dbaagent.llm.api.LlmEmbeddingProvider; import com.dbaagent.llm.api.LlmErrorCategory; import com.dbaagent.llm.api.LlmNotConfiguredException; +import com.dbaagent.model.LlmUsageRole; +import com.dbaagent.service.llm.LlmUsageRecorder; import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Value; @@ -37,6 +39,9 @@ public class EmbeddingService { private static final int MIN_EMBEDDING_CHARS = 1_000; private static final long RETRY_BACKOFF_MS = 750L; + /** See {@link #recordEmbeddingUsage}: schema text is denser than prose. */ + private static final int ESTIMATED_CHARS_PER_TOKEN = 3; + /** * Ceiling on a provider-supplied {@code Retry-After}. Brain init embeds documents on * a bounded worker pool, so a long sleep here stalls the whole stage; 30s is longer @@ -61,6 +66,18 @@ public class EmbeddingService { private final ConcurrentHashMap dimensionsBySignature = new ConcurrentHashMap<>(); + /** + * Usage accounting. Injected through a setter rather than the constructor so the three + * existing test seams below keep working unchanged, and so a null recorder (any test + * constructing this directly) simply means "do not account". + */ + private LlmUsageRecorder usageRecorder; + + @Autowired(required = false) + public void setUsageRecorder(LlmUsageRecorder usageRecorder) { + this.usageRecorder = usageRecorder; + } + @Autowired public EmbeddingService( LlmConfigResolver resolver, @@ -116,10 +133,21 @@ private String truncate(String text, int budget) { public List createEmbedding(String text) { LlmCredentials credentials = requireCredentials(); LlmEmbeddingProvider provider = registry.embeddingProvider(credentials.providerId()); - return embedWithShrink(provider, credentials, - budget -> provider.embed(truncate(text, budget), credentials), - List.of(), - text == null ? 0 : text.length()); + long startedAt = System.nanoTime(); + try { + List result = embedWithShrink(provider, credentials, + budget -> provider.embed(truncate(text, budget), credentials), + List.of(), + text == null ? 0 : text.length()); + // Fail-open returns an empty vector for a call that actually failed. Recording + // that as a success would put phantom tokens on the ledger for work the + // provider never did. + recordEmbeddingUsage(credentials, charsOf(text), startedAt, !result.isEmpty(), null); + return result; + } catch (RuntimeException e) { + recordEmbeddingUsage(credentials, charsOf(text), startedAt, false, e); + throw e; + } } /** @@ -194,11 +222,72 @@ public List> createEmbeddings(List texts) { // before the batch fits. int longest = texts.stream().filter(java.util.Objects::nonNull) .mapToInt(String::length).max().orElse(0); - return embedWithShrink(provider, credentials, - budget -> provider.embedBatch( - texts.stream().map(t -> truncate(t, budget)).toList(), credentials), - Collections.nCopies(texts.size(), List.of()), - longest); + long totalChars = texts.stream().filter(java.util.Objects::nonNull) + .mapToLong(String::length).sum(); + long startedAt = System.nanoTime(); + try { + List> result = embedWithShrink(provider, credentials, + budget -> provider.embedBatch( + texts.stream().map(t -> truncate(t, budget)).toList(), credentials), + Collections.nCopies(texts.size(), List.of()), + longest); + boolean succeeded = result.stream().anyMatch(v -> !v.isEmpty()); + recordEmbeddingUsage(credentials, totalChars, startedAt, succeeded, null); + return result; + } catch (RuntimeException e) { + recordEmbeddingUsage(credentials, totalChars, startedAt, false, e); + throw e; + } + } + + private static long charsOf(String text) { + return text == null ? 0L : text.length(); + } + + /** + * Records one embedding call with estimated token counts. + * + *

{@link LlmEmbeddingProvider} returns vectors and nothing else — no usage block — + * so unlike the chat path there is no metered figure to store. Rather than leave + * embeddings out of the ledger (they are a real and, during a Brain rebuild, large + * cost) the count is derived from input length and the row is flagged + * {@code estimated}, so a reconciliation against the vendor invoice can tell the two + * halves apart. + * + *

The divisor is deliberately 3, not the widely-quoted 4. This service embeds + * schema and relationship documents — dense identifiers, underscores, punctuation — + * which tokenize closer to 2-3 characters per token, a fact {@link #truncate} already + * documents from the same content. Four would understate the bill on exactly the + * workload that dominates it. + * + *

Never throws: accounting must not be able to fail an indexing run. + */ + private void recordEmbeddingUsage(LlmCredentials credentials, long chars, + long startedAt, boolean succeeded, RuntimeException failure) { + if (usageRecorder == null) { + return; + } + try { + long promptTokens = succeeded ? Math.max(0, chars) / ESTIMATED_CHARS_PER_TOKEN : 0; + String errorCategory = failure == null ? null + : String.valueOf(registry.embeddingProvider(credentials.providerId()) + .classify(failure)); + + usageRecorder.record(new LlmUsageRecorder.Call( + LlmUsageRole.EMBEDDING, + credentials.providerId(), + credentials.getOrDefault("model", "unknown"), + promptTokens, + 0, + promptTokens, + 0, + true, + (System.nanoTime() - startedAt) / 1_000_000L, + succeeded, + errorCategory)); + } catch (RuntimeException e) { + log.warn("Could not record embedding usage; the embedding itself was unaffected", e); + } } /** diff --git a/backend/src/main/java/com/dbaagent/service/llm/LlmPricingService.java b/backend/src/main/java/com/dbaagent/service/llm/LlmPricingService.java new file mode 100644 index 0000000..e56a55d --- /dev/null +++ b/backend/src/main/java/com/dbaagent/service/llm/LlmPricingService.java @@ -0,0 +1,206 @@ +package com.dbaagent.service.llm; + +import com.dbaagent.repository.SystemConfigRepository; +import com.dbaagent.service.SystemConfigService; +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Service; + +import java.math.BigDecimal; +import java.math.RoundingMode; +import java.util.List; +import java.util.Optional; + +/** + * Turns token counts into a USD estimate using rates held in {@code system_config}. + * + *

Rates live in configuration rather than in a constant table on purpose: vendors + * reprice, self-hosters run their own models at their own cost, and an operator who + * discovers their spend numbers are wrong must be able to correct them without waiting + * for a release. Keys are + * + * 

+ *   llm.pricing.<model>.input-per-1m
+ *   llm.pricing.<model>.output-per-1m
+ *   llm.pricing.<model>.cached-input-per-1m   (optional; defaults to input rate)
+ *
+ * + *

There are deliberately no built-in default prices. A stale bundled + * price list is worse than none: it produces confident wrong totals that nobody thinks to + * check, and this repo has already been bitten once by a credential default shipped in a + * properties file. An unpriced model yields {@link Optional#empty()}, which the caller + * stores as a null cost and the UI reports as unpriced. + */ +@Service +@Slf4j +public class LlmPricingService { + + private static final BigDecimal PER_MILLION = new BigDecimal("1000000"); + private static final int COST_SCALE = 6; + + static final String KEY_PREFIX = "llm.pricing."; + static final String INPUT = "input-per-1m"; + static final String OUTPUT = "output-per-1m"; + static final String CACHED_INPUT = "cached-input-per-1m"; + + private final SystemConfigService systemConfig; + private final SystemConfigRepository configRepository; + + public LlmPricingService(SystemConfigService systemConfig, + SystemConfigRepository configRepository) { + this.systemConfig = systemConfig; + this.configRepository = configRepository; + } + + /** + * Cost for one call, or empty when the model has no configured input rate. + * + *

{@code cachedPromptTokens} is billed at the cached rate and the remainder of the + * prompt at the full input rate. A provider that reports more cached tokens than + * prompt tokens (or a caller that miscounts) would otherwise produce a negative + * fresh-token count and undercharge, so the split is clamped at zero. + */ + public Optional estimateCost( + String model, long promptTokens, long completionTokens, long cachedPromptTokens) { + if (model == null || model.isBlank()) { + return Optional.empty(); + } + + Optional inputRate = rate(model, INPUT); + if (inputRate.isEmpty()) { + return Optional.empty(); + } + + // An output rate is optional so an embedding model, which has no completion side, + // needs only one key configured. + BigDecimal outputRate = rate(model, OUTPUT).orElse(BigDecimal.ZERO); + BigDecimal cachedRate = rate(model, CACHED_INPUT).orElse(inputRate.get()); + + long cached = Math.max(0, Math.min(cachedPromptTokens, promptTokens)); + long fresh = Math.max(0, promptTokens - cached); + + BigDecimal cost = perMillion(fresh, inputRate.get()) + .add(perMillion(cached, cachedRate)) + .add(perMillion(Math.max(0, completionTokens), outputRate)); + + return Optional.of(cost.setScale(COST_SCALE, RoundingMode.HALF_UP)); + } + + /** The three rates configured for one model; a null field means "not set". */ + public record ModelRates( + String model, + BigDecimal inputPer1m, + BigDecimal outputPer1m, + BigDecimal cachedInputPer1m) { + + public boolean isPriced() { + return inputPer1m != null; + } + } + + /** + * Model names that have at least one pricing key written, parsed back out of the key + * namespace. + * + *

A model name may itself contain dots ({@code gpt-5.4}), so the suffix is stripped + * from the end rather than splitting on the first separator — the naive split yields + * {@code gpt-5} and silently loses the row. + */ + public List configuredModels() { + return configRepository.findKeysByPrefix(KEY_PREFIX).stream() + .map(LlmPricingService::modelFromKey) + .filter(m -> m != null && !m.isBlank()) + .distinct() + .sorted() + .toList(); + } + + private static String modelFromKey(String key) { + String remainder = key.substring(KEY_PREFIX.length()); + for (String suffix : List.of(INPUT, OUTPUT, CACHED_INPUT)) { + if (remainder.endsWith("." + suffix)) { + return remainder.substring(0, remainder.length() - suffix.length() - 1); + } + } + return null; + } + + public ModelRates ratesFor(String model) { + String normalized = normalize(model); + return new ModelRates( + normalized, + rate(normalized, INPUT).orElse(null), + rate(normalized, OUTPUT).orElse(null), + rate(normalized, CACHED_INPUT).orElse(null)); + } + + /** + * Writes the rates for one model. A null or blank field clears that rate + * rather than leaving the previous value in place — the editor sends the whole set, + * so an omitted field means the operator emptied the box. + * + *

Clearing writes an empty string instead of deleting the row: {@link #rate} + * already treats blank as absent, and there is no delete on + * {@link SystemConfigService}. Adding one for this would widen a shared service for a + * single caller's convenience. + * + *

Rejects a negative rate outright. {@link #rate} also ignores one on read, but a + * value that silently does nothing after the UI reported it saved is worse than an + * error at the point of entry. + */ + public ModelRates updateRates(String model, BigDecimal input, BigDecimal output, + BigDecimal cachedInput) { + String normalized = normalize(model); + if (normalized.isBlank()) { + throw new IllegalArgumentException("A model name is required"); + } + write(normalized, INPUT, input); + write(normalized, OUTPUT, output); + write(normalized, CACHED_INPUT, cachedInput); + return ratesFor(normalized); + } + + private void write(String model, String suffix, BigDecimal value) { + if (value != null && value.signum() < 0) { + throw new IllegalArgumentException( + "Rate for " + model + " " + suffix + " cannot be negative"); + } + systemConfig.set(key(model, suffix), + value == null ? "" : value.toPlainString(), + false, + "USD per 1M tokens"); + } + + private static String normalize(String model) { + return model == null ? "" : model.trim().toLowerCase(); + } + + private static String key(String model, String suffix) { + return KEY_PREFIX + model + "." + suffix; + } + + private BigDecimal perMillion(long tokens, BigDecimal ratePerMillion) { + return BigDecimal.valueOf(tokens) + .multiply(ratePerMillion) + .divide(PER_MILLION, COST_SCALE + 4, RoundingMode.HALF_UP); + } + + /** + * Reads one rate. A malformed value is treated as absent rather than propagated: this + * runs on the accounting path behind every model call, and a typo in a config row must + * not turn into a failed chat turn. + */ + private Optional rate(String model, String suffix) { + String key = key(normalize(model), suffix); + return systemConfig.get(key) + .filter(v -> !v.isBlank()) + .flatMap(v -> { + try { + BigDecimal parsed = new BigDecimal(v.trim()); + return parsed.signum() < 0 ? Optional.empty() : Optional.of(parsed); + } catch (NumberFormatException e) { + log.warn("Ignoring unparseable LLM rate {}={}", key, v); + return Optional.empty(); + } + }); + } +} diff --git a/backend/src/main/java/com/dbaagent/service/llm/LlmUsageAttributionFilter.java b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageAttributionFilter.java new file mode 100644 index 0000000..000e3b0 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageAttributionFilter.java @@ -0,0 +1,139 @@ +package com.dbaagent.service.llm; + +import jakarta.servlet.FilterChain; +import jakarta.servlet.ServletException; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import org.springframework.stereotype.Component; +import org.springframework.web.filter.OncePerRequestFilter; + +import java.io.IOException; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Labels every request with the feature it belongs to, so usage rows written deep inside + * the provider funnels can say which product surface spent the money. + * + *

Doing this in one filter rather than annotating call sites is deliberate. The + * alternative — wrapping each of the dozen services that reach a model — would have to be + * repeated by whoever adds the thirteenth, and a missed one is invisible: it records + * {@code unknown} and nobody notices until the breakdown is asked a question it cannot + * answer. A URI is a fact the request already carries. + * + *

Background work (scheduled alerts, boot-time indexing) has no request and so is not + * covered here; those paths declare themselves with {@link LlmUsageContext#with} directly. + */ +@Component +public class LlmUsageAttributionFilter extends OncePerRequestFilter { + + /** + * Ordered longest-prefix-first, because {@code /api/dashboards/generate} must not be + * claimed by {@code /api/dashboards}. A {@link LinkedHashMap} preserves that order; + * a plain map would make the winner depend on hash iteration. + */ + private static final Map FEATURE_BY_PREFIX = new LinkedHashMap<>(); + + static { + FEATURE_BY_PREFIX.put("/api/dashboards/generate", "dashboard-generate"); + FEATURE_BY_PREFIX.put("/api/dashboards/query", "dashboard-query"); + FEATURE_BY_PREFIX.put("/api/saved-dashboards", "dashboard"); + FEATURE_BY_PREFIX.put("/api/dashboards", "dashboard"); + FEATURE_BY_PREFIX.put("/api/agent", "agent"); + FEATURE_BY_PREFIX.put("/api/chat", "chat"); + FEATURE_BY_PREFIX.put("/api/brain", "brain"); + FEATURE_BY_PREFIX.put("/api/code-scan", "code-scan"); + FEATURE_BY_PREFIX.put("/api/training", "training"); + FEATURE_BY_PREFIX.put("/api/explain", "explain"); + FEATURE_BY_PREFIX.put("/api/index-advisor", "index-advisor"); + FEATURE_BY_PREFIX.put("/api/index-recommendations", "index-advisor"); + FEATURE_BY_PREFIX.put("/api/slow-quer", "slow-query"); + FEATURE_BY_PREFIX.put("/api/playbooks", "playbook"); + FEATURE_BY_PREFIX.put("/api/company-knowledge", "company-knowledge"); + FEATURE_BY_PREFIX.put("/api/schema", "schema"); + FEATURE_BY_PREFIX.put("/api/llm/v1", "llm-proxy"); + FEATURE_BY_PREFIX.put("/api/mcp", "mcp"); + } + + /** {@code /connection/{id}} and {@code /connections/{id}} path forms. */ + private static final Pattern CONNECTION_IN_PATH = + Pattern.compile("/connections?/([^/?]+)"); + + /** + * Segments that follow {@code /connections/} but name an action rather than a + * connection. Without this, {@code POST /api/connections/test} would attribute usage + * to a connection called "test" — a plausible-looking id that belongs to nothing, and + * therefore worse than recording no connection at all. + */ + private static final Set NOT_A_CONNECTION_ID = + Set.of("test", "new", "create", "search", "all", "list"); + + @Override + protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, + FilterChain chain) throws ServletException, IOException { + String feature = featureFor(request.getRequestURI()); + String connectionId = connectionFor(request); + try { + LlmUsageContext.with(feature, connectionId, () -> { + try { + chain.doFilter(request, response); + } catch (IOException | ServletException e) { + throw new FilterFailure(e); + } + return null; + }); + } catch (FilterFailure wrapper) { + // Unwrap so the container sees the exception the chain actually threw; a + // wrapped one would break error handling further up. + Throwable cause = wrapper.getCause(); + if (cause instanceof IOException io) { + throw io; + } + throw (ServletException) cause; + } + } + + /** Lets a checked exception cross the {@code Supplier} boundary unchanged. */ + private static final class FilterFailure extends RuntimeException { + FilterFailure(Throwable cause) { + super(cause); + } + } + + private static String featureFor(String uri) { + if (uri == null) { + return LlmUsageContext.UNKNOWN_FEATURE; + } + for (Map.Entry entry : FEATURE_BY_PREFIX.entrySet()) { + if (uri.startsWith(entry.getKey())) { + return entry.getValue(); + } + } + return LlmUsageContext.UNKNOWN_FEATURE; + } + + /** + * Best-effort connection id, from the path or a query parameter. The request body is + * deliberately not read: it is consumable once, and draining it here to label a usage + * row would break every controller that expects to parse it. + */ + private static String connectionFor(HttpServletRequest request) { + String param = request.getParameter("connectionId"); + if (param != null && !param.isBlank()) { + return param; + } + String uri = request.getRequestURI(); + if (uri == null) { + return null; + } + Matcher matcher = CONNECTION_IN_PATH.matcher(uri); + if (!matcher.find()) { + return null; + } + String candidate = matcher.group(1); + return NOT_A_CONNECTION_ID.contains(candidate.toLowerCase()) ? null : candidate; + } +} diff --git a/backend/src/main/java/com/dbaagent/service/llm/LlmUsageContext.java b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageContext.java new file mode 100644 index 0000000..16e3068 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageContext.java @@ -0,0 +1,79 @@ +package com.dbaagent.service.llm; + +import java.util.function.Supplier; + +/** + * Which feature is currently calling the model, and on whose connection. + * + *

The recording funnels sit at the provider boundary, where the caller is no longer + * visible — {@code RefreshableChatModel} sees a {@code Prompt} and nothing else. Rather + * than thread a parameter through every intermediate signature (Spring AI's + * {@code ChatModel} contract has no room for one anyway), features declare themselves + * around the call, exactly as {@code QueryActorContextHolder} already does for the SQL + * actor. + * + *

Attribution is best-effort by construction: an un-annotated caller records + * {@code unknown} rather than failing. Accounting must never be able to break the feature + * it is measuring. + */ +public final class LlmUsageContext { + + public static final String UNKNOWN_FEATURE = "unknown"; + + public record Scope(String feature, String connectionId) {} + + /** + * Inheritable, and that is not a detail. Chat fans its work out across + * {@code CompletableFuture.supplyAsync} and reactive schedulers, so a plain + * {@link ThreadLocal} is invisible by the time the model is actually called — verified + * live, not reasoned about: every row landed with {@code feature = 'unknown'} while a + * single-threaded unit test of the same filter passed. An + * {@link InheritableThreadLocal} is copied into threads created from the request + * thread, which is what the async fan-out does. + * + *

It is still not total. A thread taken from a pool that was created *before* this + * scope was set inherits nothing, so some background work will read {@code unknown} — + * acceptable for attribution metadata, which is why {@link #currentFeature} degrades + * instead of failing. Anything that must be labelled exactly declares its own scope, + * as {@code DashboardAlertService} does. + */ + private static final ThreadLocal CURRENT = new InheritableThreadLocal<>(); + + private LlmUsageContext() { + } + + public static Scope current() { + return CURRENT.get(); + } + + public static String currentFeature() { + Scope scope = CURRENT.get(); + return scope == null || scope.feature() == null ? UNKNOWN_FEATURE : scope.feature(); + } + + public static String currentConnectionId() { + Scope scope = CURRENT.get(); + return scope == null ? null : scope.connectionId(); + } + + public static T with(String feature, String connectionId, Supplier supplier) { + Scope previous = CURRENT.get(); + try { + CURRENT.set(new Scope(feature, connectionId)); + return supplier.get(); + } finally { + if (previous == null) { + CURRENT.remove(); + } else { + CURRENT.set(previous); + } + } + } + + public static void with(String feature, String connectionId, Runnable runnable) { + with(feature, connectionId, () -> { + runnable.run(); + return null; + }); + } +} diff --git a/backend/src/main/java/com/dbaagent/service/llm/LlmUsageQueryService.java b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageQueryService.java new file mode 100644 index 0000000..399f9c8 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageQueryService.java @@ -0,0 +1,114 @@ +package com.dbaagent.service.llm; + +import com.dbaagent.dto.LlmUsageDailyPoint; +import com.dbaagent.dto.LlmUsageGroup; +import com.dbaagent.dto.LlmUsageTotals; +import com.dbaagent.model.LlmUsage; +import com.dbaagent.repository.LlmUsageRepository; +import org.springframework.data.domain.Page; +import org.springframework.data.domain.PageRequest; +import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; + +import java.time.LocalDateTime; +import java.util.Comparator; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; + +/** Read side of LLM accounting. */ +@Service +@Transactional(readOnly = true) +public class LlmUsageQueryService { + + /** Caps the reporting window so a hand-edited query string cannot scan the whole table. */ + private static final int MAX_WINDOW_DAYS = 365; + private static final int MAX_PAGE_SIZE = 200; + + private final LlmUsageRepository repository; + private final LlmPricingService pricing; + + public LlmUsageQueryService(LlmUsageRepository repository, LlmPricingService pricing) { + this.repository = repository; + this.pricing = pricing; + } + + public record Summary( + int windowDays, + LlmUsageTotals totals, + List byFeature, + List byUser, + List byModel, + List daily, + List unpricedModels) {} + + public Summary summary(int requestedDays) { + int days = clampDays(requestedDays); + LocalDateTime since = LocalDateTime.now().minusDays(days); + + LlmUsageTotals totals = repository.totalsSince(since); + return new Summary( + days, + totals == null ? LlmUsageTotals.empty() : totals, + repository.byFeatureSince(since), + repository.byUserSince(since), + repository.byModelSince(since), + repository.dailySince(since), + repository.unpricedModelsSince(since)); + } + + public Page recent(int requestedDays, int page, int size) { + LocalDateTime since = LocalDateTime.now().minusDays(clampDays(requestedDays)); + int bounded = Math.min(Math.max(size, 1), MAX_PAGE_SIZE); + return repository.findByCreatedAtGreaterThanEqualOrderByCreatedAtDesc( + since, PageRequest.of(Math.max(page, 0), bounded)); + } + + /** One editable row in the pricing table. */ + public record PricingRow( + String model, + java.math.BigDecimal inputPer1m, + java.math.BigDecimal outputPer1m, + java.math.BigDecimal cachedInputPer1m, + boolean priced, + boolean seenInUsage) {} + + /** + * Every model worth pricing: those the ledger has recorded, plus any model that has a + * rate configured but has not been called (so a rate set ahead of a rollout, or left + * behind by a model that was retired, is still visible and removable). + * + *

Unpriced models sort first — they are the ones costing money the totals do not + * show, and the whole point of this screen is to close that gap. + */ + public List pricing() { + List seen = repository.distinctModels(); + Set models = new LinkedHashSet<>(seen); + models.addAll(pricing.configuredModels()); + + return models.stream() + .map(model -> { + LlmPricingService.ModelRates rates = pricing.ratesFor(model); + return new PricingRow( + rates.model(), + rates.inputPer1m(), + rates.outputPer1m(), + rates.cachedInputPer1m(), + rates.isPriced(), + seen.contains(model)); + }) + .sorted(Comparator.comparing(PricingRow::priced) + .thenComparing(PricingRow::model)) + .toList(); + } + + @Transactional + public int purgeOlderThan(int days) { + return repository.deleteByCreatedAtBefore( + LocalDateTime.now().minusDays(Math.max(days, 1))); + } + + private int clampDays(int days) { + return Math.min(Math.max(days, 1), MAX_WINDOW_DAYS); + } +} diff --git a/backend/src/main/java/com/dbaagent/service/llm/LlmUsageRecorder.java b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageRecorder.java new file mode 100644 index 0000000..15baec1 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageRecorder.java @@ -0,0 +1,117 @@ +package com.dbaagent.service.llm; + +import com.dbaagent.model.LlmUsage; +import com.dbaagent.model.LlmUsageRole; +import com.dbaagent.service.QueryActorContextHolder; +import lombok.extern.slf4j.Slf4j; +import org.springframework.security.core.Authentication; +import org.springframework.security.core.context.SecurityContextHolder; +import org.springframework.stereotype.Service; + +import java.math.BigDecimal; + +/** + * Writes one {@link LlmUsage} row per model call. + * + *

This class must never break the call it is measuring. Accounting is + * strictly secondary to the feature: every public method swallows its own failures and + * logs them. A full disk, a migration not yet applied, or a bug in this file must cost the + * operator their spend numbers, never their chat. + * + *

Rows are written in their own transaction, through {@link LlmUsageWriter}. The caller + * is often mid-transaction and may go on to fail — a chat turn that rolls back still spent + * the tokens, and a ledger that discards exactly the calls that failed would understate + * spend precisely where an operator is investigating. + */ +@Service +@Slf4j +public class LlmUsageRecorder { + + private final LlmUsageWriter writer; + private final LlmPricingService pricing; + + public LlmUsageRecorder(LlmUsageWriter writer, LlmPricingService pricing) { + this.writer = writer; + this.pricing = pricing; + } + + /** A completed call, ready to be priced and stored. */ + public record Call( + LlmUsageRole role, + String providerId, + String model, + long promptTokens, + long completionTokens, + long totalTokens, + long cachedPromptTokens, + boolean estimated, + long latencyMs, + boolean succeeded, + String errorCategory) {} + + public void record(Call call) { + try { + writer.write(build(call)); + } catch (RuntimeException e) { + log.warn("Could not record LLM usage for {} {}; the call itself was unaffected", + call.role(), call.model(), e); + } + } + + private LlmUsage build(Call call) { + LlmUsage row = new LlmUsage(); + row.setRoleEnum(call.role()); + row.setProviderId(call.providerId()); + row.setModel(call.model()); + row.setFeature(LlmUsageContext.currentFeature()); + row.setConnectionId(LlmUsageContext.currentConnectionId()); + row.setUsername(resolveUsername()); + row.setPromptTokens(Math.max(0, call.promptTokens())); + row.setCompletionTokens(Math.max(0, call.completionTokens())); + row.setCachedPromptTokens(Math.max(0, call.cachedPromptTokens())); + + // Providers vary on whether total is reported; derive it when it is missing rather + // than storing a zero that would silently drop the call out of every token sum. + long total = call.totalTokens() > 0 + ? call.totalTokens() + : row.getPromptTokens() + row.getCompletionTokens(); + row.setTotalTokens(total); + + row.setEstimated(call.estimated()); + row.setLatencyMs(call.latencyMs()); + row.setSucceeded(call.succeeded()); + row.setErrorCategory(call.errorCategory()); + row.setEstimatedCostUsd(price(call, row)); + return row; + } + + private BigDecimal price(Call call, LlmUsage row) { + return pricing.estimateCost( + call.model(), + row.getPromptTokens(), + row.getCompletionTokens(), + row.getCachedPromptTokens()).orElse(null); + } + + /** + * Who to bill. The SQL actor context wins over the security context because it is + * already the product's answer to "who is really acting" — under an admin's View as, + * the security principal is the admin while the actor is the target user, and usage + * belongs to whoever the work was done for. + * + *

Null for background work with no actor at all, which is a real state (scheduled + * alert evaluation, boot-time indexing) and not an error. + */ + private String resolveUsername() { + String actor = QueryActorContextHolder.currentUsername(); + if (actor != null && !actor.isBlank()) { + return actor; + } + Authentication auth = SecurityContextHolder.getContext().getAuthentication(); + if (auth == null || !auth.isAuthenticated()) { + return null; + } + String name = auth.getName(); + return (name == null || name.isBlank() || "anonymousUser".equals(name)) ? null : name; + } +} diff --git a/backend/src/main/java/com/dbaagent/service/llm/LlmUsageWriter.java b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageWriter.java new file mode 100644 index 0000000..6748908 --- /dev/null +++ b/backend/src/main/java/com/dbaagent/service/llm/LlmUsageWriter.java @@ -0,0 +1,31 @@ +package com.dbaagent.service.llm; + +import com.dbaagent.model.LlmUsage; +import com.dbaagent.repository.LlmUsageRepository; +import org.springframework.stereotype.Component; +import org.springframework.transaction.annotation.Propagation; +import org.springframework.transaction.annotation.Transactional; + +/** + * The transactional half of usage recording, kept in its own bean deliberately. + * + *

{@code REQUIRES_NEW} is applied by a Spring proxy, and a proxy is bypassed on + * self-invocation — so annotating a private method that {@link LlmUsageRecorder} calls + * through {@code this} would silently join the caller's transaction instead of starting + * its own, and the row would roll back with a failed chat turn. That is exactly the + * failure this propagation exists to prevent, and it would leave no trace. + */ +@Component +public class LlmUsageWriter { + + private final LlmUsageRepository repository; + + public LlmUsageWriter(LlmUsageRepository repository) { + this.repository = repository; + } + + @Transactional(propagation = Propagation.REQUIRES_NEW) + public void write(LlmUsage row) { + repository.save(row); + } +} diff --git a/backend/src/main/resources/db/migration/V119__create_llm_usage.sql b/backend/src/main/resources/db/migration/V119__create_llm_usage.sql new file mode 100644 index 0000000..cbcb4a8 --- /dev/null +++ b/backend/src/main/resources/db/migration/V119__create_llm_usage.sql @@ -0,0 +1,25 @@ +CREATE TABLE llm_usage ( + id BIGSERIAL PRIMARY KEY, + role VARCHAR(16) NOT NULL, + provider_id VARCHAR(255) NOT NULL, + model VARCHAR(255) NOT NULL, + feature VARCHAR(64) NOT NULL, + username VARCHAR(255), + connection_id VARCHAR(255), + prompt_tokens BIGINT NOT NULL DEFAULT 0, + completion_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + cached_prompt_tokens BIGINT NOT NULL DEFAULT 0, + estimated_cost_usd NUMERIC(12, 6), + estimated BOOLEAN NOT NULL DEFAULT false, + latency_ms BIGINT, + succeeded BOOLEAN NOT NULL DEFAULT true, + error_category VARCHAR(64), + created_at TIMESTAMP NOT NULL DEFAULT now() +); + +-- Every rollup filters on a time window first, so created_at leads each index. +CREATE INDEX idx_llm_usage_created_at ON llm_usage(created_at DESC); +CREATE INDEX idx_llm_usage_username ON llm_usage(username, created_at DESC); +CREATE INDEX idx_llm_usage_feature ON llm_usage(feature, created_at DESC); +CREATE INDEX idx_llm_usage_connection ON llm_usage(connection_id, created_at DESC); diff --git a/backend/src/test/java/com/dbaagent/config/RefreshableChatModelUsageTest.java b/backend/src/test/java/com/dbaagent/config/RefreshableChatModelUsageTest.java new file mode 100644 index 0000000..811875c --- /dev/null +++ b/backend/src/test/java/com/dbaagent/config/RefreshableChatModelUsageTest.java @@ -0,0 +1,229 @@ +package com.dbaagent.config; + +import com.dbaagent.llm.LlmConfigResolver; +import com.dbaagent.llm.LlmProviderRegistry; +import com.dbaagent.llm.api.LlmChatProvider; +import com.dbaagent.llm.api.LlmCredentials; +import com.dbaagent.llm.api.LlmErrorCategory; +import com.dbaagent.model.LlmUsageRole; +import com.dbaagent.service.llm.LlmUsageRecorder; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.metadata.DefaultUsage; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Token accounting on the chat funnel. + * + *

The streaming cases are the point of this file. Usage on a stream arrives on one late + * chunk while every earlier chunk carries none, so the two ways to get it wrong — a row + * per chunk, or summing partials into a running total — are both easy to write and + * invisible in production until the invoice disagrees. + */ +class RefreshableChatModelUsageTest { + + private static final Map BASE = + Map.of("endpoint", "https://x.invalid/", "model", "gpt-4o"); + + private final LlmConfigResolver resolver = mock(LlmConfigResolver.class); + private final LlmProviderRegistry registry = mock(LlmProviderRegistry.class); + private final LlmChatProvider provider = mock(LlmChatProvider.class); + private final ChatModel delegate = mock(ChatModel.class); + private final LlmUsageRecorder recorder = mock(LlmUsageRecorder.class); + + private RefreshableChatModel model() { + when(resolver.resolveChat()).thenReturn(new LlmCredentials("openai", BASE)); + when(registry.chatProvider("openai")).thenReturn(provider); + when(provider.delegate(any())).thenReturn(delegate); + return new RefreshableChatModel(resolver, registry, recorder); + } + + private static ChatResponse withUsage(String text, int prompt, int completion) { + return new ChatResponse( + List.of(new Generation(new AssistantMessage(text))), + ChatResponseMetadata.builder() + .usage(new DefaultUsage(prompt, completion)) + .model("gpt-4o-2024-11-20") + .build()); + } + + private static ChatResponse noUsage(String text) { + return new ChatResponse(List.of(new Generation(new AssistantMessage(text)))); + } + + private LlmUsageRecorder.Call captureOne() { + ArgumentCaptor captor = + ArgumentCaptor.forClass(LlmUsageRecorder.Call.class); + verify(recorder).record(captor.capture()); + return captor.getValue(); + } + + // ── Non-streaming ───────────────────────────────────────────────────────── + + @Test + void recordsMeteredTokensFromTheResponse() { + when(delegate.call(any(Prompt.class))).thenReturn(withUsage("hi", 1_200, 300)); + + model().call(new Prompt("hi")); + + LlmUsageRecorder.Call call = captureOne(); + assertThat(call.role()).isEqualTo(LlmUsageRole.CHAT); + assertThat(call.promptTokens()).isEqualTo(1_200); + assertThat(call.completionTokens()).isEqualTo(300); + assertThat(call.succeeded()).isTrue(); + // Chat counts come from the provider, so they are never flagged as estimates. + assertThat(call.estimated()).isFalse(); + } + + /** + * The served model wins over the configured one: an alias resolves to a dated + * snapshot, and the snapshot is what was billed. + */ + @Test + void prefersTheServedModelOverTheConfiguredAlias() { + when(delegate.call(any(Prompt.class))).thenReturn(withUsage("hi", 10, 10)); + + model().call(new Prompt("hi")); + + assertThat(captureOne().model()).isEqualTo("gpt-4o-2024-11-20"); + } + + @Test + void fallsBackToTheConfiguredModelWhenTheResponseNamesNone() { + when(delegate.call(any(Prompt.class))).thenReturn(noUsage("hi")); + + model().call(new Prompt("hi")); + + assertThat(captureOne().model()).isEqualTo("gpt-4o"); + } + + @Test + void recordsAFailedCallSoRetryLoopsStayVisible() { + when(delegate.call(any(Prompt.class))).thenThrow(new RuntimeException("429 rate limited")); + when(provider.classify(any())).thenReturn(LlmErrorCategory.RATE_LIMIT); + + assertThatThrownBy(() -> model().call(new Prompt("hi"))) + .isInstanceOf(RuntimeException.class); + + LlmUsageRecorder.Call call = captureOne(); + assertThat(call.succeeded()).isFalse(); + assertThat(call.errorCategory()).isEqualTo("RATE_LIMIT"); + } + + /** Accounting must never convert a working call into a failed one. */ + @Test + void aRecorderFailureDoesNotBreakTheCall() { + when(delegate.call(any(Prompt.class))).thenReturn(withUsage("hi", 10, 10)); + org.mockito.Mockito.doThrow(new RuntimeException("ledger down")) + .when(recorder).record(any()); + + assertThat(model().call(new Prompt("hi"))).isNotNull(); + } + + @Test + void worksWithNoRecorderAtAll() { + when(resolver.resolveChat()).thenReturn(new LlmCredentials("openai", BASE)); + when(registry.chatProvider("openai")).thenReturn(provider); + when(provider.delegate(any())).thenReturn(delegate); + when(delegate.call(any(Prompt.class))).thenReturn(noUsage("hi")); + + assertThat(new RefreshableChatModel(resolver, registry).call(new Prompt("hi"))) + .isNotNull(); + } + + // ── Streaming ───────────────────────────────────────────────────────────── + + /** + * One row per stream, not one per chunk. A 200-chunk answer recorded per chunk would + * report 200 calls and, if partials were summed, a wildly inflated token total. + */ + @Test + void recordsExactlyOneRowForAWholeStream() { + when(delegate.stream(any(Prompt.class))).thenReturn(Flux.just( + noUsage("Hel"), noUsage("lo"), noUsage(" wor"), + withUsage("ld", 900, 120))); + + assertThat(model().stream(new Prompt("hi")).collectList().block()).hasSize(4); + + LlmUsageRecorder.Call call = captureOne(); + assertThat(call.promptTokens()).isEqualTo(900); + assertThat(call.completionTokens()).isEqualTo(120); + } + + /** + * Providers that report a cumulative running total must not have their partials added + * to the final figure. + */ + @Test + void takesTheLastReportedUsageRatherThanSummingChunks() { + when(delegate.stream(any(Prompt.class))).thenReturn(Flux.just( + withUsage("a", 900, 10), + withUsage("b", 900, 60), + withUsage("c", 900, 120))); + + assertThat(model().stream(new Prompt("hi")).collectList().block()).hasSize(3); + + LlmUsageRecorder.Call call = captureOne(); + assertThat(call.promptTokens()).isEqualTo(900); + assertThat(call.completionTokens()).isEqualTo(120); + verify(recorder, times(1)).record(any()); + } + + /** A stream that dies partway still consumed prompt tokens. */ + @Test + void recordsAStreamThatFailsAfterEmitting() { + when(delegate.stream(any(Prompt.class))).thenReturn(Flux.concat( + Flux.just(withUsage("partial", 500, 5)), + Flux.error(new RuntimeException("connection reset")))); + when(provider.classify(any())).thenReturn(LlmErrorCategory.TRANSIENT); + + assertThatThrownBy(() -> model().stream(new Prompt("hi")).collectList().block()) + .isInstanceOf(RuntimeException.class); + + LlmUsageRecorder.Call call = captureOne(); + assertThat(call.promptTokens()).isEqualTo(500); + assertThat(call.succeeded()).isFalse(); + } + + /** A cancelled dashboard build is exactly the silent spend an operator wants logged. */ + @Test + void recordsACancelledStream() { + when(delegate.stream(any(Prompt.class))).thenReturn(Flux.just( + withUsage("a", 400, 10), noUsage("b"), noUsage("c"))); + + // take(1) cancels the source after the first element, which is what a user + // navigating away mid-answer does to the stream. + assertThat(model().stream(new Prompt("hi")).take(1).collectList().block()).hasSize(1); + + assertThat(captureOne().promptTokens()).isEqualTo(400); + } + + @Test + void recordsAStreamThatReportsNoUsageAtAll() { + when(delegate.stream(any(Prompt.class))).thenReturn(Flux.just(noUsage("a"), noUsage("b"))); + + assertThat(model().stream(new Prompt("hi")).collectList().block()).hasSize(2); + + // Still one row: the call happened and is worth counting, even unmetered. + LlmUsageRecorder.Call call = captureOne(); + assertThat(call.promptTokens()).isZero(); + assertThat(call.succeeded()).isTrue(); + } +} diff --git a/backend/src/test/java/com/dbaagent/llm/openai/ResponsesApiChatModelUsageTest.java b/backend/src/test/java/com/dbaagent/llm/openai/ResponsesApiChatModelUsageTest.java new file mode 100644 index 0000000..d1b0e23 --- /dev/null +++ b/backend/src/test/java/com/dbaagent/llm/openai/ResponsesApiChatModelUsageTest.java @@ -0,0 +1,151 @@ +package com.dbaagent.llm.openai; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.prompt.Prompt; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Token counts parsed off a provider response. + * + *

{@code buildMetadata} previously discarded the {@code usage} block entirely. That was + * invisible while nothing read it and became a silent zero the moment usage accounting + * started summing it: real, billed calls recorded 0 tokens and $0.00. These cases pin both + * vendor dialects so the next refactor cannot quietly restore the stub. + */ +class ResponsesApiChatModelUsageTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static ResponsesApiChatModel model(boolean useResponsesApi) { + return new ResponsesApiChatModel("https://api.openai.com/v1", "test-key", "gpt-4o", + "2025-03-01-preview", 1.0, useResponsesApi); + } + + private static ChatResponseMetadata metadataFrom(String usageJson) throws Exception { + return model(false).buildMetadata(MAPPER.readTree(usageJson)); + } + + @Test + void readsChatCompletionsTokenNames() throws Exception { + ChatResponseMetadata metadata = metadataFrom(""" + {"prompt_tokens": 1200, "completion_tokens": 300, "total_tokens": 1500} + """); + + assertThat(metadata.getUsage().getPromptTokens()).isEqualTo(1200); + assertThat(metadata.getUsage().getCompletionTokens()).isEqualTo(300); + assertThat(metadata.getUsage().getTotalTokens()).isEqualTo(1500); + } + + /** The Responses API names the same numbers differently; both must work. */ + @Test + void readsResponsesApiTokenNames() throws Exception { + ChatResponseMetadata metadata = metadataFrom(""" + {"input_tokens": 800, "output_tokens": 120, "total_tokens": 920} + """); + + assertThat(metadata.getUsage().getPromptTokens()).isEqualTo(800); + assertThat(metadata.getUsage().getCompletionTokens()).isEqualTo(120); + assertThat(metadata.getUsage().getTotalTokens()).isEqualTo(920); + } + + @Test + void derivesTotalWhenTheProviderOmitsIt() throws Exception { + ChatResponseMetadata metadata = metadataFrom(""" + {"prompt_tokens": 40, "completion_tokens": 60} + """); + + assertThat(metadata.getUsage().getTotalTokens()).isEqualTo(100); + } + + /** Cached input bills below fresh input, so it has to survive parsing. */ + @Test + void readsCachedPromptTokensFromEitherParent() throws Exception { + assertThat(metadataFrom(""" + {"prompt_tokens": 1000, "completion_tokens": 0, + "prompt_tokens_details": {"cached_tokens": 600}} + """).getUsage().getCacheReadInputTokens()).isEqualTo(600L); + + assertThat(metadataFrom(""" + {"input_tokens": 1000, "output_tokens": 0, + "input_tokens_details": {"cached_tokens": 250}} + """).getUsage().getCacheReadInputTokens()).isEqualTo(250L); + } + + @Test + void missingUsageBlockYieldsEmptyMetadataRatherThanThrowing() throws Exception { + assertThat(model(false).buildMetadata(MAPPER.readTree("{}").path("usage")).getUsage()) + .satisfiesAnyOf( + usage -> assertThat(usage).isNull(), + usage -> assertThat(usage.getTotalTokens()).isEqualTo(0)); + } + + // ── Streaming ───────────────────────────────────────────────────────────── + + /** + * Chat Completions omits usage from a stream unless asked. Without this flag a + * streamed turn is unbillable — the tokens are spent and nothing reports them. + */ + @Test + void streamingRequestsUsageOnChatCompletions() { + String body = model(false).buildRequestBody(new Prompt("hi"), true); + + assertThat(body).contains("\"stream_options\"").contains("\"include_usage\":true"); + } + + /** A non-streaming call reports usage anyway, so the flag would be noise. */ + @Test + void nonStreamingRequestDoesNotAskForStreamUsage() { + assertThat(model(false).buildRequestBody(new Prompt("hi"), false)) + .doesNotContain("stream_options"); + } + + @Test + void extractsUsageFromAFinalChatCompletionsChunk() { + ChatResponseMetadata metadata = model(false).extractStreamUsage(""" + {"choices": [], "usage": {"prompt_tokens": 900, "completion_tokens": 120, + "total_tokens": 1020}} + """); + + assertThat(metadata).isNotNull(); + assertThat(metadata.getUsage().getTotalTokens()).isEqualTo(1020); + } + + /** The Responses API nests it under the completion event's response object. */ + @Test + void extractsUsageFromAResponsesCompletedEvent() { + ChatResponseMetadata metadata = model(true).extractStreamUsage(""" + {"type": "response.completed", + "response": {"usage": {"input_tokens": 500, "output_tokens": 25, + "total_tokens": 525}}} + """); + + assertThat(metadata).isNotNull(); + assertThat(metadata.getUsage().getTotalTokens()).isEqualTo(525); + } + + /** An ordinary text delta carries no usage and must not produce a phantom row. */ + @Test + void ordinaryDeltaEventYieldsNoUsage() { + assertThat(model(false).extractStreamUsage(""" + {"choices": [{"delta": {"content": "Hel"}}]} + """)).isNull(); + } + + /** A malformed event must cost the usage row, never the user's answer. */ + @Test + void malformedEventYieldsNullRatherThanThrowing() { + assertThat(model(false).extractStreamUsage("not json at all")).isNull(); + assertThat(model(false).extractStreamUsage("")).isNull(); + } + + /** A zero-filled usage block is not worth emitting as a chunk. */ + @Test + void zeroUsageIsTreatedAsNoUsage() { + assertThat(model(false).extractStreamUsage(""" + {"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}} + """)).isNull(); + } +} diff --git a/backend/src/test/java/com/dbaagent/service/llm/LlmPricingServiceTest.java b/backend/src/test/java/com/dbaagent/service/llm/LlmPricingServiceTest.java new file mode 100644 index 0000000..2ca4650 --- /dev/null +++ b/backend/src/test/java/com/dbaagent/service/llm/LlmPricingServiceTest.java @@ -0,0 +1,255 @@ +package com.dbaagent.service.llm; + +import com.dbaagent.repository.SystemConfigRepository; +import com.dbaagent.service.SystemConfigService; +import org.junit.jupiter.api.Test; + +import java.math.BigDecimal; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class LlmPricingServiceTest { + + private final Map config = new HashMap<>(); + private final LlmPricingService pricing = + new LlmPricingService(stubConfig(), stubRepository()); + + private SystemConfigService stubConfig() { + SystemConfigService svc = mock(SystemConfigService.class); + when(svc.get(anyString())).thenAnswer( + inv -> Optional.ofNullable(config.get(inv.getArgument(0, String.class)))); + // Writes land in the same map the reads come from, so a save is observable through + // the public read path rather than by verifying a mock interaction. + org.mockito.Mockito.doAnswer(inv -> { + config.put(inv.getArgument(0), inv.getArgument(1)); + return null; + }).when(svc).set(anyString(), org.mockito.ArgumentMatchers.any(), + org.mockito.ArgumentMatchers.anyBoolean(), + org.mockito.ArgumentMatchers.any()); + return svc; + } + + private SystemConfigRepository stubRepository() { + SystemConfigRepository repo = mock(SystemConfigRepository.class); + when(repo.findKeysByPrefix(anyString())).thenAnswer(inv -> { + String prefix = inv.getArgument(0, String.class); + return config.keySet().stream().filter(k -> k.startsWith(prefix)).sorted().toList(); + }); + return repo; + } + + private void rate(String model, String suffix, String value) { + config.put("llm.pricing." + model + "." + suffix, value); + } + + @Test + void pricesInputAndOutputSeparately() { + rate("gpt-4o", "input-per-1m", "2.50"); + rate("gpt-4o", "output-per-1m", "10.00"); + + // 1M input at $2.50 + 1M output at $10.00 + assertThat(pricing.estimateCost("gpt-4o", 1_000_000, 1_000_000, 0)) + .contains(new BigDecimal("12.500000")); + } + + @Test + void chargesCachedPromptTokensAtTheCachedRate() { + rate("gpt-4o", "input-per-1m", "2.50"); + rate("gpt-4o", "output-per-1m", "10.00"); + rate("gpt-4o", "cached-input-per-1m", "1.25"); + + // 400k fresh at $2.50/M = $1.00, 600k cached at $1.25/M = $0.75 + assertThat(pricing.estimateCost("gpt-4o", 1_000_000, 0, 600_000)) + .contains(new BigDecimal("1.750000")); + } + + @Test + void cachedRateDefaultsToTheInputRateWhenUnset() { + rate("gpt-4o", "input-per-1m", "2.00"); + + // With no cached rate configured, cached tokens must not become free. + assertThat(pricing.estimateCost("gpt-4o", 1_000_000, 0, 500_000)) + .contains(new BigDecimal("2.000000")); + } + + /** + * A provider reporting more cached tokens than prompt tokens must not produce a + * negative fresh count, which would subtract from the bill. + */ + @Test + void clampsCachedTokensThatExceedThePromptCount() { + rate("m", "input-per-1m", "10.00"); + rate("m", "cached-input-per-1m", "1.00"); + + // Cached is clamped to the 1,000 prompt tokens, so all 1,000 bill at the cached + // $1.00/M — $0.001. The bug this guards against is the unclamped arithmetic, where + // fresh = 1,000 - 9,999,999 goes negative and *subtracts* from the bill. + assertThat(pricing.estimateCost("m", 1_000, 0, 9_999_999)) + .contains(new BigDecimal("0.001000")); + } + + @Test + void unpricedModelYieldsEmptyRatherThanZero() { + assertThat(pricing.estimateCost("some-local-llama", 1_000_000, 1_000_000, 0)).isEmpty(); + } + + /** An embedding model has no completion side, so one key must be enough. */ + @Test + void outputRateIsOptional() { + rate("text-embedding-3-large", "input-per-1m", "0.13"); + + assertThat(pricing.estimateCost("text-embedding-3-large", 1_000_000, 0, 0)) + .contains(new BigDecimal("0.130000")); + } + + @Test + void malformedRateIsTreatedAsUnpricedRatherThanThrowing() { + rate("m", "input-per-1m", "not-a-number"); + + assertThat(pricing.estimateCost("m", 1_000, 1_000, 0)).isEmpty(); + } + + @Test + void negativeRateIsRejected() { + rate("m", "input-per-1m", "-5.00"); + + assertThat(pricing.estimateCost("m", 1_000, 0, 0)).isEmpty(); + } + + @Test + void modelLookupIsCaseInsensitive() { + rate("gpt-4o", "input-per-1m", "2.50"); + + assertThat(pricing.estimateCost("GPT-4o", 1_000_000, 0, 0)) + .contains(new BigDecimal("2.500000")); + } + + @Test + void blankModelIsUnpriced() { + assertThat(pricing.estimateCost(" ", 1_000, 0, 0)).isEmpty(); + assertThat(pricing.estimateCost(null, 1_000, 0, 0)).isEmpty(); + } + + /** Sub-cent calls must not round away to zero at the stored scale. */ + @Test + void keepsPrecisionForSmallCalls() { + rate("gpt-4o", "input-per-1m", "2.50"); + + assertThat(pricing.estimateCost("gpt-4o", 1_000, 0, 0)) + .contains(new BigDecimal("0.002500")); + } + + // ── Editing rates ───────────────────────────────────────────────────────── + + @Test + void savedRatesAreUsedByTheNextCostCalculation() { + pricing.updateRates("gpt-4o", new BigDecimal("2.50"), new BigDecimal("10.00"), null); + + assertThat(pricing.estimateCost("gpt-4o", 1_000_000, 1_000_000, 0)) + .contains(new BigDecimal("12.500000")); + } + + @Test + void ratesForReportsWhatWasSaved() { + pricing.updateRates("gpt-4o", new BigDecimal("2.50"), new BigDecimal("10.00"), + new BigDecimal("1.25")); + + LlmPricingService.ModelRates rates = pricing.ratesFor("gpt-4o"); + assertThat(rates.inputPer1m()).isEqualByComparingTo("2.50"); + assertThat(rates.outputPer1m()).isEqualByComparingTo("10.00"); + assertThat(rates.cachedInputPer1m()).isEqualByComparingTo("1.25"); + assertThat(rates.isPriced()).isTrue(); + } + + /** Emptying a field must clear the rate, not silently keep the old one. */ + @Test + void aNullFieldClearsThatRate() { + pricing.updateRates("gpt-4o", new BigDecimal("2.50"), new BigDecimal("10.00"), null); + pricing.updateRates("gpt-4o", new BigDecimal("2.50"), null, null); + + assertThat(pricing.ratesFor("gpt-4o").outputPer1m()).isNull(); + // Output now contributes nothing, so only the input side is billed. + assertThat(pricing.estimateCost("gpt-4o", 1_000_000, 1_000_000, 0)) + .contains(new BigDecimal("2.500000")); + } + + @Test + void clearingTheInputRateMakesTheModelUnpricedAgain() { + pricing.updateRates("gpt-4o", new BigDecimal("2.50"), null, null); + pricing.updateRates("gpt-4o", null, null, null); + + assertThat(pricing.ratesFor("gpt-4o").isPriced()).isFalse(); + assertThat(pricing.estimateCost("gpt-4o", 1_000, 0, 0)).isEmpty(); + } + + @Test + void aNegativeRateIsRejectedAtWriteTime() { + assertThatThrownBy(() -> + pricing.updateRates("gpt-4o", new BigDecimal("-1"), null, null)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("negative"); + } + + @Test + void aBlankModelNameIsRejected() { + assertThatThrownBy(() -> pricing.updateRates(" ", new BigDecimal("1"), null, null)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void modelNamesAreNormalisedOnWriteSoLookupMatches() { + pricing.updateRates(" GPT-4o ", new BigDecimal("2.50"), null, null); + + assertThat(pricing.estimateCost("gpt-4o", 1_000_000, 0, 0)) + .contains(new BigDecimal("2.500000")); + } + + /** + * A model name can itself contain dots. Splitting the key on the first '.' after the + * prefix would report "gpt-5" and lose the real row. + */ + @Test + void configuredModelsHandlesDottedModelNames() { + pricing.updateRates("gpt-5.4", new BigDecimal("1.25"), new BigDecimal("10.00"), null); + + assertThat(pricing.configuredModels()).containsExactly("gpt-5.4"); + } + + @Test + void configuredModelsListsEachModelOnceAcrossItsThreeKeys() { + pricing.updateRates("gpt-4o", new BigDecimal("2.50"), new BigDecimal("10.00"), + new BigDecimal("1.25")); + pricing.updateRates("text-embedding-3-large", new BigDecimal("0.13"), null, null); + + assertThat(pricing.configuredModels()) + .containsExactly("gpt-4o", "text-embedding-3-large"); + } + + @Test + void configuredModelsIsEmptyWhenNothingIsPriced() { + assertThat(pricing.configuredModels()).isEmpty(); + } + + /** + * Self-hosted ids look like {@code meta-llama/Llama-3-8b}. Such a name cannot travel + * in the URL path — Spring Security's StrictHttpFirewall rejects an encoded slash + * before any controller runs — so the controller accepts it in the body instead. The + * service layer must handle it like any other name. + */ + @Test + void handlesModelNamesContainingASlash() { + pricing.updateRates("meta-llama/Llama-3-8b", new BigDecimal("0.05"), + new BigDecimal("0.08"), null); + + assertThat(pricing.configuredModels()).containsExactly("meta-llama/llama-3-8b"); + assertThat(pricing.estimateCost("meta-llama/Llama-3-8b", 1_000_000, 1_000_000, 0)) + .contains(new BigDecimal("0.130000")); + } +} diff --git a/backend/src/test/java/com/dbaagent/service/llm/LlmUsageAttributionFilterTest.java b/backend/src/test/java/com/dbaagent/service/llm/LlmUsageAttributionFilterTest.java new file mode 100644 index 0000000..92afd34 --- /dev/null +++ b/backend/src/test/java/com/dbaagent/service/llm/LlmUsageAttributionFilterTest.java @@ -0,0 +1,115 @@ +package com.dbaagent.service.llm; + +import org.junit.jupiter.api.Test; +import org.springframework.mock.web.MockHttpServletRequest; +import org.springframework.mock.web.MockHttpServletResponse; + +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class LlmUsageAttributionFilterTest { + + private final LlmUsageAttributionFilter filter = new LlmUsageAttributionFilter(); + + /** Captures what the context looked like from inside the chain. */ + private LlmUsageContext.Scope scopeSeenBy(String method, String uri) throws Exception { + List seen = new ArrayList<>(); + MockHttpServletRequest request = new MockHttpServletRequest(method, uri); + request.setRequestURI(uri); + filter.doFilter(request, new MockHttpServletResponse(), + (req, res) -> seen.add(LlmUsageContext.current())); + return seen.get(0); + } + + @Test + void labelsChatRequests() throws Exception { + assertThat(scopeSeenBy("POST", "/api/chat/ask").feature()).isEqualTo("chat"); + } + + /** + * The longest prefix must win. Ordinary map iteration would let {@code /api/dashboards} + * claim this and collapse the most expensive feature in the product into the generic + * bucket. + */ + @Test + void prefersTheMoreSpecificPrefix() throws Exception { + assertThat(scopeSeenBy("POST", "/api/dashboards/generate/stream").feature()) + .isEqualTo("dashboard-generate"); + assertThat(scopeSeenBy("GET", "/api/dashboards/123").feature()) + .isEqualTo("dashboard"); + } + + @Test + void extractsTheConnectionIdFromThePath() throws Exception { + assertThat(scopeSeenBy("GET", "/api/brain/notes/connection/conn-42").connectionId()) + .isEqualTo("conn-42"); + } + + @Test + void extractsTheConnectionIdFromAQueryParameter() throws Exception { + MockHttpServletRequest request = new MockHttpServletRequest("POST", "/api/chat/ask"); + request.setRequestURI("/api/chat/ask"); + request.setParameter("connectionId", "conn-9"); + + List seen = new ArrayList<>(); + filter.doFilter(request, new MockHttpServletResponse(), + (req, res) -> seen.add(LlmUsageContext.currentConnectionId())); + + assertThat(seen).containsExactly("conn-9"); + } + + /** "test" is an action, not an id; labelling it as one invents a connection. */ + @Test + void doesNotMistakeAnActionSegmentForAConnectionId() throws Exception { + assertThat(scopeSeenBy("POST", "/api/connections/test").connectionId()).isNull(); + } + + @Test + void unmappedPathsRecordUnknownRatherThanFailing() throws Exception { + assertThat(scopeSeenBy("GET", "/api/some-new-feature").feature()) + .isEqualTo(LlmUsageContext.UNKNOWN_FEATURE); + } + + /** The context must not leak into whatever the thread handles next. */ + @Test + void clearsTheContextAfterTheRequest() throws Exception { + scopeSeenBy("POST", "/api/chat/ask"); + + assertThat(LlmUsageContext.current()).isNull(); + } + + @Test + void clearsTheContextEvenWhenTheChainThrows() { + MockHttpServletRequest request = new MockHttpServletRequest("POST", "/api/chat/ask"); + request.setRequestURI("/api/chat/ask"); + + assertThatThrownBy(() -> filter.doFilter(request, new MockHttpServletResponse(), + (req, res) -> { + throw new IllegalStateException("handler blew up"); + })).isInstanceOf(IllegalStateException.class); + + assertThat(LlmUsageContext.current()).isNull(); + } + + /** + * A checked exception from the chain must reach the container as itself. The context + * is carried through a {@code Supplier}, which cannot throw one, so it is wrapped and + * must be unwrapped again — a bug here would turn every servlet error into an opaque + * runtime failure. + */ + @Test + void propagatesCheckedExceptionsUnwrapped() { + MockHttpServletRequest request = new MockHttpServletRequest("POST", "/api/chat/ask"); + request.setRequestURI("/api/chat/ask"); + + assertThatThrownBy(() -> filter.doFilter(request, new MockHttpServletResponse(), + (req, res) -> { + throw new java.io.IOException("socket closed"); + })) + .isInstanceOf(java.io.IOException.class) + .hasMessage("socket closed"); + } +} diff --git a/backend/src/test/java/com/dbaagent/service/llm/LlmUsageRecorderTest.java b/backend/src/test/java/com/dbaagent/service/llm/LlmUsageRecorderTest.java new file mode 100644 index 0000000..240ab3a --- /dev/null +++ b/backend/src/test/java/com/dbaagent/service/llm/LlmUsageRecorderTest.java @@ -0,0 +1,221 @@ +package com.dbaagent.service.llm; + +import com.dbaagent.model.LlmUsage; +import com.dbaagent.model.LlmUsageRole; +import com.dbaagent.service.QueryActorContextHolder; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.springframework.security.authentication.UsernamePasswordAuthenticationToken; +import org.springframework.security.core.context.SecurityContextHolder; + +import java.math.BigDecimal; +import java.util.List; +import java.util.Optional; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class LlmUsageRecorderTest { + + private final LlmUsageWriter writer = mock(LlmUsageWriter.class); + private final LlmPricingService pricing = mock(LlmPricingService.class); + private final LlmUsageRecorder recorder = new LlmUsageRecorder(writer, pricing); + + @AfterEach + void clearContexts() { + SecurityContextHolder.clearContext(); + } + + private static LlmUsageRecorder.Call chatCall() { + return new LlmUsageRecorder.Call( + LlmUsageRole.CHAT, "openai", "gpt-4o", + 1_000, 500, 1_500, 0, false, 42L, true, null); + } + + private LlmUsage captureWritten() { + ArgumentCaptor captor = ArgumentCaptor.forClass(LlmUsage.class); + verify(writer).write(captor.capture()); + return captor.getValue(); + } + + @Test + void storesTokenCountsAndPricedCost() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.of(new BigDecimal("0.007500"))); + + recorder.record(chatCall()); + + LlmUsage row = captureWritten(); + assertThat(row.getRole()).isEqualTo("CHAT"); + assertThat(row.getModel()).isEqualTo("gpt-4o"); + assertThat(row.getPromptTokens()).isEqualTo(1_000); + assertThat(row.getCompletionTokens()).isEqualTo(500); + assertThat(row.getTotalTokens()).isEqualTo(1_500); + assertThat(row.getEstimatedCostUsd()).isEqualByComparingTo("0.007500"); + assertThat(row.isEstimated()).isFalse(); + } + + /** + * An unpriced model must store null, not zero — a zero would silently understate every + * total that sums this column, with nothing to indicate the gap. + */ + @Test + void unpricedModelStoresNullCostNotZero() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + + recorder.record(chatCall()); + + assertThat(captureWritten().getEstimatedCostUsd()).isNull(); + } + + @Test + void derivesTotalWhenTheProviderDidNotReportOne() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + + recorder.record(new LlmUsageRecorder.Call( + LlmUsageRole.CHAT, "openai", "gpt-4o", + 700, 300, 0, 0, false, 1L, true, null)); + + assertThat(captureWritten().getTotalTokens()).isEqualTo(1_000); + } + + @Test + void attributesUsageToTheSqlActorRatherThanTheSecurityPrincipal() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + SecurityContextHolder.getContext().setAuthentication( + new UsernamePasswordAuthenticationToken("admin", "x", List.of())); + + QueryActorContextHolder.withActor("analyst", () -> { + recorder.record(chatCall()); + return null; + }); + + // Under "View as", the principal is the admin but the work is done for the target + // user, and the spend belongs to them. + assertThat(captureWritten().getUsername()).isEqualTo("analyst"); + } + + @Test + void fallsBackToTheSecurityPrincipalWithNoActor() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + SecurityContextHolder.getContext().setAuthentication( + new UsernamePasswordAuthenticationToken("dba", "x", List.of())); + + recorder.record(chatCall()); + + assertThat(captureWritten().getUsername()).isEqualTo("dba"); + } + + @Test + void backgroundWorkRecordsNullUsernameRatherThanFailing() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + + recorder.record(chatCall()); + + assertThat(captureWritten().getUsername()).isNull(); + } + + @Test + void capturesTheDeclaredFeatureAndConnection() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + + LlmUsageContext.with("dashboard-generate", "conn-7", () -> { + recorder.record(chatCall()); + return null; + }); + + LlmUsage row = captureWritten(); + assertThat(row.getFeature()).isEqualTo("dashboard-generate"); + assertThat(row.getConnectionId()).isEqualTo("conn-7"); + } + + @Test + void unattributedCallRecordsUnknownFeatureRatherThanFailing() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + + recorder.record(chatCall()); + + assertThat(captureWritten().getFeature()).isEqualTo("unknown"); + } + + /** + * The property this whole class is built around: accounting is secondary to the + * feature, so a broken ledger must never surface to the caller. + */ + @Test + void aWriteFailureNeverPropagatesToTheCaller() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + doThrow(new RuntimeException("table llm_usage does not exist")) + .when(writer).write(any()); + + assertThatCode(() -> recorder.record(chatCall())).doesNotThrowAnyException(); + } + + @Test + void aPricingFailureNeverPropagatesToTheCaller() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenThrow(new RuntimeException("config backend down")); + + assertThatCode(() -> recorder.record(chatCall())).doesNotThrowAnyException(); + } + + @Test + void negativeTokenCountsAreClampedToZero() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + + recorder.record(new LlmUsageRecorder.Call( + LlmUsageRole.CHAT, "openai", "gpt-4o", + -5, -5, 0, -5, false, 1L, true, null)); + + LlmUsage row = captureWritten(); + assertThat(row.getPromptTokens()).isZero(); + assertThat(row.getCompletionTokens()).isZero(); + assertThat(row.getCachedPromptTokens()).isZero(); + } + + @Test + void failedCallsAreRecordedWithTheirErrorCategory() { + when(pricing.estimateCost(anyString(), anyLong(), anyLong(), anyLong())) + .thenReturn(Optional.empty()); + + recorder.record(new LlmUsageRecorder.Call( + LlmUsageRole.CHAT, "openai", "gpt-4o", + 1_000, 0, 1_000, 0, false, 5L, false, "RATE_LIMIT")); + + LlmUsage row = captureWritten(); + assertThat(row.isSucceeded()).isFalse(); + assertThat(row.getErrorCategory()).isEqualTo("RATE_LIMIT"); + } + + /** Nested scopes must restore the outer one rather than clearing it. */ + @Test + void usageContextRestoresTheEnclosingScope() { + LlmUsageContext.with("outer", "conn-a", () -> { + LlmUsageContext.with("inner", "conn-b", () -> { + assertThat(LlmUsageContext.currentFeature()).isEqualTo("inner"); + return null; + }); + assertThat(LlmUsageContext.currentFeature()).isEqualTo("outer"); + assertThat(LlmUsageContext.currentConnectionId()).isEqualTo("conn-a"); + return null; + }); + assertThat(LlmUsageContext.currentFeature()).isEqualTo(LlmUsageContext.UNKNOWN_FEATURE); + } +} diff --git a/src/components/SettingsModal.js b/src/components/SettingsModal.js index 4b5f83a..4b6a78b 100644 --- a/src/components/SettingsModal.js +++ b/src/components/SettingsModal.js @@ -1,7 +1,7 @@ "use client"; import { useState, useEffect } from "react"; -import { X, Users, Settings, Shield, Activity, KeyRound } from "lucide-react"; +import { X, Users, Settings, Shield, Activity, KeyRound, Coins } from "lucide-react"; import { useAuth } from "@/hooks/useAuth"; import { PERMISSIONS } from "@/lib/permissions"; import AdminWorkspaceSettings from "@/components/settings/AdminWorkspaceSettings"; @@ -9,6 +9,7 @@ import SlackAccessCodePanel from "@/components/settings/SlackAccessCodePanel"; import McpTokensPanel from "@/components/settings/McpTokensPanel"; import UsersTab from "./tabs/admin/UsersTab"; import AuditLogsTab from "./tabs/admin/AuditLogsTab"; +import LlmUsageTab from "./tabs/admin/LlmUsageTab"; import styles from "./SettingsModal.module.css"; export default function SettingsModal({ isOpen, onClose }) { @@ -56,6 +57,12 @@ export default function SettingsModal({ isOpen, onClose }) { label: "Audit Logs", description: "Review editor and security activity", }, + { + id: "llm-usage", + icon: Coins, + label: "AI Usage & Cost", + description: "Track model spend by feature and user", + }, ] : []), { @@ -86,6 +93,8 @@ export default function SettingsModal({ isOpen, onClose }) { return isAdmin ? : null; case "audit-logs": return isAdmin ? : null; + case "llm-usage": + return isAdmin ? : null; case "mcp-tokens": return ; case "general": diff --git a/src/components/tabs/admin/LlmPricingPanel.jsx b/src/components/tabs/admin/LlmPricingPanel.jsx new file mode 100644 index 0000000..4b413a3 --- /dev/null +++ b/src/components/tabs/admin/LlmPricingPanel.jsx @@ -0,0 +1,308 @@ +'use client' + +import { useEffect, useState } from 'react' +import { AlertCircle, Check, Plus, X } from 'lucide-react' +import { useLlmPricing, useUpdateLlmPricing } from '@/lib/hooks/queries/useLlmUsage' + +const FIELDS = [ + { key: 'inputPer1m', label: 'Input', hint: 'Required' }, + { key: 'outputPer1m', label: 'Output', hint: 'Blank for embeddings' }, + { key: 'cachedInputPer1m', label: 'Cached input', hint: 'Defaults to input' }, +] + +/** '' for a null rate so an empty box round-trips as "not set" rather than as 0. */ +function toForm(row) { + return { + inputPer1m: row.inputPer1m ?? '', + outputPer1m: row.outputPer1m ?? '', + cachedInputPer1m: row.cachedInputPer1m ?? '', + } +} + +function toPayload(form) { + const parse = (v) => { + const trimmed = String(v ?? '').trim() + if (trimmed === '') return null + const n = Number(trimmed) + return Number.isFinite(n) ? n : null + } + return { + inputPer1m: parse(form.inputPer1m), + outputPer1m: parse(form.outputPer1m), + cachedInputPer1m: parse(form.cachedInputPer1m), + } +} + +function isInvalid(form) { + return FIELDS.some(({ key }) => { + const raw = String(form[key] ?? '').trim() + if (raw === '') return false + const n = Number(raw) + return !Number.isFinite(n) || n < 0 + }) +} + +function RateInput({ value, onChange, label, disabled }) { + return ( +

+ + $ + + onChange(e.target.value)} + placeholder="—" + className="w-full pl-5 pr-2 py-1.5 text-sm text-right tabular-nums border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-gray-900 focus:border-transparent disabled:bg-gray-50" + /> +
+ + ) +} + +function PricingRow({ row, onSave, saving, justSaved }) { + const [form, setForm] = useState(() => toForm(row)) + + // Compared as strings because the inputs hold strings; a numeric compare would call + // "2.50" a change from 2.5 and leave Save enabled on an untouched row. + const dirty = FIELDS.some( + ({ key }) => String(form[key] ?? '') !== String(toForm(row)[key] ?? ''), + ) + const invalid = isInvalid(form) + + const save = () => onSave(row.model, toPayload(form)) + + return ( + + +
+ {row.model} + {!row.priced ? ( + + Unpriced + + ) : null} + {!row.seenInUsage ? ( + + Unused + + ) : null} +
+ + {FIELDS.map(({ key, label }) => ( + + setForm((f) => ({ ...f, [key]: v }))} + /> + + ))} + + {dirty ? ( +
+ + +
+ ) : justSaved ? ( + + Saved + + ) : null} + + + ) +} + +export default function LlmPricingPanel() { + const { data: rows = [], isLoading, isError, error } = useLlmPricing() + const updatePricing = useUpdateLlmPricing() + const [newModel, setNewModel] = useState('') + const [adding, setAdding] = useState(false) + // Held here, not in the row: a save refetches the list and the row is remounted with + // its new values, which would discard a flag owned by the row itself. + const [savedModel, setSavedModel] = useState(null) + + useEffect(() => { + if (!savedModel) return undefined + const timer = setTimeout(() => setSavedModel(null), 4000) + return () => clearTimeout(timer) + }, [savedModel]) + + // The rejection is caught deliberately. mutateAsync rejects on failure, and an + // uncaught rejection here propagated out of the row's click handler instead of + // letting the isError banner render — a save against a broken backend showed the + // user nothing at all. The mutation's own error state is what surfaces the message. + const save = async (model, rates) => { + try { + const result = await updatePricing.mutateAsync({ model, rates }) + setSavedModel(result?.model ?? model) + return true + } catch { + setSavedModel(null) + return false + } + } + + const addModel = async () => { + const model = newModel.trim() + if (!model) return + // Written with no rates: the row appears immediately as Unpriced and is filled in + // through the same inputs as every other row, so there is only one editing path. + const ok = await save(model, { + inputPer1m: null, + outputPer1m: null, + cachedInputPer1m: null, + }) + // The form stays open on failure so the typed name is not lost and the error is + // visible next to what caused it. + if (ok) { + setNewModel('') + setAdding(false) + } + } + + return ( +
+
+
+

Model pricing

+

+ USD per 1M tokens. Rates apply to calls made from now on — editing a rate does + not re-cost calls already recorded. +

+
+ {adding ? ( +
+ setNewModel(e.target.value)} + onKeyDown={(e) => { + if (e.key === 'Enter') addModel() + if (e.key === 'Escape') { + setNewModel('') + setAdding(false) + } + }} + placeholder="model name" + className="w-48 px-2 py-1.5 text-sm font-mono border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-gray-900" + /> + + +
+ ) : ( + + )} +
+ + {updatePricing.isError ? ( +
+ + + {updatePricing.error?.message || 'Could not save the rate.'} + +
+ ) : null} + +
+ {isLoading ? ( +

Loading rates…

+ ) : isError ? ( +

+ {error?.message || 'Could not load pricing.'} +

+ ) : rows.length === 0 ? ( +

+ No models recorded yet. Rates can be added ahead of the first call. +

+ ) : ( + + + + + {FIELDS.map(({ key, label, hint }) => ( + + ))} + + + + {rows.map((row) => ( + // Keyed by model *and* its saved values so the row's local form state is + // rebuilt after a save; keying on model alone would keep showing the + // pre-save draft and leave the row looking permanently dirty. + + ))} + +
+ Model + + {label} + + {hint} + + +
+ )} +
+
+ ) +} diff --git a/src/components/tabs/admin/LlmUsageTab.jsx b/src/components/tabs/admin/LlmUsageTab.jsx new file mode 100644 index 0000000..0b71171 --- /dev/null +++ b/src/components/tabs/admin/LlmUsageTab.jsx @@ -0,0 +1,306 @@ +'use client' + +import { useMemo, useState } from 'react' +import { + AlertCircle, + Coins, + RefreshCw, + TriangleAlert, +} from 'lucide-react' +import { + Area, + AreaChart, + CartesianGrid, + ResponsiveContainer, + Tooltip, + XAxis, + YAxis, +} from 'recharts' +import { useLlmUsageSummary } from '@/lib/hooks/queries/useLlmUsage' +import LlmPricingPanel from './LlmPricingPanel' +import { useAuth } from '@/hooks/useAuth' + +const WINDOWS = [ + { days: 7, label: '7 days' }, + { days: 30, label: '30 days' }, + { days: 90, label: '90 days' }, +] + +function formatUsd(value) { + const amount = Number(value ?? 0) + if (!Number.isFinite(amount)) return '$0.00' + // Sub-cent totals are real on small installs; showing "$0.00" for them reads as + // "nothing was spent" when the truth is "not much was spent". + if (amount > 0 && amount < 0.01) return '<$0.01' + return `$${amount.toLocaleString(undefined, { + minimumFractionDigits: 2, + maximumFractionDigits: 2, + })}` +} + +function formatTokens(value) { + const tokens = Number(value ?? 0) + if (!Number.isFinite(tokens)) return '0' + if (tokens >= 1_000_000) return `${(tokens / 1_000_000).toFixed(1)}M` + if (tokens >= 1_000) return `${(tokens / 1_000).toFixed(1)}K` + return String(tokens) +} + +function StatCard({ label, value, hint }) { + return ( +
+ + {label} + + + {value} + + {hint ? {hint} : null} +
+ ) +} + +function Breakdown({ title, rows, emptyLabel }) { + const max = Math.max(...rows.map((r) => Number(r.costUsd ?? 0)), 0) + + return ( +
+

{title}

+ {rows.length === 0 ? ( +

{emptyLabel}

+ ) : ( +
    + {rows.slice(0, 8).map((row) => ( +
  • +
    + + {row.key} + + + {formatUsd(row.costUsd)} + + {formatTokens(row.totalTokens)} tok + + +
    +
    +
    0 ? `${(Number(row.costUsd ?? 0) / max) * 100}%` : '0%', + }} + /> +
    +
    • + ))} +
+ )} +
+ ) +} + +export default function LlmUsageTab() { + const { isAdmin } = useAuth() + const [days, setDays] = useState(30) + const { data, isLoading, isError, error, refetch, isFetching } = useLlmUsageSummary(days) + + const daily = useMemo( + () => + (data?.daily ?? []).map((point) => ({ + day: point.day, + cost: Number(point.costUsd ?? 0), + })), + [data], + ) + + if (!isAdmin) { + return ( +
+ +

Admin Access Required

+

+ You need administrator privileges to view LLM spend. +

+
+ ) + } + + const totals = data?.totals + const unpriced = data?.unpricedModels ?? [] + + return ( +
+
+
+ +
+

AI Usage & Cost

+

+ What DeepSQL spent on model calls, by feature, user, and model. +

+
+
+
+
+ {WINDOWS.map((w) => ( + + ))} +
+ +
+
+ + {isError ? ( +
+ + + {error?.message || 'Could not load usage data.'} + +
+ ) : null} + +
+ {isLoading ? ( +

Loading usage…

+ ) : ( + <> + {/* + An unpriced model is an operator gap, not a bug: the ledger records tokens + for every call but can only cost the models that have a configured rate, so + a partial total is surfaced as partial rather than presented as complete. + */} + {unpriced.length > 0 ? ( +
+ +
+

+ {unpriced.length === 1 + ? '1 model has no configured price, so the totals below understate real spend.' + : `${unpriced.length} models have no configured price, so the totals below understate real spend.`} +

+

+ Set a rate in Model pricing below + for: {unpriced.join(', ')} +

+
+
+ ) : null} + +
+ 0 + ? `${totals.unpricedCalls.toLocaleString()} unpriced calls excluded` + : `Last ${days} days` + } + /> + 0 + ? `${totals.failedCalls.toLocaleString()} failed` + : 'All succeeded' + } + /> + + +
+ +
+

Daily spend

+ {daily.length === 0 ? ( +

+ No model calls recorded in this window. +

+ ) : ( + // Height is given in pixels rather than as a percentage of the + // parent. ResponsiveContainer measures its parent on mount, which + // inside this modal reports -1 for one frame; an explicit height means + // it recovers on the next measurement instead of collapsing. +
+ + + + + `$${Number(v).toFixed(2)}`} + /> + [formatUsd(value), 'Cost']} + contentStyle={{ fontSize: 12, borderRadius: 6 }} + /> + + + +
+ )} +
+ +
+ + + +
+ + + + )} +
+
+ ) +} diff --git a/src/lib/api/client.js b/src/lib/api/client.js index b859c9f..214d410 100644 --- a/src/lib/api/client.js +++ b/src/lib/api/client.js @@ -479,6 +479,47 @@ export const adminAPI = { }, }; +// LLM usage and cost accounting (ADMIN only) +export const llmUsageAPI = { + getSummary: async (days = 30) => { + const response = await apiClient.get("/api/admin/llm-usage/summary", { + params: { days }, + }); + return response.data; + }, + + getRecent: async ({ days = 30, page = 0, size = 50 } = {}) => { + const response = await apiClient.get("/api/admin/llm-usage/recent", { + params: { days, page, size }, + }); + return response.data; + }, + + purge: async (olderThanDays) => { + const response = await apiClient.delete("/api/admin/llm-usage/purge", { + params: { olderThanDays }, + }); + return response.data; + }, + + getPricing: async () => { + const response = await apiClient.get("/api/admin/llm-usage/pricing"); + return response.data; + }, + + // A dotted name (gpt-5.4) travels fine in the path, but a slash cannot: Spring + // Security's StrictHttpFirewall rejects %2F with a bare 400 before the controller + // runs. Self-hosted ids look like meta-llama/Llama-3-8b, so those go in the body + // instead, against the pathless route. + updatePricing: async (model, rates) => { + const url = model.includes("/") + ? "/api/admin/llm-usage/pricing" + : `/api/admin/llm-usage/pricing/${encodeURIComponent(model)}`; + const response = await apiClient.put(url, { ...rates, model }); + return response.data; + }, +}; + export const slackLinkAPI = { getCurrentLinkCode: async () => { const response = await apiClient.get('/api/slack/link/code') diff --git a/src/lib/hooks/queries/useLlmUsage.js b/src/lib/hooks/queries/useLlmUsage.js new file mode 100644 index 0000000..1827ca3 --- /dev/null +++ b/src/lib/hooks/queries/useLlmUsage.js @@ -0,0 +1,46 @@ +/** + * TanStack Query hooks for LLM usage and cost accounting (admin only) + */ + +import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query' +import { llmUsageAPI } from '@/lib/api/client' +import { queryKeys } from '@/lib/queryKeys' + +export function useLlmUsageSummary(days = 30) { + return useQuery({ + queryKey: queryKeys.llmUsage.summary(days), + queryFn: () => llmUsageAPI.getSummary(days), + }) +} + +export function useLlmUsageRecent({ days = 30, page = 0, size = 50 } = {}) { + return useQuery({ + queryKey: queryKeys.llmUsage.recent(days, page, size), + queryFn: () => llmUsageAPI.getRecent({ days, page, size }), + }) +} + +export function useLlmPricing() { + return useQuery({ + queryKey: queryKeys.llmUsage.pricing(), + queryFn: () => llmUsageAPI.getPricing(), + }) +} + +export function useUpdateLlmPricing() { + const queryClient = useQueryClient() + return useMutation({ + mutationFn: ({ model, rates }) => llmUsageAPI.updatePricing(model, rates), + // A rate change alters what future calls cost, so the summary is refetched too — + // not just the pricing list the form is bound to. + onSuccess: () => queryClient.invalidateQueries({ queryKey: ['llmUsage'] }), + }) +} + +export function usePurgeLlmUsage() { + const queryClient = useQueryClient() + return useMutation({ + mutationFn: (olderThanDays) => llmUsageAPI.purge(olderThanDays), + onSuccess: () => queryClient.invalidateQueries({ queryKey: ['llmUsage'] }), + }) +} diff --git a/src/lib/queryKeys.js b/src/lib/queryKeys.js index edcbd42..20a54e1 100644 --- a/src/lib/queryKeys.js +++ b/src/lib/queryKeys.js @@ -411,6 +411,12 @@ export const queryKeys = { }, // ==================== Saved Items ==================== + llmUsage: { + summary: (days) => ["llmUsage", "summary", days], + recent: (days, page, size) => ["llmUsage", "recent", days, page, size], + pricing: () => ["llmUsage", "pricing"], + }, + savedQueries: { all: (connectionId) => ["savedQueries", connectionId], detail: (id) => ["savedQueries", "detail", id],