diff --git a/heicode/controller/agent_template_handlers.go b/heicode/controller/agent_template_handlers.go index 8972fc7f..8d926097 100644 --- a/heicode/controller/agent_template_handlers.go +++ b/heicode/controller/agent_template_handlers.go @@ -444,6 +444,33 @@ func HeicodeGetAgentStatus(c *gin.Context) { }) } +// HeicodeGetAgentUsage: GET /api/heicode/agents/:deployment_id/usage +// Per-agent model-usage rollup (#9): the agent calls HM /v1/* with its minted +// token named "agent:", so its consumption is the sum of consume +// logs under that token name. Optional ?start=&end= unix-second window. Returns +// raw quota + quota_per_unit (caller converts, same contract as /api/heicode/self). +func HeicodeGetAgentUsage(c *gin.Context) { + row, ok := findUserTemplateAgent(c) + if !ok { + return + } + start, _ := strconv.ParseInt(strings.TrimSpace(c.Query("start")), 10, 64) + end, _ := strconv.ParseInt(strings.TrimSpace(c.Query("end")), 10, 64) + usage, err := model.SumAgentUsage(c.GetInt("id"), "agent:"+row.DeploymentID, start, end) + if err != nil { + agentError(c, "DEPLOYMENT_CONFLICT", "failed to aggregate agent usage") + return + } + common.ApiSuccess(c, gin.H{ + "agent_id": row.DeploymentID, + "quota": usage.Quota, + "prompt_tokens": usage.PromptTokens, + "completion_tokens": usage.CompletionTokens, + "call_count": usage.CallCount, + "quota_per_unit": common.QuotaPerUnit, + }) +} + // HeicodeStopAgent: POST /api/heicode/agents/:deployment_id/stop func HeicodeStopAgent(c *gin.Context) { row, ok := findUserTemplateAgent(c) diff --git a/heicode/model/agent_usage_test.go b/heicode/model/agent_usage_test.go new file mode 100644 index 00000000..af0c1165 --- /dev/null +++ b/heicode/model/agent_usage_test.go @@ -0,0 +1,41 @@ +package model + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// SumAgentUsage rolls up a single agent token's consume logs (issue #9). logs is +// migrated by the package TestMain; LOG_DB == DB in tests. +func TestSumAgentUsage(t *testing.T) { + tok := "agent:dep_usage_test" + mk := func(uid int, tname string, typ, quota, pt, ct int) { + require.NoError(t, LOG_DB.Create(&Log{ + UserId: uid, TokenName: tname, Type: typ, + Quota: quota, PromptTokens: pt, CompletionTokens: ct, CreatedAt: 1700000000, + }).Error) + } + mk(970001, tok, LogTypeConsume, 100, 10, 5) + mk(970001, tok, LogTypeConsume, 200, 20, 15) + mk(970001, "agent:other", LogTypeConsume, 999, 99, 99) // different token — excluded + mk(970001, tok, LogTypeManage, 500, 0, 0) // non-consume — excluded + mk(970002, tok, LogTypeConsume, 777, 7, 7) // different user — excluded by user filter + + u, err := SumAgentUsage(970001, tok, 0, 0) + require.NoError(t, err) + require.EqualValues(t, 300, u.Quota) + require.EqualValues(t, 30, u.PromptTokens) + require.EqualValues(t, 20, u.CompletionTokens) + require.EqualValues(t, 2, u.CallCount) + + // time window excludes the rows (all stamped at 1700000000) + windowed, err := SumAgentUsage(970001, tok, 1800000000, 1900000000) + require.NoError(t, err) + require.EqualValues(t, 0, windowed.CallCount) + + // empty token name -> zero, no error + z, err := SumAgentUsage(970001, "", 0, 0) + require.NoError(t, err) + require.EqualValues(t, 0, z.CallCount) +} diff --git a/heicode/model/log.go b/heicode/model/log.go index 32693eea..da47471c 100644 --- a/heicode/model/log.go +++ b/heicode/model/log.go @@ -432,6 +432,41 @@ type Stat struct { Tpm int `json:"tpm"` } +// AgentUsage is the per-deployed-agent usage rollup (issue #9): a template agent +// calls HM /v1/* with its own minted token named "agent:", so its +// consumption is the sum of consume logs under that token name. +type AgentUsage struct { + Quota int64 `json:"quota"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + CallCount int64 `json:"call_count"` +} + +// SumAgentUsage aggregates consume-log usage for a single agent token (and user, +// when userId>0) over an optional time window. COALESCE keeps SUM non-null on +// empty sets; COUNT/SUM are portable across SQLite/MySQL/PostgreSQL. +func SumAgentUsage(userId int, tokenName string, startTimestamp int64, endTimestamp int64) (AgentUsage, error) { + var u AgentUsage + if tokenName == "" { + return u, nil + } + tx := LOG_DB.Table("logs"). + Select("COALESCE(SUM(quota),0) AS quota, COALESCE(SUM(prompt_tokens),0) AS prompt_tokens, COALESCE(SUM(completion_tokens),0) AS completion_tokens, COUNT(*) AS call_count"). + Where("type = ?", LogTypeConsume). + Where("token_name = ?", tokenName) + if userId > 0 { + tx = tx.Where("user_id = ?", userId) + } + if startTimestamp != 0 { + tx = tx.Where("created_at >= ?", startTimestamp) + } + if endTimestamp != 0 { + tx = tx.Where("created_at <= ?", endTimestamp) + } + err := tx.Scan(&u).Error + return u, err +} + func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string) (stat Stat, err error) { tx := LOG_DB.Table("logs").Select("sum(quota) quota") diff --git a/heicode/router/api-router.go b/heicode/router/api-router.go index bab035f4..8ac88abb 100644 --- a/heicode/router/api-router.go +++ b/heicode/router/api-router.go @@ -535,6 +535,7 @@ func SetApiRouter(router *gin.Engine) { heicodeAgentRoute.GET("/agents", controller.HeicodeListAgents) heicodeAgentRoute.GET("/agents/:deployment_id", controller.HeicodeGetAgent) heicodeAgentRoute.GET("/agents/:deployment_id/status", controller.HeicodeGetAgentStatus) + heicodeAgentRoute.GET("/agents/:deployment_id/usage", controller.HeicodeGetAgentUsage) heicodeAgentRoute.POST("/agents/:deployment_id/stop", controller.HeicodeStopAgent) heicodeAgentRoute.DELETE("/agents/:deployment_id", controller.HeicodeDeleteAgent) }