aiflowy-modules/aiflowy-module-ai/src/main/java/tech/aiflowy/ai/entity/AiWorkflowFunction.java
@@ -1,5 +1,7 @@
package tech.aiflowy.ai.entity;
import tech.aiflowy.ai.service.AiKnowledgeService;
import tech.aiflowy.ai.service.AiLlmService;
import tech.aiflowy.ai.service.AiWorkflowService;
import tech.aiflowy.common.util.SpringContextUtil;
import com.agentsflex.core.chain.Chain;
@@ -58,13 +60,58 @@
        AiWorkflowService service = SpringContextUtil.getBean(AiWorkflowService.class);
        AiWorkflow workflow = service.getById(this.workflowId);
        if (workflow != null) {
            Chain chain = workflow.toTinyflow().toChain();
            Tinyflow tinyflow = workflow.toTinyflow();
            setLlmProvider(tinyflow);
            setKnowledgeProvider(tinyflow);
            Chain chain = tinyflow.toChain();
            return chain.executeForResult(argsMap);
        } else {
            throw new RuntimeException("can not find the workflow by id: " + this.workflowId);
        }
    }
    private void setLlmProvider( Tinyflow tinyflow){
        AiLlmService aiLlmService = SpringContextUtil.getBean(AiLlmService.class);
        tinyflow.setLlmProvider(new LlmProvider() {
            @Override
            public Llm getLlm(Object id) {
                AiLlm aiLlm = aiLlmService.getById(new BigInteger(id.toString()));
                return aiLlm.toLlm();
            }
        });
    }
    private void setKnowledgeProvider( Tinyflow tinyflow){
        AiLlmService aiLlmService = SpringContextUtil.getBean(AiLlmService.class);
        AiKnowledgeService aiKnowledgeService= SpringContextUtil.getBean(AiKnowledgeService.class);
        tinyflow.setKnowledgeProvider(new KnowledgeProvider() {
            @Override
            public Knowledge getKnowledge(Object o) {
                AiKnowledge aiKnowledge = aiKnowledgeService.getById(new BigInteger(o.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;
                    }
                };
            }
        });
    }
    @Override
    public String toString() {
        return "AiWorkflowFunction{" +