340 lines
11 KiB
Go
340 lines
11 KiB
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()
|
||
h.db.Create(&[]models.ChatHistory{
|
||
{UserID: uid, Role: "user", Content: lastUser.Content, CreatedAt: now},
|
||
{UserID: uid, 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()
|
||
h.db.Create(&[]models.ChatHistory{
|
||
{UserID: uid, Role: "user", Content: lastUser.Content, CreatedAt: now},
|
||
{UserID: uid, 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: "连接成功"})
|
||
}
|