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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user