001package dev.tinyflow.agentsflex.provider;
002
003import com.agentsflex.core.message.AiMessage;
004import com.agentsflex.core.message.SystemMessage;
005import com.agentsflex.core.model.chat.ChatModel;
006import com.agentsflex.core.model.chat.StreamResponseListener;
007import com.agentsflex.core.model.chat.response.AiMessageResponse;
008import com.agentsflex.core.model.client.StreamContext;
009import com.agentsflex.core.prompt.SimplePrompt;
010import com.agentsflex.core.util.Maps;
011import dev.tinyflow.core.chain.Chain;
012import dev.tinyflow.core.llm.Llm;
013import dev.tinyflow.core.node.LlmNode;
014import dev.tinyflow.core.util.StringUtil;
015
016import java.util.List;
017import java.util.concurrent.CompletableFuture;
018import java.util.concurrent.ExecutionException;
019import java.util.concurrent.atomic.AtomicReference;
020
021public class AgentsFlexLlm implements Llm {
022
023    public static final String METADATA_KEY_LLM_ID = "llm_id";
024    public static final String METADATA_KEY_CHAIN_NODE_ID = "chain_node_id";
025    public static final String METADATA_KEY_CHAIN_STATE_ID = "chain_state_id";
026    public static final String AGENTS_FLEX_STREAM_LISTENER_NAME = "__agents_flex_stream_listener";
027
028    private ChatModel chatModel;
029
030    public ChatModel getChatModel() {
031        return chatModel;
032    }
033
034    public void setChatModel(ChatModel chatModel) {
035        this.chatModel = chatModel;
036    }
037
038    @Override
039    public String chat(MessageInfo messageInfo, ChatOptions options, LlmNode llmNode, Chain chain) {
040
041        SimplePrompt prompt = new SimplePrompt(messageInfo.getMessage());
042
043        // 系统提示词
044        if (StringUtil.hasText(messageInfo.getSystemMessage())) {
045            prompt.setSystemMessage(SystemMessage.of(messageInfo.getSystemMessage()));
046        }
047
048        //  图片
049        List<String> images = messageInfo.getImages();
050        if (images != null && !images.isEmpty()) {
051            for (String image : images) {
052                prompt.addImageUrl(image);
053            }
054        }
055
056
057        com.agentsflex.core.model.chat.ChatOptions chatOptions = new com.agentsflex.core.model.chat.ChatOptions();
058        chatOptions.setSeed(options.getSeed());
059        chatOptions.setTemperature(options.getTemperature());
060        chatOptions.setTopP(options.getTopP());
061        chatOptions.setTopK(options.getTopK());
062        chatOptions.setMaxTokens(options.getMaxTokens());
063        chatOptions.setStop(options.getStop());
064
065        chatOptions.putMetadata(METADATA_KEY_LLM_ID, llmNode.getLlmId());
066        chatOptions.putMetadata(METADATA_KEY_CHAIN_NODE_ID, llmNode.getId());
067        chatOptions.putMetadata(METADATA_KEY_CHAIN_STATE_ID, chain.getStateInstanceId());
068
069        String outType = llmNode.getOutType();
070        if ("json".equalsIgnoreCase(outType)) {
071            String jsonSchema = llmNode.getJsonSchema();
072            if (StringUtil.hasText(jsonSchema)) {
073                chatOptions.setResponseFormat(Maps.of("type", "json_schema").set("json_schema", jsonSchema));
074            } else {
075                chatOptions.setResponseFormat(Maps.of("type", "json_object"));
076            }
077        }
078
079
080        StreamResponseListener listener = chain.getEventManager().getOtherListener(AGENTS_FLEX_STREAM_LISTENER_NAME);
081        // 流式执行
082        if (listener != null) {
083            CompletableFuture<String> future = new CompletableFuture<>();
084            AtomicReference<Throwable> errorRef = new AtomicReference<>();
085            chatModel.chatStream(prompt, new StreamResponseListener() {
086                @Override
087                public void onStart(StreamContext context) {
088                    listener.onStart(context);
089                }
090
091                @Override
092                public void onMessage(StreamContext context, AiMessageResponse response) {
093                    listener.onMessage(context, response);
094                }
095
096                @Override
097                public void onFailure(StreamContext context, Throwable throwable) {
098                    try {
099                        listener.onFailure(context, throwable);
100                    } finally {
101                        errorRef.set(throwable);
102                    }
103                }
104
105                @Override
106                public void onStop(StreamContext context) {
107                    try {
108                        listener.onStop(context);
109                    } finally {
110                        Throwable error = errorRef.get();
111                        if (error != null) {
112                            future.completeExceptionally(error);
113                        } else {
114                            future.complete(context.getFullMessage().getFullContent());
115                        }
116                    }
117                }
118            }, chatOptions);
119
120            try {
121                return future.get();
122            } catch (InterruptedException e) {
123                Thread.currentThread().interrupt();
124                throw new RuntimeException(e);
125            } catch (ExecutionException e) {
126                Throwable cause = e.getCause();
127                if (cause instanceof RuntimeException) {
128                    throw (RuntimeException) cause;
129                }
130                throw new RuntimeException(cause);
131            }
132        }
133        // 非流式执行
134        else {
135            AiMessageResponse response = chatModel.chat(prompt, chatOptions);
136            if (response == null) {
137                throw new RuntimeException("AgentsFlexLlm can not get response!");
138            }
139
140            if (response.isError()) {
141                throw new RuntimeException("AgentsFlexLlm error: " + response.getErrorMessage());
142            }
143
144            AiMessage aiMessage = response.getMessage();
145            if (aiMessage != null) {
146                return aiMessage.getContent();
147            }
148
149            throw new RuntimeException("AgentsFlexLlm can not get aiMessage!");
150        }
151
152
153    }
154}