feat(llm): usage hook to surface token counts (ADR-030, #15)

WithUsageHook callback fires with model + prompt/completion tokens parsed from the
response usage block. Keeps the copied stdlib-only llm package decoupled from
metrics (ADR-004) — the caller wires it to internal/metrics. TDD covered.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-12 08:36:02 +02:00
co-authored by Claude Opus 4.8
parent a5a8cf6f6d
commit 9c2a04406b
2 changed files with 40 additions and 0 deletions
+24
View File
@@ -85,6 +85,30 @@ func TestClient_WithMaxTokens(t *testing.T) {
}
}
// TestClient_UsageHookRecordsTokens: the usage hook fires with the model and the
// prompt/completion token counts parsed from the response usage block.
func TestClient_UsageHookRecordsTokens(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
_ = json.NewEncoder(w).Encode(map[string]any{
"choices": []map[string]any{{"message": map[string]any{"content": "ok"}}},
"usage": map[string]any{"prompt_tokens": 123, "completion_tokens": 45},
})
}))
defer srv.Close()
var gotModel string
var gotPrompt, gotCompletion int
c := New(srv.URL, "", "test-model", 10*time.Second, WithUsageHook(func(model string, p, comp int) {
gotModel, gotPrompt, gotCompletion = model, p, comp
}))
if _, err := c.Complete(context.Background(), "sys", "user"); err != nil {
t.Fatalf("Complete: %v", err)
}
if gotModel != "test-model" || gotPrompt != 123 || gotCompletion != 45 {
t.Errorf("usage hook got (%q, %d, %d), want (test-model, 123, 45)", gotModel, gotPrompt, gotCompletion)
}
}
func TestClient_ReturnsErrorOnNon200(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Error(w, "overloaded", http.StatusServiceUnavailable)