aiops/backend/internal/api/handlers/llm_chat.go

package handlers

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"ezops/internal/api/middleware"
	"ezops/internal/models"

	"github.com/gin-gonic/gin"
)

// Chat AI 对话代理（OpenAI 兼容接口） // Logan.Li
// POST /api/v1/aiops/chat
func (h *LLMHandler) Chat(c *gin.Context) {
	var cfg models.LLMConfig
	if err := h.db.First(&cfg).Error; err != nil || !cfg.Enabled {
		c.JSON(http.StatusOK, models.Response{Code: 400, Message: "LLM 服务未配置或未启用，请在「模型配置」中完成设置"})
		return
	}

	var req models.ChatRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, models.Response{Code: 400, Message: "请求参数错误: " + err.Error()})
		return
	}

	// 组装系统提示词
	systemPrompt := cfg.SystemPrompt
	if strings.TrimSpace(systemPrompt) == "" {
		systemPrompt = defaultSystemPrompt
	}
	messages := append([]models.ChatMessage{{Role: "system", Content: systemPrompt}}, req.Messages...)

	// 构造 OpenAI 兼容请求
	payload := openaiChatRequest{
		Model:       cfg.Model,
		Messages:    messages,
		Temperature: cfg.Temperature,
		MaxTokens:   cfg.MaxTokens,
	}
	body, _ := json.Marshal(payload)

	endpoint := resolveEndpoint(cfg.Provider, cfg.Endpoint)
	url := endpoint + "/chat/completions"

	httpReq, err := http.NewRequestWithContext(c.Request.Context(), http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		c.JSON(http.StatusOK, models.Response{Code: 500, Message: "构造请求失败: " + err.Error()})
		return
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)

	resp, err := h.client.Do(httpReq)
	if err != nil {
		c.JSON(http.StatusOK, models.Response{Code: 502, Message: "LLM 服务连接失败，请检查网络或 Endpoint 配置: " + err.Error()})
		return
	}
	defer resp.Body.Close()

	raw, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusOK {
		// 提取上游错误信息
		msg := fmt.Sprintf("LLM 服务返回错误 (HTTP %d)", resp.StatusCode)
		if upErr := parseOpenAIError(raw); upErr != "" {
			msg += ": " + upErr
		}
		c.JSON(http.StatusOK, models.Response{Code: resp.StatusCode, Message: msg})
		return
	}

	var chatResp openaiChatResponse
	if err := json.Unmarshal(raw, &chatResp); err != nil || len(chatResp.Choices) == 0 {
		c.JSON(http.StatusOK, models.Response{Code: 500, Message: "解析 LLM 响应失败"})
		return
	}
	replyContent := chatResp.Choices[0].Message.Content

	// 持久化对话历史（按用户隔离，成对写入 user + assistant） // Logan.Li
	// 仅在成功时入库；错误气泡为瞬时 UI 不持久化
	uid := middleware.GetUserID(c)
	if uid != 0 && len(req.Messages) > 0 {
		// 取最后一条用户消息（前端已附加当前提问）
		lastUser := req.Messages[len(req.Messages)-1]
		now := time.Now()
		// session_id：前端传入或后端生成（兼容旧版客户端） // Logan.Li
		sid := req.SessionID
		if sid == "" {
			sid = fmt.Sprintf("sess-%d", now.UnixMilli())
		}
		h.db.Create(&[]models.ChatHistory{
			{UserID: uid, SessionID: sid, Role: "user", Content: lastUser.Content, CreatedAt: now},
			{UserID: uid, SessionID: sid, Role: "assistant", Content: replyContent, Model: cfg.Model, CreatedAt: now.Add(time.Millisecond)},
		})
	}

	c.JSON(http.StatusOK, models.Response{Code: 200, Data: models.ChatResponse{
		Content:    replyContent,
		Model:      cfg.Model,
		TokensUsed: chatResp.Usage.TotalTokens,
	}})
}

// ChatStream AI 流式对话代理（SSE） // Logan.Li
// POST /api/v1/aiops/chat/stream
// 客户端用 EventSource/fetch-stream 消费，每条 SSE data 为 {"content": "..."}
// 出错时发送 event: error，data 为 {"message": "..."}
// 完成后服务端持久化 user + assistant 两条历史，event: done 结束
func (h *LLMHandler) ChatStream(c *gin.Context) {
	var cfg models.LLMConfig
	if err := h.db.First(&cfg).Error; err != nil || !cfg.Enabled {
		writeSSEError(c, 400, "LLM 服务未配置或未启用，请在「模型配置」中完成设置")
		return
	}

	var req models.ChatRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		writeSSEError(c, 400, "请求参数错误: "+err.Error())
		return
	}
	if len(req.Messages) == 0 {
		writeSSEError(c, 400, "消息不能为空")
		return
	}

	// 设置 SSE 响应头 // Logan.Li
	c.Writer.Header().Set("Content-Type", "text/event-stream")
	c.Writer.Header().Set("Cache-Control", "no-cache")
	c.Writer.Header().Set("Connection", "keep-alive")
	c.Writer.Header().Set("X-Accel-Buffering", "no") // Nginx 关闭缓冲
	c.Writer.WriteHeader(http.StatusOK)
	// 注册写入心跳（立即 flush 一个注释行让客户端确认连接建立）
	c.Writer.Flush()

	// 组装系统提示词 + 上下文
	systemPrompt := cfg.SystemPrompt
	if strings.TrimSpace(systemPrompt) == "" {
		systemPrompt = defaultSystemPrompt
	}
	messages := append([]models.ChatMessage{{Role: "system", Content: systemPrompt}}, req.Messages...)
	// 构造 OpenAI 兼容流式请求
	payload := struct {
		Model       string               `json:"model"`
		Messages    []models.ChatMessage `json:"messages"`
		Temperature float64              `json:"temperature"`
		MaxTokens   int                  `json:"max_tokens,omitempty"`
		Stream      bool                 `json:"stream"`
	}{
		Model:       cfg.Model,
		Messages:    messages,
		Temperature: cfg.Temperature,
		MaxTokens:   cfg.MaxTokens,
		Stream:      true,
	}
	body, _ := json.Marshal(payload)

	endpoint := resolveEndpoint(cfg.Provider, cfg.Endpoint)
	url := endpoint + "/chat/completions"

	httpReq, err := http.NewRequestWithContext(c.Request.Context(), http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		writeSSEError(c, 500, "构造请求失败: "+err.Error())
		return
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	httpReq.Header.Set("Accept", "text/event-stream")

	// 流式请求不能用全局 120s 超时客户端，构造无超时 client 让 ctx 控制 // Logan.Li
	streamClient := &http.Client{Timeout: 0}
	resp, err := streamClient.Do(httpReq)
	if err != nil {
		writeSSEError(c, 502, "LLM 服务连接失败，请检查网络或 Endpoint 配置: "+err.Error())
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(resp.Body)
		msg := fmt.Sprintf("LLM 服务返回错误 (HTTP %d)", resp.StatusCode)
		if e := parseOpenAIError(raw); e != "" {
			msg += ": " + e
		}
		writeSSEError(c, resp.StatusCode, msg)
		return
	}

	// 逐行解析上游 SSE，转发增量内容给前端 // Logan.Li
	var contentBuilder strings.Builder
	reader := bufio.NewReader(resp.Body)
	clientGone := false
	for {
		// 客户端断开时停止读取（c.Request.Context().Done() 在客户端断开时触发） // Logan.Li
		select {
		case <-c.Request.Context().Done():
			clientGone = true
		default:
		}
		if clientGone {
			break
		}

		line, readErr := reader.ReadString('\n')
		line = strings.TrimSpace(line)
		if line != "" {
			// 仅处理 data: 开头的行
			if strings.HasPrefix(line, "data:") {
				data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
				if data == "[DONE]" {
					// 上游结束，跳出循环准备持久化 + 发送 [DONE]
					break
				}
				// 解析 {"choices":[{"delta":{"content":"..."}}]}
				var chunk struct {
					Choices []struct {
						Delta struct {
							Content string `json:"content"`
						} `json:"delta"`
					} `json:"choices"`
				}
				if json.Unmarshal([]byte(data), &chunk) == nil && len(chunk.Choices) > 0 {
					delta := chunk.Choices[0].Delta.Content
					if delta != "" {
						contentBuilder.WriteString(delta)
						// 转发给前端：data: {"content":"..."}
						writeSSEChunk(c, delta)
					}
				}
			}
		}
		if readErr != nil {
			if readErr != io.EOF {
				// 读取出错（非正常结束）
				if contentBuilder.Len() == 0 {
					writeSSEError(c, 502, "LLM 流式响应读取失败: "+readErr.Error())
					return
				}
				// 已有部分内容则正常结束
			}
			break
		}
	}

	fullContent := contentBuilder.String()

	// 持久化对话历史（仅在成功且非空时入库；错误/中断不入库避免污染历史） // Logan.Li
	if fullContent != "" && !clientGone {
		uid := middleware.GetUserID(c)
		if uid != 0 {
			lastUser := req.Messages[len(req.Messages)-1]
			now := time.Now()
			// session_id：前端传入或后端生成（兼容旧版客户端） // Logan.Li
			sid := req.SessionID
			if sid == "" {
				sid = fmt.Sprintf("sess-%d", now.UnixMilli())
			}
			h.db.Create(&[]models.ChatHistory{
				{UserID: uid, SessionID: sid, Role: "user", Content: lastUser.Content, CreatedAt: now},
				{UserID: uid, SessionID: sid, Role: "assistant", Content: fullContent, Model: cfg.Model, CreatedAt: now.Add(time.Millisecond)},
			})
		}
		// 发送结束标志
		c.Writer.WriteString("event: done\ndata: {\"model\":\"" + cfg.Model + "\"}\n\n")
		c.Writer.Flush()
	} else if !clientGone {
		// 无内容时通知前端结束（避免前端无限等待）
		c.Writer.WriteString("event: done\ndata: {}\n\n")
		c.Writer.Flush()
	}
}

// writeSSEChunk 写入一条内容增量 // Logan.Li
func writeSSEChunk(c *gin.Context, content string) {
	// 转义为 JSON 字符串避免换行符破坏 SSE 协议
	b, _ := json.Marshal(map[string]string{"content": content})
	c.Writer.WriteString("data: " + string(b) + "\n\n")
	c.Writer.Flush()
}

// writeSSEError 发送一个 error 事件（写入后立即结束流） // Logan.Li
func writeSSEError(c *gin.Context, code int, message string) {
	// 若 header 尚未写入（连接前错误），需先设置 SSE 头
	if !c.Writer.Written() {
		c.Writer.Header().Set("Content-Type", "text/event-stream")
		c.Writer.Header().Set("Cache-Control", "no-cache")
		c.Writer.Header().Set("Connection", "keep-alive")
		c.Writer.WriteHeader(http.StatusOK)
	}
	b, _ := json.Marshal(map[string]any{"code": code, "message": message})
	c.Writer.WriteString("event: error\ndata: " + string(b) + "\n\n")
	c.Writer.Flush()
}

// parseOpenAIError 从上游错误响应中提取 message 字段 // Logan.Li
func parseOpenAIError(raw []byte) string {
	var m struct {
		Error struct {
			Message string `json:"message"`
		} `json:"error"`
	}
	if json.Unmarshal(raw, &m) == nil && m.Error.Message != "" {
		return m.Error.Message
	}
	return ""
}

// TestLLMConnection 测试 LLM 连接（发送一条 ping 消息） // Logan.Li
// POST /api/v1/aiops/llm-config/test
func (h *LLMHandler) TestLLMConnection(c *gin.Context) {
	var cfg models.LLMConfig
	if err := h.db.First(&cfg).Error; err != nil {
		c.JSON(http.StatusOK, models.Response{Code: 400, Message: "请先保存配置"})
		return
	}
	payload := openaiChatRequest{
		Model:       cfg.Model,
		Messages:    []models.ChatMessage{{Role: "user", Content: "ping"}},
		Temperature: 0,
		MaxTokens:   16,
	}
	body, _ := json.Marshal(payload)
	endpoint := resolveEndpoint(cfg.Provider, cfg.Endpoint)
	httpReq, _ := http.NewRequestWithContext(c.Request.Context(), http.MethodPost, endpoint+"/chat/completions", bytes.NewReader(body))
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	resp, err := h.client.Do(httpReq)
	if err != nil {
		c.JSON(http.StatusOK, models.Response{Code: 502, Message: "连接失败: " + err.Error()})
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(resp.Body)
		msg := fmt.Sprintf("HTTP %d", resp.StatusCode)
		if e := parseOpenAIError(raw); e != "" {
			msg += ": " + e
		}
		c.JSON(http.StatusOK, models.Response{Code: resp.StatusCode, Message: msg})
		return
	}
	c.JSON(http.StatusOK, models.Response{Code: 200, Message: "连接成功"})
}