package tech.easyflow.ai.entity;

import com.easyagents.rag.retrieval.RetrievalMode;
import tech.easyflow.ai.entity.base.BotDocumentCollectionBase;
import tech.easyflow.ai.rag.KnowledgeRetrievalModes;
import com.mybatisflex.annotation.RelationOneToOne;
import com.mybatisflex.annotation.Table;

import java.util.HashMap;
import java.util.Map;

/**
 *  实体类。
 *
 * @author michael
 * @since 2024-08-28
 */

@Table("tb_bot_document_collection")
public class BotDocumentCollection extends BotDocumentCollectionBase {

    public static final String OPTION_KEY_RETRIEVAL_MODE = "retrievalMode";

    @RelationOneToOne(selfField = "documentCollectionId", targetField = "id")
    private DocumentCollection knowledge;

    public DocumentCollection getKnowledge() {
        return knowledge;
    }

    public void setKnowledge(DocumentCollection knowledge) {
        this.knowledge = knowledge;
    }

    public RetrievalMode getRetrievalMode() {
        Map<String, Object> options = getOptions();
        if (options == null) {
            return RetrievalMode.HYBRID;
        }
        Object value = options.get(OPTION_KEY_RETRIEVAL_MODE);
        return KnowledgeRetrievalModes.parse(value == null ? null : String.valueOf(value));
    }

    public void setRetrievalMode(RetrievalMode retrievalMode) {
        Map<String, Object> options = getOptions() == null
                ? new HashMap<>()
                : new HashMap<>(getOptions());
        options.put(OPTION_KEY_RETRIEVAL_MODE, (retrievalMode == null ? RetrievalMode.HYBRID : retrievalMode).name());
        setOptions(options);
    }
}
