package tech.aiflowy.ai.utils;
|
|
import com.agentsflex.core.chain.Chain;
|
import com.agentsflex.core.document.Document;
|
import com.agentsflex.core.llm.Llm;
|
import com.agentsflex.core.store.DocumentStore;
|
import com.agentsflex.core.store.SearchWrapper;
|
import com.agentsflex.core.store.StoreOptions;
|
import dev.tinyflow.core.Tinyflow;
|
import dev.tinyflow.core.knowledge.Knowledge;
|
import dev.tinyflow.core.node.KnowledgeNode;
|
import dev.tinyflow.core.parser.ChainParser;
|
import dev.tinyflow.core.provider.KnowledgeProvider;
|
import dev.tinyflow.core.provider.LlmProvider;
|
import dev.tinyflow.core.provider.SearchEngineProvider;
|
import dev.tinyflow.core.searchengine.SearchEngine;
|
import org.springframework.stereotype.Component;
|
import tech.aiflowy.ai.config.BochaaiProps;
|
import tech.aiflowy.ai.entity.AiKnowledge;
|
import tech.aiflowy.ai.entity.AiLlm;
|
import tech.aiflowy.ai.node.*;
|
import tech.aiflowy.ai.provider.BochaaiSearchEngine;
|
import tech.aiflowy.ai.service.AiKnowledgeService;
|
import tech.aiflowy.ai.service.AiLlmService;
|
import tech.aiflowy.common.filestorage.FileStorageService;
|
|
import javax.annotation.Resource;
|
import java.math.BigInteger;
|
import java.util.List;
|
|
@Component
|
public class TinyFlowConfigService {
|
|
@Resource
|
private ReaderManager readerManager;
|
@Resource(name = "default")
|
private FileStorageService storageService;
|
@Resource
|
private AiLlmService aiLlmService;
|
@Resource
|
private BochaaiProps bochaaiProps;
|
@Resource
|
private AiKnowledgeService aiKnowledgeService;
|
|
public void initProvidersAndNodeParsers(Tinyflow tinyflow) {
|
setExtraNodeParser(tinyflow);
|
setSearchEngineProvider(tinyflow);
|
setLlmProvider(tinyflow);
|
setKnowledgeProvider(tinyflow);
|
}
|
|
public void setExtraNodeParser(Tinyflow tinyflow) {
|
|
ReadDocService readService = readerManager.getReader();
|
// 文档解析
|
DocNodeParser docNodeParser = new DocNodeParser(readService);
|
// 文件生成
|
MakeFileNodeParser makeFileNodeParser = new MakeFileNodeParser(storageService);
|
// 插件
|
PluginToolNodeParser pluginToolNodeParser = new PluginToolNodeParser();
|
|
ChainParser chainParser = tinyflow.getChainParser();
|
chainParser.addNodeParser(docNodeParser.getNodeName(), docNodeParser);
|
chainParser.addNodeParser(makeFileNodeParser.getNodeName(), makeFileNodeParser);
|
chainParser.addNodeParser(pluginToolNodeParser.getNodeName(), pluginToolNodeParser);
|
}
|
|
public void setSearchEngineProvider(Tinyflow tinyflow) {
|
tinyflow.setSearchEngineProvider(new SearchEngineProvider() {
|
@Override
|
public SearchEngine getSearchEngine(Object id) {
|
BochaaiSearchEngine searchEngine = new BochaaiSearchEngine();
|
searchEngine.setApiKey(bochaaiProps.getApiKey());
|
return searchEngine;
|
}
|
});
|
}
|
|
public void setLlmProvider(Tinyflow tinyflow) {
|
tinyflow.setLlmProvider(new LlmProvider() {
|
@Override
|
public Llm getLlm(Object id) {
|
AiLlm aiLlm = aiLlmService.getById(new BigInteger(id.toString()));
|
return aiLlm.toLlm();
|
}
|
});
|
}
|
|
public void setKnowledgeProvider(Tinyflow tinyflow) {
|
tinyflow.setKnowledgeProvider(new KnowledgeProvider() {
|
@Override
|
public Knowledge getKnowledge(Object id) {
|
AiKnowledge aiKnowledge = aiKnowledgeService.getById(new BigInteger(id.toString()));
|
return new Knowledge() {
|
@Override
|
public List<Document> search(String keyword, int limit, KnowledgeNode knowledgeNode, Chain chain) {
|
DocumentStore documentStore = aiKnowledge.toDocumentStore();
|
if (documentStore == null) {
|
return null;
|
}
|
AiLlm aiLlm = aiLlmService.getById(aiKnowledge.getVectorEmbedLlmId());
|
if (aiLlm == null) {
|
return null;
|
}
|
documentStore.setEmbeddingModel(aiLlm.toLlm());
|
SearchWrapper wrapper = new SearchWrapper();
|
wrapper.setMaxResults(Integer.valueOf(limit));
|
wrapper.setText(keyword);
|
StoreOptions options = StoreOptions.ofCollectionName(aiKnowledge.getVectorStoreCollection());
|
|
List<Document> results = documentStore.search(wrapper, options);
|
return results;
|
}
|
};
|
}
|
});
|
}
|
}
|
