Add token usage tracking and UI refinements

This commit is contained in:
temp
2026-08-24 00:05:05 +08:00
parent a34cab431a
commit e218316c55
23 changed files with 1285 additions and 93 deletions
+2
View File
@@ -1027,9 +1027,11 @@ func setupRoutes(
protected.DELETE("/batch-tasks/:queueId/tasks/:taskId", agentHandler.DeleteBatchTask)
// 对话历史
protected.GET("/usage/tokens", conversationHandler.GetTokenUsageStats)
protected.POST("/conversations", conversationHandler.CreateConversation)
protected.GET("/conversations", conversationHandler.ListConversations)
protected.GET("/conversations/:id", conversationHandler.GetConversation)
protected.GET("/conversations/:id/token-usage", conversationHandler.GetConversationTokenUsageStats)
protected.GET("/conversations/:id/plan-tasks", conversationHandler.GetConversationPlanTasks)
protected.GET("/messages/:id/process-details", conversationHandler.GetMessageProcessDetails)
protected.GET("/process-details/:id", conversationHandler.GetProcessDetail)
+2
View File
@@ -1350,6 +1350,8 @@ func (db *DB) AddProcessDetailWithID(messageID, conversationID, eventType, messa
return "", fmt.Errorf("添加过程详情失败: %w", err)
}
db.maybeRecordModelTokenUsage(messageID, conversationID, id, eventType, data)
return id, nil
}
+38
View File
@@ -216,6 +216,32 @@ func (db *DB) initTables() error {
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE
);`
// 创建模型 Token 用量表:process_details 负责时间线回放,本表负责结构化聚合统计。
createModelTokenUsageTable := `
CREATE TABLE IF NOT EXISTS model_token_usage (
id TEXT PRIMARY KEY,
process_detail_id TEXT NOT NULL UNIQUE,
message_id TEXT NOT NULL,
conversation_id TEXT NOT NULL,
project_id TEXT,
source TEXT NOT NULL DEFAULT '',
orchestration TEXT NOT NULL DEFAULT '',
reason TEXT NOT NULL DEFAULT '',
model TEXT NOT NULL DEFAULT '',
model_calls INTEGER NOT NULL DEFAULT 0,
prompt_tokens INTEGER NOT NULL DEFAULT 0,
completion_tokens INTEGER NOT NULL DEFAULT 0,
total_tokens INTEGER NOT NULL DEFAULT 0,
cached_tokens INTEGER NOT NULL DEFAULT 0,
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
created_at DATETIME NOT NULL,
updated_at DATETIME NOT NULL,
FOREIGN KEY (process_detail_id) REFERENCES process_details(id) ON DELETE CASCADE,
FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE,
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE,
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL
);`
// 创建工具执行记录表
createToolExecutionsTable := `
CREATE TABLE IF NOT EXISTS tool_executions (
@@ -719,6 +745,10 @@ func (db *DB) initTables() error {
CREATE INDEX IF NOT EXISTS idx_conversations_updated_at ON conversations(updated_at);
CREATE INDEX IF NOT EXISTS idx_process_details_message_id ON process_details(message_id);
CREATE INDEX IF NOT EXISTS idx_process_details_conversation_id ON process_details(conversation_id);
CREATE INDEX IF NOT EXISTS idx_model_token_usage_created_at ON model_token_usage(created_at);
CREATE INDEX IF NOT EXISTS idx_model_token_usage_conversation ON model_token_usage(conversation_id);
CREATE INDEX IF NOT EXISTS idx_model_token_usage_project ON model_token_usage(project_id);
CREATE INDEX IF NOT EXISTS idx_model_token_usage_model ON model_token_usage(model);
CREATE INDEX IF NOT EXISTS idx_tool_executions_tool_name ON tool_executions(tool_name);
CREATE INDEX IF NOT EXISTS idx_tool_executions_start_time ON tool_executions(start_time);
CREATE INDEX IF NOT EXISTS idx_tool_executions_status ON tool_executions(status);
@@ -806,6 +836,10 @@ func (db *DB) initTables() error {
return fmt.Errorf("创建process_details表失败: %w", err)
}
if _, err := db.Exec(createModelTokenUsageTable); err != nil {
return fmt.Errorf("创建model_token_usage表失败: %w", err)
}
if _, err := db.Exec(createToolExecutionsTable); err != nil {
return fmt.Errorf("创建tool_executions表失败: %w", err)
}
@@ -981,6 +1015,10 @@ func (db *DB) initTables() error {
if _, err := db.Exec(createIndexes); err != nil {
return fmt.Errorf("创建索引失败: %w", err)
}
if err := db.BackfillModelTokenUsageFromProcessDetails(); err != nil {
return fmt.Errorf("回填模型Token用量失败: %w", err)
}
db.logger.Debug("数据库表初始化完成")
return nil
}
+485
View File
@@ -0,0 +1,485 @@
package database
import (
"database/sql"
"encoding/json"
"fmt"
"math"
"strconv"
"strings"
"time"
"github.com/google/uuid"
"go.uber.org/zap"
)
const modelTokenUsageEventType = "eino_usage_summary"
// ModelTokenUsage records one model-usage summary emitted by an Agent run.
type ModelTokenUsage struct {
ID string `json:"id"`
ProcessDetailID string `json:"processDetailId"`
MessageID string `json:"messageId"`
ConversationID string `json:"conversationId"`
ProjectID string `json:"projectId,omitempty"`
Source string `json:"source"`
Orchestration string `json:"orchestration"`
Reason string `json:"reason"`
Model string `json:"model,omitempty"`
ModelCalls int64 `json:"modelCalls"`
PromptTokens int64 `json:"promptTokens"`
CompletionTokens int64 `json:"completionTokens"`
TotalTokens int64 `json:"totalTokens"`
CachedTokens int64 `json:"cachedTokens"`
ReasoningTokens int64 `json:"reasoningTokens"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
// ModelTokenUsageSummary is the aggregate shape used by dashboard and APIs.
type ModelTokenUsageSummary struct {
Events int64 `json:"events"`
ModelCalls int64 `json:"modelCalls"`
PromptTokens int64 `json:"promptTokens"`
CompletionTokens int64 `json:"completionTokens"`
TotalTokens int64 `json:"totalTokens"`
CachedTokens int64 `json:"cachedTokens"`
ReasoningTokens int64 `json:"reasoningTokens"`
}
// ModelTokenUsageBreakdown is a grouped aggregate row.
type ModelTokenUsageBreakdown struct {
Key string `json:"key"`
Label string `json:"label,omitempty"`
Events int64 `json:"events"`
ModelCalls int64 `json:"modelCalls"`
PromptTokens int64 `json:"promptTokens"`
CompletionTokens int64 `json:"completionTokens"`
TotalTokens int64 `json:"totalTokens"`
CachedTokens int64 `json:"cachedTokens"`
ReasoningTokens int64 `json:"reasoningTokens"`
}
// ModelTokenUsageStats is a compact API response for usage dashboards.
type ModelTokenUsageStats struct {
Summary ModelTokenUsageSummary `json:"summary"`
Today ModelTokenUsageSummary `json:"today"`
ByDay []ModelTokenUsageBreakdown `json:"byDay"`
ByModel []ModelTokenUsageBreakdown `json:"byModel"`
ByOrchestration []ModelTokenUsageBreakdown `json:"byOrchestration"`
Recent []ModelTokenUsage `json:"recent"`
}
// ModelTokenUsageFilter scopes usage queries.
type ModelTokenUsageFilter struct {
ConversationID string
ProjectID string
Since time.Time
Until time.Time
Days int
Access RBACListAccess
Limit int
}
func modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID string, data interface{}) (ModelTokenUsage, bool) {
m := mapFromUsageData(data)
if len(m) == 0 {
return ModelTokenUsage{}, false
}
usage := ModelTokenUsage{
ID: uuid.New().String(),
ProcessDetailID: strings.TrimSpace(processDetailID),
MessageID: strings.TrimSpace(messageID),
ConversationID: strings.TrimSpace(conversationID),
Source: strings.TrimSpace(fmt.Sprint(m["source"])),
Orchestration: strings.TrimSpace(fmt.Sprint(m["orchestration"])),
Reason: strings.TrimSpace(fmt.Sprint(m["reason"])),
Model: strings.TrimSpace(fmt.Sprint(m["model"])),
ModelCalls: usageInt64(m["modelCalls"]),
PromptTokens: usageInt64(m["promptTokens"]),
CompletionTokens: usageInt64(m["completionTokens"]),
TotalTokens: usageInt64(m["totalTokens"]),
CachedTokens: usageInt64(m["cachedTokens"]),
ReasoningTokens: usageInt64(m["reasoningTokens"]),
}
if usage.TotalTokens == 0 && (usage.PromptTokens > 0 || usage.CompletionTokens > 0) {
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
}
if usage.ProcessDetailID == "" || usage.MessageID == "" || usage.ConversationID == "" {
return ModelTokenUsage{}, false
}
if usage.ModelCalls == 0 && usage.TotalTokens == 0 && usage.PromptTokens == 0 && usage.CompletionTokens == 0 && usage.CachedTokens == 0 && usage.ReasoningTokens == 0 {
return ModelTokenUsage{}, false
}
return usage, true
}
func mapFromUsageData(data interface{}) map[string]interface{} {
switch v := data.(type) {
case nil:
return nil
case map[string]interface{}:
return v
case string:
var m map[string]interface{}
if err := json.Unmarshal([]byte(v), &m); err == nil {
return m
}
case []byte:
var m map[string]interface{}
if err := json.Unmarshal(v, &m); err == nil {
return m
}
default:
raw, err := json.Marshal(v)
if err == nil {
var m map[string]interface{}
if err := json.Unmarshal(raw, &m); err == nil {
return m
}
}
}
return nil
}
func usageInt64(v interface{}) int64 {
switch n := v.(type) {
case int:
return int64(n)
case int8:
return int64(n)
case int16:
return int64(n)
case int32:
return int64(n)
case int64:
return n
case uint:
return int64(n)
case uint8:
return int64(n)
case uint16:
return int64(n)
case uint32:
return int64(n)
case uint64:
if n > math.MaxInt64 {
return math.MaxInt64
}
return int64(n)
case float32:
return int64(n)
case float64:
return int64(n)
case json.Number:
i, _ := n.Int64()
return i
case string:
i, _ := strconv.ParseInt(strings.TrimSpace(n), 10, 64)
return i
default:
i, _ := strconv.ParseInt(strings.TrimSpace(fmt.Sprint(v)), 10, 64)
return i
}
}
func (db *DB) maybeRecordModelTokenUsage(messageID, conversationID, processDetailID, eventType string, data interface{}) {
if db == nil || eventType != modelTokenUsageEventType {
return
}
usage, ok := modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID, data)
if !ok {
return
}
if err := db.UpsertModelTokenUsage(usage); err != nil && db.logger != nil {
db.logger.Warn("保存模型Token用量失败",
zap.String("processDetailId", processDetailID),
zap.String("conversationId", conversationID),
zap.Error(err))
}
}
// UpsertModelTokenUsage persists usage with process_detail_id idempotency.
func (db *DB) UpsertModelTokenUsage(usage ModelTokenUsage) error {
if db == nil {
return fmt.Errorf("database is nil")
}
now := time.Now()
createdAt := usage.CreatedAt
if createdAt.IsZero() {
createdAt = now
}
if usage.ID == "" {
usage.ID = uuid.New().String()
}
var projectID sql.NullString
if err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, usage.ConversationID).Scan(&projectID); err != nil && err != sql.ErrNoRows {
return fmt.Errorf("查询对话项目失败: %w", err)
}
projectValue := interface{}(nil)
if projectID.Valid && strings.TrimSpace(projectID.String) != "" {
projectValue = strings.TrimSpace(projectID.String)
}
_, err := db.Exec(`
INSERT INTO model_token_usage (
id, process_detail_id, message_id, conversation_id, project_id,
source, orchestration, reason, model, model_calls,
prompt_tokens, completion_tokens, total_tokens, cached_tokens, reasoning_tokens,
created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(process_detail_id) DO UPDATE SET
message_id = excluded.message_id,
conversation_id = excluded.conversation_id,
project_id = excluded.project_id,
source = excluded.source,
orchestration = excluded.orchestration,
reason = excluded.reason,
model = excluded.model,
model_calls = excluded.model_calls,
prompt_tokens = excluded.prompt_tokens,
completion_tokens = excluded.completion_tokens,
total_tokens = excluded.total_tokens,
cached_tokens = excluded.cached_tokens,
reasoning_tokens = excluded.reasoning_tokens,
created_at = excluded.created_at,
updated_at = excluded.updated_at`,
usage.ID, usage.ProcessDetailID, usage.MessageID, usage.ConversationID, projectValue,
usage.Source, usage.Orchestration, usage.Reason, usage.Model, usage.ModelCalls,
usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens, usage.CachedTokens, usage.ReasoningTokens,
createdAt, now,
)
if err != nil {
return fmt.Errorf("写入模型Token用量失败: %w", err)
}
return nil
}
// BackfillModelTokenUsageFromProcessDetails makes existing timeline usage events queryable.
func (db *DB) BackfillModelTokenUsageFromProcessDetails() error {
if db == nil {
return nil
}
rows, err := db.Query(`
SELECT pd.id, pd.message_id, pd.conversation_id, pd.data, pd.created_at
FROM process_details pd
LEFT JOIN model_token_usage mtu ON mtu.process_detail_id = pd.id
WHERE pd.event_type = ?
AND (mtu.id IS NULL OR mtu.created_at != pd.created_at)`, modelTokenUsageEventType)
if err != nil {
return fmt.Errorf("查询历史模型Token用量失败: %w", err)
}
defer rows.Close()
for rows.Next() {
var processDetailID, messageID, conversationID string
var data sql.NullString
var createdAt string
if err := rows.Scan(&processDetailID, &messageID, &conversationID, &data, &createdAt); err != nil {
return fmt.Errorf("扫描历史模型Token用量失败: %w", err)
}
if !data.Valid {
continue
}
usage, ok := modelTokenUsageFromProcessDetail(messageID, conversationID, processDetailID, data.String)
if !ok {
continue
}
usage.CreatedAt = parseModelTokenUsageTime(createdAt)
if err := db.UpsertModelTokenUsage(usage); err != nil {
return err
}
}
if err := rows.Err(); err != nil {
return fmt.Errorf("遍历历史模型Token用量失败: %w", err)
}
return nil
}
func (db *DB) GetModelTokenUsageStats(filter ModelTokenUsageFilter) (*ModelTokenUsageStats, error) {
if db == nil {
return nil, fmt.Errorf("database is nil")
}
if filter.Days <= 0 {
filter.Days = 7
}
if filter.Limit <= 0 {
filter.Limit = 10
}
where, args := buildModelTokenUsageWhere(filter, "mtu", "c")
summary, err := db.queryModelTokenUsageSummary("SELECT "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where, args...)
if err != nil {
return nil, err
}
todayFilter := filter
now := time.Now()
todayFilter.Since = time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location())
todayWhere, todayArgs := buildModelTokenUsageWhere(todayFilter, "mtu", "c")
today, err := db.queryModelTokenUsageSummary("SELECT "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+todayWhere, todayArgs...)
if err != nil {
return nil, err
}
byDay, err := db.queryModelTokenUsageBreakdown(
"SELECT date(mtu.created_at) AS k, date(mtu.created_at) AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY date(mtu.created_at) ORDER BY k DESC LIMIT ?",
append(args, filter.Days)...,
)
if err != nil {
return nil, err
}
byModel, err := db.queryModelTokenUsageBreakdown(
"SELECT COALESCE(NULLIF(TRIM(mtu.model), ''), 'unknown') AS k, COALESCE(NULLIF(TRIM(mtu.model), ''), 'Unknown') AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY k ORDER BY SUM(mtu.total_tokens) DESC LIMIT ?",
append(args, filter.Limit)...,
)
if err != nil {
return nil, err
}
byOrch, err := db.queryModelTokenUsageBreakdown(
"SELECT COALESCE(NULLIF(TRIM(mtu.orchestration), ''), 'unknown') AS k, COALESCE(NULLIF(TRIM(mtu.orchestration), ''), 'Unknown') AS label, "+modelTokenUsageSummarySelect("mtu")+" FROM model_token_usage mtu JOIN conversations c ON c.id = mtu.conversation_id"+where+" GROUP BY k ORDER BY SUM(mtu.total_tokens) DESC LIMIT ?",
append(args, filter.Limit)...,
)
if err != nil {
return nil, err
}
recent, err := db.ListModelTokenUsage(filter)
if err != nil {
return nil, err
}
return &ModelTokenUsageStats{
Summary: summary,
Today: today,
ByDay: byDay,
ByModel: byModel,
ByOrchestration: byOrch,
Recent: recent,
}, nil
}
func modelTokenUsageSummarySelect(alias string) string {
p := ""
if alias != "" {
p = alias + "."
}
return fmt.Sprintf(`COUNT(%sid),
COALESCE(SUM(%smodel_calls), 0),
COALESCE(SUM(%sprompt_tokens), 0),
COALESCE(SUM(%scompletion_tokens), 0),
COALESCE(SUM(%stotal_tokens), 0),
COALESCE(SUM(%scached_tokens), 0),
COALESCE(SUM(%sreasoning_tokens), 0)`, p, p, p, p, p, p, p)
}
func buildModelTokenUsageWhere(filter ModelTokenUsageFilter, usageAlias, convAlias string) (string, []interface{}) {
where := " WHERE 1=1"
args := []interface{}{}
uPrefix := ""
if usageAlias != "" {
uPrefix = usageAlias + "."
}
if cid := strings.TrimSpace(filter.ConversationID); cid != "" {
where += " AND " + uPrefix + "conversation_id = ?"
args = append(args, cid)
}
where, args = appendConversationProjectFilter(where, args, filter.ProjectID, usageAlias)
if !filter.Since.IsZero() {
where += " AND " + uPrefix + "created_at >= ?"
args = append(args, filter.Since)
}
if !filter.Until.IsZero() {
where += " AND " + uPrefix + "created_at <= ?"
args = append(args, filter.Until)
}
where, args = appendConversationAccessFilter(where, args, filter.Access.UserID, filter.Access.Scope, convAlias)
return where, args
}
func (db *DB) queryModelTokenUsageSummary(query string, args ...interface{}) (ModelTokenUsageSummary, error) {
var s ModelTokenUsageSummary
err := db.QueryRow(query, args...).Scan(
&s.Events, &s.ModelCalls, &s.PromptTokens, &s.CompletionTokens,
&s.TotalTokens, &s.CachedTokens, &s.ReasoningTokens,
)
if err != nil {
return s, fmt.Errorf("查询模型Token用量汇总失败: %w", err)
}
return s, nil
}
func (db *DB) queryModelTokenUsageBreakdown(query string, args ...interface{}) ([]ModelTokenUsageBreakdown, error) {
rows, err := db.Query(query, args...)
if err != nil {
return nil, fmt.Errorf("查询模型Token用量分组失败: %w", err)
}
defer rows.Close()
out := []ModelTokenUsageBreakdown{}
for rows.Next() {
var row ModelTokenUsageBreakdown
if err := rows.Scan(
&row.Key, &row.Label, &row.Events, &row.ModelCalls, &row.PromptTokens,
&row.CompletionTokens, &row.TotalTokens, &row.CachedTokens, &row.ReasoningTokens,
); err != nil {
return nil, fmt.Errorf("扫描模型Token用量分组失败: %w", err)
}
out = append(out, row)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("遍历模型Token用量分组失败: %w", err)
}
return out, nil
}
func (db *DB) ListModelTokenUsage(filter ModelTokenUsageFilter) ([]ModelTokenUsage, error) {
if filter.Limit <= 0 {
filter.Limit = 20
}
if filter.Limit > 500 {
filter.Limit = 500
}
where, args := buildModelTokenUsageWhere(filter, "mtu", "c")
args = append(args, filter.Limit)
rows, err := db.Query(`
SELECT mtu.id, mtu.process_detail_id, mtu.message_id, mtu.conversation_id,
COALESCE(mtu.project_id, ''), mtu.source, mtu.orchestration, mtu.reason, mtu.model,
mtu.model_calls, mtu.prompt_tokens, mtu.completion_tokens, mtu.total_tokens,
mtu.cached_tokens, mtu.reasoning_tokens, mtu.created_at, mtu.updated_at
FROM model_token_usage mtu
JOIN conversations c ON c.id = mtu.conversation_id`+where+`
ORDER BY mtu.created_at DESC, mtu.rowid DESC
LIMIT ?`, args...)
if err != nil {
return nil, fmt.Errorf("查询模型Token用量明细失败: %w", err)
}
defer rows.Close()
out := []ModelTokenUsage{}
for rows.Next() {
var u ModelTokenUsage
var createdAt, updatedAt string
if err := rows.Scan(
&u.ID, &u.ProcessDetailID, &u.MessageID, &u.ConversationID, &u.ProjectID,
&u.Source, &u.Orchestration, &u.Reason, &u.Model, &u.ModelCalls,
&u.PromptTokens, &u.CompletionTokens, &u.TotalTokens, &u.CachedTokens,
&u.ReasoningTokens, &createdAt, &updatedAt,
); err != nil {
return nil, fmt.Errorf("扫描模型Token用量明细失败: %w", err)
}
u.CreatedAt = parseModelTokenUsageTime(createdAt)
u.UpdatedAt = parseModelTokenUsageTime(updatedAt)
out = append(out, u)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("遍历模型Token用量明细失败: %w", err)
}
return out, nil
}
func parseModelTokenUsageTime(s string) time.Time {
for _, layout := range []string{
"2006-01-02 15:04:05.999999999-07:00",
"2006-01-02 15:04:05.999999-07:00",
"2006-01-02 15:04:05",
time.RFC3339Nano,
time.RFC3339,
} {
if t, err := time.Parse(layout, strings.TrimSpace(s)); err == nil {
return t
}
}
return time.Time{}
}
@@ -0,0 +1,85 @@
package database
import (
"path/filepath"
"testing"
"go.uber.org/zap"
)
func TestModelTokenUsagePersistsFromUsageProcessDetail(t *testing.T) {
db := newModelTokenUsageTestDB(t)
conv, err := db.CreateConversation("usage", ConversationCreateMeta{})
if err != nil {
t.Fatalf("CreateConversation: %v", err)
}
msg, err := db.AddMessage(conv.ID, "assistant", "done", nil)
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
if err := db.AddProcessDetail(msg.ID, conv.ID, modelTokenUsageEventType, "usage", map[string]interface{}{
"source": "eino",
"orchestration": "deep",
"reason": "final",
"model": "gpt-test",
"modelCalls": 2,
"promptTokens": 10,
"completionTokens": 3,
"totalTokens": 13,
"cachedTokens": 4,
"reasoningTokens": 1,
}); err != nil {
t.Fatalf("AddProcessDetail: %v", err)
}
stats, err := db.GetModelTokenUsageStats(ModelTokenUsageFilter{})
if err != nil {
t.Fatalf("GetModelTokenUsageStats: %v", err)
}
if stats.Summary.Events != 1 || stats.Summary.ModelCalls != 2 || stats.Summary.TotalTokens != 13 || stats.Summary.CachedTokens != 4 || stats.Summary.ReasoningTokens != 1 {
t.Fatalf("summary = %#v", stats.Summary)
}
if len(stats.ByModel) != 1 || stats.ByModel[0].Key != "gpt-test" || stats.ByModel[0].TotalTokens != 13 {
t.Fatalf("by model = %#v", stats.ByModel)
}
}
func TestModelTokenUsageBackfillIsIdempotent(t *testing.T) {
db := newModelTokenUsageTestDB(t)
conv, err := db.CreateConversation("usage", ConversationCreateMeta{})
if err != nil {
t.Fatalf("CreateConversation: %v", err)
}
msg, err := db.AddMessage(conv.ID, "assistant", "done", nil)
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
if err := db.AddProcessDetail(msg.ID, conv.ID, modelTokenUsageEventType, "usage", map[string]interface{}{
"source": "eino", "modelCalls": 1, "promptTokens": 7, "completionTokens": 5, "totalTokens": 12,
}); err != nil {
t.Fatalf("AddProcessDetail: %v", err)
}
if err := db.BackfillModelTokenUsageFromProcessDetails(); err != nil {
t.Fatalf("Backfill 1: %v", err)
}
if err := db.BackfillModelTokenUsageFromProcessDetails(); err != nil {
t.Fatalf("Backfill 2: %v", err)
}
stats, err := db.GetModelTokenUsageStats(ModelTokenUsageFilter{})
if err != nil {
t.Fatalf("GetModelTokenUsageStats: %v", err)
}
if stats.Summary.Events != 1 || stats.Summary.TotalTokens != 12 {
t.Fatalf("summary after backfill = %#v", stats.Summary)
}
}
func newModelTokenUsageTestDB(t *testing.T) *DB {
t.Helper()
db, err := NewDB(filepath.Join(t.TempDir(), "usage.db"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
return db
}
+92
View File
@@ -0,0 +1,92 @@
package handler
import (
"net/http"
"strconv"
"strings"
"time"
"cyberstrike-ai/internal/database"
"cyberstrike-ai/internal/security"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
// GetTokenUsageStats returns model token usage aggregates for dashboard views.
func (h *ConversationHandler) GetTokenUsageStats(c *gin.Context) {
filter := tokenUsageFilterFromQuery(c)
if session, ok := security.CurrentSession(c); ok {
filter.Access = database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
}
stats, err := h.db.GetModelTokenUsageStats(filter)
if err != nil {
h.logger.Error("获取Token用量统计失败", zap.Error(err))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, stats)
}
// GetConversationTokenUsageStats returns token usage scoped to one conversation.
func (h *ConversationHandler) GetConversationTokenUsageStats(c *gin.Context) {
filter := tokenUsageFilterFromQuery(c)
filter.ConversationID = strings.TrimSpace(c.Param("id"))
if session, ok := security.CurrentSession(c); ok {
filter.Access = database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
}
stats, err := h.db.GetModelTokenUsageStats(filter)
if err != nil {
h.logger.Error("获取对话Token用量统计失败", zap.Error(err), zap.String("conversationId", filter.ConversationID))
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
c.JSON(http.StatusOK, stats)
}
func tokenUsageFilterFromQuery(c *gin.Context) database.ModelTokenUsageFilter {
days, _ := strconv.Atoi(strings.TrimSpace(c.DefaultQuery("days", "7")))
if days <= 0 {
days = 7
}
if days > 365 {
days = 365
}
limit, _ := strconv.Atoi(strings.TrimSpace(c.DefaultQuery("limit", "10")))
if limit <= 0 {
limit = 10
}
if limit > 500 {
limit = 500
}
filter := database.ModelTokenUsageFilter{
ConversationID: strings.TrimSpace(c.Query("conversation_id")),
ProjectID: strings.TrimSpace(c.Query("project_id")),
Days: days,
Limit: limit,
}
if since := parseTokenUsageQueryTime(c.Query("since")); !since.IsZero() {
filter.Since = since
} else if days > 0 {
now := time.Now()
start := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()).AddDate(0, 0, -(days - 1))
filter.Since = start
}
if until := parseTokenUsageQueryTime(c.Query("until")); !until.IsZero() {
filter.Until = until
}
return filter
}
func parseTokenUsageQueryTime(raw string) time.Time {
raw = strings.TrimSpace(raw)
if raw == "" {
return time.Time{}
}
for _, layout := range []string{time.RFC3339Nano, time.RFC3339, "2006-01-02"} {
if t, err := time.Parse(layout, raw); err == nil {
return t
}
}
return time.Time{}
}
@@ -371,5 +371,9 @@ func (s *einoRunRuntimeSession) emitUsageSummary(reason string) bool {
if s == nil || s.usage == nil {
return false
}
return s.usage.EmitOnce(s.conversationID, s.orchMode, reason, s.progress, s.logger)
modelName := ""
if s.args != nil {
modelName = s.args.ModelName
}
return s.usage.EmitOnce(s.conversationID, s.orchMode, reason, modelName, s.progress, s.logger)
}
@@ -61,6 +61,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
conversationID string,
orchestration string,
reason string,
modelName string,
progress func(eventType, message string, data interface{}),
logger *zap.Logger,
) bool {
@@ -81,6 +82,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
"source": "eino",
"orchestration": orchestration,
"reason": reason,
"model": modelName,
"modelCalls": s.ModelCalls,
"promptTokens": s.PromptTokens,
"completionTokens": s.CompletionTokens,
@@ -96,6 +98,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
zap.String("conversationId", conversationID),
zap.String("orchestration", orchestration),
zap.String("reason", reason),
zap.String("model", modelName),
zap.Int("modelCalls", s.ModelCalls),
zap.Int("promptTokens", s.PromptTokens),
zap.Int("completionTokens", s.CompletionTokens),
@@ -49,16 +49,16 @@ func TestEinoRunUsageAccumulatorEmitOnce(t *testing.T) {
}
}
if !acc.EmitOnce("conv-1", "deep", "final", progress, nil) {
if !acc.EmitOnce("conv-1", "deep", "final", "gpt-test", progress, nil) {
t.Fatal("first emit should return true")
}
if acc.EmitOnce("conv-1", "deep", "partial", progress, nil) {
if acc.EmitOnce("conv-1", "deep", "partial", "gpt-test", progress, nil) {
t.Fatal("second emit should return false")
}
if len(events) != 1 {
t.Fatalf("events = %#v, want one usage summary", events)
}
if events[0]["conversationId"] != "conv-1" || events[0]["orchestration"] != "deep" || events[0]["reason"] != "final" || events[0]["totalTokens"] != 3 {
if events[0]["conversationId"] != "conv-1" || events[0]["orchestration"] != "deep" || events[0]["reason"] != "final" || events[0]["model"] != "gpt-test" || events[0]["totalTokens"] != 3 {
t.Fatalf("event = %#v", events[0])
}
}
+2
View File
@@ -118,6 +118,8 @@ func permissionForRequest(method, fullPath string) string {
return "hitl:write"
case strings.HasPrefix(path, "/agent-loop"), strings.HasPrefix(path, "/batch-tasks"):
return crudPermission(method, "tasks")
case path == "/usage/tokens":
return "dashboard:read"
case strings.HasPrefix(path, "/conversations"), strings.HasPrefix(path, "/messages"), strings.HasPrefix(path, "/process-details"):
return crudPermission(method, "chat")
case strings.HasPrefix(path, "/groups"):
@@ -119,6 +119,12 @@ func TestRBACResourcePickerRequiresWritePermission(t *testing.T) {
}
}
func TestRBACMiddlewareMapsTokenUsageStatsToDashboardRead(t *testing.T) {
if got := permissionForRequest(http.MethodGet, "/api/usage/tokens"); got != "dashboard:read" {
t.Fatalf("token usage permission = %q, want dashboard:read", got)
}
}
func TestMCPInvocationPermissionIsSeparateFromMCPAdministration(t *testing.T) {
if got := permissionForRequest(http.MethodPost, "/api/mcp"); got != "mcp:execute" {
t.Fatalf("MCP invocation permission = %q, want mcp:execute", got)