LokiStack/backend/services/ai_provider.go

package services

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"lokistack/repository"
	"net/http"
	"strings"
	"time"

	"lokistack/models"
)

// OpenAI 兼容 Chat Completions 客户端。
//
// 只依赖 /chat/completions 这一个端点，所以 DeepSeek、通义、
// 内网 vLLM、Ollama(/v1)、各类网关都能直接用，不需要按厂商写适配。

type aiMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type aiChatRequest struct {
	Model       string      `json:"model"`
	Messages    []aiMessage `json:"messages"`
	Temperature float64     `json:"temperature"`
	MaxTokens   int         `json:"max_tokens,omitempty"`
	Stream      bool        `json:"stream"`
}

type aiChatResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
	Error *struct {
		Message string `json:"message"`
		Type    string `json:"type"`
	} `json:"error"`
}

// aiChat 调一次对话补全。
//
// 错误信息必须具体：把上游返回的 message 原样带出来，不要包装成
// "调用失败" —— 常见原因（key 错、模型名不存在、baseURL 路径少了 /v1、
// 内网模型没起）都在上游 message 里，藏起来只会让用户去猜。
func aiChat(ctx context.Context, cfg models.AIConfig, messages []aiMessage) (string, int, int, error) {
	if strings.TrimSpace(cfg.BaseURL) == "" || strings.TrimSpace(cfg.Model) == "" {
		return "", 0, 0, fmt.Errorf("未配置模型：请先在「智能运维 → 模型配置」填写 Base URL 与模型名")
	}
	url := strings.TrimSuffix(cfg.BaseURL, "/") + "/chat/completions"
	body, err := json.Marshal(aiChatRequest{
		Model:       cfg.Model,
		Messages:    messages,
		Temperature: cfg.Temperature,
		MaxTokens:   cfg.MaxTokens,
		Stream:      false,
	})
	if err != nil {
		return "", 0, 0, err
	}

	timeout := time.Duration(cfg.TimeoutSec) * time.Second
	if timeout <= 0 {
		timeout = 120 * time.Second
	}
	reqCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return "", 0, 0, fmt.Errorf("构造请求失败: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	if cfg.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		if reqCtx.Err() == context.DeadlineExceeded {
			return "", 0, 0, fmt.Errorf("模型响应超时（%d 秒）：可调大「超时」或换更快的模型", cfg.TimeoutSec)
		}
		return "", 0, 0, fmt.Errorf("连接模型服务失败: %w", err)
	}
	defer resp.Body.Close()

	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode >= 400 {
		var parsed aiChatResponse
		if json.Unmarshal(raw, &parsed) == nil && parsed.Error != nil && parsed.Error.Message != "" {
			return "", 0, 0, fmt.Errorf("模型返回 %d：%s", resp.StatusCode, parsed.Error.Message)
		}
		return "", 0, 0, fmt.Errorf("模型返回 %d：%s", resp.StatusCode, truncateStr(string(raw), 300))
	}

	var parsed aiChatResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return "", 0, 0, fmt.Errorf("解析模型响应失败（可能不是 OpenAI 兼容接口）: %s", truncateStr(string(raw), 200))
	}
	if len(parsed.Choices) == 0 {
		return "", 0, 0, fmt.Errorf("模型没有返回任何内容（finish_reason 为空）")
	}
	return parsed.Choices[0].Message.Content, parsed.Usage.PromptTokens, parsed.Usage.CompletionTokens, nil
}

// TestAIConnection 连通性测试：发一句最小请求，验证 baseURL/key/model 三件套。
func TestAIConnection(cfg models.AIConfig) (string, error) {
	cfg.MaxTokens = 32
	cfg.TimeoutSec = 30
	reply, _, _, err := aiChat(context.Background(), cfg, []aiMessage{
		{Role: "user", Content: "只回复两个字：正常"},
	})
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(reply), nil
}

// ChatMessage 浮窗助手的多轮消息(user/assistant 交替)
type ChatMessage struct {
	Role    string `json:"role" binding:"required,oneof=user assistant"`
	Content string `json:"content" binding:"required"`
}

// RunAIChat 全站 AI 助手浮窗:多轮对话 + 当前页面上下文注入(轻量信号摘要)。
// 与整页分析共享配置/配额/脱敏/Pro 闸门;context 里的 sourceId 可选,带了就注入日志摘要。
func RunAIChat(ctx context.Context, cfg models.AIConfig, src *models.LogSource, userID uint, admin bool, history []ChatMessage, page string, rangeMinutes int) (string, error) {
	if !cfg.Enabled {
		return "", fmt.Errorf("AI 未启用:请先在「智能运维 → 模型配置」中开启并保存")
	}
	if cfg.DailyTokenCap > 0 && AITokenUsedToday(userID) >= cfg.DailyTokenCap {
		return "", fmt.Errorf("今日 AI 用量已达上限,明天再试或调整上限")
	}

	system := "你是 LokiStack 日志平台的运维助手。回答简明、面向排查行动;不确定就说不确定。"
	if src != nil && rangeMinutes > 0 {
		end := time.Now()
		from := end.Add(-time.Duration(rangeMinutes) * time.Minute)
		if bundle, err := CollectAISignals(ctx, src, userID, admin, from, end, ""); err == nil {
			system += "\n\n当前用户正在查看 " + page + " 页,近 " + fmt.Sprintf("%d", rangeMinutes) + " 分钟的日志概况:总量 " +
				fmt.Sprintf("%d", bundle.TotalEntries) + " 条(估算)"
			if len(bundle.Apps) > 0 {
				system += ",Top 应用:"
				for i, a := range bundle.Apps {
					if i >= 3 {
						break
					}
					system += fmt.Sprintf(" %s(%d)", a.App, a.Entries)
				}
			}
			if len(bundle.Alerts) > 0 {
				system += ";存在告警命中"
			}
		}
	}

	msgs := []aiMessage{{Role: "system", Content: system}}
	for _, m := range history {
		msgs = append(msgs, aiMessage{Role: m.Role, Content: m.Content})
	}
	// 脱敏:用户消息里的敏感串同样替换
	if cfg.RedactExtra != nil {
		for i := range msgs {
			msgs[i].Content = RedactTextOnly(msgs[i].Content, cfg.RedactExtra)
		}
	}
	reply, promptTokens, completionTokens, err := aiChat(ctx, cfg, msgs)
	if err != nil {
		return "", err
	}
	// 计入每日配额:落一条 chat 类型用量记录
	db := repository.GetDB()
	if db != nil {
		db.Lock()
		db.AIAnalysisRuns = append(db.AIAnalysisRuns, models.AIAnalysisRun{
			UserID: userID, Username: "chat", SourceID: 0, SourceName: page,
			Status: "chat", PromptTokens: promptTokens, CompletionTokens: completionTokens,
			CreatedAt: time.Now(),
		})
		db.RequestSave()
		db.Unlock()
	}
	return reply, nil
}