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
+16
View File
@@ -32,6 +32,7 @@ type Client struct {
model string model string
maxTokens int maxTokens int
httpClient *http.Client httpClient *http.Client
usageHook func(model string, prompt, completion int)
} }
// Option configures a Client at construction. Variadic so the existing 4-arg // Option configures a Client at construction. Variadic so the existing 4-arg
@@ -50,6 +51,14 @@ func WithMaxTokens(n int) Option {
} }
} }
// WithUsageHook registers a callback fired after a successful completion with the
// model and the prompt/completion token counts from the response usage block. It
// keeps this copied, stdlib-only package (ADR-004) decoupled from metrics: the
// caller wires it to internal/metrics, the client imports nothing. nil is ignored.
func WithUsageHook(fn func(model string, prompt, completion int)) Option {
return func(c *Client) { c.usageHook = fn }
}
// New constructs a Client. // New constructs a Client.
func New(baseURL, apiKey, model string, timeout time.Duration, opts ...Option) *Client { func New(baseURL, apiKey, model string, timeout time.Duration, opts ...Option) *Client {
c := &Client{ c := &Client{
@@ -81,6 +90,10 @@ type chatResponse struct {
Choices []struct { Choices []struct {
Message message `json:"message"` Message message `json:"message"`
} `json:"choices"` } `json:"choices"`
Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
} `json:"usage"`
} }
// Complete sends a system + user message and returns the assistant's reply. // Complete sends a system + user message and returns the assistant's reply.
@@ -152,5 +165,8 @@ func (c *Client) Complete(ctx context.Context, system, user string) (string, err
if len(cr.Choices) == 0 { if len(cr.Choices) == 0 {
return "", fmt.Errorf("LLM returned no choices") return "", fmt.Errorf("LLM returned no choices")
} }
if c.usageHook != nil {
c.usageHook(c.model, cr.Usage.PromptTokens, cr.Usage.CompletionTokens)
}
return cr.Choices[0].Message.Content, nil return cr.Choices[0].Message.Content, nil
} }
+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) { func TestClient_ReturnsErrorOnNon200(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Error(w, "overloaded", http.StatusServiceUnavailable) http.Error(w, "overloaded", http.StatusServiceUnavailable)