aiops/backend/controller/llm_models.go

package controller

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"ezops/model"

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

const maxLLMModelListBytes = 2 << 20

type llmModelListItem struct {
	ID    string `json:"id"`
	Name  string `json:"name"`
	Model string `json:"model"`
}

// resolveModelListCredential 解析草稿配置，并在前端回传脱敏值时复用已保存密钥。 // Logan.Li
func (h *LLMHandler) resolveModelListCredential(req model.SaveLLMConfigRequest) (string, string, string) {
	provider := strings.TrimSpace(req.Provider)
	if provider == "" {
		return "", "", "请选择提供商"
	}
	endpoint := resolveEndpoint(provider, strings.TrimSpace(req.Endpoint))
	if endpoint == "" {
		return "", "", "请填写 Endpoint"
	}
	apiKey := strings.TrimSpace(req.APIKey)
	if strings.Contains(apiKey, "****") {
		var saved model.LLMConfig
		if err := h.db.Take(&saved).Error; err != nil || !canReuseMaskedAPIKey(provider, req.Endpoint, saved) {
			return "", "", "提供商或 Endpoint 已变化，请重新填写 API Key"
		}
		apiKey = saved.APIKey
	}
	if apiKey == "" && !isLocalLLMProvider(provider) {
		return "", "", "请填写 API Key"
	}
	return endpoint, apiKey, ""
}

// buildLLMModelListURL 将基础地址转换为兼容协议的模型列表地址。 // Logan.Li
func buildLLMModelListURL(provider, endpoint string) (string, error) {
	base := strings.TrimRight(strings.TrimSpace(endpoint), "/")
	if resolveProtocol(provider) == "anthropic" && !strings.HasSuffix(base, "/v1") {
		base += "/v1"
	}
	parsed, err := url.Parse(base + "/models")
	if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") {
		return "", fmt.Errorf("Endpoint 必须是有效的 HTTP 或 HTTPS 地址")
	}
	return parsed.String(), nil
}

// decodeLLMModelItems 兼容 OpenAI/New API 的 data 数组及部分本地服务的 models 数组。 // Logan.Li
func decodeLLMModelItems(raw json.RawMessage) []string {
	var items []llmModelListItem
	if err := json.Unmarshal(raw, &items); err == nil {
		models := make([]string, 0, len(items))
		for _, item := range items {
			name := strings.TrimSpace(item.ID)
			if name == "" {
				name = strings.TrimSpace(item.Name)
			}
			if name == "" {
				name = strings.TrimSpace(item.Model)
			}
			if name != "" {
				models = append(models, name)
			}
		}
		return models
	}
	var names []string
	_ = json.Unmarshal(raw, &names)
	return names
}

// parseLLMModelList 提取、去重并排序上游返回的模型 ID。 // Logan.Li
func parseLLMModelList(raw []byte) ([]string, error) {
	var payload map[string]json.RawMessage
	if err := json.Unmarshal(raw, &payload); err != nil {
		return nil, fmt.Errorf("解析模型列表失败: %w", err)
	}
	models := decodeLLMModelItems(payload["data"])
	models = append(models, decodeLLMModelItems(payload["models"])...)
	unique := make(map[string]struct{}, len(models))
	result := make([]string, 0, len(models))
	for _, item := range models {
		name := strings.TrimSpace(item)
		if name == "" {
			continue
		}
		if _, exists := unique[name]; exists {
			continue
		}
		unique[name] = struct{}{}
		result = append(result, name)
	}
	sort.Strings(result)
	return result, nil
}

// requestLLMModelList 请求上游模型目录，禁止跨域重定向以免凭证被转发。 // Logan.Li
func requestLLMModelList(ctx context.Context, provider, endpoint, apiKey string) ([]string, error) {
	requestURL, err := buildLLMModelListURL(provider, endpoint)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil)
	if err != nil {
		return nil, fmt.Errorf("构造模型列表请求失败: %w", err)
	}
	req.Header.Set("Accept", "application/json")
	if apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}
	if resolveProtocol(provider) == "anthropic" {
		req.Header.Set("x-api-key", apiKey)
		req.Header.Set("anthropic-version", "2023-06-01")
	}
	client := &http.Client{
		Timeout: 20 * time.Second,
		CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("读取模型列表失败: %w", err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, maxLLMModelListBytes+1))
	if err != nil {
		return nil, fmt.Errorf("读取模型列表响应失败: %w", err)
	}
	if len(raw) > maxLLMModelListBytes {
		return nil, fmt.Errorf("模型列表响应超过 2 MiB 限制")
	}
	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
		message := []rune(strings.TrimSpace(string(raw)))
		if len(message) > 1000 {
			message = message[:1000]
		}
		return nil, fmt.Errorf("模型服务返回 HTTP %d: %s", resp.StatusCode, string(message))
	}
	return parseLLMModelList(raw)
}

// ListLLMModels 从当前草稿 Endpoint 读取可用模型，便于 New API/One API 和本地服务配置。 // Logan.Li
// POST /api/v1/aiops/llm-config/models // Logan.Li
func (h *LLMHandler) ListLLMModels(c *gin.Context) {
	var req model.SaveLLMConfigRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "请求参数错误: " + err.Error()})
		return
	}
	endpoint, apiKey, errMsg := h.resolveModelListCredential(req)
	if errMsg != "" {
		c.JSON(http.StatusOK, model.Response{Code: 400, Message: errMsg})
		return
	}
	models, err := requestLLMModelList(c.Request.Context(), req.Provider, endpoint, apiKey)
	if err != nil {
		c.JSON(http.StatusOK, model.Response{Code: 502, Message: err.Error()})
		return
	}
	c.JSON(http.StatusOK, model.Response{Code: 200, Data: gin.H{
		"models": models,
		"total":  len(models),
	}})
}