feat: 增加token 统计

This commit is contained in:
ViperEkura 2026-04-13 23:04:41 +08:00
parent 30fc1779f4
commit 6f9bff1f1f
6 changed files with 95 additions and 10 deletions

View File

@ -26,7 +26,12 @@
</template>
</div>
<div class="message-footer">
<span class="token-count" v-if="tokenCount">{{ tokenCount }} tokens</span>
<span class="token-info" v-if="usage">
<span class="token-item" v-if="usage.prompt_tokens">输入: {{ usage.prompt_tokens }}</span>
<span class="token-item" v-if="usage.completion_tokens">输出: {{ usage.completion_tokens }}</span>
<span class="token-item total" v-if="usage.total_tokens">总计: {{ usage.total_tokens }}</span>
</span>
<span class="token-count" v-else-if="tokenCount">{{ tokenCount }} tokens</span>
<span class="message-time">{{ formatTime(createdAt) }}</span>
<button v-if="role === 'assistant'" class="ghost-btn accent" @click="copyContent" title="复制">
<span v-html="copyIcon"></span>
@ -50,6 +55,7 @@ const props = defineProps({
toolCalls: { type: Array, default: () => [] },
processSteps: { type: Array, default: () => [] },
tokenCount: { type: Number, default: 0 },
usage: { type: Object, default: null },
createdAt: { type: String, default: '' },
deletable: { type: Boolean, default: false },
attachments: { type: Array, default: () => [] },
@ -131,6 +137,25 @@ const trashIcon = `<svg viewBox="0 0 24 24" width="14" height="14" fill="none" s
font-size: 12px;
}
.token-info {
display: flex;
align-items: center;
gap: 6px;
font-size: 11px;
color: var(--text-tertiary);
}
.token-item {
padding: 2px 6px;
background: var(--bg-code);
border-radius: 4px;
}
.token-item.total {
font-weight: 600;
color: var(--accent-primary);
}
.token-count,
.message-time {
font-size: 12px;

View File

@ -276,6 +276,8 @@ const alertIcon = `<svg viewBox="0 0 24 24" width="16" height="16" fill="none" s
cursor: pointer;
font-size: 13px;
transition: background 0.15s;
width: 100%;
box-sizing: border-box;
}
.thinking .step-header:hover,
@ -370,8 +372,9 @@ const alertIcon = `<svg viewBox="0 0 24 24" width="16" height="16" fill="none" s
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
flex: 1;
flex: 1 1 auto;
min-width: 0;
max-width: 300px;
}
.arrow.open {

View File

@ -41,6 +41,7 @@
:tool-calls="msg.tool_calls"
:process-steps="msg.process_steps"
:token-count="msg.token_count"
:usage="msg.usage"
:created-at="msg.created_at"
:deletable="msg.role === 'user'"
:attachments="msg.attachments"
@ -233,13 +234,15 @@ const sendMessage = async () => {
streamingMessage.value.process_steps.push(step)
}
},
onDone: () => {
onDone: (data) => {
//
autoScroll.value = true
if (streamingMessage.value) {
messages.value.push({
...streamingMessage.value,
created_at: new Date().toISOString()
created_at: new Date().toISOString(),
token_count: data.token_count,
usage: data.usage
})
//

View File

@ -167,6 +167,7 @@ class Message(Base):
role: Mapped[str] = mapped_column(String(16), nullable=False) # user, assistant, system, tool
content: Mapped[str] = mapped_column(Text, nullable=False, default="")
token_count: Mapped[int] = mapped_column(Integer, default=0)
usage: Mapped[Optional[str]] = mapped_column(Text, nullable=True) # JSON string for usage info
created_at: Mapped[datetime] = mapped_column(DateTime, default=local_now)
# Relationships
@ -184,6 +185,13 @@ class Message(Base):
"created_at": self.created_at.isoformat() if self.created_at else None
}
# Parse usage JSON
if self.usage:
try:
result["usage"] = json.loads(self.usage)
except json.JSONDecodeError:
result["usage"] = None
# Parse content JSON
try:
content_obj = json.loads(self.content) if self.content else {}

View File

@ -43,6 +43,7 @@ def list_conversations(
db: Session = Depends(get_db)
):
"""Get conversation list"""
import json
query = db.query(Conversation).filter(Conversation.user_id == current_user.id)
result = paginate(query.order_by(Conversation.updated_at.desc()), page, page_size)
@ -56,6 +57,23 @@ def list_conversations(
).order_by(Message.created_at).first()
if first_msg:
conv_dict['first_message'] = first_msg.content[:50] + ('...' if len(first_msg.content) > 50 else '')
# Calculate total tokens from all assistant messages in this conversation
assistant_messages = db.query(Message).filter(
Message.conversation_id == c.id,
Message.role == 'assistant'
).all()
total_tokens = 0
for msg in assistant_messages:
total_tokens += msg.token_count or 0
# Also try to get usage from the usage field
if msg.usage:
try:
usage_obj = json.loads(msg.usage)
total_tokens = usage_obj.get("total_tokens", total_tokens)
except:
pass
conv_dict['token_count'] = total_tokens
items.append(conv_dict)
return success_response(data={

View File

@ -1,6 +1,7 @@
"""Chat service module"""
import json
import uuid
import logging
from typing import List, Dict, Any, AsyncGenerator, Optional
from luxx.models import Conversation, Message
@ -9,7 +10,7 @@ from luxx.tools.core import registry
from luxx.services.llm_client import LLMClient
from luxx.config import config
logger = logging.getLogger(__name__)
# Maximum iterations to prevent infinite loops
MAX_ITERATIONS = 10
@ -130,6 +131,13 @@ class ChatService:
all_tool_results = []
step_index = 0
# Token usage tracking
total_usage = {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0
}
# Global step IDs for thinking and text (persist across iterations)
thinking_step_id = None
thinking_step_idx = None
@ -186,6 +194,13 @@ class ChatService:
yield _sse_event("error", {"content": f"Failed to parse response: {data_str}"})
return
# 提取 API 返回的 usage 信息
if "usage" in chunk:
usage = chunk["usage"]
total_usage["prompt_tokens"] = usage.get("prompt_tokens", 0)
total_usage["completion_tokens"] = usage.get("completion_tokens", 0)
total_usage["total_tokens"] = usage.get("total_tokens", 0)
# Check for error in response
if "error" in chunk:
error_msg = chunk["error"].get("message", str(chunk["error"]))
@ -362,18 +377,26 @@ class ChatService:
# No tool calls - final iteration, save message
msg_id = str(uuid.uuid4())
# 使用 API 返回的真实 completion_tokens如果 API 没返回则降级使用估算值
actual_token_count = total_usage.get("completion_tokens", 0) or len(full_content) // 4
logger.info(f"[TOKEN] total_usage: {total_usage}, actual_token_count: {actual_token_count}")
self._save_message(
conversation.id,
msg_id,
full_content,
all_tool_calls,
all_tool_results,
all_steps
all_steps,
actual_token_count,
total_usage
)
yield _sse_event("done", {
"message_id": msg_id,
"token_count": len(full_content) // 4
"token_count": actual_token_count,
"usage": total_usage
})
return
@ -386,7 +409,9 @@ class ChatService:
full_content,
all_tool_calls,
all_tool_results,
all_steps
all_steps,
actual_token_count,
total_usage
)
yield _sse_event("error", {"content": "Exceeded maximum tool call iterations"})
@ -400,7 +425,9 @@ class ChatService:
full_content: str,
all_tool_calls: list,
all_tool_results: list,
all_steps: list
all_steps: list,
token_count: int = 0,
usage: dict = None
):
"""Save the assistant message to database."""
from luxx.database import SessionLocal
@ -420,7 +447,8 @@ class ChatService:
conversation_id=conversation_id,
role="assistant",
content=json.dumps(content_json, ensure_ascii=False),
token_count=len(full_content) // 4
token_count=token_count,
usage=json.dumps(usage) if usage else None
)
db.add(msg)
db.commit()