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
This commit is contained in:
swordandart 2026-03-11 14:36:49 +08:00
parent 270814c3c9
commit 608574a701
10 changed files with 431 additions and 11 deletions

View file

@ -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 ""
}

View file

@ -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(),
)

View file

@ -40,6 +40,7 @@ func TestNewPicoclawCommand(t *testing.T) {
"migrate",
"onboard",
"skills",
"stats",
"status",
"version",
}

View file

@ -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) {

View file

@ -14,5 +14,6 @@ func BuiltinDefinitions() []Definition {
checkCommand(),
stopCommand(),
clearCommand(),
statsCommand(),
}
}

39
pkg/commands/cmd_stats.go Normal file
View file

@ -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))
}

View file

@ -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
}

204
pkg/llmlog/stats.go Normal file
View file

@ -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 (AllUntil %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()
}

View file

@ -18,7 +18,9 @@ type CallRecord struct {
Messages []Message `json:"messages"`
Response string `json:"response"`
Duration time.Duration `json:"duration"`
TokensUsed int `json:"tokens_used,omitempty"`
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"`
}

View file

@ -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
}
}