mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-28 21:50:43 +02:00
Add token usage tracking and UI refinements
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user