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