From 608574a701cf32a39bcbd5c42c72ba6fc2be98e6 Mon Sep 17 00:00:00 2001 From: swordandart <36768075+swordandart@users.noreply.github.com> Date: Wed, 11 Mar 2026 14:36:49 +0800 Subject: [PATCH] feat: add /stats command for session statistics - Add /stats command to display session metrics (messages, tokens, cost) - Refactor stats calculation into separate module - Update logging provider to track token usage - Add stats display on session end --- cmd/picoclaw/internal/stats/command.go | 118 ++++++++++++++ cmd/picoclaw/main.go | 2 + cmd/picoclaw/main_test.go | 1 + pkg/agent/loop.go | 43 +++++- pkg/commands/builtin.go | 1 + pkg/commands/cmd_stats.go | 39 +++++ pkg/commands/runtime.go | 10 ++ pkg/llmlog/stats.go | 204 +++++++++++++++++++++++++ pkg/llmlog/types.go | 20 +-- pkg/providers/logging_provider.go | 4 +- 10 files changed, 431 insertions(+), 11 deletions(-) create mode 100644 cmd/picoclaw/internal/stats/command.go create mode 100644 pkg/commands/cmd_stats.go create mode 100644 pkg/llmlog/stats.go diff --git a/cmd/picoclaw/internal/stats/command.go b/cmd/picoclaw/internal/stats/command.go new file mode 100644 index 000000000..0c0d12b9a --- /dev/null +++ b/cmd/picoclaw/internal/stats/command.go @@ -0,0 +1,118 @@ +package stats + +import ( + "fmt" + "os" + "path/filepath" + + "github.com/spf13/cobra" + + "github.com/sipeed/picoclaw/cmd/picoclaw/internal" + "github.com/sipeed/picoclaw/pkg/llmlog" +) + +func NewStatsCommand() *cobra.Command { + var days int + var model string + var logDir string + + cmd := &cobra.Command{ + Use: "stats", + Aliases: []string{"stat"}, + Short: "Show token usage statistics", + Long: "Display token usage statistics from LLM call logs", + Example: ` picoclaw stats # Show all time statistics + picoclaw stats --days 7 # Show last 7 days statistics + picoclaw stats --days 30 # Show last 30 days statistics + picoclaw stats --model gpt-4o # Filter by model + picoclaw stats --log-dir /path/to/logs # Specify log directory`, + Run: func(cmd *cobra.Command, args []string) { + // Find log directory + resolvedLogDir := findLogDir(logDir) + if resolvedLogDir == "" { + fmt.Println("❌ 无法找到日志目录") + fmt.Println("提示: 使用 --log-dir 指定日志目录，或在配置文件中设置 tools.llm_call_log.log_dir") + os.Exit(1) + } + + // Calculate statistics + result, err := llmlog.CalculateStats(resolvedLogDir, days, model) + if err != nil { + fmt.Printf("❌ Stats calculation failed: %v\n", err) + os.Exit(1) + } + + // Print result + fmt.Println(llmlog.FormatStatsTable(result)) + }, + } + + cmd.Flags().IntVarP(&days, "days", "d", 0, "Number of days to include (0 means all time)") + cmd.Flags().StringVarP(&model, "model", "m", "", "Filter by specific model") + cmd.Flags().StringVarP(&logDir, "log-dir", "l", "", "Log directory path (overrides config)") + + return cmd +} + +// findLogDir finds the LLM log directory with the following priority: +// 1. Explicitly provided logDir parameter +// 2. Config file: tools.llm_call_log.log_dir +// 3. Environment variable: PICOCLAW_TOOLS_LLM_CALL_LOG_DIR +// 4. Default locations: ~/.picoclaw/logs/llmcall, ~/picoclaw/workspace/logs/llmcall +func findLogDir(logDir string) string { + // 1. Check explicit parameter + if logDir != "" { + if _, err := os.Stat(logDir); err == nil { + return logDir + } + // If explicitly provided but doesn't exist, still return it + // (error will be reported later) + return logDir + } + + // 2. Check config file and environment variable + cfg, err := internal.LoadConfig() + if err == nil && cfg.Tools.LLMCallLog.LogDir != "" { + if _, err := os.Stat(cfg.Tools.LLMCallLog.LogDir); err == nil { + return cfg.Tools.LLMCallLog.LogDir + } + } + + // 3. Check environment variable directly (in case config loading failed) + if envDir := os.Getenv("PICOCLAW_TOOLS_LLM_CALL_LOG_DIR"); envDir != "" { + if _, err := os.Stat(envDir); err == nil { + return envDir + } + } + + // 4. Check default locations + homeDir, err := os.UserHomeDir() + if err != nil { + return "" + } + + // Default: ~/.picoclaw/logs/llmcall + defaultDir := filepath.Join(homeDir, ".picoclaw", "logs", "llmcall") + if _, err := os.Stat(defaultDir); err == nil { + return defaultDir + } + + // Alternative: workspace logs + workspaceLogDir := filepath.Join(homeDir, "picoclaw", "workspace", "logs", "llmcall") + if _, err := os.Stat(workspaceLogDir); err == nil { + return workspaceLogDir + } + + // Legacy: check old directory names for backward compatibility + legacyDir := filepath.Join(homeDir, ".picoclaw", "logs", "llm") + if _, err := os.Stat(legacyDir); err == nil { + return legacyDir + } + + legacyWorkspaceDir := filepath.Join(homeDir, "picoclaw", "workspace", "logs", "llm") + if _, err := os.Stat(legacyWorkspaceDir); err == nil { + return legacyWorkspaceDir + } + + return "" +} \ No newline at end of file diff --git a/cmd/picoclaw/main.go b/cmd/picoclaw/main.go index fe4de8ecc..193b5345c 100644 --- a/cmd/picoclaw/main.go +++ b/cmd/picoclaw/main.go @@ -20,6 +20,7 @@ import ( "github.com/sipeed/picoclaw/cmd/picoclaw/internal/migrate" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/onboard" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/skills" + "github.com/sipeed/picoclaw/cmd/picoclaw/internal/stats" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/status" "github.com/sipeed/picoclaw/cmd/picoclaw/internal/version" ) @@ -42,6 +43,7 @@ func NewPicoclawCommand() *cobra.Command { cron.NewCronCommand(), migrate.NewMigrateCommand(), skills.NewSkillsCommand(), + stats.NewStatsCommand(), version.NewVersionCommand(), ) diff --git a/cmd/picoclaw/main_test.go b/cmd/picoclaw/main_test.go index 3740ba358..f73fc5475 100644 --- a/cmd/picoclaw/main_test.go +++ b/cmd/picoclaw/main_test.go @@ -40,6 +40,7 @@ func TestNewPicoclawCommand(t *testing.T) { "migrate", "onboard", "skills", + "stats", "status", "version", } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 09130da0b..c309453e9 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -52,6 +52,10 @@ type AgentLoop struct { // Task cancellation support currentCancel context.CancelFunc currentCancelMu sync.Mutex + + // Token statistics for current session + tokenStats *commands.TokenStats + tokenStatsMu sync.RWMutex } // processOptions configures how a message is processed @@ -432,6 +436,34 @@ func (al *AgentLoop) CancelCurrentTask() bool { return true } +// updateTokenStats updates the token statistics for the current session. +func (al *AgentLoop) updateTokenStats(model string, promptTokens, completionTokens int) { + al.tokenStatsMu.Lock() + defer al.tokenStatsMu.Unlock() + + if al.tokenStats == nil { + al.tokenStats = &commands.TokenStats{Model: model} + } + al.tokenStats.PromptTokens += promptTokens + al.tokenStats.CompletionTokens += completionTokens + al.tokenStats.TotalTokens += promptTokens + completionTokens + al.tokenStats.CallCount++ +} + +// getTokenStats returns the current session's token statistics. +func (al *AgentLoop) getTokenStats() *commands.TokenStats { + al.tokenStatsMu.RLock() + defer al.tokenStatsMu.RUnlock() + return al.tokenStats +} + +// resetTokenStats resets the token statistics for a new session. +func (al *AgentLoop) resetTokenStats() { + al.tokenStatsMu.Lock() + defer al.tokenStatsMu.Unlock() + al.tokenStats = nil +} + func (al *AgentLoop) RegisterTool(tool tools.Tool) { for _, agentID := range al.registry.ListAgentIDs() { if agent, ok := al.registry.GetAgent(agentID); ok { @@ -1165,6 +1197,12 @@ func (al *AgentLoop) runLLMIteration( "target_channel": al.targetReasoningChannelID(opts.Channel), "channel": opts.Channel, }) + + // Update token statistics for current session + if response.Usage != nil { + al.updateTokenStats(activeModel, response.Usage.PromptTokens, response.Usage.CompletionTokens) + } + // Check if no tool calls - then check reasoning content if any if len(response.ToolCalls) == 0 { finalContent = response.Content @@ -1871,7 +1909,9 @@ func (al *AgentLoop) handleCommandAsync(ctx context.Context, msg bus.InboundMess response, handled := al.handleCommand(ctx, msg, agent, opts) if !handled { - // Command not recognized or passed through, ignore + // Command not recognized or passed through (e.g., strict command with extra args) + // Process as normal message through the agent + al.processMessageAsync(ctx, msg) return } @@ -1955,6 +1995,7 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt return nil }, CancelCurrentTask: al.CancelCurrentTask, + GetTokenStats: al.getTokenStats, } if agent != nil { rt.GetModelInfo = func() (string, string) { diff --git a/pkg/commands/builtin.go b/pkg/commands/builtin.go index a28e3fe78..c86eafb5c 100644 --- a/pkg/commands/builtin.go +++ b/pkg/commands/builtin.go @@ -14,5 +14,6 @@ func BuiltinDefinitions() []Definition { checkCommand(), stopCommand(), clearCommand(), + statsCommand(), } } diff --git a/pkg/commands/cmd_stats.go b/pkg/commands/cmd_stats.go new file mode 100644 index 000000000..d79c77410 --- /dev/null +++ b/pkg/commands/cmd_stats.go @@ -0,0 +1,39 @@ +package commands + +import ( + "context" + "fmt" +) + +// statsCommand returns the definition for the /stats command +func statsCommand() Definition { + return Definition{ + Name: "stats", + Description: "Show token usage statistics for current session", + Handler: handleStats, + Strict: true, + } +} + +func handleStats(ctx context.Context, req Request, rt *Runtime) error { + if rt.GetTokenStats == nil { + return req.Reply("Token stats not available") + } + + stats := rt.GetTokenStats() + if stats == nil { + return req.Reply("No token usage data for current session") + } + + return req.Reply(fmt.Sprintf(`Token Stats (Current Session) +Model: %s +Prompt: %d +Completion: %d +Total: %d +Calls: %d`, + stats.Model, + stats.PromptTokens, + stats.CompletionTokens, + stats.TotalTokens, + stats.CallCount)) +} \ No newline at end of file diff --git a/pkg/commands/runtime.go b/pkg/commands/runtime.go index c07d46d87..a7ef1f840 100644 --- a/pkg/commands/runtime.go +++ b/pkg/commands/runtime.go @@ -2,6 +2,15 @@ package commands import "github.com/sipeed/picoclaw/pkg/config" +// TokenStats represents token usage statistics for the current session +type TokenStats struct { + Model string + PromptTokens int + CompletionTokens int + TotalTokens int + CallCount int +} + // Runtime provides runtime dependencies to command handlers. It is constructed // per-request by the agent loop so that per-request state (like session scope) // can coexist with long-lived callbacks (like GetModelInfo). @@ -15,4 +24,5 @@ type Runtime struct { SwitchChannel func(value string) error CancelCurrentTask func() bool // Cancel the currently running task, returns true if a task was cancelled ClearHistory func() error + GetTokenStats func() *TokenStats // Get token usage statistics for current session } diff --git a/pkg/llmlog/stats.go b/pkg/llmlog/stats.go new file mode 100644 index 000000000..7d0543f36 --- /dev/null +++ b/pkg/llmlog/stats.go @@ -0,0 +1,204 @@ +package llmlog + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "time" +) + +// TokenStats represents token usage statistics for a model +type TokenStats struct { + Model string + PromptTokens int + CompletionTokens int + TotalTokens int + CallCount int + ErrorCount int +} + +// StatsResult represents the result of a stats query +type StatsResult struct { + ByModel map[string]*TokenStats + Total *TokenStats + StartDay string + EndDay string +} + +// CalculateStats calculates token usage statistics from log files +// If days is 0, it calculates all available logs +// If modelFilter is empty, it includes all models +// Supports filename formats: "2006-01-02.jsonl" and "llmcall_2006-01-02.jsonl" +func CalculateStats(logDir string, days int, modelFilter string) (*StatsResult, error) { + result := &StatsResult{ + ByModel: make(map[string]*TokenStats), + Total: &TokenStats{ + Model: "Total", + }, + } + + now := time.Now() + var startDate time.Time + if days > 0 { + startDate = now.AddDate(0, 0, -days) + result.StartDay = startDate.Format("2006-01-02") + } else { + startDate = time.Time{} // zero time means all + } + result.EndDay = now.Format("2006-01-02") + + // Find all log files + files, err := filepath.Glob(filepath.Join(logDir, "*.jsonl")) + if err != nil { + return nil, fmt.Errorf("failed to list log files: %w", err) + } + + for _, file := range files { + // Parse date from filename + // Supported formats: "2006-01-02.jsonl" and "llmcall_2006-01-02.jsonl" + filename := filepath.Base(file) + dateStr := strings.TrimSuffix(filename, ".jsonl") + + // Handle "llmcall_" prefix + if strings.HasPrefix(dateStr, "llmcall_") { + dateStr = strings.TrimPrefix(dateStr, "llmcall_") + } + + fileDate, err := time.Parse("2006-01-02", dateStr) + if err != nil { + continue // skip files with invalid names + } + + // Check if file is within date range + if !startDate.IsZero() && fileDate.Before(startDate) { + continue + } + + // Process file + if err := processLogFile(file, modelFilter, result); err != nil { + // Log error but continue processing other files + continue + } + } + + // Sort models by total tokens for consistent output + return result, nil +} + +func processLogFile(filePath string, modelFilter string, result *StatsResult) error { + file, err := os.Open(filePath) + if err != nil { + return err + } + defer file.Close() + + decoder := json.NewDecoder(file) + for decoder.More() { + var record CallRecord + if err := decoder.Decode(&record); err != nil { + continue // skip malformed records + } + + // Filter by model if specified + if modelFilter != "" && record.Model != modelFilter { + continue + } + + // Get or create stats for this model + stats, exists := result.ByModel[record.Model] + if !exists { + stats = &TokenStats{Model: record.Model} + result.ByModel[record.Model] = stats + } + + // Update stats + stats.PromptTokens += record.PromptTokens + stats.CompletionTokens += record.CompletionTokens + stats.TotalTokens += record.TotalTokens + stats.CallCount++ + if record.Error != "" { + stats.ErrorCount++ + } + + // Update totals + result.Total.PromptTokens += record.PromptTokens + result.Total.CompletionTokens += record.CompletionTokens + result.Total.TotalTokens += record.TotalTokens + result.Total.CallCount++ + if record.Error != "" { + result.Total.ErrorCount++ + } + } + + return nil +} + +// FormatStatsTable formats the stats result as a table string +func FormatStatsTable(result *StatsResult) string { + if len(result.ByModel) == 0 { + return "📊 no stats data found" + } + + var sb strings.Builder + + if result.StartDay != "" { + sb.WriteString(fmt.Sprintf("📊 Token Count (%s ~ %s)\n\n", result.StartDay, result.EndDay)) + } else { + sb.WriteString(fmt.Sprintf("📊 Token Count (All，Until %s)\n\n", result.EndDay)) + } + + // Sort models by total tokens descending + models := make([]string, 0, len(result.ByModel)) + for model := range result.ByModel { + models = append(models, model) + } + sort.Slice(models, func(i, j int) bool { + return result.ByModel[models[i]].TotalTokens > result.ByModel[models[j]].TotalTokens + }) + + // Table header + sb.WriteString("┌──────────────────────┬──────────────┬─────────────────┬──────────────┬──────────┐\n") + sb.WriteString("│ Model │ Prompt Tokens│Completion Tokens│ Total Tokens │Call Count│\n") + sb.WriteString("├──────────────────────┼──────────────┼─────────────────┼──────────────┼──────────┤\n") + + // Table rows + for _, model := range models { + stats := result.ByModel[model] + // Truncate model name if too long + displayModel := model + if len(displayModel) > 20 { + displayModel = displayModel[:17] + "..." + } + sb.WriteString(fmt.Sprintf("│ %-20s │ %12d │ %15d │ %12d │ %8d │\n", + displayModel, stats.PromptTokens, stats.CompletionTokens, stats.TotalTokens, stats.CallCount)) + } + + // Total row + sb.WriteString("├──────────────────────┼──────────────┼─────────────────┼──────────────┼──────────┤\n") + sb.WriteString(fmt.Sprintf("│ %-20s │ %12d │ %15d │ %12d │ %8d │\n", + "Total", result.Total.PromptTokens, result.Total.CompletionTokens, result.Total.TotalTokens, result.Total.CallCount)) + sb.WriteString("└──────────────────────┴──────────────┴─────────────────┴──────────────┴──────────┘\n") + + return sb.String() +} + +// FormatSimpleStats formats stats for a single model (used in runtime /stats command) +func FormatSimpleStats(stats *TokenStats) string { + var sb strings.Builder + + sb.WriteString(fmt.Sprintf("📊 Token Stats (Current Session)\n")) + sb.WriteString(fmt.Sprintf("Model: %s\n", stats.Model)) + sb.WriteString("┌─────────────────────┬───────────┐\n") + sb.WriteString("│ Metric │ Value │\n") + sb.WriteString("├─────────────────────┼───────────┤\n") + sb.WriteString(fmt.Sprintf("│ %-19s │ %9d │\n", "Prompt Tokens", stats.PromptTokens)) + sb.WriteString(fmt.Sprintf("│ %-19s │ %9d │\n", "Completion Tokens", stats.CompletionTokens)) + sb.WriteString(fmt.Sprintf("│ %-19s │ %9d │\n", "Total Tokens", stats.TotalTokens)) + sb.WriteString(fmt.Sprintf("│ %-19s │ %9d │\n", "Call Count", stats.CallCount)) + sb.WriteString("└─────────────────────┴───────────┘\n") + + return sb.String() +} \ No newline at end of file diff --git a/pkg/llmlog/types.go b/pkg/llmlog/types.go index 2cea5bf17..38bcf1584 100644 --- a/pkg/llmlog/types.go +++ b/pkg/llmlog/types.go @@ -12,15 +12,17 @@ type Message struct { // CallRecord 表示一次 LLM 调用记录 type CallRecord struct { - Timestamp time.Time `json:"timestamp"` - Model string `json:"model"` - Provider string `json:"provider"` - Messages []Message `json:"messages"` - Response string `json:"response"` - Duration time.Duration `json:"duration"` - TokensUsed int `json:"tokens_used,omitempty"` - Error string `json:"error,omitempty"` - IsStreaming bool `json:"is_streaming"` + Timestamp time.Time `json:"timestamp"` + Model string `json:"model"` + Provider string `json:"provider"` + Messages []Message `json:"messages"` + Response string `json:"response"` + Duration time.Duration `json:"duration"` + PromptTokens int `json:"prompt_tokens,omitempty"` + CompletionTokens int `json:"completion_tokens,omitempty"` + TotalTokens int `json:"total_tokens,omitempty"` + Error string `json:"error,omitempty"` + IsStreaming bool `json:"is_streaming"` } // Logger 定义 LLM 调用日志记录器接口 diff --git a/pkg/providers/logging_provider.go b/pkg/providers/logging_provider.go index 02e8a07d2..6cb991b5c 100644 --- a/pkg/providers/logging_provider.go +++ b/pkg/providers/logging_provider.go @@ -66,7 +66,9 @@ func (p *LoggingProvider) Chat( } else { record.Response = resp.Content if resp.Usage != nil { - record.TokensUsed = resp.Usage.TotalTokens + record.PromptTokens = resp.Usage.PromptTokens + record.CompletionTokens = resp.Usage.CompletionTokens + record.TotalTokens = resp.Usage.TotalTokens } }