ragflow/internal/service/kb.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package service

import (
	"errors"
	"fmt"
	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/model"
	"ragflow/internal/utility"
	"strings"
	"time"

	"github.com/google/uuid"
)

// KnowledgebaseService service class for managing dataset operations
type KnowledgebaseService struct {
	kbDAO         *dao.KnowledgebaseDAO
	userTenantDAO *dao.UserTenantDAO
	userDAO       *dao.UserDAO
	tenantDAO     *dao.TenantDAO
	connectorDAO  *dao.ConnectorDAO
}

// NewKnowledgebaseService creates a new knowledge base service
func NewKnowledgebaseService() *KnowledgebaseService {
	return &KnowledgebaseService{
		kbDAO:         dao.NewKnowledgebaseDAO(),
		userTenantDAO: dao.NewUserTenantDAO(),
		userDAO:       dao.NewUserDAO(),
		tenantDAO:     dao.NewTenantDAO(),
		connectorDAO:  dao.NewConnectorDAO(),
	}
}

// CreateKBRequest represents the request for creating a knowledge base
type CreateKBRequest struct {
	Name         string                 `json:"name" binding:"required"`
	ParserID     *string                `json:"parser_id,omitempty"`
	Description  *string                `json:"description,omitempty"`
	Language     *string                `json:"language,omitempty"`
	Permission   *string                `json:"permission,omitempty"`
	Avatar       *string                `json:"avatar,omitempty"`
	ParserConfig map[string]interface{} `json:"parser_config,omitempty"`
}

// CreateKBResponse represents the response for creating a knowledge base
type CreateKBResponse struct {
	KBID string `json:"kb_id"`
}

// UpdateKBRequest represents the request for updating a knowledge base
type UpdateKBRequest struct {
	KBID         string                 `json:"kb_id" binding:"required"`
	Name         string                 `json:"name" binding:"required"`
	Description  *string                `json:"description"`
	ParserID     string                 `json:"parser_id" binding:"required"`
	Permission   *string                `json:"permission,omitempty"`
	Language     *string                `json:"language,omitempty"`
	Avatar       *string                `json:"avatar,omitempty"`
	Pagerank     *int64                 `json:"pagerank,omitempty"`
	ParserConfig map[string]interface{} `json:"parser_config,omitempty"`
	Connectors   []string               `json:"connectors,omitempty"`
}

// UpdateMetadataSettingRequest represents the request for updating metadata settings
type UpdateMetadataSettingRequest struct {
	KBID           string                 `json:"kb_id" binding:"required"`
	Metadata       map[string]interface{} `json:"metadata" binding:"required"`
	EnableMetadata *bool                  `json:"enable_metadata,omitempty"`
}

// ListKbsRequest represents the request for listing knowledge bases
type ListKbsRequest struct {
	Keywords *string   `json:"keywords,omitempty"`
	Page     *int      `json:"page,omitempty"`
	PageSize *int      `json:"page_size,omitempty"`
	ParserID *string   `json:"parser_id,omitempty"`
	Orderby  *string   `json:"orderby,omitempty"`
	Desc     *bool     `json:"desc,omitempty"`
	OwnerIDs *[]string `json:"owner_ids,omitempty"`
}

// ListKbsResponse represents the response for listing knowledge bases
type ListKbsResponse struct {
	KBs   []map[string]interface{} `json:"kbs"`
	Total int64                    `json:"total"`
}

// CreateKB creates a new knowledge base
// This matches the Python create endpoint in kb_app.py
func (s *KnowledgebaseService) CreateKB(req *CreateKBRequest, tenantID string) (*CreateKBResponse, common.ErrorCode, error) {
	// Validate name is a string
	if !isValidString(req.Name) {
		return nil, common.CodeDataError, errors.New("Dataset name must be string.")
	}

	// Trim and validate name
	name := strings.TrimSpace(req.Name)
	if name == "" {
		return nil, common.CodeDataError, errors.New("Dataset name can't be empty.")
	}

	// Check name length (using UTF-8 byte length like Python)
	if len(name) > model.DatasetNameLimit {
		return nil, common.CodeDataError, fmt.Errorf("Dataset name length is %d which is large than %d", len(name), model.DatasetNameLimit)
	}

	// Verify tenant exists
	tenant, err := s.tenantDAO.GetByID(tenantID)
	if err != nil {
		return nil, common.CodeDataError, errors.New("Tenant not found.")
	}

	// Deduplicate name within tenant
	duplicateName := s.kbDAO.DuplicateName(name, tenantID)

	// Get parser ID (default to "naive")
	parserID := "naive"
	if req.ParserID != nil && *req.ParserID != "" {
		parserID = *req.ParserID
	}

	// Get parser config with defaults
	parserConfig := getParserConfig(parserID, req.ParserConfig)
	parserConfig["llm_id"] = tenant.LLMID

	// Generate KB ID
	kbID := strings.ReplaceAll(uuid.New().String(), "-", "")

	// Create knowledge base model
	now := time.Now().Unix()
	nowDate := time.Now().Truncate(time.Second)
	kb := &model.Knowledgebase{
		ID:           kbID,
		Name:         duplicateName,
		TenantID:     tenantID,
		CreatedBy:    tenantID,
		ParserID:     parserID,
		ParserConfig: parserConfig,
		Permission:   "me",
		EmbdID:       "",
	}
	kb.CreateTime = &now
	kb.UpdateTime = &now
	kb.CreateDate = &nowDate
	kb.UpdateDate = &nowDate
	status := string(model.StatusValid)
	kb.Status = &status

	// Set optional fields
	if req.Description != nil {
		kb.Description = req.Description
	}
	if req.Language != nil {
		kb.Language = req.Language
	}
	if req.Permission != nil {
		kb.Permission = *req.Permission
	}
	if req.Avatar != nil {
		kb.Avatar = req.Avatar
	}

	// Create in database
	if err := s.kbDAO.Create(kb); err != nil {
		return nil, common.CodeServerError, fmt.Errorf("failed to create knowledge base: %w", err)
	}

	return &CreateKBResponse{KBID: kbID}, common.CodeSuccess, nil
}

// UpdateKB updates an existing knowledge base
// This matches the Python update endpoint in kb_app.py
func (s *KnowledgebaseService) UpdateKB(req *UpdateKBRequest, userID string) (map[string]interface{}, common.ErrorCode, error) {
	// Validate name is a string
	if !isValidString(req.Name) {
		return nil, common.CodeDataError, errors.New("Dataset name must be string.")
	}

	// Trim and validate name
	name := strings.TrimSpace(req.Name)
	if name == "" {
		return nil, common.CodeDataError, errors.New("Dataset name can't be empty.")
	}

	// Check name length
	if len(name) > model.DatasetNameLimit {
		return nil, common.CodeDataError, fmt.Errorf("Dataset name length is %d which is large than %d", len(name), model.DatasetNameLimit)
	}

	// Check authorization
	if !s.kbDAO.Accessible4Deletion(req.KBID, userID) {
		return nil, common.CodeAuthenticationError, errors.New("No authorization.")
	}

	// Verify ownership
	kbs, err := s.kbDAO.Query(map[string]interface{}{"created_by": userID, "id": req.KBID})
	if err != nil || len(kbs) == 0 {
		return nil, common.CodeOperatingError, errors.New("only owner of dataset authorized for this operation")
	}

	// Get existing KB
	kb, err := s.kbDAO.GetByID(req.KBID)
	if err != nil {
		return nil, common.CodeDataError, errors.New("can't find this dataset")
	}

	// Check for duplicate name
	if strings.ToLower(name) != strings.ToLower(kb.Name) {
		existingKB, _ := s.kbDAO.GetByName(name, userID)
		if existingKB != nil {
			return nil, common.CodeDataError, errors.New("duplicated dataset name")
		}
	}

	// Build updates
	updates := map[string]interface{}{
		"name":      name,
		"parser_id": req.ParserID,
	}

	if req.Description != nil {
		updates["description"] = *req.Description
	}
	if req.Permission != nil {
		updates["permission"] = *req.Permission
	}
	if req.Language != nil {
		updates["language"] = *req.Language
	}
	if req.Avatar != nil {
		updates["avatar"] = *req.Avatar
	}
	if req.Pagerank != nil {
		updates["pagerank"] = *req.Pagerank
	}
	if req.ParserConfig != nil {
		updates["parser_config"] = req.ParserConfig
	}

	now := time.Now().Unix()
	nowDate := time.Now().Truncate(time.Second)
	updates["update_time"] = now
	updates["update_date"] = nowDate

	// Update in database
	if err := s.kbDAO.UpdateByID(req.KBID, updates); err != nil {
		return nil, common.CodeServerError, fmt.Errorf("failed to update knowledge base: %w", err)
	}

	// Get updated KB
	updatedKB, err := s.kbDAO.GetByID(req.KBID)
	if err != nil {
		return nil, common.CodeDataError, errors.New("database error (knowledgebase rename)")
	}

	result := updatedKB.ToMap()
	result["connectors"] = req.Connectors

	return result, common.CodeSuccess, nil
}

// UpdateMetadataSetting updates the metadata settings for a knowledge base
func (s *KnowledgebaseService) UpdateMetadataSetting(req *UpdateMetadataSettingRequest) (map[string]interface{}, common.ErrorCode, error) {
	kb, err := s.kbDAO.GetByID(req.KBID)
	if err != nil {
		return nil, common.CodeDataError, errors.New("database error (knowledgebase not found)")
	}

	parserConfig := kb.ParserConfig
	if parserConfig == nil {
		parserConfig = make(map[string]interface{})
	}

	parserConfig["metadata"] = req.Metadata
	enableMetadata := true
	if req.EnableMetadata != nil {
		enableMetadata = *req.EnableMetadata
	}
	parserConfig["enable_metadata"] = enableMetadata

	if err := s.kbDAO.UpdateParserConfig(req.KBID, parserConfig); err != nil {
		return nil, common.CodeServerError, fmt.Errorf("failed to update metadata setting: %w", err)
	}

	result := kb.ToMap()
	result["parser_config"] = parserConfig

	return result, common.CodeSuccess, nil
}

// GetDetail retrieves detailed information about a knowledge base
// This matches the Python kb_detail endpoint in kb_app.py
func (s *KnowledgebaseService) GetDetail(kbID, userID string) (*model.KnowledgebaseDetail, common.ErrorCode, error) {
	// Check authorization
	if !s.kbDAO.Accessible(kbID, userID) {
		return nil, common.CodeOperatingError, errors.New("only owner of dataset authorized for this operation")
	}

	// Get detail
	detail, err := s.kbDAO.GetDetail(kbID)
	if err != nil {
		return nil, common.CodeDataError, errors.New("can't find this dataset")
	}

	// Set connectors (empty for now)
	detail.Connectors = []string{}

	return detail, common.CodeSuccess, nil
}

// ListKbs lists knowledge bases with pagination and filtering
// This matches the Python list endpoint in kb_app.py
func (s *KnowledgebaseService) ListKbs(keywords string, page int, pageSize int, parserID string, orderby string, desc bool, ownerIDs []string, userID string) (*ListKbsResponse, common.ErrorCode, error) {
	var kbs []*model.KnowledgebaseListItem
	var total int64
	var err error

	if len(ownerIDs) > 0 {
		// List by owner IDs
		kbs, total, err = s.kbDAO.GetByTenantIDs(ownerIDs, userID, page, pageSize, orderby, desc, keywords, parserID)
	} else {
		// Get tenant IDs for user
		tenantIDs, err := s.userTenantDAO.GetTenantIDsByUserID(userID)
		if err != nil {
			return nil, common.CodeServerError, err
		}

		kbs, total, err = s.kbDAO.GetByTenantIDs(tenantIDs, userID, page, pageSize, orderby, desc, keywords, parserID)
	}

	if err != nil {
		return nil, common.CodeServerError, err
	}

	// Convert to map slice
	kbMaps := make([]map[string]interface{}, len(kbs))
	for i, kb := range kbs {
		kbMaps[i] = map[string]interface{}{
			"id":            kb.ID,
			"avatar":        kb.Avatar,
			"name":          kb.Name,
			"language":      kb.Language,
			"description":   kb.Description,
			"tenant_id":     kb.TenantID,
			"permission":    kb.Permission,
			"doc_num":       kb.DocNum,
			"token_num":     kb.TokenNum,
			"chunk_num":     kb.ChunkNum,
			"parser_id":     kb.ParserID,
			"embd_id":       kb.EmbdID,
			"nickname":      kb.Nickname,
			"tenant_avatar": kb.TenantAvatar,
			"update_time":   kb.UpdateTime,
		}
	}

	return &ListKbsResponse{
		KBs:   kbMaps,
		Total: total,
	}, common.CodeSuccess, nil
}

// DeleteKB soft deletes a knowledge base
// This matches the Python rm endpoint in kb_app.py
func (s *KnowledgebaseService) DeleteKB(kbID, userID string) (common.ErrorCode, error) {
	// Check authorization
	if !s.kbDAO.Accessible4Deletion(kbID, userID) {
		return common.CodeAuthenticationError, errors.New("No authorization.")
	}

	// Verify ownership
	kbs, err := s.kbDAO.Query(map[string]interface{}{"created_by": userID, "id": kbID})
	if err != nil || len(kbs) == 0 {
		return common.CodeOperatingError, errors.New("only owner of dataset authorized for this operation")
	}

	// Soft delete
	if err := s.kbDAO.Delete(kbID); err != nil {
		return common.CodeServerError, fmt.Errorf("database error (knowledgebase removal): %w", err)
	}

	return common.CodeSuccess, nil
}

// Accessible checks if a knowledge base is accessible by a user
func (s *KnowledgebaseService) Accessible(kbID, userID string) bool {
	return s.kbDAO.Accessible(kbID, userID)
}

// GetByID retrieves a knowledge base by ID
func (s *KnowledgebaseService) GetByID(kbID string) (*model.Knowledgebase, error) {
	return s.kbDAO.GetByID(kbID)
}

// GetKBIDsByTenantID retrieves all knowledge base IDs for a tenant
func (s *KnowledgebaseService) GetKBIDsByTenantID(tenantID string) ([]string, error) {
	return s.kbDAO.GetKBIDsByTenantID(tenantID)
}

// isValidString checks if a value is a non-empty string
func isValidString(v interface{}) bool {
	str, ok := v.(string)
	return ok && str != ""
}

// getParserConfig returns the parser configuration with defaults
// This matches the Python get_parser_config function
func getParserConfig(parserID string, customConfig map[string]interface{}) map[string]interface{} {
	config := map[string]interface{}{
		"pages":              [][]int{{1, 1000000}},
		"table_context_size": 0,
		"image_context_size": 0,
	}

	switch parserID {
	case "table":
		config["layout_recognize"] = false
		config["chunk_token_num"] = 128
		config["delimiter"] = "\n!?;。；！？"
		config["html4excel"] = false
	case "naive":
		config["chunk_token_num"] = 128
		config["delimiter"] = "\n!?;。；！？"
		config["html4excel"] = false
	default:
		config["raptor"] = map[string]interface{}{
			"use_raptor": false,
		}
	}

	// Merge custom config if provided
	if customConfig != nil {
		config = mergeParserConfig(config, customConfig)
	}

	return config
}

// mergeParserConfig merges two parser configurations
func mergeParserConfig(base, override map[string]interface{}) map[string]interface{} {
	result := make(map[string]interface{})
	for k, v := range base {
		result[k] = v
	}

	for k, v := range override {
		if existing, ok := result[k]; ok {
			if existingMap, ok := existing.(map[string]interface{}); ok {
				if newMap, ok := v.(map[string]interface{}); ok {
					result[k] = mergeParserConfig(existingMap, newMap)
					continue
				}
			}
		}
		result[k] = v
	}

	return result
}

// GenerateUUID generates a UUID string without dashes
func GenerateUUID() string {
	return strings.ReplaceAll(uuid.New().String(), "-", "")
}

// GetUserByToken gets user by authorization token
func (s *KnowledgebaseService) GetUserByToken(authorization string) (*model.User, common.ErrorCode, error) {
	userService := NewUserService()
	return userService.GetUserByToken(authorization)
}

// GetUserByID gets user by ID
func (s *KnowledgebaseService) GetUserByID(id string) (*model.User, error) {
	return s.userDAO.GetByAccessToken(id)
}

// GetTenantIDsByUserID gets tenant IDs for a user
func (s *KnowledgebaseService) GetTenantIDsByUserID(userID string) ([]string, error) {
	return s.userTenantDAO.GetTenantIDsByUserID(userID)
}

// GetConnectorsByTenantID gets connectors for a tenant
func (s *KnowledgebaseService) GetConnectorsByTenantID(tenantID string) ([]*dao.ConnectorListItem, error) {
	return s.connectorDAO.ListByTenantID(tenantID)
}

// GetKBList retrieves knowledge bases with ID and name filtering
func (s *KnowledgebaseService) GetKBList(tenantIDs []string, userID string, page, pageSize int, orderby string, desc bool, id, name string) ([]*model.Knowledgebase, int64, common.ErrorCode, error) {
	kbs, total, err := s.kbDAO.GetList(tenantIDs, userID, page, pageSize, orderby, desc, id, name)
	if err != nil {
		return nil, 0, common.CodeServerError, err
	}
	return kbs, total, common.CodeSuccess, nil
}

// GetKBByIDAndUserID retrieves a knowledge base by ID and user ID
func (s *KnowledgebaseService) GetKBByIDAndUserID(kbID, userID string) ([]*model.Knowledgebase, error) {
	return s.kbDAO.GetKBByIDAndUserID(kbID, userID)
}

// GetKBByNameAndUserID retrieves a knowledge base by name and user ID
func (s *KnowledgebaseService) GetKBByNameAndUserID(kbName, userID string) ([]*model.Knowledgebase, error) {
	return s.kbDAO.GetKBByNameAndUserID(kbName, userID)
}

// AtomicIncreaseDocNumByID atomically increments the document count
func (s *KnowledgebaseService) AtomicIncreaseDocNumByID(kbID string) error {
	return s.kbDAO.AtomicIncreaseDocNumByID(kbID)
}

// DecreaseDocumentNum decreases document, chunk, and token counts
func (s *KnowledgebaseService) DecreaseDocumentNum(kbID string, docNum, chunkNum, tokenNum int64) error {
	return s.kbDAO.DecreaseDocumentNum(kbID, docNum, chunkNum, tokenNum)
}

// UpdateParserConfig updates the parser configuration
func (s *KnowledgebaseService) UpdateParserConfig(id string, config map[string]interface{}) error {
	return s.kbDAO.UpdateParserConfig(id, config)
}

// DeleteFieldMap removes the field_map from parser_config
func (s *KnowledgebaseService) DeleteFieldMap(id string) error {
	return s.kbDAO.DeleteFieldMap(id)
}

// GetFieldMap retrieves field mappings from multiple knowledge bases
func (s *KnowledgebaseService) GetFieldMap(ids []string) (map[string]interface{}, error) {
	return s.kbDAO.GetFieldMap(ids)
}

// GetAllIDs retrieves all knowledge base IDs
func (s *KnowledgebaseService) GetAllIDs() ([]string, error) {
	return s.kbDAO.GetAllIDs()
}

// ExtractAccessToken extracts access token from authorization header
func ExtractAccessToken(authorization, secretKey string) (string, error) {
	return utility.ExtractAccessToken(authorization, secretKey)
}