agent-management/backend/src/main/java/com/agent/management/rag/kb/RagKnowledgeBaseConfigService.java

package com.agent.management.rag.kb;

import com.agent.management.model.entity.RagKnowledgeSourceBinding;
import com.agent.management.repository.RagKnowledgeSourceBindingRepository;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.Map;

@Service
@RequiredArgsConstructor
public class RagKnowledgeBaseConfigService {
    private final RagKnowledgeSourceBindingRepository bindings;
    private final ObjectMapper mapper;

    public List<RagKnowledgeSourceBinding> listBindings(Long knowledgeBaseId) {
        return bindings.findByKnowledgeBaseIdOrderByPriorityAsc(knowledgeBaseId);
    }

    public RagKnowledgeSourceBinding updateConfig(Long knowledgeBaseId, Long bindingId, Map<String, Object> config) {
        RagKnowledgeSourceBinding binding = bindings.findById(bindingId)
                .filter(item -> knowledgeBaseId.equals(item.getKnowledgeBaseId()))
                .orElseThrow(() -> new IllegalArgumentException("knowledge base binding does not exist: " + bindingId));
        validateAuthorization(binding, config);
        try {
            binding.setConfigJson(mapper.writeValueAsString(config == null ? Map.of() : config));
            return bindings.save(binding);
        } catch (Exception e) {
            throw new IllegalArgumentException("invalid binding config: " + e.getMessage(), e);
        }
    }

    private static void validateAuthorization(RagKnowledgeSourceBinding binding, Map<String,Object> config) {
        if (config == null || !"AUTO_GENERATE".equals(String.valueOf(config.get("retrievalMode")))) return;
        String key = binding.getSourceType() == com.agent.management.rag.model.RagSourceType.GRAPH ? "allowedLabels"
                : binding.getSourceType() == com.agent.management.rag.model.RagSourceType.STRUCTURED_DATA ? "allowedTables" : null;
        if (key != null && (!(config.get(key) instanceof List<?> values) || values.isEmpty()))
            throw new IllegalArgumentException("automatic generation requires a non-empty explicit " + key);
    }
}