wizardchen
4a9977b11f
- Introduced a loading state (`loadingChunks`) to manage document chunk loading more effectively. - Updated the scroll handling logic to prevent multiple requests while loading chunks. - Adjusted the pagination logic to ensure correct page increments based on the total document count. This update improves the user experience by ensuring smoother document loading and better state management during scrolling interactions.
315 lines
11 KiB
Go
315 lines
11 KiB
Go
package session
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/Tencent/WeKnora/internal/event"
|
|
"github.com/Tencent/WeKnora/internal/logger"
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// convertMentionedItems converts MentionedItemRequest slice to types.MentionedItems
|
|
func convertMentionedItems(items []MentionedItemRequest) types.MentionedItems {
|
|
if len(items) == 0 {
|
|
return nil
|
|
}
|
|
result := make(types.MentionedItems, len(items))
|
|
for i, item := range items {
|
|
result[i] = types.MentionedItem{
|
|
ID: item.ID,
|
|
Name: item.Name,
|
|
Type: item.Type,
|
|
KBType: item.KBType,
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
// setSSEHeaders sets the standard Server-Sent Events headers
|
|
func setSSEHeaders(c *gin.Context) {
|
|
c.Header("Content-Type", "text/event-stream")
|
|
c.Header("Cache-Control", "no-cache")
|
|
c.Header("Connection", "keep-alive")
|
|
c.Header("X-Accel-Buffering", "no")
|
|
}
|
|
|
|
// buildStreamResponse constructs a StreamResponse from a StreamEvent
|
|
func buildStreamResponse(evt interfaces.StreamEvent, requestID string) *types.StreamResponse {
|
|
response := &types.StreamResponse{
|
|
ID: requestID,
|
|
ResponseType: evt.Type,
|
|
Content: evt.Content,
|
|
Done: evt.Done,
|
|
Data: evt.Data,
|
|
}
|
|
|
|
// Extract session_id and assistant_message_id for agent_query events
|
|
if evt.Type == types.ResponseTypeAgentQuery {
|
|
if sid, ok := evt.Data["session_id"].(string); ok {
|
|
response.SessionID = sid
|
|
}
|
|
if amid, ok := evt.Data["assistant_message_id"].(string); ok {
|
|
response.AssistantMessageID = amid
|
|
}
|
|
}
|
|
|
|
// Special handling for references event
|
|
if evt.Type == types.ResponseTypeReferences {
|
|
refsData := evt.Data["references"]
|
|
if refs, ok := refsData.(types.References); ok {
|
|
response.KnowledgeReferences = refs
|
|
} else if refs, ok := refsData.([]*types.SearchResult); ok {
|
|
response.KnowledgeReferences = types.References(refs)
|
|
} else if refs, ok := refsData.([]interface{}); ok {
|
|
// Handle case where data was serialized/deserialized (e.g., from Redis)
|
|
searchResults := make([]*types.SearchResult, 0, len(refs))
|
|
for _, ref := range refs {
|
|
if refMap, ok := ref.(map[string]interface{}); ok {
|
|
sr := &types.SearchResult{
|
|
ID: getString(refMap, "id"),
|
|
Content: getString(refMap, "content"),
|
|
KnowledgeID: getString(refMap, "knowledge_id"),
|
|
ChunkIndex: int(getFloat64(refMap, "chunk_index")),
|
|
KnowledgeTitle: getString(refMap, "knowledge_title"),
|
|
StartAt: int(getFloat64(refMap, "start_at")),
|
|
EndAt: int(getFloat64(refMap, "end_at")),
|
|
Seq: int(getFloat64(refMap, "seq")),
|
|
Score: getFloat64(refMap, "score"),
|
|
ChunkType: getString(refMap, "chunk_type"),
|
|
ParentChunkID: getString(refMap, "parent_chunk_id"),
|
|
ImageInfo: getString(refMap, "image_info"),
|
|
KnowledgeFilename: getString(refMap, "knowledge_filename"),
|
|
KnowledgeSource: getString(refMap, "knowledge_source"),
|
|
KnowledgeBaseID: getString(refMap, "knowledge_base_id"),
|
|
}
|
|
searchResults = append(searchResults, sr)
|
|
}
|
|
}
|
|
response.KnowledgeReferences = types.References(searchResults)
|
|
}
|
|
}
|
|
|
|
return response
|
|
}
|
|
|
|
// sendCompletionEvent sends a final completion event to the client
|
|
// NOTE: This is now a no-op because:
|
|
// 1. The 'complete' event from handleComplete already signals stream completion
|
|
// 2. Sending an extra empty 'answer' event with done:true causes frontend issues
|
|
// (multiple done events can confuse state management)
|
|
// The frontend should use 'complete' response_type to detect stream completion
|
|
func sendCompletionEvent(c *gin.Context, requestID string) {
|
|
// Intentionally empty - completion is signaled by the 'complete' event
|
|
// which is already sent before this function is called
|
|
}
|
|
|
|
// createAgentQueryEvent creates a standard agent query event
|
|
func createAgentQueryEvent(sessionID, assistantMessageID string) interfaces.StreamEvent {
|
|
return interfaces.StreamEvent{
|
|
ID: fmt.Sprintf("query-%d", time.Now().UnixNano()),
|
|
Type: types.ResponseTypeAgentQuery,
|
|
Content: "",
|
|
Done: true,
|
|
Timestamp: time.Now(),
|
|
Data: map[string]interface{}{
|
|
"session_id": sessionID,
|
|
"assistant_message_id": assistantMessageID,
|
|
},
|
|
}
|
|
}
|
|
|
|
// createUserMessage creates a user message
|
|
func (h *Handler) createUserMessage(ctx context.Context, sessionID, query, requestID string, mentionedItems types.MentionedItems) error {
|
|
_, err := h.messageService.CreateMessage(ctx, &types.Message{
|
|
SessionID: sessionID,
|
|
Role: "user",
|
|
Content: query,
|
|
RequestID: requestID,
|
|
CreatedAt: time.Now(),
|
|
IsCompleted: true,
|
|
MentionedItems: mentionedItems,
|
|
})
|
|
return err
|
|
}
|
|
|
|
// createAssistantMessage creates an assistant message
|
|
func (h *Handler) createAssistantMessage(ctx context.Context, assistantMessage *types.Message) (*types.Message, error) {
|
|
assistantMessage.CreatedAt = time.Now()
|
|
return h.messageService.CreateMessage(ctx, assistantMessage)
|
|
}
|
|
|
|
// setupStreamHandler creates and subscribes a stream handler
|
|
func (h *Handler) setupStreamHandler(
|
|
ctx context.Context,
|
|
sessionID, assistantMessageID, requestID string,
|
|
assistantMessage *types.Message,
|
|
eventBus *event.EventBus,
|
|
) *AgentStreamHandler {
|
|
streamHandler := NewAgentStreamHandler(
|
|
ctx, sessionID, assistantMessageID, requestID,
|
|
assistantMessage, h.streamManager, eventBus,
|
|
)
|
|
streamHandler.Subscribe()
|
|
return streamHandler
|
|
}
|
|
|
|
// setupStopEventHandler registers a stop event handler
|
|
func (h *Handler) setupStopEventHandler(
|
|
eventBus *event.EventBus,
|
|
sessionID string,
|
|
sessionTenantID uint64,
|
|
assistantMessage *types.Message,
|
|
cancel context.CancelFunc,
|
|
) {
|
|
eventBus.On(event.EventStop, func(ctx context.Context, evt event.Event) error {
|
|
logger.Infof(ctx, "Received stop event, cancelling async operations for session: %s", sessionID)
|
|
cancel()
|
|
assistantMessage.Content = "用户停止了本次对话"
|
|
// Use session's tenant for message update (ctx may have effectiveTenantID when using shared agent)
|
|
updateCtx := context.WithValue(ctx, types.TenantIDContextKey, sessionTenantID)
|
|
h.completeAssistantMessage(updateCtx, assistantMessage)
|
|
return nil
|
|
})
|
|
}
|
|
|
|
// writeAgentQueryEvent writes an agent query event to the stream manager
|
|
func (h *Handler) writeAgentQueryEvent(ctx context.Context, sessionID, assistantMessageID string) {
|
|
agentQueryEvent := createAgentQueryEvent(sessionID, assistantMessageID)
|
|
if err := h.streamManager.AppendEvent(ctx, sessionID, assistantMessageID, agentQueryEvent); err != nil {
|
|
logger.ErrorWithFields(ctx, err, map[string]interface{}{
|
|
"session_id": sessionID,
|
|
"message_id": assistantMessageID,
|
|
})
|
|
// Non-fatal error, continue
|
|
}
|
|
}
|
|
|
|
// getRequestID gets the request ID from gin context
|
|
func getRequestID(c *gin.Context) string {
|
|
return c.GetString(types.RequestIDContextKey.String())
|
|
}
|
|
|
|
// Helper function for type assertion with default value
|
|
func getString(m map[string]interface{}, key string) string {
|
|
if val, ok := m[key].(string); ok {
|
|
return val
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func getFloat64(m map[string]interface{}, key string) float64 {
|
|
if val, ok := m[key].(float64); ok {
|
|
return val
|
|
}
|
|
if val, ok := m[key].(int); ok {
|
|
return float64(val)
|
|
}
|
|
return 0.0
|
|
}
|
|
|
|
// createDefaultSummaryConfig creates a default summary configuration from config
|
|
// It prioritizes tenant-level ConversationConfig, then falls back to config.yaml defaults
|
|
func (h *Handler) createDefaultSummaryConfig(ctx context.Context) *types.SummaryConfig {
|
|
// Try to get tenant from context
|
|
tenant, _ := types.TenantInfoFromContext(ctx)
|
|
|
|
// Initialize with config.yaml defaults
|
|
cfg := &types.SummaryConfig{
|
|
MaxTokens: h.config.Conversation.Summary.MaxTokens,
|
|
TopP: h.config.Conversation.Summary.TopP,
|
|
TopK: h.config.Conversation.Summary.TopK,
|
|
FrequencyPenalty: h.config.Conversation.Summary.FrequencyPenalty,
|
|
PresencePenalty: h.config.Conversation.Summary.PresencePenalty,
|
|
RepeatPenalty: h.config.Conversation.Summary.RepeatPenalty,
|
|
Prompt: h.config.Conversation.Summary.Prompt,
|
|
ContextTemplate: h.config.Conversation.Summary.ContextTemplate,
|
|
NoMatchPrefix: h.config.Conversation.Summary.NoMatchPrefix,
|
|
Temperature: h.config.Conversation.Summary.Temperature,
|
|
Seed: h.config.Conversation.Summary.Seed,
|
|
MaxCompletionTokens: h.config.Conversation.Summary.MaxCompletionTokens,
|
|
}
|
|
|
|
// Override with tenant-level conversation config if available
|
|
if tenant != nil && tenant.ConversationConfig != nil {
|
|
// Use custom prompt if provided
|
|
if tenant.ConversationConfig.Prompt != "" {
|
|
cfg.Prompt = tenant.ConversationConfig.Prompt
|
|
}
|
|
|
|
// Use custom context template if provided
|
|
if tenant.ConversationConfig.ContextTemplate != "" {
|
|
cfg.ContextTemplate = tenant.ConversationConfig.ContextTemplate
|
|
}
|
|
if tenant.ConversationConfig.Temperature > 0 {
|
|
cfg.Temperature = tenant.ConversationConfig.Temperature
|
|
}
|
|
if tenant.ConversationConfig.MaxCompletionTokens > 0 {
|
|
cfg.MaxCompletionTokens = tenant.ConversationConfig.MaxCompletionTokens
|
|
}
|
|
}
|
|
|
|
return cfg
|
|
}
|
|
|
|
// fillSummaryConfigDefaults fills missing fields in summary config with defaults
|
|
// It prioritizes tenant-level ConversationConfig, then falls back to config.yaml defaults
|
|
func (h *Handler) fillSummaryConfigDefaults(ctx context.Context, config *types.SummaryConfig) {
|
|
// Try to get tenant from context
|
|
tenant, _ := types.TenantInfoFromContext(ctx)
|
|
|
|
// Determine default values: tenant config first, then config.yaml
|
|
var defaultPrompt, defaultContextTemplate, defaultNoMatchPrefix string
|
|
var defaultTemperature float64
|
|
var defaultMaxCompletionTokens int
|
|
|
|
if tenant != nil && tenant.ConversationConfig != nil {
|
|
// Use custom prompt if provided
|
|
if tenant.ConversationConfig.Prompt != "" {
|
|
defaultPrompt = tenant.ConversationConfig.Prompt
|
|
}
|
|
|
|
// Use custom context template if provided
|
|
if tenant.ConversationConfig.ContextTemplate != "" {
|
|
defaultContextTemplate = tenant.ConversationConfig.ContextTemplate
|
|
}
|
|
defaultTemperature = tenant.ConversationConfig.Temperature
|
|
defaultMaxCompletionTokens = tenant.ConversationConfig.MaxCompletionTokens
|
|
}
|
|
|
|
// Fall back to config.yaml if tenant config is empty
|
|
if defaultPrompt == "" {
|
|
defaultPrompt = h.config.Conversation.Summary.Prompt
|
|
}
|
|
if defaultContextTemplate == "" {
|
|
defaultContextTemplate = h.config.Conversation.Summary.ContextTemplate
|
|
}
|
|
if defaultTemperature == 0 {
|
|
defaultTemperature = h.config.Conversation.Summary.Temperature
|
|
}
|
|
if defaultMaxCompletionTokens == 0 {
|
|
defaultMaxCompletionTokens = h.config.Conversation.Summary.MaxCompletionTokens
|
|
}
|
|
defaultNoMatchPrefix = h.config.Conversation.Summary.NoMatchPrefix
|
|
|
|
// Fill missing fields
|
|
if config.Prompt == "" {
|
|
config.Prompt = defaultPrompt
|
|
}
|
|
if config.ContextTemplate == "" {
|
|
config.ContextTemplate = defaultContextTemplate
|
|
}
|
|
if config.Temperature == 0 {
|
|
config.Temperature = defaultTemperature
|
|
}
|
|
if config.MaxCompletionTokens == 0 {
|
|
config.MaxCompletionTokens = defaultMaxCompletionTokens
|
|
}
|
|
if config.NoMatchPrefix == "" {
|
|
config.NoMatchPrefix = defaultNoMatchPrefix
|
|
}
|
|
}
|