diff --git a/internal/adapters/llm/client.go b/internal/adapters/llm/client.go index 430876e..f2fc1a6 100644 --- a/internal/adapters/llm/client.go +++ b/internal/adapters/llm/client.go @@ -32,6 +32,7 @@ type Client struct { model string maxTokens int httpClient *http.Client + usageHook func(model string, prompt, completion int) } // 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. func New(baseURL, apiKey, model string, timeout time.Duration, opts ...Option) *Client { c := &Client{ @@ -81,6 +90,10 @@ type chatResponse struct { Choices []struct { Message message `json:"message"` } `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. @@ -152,5 +165,8 @@ func (c *Client) Complete(ctx context.Context, system, user string) (string, err if len(cr.Choices) == 0 { 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 } diff --git a/internal/adapters/llm/client_test.go b/internal/adapters/llm/client_test.go index 1ec7b74..186d358 100644 --- a/internal/adapters/llm/client_test.go +++ b/internal/adapters/llm/client_test.go @@ -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)