Add debug stats

internal/llm/client_completions.go

@@ -17,19 +17,27 @@ import (
 // maximum number of iterations.
 var ErrMaxIterations = fmt.Errorf("max iterations exceeded")
 
+// CompletionsResult bundles the messages generated during a completions
+// request along with usage and timings from the final API response.
+type CompletionsResult struct {
+	Messages []Message
+	Usage    *Usage
+	Timings  *Timings
+}
+
 // Completions sends messages to the LLM using the specified model.
 // If systemMessage is non-empty, it overrides the configured system message.
 // maxIterations caps the number of agent loop iterations.
 // A value of 0 means unlimited (no cap).
 // Returns all messages generated during the completions request, including
-// tool calls and tool results.
+// tool calls and tool results, along with usage and timings from the final response.
 func (c *Client) Completions(
 	ctx context.Context,
 	model string,
 	systemMessage string,
 	maxIterations int,
 	messages []Message,
-) ([]Message, error) {
+) (*CompletionsResult, error) {
 	ctx, cancel := context.WithTimeout(ctx, c.timeout)
 	defer cancel()
 
@@ -57,16 +65,29 @@ func (c *Client) Completions(
 	// generated during this call (the suffix of the conversation history).
 	start := len(messages)
 	// Loop for tool calls. maxIterations == 0 means unlimited.
+	var finalUsage *Usage
+	var finalTimings *Timings
 	for i := 0; maxIterations == 0 || i < maxIterations; i++ {
 		cr, err := c.completions(ctx, model, messages)
 		if err != nil {
 			return nil, fmt.Errorf("completions: %w", err)
 		}
+		// Capture usage and timings from this iteration.
+		if cr.Usage != nil {
+			finalUsage = cr.Usage
+		}
+		if cr.Timings != nil {
+			finalTimings = cr.Timings
+		}
 		// Use first choice for tool call loop.
 		message := &cr.Choices[0].Message
 		messages = append(messages, *message)
 		if len(message.ToolCalls) == 0 {
-			return messages[start:], nil
+			return &CompletionsResult{
+				Messages: messages[start:],
+				Usage:    finalUsage,
+				Timings:  finalTimings,
+			}, nil
 		}
 		toolResults, err := c.callTools(ctx, *message)
 		if err != nil {
@@ -75,7 +96,11 @@ func (c *Client) Completions(
 		messages = append(messages, toolResults...)
 	}
 
-	return messages[start:], ErrMaxIterations
+	return &CompletionsResult{
+		Messages: messages[start:],
+		Usage:    finalUsage,
+		Timings:  finalTimings,
+	}, ErrMaxIterations
 }
 
 // completions sends a single non-streaming request and returns the
@@ -185,6 +210,10 @@ type StreamEvent struct {
 	ReasoningDelta string
 	// Message is the complete chat completion message.
 	Message Message
+	// Usage is token usage info from the final chunk, sent on the done event.
+	Usage *Usage
+	// Timings is llama.cpp timing info from the final chunk, sent on the done event.
+	Timings *Timings
 }
 
 const (
@@ -273,13 +302,15 @@ type streamAccum struct {
 
 // completionsStream sends a single streaming request, accumulates the
 // response, and returns all assistant messages from the response choices.
-// Sends each message to the events channel.
+// Sends each message to the events channel. Usage and timings are included
+// on the done event for the last message.
 func (c *Client) completionsStream(ctx context.Context, messages []Message, model string, events chan<- StreamEvent) ([]Message, error) {
 	body, err := json.Marshal(ChatRequest{
-		Model:    model,
-		Messages: messages,
-		Stream:   true,
-		Tools:    c.tools,
+		Model:         model,
+		Messages:      messages,
+		Stream:        true,
+		Tools:         c.tools,
+		StreamOptions: &StreamOptions{IncludeUsage: true},
 	})
 	if err != nil {
 		return nil, fmt.Errorf(
@@ -361,6 +392,8 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
 
 	// Accumulate the streamed response per choice index.
 	accums := make(map[int]*streamAccum)
+	var usage *Usage
+	var timings *Timings
 
 	scanner := bufio.NewScanner(res.Body)
 	for scanner.Scan() {
@@ -437,6 +470,13 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
 				}
 			}
 		}
+		// Capture usage from the final chunk (has null choices).
+		if chunk.Usage != nil {
+			usage = chunk.Usage
+		}
+		if chunk.Timings != nil {
+			timings = chunk.Timings
+		}
 	}
 	res.Body.Close()
 
@@ -454,7 +494,7 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
 	slices.Sort(indices)
 
 	msgs := make([]Message, 0, len(accums))
-	for _, idx := range indices {
+	for i, idx := range indices {
 		acc := accums[idx]
 
 		// Content and refusal are mutually exclusive in the API.
@@ -478,9 +518,19 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
 			ReasoningContent: acc.reasoning.String(),
 			ToolCalls:        acc.toolCalls,
 		}
+
+		// Attach usage and timings to the done event of the last message.
+		var evUsage *Usage
+		var evTimings *Timings
+		if i == len(indices)-1 {
+			evUsage = usage
+			evTimings = timings
+		}
 		events <- StreamEvent{
 			Type:    streamEventDone,
 			Message: message,
+			Usage:   evUsage,
+			Timings: evTimings,
 		}
 		msgs = append(msgs, message)
 	}