admin/MonkeyCode
1
0
Fork 0
mirror of https://github.com/chaitin/MonkeyCode.git synced 2026-02-02 14:53:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #155 from yokowu/fix-chat-stat

Browse Source 
fix: 修复输出token统计问题
This commit is contained in:
Yoko
2025-07-25 18:17:52 +08:00
committed by GitHub
parent 9036c75ce3 c9996a86a8
commit 341e1da36e
1 changed files with 24 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
backend/internal/proxy/recorder.go
View File

@@ -18,7 +18,6 @@ import (
"github.com/chaitin/MonkeyCode/backend/consts"
"github.com/chaitin/MonkeyCode/backend/domain"
"github.com/chaitin/MonkeyCode/backend/pkg/diff"
"github.com/chaitin/MonkeyCode/backend/pkg/promptparser"
)
type Recorder struct {
@@ -219,8 +218,12 @@ func (r *Recorder) handleJson(rc *domain.RecordParam) {
}
if len(resp.Choices) > 0 {
rc.Completion = resp.Choices[0].Message.Content
rc.InputTokens = int64(resp.Usage.PromptTokens)
rc.OutputTokens = int64(resp.Usage.CompletionTokens)
if input := resp.Usage.PromptTokens; input > 0 {
rc.InputTokens = int64(input)
}
if output := resp.Usage.CompletionTokens; output > 0 {
rc.OutputTokens = int64(output)
}
}
case consts.ModelTypeCoder:
@@ -232,8 +235,12 @@ func (r *Recorder) handleJson(rc *domain.RecordParam) {
if rc.TaskID == "" {
rc.TaskID = resp.ID
}
rc.InputTokens = int64(resp.Usage.PromptTokens)
rc.OutputTokens = int64(resp.Usage.CompletionTokens)
if input := resp.Usage.PromptTokens; input > 0 {
rc.InputTokens = int64(input)
}
if output := resp.Usage.CompletionTokens; output > 0 {
rc.OutputTokens = int64(output)
}
if len(resp.Choices) > 0 {
rc.Completion = resp.Choices[0].Text
rc.CodeLines = int64(strings.Count(resp.Choices[0].Text, "\n"))
@@ -290,8 +297,12 @@ func (r *Recorder) processSSELine(ctx context.Context, line string, rc *domain.R
return nil
}
if resp.Usage != nil {
rc.InputTokens = int64(resp.Usage.PromptTokens)
rc.OutputTokens += int64(resp.Usage.CompletionTokens)
if input := resp.Usage.PromptTokens; input > 0 {
rc.InputTokens = int64(input)
}
if output := resp.Usage.CompletionTokens; output > 0 {
rc.OutputTokens = int64(output)
}
}
if len(resp.Choices) > 0 {
content := resp.Choices[0].Delta.Content
@@ -310,8 +321,12 @@ func (r *Recorder) processSSELine(ctx context.Context, line string, rc *domain.R
if rc.TaskID == "" {
rc.TaskID = resp.ID
}
rc.InputTokens = int64(resp.Usage.PromptTokens)
rc.OutputTokens += int64(resp.Usage.CompletionTokens)
if input := resp.Usage.PromptTokens; input > 0 {
rc.InputTokens = int64(input)
}
if output := resp.Usage.CompletionTokens; output > 0 {
rc.OutputTokens = int64(output)
}
if len(resp.Choices) > 0 {
rc.Completion += resp.Choices[0].Text
rc.CodeLines += int64(strings.Count(resp.Choices[0].Text, "\n"))
@@ -345,36 +360,3 @@ func (r *Recorder) Read(p []byte) (n int, err error) {
}
return
}
func (r *Recorder) getPrompt(ctx context.Context, req *openai.ChatCompletionRequest) string {
prompt := ""
parse := promptparser.New(promptparser.KindTask)
for _, message := range req.Messages {
if message.Role == "system" {
continue
}
if strings.Contains(message.Content, "<task>") ||
strings.Contains(message.Content, "<feedback>") ||
strings.Contains(message.Content, "<user_message>") {
if info, err := parse.Parse(message.Content); err == nil {
prompt = info.Prompt
} else {
r.logger.With("message", message.Content).WarnContext(ctx, "解析Prompt失败", "error", err)
}
}
for _, m := range message.MultiContent {
if strings.Contains(m.Text, "<task>") ||
strings.Contains(m.Text, "<feedback>") ||
strings.Contains(m.Text, "<user_message>") {
if info, err := parse.Parse(m.Text); err == nil {
prompt = info.Prompt
} else {
r.logger.With("message", m.Text).WarnContext(ctx, "解析Prompt失败", "error", err)
}
}
}
}
return prompt
}