adamdottv 9 месяцев назад
Родитель
Сommit
051d7d7936
1 измененных файлов с 38 добавлено и 2 удалено
  1. 38 2
      internal/llm/provider/provider.go

+ 38 - 2
internal/llm/provider/provider.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/opencode-ai/opencode/internal/llm/models"
 	"github.com/opencode-ai/opencode/internal/llm/tools"
+	"github.com/opencode-ai/opencode/internal/logging"
 	"github.com/opencode-ai/opencode/internal/message"
 )
 
@@ -163,7 +164,17 @@ func (p *baseProvider[C]) cleanMessages(messages []message.Message) (cleaned []m
 
 func (p *baseProvider[C]) SendMessages(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
 	messages = p.cleanMessages(messages)
-	return p.client.send(ctx, messages, tools)
+	response, err := p.client.send(ctx, messages, tools)
+	if err == nil && response != nil {
+		logging.Debug("API request token usage", 
+			"model", p.options.model.Name,
+			"input_tokens", response.Usage.InputTokens,
+			"output_tokens", response.Usage.OutputTokens,
+			"cache_creation_tokens", response.Usage.CacheCreationTokens,
+			"cache_read_tokens", response.Usage.CacheReadTokens,
+			"total_tokens", response.Usage.InputTokens + response.Usage.OutputTokens)
+	}
+	return response, err
 }
 
 func (p *baseProvider[C]) Model() models.Model {
@@ -176,7 +187,32 @@ func (p *baseProvider[C]) MaxTokens() int64 {
 
 func (p *baseProvider[C]) StreamResponse(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {
 	messages = p.cleanMessages(messages)
-	return p.client.stream(ctx, messages, tools)
+	eventChan := p.client.stream(ctx, messages, tools)
+	
+	// Create a new channel to intercept events
+	wrappedChan := make(chan ProviderEvent)
+	
+	go func() {
+		defer close(wrappedChan)
+		
+		for event := range eventChan {
+			// Pass the event through
+			wrappedChan <- event
+			
+			// Log token usage when we get the complete event
+			if event.Type == EventComplete && event.Response != nil {
+				logging.Debug("API streaming request token usage", 
+					"model", p.options.model.Name,
+					"input_tokens", event.Response.Usage.InputTokens,
+					"output_tokens", event.Response.Usage.OutputTokens,
+					"cache_creation_tokens", event.Response.Usage.CacheCreationTokens,
+					"cache_read_tokens", event.Response.Usage.CacheReadTokens,
+					"total_tokens", event.Response.Usage.InputTokens + event.Response.Usage.OutputTokens)
+			}
+		}
+	}()
+	
+	return wrappedChan
 }
 
 func WithAPIKey(apiKey string) ProviderClientOption {