PNArchives/Memoh
1
0
Fork 0
mirror of https://github.com/memohai/Memoh.git synced 2026-04-27 07:16:19 +09:00
Code Issues Packages Projects Releases Wiki Activity

feat: token usage state (#153)

Browse Source 
* feat: token usage state

* fix: typo
This commit is contained in:
Acbox Liu
2026-03-01 02:19:07 +08:00
committed by GitHub
parent 443ede30b4
commit 0cdf822603
33 changed files with 1916 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/heartbeat/service.go
+7 -5
View File
@@ -114,7 +114,7 @@ func (s *Service) runHeartbeat(ctx context.Context, cfg Config) {
token, err := s.generateTriggerToken(cfg.OwnerUserID)
if err != nil {
s.completeLog(ctx, logRow.ID, "error", "", err.Error(), nil)
s.completeLog(ctx, logRow.ID, "error", "", err.Error(), nil, pgtype.UUID{})
s.logger.Error("generate trigger token failed", slog.String("bot_id", cfg.BotID), slog.Any("error", err))
return
}
@@ -125,22 +125,24 @@ func (s *Service) runHeartbeat(ctx context.Context, cfg Config) {
OwnerUserID: cfg.OwnerUserID,
}, token)
if err != nil {
s.completeLog(ctx, logRow.ID, "error", "", err.Error(), nil)
s.completeLog(ctx, logRow.ID, "error", "", err.Error(), nil, pgtype.UUID{})
s.logger.Error("heartbeat trigger failed", slog.String("bot_id", cfg.BotID), slog.Any("error", err))
return
}
s.completeLog(ctx, logRow.ID, result.Status, result.Text, "", result.UsageBytes)
modelID := db.ParseUUIDOrEmpty(result.ModelID)
s.completeLog(ctx, logRow.ID, result.Status, result.Text, "", result.UsageBytes, modelID)
s.logger.Info("heartbeat completed", slog.String("bot_id", cfg.BotID), slog.String("status", result.Status))
}
func (s *Service) completeLog(ctx context.Context, logID pgtype.UUID, status, resultText, errorMessage string, usageBytes []byte) {
func (s *Service) completeLog(ctx context.Context, logID pgtype.UUID, status, resultText, errorMessage string, usageBytes []byte, modelID pgtype.UUID) {
_, err := s.queries.CompleteHeartbeatLog(ctx, sqlc.CompleteHeartbeatLogParams{
ID: logID,
Status: status,
ResultText: resultText,
ErrorMessage: errorMessage,
Usage: usageBytes,
ModelID: modelID,
})
if err != nil {
s.logger.Error("complete heartbeat log failed", slog.Any("error", err))
@@ -222,7 +224,7 @@ func (s *Service) removeJob(botID string) {
}
}
func toLog(row sqlc.BotHeartbeatLog) Log {
func toLog(row sqlc.ListHeartbeatLogsByBotRow) Log {
l := Log{
ID: row.ID.String(),
BotID: row.BotID.String(),
internal/heartbeat/trigger.go
+1
View File
@@ -13,6 +13,7 @@ type TriggerResult struct {
Text string
Usage any
UsageBytes []byte
ModelID string
}
type Triggerer interface {