diff --git a/internal/agent/engine.go b/internal/agent/engine.go index a8710a47..6d8f9240 100644 --- a/internal/agent/engine.go +++ b/internal/agent/engine.go @@ -48,6 +48,7 @@ type AgentEngine struct { tokenEstimator *agenttoken.Estimator // Token estimator for context window management memoryConsolidator *agentmemory.Consolidator // Memory consolidator for LLM-powered summarization (optional) lastUsage types.TokenUsage // Token usage from the most recent LLM call + totalUsage types.TokenUsage // Aggregate token usage across all LLM rounds in this execution lastSentMsgCount int // Number of messages sent in the most recent LLM call } @@ -183,6 +184,17 @@ func (e *AgentEngine) estimateCurrentTokens(messages []chat.Message) int { return e.tokenEstimator.EstimateMessages(messages) } +func (e *AgentEngine) recordTokenUsage(usage types.TokenUsage) { + if usage.TotalTokens <= 0 { + return + } + e.lastUsage = usage + e.totalUsage.PromptTokens += usage.PromptTokens + e.totalUsage.CompletionTokens += usage.CompletionTokens + e.totalUsage.TotalTokens += usage.TotalTokens + e.totalUsage.CachedTokens += usage.CachedTokens +} + // Execute executes the agent with conversation history and streaming output // All events are emitted to EventBus and handled by subscribers (like Handler layer) func (e *AgentEngine) Execute( @@ -191,6 +203,9 @@ func (e *AgentEngine) Execute( llmContext []chat.Message, imageURLs ...[]string, ) (*types.AgentState, error) { + e.lastUsage = types.TokenUsage{} + e.totalUsage = types.TokenUsage{} + e.lastSentMsgCount = 0 logger.Infof(ctx, "[Agent] Starting execution: session=%s, message=%s, query_len=%d, context_msgs=%d", sessionID, messageID, len(query), len(llmContext)) // Ensure tools are cleaned up after execution @@ -540,7 +555,7 @@ func (e *AgentEngine) runReActIteration( } response = resp if response.Usage.TotalTokens > 0 { - e.lastUsage = response.Usage + e.recordTokenUsage(response.Usage) logger.Debugf(ctx, "[Agent][Round-%d] Usage: prompt=%d, completion=%d, total=%d", round, response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens) diff --git a/internal/agent/finalize.go b/internal/agent/finalize.go index c80344b0..ef5e6a4a 100644 --- a/internal/agent/finalize.go +++ b/internal/agent/finalize.go @@ -169,6 +169,11 @@ func (e *AgentEngine) emitCompletionEvent( for _, ref := range state.KnowledgeRefs { knowledgeRefsInterface = append(knowledgeRefsInterface, ref) } + var tokenUsage *types.TokenUsage + if e.totalUsage.TotalTokens > 0 { + usage := e.totalUsage + tokenUsage = &usage + } e.eventBus.Emit(ctx, event.Event{ ID: generateEventID("complete"), @@ -181,6 +186,7 @@ func (e *AgentEngine) emitCompletionEvent( TotalSteps: len(state.RoundSteps), TotalDurationMs: time.Since(startTime).Milliseconds(), MessageID: messageID, // Include message ID for proper message update + TokenUsage: tokenUsage, }, }) diff --git a/internal/agent/usage_audit_test.go b/internal/agent/usage_audit_test.go new file mode 100644 index 00000000..e1c88115 --- /dev/null +++ b/internal/agent/usage_audit_test.go @@ -0,0 +1,30 @@ +package agent + +import ( + "testing" + + "github.com/Tencent/WeKnora/internal/types" +) + +func TestRecordTokenUsageAggregatesAgentRounds(t *testing.T) { + engine := &AgentEngine{} + engine.recordTokenUsage(types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120, CachedTokens: 40}) + engine.recordTokenUsage(types.TokenUsage{PromptTokens: 140, CompletionTokens: 30, TotalTokens: 170, CachedTokens: 60}) + + wantTotal := types.TokenUsage{PromptTokens: 240, CompletionTokens: 50, TotalTokens: 290, CachedTokens: 100} + if engine.totalUsage != wantTotal { + t.Fatalf("totalUsage = %#v, want %#v", engine.totalUsage, wantTotal) + } + wantLast := types.TokenUsage{PromptTokens: 140, CompletionTokens: 30, TotalTokens: 170, CachedTokens: 60} + if engine.lastUsage != wantLast { + t.Fatalf("lastUsage = %#v, want %#v", engine.lastUsage, wantLast) + } +} + +func TestRecordTokenUsageIgnoresEmptyUsage(t *testing.T) { + engine := &AgentEngine{} + engine.recordTokenUsage(types.TokenUsage{}) + if engine.totalUsage != (types.TokenUsage{}) || engine.lastUsage != (types.TokenUsage{}) { + t.Fatalf("empty usage changed state: total=%#v last=%#v", engine.totalUsage, engine.lastUsage) + } +} diff --git a/internal/application/service/chat_pipeline/chat_completion_stream.go b/internal/application/service/chat_pipeline/chat_completion_stream.go index 9f519568..336eac22 100644 --- a/internal/application/service/chat_pipeline/chat_completion_stream.go +++ b/internal/application/service/chat_pipeline/chat_completion_stream.go @@ -186,8 +186,9 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context, Type: types.EventType(event.EventAgentFinalAnswer), SessionID: chatManage.SessionID, Data: event.AgentFinalAnswerData{ - Content: response.Content, - Done: response.Done, + Content: response.Content, + Done: response.Done, + TokenUsage: response.Usage, }, }) } diff --git a/internal/event/event_data.go b/internal/event/event_data.go index 09353da5..15a3d6ac 100644 --- a/internal/event/event_data.go +++ b/internal/event/event_data.go @@ -1,5 +1,7 @@ package event +import "github.com/Tencent/WeKnora/internal/types" + // EventData contains common event data structures for different stages // QueryData represents query-related event data @@ -142,6 +144,7 @@ type AgentCompleteData struct { TotalDurationMs int64 `json:"total_duration_ms"` MessageID string `json:"message_id,omitempty"` // Assistant message ID RequestID string `json:"request_id,omitempty"` + TokenUsage *types.TokenUsage `json:"token_usage,omitempty"` Extra map[string]interface{} `json:"extra,omitempty"` } @@ -184,9 +187,10 @@ type AgentReferencesData struct { // AgentFinalAnswerData represents final answer streaming data type AgentFinalAnswerData struct { - Content string `json:"content"` - Done bool `json:"done"` - IsFallback bool `json:"is_fallback,omitempty"` // True when response is a fallback (no knowledge base match) + Content string `json:"content"` + Done bool `json:"done"` + IsFallback bool `json:"is_fallback,omitempty"` // True when response is a fallback (no knowledge base match) + TokenUsage *types.TokenUsage `json:"token_usage,omitempty"` } // AgentReflectionData represents agent reflection data diff --git a/internal/handler/session/agent_stream_handler.go b/internal/handler/session/agent_stream_handler.go index cde73e89..880ddb0d 100644 --- a/internal/handler/session/agent_stream_handler.go +++ b/internal/handler/session/agent_stream_handler.go @@ -475,6 +475,10 @@ func (h *AgentStreamHandler) handleFinalAnswer(ctx context.Context, evt event.Ev if data.IsFallback { h.assistantMessage.IsFallback = true } + if data.TokenUsage != nil { + usage := *data.TokenUsage + h.assistantMessage.TokenUsage = &usage + } // Calculate duration if done var metadata map[string]interface{} @@ -604,6 +608,10 @@ func (h *AgentStreamHandler) handleComplete(ctx context.Context, evt event.Event // h.assistantMessage.Content = data.FinalAnswer h.assistantMessage.IsCompleted = true h.assistantMessage.AgentDurationMs = data.TotalDurationMs + if data.TokenUsage != nil { + usage := *data.TokenUsage + h.assistantMessage.TokenUsage = &usage + } // Update knowledge references if provided if len(data.KnowledgeRefs) > 0 { diff --git a/internal/handler/session/qa.go b/internal/handler/session/qa.go index 845c9270..929c5ecf 100644 --- a/internal/handler/session/qa.go +++ b/internal/handler/session/qa.go @@ -718,6 +718,10 @@ func (h *Handler) executeQA(reqCtx *qaRequestContext, mode qaMode, generateTitle return nil } streamCtx.assistantMessage.Content += data.Content + if data.TokenUsage != nil { + usage := *data.TokenUsage + streamCtx.assistantMessage.TokenUsage = &usage + } if data.IsFallback { streamCtx.assistantMessage.IsFallback = true } diff --git a/internal/im/service.go b/internal/im/service.go index 31b17a07..24e99401 100644 --- a/internal/im/service.go +++ b/internal/im/service.go @@ -521,8 +521,47 @@ func buildIMLastRequestState(agentID string, customAgent *types.CustomAgent, kbI return state } -func createIMUserMessagePayload(sessionID, content, requestID string) *types.Message { - return &types.Message{ +type imMessageAuditContext struct { + platform string + channelID string + messageID string + userID string + chatID string + threadID string + knowledgeBaseIDs types.MessageKnowledgeBaseIDs +} + +func newIMMessageAuditContext(msg *IncomingMessage, channelID string, kbIDs []string) *imMessageAuditContext { + audit := &imMessageAuditContext{ + channelID: channelID, + knowledgeBaseIDs: append(types.MessageKnowledgeBaseIDs(nil), kbIDs...), + } + if msg == nil { + return audit + } + audit.platform = string(msg.Platform) + audit.messageID = msg.MessageID + audit.userID = msg.UserID + audit.chatID = msg.ChatID + audit.threadID = msg.ThreadID + return audit +} + +func applyIMMessageAudit(message *types.Message, audit *imMessageAuditContext) { + if message == nil || audit == nil { + return + } + message.SourcePlatform = audit.platform + message.SourceChannelID = audit.channelID + message.SourceMessageID = audit.messageID + message.SourceUserID = audit.userID + message.SourceChatID = audit.chatID + message.SourceThreadID = audit.threadID + message.SearchedKnowledgeBaseIDs = append(types.MessageKnowledgeBaseIDs(nil), audit.knowledgeBaseIDs...) +} + +func createIMUserMessagePayload(sessionID, content, requestID string, audit *imMessageAuditContext) *types.Message { + message := &types.Message{ SessionID: sessionID, Role: "user", Content: content, @@ -531,10 +570,12 @@ func createIMUserMessagePayload(sessionID, content, requestID string) *types.Mes IsCompleted: true, Channel: "im", } + applyIMMessageAudit(message, audit) + return message } -func createIMAssistantMessagePayload(sessionID, requestID string) *types.Message { - return &types.Message{ +func createIMAssistantMessagePayload(sessionID, requestID string, audit *imMessageAuditContext) *types.Message { + message := &types.Message{ SessionID: sessionID, Role: "assistant", RequestID: requestID, @@ -542,6 +583,8 @@ func createIMAssistantMessagePayload(sessionID, requestID string) *types.Message IsCompleted: false, Channel: "im", } + applyIMMessageAudit(message, audit) + return message } func collectIMKnowledgeReferences(dst *[]*types.SearchResult, refs interface{}) { @@ -577,6 +620,10 @@ func applyIMCompleteDataToMessage(msg *types.Message, data event.AgentCompleteDa } msg.IsCompleted = true msg.AgentDurationMs = data.TotalDurationMs + if data.TokenUsage != nil { + usage := *data.TokenUsage + msg.TokenUsage = &usage + } if len(data.KnowledgeRefs) > 0 { refs := make([]*types.SearchResult, 0, len(data.KnowledgeRefs)) collectIMKnowledgeReferences(&refs, data.KnowledgeRefs) @@ -589,6 +636,14 @@ func applyIMCompleteDataToMessage(msg *types.Message, data event.AgentCompleteDa } } +func applyIMFinalAnswerDataToMessage(msg *types.Message, data event.AgentFinalAnswerData) { + if msg == nil || data.TokenUsage == nil { + return + } + usage := *data.TokenUsage + msg.TokenUsage = &usage +} + // waitForIMAgentComplete blocks until EventAgentComplete, ctx cancellation, or timeout. func waitForIMAgentComplete(ctx context.Context, completeDone <-chan struct{}, sessionID string) { timer := time.NewTimer(agentCompleteWaitTimeout) @@ -1417,6 +1472,11 @@ func (s *Service) executeQARequest(req *qaRequest) { // kbIDs is left empty so the QA pipeline resolves them from the agent config. var kbIDs []string + auditedKBIDs := kbIDs + if len(auditedKBIDs) == 0 && req.agent != nil { + auditedKBIDs = req.agent.Config.KnowledgeBases + } + audit := newIMMessageAuditContext(req.msg, req.channelID, auditedKBIDs) // Determine output mode from channel config. streamDisabled := req.channel.OutputMode == "full" @@ -1424,7 +1484,7 @@ func (s *Service) executeQARequest(req *qaRequest) { // If the adapter supports streaming and output is not "full", use streaming. if !streamDisabled { if streamer, ok := req.adapter.(StreamSender); ok { - if err := s.handleMessageStream(ctx, req.msg, req.session, req.agent, kbIDs, streamer, req.adapter, req.userKey, req.tenant); err != nil { + if err := s.handleMessageStream(ctx, req.msg, req.session, req.agent, kbIDs, streamer, req.adapter, req.userKey, req.tenant, audit); err != nil { logger.Errorf(ctx, "[IM] Stream QA failed: %v", err) } return @@ -1432,7 +1492,7 @@ func (s *Service) executeQARequest(req *qaRequest) { } // Non-streaming fallback: collect full answer then send. - answer, err := s.runQA(ctx, req.session, req.msg.Content, req.agent, kbIDs, req.userKey, req.msg.Quote) + answer, err := s.runQA(ctx, req.session, req.msg.Content, req.agent, kbIDs, req.userKey, req.msg.Quote, audit) if err != nil { logger.Errorf(ctx, "[IM] QA failed: %v, sending fallback reply", err) answer = "抱歉，处理您的问题时出现了异常，请稍后再试。" @@ -1855,12 +1915,12 @@ func briefToolSummary(output string) string { // handleMessageStream runs the QA pipeline and streams answer chunks to the IM platform // in real-time via the StreamSender interface. Chunks are batched at streamFlushInterval // to avoid API rate-limiting. -func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, streamer StreamSender, adapter Adapter, userKey string, tenant *types.Tenant) error { +func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, streamer StreamSender, adapter Adapter, userKey string, tenant *types.Tenant, audit *imMessageAuditContext) error { // Start the stream on the IM platform (e.g., create Feishu streaming card) streamID, err := streamer.StartStream(ctx, msg) if err != nil { logger.Warnf(ctx, "[IM] StartStream failed, falling back to non-streaming: %v", err) - return s.fallbackNonStream(ctx, msg, session, customAgent, kbIDs, adapter, userKey, tenant) + return s.fallbackNonStream(ctx, msg, session, customAgent, kbIDs, adapter, userKey, tenant, audit) } // Prepare the QA pipeline @@ -1951,6 +2011,7 @@ func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage, } bufMu.Lock() + applyIMFinalAnswerDataToMessage(assistantMsg, data) if useAgent && !agentDone { if data.Content != "" { agentLiveAnswer.WriteString(data.Content) @@ -2104,13 +2165,13 @@ func (s *Service) handleMessageStream(ctx context.Context, msg *IncomingMessage, requestID := uuid.New().String() // Create user message - userMsg, err := s.messageService.CreateMessage(qaCtx, createIMUserMessagePayload(session.ID, msg.Content, requestID)) + userMsg, err := s.messageService.CreateMessage(qaCtx, createIMUserMessagePayload(session.ID, msg.Content, requestID, audit)) if err != nil { return fmt.Errorf("create user message: %w", err) } // Create placeholder assistant message - assistantMsg, err = s.messageService.CreateMessage(qaCtx, createIMAssistantMessagePayload(session.ID, requestID)) + assistantMsg, err = s.messageService.CreateMessage(qaCtx, createIMAssistantMessagePayload(session.ID, requestID, audit)) if err != nil { return fmt.Errorf("create assistant message: %w", err) } @@ -2246,8 +2307,8 @@ loop: } // fallbackNonStream is used when streaming initialization fails. -func (s *Service) fallbackNonStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, adapter Adapter, userKey string, tenant *types.Tenant) error { - answer, err := s.runQA(ctx, session, msg.Content, customAgent, kbIDs, userKey, msg.Quote) +func (s *Service) fallbackNonStream(ctx context.Context, msg *IncomingMessage, session *types.Session, customAgent *types.CustomAgent, kbIDs []string, adapter Adapter, userKey string, tenant *types.Tenant, audit *imMessageAuditContext) error { + answer, err := s.runQA(ctx, session, msg.Content, customAgent, kbIDs, userKey, msg.Quote, audit) if err != nil { logger.Errorf(ctx, "[IM] QA fallback failed: %v", err) answer = "抱歉，处理您的问题时出现了异常，请稍后再试。" @@ -2257,7 +2318,7 @@ func (s *Service) fallbackNonStream(ctx context.Context, msg *IncomingMessage, s } // runQA executes the WeKnora QA pipeline and returns the full answer text. -func (s *Service) runQA(ctx context.Context, session *types.Session, query string, customAgent *types.CustomAgent, kbIDs []string, userKey string, quote *QuotedMessage) (string, error) { +func (s *Service) runQA(ctx context.Context, session *types.Session, query string, customAgent *types.CustomAgent, kbIDs []string, userKey string, quote *QuotedMessage, audit *imMessageAuditContext) (string, error) { // Cancellable context (no hard deadline): each agent round has its own // LLMCallTimeout. The context can still be cancelled by /stop. ctx, cancel := context.WithCancel(ctx) @@ -2268,6 +2329,7 @@ func (s *Service) runQA(ctx context.Context, session *types.Session, query strin // Thread-safe answer collection var answerMu sync.Mutex var answerBuilder strings.Builder + var assistantMsg *types.Message var qaErr error done := make(chan struct{}) completeDone := make(chan struct{}) @@ -2282,6 +2344,7 @@ func (s *Service) runQA(ctx context.Context, session *types.Session, query strin return nil } answerMu.Lock() + applyIMFinalAnswerDataToMessage(assistantMsg, data) answerBuilder.WriteString(data.Content) answerMu.Unlock() if data.Done { @@ -2311,13 +2374,13 @@ func (s *Service) runQA(ctx context.Context, session *types.Session, query strin requestID := uuid.New().String() // Create user message so it appears in conversation history - userMsg, err := s.messageService.CreateMessage(ctx, createIMUserMessagePayload(session.ID, query, requestID)) + userMsg, err := s.messageService.CreateMessage(ctx, createIMUserMessagePayload(session.ID, query, requestID, audit)) if err != nil { return "", fmt.Errorf("create user message: %w", err) } // Create a placeholder assistant message - assistantMsg, err := s.messageService.CreateMessage(ctx, createIMAssistantMessagePayload(session.ID, requestID)) + assistantMsg, err = s.messageService.CreateMessage(ctx, createIMAssistantMessagePayload(session.ID, requestID, audit)) if err != nil { return "", fmt.Errorf("create assistant message: %w", err) } diff --git a/internal/im/session_test.go b/internal/im/session_test.go index 513cd9e0..43db6a75 100644 --- a/internal/im/session_test.go +++ b/internal/im/session_test.go @@ -202,8 +202,16 @@ func TestBuildIMLastRequestStateKeepsExplicitKBs(t *testing.T) { } func TestCreateIMMessagePayloadsShareRequestShape(t *testing.T) { - userMsg := createIMUserMessagePayload("session-1", "hello", "request-1") - assistantMsg := createIMAssistantMessagePayload("session-1", "request-1") + incoming := &IncomingMessage{ + Platform: PlatformDingtalk, + UserID: "platform-user-1", + ChatID: "platform-chat-1", + ThreadID: "platform-thread-1", + MessageID: "platform-message-1", + } + audit := newIMMessageAuditContext(incoming, "channel-1", []string{"kb-1", "kb-2"}) + userMsg := createIMUserMessagePayload("session-1", "hello", "request-1", audit) + assistantMsg := createIMAssistantMessagePayload("session-1", "request-1", audit) if userMsg.SessionID != "session-1" || assistantMsg.SessionID != "session-1" { t.Fatalf("SessionID mismatch: user=%q assistant=%q", userMsg.SessionID, assistantMsg.SessionID) @@ -229,6 +237,20 @@ func TestCreateIMMessagePayloadsShareRequestShape(t *testing.T) { if assistantMsg.Content != "" { t.Fatalf("assistant placeholder content = %q, want empty", assistantMsg.Content) } + for _, message := range []*types.Message{userMsg, assistantMsg} { + if message.SourcePlatform != "dingtalk" || message.SourceChannelID != "channel-1" { + t.Fatalf("source channel metadata = platform %q channel %q", message.SourcePlatform, message.SourceChannelID) + } + if message.SourceMessageID != "platform-message-1" || message.SourceUserID != "platform-user-1" { + t.Fatalf("source actor metadata = message %q user %q", message.SourceMessageID, message.SourceUserID) + } + if message.SourceChatID != "platform-chat-1" || message.SourceThreadID != "platform-thread-1" { + t.Fatalf("source conversation metadata = chat %q thread %q", message.SourceChatID, message.SourceThreadID) + } + if !reflect.DeepEqual(message.SearchedKnowledgeBaseIDs, types.MessageKnowledgeBaseIDs{"kb-1", "kb-2"}) { + t.Fatalf("searched knowledge bases = %#v", message.SearchedKnowledgeBaseIDs) + } + } } func TestApplyIMCompleteDataToMessage(t *testing.T) { @@ -239,11 +261,13 @@ func TestApplyIMCompleteDataToMessage(t *testing.T) { ReasoningContent: "thinking", }} + usage := &types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120, CachedTokens: 40} applyIMCompleteDataToMessage(msg, event.AgentCompleteData{ MessageID: "assistant-1", TotalDurationMs: 1234, KnowledgeRefs: []interface{}{ref}, AgentSteps: steps, + TokenUsage: usage, }) if !msg.IsCompleted { @@ -258,6 +282,20 @@ func TestApplyIMCompleteDataToMessage(t *testing.T) { if len(msg.AgentSteps) != 1 || msg.AgentSteps[0].ReasoningContent != "thinking" { t.Fatalf("AgentSteps = %#v, want one thinking step", msg.AgentSteps) } + if !reflect.DeepEqual(msg.TokenUsage, usage) { + t.Fatalf("TokenUsage = %#v, want %#v", msg.TokenUsage, usage) + } +} + +func TestApplyIMFinalAnswerUsageToMessage(t *testing.T) { + msg := &types.Message{ID: "assistant-1", Role: "assistant"} + usage := &types.TokenUsage{PromptTokens: 80, CompletionTokens: 16, TotalTokens: 96} + + applyIMFinalAnswerDataToMessage(msg, event.AgentFinalAnswerData{Done: true, TokenUsage: usage}) + + if !reflect.DeepEqual(msg.TokenUsage, usage) { + t.Fatalf("TokenUsage = %#v, want %#v", msg.TokenUsage, usage) + } } func TestPickIMStoredAnswerPrefersFirstNonEmpty(t *testing.T) { diff --git a/internal/types/message.go b/internal/types/message.go index c4bdd262..f4417c44 100644 --- a/internal/types/message.go +++ b/internal/types/message.go @@ -150,6 +150,10 @@ func (m *MessageAttachments) Scan(value interface{}) error { // MentionedItems is a slice of MentionedItem for database storage type MentionedItems []MentionedItem +// MessageKnowledgeBaseIDs stores the knowledge-base scope searched for a message. +// It is kept even when retrieval returns no references so the request remains auditable. +type MessageKnowledgeBaseIDs []string + // Value implements the driver.Valuer interface for database serialization func (m MentionedItems) Value() (driver.Value, error) { if m == nil { @@ -216,6 +220,22 @@ type Message struct { RenderedContent string `json:"-" gorm:"type:text;column:rendered_content;default:''"` // Channel indicates the source channel of this message (e.g., "web", "api", "im") Channel string `json:"channel,omitempty" gorm:"type:varchar(50);default:''"` + // SourcePlatform is the concrete external platform for IM messages (for example, dingtalk). + SourcePlatform string `json:"source_platform,omitempty" gorm:"type:varchar(20);default:''"` + // SourceChannelID identifies the configured WeKnora IM channel that received the message. + SourceChannelID string `json:"source_channel_id,omitempty" gorm:"type:varchar(36);default:''"` + // SourceMessageID is the platform message identifier used for end-to-end correlation. + SourceMessageID string `json:"source_message_id,omitempty" gorm:"type:varchar(128);default:''"` + // SourceUserID records the platform user for this individual message. + SourceUserID string `json:"source_user_id,omitempty" gorm:"type:varchar(128);default:''"` + // SourceChatID records the platform conversation or group identifier. + SourceChatID string `json:"source_chat_id,omitempty" gorm:"type:varchar(128);default:''"` + // SourceThreadID records the platform thread identifier when present. + SourceThreadID string `json:"source_thread_id,omitempty" gorm:"type:varchar(128);default:''"` + // SearchedKnowledgeBaseIDs records the request scope even when retrieval has zero hits. + SearchedKnowledgeBaseIDs MessageKnowledgeBaseIDs `json:"searched_knowledge_base_ids,omitempty" gorm:"serializer:json;type:jsonb;column:searched_knowledge_base_ids"` + // TokenUsage records answer-generation model usage for the assistant message. + TokenUsage *TokenUsage `json:"token_usage,omitempty" gorm:"serializer:json;type:jsonb;column:token_usage"` // KnowledgeID links this message to a Knowledge entry in the chat history knowledge base // Used for vector search indexing: when set, the message content has been indexed as a Knowledge passage KnowledgeID string `json:"knowledge_id,omitempty" gorm:"type:varchar(36);index"` diff --git a/internal/types/message_audit_migration_test.go b/internal/types/message_audit_migration_test.go new file mode 100644 index 00000000..2b807b82 --- /dev/null +++ b/internal/types/message_audit_migration_test.go @@ -0,0 +1,57 @@ +package types + +import ( + "os" + "strings" + "testing" +) + +func TestMessageAuditMigrationCoversModelFields(t *testing.T) { + upPath := "../../migrations/versioned/000064_message_audit.up.sql" + downPath := "../../migrations/versioned/000064_message_audit.down.sql" + + up, err := os.ReadFile(upPath) + if err != nil { + t.Fatalf("read up migration: %v", err) + } + down, err := os.ReadFile(downPath) + if err != nil { + t.Fatalf("read down migration: %v", err) + } + + columns := []string{ + "source_platform", + "source_channel_id", + "source_message_id", + "source_user_id", + "source_chat_id", + "source_thread_id", + "searched_knowledge_base_ids", + "token_usage", + } + upSQL := strings.ToLower(string(up)) + downSQL := strings.ToLower(string(down)) + for _, column := range columns { + if !strings.Contains(upSQL, column) { + t.Errorf("up migration missing column %q", column) + } + if !strings.Contains(downSQL, column) { + t.Errorf("down migration missing column %q", column) + } + } + + indexes := []string{ + "idx_messages_request_role", + "idx_messages_source_channel", + "idx_messages_source_message", + "idx_messages_source_user", + } + for _, index := range indexes { + if !strings.Contains(upSQL, index) { + t.Errorf("up migration missing index %q", index) + } + if !strings.Contains(downSQL, index) { + t.Errorf("down migration missing index %q", index) + } + } +} diff --git a/internal/types/message_audit_test.go b/internal/types/message_audit_test.go new file mode 100644 index 00000000..03aeedf6 --- /dev/null +++ b/internal/types/message_audit_test.go @@ -0,0 +1,51 @@ +package types + +import ( + "encoding/json" + "reflect" + "testing" +) + +func TestMessageCarriesAuditMetadata(t *testing.T) { + original := Message{ + ID: "message-1", + SourcePlatform: "dingtalk", + SourceChannelID: "channel-1", + SourceMessageID: "external-message-1", + SourceUserID: "user-1", + SourceChatID: "chat-1", + SourceThreadID: "thread-1", + SearchedKnowledgeBaseIDs: MessageKnowledgeBaseIDs{"kb-1", "kb-2"}, + TokenUsage: &TokenUsage{ + PromptTokens: 100, + CompletionTokens: 20, + TotalTokens: 120, + CachedTokens: 40, + }, + } + + payload, err := json.Marshal(original) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + + var decoded Message + if err := json.Unmarshal(payload, &decoded); err != nil { + t.Fatalf("json.Unmarshal() error = %v", err) + } + + if decoded.SourcePlatform != original.SourcePlatform || + decoded.SourceChannelID != original.SourceChannelID || + decoded.SourceMessageID != original.SourceMessageID || + decoded.SourceUserID != original.SourceUserID || + decoded.SourceChatID != original.SourceChatID || + decoded.SourceThreadID != original.SourceThreadID { + t.Fatalf("source metadata was not preserved: got %+v", decoded) + } + if !reflect.DeepEqual(decoded.SearchedKnowledgeBaseIDs, original.SearchedKnowledgeBaseIDs) { + t.Fatalf("searched knowledge bases = %#v, want %#v", decoded.SearchedKnowledgeBaseIDs, original.SearchedKnowledgeBaseIDs) + } + if !reflect.DeepEqual(decoded.TokenUsage, original.TokenUsage) { + t.Fatalf("token usage = %#v, want %#v", decoded.TokenUsage, original.TokenUsage) + } +} diff --git a/migrations/sqlite/000000_init.up.sql b/migrations/sqlite/000000_init.up.sql index 9a98d57b..5efe350e 100644 --- a/migrations/sqlite/000000_init.up.sql +++ b/migrations/sqlite/000000_init.up.sql @@ -168,6 +168,14 @@ CREATE TABLE IF NOT EXISTS messages ( is_completed BOOLEAN NOT NULL DEFAULT 0, is_fallback BOOLEAN NOT NULL DEFAULT 0, channel VARCHAR(50) NOT NULL DEFAULT '', + source_platform VARCHAR(20) NOT NULL DEFAULT '', + source_channel_id VARCHAR(36) NOT NULL DEFAULT '', + source_message_id VARCHAR(128) NOT NULL DEFAULT '', + source_user_id VARCHAR(128) NOT NULL DEFAULT '', + source_chat_id VARCHAR(128) NOT NULL DEFAULT '', + source_thread_id VARCHAR(128) NOT NULL DEFAULT '', + searched_knowledge_base_ids TEXT NOT NULL DEFAULT '[]', + token_usage TEXT, agent_duration_ms INTEGER DEFAULT 0, knowledge_id VARCHAR(36), created_at DATETIME DEFAULT CURRENT_TIMESTAMP, @@ -177,6 +185,10 @@ CREATE TABLE IF NOT EXISTS messages ( CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages(session_id); CREATE INDEX IF NOT EXISTS idx_messages_knowledge_id ON messages(knowledge_id); +CREATE INDEX IF NOT EXISTS idx_messages_request_role ON messages(request_id, role); +CREATE INDEX IF NOT EXISTS idx_messages_source_channel ON messages(source_platform, source_channel_id); +CREATE INDEX IF NOT EXISTS idx_messages_source_message ON messages(source_platform, source_message_id); +CREATE INDEX IF NOT EXISTS idx_messages_source_user ON messages(source_platform, source_user_id); CREATE TABLE IF NOT EXISTS chunks ( id VARCHAR(36) PRIMARY KEY, diff --git a/migrations/versioned/000064_message_audit.down.sql b/migrations/versioned/000064_message_audit.down.sql new file mode 100644 index 00000000..9aa3af5a --- /dev/null +++ b/migrations/versioned/000064_message_audit.down.sql @@ -0,0 +1,18 @@ +-- Migration: 000064_message_audit (rollback) +DO $$ BEGIN RAISE NOTICE '[Migration 000064 rollback] Removing message audit metadata'; END $$; + +DROP INDEX IF EXISTS idx_messages_source_user; +DROP INDEX IF EXISTS idx_messages_source_message; +DROP INDEX IF EXISTS idx_messages_source_channel; +DROP INDEX IF EXISTS idx_messages_request_role; + +ALTER TABLE messages DROP COLUMN IF EXISTS token_usage; +ALTER TABLE messages DROP COLUMN IF EXISTS searched_knowledge_base_ids; +ALTER TABLE messages DROP COLUMN IF EXISTS source_thread_id; +ALTER TABLE messages DROP COLUMN IF EXISTS source_chat_id; +ALTER TABLE messages DROP COLUMN IF EXISTS source_user_id; +ALTER TABLE messages DROP COLUMN IF EXISTS source_message_id; +ALTER TABLE messages DROP COLUMN IF EXISTS source_channel_id; +ALTER TABLE messages DROP COLUMN IF EXISTS source_platform; + +DO $$ BEGIN RAISE NOTICE '[Migration 000064 rollback] Message audit metadata removed'; END $$; diff --git a/migrations/versioned/000064_message_audit.up.sql b/migrations/versioned/000064_message_audit.up.sql new file mode 100644 index 00000000..92f5f0aa --- /dev/null +++ b/migrations/versioned/000064_message_audit.up.sql @@ -0,0 +1,35 @@ +-- Migration: 000064_message_audit +-- Description: Persist per-message source metadata, retrieval scope, and model token usage. +DO $$ BEGIN RAISE NOTICE '[Migration 000064] Adding message audit metadata'; END $$; + +ALTER TABLE messages ADD COLUMN IF NOT EXISTS source_platform VARCHAR(20) NOT NULL DEFAULT ''; +ALTER TABLE messages ADD COLUMN IF NOT EXISTS source_channel_id VARCHAR(36) NOT NULL DEFAULT ''; +ALTER TABLE messages ADD COLUMN IF NOT EXISTS source_message_id VARCHAR(128) NOT NULL DEFAULT ''; +ALTER TABLE messages ADD COLUMN IF NOT EXISTS source_user_id VARCHAR(128) NOT NULL DEFAULT ''; +ALTER TABLE messages ADD COLUMN IF NOT EXISTS source_chat_id VARCHAR(128) NOT NULL DEFAULT ''; +ALTER TABLE messages ADD COLUMN IF NOT EXISTS source_thread_id VARCHAR(128) NOT NULL DEFAULT ''; +ALTER TABLE messages ADD COLUMN IF NOT EXISTS searched_knowledge_base_ids JSONB NOT NULL DEFAULT '[]'::jsonb; +ALTER TABLE messages ADD COLUMN IF NOT EXISTS token_usage JSONB; + +CREATE INDEX IF NOT EXISTS idx_messages_request_role + ON messages (request_id, role); +CREATE INDEX IF NOT EXISTS idx_messages_source_channel + ON messages (source_platform, source_channel_id) + WHERE source_platform != ''; +CREATE INDEX IF NOT EXISTS idx_messages_source_message + ON messages (source_platform, source_message_id) + WHERE source_message_id != ''; +CREATE INDEX IF NOT EXISTS idx_messages_source_user + ON messages (source_platform, source_user_id) + WHERE source_user_id != ''; + +COMMENT ON COLUMN messages.source_platform IS 'Concrete source platform such as dingtalk; channel remains im for compatibility'; +COMMENT ON COLUMN messages.source_channel_id IS 'WeKnora im_channels.id that received the message'; +COMMENT ON COLUMN messages.source_message_id IS 'External platform message identifier'; +COMMENT ON COLUMN messages.source_user_id IS 'External platform user identifier for this individual message'; +COMMENT ON COLUMN messages.source_chat_id IS 'External platform conversation or group identifier'; +COMMENT ON COLUMN messages.source_thread_id IS 'External platform thread identifier when available'; +COMMENT ON COLUMN messages.searched_knowledge_base_ids IS 'Knowledge-base search scope, retained even when retrieval has zero hits'; +COMMENT ON COLUMN messages.token_usage IS 'Answer-generation token usage: prompt, completion, total, and cached tokens'; + +DO $$ BEGIN RAISE NOTICE '[Migration 000064] Message audit metadata ready'; END $$;