Enterprise-Knowledge-Base/internal/agent/think.go

package agent

import (
	"context"
	"fmt"
	"time"

	agenttools "github.com/Tencent/WeKnora/internal/agent/tools"
	"github.com/Tencent/WeKnora/internal/common"
	"github.com/Tencent/WeKnora/internal/event"
	"github.com/Tencent/WeKnora/internal/logger"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
)

// streamLLMResult holds accumulated output from a streaming LLM call.
type streamLLMResult struct {
	Content   string
	ToolCalls []types.LLMToolCall
	Usage     *types.TokenUsage
}

// streamLLMToEventBus streams LLM response through EventBus (generic method)
// emitFunc: callback to emit each chunk event
func (e *AgentEngine) streamLLMToEventBus(
	ctx context.Context,
	messages []chat.Message,
	opts *chat.ChatOptions,
	emitFunc func(chunk *types.StreamResponse, fullContent string),
) (*streamLLMResult, error) {
	logger.Debugf(ctx, "[Agent][Stream] Starting LLM stream with %d messages", len(messages))

	llmCtx, llmCancel := context.WithTimeout(ctx, e.getLLMCallTimeout())
	defer llmCancel()

	stream, err := e.chatModel.ChatStream(llmCtx, messages, opts)
	if err != nil {
		logger.Errorf(ctx, "[Agent][Stream] Failed to start LLM stream: %v", err)
		return nil, err
	}

	result := &streamLLMResult{}
	chunkCount := 0

	for chunk := range stream {
		chunkCount++

		if chunk.Content != "" {
			isExtracted := chunk.Data != nil && chunk.Data["source"] != nil
			if !isExtracted {
				result.Content += chunk.Content
			}
		}

		if len(chunk.ToolCalls) > 0 {
			result.ToolCalls = chunk.ToolCalls
		}

		if chunk.Usage != nil {
			result.Usage = chunk.Usage
		}

		if emitFunc != nil {
			emitFunc(&chunk, result.Content)
		}
	}

	return result, nil
}

// streamReflectionToEventBus streams reflection process through EventBus
// Note: Reflection is now handled through the think tool in main loop
func (e *AgentEngine) streamReflectionToEventBus(
	ctx context.Context,
	toolCallID string,
	toolName string,
	result string,
	iteration int,
	sessionID string,
) (string, error) {
	// Simplified reflection without BuildReflectionPrompt
	reflectionPrompt := fmt.Sprintf(`Evaluate the result of calling tool %s and decide the next action.

Tool returned: %s

Think:
1. Does the result satisfy the requirement?
2. What should be done next?`, toolName, result)

	messages := []chat.Message{
		{Role: "user", Content: reflectionPrompt},
	}

	// Generate a single ID for this entire reflection stream
	reflectionID := generateEventID("reflection")

	llmResult, err := e.streamLLMToEventBus(
		ctx,
		messages,
		&chat.ChatOptions{Temperature: 0.5},
		func(chunk *types.StreamResponse, fullContent string) {
			if chunk.Content != "" {
				e.eventBus.Emit(ctx, event.Event{
					ID:        reflectionID,
					Type:      event.EventAgentReflection,
					SessionID: sessionID,
					Data: event.AgentReflectionData{
						ToolCallID: toolCallID,
						Content:    chunk.Content,
						Iteration:  iteration,
						Done:       chunk.Done,
					},
				})
			}
		},
	)
	if err != nil {
		logger.Warnf(ctx, "Reflection failed: %v", err)
		return "", err
	}

	return llmResult.Content, nil
}

// streamThinkingToEventBus streams the thinking process through EventBus
func (e *AgentEngine) streamThinkingToEventBus(
	ctx context.Context,
	messages []chat.Message,
	tools []chat.Tool,
	iteration int,
	sessionID string,
) (*types.ChatResponse, error) {
	logger.Debugf(ctx, "[Agent][Thinking] Iteration-%d: temp=%.2f, tools=%d, thinking=%v",
		iteration+1, e.config.Temperature, len(tools), e.config.Thinking)

	opts := &chat.ChatOptions{
		Temperature: e.config.Temperature,
		Tools:       tools,
		Thinking:    e.config.Thinking,
	}

	pendingToolCalls := make(map[string]bool)
	thinkingToolIDs := make(map[string]string) // tool_call_id -> event ID for thinking tool streams

	// Generate IDs for this stream
	thinkingID := generateEventID("thinking")
	answerID := generateEventID("answer")

	llmResult, err := e.streamLLMToEventBus(
		ctx,
		messages,
		opts,
		func(chunk *types.StreamResponse, fullContent string) {
			if chunk.ResponseType == types.ResponseTypeToolCall && chunk.Data != nil {
				toolCallID, _ := chunk.Data["tool_call_id"].(string)
				toolName, _ := chunk.Data["tool_name"].(string)

				if toolCallID != "" && toolName != "" && !pendingToolCalls[toolCallID] {
					pendingToolCalls[toolCallID] = true
					e.eventBus.Emit(ctx, event.Event{
						ID:        fmt.Sprintf("%s-tool-call-pending", toolCallID),
						Type:      event.EventAgentToolCall,
						SessionID: sessionID,
						Data: event.AgentToolCallData{
							ToolCallID: toolCallID,
							ToolName:   toolName,
							Iteration:  iteration,
						},
					})
				}
			}

			// Handle final_answer tool's streaming answer content
			if chunk.ResponseType == types.ResponseTypeAnswer {
				if source, _ := chunk.Data["source"].(string); source == "final_answer_tool" {
					e.eventBus.Emit(ctx, event.Event{
						ID:        answerID,
						Type:      event.EventAgentFinalAnswer,
						SessionID: sessionID,
						Data: event.AgentFinalAnswerData{
							Content: chunk.Content,
							Done:    false,
						},
					})
					return
				}
			}

			// Handle thinking tool's streaming thought content
			if chunk.ResponseType == types.ResponseTypeThinking && chunk.Data != nil {
				if source, _ := chunk.Data["source"].(string); source == "thinking_tool" {
					toolCallID, _ := chunk.Data["tool_call_id"].(string)
					eventID, exists := thinkingToolIDs[toolCallID]
					if !exists {
						eventID = generateEventID("thinking-tool")
						thinkingToolIDs[toolCallID] = eventID
					}
					e.eventBus.Emit(ctx, event.Event{
						ID:        eventID,
						Type:      event.EventAgentThought,
						SessionID: sessionID,
						Data: event.AgentThoughtData{
							Content:   chunk.Content,
							Iteration: iteration,
							Done:      false,
						},
					})
					return
				}
			}

			if chunk.Content != "" {
				e.eventBus.Emit(ctx, event.Event{
					ID:        thinkingID, // Same ID for all chunks in this stream
					Type:      event.EventAgentThought,
					SessionID: sessionID,
					Data: event.AgentThoughtData{
						Content:   chunk.Content,
						Iteration: iteration,
						Done:      chunk.Done,
					},
				})
			}
		},
	)
	if err != nil {
		logger.Errorf(ctx, "[Agent][Thinking] Iteration-%d failed: %v", iteration+1, err)
		return nil, err
	}

	logger.Debugf(ctx, "[Agent][Thinking] Iteration-%d completed: content=%d chars, tool_calls=%d",
		iteration+1, len(llmResult.Content), len(llmResult.ToolCalls))

	fullContent := agenttools.StripThinkBlocks(llmResult.Content)

	resp := &types.ChatResponse{
		Content:      fullContent,
		ToolCalls:    llmResult.ToolCalls,
		FinishReason: "stop",
	}
	if llmResult.Usage != nil {
		resp.Usage = *llmResult.Usage
	}
	return resp, nil
}

// callLLMWithRetry logs messages, sanitizes them, calls the LLM with retry on transient errors,
// and handles graceful degradation when prior tool results exist.
// Returns nil response (with state.IsComplete=true) when graceful degradation succeeds.
// Returns a non-nil error only when the call fails irrecoverably.
func (e *AgentEngine) callLLMWithRetry(
	ctx context.Context, messages []chat.Message, tools []chat.Tool,
	state *types.AgentState, query string, iteration int, sessionID string,
) (*types.ChatResponse, error) {
	round := iteration + 1

	// Log message summary at Debug level (per-message detail is noisy at Info)
	logger.Infof(ctx, "[Agent][Round-%d] Calling LLM: %d messages, %d tools",
		round, len(messages), len(tools))
	for i, msg := range messages {
		if msg.Role == "tool" {
			logger.Debugf(ctx, "[Agent][Round-%d] msg[%d]: role=tool, name=%s, len=%d",
				round, i, msg.Name, len(msg.Content))
		} else if len(msg.ToolCalls) > 0 {
			tcNames := make([]string, len(msg.ToolCalls))
			for j, tc := range msg.ToolCalls {
				tcNames[j] = tc.Function.Name
			}
			logger.Debugf(ctx, "[Agent][Round-%d] msg[%d]: role=%s, len=%d, tool_calls=%v",
				round, i, msg.Role, len(msg.Content), tcNames)
		} else {
			preview := msg.Content
			if len(preview) > 100 {
				preview = preview[:100] + "..."
			}
			logger.Debugf(ctx, "[Agent][Round-%d] msg[%d]: role=%s, len=%d, content=%s",
				round, i, msg.Role, len(msg.Content), preview)
		}
	}
	common.PipelineInfo(ctx, "Agent", "think_start", map[string]interface{}{
		"iteration": iteration,
		"round":     round,
		"tool_cnt":  len(tools),
	})

	// Sanitize messages before sending to LLM (fix consecutive roles, orphaned tool results)
	messages = agenttools.SanitizeMessages(messages)

	response, err := e.streamThinkingToEventBus(ctx, messages, tools, iteration, sessionID)
	if err != nil && isTransientError(err) {
		// Retry transient errors (timeout, rate limit, server errors) up to maxLLMRetries times
		for retry := 1; retry <= maxLLMRetries; retry++ {
			retryDelay := time.Duration(retry) * time.Second
			logger.Warnf(ctx, "[Agent][Round-%d] LLM transient error (attempt %d/%d), retrying in %v: %v",
				round, retry, maxLLMRetries, retryDelay, err)
			time.Sleep(retryDelay)

			response, err = e.streamThinkingToEventBus(ctx, messages, tools, iteration, sessionID)
			if err == nil || !isTransientError(err) {
				break
			}
		}
	}
	if err != nil {
		logger.Errorf(ctx, "[Agent][Round-%d] LLM call failed: %v", round, err)
		common.PipelineError(ctx, "Agent", "think_failed", map[string]interface{}{
			"iteration": iteration,
			"error":     err.Error(),
		})

		// Graceful degradation: if we have tool results from previous rounds,
		// try to synthesize a final answer from them instead of losing everything.
		if totalTC := countTotalToolCalls(state.RoundSteps); totalTC > 0 {
			logger.Warnf(ctx, "[Agent] LLM failed but have %d steps with %d tool calls — "+
				"attempting final answer synthesis from existing results",
				len(state.RoundSteps), totalTC)
			common.PipelineWarn(ctx, "Agent", "llm_failed_synthesizing", map[string]interface{}{
				"steps":      len(state.RoundSteps),
				"tool_calls": totalTC,
			})
			if synthErr := e.streamFinalAnswerToEventBus(ctx, query, state, sessionID); synthErr != nil {
				logger.Errorf(ctx, "[Agent] Final answer synthesis also failed: %v", synthErr)
				return nil, fmt.Errorf("LLM call failed: %w (synthesis also failed: %v)", err, synthErr)
			}
			state.IsComplete = true
			return nil, nil // graceful degradation succeeded
		}

		return nil, fmt.Errorf("LLM call failed: %w", err)
	}

	common.PipelineInfo(ctx, "Agent", "think_result", map[string]interface{}{
		"iteration":     iteration,
		"finish_reason": response.FinishReason,
		"tool_calls":    len(response.ToolCalls),
		"content_len":   len(response.Content),
	})

	// Log LLM response summary
	if len(response.ToolCalls) > 0 {
		tcNames := make([]string, len(response.ToolCalls))
		for i, tc := range response.ToolCalls {
			tcNames[i] = tc.Function.Name
		}
		logger.Infof(ctx, "[Agent][Round-%d] LLM responded: finish=%s, content=%d chars, tools=%v",
			round, response.FinishReason, len(response.Content), tcNames)
	} else {
		logger.Infof(ctx, "[Agent][Round-%d] LLM responded: finish=%s, content=%d chars",
			round, response.FinishReason, len(response.Content))
	}
	if response.Content != "" {
		preview := response.Content
		if len(preview) > 300 {
			preview = preview[:300] + "..."
		}
		logger.Debugf(ctx, "[Agent][Round-%d] LLM content preview:\n%s", round, preview)
	}

	return response, nil
}