Skip to content

Commit 647382b

Browse files
notSumit25claude
andcommitted
feat: LLM usage and cost accounting, with an in-app pricing editor
Every model call now writes one `llm_usage` row (V119), surfaced at Settings -> AI Usage & Cost: spend by feature, user and model, a daily trend, and an editable per-model rate table. Capture sits at the two provider funnels -- RefreshableChatModel and EmbeddingService -- rather than in feature code, so a new caller is accounted for without touching it. LlmUsageRecorder swallows its own failures by design: accounting must never break the call it measures. Rows are written through LlmUsageWriter in their own transaction, because self-invoking a @transactional method bypasses the Spring proxy and the row would roll back with a failed chat turn. Rates live in system_config with no bundled defaults -- a stale price list produces confident wrong totals nobody checks. An unpriced model stores a NULL cost, never 0, and the UI reports the gap instead of silently understating spend. Two pre-existing bugs surfaced while wiring this up: - ResponsesApiChatModel.buildMetadata discarded the provider's `usage` block entirely. Harmless while nothing read it, a silent zero the moment accounting summed it: real billed calls recorded 0 tokens and $0.00. It now reads both vendor dialects, and the streaming path requests usage (stream_options.include_usage) and emits the late usage event so a streamed turn is billable at all. - Attribution cannot ride a ThreadLocal. Chat returns a Flux and does its model work later on a CompletableFuture, so the servlet is gone by then; every row read feature=unknown while a single-threaded unit test of the filter passed. ChatService now re-establishes the scope where it already re-establishes the SQL actor. Verified against the running stack, not inferred: real Azure OpenAI calls at three configured rates, costs matching to six decimals with no restart; a DEVELOPER 403'd on all endpoints with the rejected write leaving rates untouched; and chat still answering with the llm_usage table deleted outright. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 1b36bda commit 647382b

33 files changed

Lines changed: 3271 additions & 25 deletions

CLAUDE.md

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -311,6 +311,82 @@ embedding model and read through another raises no error — retrieval just degr
311311
(pgvector's `text`-column fallback has no dimension constraint and cosine similarity still
312312
returns a number).
313313

314+
## LLM Usage & Cost Accounting
315+
316+
Every model call writes one `llm_usage` row (`V119__create_llm_usage.sql`), surfaced at
317+
**Settings → AI Usage & Cost** (`LlmUsageTab.jsx`) via `GET /admin/llm-usage/summary`.
318+
Admin-only and not connection-scoped, so it carries a class-level `@PreAuthorize` rather
319+
than an `assertCan*` call — the case `Endpoint Authorization Rules` describes as
320+
"an endpoint with no connection scope at all is admin-only".
321+
322+
- **Recording happens at the two provider funnels, never in feature code.**
323+
`RefreshableChatModel.call/stream` and `EmbeddingService.createEmbedding(s)` are the only
324+
call sites, so a thirteenth service that reaches a model is accounted for without
325+
touching it. `LlmUsageRecorder` swallows its own failures by design: accounting must
326+
never be able to break the call it measures.
327+
- **Rows are written in their own transaction** (`LlmUsageWriter`, `REQUIRES_NEW`). It is a
328+
separate bean on purpose — self-invoking a `@Transactional` method bypasses the Spring
329+
proxy, and the row would then roll back with a failed chat turn, which is exactly the
330+
case it exists to record. Same trap `McpTokenRepository.deleteByUserId` documents.
331+
- **`estimated` distinguishes metered from derived counts.** Chat providers return real
332+
token counts; `LlmEmbeddingProvider` returns vectors only, so embedding tokens are
333+
derived from input length at 3 chars/token (schema text is denser than prose — see
334+
`EmbeddingService.truncate`). Both land in the same columns so one spend total is
335+
possible, but a vendor-invoice reconciliation can tell the halves apart.
336+
- **An unpriced model stores `NULL` cost, never `0`.** `LlmPricingService` reads rates from
337+
`system_config` (`llm.pricing.<model>.{input,output,cached-input}-per-1m`) and ships **no
338+
default prices** — a stale bundled price list produces confident wrong totals nobody
339+
thinks to check. The UI reports unpriced calls rather than silently understating spend.
340+
- **Rates are edited in the UI**, in the Model pricing panel of the same tab
341+
(`LlmPricingPanel.jsx``GET|PUT /admin/llm-usage/pricing`). The list is every model the
342+
ledger has seen plus every model with a rate configured, unpriced first. Writes take
343+
effect on the next call with no restart, since `LlmPricingService` reads `system_config`
344+
per call. Three details are load-bearing:
345+
- **A cleared field writes `""`, not a deleted row.** `rate()` already treats blank as
346+
absent, and `SystemConfigService` has no delete — adding one for a single caller would
347+
widen a shared service. Sending the whole set means an emptied box genuinely clears
348+
that rate rather than leaving the old value behind.
349+
- **A model name can contain dots** (`gpt-5.4`), so `configuredModels()` strips the known
350+
suffix from the *end* of the key. Splitting on the first `.` after the prefix reports
351+
`gpt-5` and loses the row. The controller mapping is `{model:.+}` for the same reason.
352+
- **A name containing a slash cannot go in the path at all**, even percent-encoded:
353+
Spring Security's default `StrictHttpFirewall` rejects `%2F` with a bare 400 before any
354+
controller runs — verified in QA, and it applies to every endpoint, not just this one.
355+
Self-hosted ids look like `meta-llama/Llama-3-8b`, so `PUT /pricing` (no path segment)
356+
takes the name in the body, and `client.js` switches to it when the name contains `/`.
357+
Relaxing the firewall would be the wrong trade for a naming convenience.
358+
- **A failed save must still say so.** Spring's default 500 body carries no `message`, so
359+
the client had nothing to display and a save against a broken config store showed the
360+
user *nothing at all* — no error, no toast, silent. The handler now returns a `message`
361+
on any non-`IllegalArgumentException` failure, and the panel catches the `mutateAsync`
362+
rejection rather than letting it escape the click handler (an uncaught rejection there
363+
is what stopped the `isError` banner from rendering). Both found by taking
364+
`system_config` away mid-save.
365+
- **A negative rate is rejected at write time**, not just ignored on read: a value that
366+
silently does nothing after the UI said "Saved" is worse than an error at entry.
367+
- **Editing a rate does not re-cost recorded calls.** `estimated_cost_usd` is a snapshot
368+
written at record time, which is what an audit trail wants — but it means a mid-window
369+
price change shows as a step in the daily chart, not a uniform restatement. The panel
370+
says so.
371+
- **Attribution cannot ride a ThreadLocal alone.** `LlmUsageAttributionFilter` labels each
372+
request from its URI, but chat returns a `Flux` and does its model work later on a
373+
`CompletableFuture` — the servlet, and any thread-local set during it, is gone by then.
374+
`ChatService` re-establishes `LlmUsageContext.with(...)` inside `runAsync` exactly where
375+
it already re-establishes `QueryActorContextHolder.withActor`; `DashboardAlertService`
376+
declares its own scope since scheduled work has no request at all. Verified live: every
377+
row read `feature = unknown` until this was added, while a single-threaded unit test of
378+
the filter passed.
379+
- **`ResponsesApiChatModel.buildMetadata` used to discard `usage` entirely**, which cost
380+
nothing while nothing read it and became a silent zero the moment accounting summed it —
381+
real billed calls recorded 0 tokens and $0.00. It now reads both dialects
382+
(`prompt_tokens`/`completion_tokens` and `input_tokens`/`output_tokens`) plus cached-token
383+
details, and the streaming path sends `stream_options.include_usage` and emits the late
384+
usage event as a text-free chunk. `RefreshableChatModel.meteredStream` records **one row
385+
per stream**, taking the last reported usage rather than summing chunks: providers that
386+
report running totals would otherwise have every partial added to the final figure.
387+
- Usage belongs to `QueryActorContextHolder` first and the security principal second, so
388+
under **View as** the spend is attributed to the target user, not the admin.
389+
314390
## Key Rules & Patterns
315391

316392
### Backend Rules
@@ -853,6 +929,10 @@ DEEPSQL_EMBEDDING_API_KEY=<key>
853929
DEEPSQL_EMBEDDING_MODEL=text-embedding-3-large
854930
# Optional chat tuning: DEEPSQL_CHAT_TEMPERATURE, DEEPSQL_CHAT_API_VERSION,
855931
# DEEPSQL_CHAT_USE_RESPONSES_API (true|false|auto).
932+
# Model prices are NOT environment variables. They live in system_config and are edited in
933+
# Settings -> AI Usage & Cost -> Model pricing; there are deliberately no defaults:
934+
# llm.pricing.<model>.input-per-1m / .output-per-1m / .cached-input-per-1m
935+
# An unpriced model still records tokens; its cost is NULL and the UI flags it.
856936

857937
# Only if using Azure AI Search instead of pgvector for the vector store.
858938
azure.search.api-key=<key>

backend/src/main/java/com/dbaagent/config/LlmConfig.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import com.dbaagent.llm.LlmProviderRegistry;
55
import com.dbaagent.llm.spring.ProviderBackedEmbeddingModel;
66
import com.dbaagent.service.EmbeddingService;
7+
import com.dbaagent.service.llm.LlmUsageRecorder;
78
import org.springframework.ai.chat.model.ChatModel;
89
import org.springframework.ai.embedding.EmbeddingModel;
910
import org.springframework.context.annotation.Bean;
@@ -24,8 +25,9 @@ public class LlmConfig {
2425

2526
@Bean
2627
@Primary
27-
public ChatModel chatModel(LlmConfigResolver resolver, LlmProviderRegistry registry) {
28-
return new RefreshableChatModel(resolver, registry);
28+
public ChatModel chatModel(LlmConfigResolver resolver, LlmProviderRegistry registry,
29+
LlmUsageRecorder usageRecorder) {
30+
return new RefreshableChatModel(resolver, registry, usageRecorder);
2931
}
3032

3133
/**

backend/src/main/java/com/dbaagent/config/RefreshableChatModel.java

Lines changed: 163 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,20 @@
77
import com.dbaagent.llm.api.LlmErrorCategory;
88
import com.dbaagent.llm.api.LlmNotConfiguredException;
99
import com.dbaagent.llm.api.UnsupportedLlmProviderException;
10+
import com.dbaagent.model.LlmUsageRole;
11+
import com.dbaagent.service.llm.LlmUsageRecorder;
1012
import lombok.extern.slf4j.Slf4j;
1113
import org.springframework.ai.chat.model.ChatModel;
1214
import org.springframework.ai.chat.model.ChatResponse;
1315
import org.springframework.ai.chat.prompt.ChatOptions;
1416
import org.springframework.ai.chat.prompt.Prompt;
17+
import org.springframework.ai.chat.metadata.Usage;
1518
import reactor.core.publisher.Flux;
1619

1720
import java.util.concurrent.atomic.AtomicBoolean;
21+
import java.util.concurrent.atomic.AtomicReference;
1822
import java.util.concurrent.locks.ReentrantLock;
23+
import java.util.function.Supplier;
1924

2025
/**
2126
* A {@link ChatModel} that resolves its credentials on every call through
@@ -52,28 +57,62 @@ private record CachedDelegate(LlmCredentials key, ChatModel model) {}
5257
private final LlmConfigResolver resolver;
5358
private final LlmProviderRegistry registry;
5459

60+
/**
61+
* Nullable so the accounting dependency stays optional. This class is constructed
62+
* directly in tests and in {@code LlmConfig}; a hard dependency would make every
63+
* existing construction site a compile error for a concern none of them care about.
64+
*/
65+
private final LlmUsageRecorder usageRecorder;
66+
5567
/** Serialises rebuilds only; the read path is a single lock-free volatile read. */
5668
private final ReentrantLock buildLock = new ReentrantLock();
5769
private volatile CachedDelegate cached;
5870

5971
public RefreshableChatModel(LlmConfigResolver resolver, LlmProviderRegistry registry) {
72+
this(resolver, registry, null);
73+
}
74+
75+
public RefreshableChatModel(LlmConfigResolver resolver, LlmProviderRegistry registry,
76+
LlmUsageRecorder usageRecorder) {
6077
this.resolver = resolver;
6178
this.registry = registry;
79+
this.usageRecorder = usageRecorder;
6280
}
6381

6482
@Override
6583
public ChatResponse call(Prompt prompt) {
6684
CachedDelegate active = resolveDelegate();
6785
try {
68-
return active.model().call(prompt);
86+
return metered(active, () -> active.model().call(prompt));
6987
} catch (RuntimeException e) {
7088
if (shouldRetryWithEnvFallback(e, active.key())) {
71-
return resolveDelegate().model().call(prompt);
89+
CachedDelegate retry = resolveDelegate();
90+
return metered(retry, () -> retry.model().call(prompt));
7291
}
7392
throw e;
7493
}
7594
}
7695

96+
/**
97+
* Runs a chat call and records what it cost.
98+
*
99+
* <p>A failed call is recorded too. Providers bill for prompt tokens on responses that
100+
* error partway, and an operator investigating a spend spike caused by a retry loop
101+
* needs to see the failures — a ledger holding only successes hides exactly the
102+
* pathology it would be consulted about.
103+
*/
104+
private ChatResponse metered(CachedDelegate active, Supplier<ChatResponse> call) {
105+
long startedAt = System.nanoTime();
106+
try {
107+
ChatResponse response = call.get();
108+
recordUsage(active, response, startedAt, null);
109+
return response;
110+
} catch (RuntimeException e) {
111+
recordUsage(active, null, startedAt, e);
112+
throw e;
113+
}
114+
}
115+
77116
/**
78117
* Streaming delegates report failures asynchronously through the sink rather than by
79118
* throwing, so the fallback hangs off the error signal as well as off a synchronous
@@ -87,15 +126,65 @@ public Flux<ChatResponse> stream(Prompt prompt) {
87126
CachedDelegate active = resolveDelegate();
88127
AtomicBoolean emitted = new AtomicBoolean(false);
89128
try {
90-
return active.model().stream(prompt)
129+
return meteredStream(active, active.model().stream(prompt)
91130
.doOnNext(chunk -> emitted.set(true))
92-
.onErrorResume(e -> resumeStream(prompt, e, active.key(), emitted.get()));
131+
.onErrorResume(e -> resumeStream(prompt, e, active.key(), emitted.get())));
93132
} catch (RuntimeException e) {
94-
return resumeStream(prompt, e, active.key(), false);
133+
return meteredStream(active, resumeStream(prompt, e, active.key(), false));
95134
}
96135
});
97136
}
98137

138+
/**
139+
* Records one row for a whole stream, not one per chunk.
140+
*
141+
* <p>Usage on a stream arrives on a single late chunk — typically the last, after the
142+
* provider has finished counting — while every earlier chunk carries either no
143+
* metadata or a zero-filled {@link Usage}. Recording per chunk would write hundreds of
144+
* rows for one call and inflate the call count enormously; summing across chunks would
145+
* be worse, because providers that report cumulative running totals would have their
146+
* final figure added on top of every partial. So the last non-zero usage seen wins, and
147+
* exactly one row is written when the stream terminates.
148+
*
149+
* <p>{@code doFinally} rather than {@code doOnComplete}: a stream that errors or is
150+
* cancelled partway still consumed prompt tokens, and a cancelled dashboard build is
151+
* precisely the kind of silent spend an operator wants on the ledger.
152+
*/
153+
private Flux<ChatResponse> meteredStream(CachedDelegate active, Flux<ChatResponse> source) {
154+
if (usageRecorder == null) {
155+
return source;
156+
}
157+
long startedAt = System.nanoTime();
158+
AtomicReference<ChatResponse> lastWithUsage = new AtomicReference<>();
159+
AtomicReference<Throwable> failure = new AtomicReference<>();
160+
AtomicBoolean recorded = new AtomicBoolean(false);
161+
162+
return source
163+
.doOnNext(chunk -> {
164+
if (hasUsage(chunk)) {
165+
lastWithUsage.set(chunk);
166+
}
167+
})
168+
.doOnError(failure::set)
169+
.doFinally(signal -> {
170+
// doFinally can fire once per subscription; a Flux that is retried or
171+
// resubscribed must not double-bill the same logical call.
172+
if (recorded.compareAndSet(false, true)) {
173+
Throwable error = failure.get();
174+
recordUsage(active, lastWithUsage.get(), startedAt,
175+
error instanceof RuntimeException re ? re : null);
176+
}
177+
});
178+
}
179+
180+
private static boolean hasUsage(ChatResponse chunk) {
181+
if (chunk == null || chunk.getMetadata() == null) {
182+
return false;
183+
}
184+
Usage usage = chunk.getMetadata().getUsage();
185+
return usage != null && zeroIfNull(usage.getTotalTokens()) > 0;
186+
}
187+
99188
/**
100189
* Reports the active delegate's options, or neutral ones when nothing is configured.
101190
*
@@ -128,6 +217,75 @@ public ChatOptions getDefaultOptions() {
128217
}
129218
}
130219

220+
/**
221+
* Records one chat call, taking token counts from the response the provider returned.
222+
*
223+
* <p>The model name comes from the response metadata when the provider reports it and
224+
* falls back to the configured model. Those can legitimately differ — an alias that
225+
* resolves to a dated snapshot, for instance — and the served model is the one that
226+
* was actually billed, so it wins.
227+
*
228+
* <p>Never throws. It runs inside the call path of every chat turn in the product;
229+
* a defect here must not become a failed conversation.
230+
*/
231+
private void recordUsage(CachedDelegate active, ChatResponse response,
232+
long startedAt, RuntimeException failure) {
233+
if (usageRecorder == null) {
234+
return;
235+
}
236+
try {
237+
long latencyMs = (System.nanoTime() - startedAt) / 1_000_000L;
238+
LlmCredentials key = active.key();
239+
240+
long prompt = 0;
241+
long completion = 0;
242+
long total = 0;
243+
long cached = 0;
244+
String model = key.getOrDefault("model", "unknown");
245+
246+
if (response != null && response.getMetadata() != null) {
247+
Usage usage = response.getMetadata().getUsage();
248+
if (usage != null) {
249+
prompt = zeroIfNull(usage.getPromptTokens());
250+
completion = zeroIfNull(usage.getCompletionTokens());
251+
total = zeroIfNull(usage.getTotalTokens());
252+
cached = zeroIfNullLong(usage.getCacheReadInputTokens());
253+
}
254+
String served = response.getMetadata().getModel();
255+
if (served != null && !served.isBlank()) {
256+
model = served;
257+
}
258+
}
259+
260+
String errorCategory = failure == null ? null
261+
: String.valueOf(registry.chatProvider(key.providerId()).classify(failure));
262+
263+
usageRecorder.record(new LlmUsageRecorder.Call(
264+
LlmUsageRole.CHAT,
265+
key.providerId(),
266+
model,
267+
prompt,
268+
completion,
269+
total,
270+
cached,
271+
false,
272+
latencyMs,
273+
failure == null,
274+
errorCategory));
275+
} catch (RuntimeException e) {
276+
log.warn("RefreshableChatModel: could not record usage; the call was unaffected", e);
277+
}
278+
}
279+
280+
private static long zeroIfNull(Integer value) {
281+
return value == null ? 0L : value.longValue();
282+
}
283+
284+
/** Cache token accessors are {@code Long} in Spring AI 2.0, unlike prompt/completion. */
285+
private static long zeroIfNullLong(Long value) {
286+
return value == null ? 0L : value;
287+
}
288+
131289
// ── Internal ──────────────────────────────────────────────────────────────
132290

133291
private CachedDelegate resolveDelegate() {

0 commit comments

Comments
 (0)