Version: 0.6.6.dev.260317

 feat(task,agent): 新增任务完成接口,并打通聊天全链路 Token 记账

-  新增“标记任务为完成”接口,并补充幂等保护,避免重复完成导致状态污染
- 📊 为聊天链路补充 Token 统计能力:
  - 流式主对话链路直接读取模型 `usage`
  - Agent 链路通过 `Eino callback + ctx` 聚合 `Generate usage`
  - 在流式场景下补齐缺失的 `usage` 数据
- 🧾 按口径 B 完成 Token 落库:
  - 用户消息 `token` 记为 `0`
  - 助手消息记录本轮总 `token`
  - 持久化时同步更新 `chat_histories.tokens_consumed`、`agent_chats.tokens_total`、`users.token_usage`
- 🔄 异步标题生成产生的 Token 通过 Outbox 事件完成账本增量调整,保证统计口径一致
- 📝 同步更新 `AGENTS.md` 与 `.gitignore`
- 📚 小幅更新 README 说明文档
This commit is contained in:
LoveLosita
2026-03-17 18:23:07 +08:00
parent 09dca9f772
commit 96be3e2a02
19 changed files with 660 additions and 36 deletions

View File

@@ -32,6 +32,11 @@ type AgentService struct {
// 这里通过依赖注入把“模型、仓储、缓存、异步持久化通道”统一交给服务层管理,
// 便于后续在单测中替换实现,或在启动流程中按环境切换配置。
func NewAgentService(aiHub *inits.AIHub, repo *dao.AgentDAO, taskRepo *dao.TaskDAO, agentRedis *dao.AgentCache, eventPublisher outboxinfra.EventPublisher) *AgentService {
// 全局注册一次 token 采集 callback
// 1. 只注册一次,避免重复处理;
// 2. 只有带 RequestTokenMeter 的请求上下文才会真正累加。
ensureTokenMeterCallbackRegistered()
return &AgentService{
AIHub: aiHub,
repo: repo,
@@ -76,7 +81,7 @@ func (s *AgentService) PersistChatHistory(ctx context.Context, payload model.Cha
// 1. 未注入事件发布器时(例如本地极简环境),直接同步写 DB。
// 这样可以保证功能不依赖 Kafka 也能跑通。
if s.eventPublisher == nil {
return s.repo.SaveChatHistory(ctx, payload.UserID, payload.ConversationID, payload.Role, payload.Message)
return s.repo.SaveChatHistory(ctx, payload.UserID, payload.ConversationID, payload.Role, payload.Message, payload.TokensConsumed)
}
// 2. 已启用异步总线时,只发布“持久化请求事件”,不在请求路径阻塞 Kafka。
// 2.1 发布成功仅代表“事件安全入队”,实际落库由消费者异步完成。
@@ -167,12 +172,23 @@ func (s *AgentService) runNormalChatFlow(
// 6. 执行真正的流式聊天。
// fullText 用于后续写 Redis/持久化outChan 用于把流片段实时推给前端。
fullText, streamErr := chat.StreamChat(ctx, selectedModel, resolvedModelName, userMessage, ifThinking, chatHistory, outChan, traceID, chatID, requestStart)
fullText, streamUsage, streamErr := chat.StreamChat(ctx, selectedModel, resolvedModelName, userMessage, ifThinking, chatHistory, outChan, traceID, chatID, requestStart)
if streamErr != nil {
pushErrNonBlocking(errChan, streamErr)
return
}
// 6.1 流式 usage 并入请求级 token 统计器:
// 6.1.1 route/quicknote/taskquery 等 Generate 调用由 callback 自动累加;
// 6.1.2 主对话 Stream usage 在这里手动补齐。
addSchemaUsageIntoRequest(ctx, streamUsage)
requestTokenSnapshot := snapshotRequestTokenMeter(ctx)
requestTotalTokens := requestTokenSnapshot.TotalTokens
if requestTotalTokens <= 0 && streamUsage != nil {
// 兜底:若 callback/meter 未生效,至少使用流式 usage 保底记账。
requestTotalTokens = normalizeUsageTotal(streamUsage.TotalTokens, streamUsage.PromptTokens, streamUsage.CompletionTokens)
}
// 7. 后置持久化(用户消息):
// 7.1 先写 Redis保证“最新会话上下文”可立即用于下一轮推理
// 7.2 再走可靠持久化入口outbox 或同步 DB
@@ -185,6 +201,8 @@ func (s *AgentService) runNormalChatFlow(
ConversationID: chatID,
Role: "user",
Message: userMessage,
// 口径B用户消息固定记 0本轮总 token 统一记在助手消息。
TokensConsumed: 0,
}); err != nil {
pushErrNonBlocking(errChan, err)
return
@@ -204,6 +222,8 @@ func (s *AgentService) runNormalChatFlow(
ConversationID: chatID,
Role: "assistant",
Message: fullText,
// 口径B助手消息记录“本轮请求总 token”。
TokensConsumed: requestTotalTokens,
}); saveErr != nil {
pushErrNonBlocking(errChan, saveErr)
}
@@ -223,13 +243,16 @@ func (s *AgentService) AgentChat(ctx context.Context, userMessage string, ifThin
outChan := make(chan string, 8)
errChan := make(chan error, 1)
// 0. 初始化“请求级 token 统计器”,用于聚合本次请求所有模型开销。
requestCtx, _ := withRequestTokenMeter(ctx)
// 1) 规范会话 ID选择模型。
chatID = normalizeConversationID(chatID)
selectedModel, resolvedModelName := s.pickChatModel(modelName)
// 2) 确保会话存在(优先缓存,必要时回源 DB 并创建)。
// 2.1 先查 Redis 会话标记,命中则可跳过 DB 存在性校验。
result, err := s.agentCache.GetConversationStatus(ctx, chatID)
result, err := s.agentCache.GetConversationStatus(requestCtx, chatID)
if err != nil {
errChan <- err
close(outChan)
@@ -238,7 +261,7 @@ func (s *AgentService) AgentChat(ctx context.Context, userMessage string, ifThin
}
if !result {
// 2.2 缓存未命中时回源 DB确认会话是否存在。
innerResult, ifErr := s.repo.IfChatExists(ctx, userID, chatID)
innerResult, ifErr := s.repo.IfChatExists(requestCtx, userID, chatID)
if ifErr != nil {
errChan <- ifErr
close(outChan)
@@ -255,7 +278,7 @@ func (s *AgentService) AgentChat(ctx context.Context, userMessage string, ifThin
}
}
// 2.4 补写 Redis 会话标记,优化下次访问。
if err = s.agentCache.SetConversationStatus(ctx, chatID); err != nil {
if err = s.agentCache.SetConversationStatus(requestCtx, chatID); err != nil {
log.Printf("设置会话状态缓存失败 chat=%s: %v", chatID, err)
}
}
@@ -269,11 +292,11 @@ func (s *AgentService) AgentChat(ctx context.Context, userMessage string, ifThin
defer close(outChan)
// 3.1 先走轻量路由,拿到统一 action。
routing := s.decideActionRouting(ctx, selectedModel, userMessage)
routing := s.decideActionRouting(requestCtx, selectedModel, userMessage)
// 3.2 chat直接走普通聊天主链路。
if routing.Action == route.ActionChat {
s.runNormalChatFlow(ctx, selectedModel, resolvedModelName, userMessage, ifThinking, userID, chatID, traceID, requestStart, outChan, errChan)
s.runNormalChatFlow(requestCtx, selectedModel, resolvedModelName, userMessage, ifThinking, userID, chatID, traceID, requestStart, outChan, errChan)
return
}
@@ -284,7 +307,7 @@ func (s *AgentService) AgentChat(ctx context.Context, userMessage string, ifThin
// 3.4 quick_note_create执行随口记 graph。
if routing.Action == route.ActionQuickNoteCreate {
quickHandled, quickState, quickErr := s.tryHandleQuickNoteWithGraph(
ctx,
requestCtx,
selectedModel,
userMessage,
userID,
@@ -301,14 +324,15 @@ func (s *AgentService) AgentChat(ctx context.Context, userMessage string, ifThin
if quickHandled {
// 3.4.1 随口记处理成功:组织最终回复并按 OpenAI 兼容格式输出。
progress.Emit("quick_note.reply.polishing", "正在结合你的话题润色回复。")
quickReply := buildQuickNoteFinalReply(ctx, selectedModel, userMessage, quickState)
quickReply := buildQuickNoteFinalReply(requestCtx, selectedModel, userMessage, quickState)
if emitErr := emitSingleAssistantCompletion(outChan, resolvedModelName, quickReply); emitErr != nil {
pushErrNonBlocking(errChan, emitErr)
return
}
// 3.4.2 对随口记回复执行统一后置持久化Redis + outbox/DB
s.persistChatAfterReply(ctx, userID, chatID, userMessage, quickReply, errChan)
requestTotalTokens := snapshotRequestTokenMeter(requestCtx).TotalTokens
s.persistChatAfterReply(requestCtx, userID, chatID, userMessage, quickReply, 0, requestTotalTokens, errChan)
// 3.4.3 随口记链路同样异步生成会话标题(仅首次写入)。
s.ensureConversationTitleAsync(userID, chatID)
return
@@ -316,18 +340,18 @@ func (s *AgentService) AgentChat(ctx context.Context, userMessage string, ifThin
// 3.4.4 路由误判或 graph 判定非随口记时,回落普通聊天,保证“能聊”。
progress.Emit("quick_note.fallback", "当前输入不是随口记请求,切换到普通对话。")
s.runNormalChatFlow(ctx, selectedModel, resolvedModelName, userMessage, ifThinking, userID, chatID, traceID, requestStart, outChan, errChan)
s.runNormalChatFlow(requestCtx, selectedModel, resolvedModelName, userMessage, ifThinking, userID, chatID, traceID, requestStart, outChan, errChan)
return
}
// 3.5 task_query执行任务查询 tool-calling。
if routing.Action == route.ActionTaskQuery {
reply, queryErr := s.runTaskQueryFlow(ctx, selectedModel, userMessage, userID, progress.Emit)
reply, queryErr := s.runTaskQueryFlow(requestCtx, selectedModel, userMessage, userID, progress.Emit)
if queryErr != nil {
// 3.5.1 任务查询失败时回退普通聊天,避免请求直接中断。
log.Printf("任务查询 tool-calling 执行失败,回退普通聊天 trace_id=%s chat_id=%s err=%v", traceID, chatID, queryErr)
progress.Emit("task_query.fallback", "任务查询暂不可用,先切回普通对话。")
s.runNormalChatFlow(ctx, selectedModel, resolvedModelName, userMessage, ifThinking, userID, chatID, traceID, requestStart, outChan, errChan)
s.runNormalChatFlow(requestCtx, selectedModel, resolvedModelName, userMessage, ifThinking, userID, chatID, traceID, requestStart, outChan, errChan)
return
}
@@ -336,13 +360,14 @@ func (s *AgentService) AgentChat(ctx context.Context, userMessage string, ifThin
pushErrNonBlocking(errChan, emitErr)
return
}
s.persistChatAfterReply(ctx, userID, chatID, userMessage, reply, errChan)
requestTotalTokens := snapshotRequestTokenMeter(requestCtx).TotalTokens
s.persistChatAfterReply(requestCtx, userID, chatID, userMessage, reply, 0, requestTotalTokens, errChan)
s.ensureConversationTitleAsync(userID, chatID)
return
}
// 3.6 未知 action 兜底:走普通聊天,保证可用性。
s.runNormalChatFlow(ctx, selectedModel, resolvedModelName, userMessage, ifThinking, userID, chatID, traceID, requestStart, outChan, errChan)
s.runNormalChatFlow(requestCtx, selectedModel, resolvedModelName, userMessage, ifThinking, userID, chatID, traceID, requestStart, outChan, errChan)
}()
return outChan, errChan

View File

@@ -10,6 +10,7 @@ import (
"github.com/LoveLosita/smartflow/backend/model"
"github.com/LoveLosita/smartflow/backend/respond"
eventsvc "github.com/LoveLosita/smartflow/backend/service/events"
"github.com/cloudwego/eino-ext/components/model/ark"
einoModel "github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/schema"
@@ -32,6 +33,9 @@ const (
conversationListDefaultPageSize = 20
// conversationListMaxPageSize 是会话列表单页上限,避免超大分页压垮数据库。
conversationListMaxPageSize = 100
// conversationTitleTokenAdjustReason 是“标题异步生成 token 账本调整”原因码。
// 用于日志和后续审计归因。
conversationTitleTokenAdjustReason = "conversation_title_async"
)
const conversationTitlePrompt = `你是 SmartFlow 的会话标题生成器。
@@ -190,7 +194,7 @@ func (s *AgentService) ensureConversationTitleAsync(userID int, chatID string) {
}
// 4. 调用模型生成标题,并做格式清洗。
generated, err := s.generateConversationTitle(ctx, history)
generated, titleTokens, err := s.generateConversationTitle(ctx, history)
if err != nil {
log.Printf("异步生成会话标题失败(模型生成失败) chat=%s err=%v", chatID, err)
return
@@ -199,6 +203,28 @@ func (s *AgentService) ensureConversationTitleAsync(userID int, chatID string) {
return
}
// 4.1 标题生成成功后,把本次异步模型 token 记账:
// 4.1.1 启用 outbox 时走 adjust 事件,异步可靠入账;
// 4.1.2 未启用 outbox 时走同步兜底,直接更新账本。
if titleTokens > 0 {
if s.eventPublisher != nil {
publishErr := eventsvc.PublishChatTokenUsageAdjustRequested(ctx, s.eventPublisher, model.ChatTokenUsageAdjustPayload{
UserID: userID,
ConversationID: chatID,
TokensDelta: titleTokens,
Reason: conversationTitleTokenAdjustReason,
TriggeredAt: time.Now(),
})
if publishErr != nil {
log.Printf("异步标题 token 记账事件发布失败 chat=%s tokens=%d err=%v", chatID, titleTokens, publishErr)
}
} else {
if adjustErr := s.repo.AdjustTokenUsage(ctx, userID, chatID, titleTokens); adjustErr != nil {
log.Printf("异步标题 token 同步记账失败 chat=%s tokens=%d err=%v", chatID, titleTokens, adjustErr)
}
}
}
// 5. 只在标题仍为空时写入,保证并发幂等。
if err = s.repo.UpdateConversationTitleIfEmpty(ctx, userID, chatID, generated); err != nil {
log.Printf("异步生成会话标题失败(写库失败) chat=%s err=%v", chatID, err)
@@ -207,17 +233,17 @@ func (s *AgentService) ensureConversationTitleAsync(userID int, chatID string) {
}
// generateConversationTitle 使用聊天模型从近期历史生成标题。
func (s *AgentService) generateConversationTitle(ctx context.Context, history []*schema.Message) (string, error) {
func (s *AgentService) generateConversationTitle(ctx context.Context, history []*schema.Message) (string, int, error) {
modelInst := s.pickTitleModel()
if modelInst == nil {
return "", fmt.Errorf("标题生成模型未初始化")
return "", 0, fmt.Errorf("标题生成模型未初始化")
}
// 1. 只取最近 N 条,降低 token 并聚焦当前会话主题。
trimmed := tailMessages(history, conversationTitleHistoryLimit)
prompt := buildConversationTitleUserPrompt(trimmed)
if strings.TrimSpace(prompt) == "" {
return "", fmt.Errorf("缺少可用历史内容")
return "", 0, fmt.Errorf("缺少可用历史内容")
}
messages := []*schema.Message{
@@ -232,12 +258,22 @@ func (s *AgentService) generateConversationTitle(ctx context.Context, history []
einoModel.WithMaxTokens(40),
)
if err != nil {
return "", err
return "", 0, err
}
if resp == nil {
return "", fmt.Errorf("标题生成模型返回为空")
return "", 0, fmt.Errorf("标题生成模型返回为空")
}
return normalizeConversationTitle(resp.Content), nil
// 2.1 标题链路的 token 从模型响应 usage 中提取;缺失则按 0 处理,不影响主流程。
titleTokens := 0
if resp.ResponseMeta != nil && resp.ResponseMeta.Usage != nil {
titleTokens = normalizeUsageTotal(
resp.ResponseMeta.Usage.TotalTokens,
resp.ResponseMeta.Usage.PromptTokens,
resp.ResponseMeta.Usage.CompletionTokens,
)
}
return normalizeConversationTitle(resp.Content), titleTokens, nil
}
// pickTitleModel 选择用于标题生成的模型。

View File

@@ -345,6 +345,8 @@ func (s *AgentService) persistChatAfterReply(
chatID string,
userMessage string,
assistantReply string,
userTokens int,
assistantTokens int,
errChan chan error,
) {
// 1. 先把用户消息写入 Redis保证会话上下文“马上可见”。
@@ -358,6 +360,7 @@ func (s *AgentService) persistChatAfterReply(
ConversationID: chatID,
Role: "user",
Message: userMessage,
TokensConsumed: userTokens,
}); err != nil {
pushErrNonBlocking(errChan, err)
return
@@ -374,6 +377,7 @@ func (s *AgentService) persistChatAfterReply(
ConversationID: chatID,
Role: "assistant",
Message: assistantReply,
TokensConsumed: assistantTokens,
}); err != nil {
pushErrNonBlocking(errChan, err)
}

View File

@@ -0,0 +1,145 @@
package agentsvc
import (
"context"
"sync"
einoCallbacks "github.com/cloudwego/eino/callbacks"
einoModel "github.com/cloudwego/eino/components/model"
"github.com/cloudwego/eino/schema"
templatecb "github.com/cloudwego/eino/utils/callbacks"
)
type requestTokenMeterCtxKey struct{}
// RequestTokenMeter 是“单次请求级”的 token 统计容器。
//
// 设计目标:
// 1. 聚合本次请求内所有模型调用 token路由/图节点/流式主对话);
// 2. 线程安全,允许在同一请求内被多个链路节点并发累加;
// 3. 最终由服务层一次性读取快照并写入持久化。
type RequestTokenMeter struct {
mu sync.Mutex
promptTokens int
completionTokens int
totalTokens int
}
// RequestTokenMeterSnapshot 是 RequestTokenMeter 的只读快照。
type RequestTokenMeterSnapshot struct {
PromptTokens int
CompletionTokens int
TotalTokens int
}
var registerTokenMeterCallbackOnce sync.Once
// ensureTokenMeterCallbackRegistered 注册一次全局 ChatModel callback。
//
// 说明:
// 1. callback 只负责“采集并累加 token”不做业务决策
// 2. 仅当 ctx 里存在 RequestTokenMeter 时才会生效;
// 3. 采用 once避免在测试/多次构造服务时重复注册。
func ensureTokenMeterCallbackRegistered() {
registerTokenMeterCallbackOnce.Do(func() {
handler := templatecb.NewHandlerHelper().
ChatModel(&templatecb.ModelCallbackHandler{
OnEnd: func(ctx context.Context, _ *einoCallbacks.RunInfo, output *einoModel.CallbackOutput) context.Context {
if output == nil || output.TokenUsage == nil {
return ctx
}
addModelUsageIntoRequest(ctx, output.TokenUsage)
return ctx
},
}).
Handler()
einoCallbacks.AppendGlobalHandlers(handler)
})
}
// withRequestTokenMeter 创建并挂载“请求级 token 统计器”。
func withRequestTokenMeter(ctx context.Context) (context.Context, *RequestTokenMeter) {
meter := &RequestTokenMeter{}
return context.WithValue(ctx, requestTokenMeterCtxKey{}, meter), meter
}
// getRequestTokenMeter 读取请求上下文中的 token 统计器。
func getRequestTokenMeter(ctx context.Context) *RequestTokenMeter {
if ctx == nil {
return nil
}
meter, _ := ctx.Value(requestTokenMeterCtxKey{}).(*RequestTokenMeter)
return meter
}
// addSchemaUsageIntoRequest 把 schema usage 累加到请求级统计器。
func addSchemaUsageIntoRequest(ctx context.Context, usage *schema.TokenUsage) {
if usage == nil {
return
}
addTokenUsageValues(ctx, usage.PromptTokens, usage.CompletionTokens, normalizeUsageTotal(usage.TotalTokens, usage.PromptTokens, usage.CompletionTokens))
}
// addModelUsageIntoRequest 把 Eino model callback usage 累加到请求级统计器。
func addModelUsageIntoRequest(ctx context.Context, usage *einoModel.TokenUsage) {
if usage == nil {
return
}
addTokenUsageValues(ctx, usage.PromptTokens, usage.CompletionTokens, normalizeUsageTotal(usage.TotalTokens, usage.PromptTokens, usage.CompletionTokens))
}
// addTokenUsageValues 统一累加 token 数值。
func addTokenUsageValues(ctx context.Context, promptTokens, completionTokens, totalTokens int) {
meter := getRequestTokenMeter(ctx)
if meter == nil {
return
}
if promptTokens < 0 {
promptTokens = 0
}
if completionTokens < 0 {
completionTokens = 0
}
if totalTokens < 0 {
totalTokens = 0
}
meter.mu.Lock()
defer meter.mu.Unlock()
meter.promptTokens += promptTokens
meter.completionTokens += completionTokens
meter.totalTokens += totalTokens
}
// snapshotRequestTokenMeter 获取请求级 token 统计快照。
func snapshotRequestTokenMeter(ctx context.Context) RequestTokenMeterSnapshot {
meter := getRequestTokenMeter(ctx)
if meter == nil {
return RequestTokenMeterSnapshot{}
}
meter.mu.Lock()
defer meter.mu.Unlock()
return RequestTokenMeterSnapshot{
PromptTokens: meter.promptTokens,
CompletionTokens: meter.completionTokens,
TotalTokens: meter.totalTokens,
}
}
// normalizeUsageTotal 统一 total token 口径。
//
// 规则:
// 1. 模型返回 total>0 时优先使用 total
// 2. total 缺失时使用 prompt+completion 回退。
func normalizeUsageTotal(totalTokens, promptTokens, completionTokens int) int {
if totalTokens > 0 {
return totalTokens
}
sum := promptTokens + completionTokens
if sum < 0 {
return 0
}
return sum
}