X-Agents/server/internal/service/knowledge_service.go

package service

import (
	"mime/multipart"
	"time"

	"github.com/google/uuid"
	"x-agents/server/internal/model"
	"x-agents/server/internal/repository"
)

type KnowledgeService struct {
	repo          *repository.KnowledgeRepository
	modelRepo     *repository.ModelRepository
	uploadService *UploadService
}

func NewKnowledgeService(repo *repository.KnowledgeRepository, modelRepo *repository.ModelRepository, uploadService *UploadService) *KnowledgeService {
	return &KnowledgeService{
		repo:          repo,
		modelRepo:     modelRepo,
		uploadService: uploadService,
	}
}

// Create 创建知识库
func (s *KnowledgeService) Create(req model.CreateKnowledgeRequest) (*model.KnowledgeBase, error) {
	// 验证 LLM 模型存在
	if _, err := s.modelRepo.FindByID(req.LLMModelID); err != nil {
		return nil, err
	}
	// 验证 Embedding 模型存在
	if _, err := s.modelRepo.FindByID(req.EmbeddingModelID); err != nil {
		return nil, err
	}

	kb := &model.KnowledgeBase{
		ID:                uuid.New().String(),
		Name:              req.Name,
		Description:       req.Description,
		LLMModelID:        req.LLMModelID,
		EmbeddingModelID:  req.EmbeddingModelID,
		ParsingConfig:     req.ParsingConfig,
		Status:            "active",
		DocumentCount:     0,
		ChunkCount:        0,
	}

	// 设置默认值
	if kb.ParsingConfig.EnablePDF == false && kb.ParsingConfig.EnablePDF != true {
		kb.ParsingConfig.EnablePDF = true
	}
	if kb.ParsingConfig.Pandoc == false && kb.ParsingConfig.Pandoc != true {
		kb.ParsingConfig.Pandoc = true
	}

	if err := s.repo.Create(kb); err != nil {
		return nil, err
	}
	return kb, nil
}

// List 获取知识库列表
func (s *KnowledgeService) List() ([]model.KnowledgeBase, error) {
	return s.repo.FindAll()
}

// GetByID 获取知识库详情
func (s *KnowledgeService) GetByID(id string) (*model.KnowledgeBase, error) {
	return s.repo.FindByID(id)
}

// Update 更新知识库
func (s *KnowledgeService) Update(id string, req model.UpdateKnowledgeRequest) error {
	updates := make(map[string]interface{})
	if req.Name != "" {
		updates["name"] = req.Name
	}
	if req.Description != "" {
		updates["description"] = req.Description
	}
	if req.LLMModelID != "" {
		// 验证模型存在
		if _, err := s.modelRepo.FindByID(req.LLMModelID); err != nil {
			return err
		}
		updates["llm_model_id"] = req.LLMModelID
	}
	if req.EmbeddingModelID != "" {
		// 验证模型存在
		if _, err := s.modelRepo.FindByID(req.EmbeddingModelID); err != nil {
			return err
		}
		updates["embedding_model_id"] = req.EmbeddingModelID
	}
	if req.ParsingConfig.Engine != "" {
		updates["parsing_config"] = req.ParsingConfig
	}
	if req.Status != "" {
		updates["status"] = req.Status
	}

	return s.repo.Update(id, updates)
}

// Delete 删除知识库
func (s *KnowledgeService) Delete(id string) error {
	// 先删除关联的文档
	if err := s.repo.DeleteDocumentsByKBID(id); err != nil {
		return err
	}
	return s.repo.Delete(id)
}

// ListDocuments 获取知识库下的文档列表
func (s *KnowledgeService) ListDocuments(kbID string, status string) ([]model.KnowledgeDocument, error) {
	return s.repo.FindDocumentsByKBID(kbID, status)
}

// UploadDocument 上传文档到知识库
func (s *KnowledgeService) UploadDocument(kbID string, file *multipart.FileHeader) (*model.KnowledgeDocument, error) {
	// 验证知识库存在
	_, err := s.repo.FindByID(kbID)
	if err != nil {
		return nil, err
	}

	// 上传文件
	result, err := s.uploadService.Upload(file)
	if err != nil {
		return nil, err
	}
	if !result.Success {
		return nil, nil
	}

	// 创建文档记录
	doc := &model.KnowledgeDocument{
		ID:               uuid.New().String(),
		KnowledgeBaseID:  kbID,
		Name:             file.Filename,
		FileKey:          result.FileKey,
		FileSize:         file.Size,
		Status:           "parsing",
		UploadedAt:       time.Now(),
	}

	if err := s.repo.CreateDocument(doc); err != nil {
		return nil, err
	}

	// 更新知识库文档数
	s.updateDocumentCount(kbID)

	return doc, nil
}

// DeleteDocument 删除文档
func (s *KnowledgeService) DeleteDocument(kbID, docID string) error {
	// 验证文档存在
	doc, err := s.repo.FindDocumentByID(docID)
	if err != nil {
		return err
	}
	if doc.KnowledgeBaseID != kbID {
		return nil
	}

	// 删除文件
	if doc.FileKey != "" {
		s.uploadService.DeleteFile(doc.FileKey + getFileExt(doc.Name))
	}

	// 删除文档记录
	if err := s.repo.DeleteDocument(docID); err != nil {
		return err
	}

	// 更新知识库文档数
	s.updateDocumentCount(kbID)

	return nil
}

// ReparseDocument 重新解析文档
func (s *KnowledgeService) ReparseDocument(kbID, docID string) error {
	// 验证文档存在
	doc, err := s.repo.FindDocumentByID(docID)
	if err != nil {
		return err
	}
	if doc.KnowledgeBaseID != kbID {
		return nil
	}

	// 重置状态为 parsing
	return s.repo.UpdateDocument(docID, map[string]interface{}{"status": "parsing"})
}

// GetDocumentPreview 获取文档预览
func (s *KnowledgeService) GetDocumentPreview(kbID, docID string, page int) (*model.DocumentPreviewResponse, error) {
	// 验证文档存在
	doc, err := s.repo.FindDocumentByID(docID)
	if err != nil {
		return nil, err
	}
	if doc.KnowledgeBaseID != kbID {
		return nil, nil
	}

	// 简单实现：返回文件 URL（实际应由 Python 服务处理）
	fileURL, err := s.uploadService.GetFileURL(doc.FileKey + getFileExt(doc.Name))
	if err != nil {
		return nil, err
	}

	return &model.DocumentPreviewResponse{
		TotalPages:  1,
		CurrentPage: page,
		Content:     fileURL,
	}, nil
}

// updateDocumentCount 更新知识库文档数
func (s *KnowledgeService) updateDocumentCount(kbID string) {
	count, _ := s.repo.CountDocumentsByKBID(kbID)
	s.repo.Update(kbID, map[string]interface{}{"document_count": int(count)})
}

func getFileExt(filename string) string {
	if len(filename) > 4 {
		return filename[len(filename)-4:]
	}
	return ""
}