Enterprise-Knowledge-Base/internal/handler/session/helpers.go

package session

import (
	"context"
	"fmt"
	"strings"
	"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"
)

// convertImageAttachments converts ImageAttachment slice to types.MessageImages
func convertImageAttachments(items []ImageAttachment) types.MessageImages {
	if len(items) == 0 {
		return nil
	}
	result := make(types.MessageImages, len(items))
	for i, item := range items {
		result[i] = types.MessageImage{
			URL:     item.URL,
			Caption: item.Caption,
		}
	}
	return result
}

// extractImageURLsAndOCRText extracts image references and concatenated analysis text.
// For LLM consumption it prefers the raw Data (data URI) when available so that
// image_resolve can skip the disk round-trip; falls back to the storage URL otherwise.
func extractImageURLsAndOCRText(images []ImageAttachment) (urls []string, ocrText string) {
	if len(images) == 0 {
		return nil, ""
	}
	urls = make([]string, 0, len(images))
	var parts []string
	for _, img := range images {
		switch {
		case img.Data != "":
			urls = append(urls, img.Data)
		case img.URL != "":
			urls = append(urls, img.URL)
		}
		if img.Caption != "" {
			parts = append(parts, img.Caption)
		}
	}
	if len(parts) > 0 {
		ocrText = strings.Join(parts, "\n")
	}
	return
}

// 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"),
						KnowledgeDescription: getString(refMap, "knowledge_description"),
						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 and returns the created message.
func (h *Handler) createUserMessage(ctx context.Context, sessionID, query, requestID string, mentionedItems types.MentionedItems, images types.MessageImages, attachments types.MessageAttachments, channel string) (*types.Message, error) {
	return h.messageService.CreateMessage(ctx, &types.Message{
		SessionID:      sessionID,
		Role:           "user",
		Content:        query,
		RequestID:      requestID,
		CreatedAt:      time.Now(),
		IsCompleted:    true,
		MentionedItems: mentionedItems,
		Images:         images,
		Attachments:    attachments,
		Channel:        channel,
	})
}

// 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,
	receivedAt time.Time,
	assistantMessage *types.Message,
	eventBus *event.EventBus,
) *AgentStreamHandler {
	streamHandler := NewAgentStreamHandler(
		ctx, sessionID, assistantMessageID, requestID, receivedAt,
		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()
		// Preserve whatever has been streamed so far; do not overwrite Content.
		// Use session's tenant for message update (ctx may have effectiveTenantID when using shared agent).
		// Use WithoutCancel so the GORM UPDATE survives the upcoming ctx.Done triggered by cancel()/client disconnect.
		updateCtx := context.WithValue(
			context.WithoutCancel(ctx),
			types.TenantIDContextKey, sessionTenantID,
		)
		h.completeAssistantMessage(updateCtx, assistantMessage, "") // empty query: stopped conversations are not indexed
		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 and fillSummaryConfigDefaults used to build
// per-session SummaryConfig from tenant-level ConversationConfig + config.yaml
// defaults. Both helpers became unreachable when the chat pipeline moved to
// CustomAgent (builtin-quick-answer / smart-reasoning) and the tenant-level
// ConversationConfig field was removed; deleting them avoids the only
// remaining references to that defunct path.