package tech.easyflow.ai.easyagentsflow.knowledge;

import com.easyagents.core.document.Document;
import com.alibaba.fastjson2.JSONObject;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.knowledge.Knowledge;
import com.easyagents.flow.core.knowledge.KnowledgeProvider;
import com.easyagents.flow.core.node.KnowledgeNode;
import org.springframework.stereotype.Component;
import tech.easyflow.ai.rag.KnowledgeRetrievalRequest;
import tech.easyflow.ai.rag.KnowledgeRetrievalModes;
import tech.easyflow.ai.service.DocumentCollectionService;

import javax.annotation.Resource;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Component
public class KnowledgeProviderImpl implements KnowledgeProvider {

    @Resource
    private DocumentCollectionService documentCollectionService;

    /**
     * 获取知识库
     * @param id 知识库id
     */
    @Override
    public Knowledge getKnowledge(Object id) {
        return new Knowledge() {
            @Override
            public List<Map<String, Object>> search(String keyword, int limit, KnowledgeNode knowledgeNode, Chain chain) {
                KnowledgeRetrievalRequest request = new KnowledgeRetrievalRequest();
                request.setKnowledgeId(new BigInteger(id.toString()));
                request.setQuery(keyword);
                request.setLimit(limit);
                request.setRetrievalMode(KnowledgeRetrievalModes.parse(knowledgeNode.getRetrievalMode()));
                request.setCallerType("WORKFLOW");
                request.setCallerId(knowledgeNode.getId());
                List<Document> documents = documentCollectionService.search(request);
                if (limit > 0 && documents.size() > limit) {
                    documents = new ArrayList<>(documents.subList(0, limit));
                }
                List<Map<String, Object>> res = new ArrayList<>();
                for (Document document : documents) {
                    res.add(JSONObject.from(document));
                }
                return res;
            }
        };
    }
}
