Enterprise-Knowledge-Base/internal/handler/session/helpers.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"
)

// 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 {
		if refs, ok := evt.Data["references"].(types.References); ok {
			response.KnowledgeReferences = refs
		}
	}

	return response
}

// sendCompletionEvent sends a final completion event to the client
func sendCompletionEvent(c *gin.Context, requestID string) {
	c.SSEvent("message", &types.StreamResponse{
		ID:           requestID,
		ResponseType: types.ResponseTypeAnswer,
		Content:      "",
		Done:         true,
	})
	c.Writer.Flush()
}

// 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) error {
	_, err := h.messageService.CreateMessage(ctx, &types.Message{
		SessionID:   sessionID,
		Role:        "user",
		Content:     query,
		RequestID:   requestID,
		CreatedAt:   time.Now(),
		IsCompleted: true,
	})
	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,
	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 = "用户停止了本次对话"
		h.completeAssistantMessage(ctx, 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
func (h *Handler) createDefaultSummaryConfig() *types.SummaryConfig {
	return &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,
	}
}

// fillSummaryConfigDefaults fills missing fields in summary config with defaults
func (h *Handler) fillSummaryConfigDefaults(config *types.SummaryConfig) {
	if config.Prompt == "" {
		config.Prompt = h.config.Conversation.Summary.Prompt
	}
	if config.ContextTemplate == "" {
		config.ContextTemplate = h.config.Conversation.Summary.ContextTemplate
	}
	if config.NoMatchPrefix == "" {
		config.NoMatchPrefix = h.config.Conversation.Summary.NoMatchPrefix
	}
}