mirror of
https://github.com/alibaba/open-code-review.git
synced 2026-07-09 17:28:58 +00:00
fix(viewer): use real API token usage in session page instead of tiktoken estimates
The session page was showing significantly lower token counts than the console summary because it used local tiktoken estimation while the console used actual API-reported usage data. Now SetResponse prefers resp.Usage when available (falling back to tiktoken), persists cache token fields to JSONL, and the viewer displays cache read/write stats alongside prompt/completion totals.
This commit is contained in:
kite
parent
f24334deb2
commit
0d530f3b73
4 changed files with 59 additions and 17 deletions
|
|
@ -67,8 +67,9 @@ type TaskRecord struct {
|
|||
fileSession *FileSession // back-reference for JSONL persistence
|
||||
}
|
||||
|
||||
// TokenUsage holds estimated token usage for a single LLM request/response cycle.
|
||||
// Calculated locally using tiktoken so it works with any model provider.
|
||||
// TokenUsage holds token usage for a single LLM request/response cycle.
|
||||
// Uses actual token counts from the API response when available,
|
||||
// falling back to local estimation via tiktoken.
|
||||
type TokenUsage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
|
|
@ -219,8 +220,8 @@ func copyMessagesForJSON(msgs []llm.Message) any {
|
|||
}
|
||||
|
||||
// SetResponse records the LLM response in the most recent TaskRecord of the given type.
|
||||
// It computes estimated token usage locally using tiktoken, independent of any API format,
|
||||
// and writes an llm_response record to the JSONL stream.
|
||||
// It uses actual token usage from the API response when available, falling back to
|
||||
// local estimation via tiktoken, and writes an llm_response record to the JSONL stream.
|
||||
func (tr *TaskRecord) SetResponse(resp *llm.ChatResponse, duration time.Duration) {
|
||||
if resp == nil || len(resp.Choices) == 0 {
|
||||
tr.SetError(fmt.Errorf("empty response"), duration)
|
||||
|
|
@ -232,16 +233,24 @@ func (tr *TaskRecord) SetResponse(resp *llm.ChatResponse, duration time.Duration
|
|||
content = *choice.Message.Content
|
||||
}
|
||||
|
||||
// Estimate tokens using tiktoken — works with any model provider.
|
||||
var promptTokens int
|
||||
for _, m := range tr.RequestMessages {
|
||||
promptTokens += llm.CountTokens(m.ExtractText())
|
||||
var promptTokens, completionTokens, cacheReadTokens, cacheWriteTokens int
|
||||
if resp.Usage != nil {
|
||||
promptTokens = int(resp.Usage.PromptTokens)
|
||||
completionTokens = int(resp.Usage.CompletionTokens)
|
||||
cacheReadTokens = int(resp.Usage.CacheReadTokens)
|
||||
cacheWriteTokens = int(resp.Usage.CacheWriteTokens)
|
||||
} else {
|
||||
for _, m := range tr.RequestMessages {
|
||||
promptTokens += llm.CountTokens(m.ExtractText())
|
||||
}
|
||||
completionTokens = llm.CountTokens(content)
|
||||
}
|
||||
completionTokens := llm.CountTokens(content)
|
||||
|
||||
usage := &TokenUsage{
|
||||
PromptTokens: promptTokens,
|
||||
CompletionTokens: completionTokens,
|
||||
CacheReadTokens: cacheReadTokens,
|
||||
CacheWriteTokens: cacheWriteTokens,
|
||||
}
|
||||
|
||||
tr.Response = &ResponseRecord{
|
||||
|
|
@ -262,7 +271,7 @@ func (tr *TaskRecord) SetResponse(resp *llm.ChatResponse, duration time.Duration
|
|||
"arguments": tc.Function.Arguments,
|
||||
})
|
||||
}
|
||||
p.WriteLLMResponse(fs.FilePath, tr.Type, content, toolCallsJSON, resp.Model, promptTokens, completionTokens, duration)
|
||||
p.WriteLLMResponse(fs.FilePath, tr.Type, content, toolCallsJSON, resp.Model, *usage, duration)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -177,7 +177,7 @@ func (jw *jsonlWriter) WriteLLMRequest(filePath string, taskType TaskType, reque
|
|||
}
|
||||
|
||||
// WriteLLMResponse writes a response entry with model, content, tool calls, usage.
|
||||
func (jw *jsonlWriter) WriteLLMResponse(filePath string, taskType TaskType, content string, toolCalls []map[string]any, model string, promptTokens, completionTokens int, duration time.Duration) string {
|
||||
func (jw *jsonlWriter) WriteLLMResponse(filePath string, taskType TaskType, content string, toolCalls []map[string]any, model string, usage TokenUsage, duration time.Duration) string {
|
||||
uuid := generateUUID()
|
||||
|
||||
jw.mu.Lock()
|
||||
|
|
@ -195,8 +195,10 @@ func (jw *jsonlWriter) WriteLLMResponse(filePath string, taskType TaskType, cont
|
|||
"tool_calls": toolCalls,
|
||||
"duration_ms": duration.Milliseconds(),
|
||||
"usage": map[string]int{
|
||||
"prompt_tokens": promptTokens,
|
||||
"completion_tokens": completionTokens,
|
||||
"prompt_tokens": usage.PromptTokens,
|
||||
"completion_tokens": usage.CompletionTokens,
|
||||
"cache_read_tokens": usage.CacheReadTokens,
|
||||
"cache_write_tokens": usage.CacheWriteTokens,
|
||||
},
|
||||
}
|
||||
jw.writeRecordLocked(rec)
|
||||
|
|
|
|||
|
|
@ -200,12 +200,14 @@ type ViewSession struct {
|
|||
Files []*FileGroup // ordered by file path
|
||||
}
|
||||
|
||||
// TokenUsageSummary aggregates estimated token counts across the session.
|
||||
// TokenUsageSummary aggregates token counts across the session.
|
||||
type TokenUsageSummary struct {
|
||||
TotalPromptTokens int
|
||||
TotalCompletionTokens int
|
||||
TotalCacheReadTokens int
|
||||
TotalCacheWriteTokens int
|
||||
RequestCount int
|
||||
FileTokenBreakdown []FileTokenUsage // populated after parsing
|
||||
FileTokenBreakdown []FileTokenUsage
|
||||
}
|
||||
|
||||
// FileTokenUsage tracks token totals for a single file within a session.
|
||||
|
|
@ -213,6 +215,8 @@ type FileTokenUsage struct {
|
|||
FilePath string
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
CacheReadTokens int
|
||||
CacheWriteTokens int
|
||||
}
|
||||
|
||||
// FileGroup aggregates records for a single file.
|
||||
|
|
@ -242,6 +246,8 @@ type TaskCard struct {
|
|||
Model string
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
CacheReadTokens int
|
||||
CacheWriteTokens int
|
||||
}
|
||||
|
||||
// ToolCallInfo summarizes a single tool call.
|
||||
|
|
@ -334,6 +340,8 @@ func LoadSession(root, encodedRepo, sessionID string) (*ViewSession, error) {
|
|||
|
||||
promptTok := 0
|
||||
completionTok := 0
|
||||
cacheReadTok := 0
|
||||
cacheWriteTok := 0
|
||||
if usage, ok := rec["usage"].(map[string]any); ok {
|
||||
if v, ok := usage["prompt_tokens"].(float64); ok {
|
||||
promptTok = int(v)
|
||||
|
|
@ -341,6 +349,12 @@ func LoadSession(root, encodedRepo, sessionID string) (*ViewSession, error) {
|
|||
if v, ok := usage["completion_tokens"].(float64); ok {
|
||||
completionTok = int(v)
|
||||
}
|
||||
if v, ok := usage["cache_read_tokens"].(float64); ok {
|
||||
cacheReadTok = int(v)
|
||||
}
|
||||
if v, ok := usage["cache_write_tokens"].(float64); ok {
|
||||
cacheWriteTok = int(v)
|
||||
}
|
||||
}
|
||||
|
||||
tt, _ := rec["taskType"].(string)
|
||||
|
|
@ -355,6 +369,8 @@ func LoadSession(root, encodedRepo, sessionID string) (*ViewSession, error) {
|
|||
card.Error = errStr
|
||||
card.PromptTokens = promptTok
|
||||
card.CompletionTokens = completionTok
|
||||
card.CacheReadTokens = cacheReadTok
|
||||
card.CacheWriteTokens = cacheWriteTok
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -453,11 +469,15 @@ func LoadSession(root, encodedRepo, sessionID string) (*ViewSession, error) {
|
|||
for _, c := range cards {
|
||||
vs.TokenUsage.TotalPromptTokens += c.PromptTokens
|
||||
vs.TokenUsage.TotalCompletionTokens += c.CompletionTokens
|
||||
vs.TokenUsage.TotalCacheReadTokens += c.CacheReadTokens
|
||||
vs.TokenUsage.TotalCacheWriteTokens += c.CacheWriteTokens
|
||||
if c.ResponseContent != "" || c.PromptTokens > 0 {
|
||||
vs.TokenUsage.RequestCount++
|
||||
}
|
||||
ft.PromptTokens += c.PromptTokens
|
||||
ft.CompletionTokens += c.CompletionTokens
|
||||
ft.CacheReadTokens += c.CacheReadTokens
|
||||
ft.CacheWriteTokens += c.CacheWriteTokens
|
||||
}
|
||||
}
|
||||
fileBreakdown = append(fileBreakdown, ft)
|
||||
|
|
|
|||
|
|
@ -47,6 +47,16 @@
|
|||
<div class="token-value">{{.Session.TokenUsage.RequestCount}}</div>
|
||||
<div class="token-label">LLM Requests</div>
|
||||
</div>
|
||||
{{if or .Session.TokenUsage.TotalCacheReadTokens .Session.TokenUsage.TotalCacheWriteTokens}}
|
||||
<div class="token-item">
|
||||
<div class="token-value">{{.Session.TokenUsage.TotalCacheReadTokens}}</div>
|
||||
<div class="token-label">Cache Read</div>
|
||||
</div>
|
||||
<div class="token-item">
|
||||
<div class="token-value">{{.Session.TokenUsage.TotalCacheWriteTokens}}</div>
|
||||
<div class="token-label">Cache Write</div>
|
||||
</div>
|
||||
{{end}}
|
||||
{{if .Session.Summary.LLMFailures}}
|
||||
<div class="token-item">
|
||||
<div class="token-value badge-error">{{.Session.Summary.LLMFailures}}</div>
|
||||
|
|
@ -56,13 +66,14 @@
|
|||
</div>
|
||||
{{with .Session.TokenUsage.FileTokenBreakdown}}
|
||||
<table class="token-table">
|
||||
<thead><tr><th>File</th><th>Prompt</th><th>Completion</th><th>Total</th></tr></thead>
|
||||
<thead><tr><th>File</th><th>Prompt</th><th>Completion</th>{{if or $.Session.TokenUsage.TotalCacheReadTokens $.Session.TokenUsage.TotalCacheWriteTokens}}<th>Cache Read</th><th>Cache Write</th>{{end}}<th>Total</th></tr></thead>
|
||||
<tbody>
|
||||
{{range .}}
|
||||
<tr>
|
||||
<td title="{{.FilePath}}">{{.FilePath | truncate 60}}</td>
|
||||
<td>{{.PromptTokens}}</td>
|
||||
<td>{{.CompletionTokens}}</td>
|
||||
{{if or $.Session.TokenUsage.TotalCacheReadTokens $.Session.TokenUsage.TotalCacheWriteTokens}}<td>{{.CacheReadTokens}}</td><td>{{.CacheWriteTokens}}</td>{{end}}
|
||||
<td><strong>{{add .PromptTokens .CompletionTokens}}</strong></td>
|
||||
</tr>
|
||||
{{end}}
|
||||
|
|
@ -101,7 +112,7 @@
|
|||
<div class="card-header">
|
||||
<span class="request-label">Request #{{.RequestNo}}</span>
|
||||
{{if .Model}}<span class="badge badge-model">{{.Model}}</span>{{end}}
|
||||
{{if or .PromptTokens .CompletionTokens}}<span class="badge badge-tokens">P:{{.PromptTokens}} C:{{.CompletionTokens}}</span>{{end}}
|
||||
{{if or .PromptTokens .CompletionTokens}}<span class="badge badge-tokens">P:{{.PromptTokens}} C:{{.CompletionTokens}}{{if or .CacheReadTokens .CacheWriteTokens}} CR:{{.CacheReadTokens}} CW:{{.CacheWriteTokens}}{{end}}</span>{{end}}
|
||||
{{if .DurationMs}}<span class="badge badge-duration">{{.DurationMs}}ms</span>{{end}}
|
||||
{{if .Error}}<span class="badge badge-error">{{.Error}}</span>{{end}}
|
||||
</div>
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue