Files

2601 lines
85 KiB
Go
Raw Permalink Normal View History

package handler
import (
"context"
"encoding/json"
stderrors "errors"
"fmt"
"io"
"math/rand"
"net/http"
"os"
"strconv"
"strings"
"sync"
"time"
"github.com/Tencent/WeKnora/internal/application/repository"
chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline"
"github.com/Tencent/WeKnora/internal/assets"
"github.com/Tencent/WeKnora/internal/config"
"github.com/Tencent/WeKnora/internal/errors"
"github.com/Tencent/WeKnora/internal/handler/dto"
"github.com/Tencent/WeKnora/internal/logger"
"github.com/Tencent/WeKnora/internal/models/asr"
"github.com/Tencent/WeKnora/internal/models/chat"
"github.com/Tencent/WeKnora/internal/models/embedding"
"github.com/Tencent/WeKnora/internal/models/provider"
"github.com/Tencent/WeKnora/internal/models/rerank"
"github.com/Tencent/WeKnora/internal/models/utils/ollama"
"github.com/Tencent/WeKnora/internal/types"
"github.com/Tencent/WeKnora/internal/types/interfaces"
"github.com/Tencent/WeKnora/internal/utils"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"github.com/ollama/ollama/api"
)
// DownloadTask 下载任务信息
type DownloadTask struct {
ID string `json:"id"`
ModelName string `json:"modelName"`
Status string `json:"status"` // pending, downloading, completed, failed
Progress float64 `json:"progress"`
Message string `json:"message"`
StartTime time.Time `json:"startTime"`
EndTime *time.Time `json:"endTime,omitempty"`
}
// 全局下载任务管理器
var (
downloadTasks = make(map[string]*DownloadTask)
tasksMutex sync.RWMutex
)
// InitializationHandler 初始化处理器
type InitializationHandler struct {
config *config.Config
tenantService interfaces.TenantService
modelService interfaces.ModelService
kbService interfaces.KnowledgeBaseService
kbRepository interfaces.KnowledgeBaseRepository
knowledgeService interfaces.KnowledgeService
ollamaService *ollama.OllamaService
documentReader interfaces.DocumentReader
pooler embedding.EmbedderPooler
storageResolver interfaces.StorageBackendResolver
}
// NewInitializationHandler 创建初始化处理器
func NewInitializationHandler(
config *config.Config,
tenantService interfaces.TenantService,
modelService interfaces.ModelService,
kbService interfaces.KnowledgeBaseService,
kbRepository interfaces.KnowledgeBaseRepository,
knowledgeService interfaces.KnowledgeService,
ollamaService *ollama.OllamaService,
documentReader interfaces.DocumentReader,
pooler embedding.EmbedderPooler,
storageResolver interfaces.StorageBackendResolver,
) *InitializationHandler {
return &InitializationHandler{
config: config,
tenantService: tenantService,
modelService: modelService,
kbService: kbService,
kbRepository: kbRepository,
knowledgeService: knowledgeService,
ollamaService: ollamaService,
documentReader: documentReader,
pooler: pooler,
storageResolver: storageResolver,
}
}
// KBModelConfigRequest 知识库模型配置请求(简化版,只传模型ID)
type KBModelConfigRequest struct {
LLMModelID string `json:"llmModelId" binding:"required"`
EmbeddingModelID string `json:"embeddingModelId"` // optional when RAG indexing is disabled
VLMConfig *types.VLMConfig `json:"vlm_config"`
ASRConfig *types.ASRConfig `json:"asr_config"`
// 文档分块配置
DocumentSplitting struct {
ChunkSize int `json:"chunkSize"`
ChunkOverlap int `json:"chunkOverlap"`
Separators []string `json:"separators"`
ParserEngineRules []types.ParserEngineRule `json:"parserEngineRules,omitempty"`
EnableParentChild bool `json:"enableParentChild"`
ParentChunkSize int `json:"parentChunkSize,omitempty"`
ChildChunkSize int `json:"childChunkSize,omitempty"`
// Strategy / TokenLimit / Languages use pointer types so the
// handler can distinguish "field absent in payload" (no change)
// from "field present with empty/zero value" (clear / disable).
// Without that distinction, users could set strategy="auto" once
// but never reset it back to legacy / unset.
Strategy *string `json:"strategy,omitempty"`
TokenLimit *int `json:"tokenLimit,omitempty"`
Languages *[]string `json:"languages,omitempty"`
TableMetadataInstructions *string `json:"tableMetadataInstructions,omitempty"`
} `json:"documentSplitting"`
// 多模态配置(仅模型相关;存储引擎在 storageProvider 中配置)
Multimodal struct {
Enabled bool `json:"enabled"`
} `json:"multimodal"`
// 存储引擎选择("local" | "minio" | "cos"),影响文档上传与文档内图片存储,参数从全局设置读取
StorageProvider string `json:"storageProvider"`
StorageBackendID string `json:"storageBackendId"`
// 知识图谱配置
NodeExtract struct {
Enabled bool `json:"enabled"`
Text string `json:"text"`
Tags []string `json:"tags"`
Nodes []types.GraphNode `json:"nodes"`
Relations []types.GraphRelation `json:"relations"`
CustomInstructions string `json:"customInstructions"`
} `json:"nodeExtract"`
// 问题生成配置
QuestionGeneration struct {
Enabled bool `json:"enabled"`
QuestionCount int `json:"questionCount"`
CustomInstructions string `json:"customInstructions"`
} `json:"questionGeneration"`
}
// InitializationRequest 初始化请求结构
type InitializationRequest struct {
LLM struct {
Source string `json:"source" binding:"required"`
ModelName string `json:"modelName" binding:"required"`
BaseURL string `json:"baseUrl"`
APIKey string `json:"apiKey"`
} `json:"llm" binding:"required"`
Embedding struct {
Source string `json:"source" binding:"required"`
ModelName string `json:"modelName" binding:"required"`
BaseURL string `json:"baseUrl"`
APIKey string `json:"apiKey"`
Dimension int `json:"dimension"` // 添加embedding维度字段
} `json:"embedding" binding:"required"`
Rerank struct {
Enabled bool `json:"enabled"`
ModelName string `json:"modelName"`
BaseURL string `json:"baseUrl"`
APIKey string `json:"apiKey"`
} `json:"rerank"`
Multimodal struct {
Enabled bool `json:"enabled"`
VLM *struct {
ModelName string `json:"modelName"`
BaseURL string `json:"baseUrl"`
APIKey string `json:"apiKey"`
InterfaceType string `json:"interfaceType"` // "ollama" or "openai"
} `json:"vlm,omitempty"`
StorageType string `json:"storageType"`
COS *struct {
SecretID string `json:"secretId"`
SecretKey string `json:"secretKey"`
Region string `json:"region"`
BucketName string `json:"bucketName"`
AppID string `json:"appId"`
PathPrefix string `json:"pathPrefix"`
} `json:"cos,omitempty"`
Minio *struct {
BucketName string `json:"bucketName"`
PathPrefix string `json:"pathPrefix"`
} `json:"minio,omitempty"`
} `json:"multimodal"`
DocumentSplitting struct {
ChunkSize int `json:"chunkSize" binding:"required,min=100,max=10000"`
ChunkOverlap int `json:"chunkOverlap" binding:"min=0"`
Separators []string `json:"separators" binding:"required,min=1"`
} `json:"documentSplitting" binding:"required"`
NodeExtract struct {
Enabled bool `json:"enabled"`
Text string `json:"text"`
Tags []string `json:"tags"`
Nodes []struct {
Name string `json:"name"`
Attributes []string `json:"attributes"`
} `json:"nodes"`
Relations []struct {
Node1 string `json:"node1"`
Node2 string `json:"node2"`
Type string `json:"type"`
} `json:"relations"`
} `json:"nodeExtract"`
QuestionGeneration struct {
Enabled bool `json:"enabled"`
QuestionCount int `json:"questionCount"`
} `json:"questionGeneration"`
}
// UpdateKBConfig godoc
// @Summary 更新知识库配置
// @Description 根据知识库ID更新模型和分块配置
// @Tags 初始化
// @Accept json
// @Produce json
// @Param kbId path string true "知识库ID"
// @Param request body KBModelConfigRequest true "配置请求"
// @Success 200 {object} map[string]interface{} "更新成功"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Failure 404 {object} errors.AppError "知识库不存在"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/config/{kbId} [put]
func (h *InitializationHandler) UpdateKBConfig(c *gin.Context) {
ctx := c.Request.Context()
kbIdStr := utils.SanitizeForLog(c.Param("kbId"))
var req KBModelConfigRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "Failed to parse KB config request", err)
c.Error(errors.NewBadRequestError(err.Error()))
return
}
// 获取知识库信息
kb, err := h.kbService.GetKnowledgeBaseByID(ctx, kbIdStr)
if err != nil || kb == nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{"kbId": utils.SanitizeForLog(kbIdStr)})
c.Error(errors.NewNotFoundError("知识库不存在"))
return
}
// 检查Embedding模型是否可以修改
if kb.EmbeddingModelID != "" && req.EmbeddingModelID != "" && kb.EmbeddingModelID != req.EmbeddingModelID {
// 检查是否已有文件
knowledgeList, err := h.knowledgeService.ListPagedKnowledgeByKnowledgeBaseID(ctx,
kbIdStr, &types.Pagination{
Page: 1,
PageSize: 1,
}, types.KnowledgeListFilter{})
if err == nil && knowledgeList != nil && knowledgeList.Total > 0 {
logger.Error(ctx, "Cannot change embedding model when files exist")
c.Error(errors.NewBadRequestError("知识库中已有文件,无法修改Embedding模型"))
return
}
}
// 从数据库获取模型详情并验证
llmModel, err := h.modelService.GetModelByID(ctx, req.LLMModelID)
if err != nil || llmModel == nil {
logger.Error(ctx, "LLM model not found")
c.Error(errors.NewBadRequestError("LLM模型不存在"))
return
}
// Embedding模型仅在需要时验证(RAG检索启用时)
if req.EmbeddingModelID != "" {
embeddingModel, err := h.modelService.GetModelByID(ctx, req.EmbeddingModelID)
if err != nil || embeddingModel == nil {
logger.Error(ctx, "Embedding model not found")
c.Error(errors.NewBadRequestError("Embedding模型不存在"))
return
}
}
// 更新知识库的模型ID
kb.SummaryModelID = req.LLMModelID
if req.EmbeddingModelID != "" {
kb.EmbeddingModelID = req.EmbeddingModelID
}
// 处理多模态模型配置
kb.VLMConfig = types.VLMConfig{}
if req.VLMConfig != nil && req.Multimodal.Enabled && req.VLMConfig.ModelID != "" {
vllmModel, err := h.modelService.GetModelByID(ctx, req.VLMConfig.ModelID)
if err != nil || vllmModel == nil {
logger.Warn(ctx, "VLM model not found")
} else {
kb.VLMConfig.Enabled = req.VLMConfig.Enabled
kb.VLMConfig.ModelID = req.VLMConfig.ModelID
}
}
if !kb.VLMConfig.Enabled {
kb.VLMConfig.ModelID = ""
}
// 处理ASR语音识别配置
kb.ASRConfig = types.ASRConfig{}
if req.ASRConfig != nil && req.ASRConfig.Enabled && req.ASRConfig.ModelID != "" {
asrModel, err := h.modelService.GetModelByID(ctx, req.ASRConfig.ModelID)
if err != nil || asrModel == nil {
logger.Warn(ctx, "ASR model not found")
} else {
kb.ASRConfig.Enabled = true
kb.ASRConfig.ModelID = req.ASRConfig.ModelID
kb.ASRConfig.Language = req.ASRConfig.Language
}
}
// 更新文档分块配置
if req.DocumentSplitting.ChunkSize > 0 {
kb.ChunkingConfig.ChunkSize = req.DocumentSplitting.ChunkSize
}
if req.DocumentSplitting.ChunkOverlap >= 0 {
kb.ChunkingConfig.ChunkOverlap = req.DocumentSplitting.ChunkOverlap
}
if len(req.DocumentSplitting.Separators) > 0 {
kb.ChunkingConfig.Separators = req.DocumentSplitting.Separators
}
kb.ChunkingConfig.ParserEngineRules = req.DocumentSplitting.ParserEngineRules
kb.ChunkingConfig.EnableParentChild = req.DocumentSplitting.EnableParentChild
if req.DocumentSplitting.ParentChunkSize > 0 {
kb.ChunkingConfig.ParentChunkSize = req.DocumentSplitting.ParentChunkSize
}
if req.DocumentSplitting.ChildChunkSize > 0 {
kb.ChunkingConfig.ChildChunkSize = req.DocumentSplitting.ChildChunkSize
}
// Pointer-based fields support clearing (empty string / 0 / empty slice
// is a valid "user picked default again" signal; absent in payload means
// "no change").
if req.DocumentSplitting.Strategy != nil {
kb.ChunkingConfig.Strategy = *req.DocumentSplitting.Strategy
}
if req.DocumentSplitting.TokenLimit != nil {
kb.ChunkingConfig.TokenLimit = *req.DocumentSplitting.TokenLimit
}
if req.DocumentSplitting.Languages != nil {
kb.ChunkingConfig.Languages = *req.DocumentSplitting.Languages
}
if req.DocumentSplitting.TableMetadataInstructions != nil {
kb.ChunkingConfig.TableMetadataInstructions = strings.TrimSpace(*req.DocumentSplitting.TableMetadataInstructions)
}
// 更新多模态配置
if req.Multimodal.Enabled {
// VLM model already set above
} else {
kb.VLMConfig.ModelID = ""
}
if req.VLMConfig != nil {
kb.VLMConfig.DescriptionLanguage = strings.TrimSpace(req.VLMConfig.DescriptionLanguage)
kb.VLMConfig.CustomInstructions = strings.TrimSpace(req.VLMConfig.CustomInstructions)
}
// Bind the concrete storage instance. Provider remains a compatibility
// projection for older clients and historical rows.
if strings.TrimSpace(req.StorageBackendID) != "" {
tenant, _ := types.TenantInfoFromContext(ctx)
backend, resolveErr := h.storageResolver.ResolveBackend(ctx, tenant, req.StorageBackendID, "")
if resolveErr != nil || backend == nil {
c.Error(errors.NewBadRequestError("Storage backend is unavailable"))
return
}
oldID := ""
if kb.StorageBackendID != nil {
oldID = *kb.StorageBackendID
}
if oldID != "" && oldID != backend.ID {
knowledgeList, listErr := h.knowledgeService.ListPagedKnowledgeByKnowledgeBaseID(ctx, kbIdStr, &types.Pagination{Page: 1, PageSize: 1}, types.KnowledgeListFilter{})
if listErr == nil && knowledgeList != nil && knowledgeList.Total > 0 {
c.Error(errors.NewBadRequestError("Storage backend cannot be changed while the knowledge base contains files; migrate storage first"))
return
}
}
kb.StorageBackendID = &backend.ID
req.StorageProvider = backend.Provider
}
// Legacy provider projection.
provider := strings.ToLower(strings.TrimSpace(req.StorageProvider))
if provider == "" {
provider = "local"
}
if !isStorageProviderAllowed(provider) {
c.Error(errors.NewBadRequestError("Storage provider is not allowed by STORAGE_ALLOW_LIST"))
return
}
oldProvider := kb.GetStorageProvider()
if oldProvider == "" {
oldProvider = "local"
}
if oldProvider != provider {
knowledgeList, err := h.knowledgeService.ListPagedKnowledgeByKnowledgeBaseID(ctx,
kbIdStr, &types.Pagination{Page: 1, PageSize: 1}, types.KnowledgeListFilter{})
if err == nil && knowledgeList != nil && knowledgeList.Total > 0 {
logger.Warn(ctx, "Storage engine changed with existing files, old files may become inaccessible")
}
}
kb.SetStorageProvider(provider)
// 更新知识图谱配置
if req.NodeExtract.Enabled {
// 转换 Nodes 和 Relations 为指针类型
nodes := make([]*types.GraphNode, len(req.NodeExtract.Nodes))
for i := range req.NodeExtract.Nodes {
nodes[i] = &req.NodeExtract.Nodes[i]
}
relations := make([]*types.GraphRelation, len(req.NodeExtract.Relations))
for i := range req.NodeExtract.Relations {
relations[i] = &req.NodeExtract.Relations[i]
}
kb.ExtractConfig = &types.ExtractConfig{
Enabled: req.NodeExtract.Enabled,
Text: req.NodeExtract.Text,
Tags: req.NodeExtract.Tags,
Nodes: nodes,
Relations: relations,
CustomInstructions: strings.TrimSpace(req.NodeExtract.CustomInstructions),
}
} else if kb.ExtractConfig != nil {
kb.ExtractConfig.Enabled = false
} else {
kb.ExtractConfig = &types.ExtractConfig{Enabled: false}
}
if err := validateExtractConfig(kb.ExtractConfig); err != nil {
logger.Error(ctx, "Invalid extract configuration", err)
c.Error(err)
return
}
// 更新问题生成配置
if req.QuestionGeneration.Enabled {
questionCount := req.QuestionGeneration.QuestionCount
if questionCount <= 0 {
questionCount = 3
}
if questionCount > 10 {
questionCount = 10
}
kb.QuestionGenerationConfig = &types.QuestionGenerationConfig{
Enabled: true,
QuestionCount: questionCount,
CustomInstructions: strings.TrimSpace(req.QuestionGeneration.CustomInstructions),
}
} else {
kb.QuestionGenerationConfig = &types.QuestionGenerationConfig{
Enabled: false,
CustomInstructions: strings.TrimSpace(req.QuestionGeneration.CustomInstructions),
}
}
types.NormalizeKnowledgeBasePromptInstructions(kb)
if err := validateKnowledgeBasePromptInstructions(kb); err != nil {
c.Error(err)
return
}
// 保存更新后的知识库
if err := h.kbRepository.UpdateKnowledgeBase(ctx, kb); err != nil {
logger.Error(ctx, "Failed to update knowledge base", err)
c.Error(errors.NewInternalServerError("更新知识库失败: " + err.Error()))
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "配置更新成功",
})
}
// InitializeByKB godoc
// @Summary 初始化知识库配置
// @Description 根据知识库ID执行完整配置更新
// @Tags 初始化
// @Accept json
// @Produce json
// @Param kbId path string true "知识库ID"
// @Param request body handler.InitializationRequest true "初始化请求"
// @Success 200 {object} map[string]interface{} "初始化成功"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/initialize/{kbId} [post]
func (h *InitializationHandler) InitializeByKB(c *gin.Context) {
ctx := c.Request.Context()
kbIdStr := utils.SanitizeForLog(c.Param("kbId"))
req, err := h.bindInitializationRequest(ctx, c)
if err != nil {
c.Error(err)
return
}
logger.Infof(
ctx,
"Starting knowledge base configuration update, kbId: %s, request: %s",
utils.SanitizeForLog(kbIdStr),
utils.SanitizeForLog(utils.ToJSON(req)),
)
kb, err := h.getKnowledgeBaseForInitialization(ctx, kbIdStr)
if err != nil {
c.Error(err)
return
}
if err := h.validateInitializationConfigs(ctx, req); err != nil {
c.Error(err)
return
}
processedModels, err := h.processInitializationModels(ctx, kb, kbIdStr, req)
if err != nil {
c.Error(err)
return
}
h.applyKnowledgeBaseInitialization(kb, req, processedModels)
if err := h.kbRepository.UpdateKnowledgeBase(ctx, kb); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{"kbId": utils.SanitizeForLog(kbIdStr)})
c.Error(errors.NewInternalServerError("更新知识库配置失败: " + err.Error()))
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "知识库配置更新成功",
"data": gin.H{
"models": processedModels,
"knowledge_base": kb,
},
})
}
func (h *InitializationHandler) bindInitializationRequest(ctx context.Context, c *gin.Context) (*InitializationRequest, error) {
var req InitializationRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "Failed to parse initialization request", err)
return nil, errors.NewBadRequestError(err.Error())
}
return &req, nil
}
func (h *InitializationHandler) getKnowledgeBaseForInitialization(ctx context.Context, kbIdStr string) (*types.KnowledgeBase, error) {
kb, err := h.kbService.GetKnowledgeBaseByID(ctx, kbIdStr)
if err != nil {
// The repo's not-found sentinel must surface as 404, not 500.
// Without this, every probe of a stale kb id from the
// initialization flow burns ops attention with a fake server
// error. See knowledgebase.go:validateAndGetKnowledgeBase.
if stderrors.Is(err, repository.ErrKnowledgeBaseNotFound) {
return nil, errors.NewNotFoundError("知识库不存在")
}
logger.ErrorWithFields(ctx, err, map[string]interface{}{"kbId": utils.SanitizeForLog(kbIdStr)})
return nil, errors.NewInternalServerError("获取知识库信息失败: " + err.Error())
}
if kb == nil {
logger.Error(ctx, "Knowledge base not found")
return nil, errors.NewNotFoundError("知识库不存在")
}
return kb, nil
}
func (h *InitializationHandler) validateInitializationConfigs(ctx context.Context, req *InitializationRequest) error {
// SSRF validation for all user-supplied BaseURLs
urlsToCheck := []struct {
label string
url string
}{
{"LLM BaseURL", req.LLM.BaseURL},
{"Embedding BaseURL", req.Embedding.BaseURL},
{"Rerank BaseURL", req.Rerank.BaseURL},
}
if req.Multimodal.VLM != nil {
urlsToCheck = append(urlsToCheck, struct {
label string
url string
}{"VLM BaseURL", req.Multimodal.VLM.BaseURL})
}
for _, u := range urlsToCheck {
if u.url != "" {
if err := utils.ValidateURLForSSRF(u.url); err != nil {
logger.Warnf(ctx, "SSRF validation failed for %s: %v", u.label, err)
return errors.NewBadRequestError(utils.FormatSSRFError(u.label, u.url, err))
}
}
}
if err := h.validateMultimodalConfig(ctx, req); err != nil {
return err
}
if err := validateRerankConfig(ctx, req); err != nil {
return err
}
return validateNodeExtractConfig(ctx, req)
}
func (h *InitializationHandler) validateMultimodalConfig(ctx context.Context, req *InitializationRequest) error {
if !req.Multimodal.Enabled {
return nil
}
storageType := strings.ToLower(req.Multimodal.StorageType)
if req.Multimodal.VLM == nil {
logger.Error(ctx, "Multimodal enabled but missing VLM configuration")
return errors.NewBadRequestError("启用多模态时需要配置VLM信息")
}
if req.Multimodal.VLM.InterfaceType == "ollama" {
req.Multimodal.VLM.BaseURL = os.Getenv("OLLAMA_BASE_URL") + "/v1"
}
if req.Multimodal.VLM.ModelName == "" || req.Multimodal.VLM.BaseURL == "" {
logger.Error(ctx, "VLM configuration incomplete")
return errors.NewBadRequestError("VLM配置不完整")
}
switch storageType {
case "cos":
if req.Multimodal.COS == nil || req.Multimodal.COS.SecretID == "" || req.Multimodal.COS.SecretKey == "" ||
req.Multimodal.COS.Region == "" || req.Multimodal.COS.BucketName == "" ||
req.Multimodal.COS.AppID == "" {
logger.Error(ctx, "COS configuration incomplete")
return errors.NewBadRequestError("COS配置不完整")
}
case "minio":
if req.Multimodal.Minio == nil || req.Multimodal.Minio.BucketName == "" ||
os.Getenv("MINIO_ACCESS_KEY_ID") == "" || os.Getenv("MINIO_SECRET_ACCESS_KEY") == "" {
logger.Error(ctx, "MinIO configuration incomplete")
return errors.NewBadRequestError("MinIO配置不完整")
}
}
return nil
}
func validateRerankConfig(ctx context.Context, req *InitializationRequest) error {
if !req.Rerank.Enabled {
return nil
}
if req.Rerank.ModelName == "" || req.Rerank.BaseURL == "" {
logger.Error(ctx, "Rerank configuration incomplete")
return errors.NewBadRequestError("Rerank配置不完整")
}
return nil
}
func validateNodeExtractConfig(ctx context.Context, req *InitializationRequest) error {
if !req.NodeExtract.Enabled {
return nil
}
if !strings.EqualFold(os.Getenv("NEO4J_ENABLE"), "true") &&
!strings.EqualFold(os.Getenv("LEGACY_GRAPH_ENABLE"), "true") {
logger.Error(ctx, "Node Extractor configuration incomplete")
return errors.NewBadRequestError("请正确配置环境变量NEO4J_ENABLE或LEGACY_GRAPH_ENABLE")
}
if req.NodeExtract.Text == "" || len(req.NodeExtract.Tags) == 0 {
logger.Error(ctx, "Node Extractor configuration incomplete")
return errors.NewBadRequestError("Node Extractor配置不完整")
}
if len(req.NodeExtract.Nodes) == 0 || len(req.NodeExtract.Relations) == 0 {
logger.Error(ctx, "Node Extractor configuration incomplete")
return errors.NewBadRequestError("请先提取实体和关系")
}
return nil
}
type modelDescriptor struct {
modelType types.ModelType
name string
source types.ModelSource
description string
baseURL string
apiKey string
dimension int
interfaceType string
}
func buildModelDescriptors(req *InitializationRequest) []modelDescriptor {
descriptors := []modelDescriptor{
{
modelType: types.ModelTypeKnowledgeQA,
name: utils.SanitizeForLog(req.LLM.ModelName),
source: types.ModelSource(req.LLM.Source),
description: "LLM Model for Knowledge QA",
baseURL: utils.SanitizeForLog(req.LLM.BaseURL),
apiKey: req.LLM.APIKey,
},
{
modelType: types.ModelTypeEmbedding,
name: utils.SanitizeForLog(req.Embedding.ModelName),
source: types.ModelSource(req.Embedding.Source),
description: "Embedding Model",
baseURL: utils.SanitizeForLog(req.Embedding.BaseURL),
apiKey: req.Embedding.APIKey,
dimension: req.Embedding.Dimension,
},
}
if req.Rerank.Enabled {
descriptors = append(descriptors, modelDescriptor{
modelType: types.ModelTypeRerank,
name: utils.SanitizeForLog(req.Rerank.ModelName),
source: types.ModelSourceRemote,
description: "Rerank Model",
baseURL: utils.SanitizeForLog(req.Rerank.BaseURL),
apiKey: req.Rerank.APIKey,
})
}
if req.Multimodal.Enabled && req.Multimodal.VLM != nil {
descriptors = append(descriptors, modelDescriptor{
modelType: types.ModelTypeVLLM,
name: utils.SanitizeForLog(req.Multimodal.VLM.ModelName),
source: types.ModelSourceRemote,
description: "VLM Model",
baseURL: utils.SanitizeForLog(req.Multimodal.VLM.BaseURL),
apiKey: req.Multimodal.VLM.APIKey,
interfaceType: req.Multimodal.VLM.InterfaceType,
})
}
return descriptors
}
func (h *InitializationHandler) processInitializationModels(
ctx context.Context,
kb *types.KnowledgeBase,
kbIdStr string,
req *InitializationRequest,
) ([]*types.Model, error) {
descriptors := buildModelDescriptors(req)
var processedModels []*types.Model
for _, descriptor := range descriptors {
model := descriptor.toModel()
existingModelID := h.findExistingModelID(kb, descriptor.modelType)
var existingModel *types.Model
if existingModelID != "" {
var err error
existingModel, err = h.modelService.GetModelByID(ctx, existingModelID)
if err != nil {
logger.Warnf(ctx, "Failed to get existing model %s: %v, will create new one", existingModelID, err)
existingModel = nil
}
}
if existingModel != nil {
existingModel.Name = model.Name
existingModel.Source = model.Source
existingModel.Description = model.Description
existingModel.Parameters = model.Parameters
existingModel.UpdatedAt = time.Now()
if err := h.modelService.UpdateModel(ctx, existingModel); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"model_id": model.ID,
"kb_id": kbIdStr,
})
return nil, errors.NewInternalServerError("更新模型失败: " + err.Error())
}
processedModels = append(processedModels, existingModel)
continue
}
if err := h.modelService.CreateModel(ctx, model); err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"model_id": model.ID,
"kb_id": kbIdStr,
})
return nil, errors.NewInternalServerError("创建模型失败: " + err.Error())
}
processedModels = append(processedModels, model)
}
return processedModels, nil
}
func (descriptor modelDescriptor) toModel() *types.Model {
model := &types.Model{
Type: descriptor.modelType,
Name: descriptor.name,
Source: descriptor.source,
Description: descriptor.description,
Parameters: types.ModelParameters{
BaseURL: descriptor.baseURL,
APIKey: descriptor.apiKey,
InterfaceType: descriptor.interfaceType,
},
IsDefault: false,
Status: types.ModelStatusActive,
}
if descriptor.modelType == types.ModelTypeEmbedding {
model.Parameters.EmbeddingParameters = types.EmbeddingParameters{
Dimension: descriptor.dimension,
}
}
return model
}
func (h *InitializationHandler) findExistingModelID(kb *types.KnowledgeBase, modelType types.ModelType) string {
switch modelType {
case types.ModelTypeEmbedding:
return kb.EmbeddingModelID
case types.ModelTypeKnowledgeQA:
return kb.SummaryModelID
case types.ModelTypeVLLM:
return kb.VLMConfig.ModelID
default:
return ""
}
}
func (h *InitializationHandler) applyKnowledgeBaseInitialization(
kb *types.KnowledgeBase,
req *InitializationRequest,
processedModels []*types.Model,
) {
embeddingModelID, llmModelID, vlmModelID := extractModelIDs(processedModels)
kb.SummaryModelID = llmModelID
kb.EmbeddingModelID = embeddingModelID
kb.ChunkingConfig = types.ChunkingConfig{
ChunkSize: req.DocumentSplitting.ChunkSize,
ChunkOverlap: req.DocumentSplitting.ChunkOverlap,
Separators: req.DocumentSplitting.Separators,
}
if req.Multimodal.Enabled {
kb.VLMConfig = types.VLMConfig{
Enabled: req.Multimodal.Enabled,
ModelID: vlmModelID,
}
switch req.Multimodal.StorageType {
case "cos":
if req.Multimodal.COS != nil {
kb.SetStorageProvider("cos")
// Legacy: also write to cos_config for backward compat with old code paths
kb.StorageConfig = types.StorageConfig{
Provider: req.Multimodal.StorageType,
BucketName: req.Multimodal.COS.BucketName,
AppID: req.Multimodal.COS.AppID,
PathPrefix: req.Multimodal.COS.PathPrefix,
SecretID: req.Multimodal.COS.SecretID,
SecretKey: req.Multimodal.COS.SecretKey,
Region: req.Multimodal.COS.Region,
}
}
case "minio":
if req.Multimodal.Minio != nil {
kb.SetStorageProvider("minio")
// Legacy: also write to cos_config for backward compat with old code paths
kb.StorageConfig = types.StorageConfig{
Provider: req.Multimodal.StorageType,
BucketName: req.Multimodal.Minio.BucketName,
PathPrefix: req.Multimodal.Minio.PathPrefix,
SecretID: os.Getenv("MINIO_ACCESS_KEY_ID"),
SecretKey: os.Getenv("MINIO_SECRET_ACCESS_KEY"),
}
}
}
} else {
kb.VLMConfig = types.VLMConfig{}
kb.SetStorageProvider("")
kb.StorageConfig = types.StorageConfig{}
}
if req.NodeExtract.Enabled {
kb.ExtractConfig = &types.ExtractConfig{
Text: req.NodeExtract.Text,
Tags: req.NodeExtract.Tags,
Nodes: make([]*types.GraphNode, 0),
Relations: make([]*types.GraphRelation, 0),
}
for _, rnode := range req.NodeExtract.Nodes {
node := &types.GraphNode{
Name: rnode.Name,
Attributes: rnode.Attributes,
}
kb.ExtractConfig.Nodes = append(kb.ExtractConfig.Nodes, node)
}
for _, relation := range req.NodeExtract.Relations {
kb.ExtractConfig.Relations = append(kb.ExtractConfig.Relations, &types.GraphRelation{
Node1: relation.Node1,
Node2: relation.Node2,
Type: relation.Type,
})
}
}
}
func extractModelIDs(processedModels []*types.Model) (embeddingModelID, llmModelID, vlmModelID string) {
for _, model := range processedModels {
if model == nil {
continue
}
switch model.Type {
case types.ModelTypeEmbedding:
embeddingModelID = model.ID
case types.ModelTypeKnowledgeQA:
llmModelID = model.ID
case types.ModelTypeVLLM:
vlmModelID = model.ID
}
}
return
}
// CheckOllamaStatus godoc
// @Summary 检查Ollama服务状态
// @Description 检查Ollama服务是否可用
// @Tags 初始化
// @Accept json
// @Produce json
// @Success 200 {object} map[string]interface{} "Ollama状态"
// @Router /initialization/ollama/status [get]
func (h *InitializationHandler) CheckOllamaStatus(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Checking Ollama service status")
// Determine Ollama base URL for display
baseURL := os.Getenv("OLLAMA_BASE_URL")
if baseURL == "" {
baseURL = "http://host.docker.internal:11434"
}
// 检查Ollama服务是否可用
err := h.ollamaService.StartService(ctx)
if err != nil {
logger.ErrorWithFields(ctx, err, nil)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"available": false,
"error": err.Error(),
"baseUrl": baseURL,
},
})
return
}
version, err := h.ollamaService.GetVersion(ctx)
if err != nil {
logger.ErrorWithFields(ctx, err, nil)
version = "unknown"
}
logger.Info(ctx, "Ollama service is available")
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"available": h.ollamaService.IsAvailable(),
"version": version,
"baseUrl": baseURL,
},
})
}
// CheckOllamaModels godoc
// @Summary 检查Ollama模型状态
// @Description 检查指定的Ollama模型是否已安装
// @Tags 初始化
// @Accept json
// @Produce json
// @Param request body object{models=[]string} true "模型名称列表"
// @Success 200 {object} map[string]interface{} "模型状态"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/ollama/models/check [post]
func (h *InitializationHandler) CheckOllamaModels(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Checking Ollama models status")
var req struct {
Models []string `json:"models" binding:"required"`
}
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "Failed to parse models check request", err)
c.Error(errors.NewBadRequestError(err.Error()))
return
}
// 检查Ollama服务是否可用
if !h.ollamaService.IsAvailable() {
err := h.ollamaService.StartService(ctx)
if err != nil {
logger.ErrorWithFields(ctx, err, nil)
c.Error(errors.NewInternalServerError("Ollama服务不可用: " + err.Error()))
return
}
}
modelStatus := make(map[string]bool)
// 检查每个模型是否存在
for _, modelName := range req.Models {
available, err := h.ollamaService.IsModelAvailable(ctx, modelName)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"model_name": modelName,
})
modelStatus[modelName] = false
} else {
modelStatus[modelName] = available
}
logger.Infof(ctx, "Model %s availability: %v", utils.SanitizeForLog(modelName), modelStatus[modelName])
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"models": modelStatus,
},
})
}
// DownloadOllamaModel godoc
// @Summary 下载Ollama模型
// @Description 异步下载指定的Ollama模型
// @Tags 初始化
// @Accept json
// @Produce json
// @Param request body object{modelName=string} true "模型名称"
// @Success 200 {object} map[string]interface{} "下载任务信息"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/ollama/models/download [post]
func (h *InitializationHandler) DownloadOllamaModel(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Starting async Ollama model download")
var req struct {
ModelName string `json:"modelName" binding:"required"`
}
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "Failed to parse model download request", err)
c.Error(errors.NewBadRequestError(err.Error()))
return
}
// 检查Ollama服务是否可用
if !h.ollamaService.IsAvailable() {
err := h.ollamaService.StartService(ctx)
if err != nil {
logger.ErrorWithFields(ctx, err, nil)
c.Error(errors.NewInternalServerError("Ollama服务不可用: " + err.Error()))
return
}
}
// 检查模型是否已存在
available, err := h.ollamaService.IsModelAvailable(ctx, req.ModelName)
if err != nil {
c.Error(errors.NewInternalServerError("检查模型状态失败: " + err.Error()))
return
}
if available {
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "模型已存在",
"data": gin.H{
"modelName": req.ModelName,
"status": "completed",
"progress": 100.0,
},
})
return
}
// 检查是否已有相同模型的下载任务
tasksMutex.RLock()
for _, task := range downloadTasks {
if task.ModelName == req.ModelName && (task.Status == "pending" || task.Status == "downloading") {
tasksMutex.RUnlock()
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "模型下载任务已存在",
"data": gin.H{
"taskId": task.ID,
"modelName": task.ModelName,
"status": task.Status,
"progress": task.Progress,
},
})
return
}
}
tasksMutex.RUnlock()
// 创建下载任务
taskID := uuid.New().String()
task := &DownloadTask{
ID: taskID,
ModelName: req.ModelName,
Status: "pending",
Progress: 0.0,
Message: "准备下载",
StartTime: time.Now(),
}
tasksMutex.Lock()
downloadTasks[taskID] = task
tasksMutex.Unlock()
// 启动异步下载
newCtx, cancel := context.WithTimeout(context.Background(), 12*time.Hour)
go func() {
defer cancel()
h.downloadModelAsync(newCtx, taskID, req.ModelName)
}()
logger.Infof(ctx, "Created download task for model, task ID: %s", taskID)
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "模型下载任务已创建",
"data": gin.H{
"taskId": taskID,
"modelName": req.ModelName,
"status": "pending",
"progress": 0.0,
},
})
}
// GetDownloadProgress godoc
// @Summary 获取下载进度
// @Description 获取Ollama模型下载任务的进度
// @Tags 初始化
// @Accept json
// @Produce json
// @Param taskId path string true "任务ID"
// @Success 200 {object} map[string]interface{} "下载进度"
// @Failure 404 {object} errors.AppError "任务不存在"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/ollama/download/progress/{taskId} [get]
func (h *InitializationHandler) GetDownloadProgress(c *gin.Context) {
taskID := c.Param("taskId")
if taskID == "" {
c.Error(errors.NewBadRequestError("任务ID不能为空"))
return
}
tasksMutex.RLock()
task, exists := downloadTasks[taskID]
tasksMutex.RUnlock()
if !exists {
c.Error(errors.NewNotFoundError("下载任务不存在"))
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": task,
})
}
// ListDownloadTasks godoc
// @Summary 列出下载任务
// @Description 列出所有Ollama模型下载任务
// @Tags 初始化
// @Accept json
// @Produce json
// @Success 200 {object} map[string]interface{} "任务列表"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/ollama/download/tasks [get]
func (h *InitializationHandler) ListDownloadTasks(c *gin.Context) {
tasksMutex.RLock()
tasks := make([]*DownloadTask, 0, len(downloadTasks))
for _, task := range downloadTasks {
tasks = append(tasks, task)
}
tasksMutex.RUnlock()
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": tasks,
})
}
// ListOllamaModels godoc
// @Summary 列出Ollama模型
// @Description 列出已安装的Ollama模型
// @Tags 初始化
// @Accept json
// @Produce json
// @Success 200 {object} map[string]interface{} "模型列表"
// @Failure 500 {object} errors.AppError "服务器错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/ollama/models [get]
func (h *InitializationHandler) ListOllamaModels(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Listing installed Ollama models")
// 确保服务可用
if !h.ollamaService.IsAvailable() {
if err := h.ollamaService.StartService(ctx); err != nil {
logger.ErrorWithFields(ctx, err, nil)
c.Error(errors.NewInternalServerError("Ollama服务不可用: " + err.Error()))
return
}
}
// 使用 ListModelsDetailed 获取包含大小等详细信息的模型列表
models, err := h.ollamaService.ListModelsDetailed(ctx)
if err != nil {
logger.ErrorWithFields(ctx, err, nil)
c.Error(errors.NewInternalServerError("获取模型列表失败: " + err.Error()))
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"models": models,
},
})
}
// downloadModelAsync 异步下载模型
func (h *InitializationHandler) downloadModelAsync(ctx context.Context,
taskID, modelName string,
) {
logger.Infof(ctx, "Starting async download for model, task: %s", taskID)
// 更新任务状态为下载中
h.updateTaskStatus(taskID, "downloading", 0.0, "开始下载模型")
// 执行下载,带进度回调
err := h.pullModelWithProgress(ctx, modelName, func(progress float64, message string) {
h.updateTaskStatus(taskID, "downloading", progress, message)
})
if err != nil {
logger.Error(ctx, "Failed to download model", err)
h.updateTaskStatus(taskID, "failed", 0.0, fmt.Sprintf("下载失败: %v", err))
return
}
// 下载成功
logger.Infof(ctx, "Model downloaded successfully, task: %s", taskID)
h.updateTaskStatus(taskID, "completed", 100.0, "下载完成")
}
// pullModelWithProgress 下载模型并提供进度回调
func (h *InitializationHandler) pullModelWithProgress(ctx context.Context,
modelName string,
progressCallback func(float64, string),
) error {
// 检查服务是否可用
if err := h.ollamaService.StartService(ctx); err != nil {
logger.ErrorWithFields(ctx, err, nil)
return err
}
// 检查模型是否已存在
available, err := h.ollamaService.IsModelAvailable(ctx, modelName)
if err != nil {
logger.Error(ctx, "Failed to check model availability", err)
return err
}
if available {
progressCallback(100.0, "模型已存在")
return nil
}
// 创建下载请求
pullReq := &api.PullRequest{
Name: modelName,
}
// 使用Ollama客户端的Pull方法,带进度回调
err = h.ollamaService.GetClient().Pull(ctx, pullReq, func(progress api.ProgressResponse) error {
progressPercent := 0.0
message := "下载中"
if progress.Total > 0 && progress.Completed > 0 {
progressPercent = float64(progress.Completed) / float64(progress.Total) * 100
message = fmt.Sprintf("下载中: %.1f%% (%s)", progressPercent, progress.Status)
} else if progress.Status != "" {
message = progress.Status
}
// 调用进度回调
progressCallback(progressPercent, message)
logger.Infof(ctx,
"Download progress: %.2f%% - %s", progressPercent, message,
)
return nil
})
if err != nil {
return fmt.Errorf("failed to pull model: %w", err)
}
return nil
}
// updateTaskStatus 更新任务状态
func (h *InitializationHandler) updateTaskStatus(
taskID, status string, progress float64, message string,
) {
tasksMutex.Lock()
defer tasksMutex.Unlock()
if task, exists := downloadTasks[taskID]; exists {
task.Status = status
task.Progress = progress
task.Message = message
if status == "completed" || status == "failed" {
now := time.Now()
task.EndTime = &now
}
}
}
// GetCurrentConfigByKB godoc
// @Summary 获取知识库配置
// @Description 根据知识库ID获取当前配置信息
// @Tags 初始化
// @Accept json
// @Produce json
// @Param kbId path string true "知识库ID"
// @Success 200 {object} map[string]interface{} "配置信息"
// @Failure 404 {object} errors.AppError "知识库不存在"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/config/{kbId} [get]
func (h *InitializationHandler) GetCurrentConfigByKB(c *gin.Context) {
ctx := c.Request.Context()
kbIdStr := utils.SanitizeForLog(c.Param("kbId"))
logger.Info(ctx, "Getting configuration for knowledge base")
// 获取指定知识库信息
kb, err := h.kbService.GetKnowledgeBaseByID(ctx, kbIdStr)
if err != nil {
// Mirror getKnowledgeBaseForInitialization above: missing /
// cross-tenant kb ids are 404, not 500.
if stderrors.Is(err, repository.ErrKnowledgeBaseNotFound) {
c.Error(errors.NewNotFoundError("知识库不存在"))
return
}
logger.Error(ctx, "Failed to get knowledge base", err)
c.Error(errors.NewInternalServerError("获取知识库信息失败: " + err.Error()))
return
}
if kb == nil {
logger.Error(ctx, "Knowledge base not found")
c.Error(errors.NewNotFoundError("知识库不存在"))
return
}
// 根据知识库的模型ID获取特定模型
var models []*types.Model
modelIDs := []string{
kb.EmbeddingModelID,
kb.SummaryModelID,
kb.VLMConfig.ModelID,
}
for _, modelID := range modelIDs {
if modelID != "" {
model, err := h.modelService.GetModelByID(ctx, modelID)
if err != nil {
logger.Warn(ctx, "Failed to get model", err)
// 如果模型不存在或获取失败,继续处理其他模型
continue
}
if model != nil {
models = append(models, model)
}
}
}
// 检查知识库是否有文件
knowledgeList, err := h.knowledgeService.ListPagedKnowledgeByKnowledgeBaseID(ctx,
kbIdStr, &types.Pagination{
Page: 1,
PageSize: 1,
}, types.KnowledgeListFilter{})
hasFiles := err == nil && knowledgeList != nil && knowledgeList.Total > 0
// 构建配置响应
config := h.buildConfigResponse(ctx, models, kb, hasFiles)
logger.Info(ctx, "Knowledge base configuration retrieved successfully")
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": config,
})
}
// buildConfigResponse 构建配置响应数据
func (h *InitializationHandler) buildConfigResponse(ctx context.Context, models []*types.Model,
kb *types.KnowledgeBase, hasFiles bool,
) map[string]interface{} {
config := map[string]interface{}{
"hasFiles": hasFiles,
}
includeIntegrationDetail := dto.CanViewIntegrationSecrets(ctx)
// 按类型分组模型
for _, model := range models {
if model == nil {
continue
}
// Hide sensitive information for builtin models and viewers.
baseURL := model.Parameters.BaseURL
if model.IsBuiltin || !includeIntegrationDetail {
baseURL = ""
}
switch model.Type {
case types.ModelTypeKnowledgeQA:
config["llm"] = map[string]interface{}{
"source": string(model.Source),
"modelName": model.Name,
"baseUrl": baseURL,
"credentials": map[string]bool{
"apiKey": model.Parameters.APIKey != "" && !model.IsBuiltin,
},
}
case types.ModelTypeEmbedding:
config["embedding"] = map[string]interface{}{
"source": string(model.Source),
"modelName": model.Name,
"baseUrl": baseURL,
"dimension": model.Parameters.EmbeddingParameters.Dimension,
"credentials": map[string]bool{
"apiKey": model.Parameters.APIKey != "" && !model.IsBuiltin,
},
}
case types.ModelTypeRerank:
config["rerank"] = map[string]interface{}{
"enabled": true,
"modelName": model.Name,
"baseUrl": baseURL,
"credentials": map[string]bool{
"apiKey": model.Parameters.APIKey != "" && !model.IsBuiltin,
},
}
case types.ModelTypeVLLM:
if config["multimodal"] == nil {
config["multimodal"] = map[string]interface{}{
"enabled": true,
}
}
multimodal := config["multimodal"].(map[string]interface{})
multimodal["vlm"] = map[string]interface{}{
"modelName": model.Name,
"baseUrl": baseURL,
"interfaceType": model.Parameters.InterfaceType,
"modelId": model.ID,
"credentials": map[string]bool{
"apiKey": model.Parameters.APIKey != "" && !model.IsBuiltin,
},
}
}
}
// 判断多模态是否启用:有VLM模型ID或有存储配置(兼容新旧字段)
storageProvider := kb.GetStorageProvider()
hasMultimodal := (kb.VLMConfig.IsEnabled() ||
kb.StorageConfig.SecretID != "" || kb.StorageConfig.BucketName != "" ||
(storageProvider != "" && storageProvider != "local"))
if config["multimodal"] == nil {
config["multimodal"] = map[string]interface{}{
"enabled": hasMultimodal,
}
} else {
config["multimodal"].(map[string]interface{})["enabled"] = hasMultimodal
}
if kb.VLMConfig.DescriptionLanguage != "" || kb.VLMConfig.CustomInstructions != "" {
if config["multimodal"] == nil {
config["multimodal"] = map[string]interface{}{
"enabled": hasMultimodal,
}
}
multimodal := config["multimodal"].(map[string]interface{})
if kb.VLMConfig.DescriptionLanguage != "" {
multimodal["descriptionLanguage"] = kb.VLMConfig.DescriptionLanguage
}
if kb.VLMConfig.CustomInstructions != "" {
multimodal["customInstructions"] = kb.VLMConfig.CustomInstructions
}
}
// 如果没有Rerank模型,设置rerank为disabled
if config["rerank"] == nil {
config["rerank"] = map[string]interface{}{
"enabled": false,
"modelName": "",
"baseUrl": "",
"credentials": map[string]bool{
"apiKey": false,
},
}
}
// 添加知识库的文档分割配置
if kb != nil {
ds := map[string]interface{}{
"chunkSize": kb.ChunkingConfig.ChunkSize,
"chunkOverlap": kb.ChunkingConfig.ChunkOverlap,
"separators": kb.ChunkingConfig.Separators,
}
if kb.ChunkingConfig.Strategy != "" {
ds["strategy"] = kb.ChunkingConfig.Strategy
}
if kb.ChunkingConfig.TokenLimit > 0 {
ds["tokenLimit"] = kb.ChunkingConfig.TokenLimit
}
if len(kb.ChunkingConfig.Languages) > 0 {
ds["languages"] = kb.ChunkingConfig.Languages
}
if kb.ChunkingConfig.TableMetadataInstructions != "" {
ds["tableMetadataInstructions"] = kb.ChunkingConfig.TableMetadataInstructions
}
config["documentSplitting"] = ds
// 添加多模态的存储配置信息(优先读新字段,兼容旧 cos_config)
effectiveProvider := kb.GetStorageProvider()
if kb.StorageConfig.SecretID != "" || (effectiveProvider != "" && effectiveProvider != "local") {
if config["multimodal"] == nil {
config["multimodal"] = map[string]interface{}{
"enabled": true,
}
}
multimodal := config["multimodal"].(map[string]interface{})
multimodal["storageType"] = effectiveProvider
switch effectiveProvider {
case "cos":
multimodal["cos"] = map[string]interface{}{
"region": kb.StorageConfig.Region,
"bucketName": kb.StorageConfig.BucketName,
"appId": kb.StorageConfig.AppID,
"pathPrefix": kb.StorageConfig.PathPrefix,
"credentials": map[string]bool{
"secretId": kb.StorageConfig.SecretID != "",
"secretKey": kb.StorageConfig.SecretKey != "",
},
}
case "minio":
multimodal["minio"] = map[string]interface{}{
"bucketName": kb.StorageConfig.BucketName,
"pathPrefix": kb.StorageConfig.PathPrefix,
}
}
}
}
if kb.ExtractConfig != nil {
nodeExtract := map[string]interface{}{
"enabled": kb.ExtractConfig.Enabled,
"text": kb.ExtractConfig.Text,
"tags": kb.ExtractConfig.Tags,
"nodes": kb.ExtractConfig.Nodes,
"relations": kb.ExtractConfig.Relations,
}
if kb.ExtractConfig.CustomInstructions != "" {
nodeExtract["customInstructions"] = kb.ExtractConfig.CustomInstructions
}
config["nodeExtract"] = nodeExtract
} else {
config["nodeExtract"] = map[string]interface{}{
"enabled": false,
}
}
if kb.QuestionGenerationConfig != nil {
config["questionGeneration"] = map[string]interface{}{
"enabled": kb.QuestionGenerationConfig.Enabled,
"questionCount": kb.QuestionGenerationConfig.QuestionCount,
"customInstructions": kb.QuestionGenerationConfig.CustomInstructions,
}
} else {
config["questionGeneration"] = map[string]interface{}{
"enabled": false,
}
}
return config
}
// ModelTestRequest 统一的"测试连接"请求体。
//
// 四种模型(chat/embedding/rerank/asr)的测试接口共享同一份结构,以便:
// - 前端只需维护一份表单 → 后端映射。
// - 后端可以直接把请求转成 *types.Model,再调用各包的 ConfigFromModel,
// 与生产路径(service.modelService.GetXxxModel)走完全相同的装配流程,
// 彻底消除过去每个测试端点手工拼 Config 的样板代码。
//
// 所有 provider/model 通用字段都在这里集中声明;若未来新增字段(比如现在的
// custom_headers),只需改一处,生产路径和测试路径会同时生效。
type ModelTestRequest struct {
Source string `json:"source"` // 为空时按需默认为 "remote"
ModelName string `json:"modelName" binding:"required"`
BaseURL string `json:"baseUrl"`
APIKey string `json:"apiKey"`
Provider string `json:"provider"`
InterfaceType string `json:"interfaceType,omitempty"`
Dimension int `json:"dimension,omitempty"`
SupportsDimensionOverride bool `json:"supportsDimensionOverride,omitempty"`
CustomHeaders map[string]string `json:"customHeaders,omitempty"`
ExtraConfig map[string]string `json:"extraConfig,omitempty"`
// AppSecret 用于 LKEAP / Volcengine Rerank 等需要第二段密钥的场景(对应模型 Parameters.AppSecret)。
AppSecret string `json:"appSecret,omitempty"`
// ModelID, when set, instructs the handler to substitute any missing
// secrets (APIKey, AppSecret via ExtraConfig) from the stored model
// record before assembling the test client. This lets the "Test
// connection" button work on existing models without making the
// frontend reload — and ship — the plaintext API key. Other fields
// (BaseURL, ModelName, etc.) on this request still override the
// stored values, so a user can validate a new endpoint against the
// existing credentials in one click.
ModelID string `json:"modelId,omitempty"`
}
// fillSecretsFromStoredModel mutates req in place: if req.ModelID is set
// and a secret field on the request is empty, the corresponding value from
// the stored (and decrypted) model is copied in. Non-empty request values
// are always preferred — they represent the user actively typing a new key
// they want to verify. Missing or inaccessible model is treated as a no-op
// (the connection test will fail downstream with a clearer "missing apiKey"
// error than we could produce here).
func (h *InitializationHandler) fillSecretsFromStoredModel(ctx context.Context, req *ModelTestRequest) {
if req == nil || req.ModelID == "" {
return
}
if req.APIKey != "" && req.AppSecret != "" {
return
}
stored, err := h.modelService.GetModelByID(ctx, req.ModelID)
if err != nil || stored == nil {
logger.Warnf(ctx, "test-connection: stored model %s not found, leaving secrets empty: %v",
utils.SanitizeForLog(req.ModelID), err)
return
}
if req.APIKey == "" {
req.APIKey = stored.Parameters.APIKey
}
if req.AppSecret == "" {
req.AppSecret = stored.Parameters.AppSecret
}
}
// RemoteModelCheckRequest 兼容旧 swagger 定义。
//
// Deprecated: 保留是为了不破坏已生成的 API 文档,新代码请直接使用 ModelTestRequest。
type RemoteModelCheckRequest = ModelTestRequest
// decryptModelAppSecret 解密模型 Parameters 中的 AppSecret(与 modelService 行为一致)。
func decryptModelAppSecret(encrypted string) string {
if encrypted == "" {
return encrypted
}
if key := utils.GetAESKey(); key != nil {
if plain, err := utils.DecryptAESGCM(encrypted, key); err == nil {
return plain
}
}
return encrypted
}
// buildTestModel 把测试连接请求转成一个临时的 *types.Model(不落库),
// 供 ConfigFromModel 使用。source 为空时按 defaultSource 兜底(chat/rerank/asr
// 默认 remote,embedding 会根据前端传入的 source 决定)。
func (h *InitializationHandler) buildTestModel(
req *ModelTestRequest, modelType types.ModelType, defaultSource types.ModelSource,
) *types.Model {
source := types.ModelSource(strings.ToLower(req.Source))
if source == "" {
source = defaultSource
}
return &types.Model{
Name: req.ModelName,
Type: modelType,
Source: source,
Parameters: types.ModelParameters{
BaseURL: req.BaseURL,
APIKey: req.APIKey,
AppSecret: req.AppSecret,
Provider: req.Provider,
InterfaceType: req.InterfaceType,
ExtraConfig: req.ExtraConfig,
CustomHeaders: req.CustomHeaders,
EmbeddingParameters: types.EmbeddingParameters{
Dimension: req.Dimension,
TruncatePromptTokens: 256,
SupportsDimensionOverride: req.SupportsDimensionOverride,
},
},
}
}
// resolveTenantWeKnoraCloudCreds 从当前空间上下文里取出 WeKnoraCloud 凭证,
// 供测试连接端点补齐 appID/appSecret。与 service.resolveWeKnoraCloudCredentials
// 对应,但因为 handler 还没有被注入 tenantService(历史原因),暂时从
// TenantInfoFromContext 读取,等效果相同。
func (h *InitializationHandler) resolveTenantWeKnoraCloudCreds(ctx context.Context) (string, string, bool) {
tenantInfo, ok := types.TenantInfoFromContext(ctx)
if !ok {
return "", "", false
}
creds := tenantInfo.Credentials.GetWeKnoraCloud()
if creds == nil {
return "", "", true
}
return creds.AppID, creds.AppSecret, true
}
// CheckRemoteModel godoc
// @Summary 检查远程模型
// @Description 检查远程API模型连接是否正常
// @Tags 初始化
// @Accept json
// @Produce json
// @Param request body RemoteModelCheckRequest true "模型检查请求"
// @Success 200 {object} map[string]interface{} "检查结果"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/remote/check [post]
func (h *InitializationHandler) CheckRemoteModel(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Checking remote model connection")
var req ModelTestRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "Failed to parse remote model check request", err)
c.Error(errors.NewBadRequestError(err.Error()))
return
}
h.fillSecretsFromStoredModel(ctx, &req)
if req.ModelName == "" || req.BaseURL == "" {
logger.Error(ctx, "Model name and base URL are required")
c.Error(errors.NewBadRequestError("模型名称和Base URL不能为空"))
return
}
if err := utils.ValidateURLForSSRF(req.BaseURL); err != nil {
logger.Warnf(ctx, "SSRF validation failed for remote model BaseURL: %v", err)
c.Error(errors.NewBadRequestError(utils.FormatSSRFError("Base URL", req.BaseURL, err)))
return
}
appID, appSecret, ok := h.resolveTenantWeKnoraCloudCreds(ctx)
if !ok {
logger.Error(ctx, "Tenant info not found")
c.Error(errors.NewBadRequestError("空间信息未找到"))
return
}
model := h.buildTestModel(&req, types.ModelTypeKnowledgeQA, types.ModelSourceRemote)
available, message := h.checkChatModelConnection(ctx, model, appID, appSecret)
logger.Infof(ctx, "Remote model check completed, available: %v, message: %s", available, message)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"available": available,
"message": message,
},
})
}
// TestEmbeddingModel godoc
// @Summary 测试Embedding模型
// @Description 测试Embedding接口是否可用并返回向量维度
// @Tags 初始化
// @Accept json
// @Produce json
// @Param request body handler.ModelTestRequest true "Embedding测试请求"
// @Success 200 {object} map[string]interface{} "测试结果"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/embedding/test [post]
func (h *InitializationHandler) TestEmbeddingModel(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Testing embedding model connectivity and functionality")
var req ModelTestRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "Failed to parse embedding test request", err)
c.Error(errors.NewBadRequestError(err.Error()))
return
}
h.fillSecretsFromStoredModel(ctx, &req)
if req.Source == "" {
req.Source = string(types.ModelSourceRemote)
}
if req.BaseURL != "" {
if err := utils.ValidateURLForSSRF(req.BaseURL); err != nil {
logger.Warnf(ctx, "SSRF validation failed for embedding BaseURL: %v", err)
c.Error(errors.NewBadRequestError(utils.FormatSSRFError("Base URL", req.BaseURL, err)))
return
}
}
// 阿里云多模态 Embedding 模型暂不支持
if strings.ToLower(req.Provider) == "aliyun" {
modelNameLower := strings.ToLower(req.ModelName)
if strings.Contains(modelNameLower, "vision") || strings.Contains(modelNameLower, "multimodal") {
logger.Infof(ctx, "Aliyun multimodal embedding model not supported: %s", req.ModelName)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"available": false,
"message": "阿里云多模态 Embedding 模型暂不支持,请使用纯文本 Embedding 模型(如 text-embedding-v4)",
"dimension": 0,
},
})
return
}
}
appID, appSecret, ok := h.resolveTenantWeKnoraCloudCreds(ctx)
if !ok {
logger.Error(ctx, "Tenant info not found")
c.Error(errors.NewBadRequestError("空间信息未找到"))
return
}
model := h.buildTestModel(&req, types.ModelTypeEmbedding, types.ModelSourceRemote)
emb, err := embedding.NewEmbedder(embedding.ConfigFromModel(model, appID, appSecret), h.pooler, h.ollamaService)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{"model": utils.SanitizeForLog(req.ModelName)})
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{`available`: false, `message`: fmt.Sprintf("创建Embedder失败: %v", err), `dimension`: 0},
})
return
}
vec, err := emb.Embed(ctx, "hello")
if err != nil {
logger.Error(ctx, "Failed to call embedder", err)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{`available`: false, `message`: fmt.Sprintf("调用Embedding失败: %v", err), `dimension`: 0},
})
return
}
logger.Infof(ctx, "Embedding test succeeded, dimension: %d", len(vec))
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{`available`: true, `message`: fmt.Sprintf("测试成功,向量维度=%d", len(vec)), `dimension`: len(vec)},
})
}
// classifyConnectionError maps an upstream error string to a short
// human-readable hint in Chinese. Callers should always combine the hint
// with the raw error message (e.g. fmt.Sprintf("%s:%v", hint, err)) so
// the operator can still see what URL / response body the SDK actually
// got — the hint is for "where to start looking", the raw error is for
// "what actually happened".
func classifyConnectionError(errMsg string) string {
switch {
case strings.Contains(errMsg, "401") || strings.Contains(errMsg, "unauthorized"):
return "认证失败,请检查API Key"
case strings.Contains(errMsg, "403") || strings.Contains(errMsg, "forbidden"):
return "权限不足,请检查API Key权限"
case strings.Contains(errMsg, "404") || strings.Contains(errMsg, "not found"):
return "API端点不存在,请检查Base URL"
case strings.Contains(errMsg, "timeout") || strings.Contains(errMsg, "context deadline exceeded"):
return "连接超时,请检查网络连接"
case strings.Contains(errMsg, "connection refused") || strings.Contains(errMsg, "no such host") || strings.Contains(errMsg, "dial tcp"):
return "无法连接到服务器,请检查Base URL"
default:
return "连接失败"
}
}
// checkChatModelConnection 使用 chat 模块做一次最小化调用来测试连通性与鉴权。
// 与生产路径走完全相同的 ConfigFromModel → NewChat 流程,因此 CustomHeaders、
// ExtraConfig、Provider 等字段都会被正确透传。
func (h *InitializationHandler) checkChatModelConnection(
ctx context.Context, model *types.Model, appID, appSecret string,
) (bool, string) {
chatInstance, err := chat.NewChat(chat.ConfigFromModel(model, appID, appSecret), h.ollamaService)
if err != nil {
return false, fmt.Sprintf("创建聊天实例失败: %v", err)
}
testMessages := []chat.Message{{Role: "user", Content: "test"}}
testOptions := &chat.ChatOptions{
MaxTokens: 1,
Thinking: &[]bool{false}[0], // for dashscope.aliyuncs qwen3-32b
}
_, err = chatInstance.Chat(ctx, testMessages, testOptions)
if err != nil {
errMsg := err.Error()
// 400 = endpoint reachable + auth ok, just a parameter mismatch
// (e.g. max_tokens vs max_completion_tokens). Treat as success.
if strings.Contains(errMsg, "status code: 400") {
return true, "连接正常,模型可用"
}
// For every other failure mode we surface a human-readable hint
// AND the upstream error verbatim. Swallowing the underlying
// message used to hide things like the actual URL the SDK
// tried, response body, etc. — making remote debugging nearly
// impossible. Format: "<hint>:<raw err>".
return false, fmt.Sprintf("%s:%v", classifyConnectionError(errMsg), err)
}
// 连接成功,模型可用
return true, "连接正常,模型可用"
}
// checkRerankModelConnection 使用 rerank 模块做一次最小化调用来测试连通性与鉴权。
// 与生产路径共用 ConfigFromModel,所有字段(CustomHeaders 等)都透传。
func (h *InitializationHandler) checkRerankModelConnection(
ctx context.Context, model *types.Model, appID, appSecret string,
) (bool, string) {
reranker, err := rerank.NewReranker(rerank.ConfigFromModel(model, appID, appSecret))
if err != nil {
return false, fmt.Sprintf("创建Reranker失败: %v", err)
}
results, err := reranker.Rerank(ctx, "ping", []string{"pong"})
if err != nil {
return false, fmt.Sprintf("重排测试失败: %v", err)
}
if len(results) > 0 {
return true, fmt.Sprintf("重排功能正常,返回%d个结果", len(results))
}
return false, "重排接口连接成功,但未返回重排结果"
}
// CheckRerankModel godoc
// @Summary 检查Rerank模型
// @Description 检查Rerank模型连接和功能是否正常
// @Tags 初始化
// @Accept json
// @Produce json
// @Param request body handler.ModelTestRequest true "Rerank检查请求"
// @Success 200 {object} map[string]interface{} "检查结果"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/rerank/check [post]
func (h *InitializationHandler) CheckRerankModel(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Checking rerank model connection and functionality")
var req ModelTestRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "Failed to parse rerank model check request", err)
c.Error(errors.NewBadRequestError(err.Error()))
return
}
h.fillSecretsFromStoredModel(ctx, &req)
if req.ModelName == "" || req.BaseURL == "" {
logger.Error(ctx, "Model name and base URL are required")
c.Error(errors.NewBadRequestError("模型名称和Base URL不能为空"))
return
}
if err := utils.ValidateURLForSSRF(req.BaseURL); err != nil {
logger.Warnf(ctx, "SSRF validation failed for rerank BaseURL: %v", err)
c.Error(errors.NewBadRequestError(utils.FormatSSRFError("Base URL", req.BaseURL, err)))
return
}
appID, appSecret, ok := h.resolveTenantWeKnoraCloudCreds(ctx)
if !ok {
logger.Error(ctx, "Tenant info not found")
c.Error(errors.NewBadRequestError("空间信息未找到"))
return
}
model := h.buildTestModel(&req, types.ModelTypeRerank, types.ModelSourceRemote)
if providerName := provider.ProviderName(model.Parameters.Provider); providerName == provider.ProviderLKEAP || providerName == provider.ProviderVolcengine {
appID = ""
appSecret = decryptModelAppSecret(model.Parameters.AppSecret)
}
available, message := h.checkRerankModelConnection(ctx, model, appID, appSecret)
logger.Infof(ctx, "Rerank model check completed, available: %v, message: %s", available, message)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"available": available,
"message": message,
},
})
}
// CheckASRModel godoc
// @Summary 检查ASR模型
// @Description 检查ASR(语音识别)模型连接是否正常,通过发送一段静默音频测试 /v1/audio/transcriptions 端点
// @Tags 初始化
// @Accept json
// @Produce json
// @Param request body handler.ModelTestRequest true "ASR检查请求"
// @Success 200 {object} map[string]interface{} "检查结果"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/asr/check [post]
func (h *InitializationHandler) CheckASRModel(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Checking ASR model connection")
var req ModelTestRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "Failed to parse ASR model check request", err)
c.Error(errors.NewBadRequestError(err.Error()))
return
}
h.fillSecretsFromStoredModel(ctx, &req)
if req.ModelName == "" || req.BaseURL == "" {
logger.Error(ctx, "Model name and base URL are required for ASR check")
c.Error(errors.NewBadRequestError("模型名称和Base URL不能为空"))
return
}
if err := utils.ValidateURLForSSRF(req.BaseURL); err != nil {
logger.Warnf(ctx, "SSRF validation failed for ASR BaseURL: %v", err)
c.Error(errors.NewBadRequestError(utils.FormatSSRFError("Base URL", req.BaseURL, err)))
return
}
// 用统一构造器生成测试用 *types.Model(ASR 不涉及 WeKnoraCloud 凭证),
// 发送一段极短的静默 WAV 音频验证 /v1/audio/transcriptions 端点可达。
model := h.buildTestModel(&req, types.ModelTypeASR, types.ModelSourceRemote)
asrInstance, err := asr.NewASR(asr.ConfigFromModel(model))
if err != nil {
logger.Errorf(ctx, "Failed to create ASR instance for check: %v", err)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"available": false,
"message": fmt.Sprintf("创建ASR实例失败: %v", err),
},
})
return
}
res, err := asrInstance.Transcribe(ctx, assets.ASRTestWAV, "asr_test.wav")
var text string
if res != nil {
text = res.Text
}
available := true
message := "ASR连接成功"
if err != nil {
errMsg := err.Error()
// Always include the raw upstream error after the hint — see
// classifyConnectionError comment for rationale.
switch {
case strings.Contains(errMsg, "401") || strings.Contains(errMsg, "Unauthorized") || strings.Contains(errMsg, "authentication"):
available = false
message = fmt.Sprintf("认证失败,请检查API Key:%s", errMsg)
case strings.Contains(errMsg, "404") || strings.Contains(errMsg, "Not Found"):
available = false
message = fmt.Sprintf("API端点不存在,请检查Base URL:%s", errMsg)
case strings.Contains(errMsg, "connection refused") || strings.Contains(errMsg, "no such host") || strings.Contains(errMsg, "dial tcp"):
available = false
message = fmt.Sprintf("无法连接到服务器,请检查Base URL:%s", errMsg)
case strings.Contains(errMsg, "model") && strings.Contains(errMsg, "not found"):
available = false
message = fmt.Sprintf("模型不存在,请检查模型名称:%s", errMsg)
default:
logger.Infof(ctx, "ASR check got non-fatal error (endpoint reachable): %v", err)
available = true
message = fmt.Sprintf("ASR端点可达(非致命错误: %s)", errMsg)
}
} else if text != "" {
message = fmt.Sprintf("ASR连接成功,转写结果: %s", text)
}
logger.Infof(ctx, "ASR model check completed, available: %v, message: %s", available, message)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"available": available,
"message": message,
},
})
}
// 使用结构体解析表单数据
type testMultimodalForm struct {
VLMModel string `form:"vlm_model"`
VLMBaseURL string `form:"vlm_base_url"`
VLMAPIKey string `form:"vlm_api_key"`
VLMInterfaceType string `form:"vlm_interface_type"`
StorageType string `form:"storage_type"`
// COS 配置
COSSecretID string `form:"cos_secret_id"`
COSSecretKey string `form:"cos_secret_key"`
COSRegion string `form:"cos_region"`
COSBucketName string `form:"cos_bucket_name"`
COSAppID string `form:"cos_app_id"`
COSPathPrefix string `form:"cos_path_prefix"`
// MinIO 配置(当存储为 minio 时)
MinioBucketName string `form:"minio_bucket_name"`
MinioPathPrefix string `form:"minio_path_prefix"`
// 文档切分配置(字符串后续自行解析,以避免类型绑定失败)
ChunkSize string `form:"chunk_size"`
ChunkOverlap string `form:"chunk_overlap"`
SeparatorsRaw string `form:"separators"`
}
// TestMultimodalFunction godoc
// @Summary 测试多模态功能
// @Description 上传图片测试多模态处理功能
// @Tags 初始化
// @Accept multipart/form-data
// @Produce json
// @Param image formData file true "测试图片"
// @Param vlm_model formData string true "VLM模型名称"
// @Param vlm_base_url formData string true "VLM Base URL"
// @Param vlm_api_key formData string false "VLM API Key"
// @Param vlm_interface_type formData string false "VLM接口类型"
// @Param storage_type formData string true "存储类型(cos/minio)"
// @Success 200 {object} map[string]interface{} "测试结果"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/multimodal/test [post]
func (h *InitializationHandler) TestMultimodalFunction(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Testing multimodal functionality")
var req testMultimodalForm
if err := c.ShouldBind(&req); err != nil {
logger.Error(ctx, "Failed to parse form data", err)
c.Error(errors.NewBadRequestError("表单参数解析失败"))
return
}
// ollama 场景自动拼接 base url
if req.VLMInterfaceType == "ollama" {
req.VLMBaseURL = os.Getenv("OLLAMA_BASE_URL") + "/v1"
}
req.StorageType = strings.ToLower(req.StorageType)
if req.VLMModel == "" || req.VLMBaseURL == "" {
logger.Error(ctx, "VLM model name and base URL are required")
c.Error(errors.NewBadRequestError("VLM模型名称和Base URL不能为空"))
return
}
// SSRF validation for VLM BaseURL
if err := utils.ValidateURLForSSRF(req.VLMBaseURL); err != nil {
logger.Warnf(ctx, "SSRF validation failed for VLM BaseURL: %v", err)
c.Error(errors.NewBadRequestError(utils.FormatSSRFError("VLM Base URL", req.VLMBaseURL, err)))
return
}
switch req.StorageType {
case "cos":
// 必填:SecretID/SecretKey/Region/BucketName/AppID;PathPrefix 可选
if req.COSSecretID == "" || req.COSSecretKey == "" ||
req.COSRegion == "" || req.COSBucketName == "" ||
req.COSAppID == "" {
logger.Error(ctx, "COS configuration is required")
c.Error(errors.NewBadRequestError("COS配置信息不能为空"))
return
}
case "minio":
if req.MinioBucketName == "" {
logger.Error(ctx, "MinIO configuration is required")
c.Error(errors.NewBadRequestError("MinIO配置信息不能为空"))
return
}
default:
logger.Error(ctx, "Invalid storage type")
c.Error(errors.NewBadRequestError("无效的存储类型"))
return
}
// 验证文件大小 — MAX_FILE_SIZE_MB env (50MB 默认)。
// 见 utils/filesize.go 注释:故意保留为部署期 env,不做 runtime setting。
maxSizeMB := utils.GetMaxFileSizeMB()
maxSize := maxSizeMB * 1024 * 1024
// 上限必须在 multipart 解析之前生效:FormFile 会先把整个 body 缓冲下来
// (超出内存部分落临时文件),之后的 header.Size 检查看到的已是收完的上传。
// nginx 的 location /api/ 仍是 MAX_FILE_SIZE;这里是直连 app 时的同一道上限。
limitUploadBody(c, maxSize)
// 获取上传的图片文件
file, header, err := c.Request.FormFile("image")
if err != nil {
if isRequestBodyTooLarge(err) {
logger.Error(ctx, "File size too large")
c.Error(errors.NewBadRequestError(fmt.Sprintf("图片文件大小不能超过%dMB", maxSizeMB)))
return
}
logger.Error(ctx, "Failed to get uploaded image", err)
c.Error(errors.NewBadRequestError("获取上传图片失败"))
return
}
defer file.Close()
// 验证文件类型
if !strings.HasPrefix(header.Header.Get("Content-Type"), "image/") {
logger.Error(ctx, "Invalid file type, only images are allowed")
c.Error(errors.NewBadRequestError("只允许上传图片文件"))
return
}
if header.Size > maxSize {
logger.Error(ctx, "File size too large")
c.Error(errors.NewBadRequestError(fmt.Sprintf("图片文件大小不能超过%dMB", maxSizeMB)))
return
}
logger.Infof(ctx, "Processing image: %s", utils.SanitizeForLog(header.Filename))
// 解析文档分割配置
chunkSizeInt32, err := strconv.ParseInt(req.ChunkSize, 10, 32)
if err != nil {
logger.Error(ctx, "Failed to parse chunk size", err)
c.Error(errors.NewBadRequestError("Failed to parse chunk size"))
return
}
chunkSize := int32(chunkSizeInt32)
if chunkSize < 100 || chunkSize > 10000 {
chunkSize = 1000
}
chunkOverlapInt32, err := strconv.ParseInt(req.ChunkOverlap, 10, 32)
if err != nil {
logger.Error(ctx, "Failed to parse chunk overlap", err)
c.Error(errors.NewBadRequestError("Failed to parse chunk overlap"))
return
}
chunkOverlap := int32(chunkOverlapInt32)
if chunkOverlap < 0 || chunkOverlap >= chunkSize {
chunkOverlap = 200
}
var separators []string
if req.SeparatorsRaw != "" {
if err := json.Unmarshal([]byte(req.SeparatorsRaw), &separators); err != nil {
separators = []string{"\n\n", "\n", "。", "!", "?", ";", ";"}
}
} else {
separators = []string{"\n\n", "\n", "。", "!", "?", ";", ";"}
}
// 读取图片文件内容
imageContent, err := io.ReadAll(file)
if err != nil {
logger.Error(ctx, "Failed to read image file", err)
c.Error(errors.NewBadRequestError("读取图片文件失败"))
return
}
// 调用多模态测试
startTime := time.Now()
result, err := h.testMultimodalWithDocReader(
ctx,
imageContent, header.Filename,
chunkSize, chunkOverlap, separators, &req,
)
processingTime := time.Since(startTime).Milliseconds()
if err != nil {
logger.Error(ctx, "Failed to test multimodal", err)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"success": false,
"message": err.Error(),
"processing_time": processingTime,
},
})
return
}
logger.Infof(ctx, "Multimodal test completed successfully in %dms", processingTime)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"success": true,
"caption": result["caption"],
"ocr": result["ocr"],
"processing_time": processingTime,
},
})
}
// testMultimodalWithDocReader uses DocumentReader.Read for document reading,
// then returns basic information about the result.
func (h *InitializationHandler) testMultimodalWithDocReader(
ctx context.Context,
imageContent []byte, filename string,
chunkSize, chunkOverlap int32, separators []string,
req *testMultimodalForm,
) (map[string]string, error) {
fileExt := ""
if idx := strings.LastIndex(filename, "."); idx != -1 {
fileExt = strings.ToLower(filename[idx+1:])
}
if h.documentReader == nil {
return nil, fmt.Errorf("DocReader service not configured")
}
requestID, _ := types.RequestIDFromContext(ctx)
readResult, err := h.documentReader.Read(ctx, &types.ReadRequest{
FileContent: imageContent,
FileName: filename,
FileType: fileExt,
RequestID: requestID,
})
if err != nil {
return nil, fmt.Errorf("调用DocReader服务失败: %v", err)
}
if readResult.Error != "" {
return nil, fmt.Errorf("DocReader服务返回错误: %s", readResult.Error)
}
result := map[string]string{
"markdown": readResult.MarkdownContent,
"caption": "",
"ocr": "",
}
return result, nil
}
// TextRelationExtractionRequest 文本关系提取请求结构
type TextRelationExtractionRequest struct {
Text string `json:"text" binding:"required"`
Tags []string `json:"tags" binding:"required"`
ModelID string `json:"model_id" binding:"required"`
}
// TextRelationExtractionResponse 文本关系提取响应结构
type TextRelationExtractionResponse struct {
Nodes []*types.GraphNode `json:"nodes"`
Relations []*types.GraphRelation `json:"relations"`
}
// ExtractTextRelations godoc
// @Summary 提取文本关系
// @Description 从文本中提取实体和关系
// @Tags 初始化
// @Accept json
// @Produce json
// @Param request body TextRelationExtractionRequest true "提取请求"
// @Success 200 {object} map[string]interface{} "提取结果"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/extract/text-relation [post]
func (h *InitializationHandler) ExtractTextRelations(c *gin.Context) {
ctx := c.Request.Context()
var req TextRelationExtractionRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "文本关系提取请求参数错误")
c.Error(errors.NewBadRequestError("文本关系提取请求参数错误"))
return
}
// 验证文本内容
if len(req.Text) == 0 {
c.Error(errors.NewBadRequestError("文本内容不能为空"))
return
}
if len(req.Text) > 5000 {
c.Error(errors.NewBadRequestError("文本内容长度不能超过5000字符"))
return
}
// 验证标签
if len(req.Tags) == 0 {
c.Error(errors.NewBadRequestError("至少需要选择一个关系标签"))
return
}
// 根据模型ID获取chat模型
chatModel, err := h.modelService.GetChatModel(ctx, req.ModelID)
if err != nil {
logger.Error(ctx, "获取模型失败", err)
c.Error(errors.NewBadRequestError("获取模型失败: " + err.Error()))
return
}
// 调用模型服务进行文本关系提取
result, err := h.extractRelationsFromText(ctx, req.Text, req.Tags, chatModel)
if err != nil {
logger.Error(ctx, "文本关系提取失败", err)
c.Error(errors.NewInternalServerError("文本关系提取失败: " + err.Error()))
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": result,
})
}
// extractRelationsFromText 从文本中提取关系
func (h *InitializationHandler) extractRelationsFromText(
ctx context.Context,
text string,
tags []string,
chatModel chat.Chat,
) (*TextRelationExtractionResponse, error) {
template := &types.PromptTemplateStructured{
Description: h.config.ExtractManager.ExtractGraph.Description,
Tags: tags,
Examples: h.config.ExtractManager.ExtractGraph.Examples,
}
extractor := chatpipeline.NewExtractor(chatModel, template)
graph, err := extractor.Extract(ctx, text)
if err != nil {
logger.Error(ctx, "文本关系提取失败", err)
return nil, err
}
extractor.RemoveUnknownRelation(ctx, graph)
result := &TextRelationExtractionResponse{
Nodes: graph.Node,
Relations: graph.Relation,
}
return result, nil
}
// FabriTextRequest is a request for generating example text
type FabriTextRequest struct {
Tags []string `json:"tags"`
ModelID string `json:"model_id" binding:"required"`
}
// FabriTextResponse is a response for generating example text
type FabriTextResponse struct {
Text string `json:"text"`
}
// FabriText godoc
// @Summary 生成示例文本
// @Description 根据标签生成示例文本
// @Tags 初始化
// @Accept json
// @Produce json
// @Param request body FabriTextRequest true "生成请求"
// @Success 200 {object} map[string]interface{} "生成的文本"
// @Failure 400 {object} errors.AppError "请求参数错误"
// @Security Bearer
// @Security ApiKeyAuth
// @Router /initialization/extract/fabri-text [post]
func (h *InitializationHandler) FabriText(c *gin.Context) {
ctx := c.Request.Context()
var req FabriTextRequest
if err := c.ShouldBindJSON(&req); err != nil {
logger.Error(ctx, "failed to parse fabri text request")
c.Error(errors.NewBadRequestError("invalid fabri text request parameters"))
return
}
chatModel, err := h.modelService.GetChatModel(ctx, req.ModelID)
if err != nil {
logger.Error(ctx, "获取模型失败", err)
c.Error(errors.NewBadRequestError("获取模型失败: " + err.Error()))
return
}
result, err := h.fabriText(ctx, req.Tags, chatModel)
if err != nil {
logger.Error(ctx, "failed to generate fabri text", err)
c.Error(errors.NewInternalServerError("failed to generate fabri text: " + err.Error()))
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": FabriTextResponse{Text: result},
})
}
// fabriText generates example text
func (h *InitializationHandler) fabriText(ctx context.Context, tags []string, chatModel chat.Chat) (string, error) {
content := h.config.ExtractManager.FabriText.WithNoTag
if len(tags) > 0 {
tagStr, _ := json.Marshal(tags)
content = fmt.Sprintf(h.config.ExtractManager.FabriText.WithTag, string(tagStr))
}
think := false
result, err := chatModel.Chat(ctx, []chat.Message{
{Role: "user", Content: content},
}, &chat.ChatOptions{
Temperature: 0.3,
MaxTokens: 4096,
Thinking: &think,
})
if err != nil {
logger.Error(ctx, "生成示例文本失败", err)
return "", err
}
return result.Content, nil
}
// FabriTagRequest is a request for generating tags
type FabriTagRequest struct{}
// FabriTagResponse is a response for generating tags
type FabriTagResponse struct {
Tags []string `json:"tags"`
}
var tagOptions = []string{
"Content", "Culture", "Person", "Event", "Time", "Location",
"Work", "Author", "Relation", "Attribute",
}
// FabriTag godoc
// @Summary 生成随机标签
// @Description 随机生成一组标签
// @Tags 初始化
// @Accept json
// @Produce json
// @Success 200 {object} map[string]interface{} "生成的标签"
// @Router /initialization/extract/fabri-tag [post]
func (h *InitializationHandler) FabriTag(c *gin.Context) {
tagRandom := RandomSelect(tagOptions, rand.Intn(len(tagOptions)-1)+1)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": FabriTagResponse{Tags: tagRandom},
})
}
// RandomSelect selects random strings
func RandomSelect(strs []string, n int) []string {
if n <= 0 {
return []string{}
}
result := make([]string, len(strs))
copy(result, strs)
rand.Shuffle(len(result), func(i, j int) {
result[i], result[j] = result[j], result[i]
})
if n > len(strs) {
n = len(strs)
}
return result[:n]
}