X-Agents/server/internal/repository/knowledge_repo.go

package repository

import (
	"x-agents/server/internal/model"

	"gorm.io/gorm"
)

type KnowledgeRepository struct {
	db *gorm.DB
}

func NewKnowledgeRepository(db *gorm.DB) *KnowledgeRepository {
	return &KnowledgeRepository{db: db}
}

// Create 创建知识库
func (r *KnowledgeRepository) Create(kb *model.KnowledgeBase) error {
	return r.db.Create(kb).Error
}

// FindByID 根据ID查询
func (r *KnowledgeRepository) FindByID(id string) (*model.KnowledgeBase, error) {
	var kb model.KnowledgeBase
	err := r.db.First(&kb, "id = ?", id).Error
	if err != nil {
		return nil, err
	}
	return &kb, nil
}

// FindAll 查询所有
func (r *KnowledgeRepository) FindAll() ([]model.KnowledgeBase, error) {
	var list []model.KnowledgeBase
	err := r.db.Order("created_at DESC").Find(&list).Error
	return list, err
}

// Update 更新知识库
func (r *KnowledgeRepository) Update(id string, updates map[string]interface{}) error {
	return r.db.Model(&model.KnowledgeBase{}).Where("id = ?", id).Updates(updates).Error
}

// Delete 删除知识库
func (r *KnowledgeRepository) Delete(id string) error {
	return r.db.Delete(&model.KnowledgeBase{}, "id = ?", id).Error
}

// ========== Document ==========

// CreateDocument 创建文档
func (r *KnowledgeRepository) CreateDocument(doc *model.KnowledgeDocument) error {
	return r.db.Create(doc).Error
}

// FindDocumentByID 根据ID查询文档
func (r *KnowledgeRepository) FindDocumentByID(id string) (*model.KnowledgeDocument, error) {
	var doc model.KnowledgeDocument
	err := r.db.First(&doc, "id = ?", id).Error
	if err != nil {
		return nil, err
	}
	return &doc, nil
}

// FindDocumentsByKBID 根据知识库ID查询文档
func (r *KnowledgeRepository) FindDocumentsByKBID(kbID string, status string) ([]model.KnowledgeDocument, error) {
	var list []model.KnowledgeDocument
	query := r.db.Where("knowledge_base_id = ?", kbID).Order("uploaded_at DESC")
	if status != "" && status != "all" {
		query = query.Where("status = ?", status)
	}
	err := query.Find(&list).Error
	return list, err
}

// UpdateDocument 更新文档
func (r *KnowledgeRepository) UpdateDocument(id string, updates map[string]interface{}) error {
	return r.db.Model(&model.KnowledgeDocument{}).Where("id = ?", id).Updates(updates).Error
}

// DeleteDocument 删除文档
func (r *KnowledgeRepository) DeleteDocument(id string) error {
	return r.db.Delete(&model.KnowledgeDocument{}, "id = ?", id).Error
}

// DeleteDocumentsByKBID 删除知识库下所有文档
func (r *KnowledgeRepository) DeleteDocumentsByKBID(kbID string) error {
	return r.db.Delete(&model.KnowledgeDocument{}, "knowledge_base_id = ?", kbID).Error
}

// CountDocumentsByKBID 统计知识库下文档数
func (r *KnowledgeRepository) CountDocumentsByKBID(kbID string) (int64, error) {
	var count int64
	err := r.db.Model(&model.KnowledgeDocument{}).Where("knowledge_base_id = ?", kbID).Count(&count).Error
	return count, err
}