001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.parser.impl;
017
018import com.agentsflex.core.message.AiMessage;
019import com.agentsflex.core.message.ToolCall;
020import com.agentsflex.core.model.chat.ChatContext;
021import com.agentsflex.core.parser.AiMessageParser;
022import com.agentsflex.core.parser.JSONArrayParser;
023import com.agentsflex.core.util.JSONUtil;
024import com.alibaba.fastjson2.JSON;
025import com.alibaba.fastjson2.JSONArray;
026import com.alibaba.fastjson2.JSONObject;
027import com.alibaba.fastjson2.JSONPath;
028
029import java.util.ArrayList;
030import java.util.Collections;
031import java.util.List;
032import java.util.Map;
033
034
035public class DefaultAiMessageParser implements AiMessageParser<JSONObject> {
036
037    private JSONPath contentPath;
038    private JSONPath deltaContentPath;
039    private JSONPath[] reasoningContentPaths;
040    private JSONPath[] deltaReasoningContentPaths;
041    private JSONPath indexPath;
042    private JSONPath totalTokensPath;
043    private JSONPath promptTokensPath;
044    private JSONPath completionTokensPath;
045    private JSONPath finishReasonPath;
046    private JSONPath stopReasonPath;
047
048    private JSONPath toolCallsJsonPath;
049    private JSONPath deltaToolCallsJsonPath;
050
051    private JSONArrayParser<List<ToolCall>> callsParser;
052
053    public JSONPath getContentPath() {
054        return contentPath;
055    }
056
057    public void setContentPath(JSONPath contentPath) {
058        this.contentPath = contentPath;
059    }
060
061    public JSONPath getDeltaContentPath() {
062        return deltaContentPath;
063    }
064
065    public void setDeltaContentPath(JSONPath deltaContentPath) {
066        this.deltaContentPath = deltaContentPath;
067    }
068
069    public JSONPath[] getReasoningContentPaths() {
070        return reasoningContentPaths;
071    }
072
073    public void setReasoningContentPaths(JSONPath[] reasoningContentPaths) {
074        this.reasoningContentPaths = reasoningContentPaths;
075    }
076
077    public JSONPath[] getDeltaReasoningContentPaths() {
078        return deltaReasoningContentPaths;
079    }
080
081    public void setDeltaReasoningContentPaths(JSONPath[] deltaReasoningContentPaths) {
082        this.deltaReasoningContentPaths = deltaReasoningContentPaths;
083    }
084
085    public JSONPath getIndexPath() {
086        return indexPath;
087    }
088
089    public void setIndexPath(JSONPath indexPath) {
090        this.indexPath = indexPath;
091    }
092
093    public JSONPath getTotalTokensPath() {
094        return totalTokensPath;
095    }
096
097    public void setTotalTokensPath(JSONPath totalTokensPath) {
098        this.totalTokensPath = totalTokensPath;
099    }
100
101    public JSONPath getPromptTokensPath() {
102        return promptTokensPath;
103    }
104
105    public void setPromptTokensPath(JSONPath promptTokensPath) {
106        this.promptTokensPath = promptTokensPath;
107    }
108
109    public JSONPath getCompletionTokensPath() {
110        return completionTokensPath;
111    }
112
113    public void setCompletionTokensPath(JSONPath completionTokensPath) {
114        this.completionTokensPath = completionTokensPath;
115    }
116
117    public JSONPath getFinishReasonPath() {
118        return finishReasonPath;
119    }
120
121    public void setFinishReasonPath(JSONPath finishReasonPath) {
122        this.finishReasonPath = finishReasonPath;
123    }
124
125    public JSONPath getStopReasonPath() {
126        return stopReasonPath;
127    }
128
129    public void setStopReasonPath(JSONPath stopReasonPath) {
130        this.stopReasonPath = stopReasonPath;
131    }
132
133    public JSONPath getToolCallsJsonPath() {
134        return toolCallsJsonPath;
135    }
136
137    public void setToolCallsJsonPath(JSONPath toolCallsJsonPath) {
138        this.toolCallsJsonPath = toolCallsJsonPath;
139    }
140
141    public JSONPath getDeltaToolCallsJsonPath() {
142        return deltaToolCallsJsonPath;
143    }
144
145    public void setDeltaToolCallsJsonPath(JSONPath deltaToolCallsJsonPath) {
146        this.deltaToolCallsJsonPath = deltaToolCallsJsonPath;
147    }
148
149    public JSONArrayParser<List<ToolCall>> getCallsParser() {
150        return callsParser;
151    }
152
153    public void setCallsParser(JSONArrayParser<List<ToolCall>> callsParser) {
154        this.callsParser = callsParser;
155    }
156
157    @Override
158    public AiMessage parse(JSONObject rootJson, ChatContext context) {
159        AiMessage aiMessage = new AiMessage();
160
161        JSONArray toolCallsJsonArray = null;
162        if (context.getOptions().isStreaming()) {
163            if (this.deltaContentPath != null) {
164                aiMessage.setContent((String) this.deltaContentPath.eval(rootJson));
165            }
166            if (this.deltaReasoningContentPaths != null) {
167                for (JSONPath deltaReasoningContentPath : this.deltaReasoningContentPaths) {
168                    String result = (String) deltaReasoningContentPath.eval(rootJson);
169                    if (result != null) {
170                        aiMessage.setReasoningContent(result);
171                        break;
172                    }
173                }
174            }
175            if (this.deltaToolCallsJsonPath != null) {
176                toolCallsJsonArray = (JSONArray) this.deltaToolCallsJsonPath.eval(rootJson);
177            }
178        } else {
179            if (this.contentPath != null) {
180                aiMessage.setContent((String) this.contentPath.eval(rootJson));
181            }
182
183            if (this.reasoningContentPaths != null) {
184                for (JSONPath reasoningContentPath : this.reasoningContentPaths) {
185                    String result = (String) reasoningContentPath.eval(rootJson);
186                    if (result != null) {
187                        aiMessage.setReasoningContent(result);
188                        break;
189                    }
190                }
191            }
192            if (this.toolCallsJsonPath != null) {
193                toolCallsJsonArray = (JSONArray) this.toolCallsJsonPath.eval(rootJson);
194            }
195        }
196
197
198        if (this.indexPath != null) {
199            aiMessage.setIndex((Integer) this.indexPath.eval(rootJson));
200        }
201
202        if (this.promptTokensPath != null) {
203            aiMessage.setPromptTokens((Integer) this.promptTokensPath.eval(rootJson));
204        }
205
206        if (this.completionTokensPath != null) {
207            aiMessage.setCompletionTokens((Integer) this.completionTokensPath.eval(rootJson));
208        }
209
210        if (this.finishReasonPath != null) {
211            aiMessage.setFinishReason((String) this.finishReasonPath.eval(rootJson));
212        }
213
214        if (this.stopReasonPath != null) {
215            aiMessage.setStopReason((String) this.stopReasonPath.eval(rootJson));
216        }
217
218        if (this.totalTokensPath != null) {
219            aiMessage.setTotalTokens((Integer) this.totalTokensPath.eval(rootJson));
220        }
221        //some LLMs like Ollama not response the total tokens
222        else if (aiMessage.getPromptTokens() != null && aiMessage.getCompletionTokens() != null) {
223            aiMessage.setTotalTokens(aiMessage.getPromptTokens() + aiMessage.getCompletionTokens());
224        }
225
226        if (toolCallsJsonArray != null && this.callsParser != null) {
227            aiMessage.setToolCalls(this.callsParser.parse(toolCallsJsonArray));
228        }
229
230        return aiMessage;
231    }
232
233
234    public static DefaultAiMessageParser getOpenAIMessageParser() {
235        DefaultAiMessageParser aiMessageParser = new DefaultAiMessageParser();
236        aiMessageParser.setContentPath(JSONUtil.getJsonPath("$.choices[0].message.content"));
237        aiMessageParser.setDeltaContentPath(JSONUtil.getJsonPath("$.choices[0].delta.content"));
238
239        // 兼容最新版本的 vllm 输出
240        aiMessageParser.setReasoningContentPaths(new JSONPath[]{
241            JSONUtil.getJsonPath("$.choices[0].message.reasoning_content"),
242            JSONUtil.getJsonPath("$.choices[0].message.reasoning"),
243        });
244
245        aiMessageParser.setDeltaReasoningContentPaths(new JSONPath[]{
246            JSONUtil.getJsonPath("$.choices[0].delta.reasoning_content"),
247            JSONUtil.getJsonPath("$.choices[0].delta.reasoning"),
248        });
249
250        aiMessageParser.setIndexPath(JSONUtil.getJsonPath("$.choices[0].index"));
251        aiMessageParser.setTotalTokensPath(JSONUtil.getJsonPath("$.usage.total_tokens"));
252        aiMessageParser.setPromptTokensPath(JSONUtil.getJsonPath("$.usage.prompt_tokens"));
253        aiMessageParser.setCompletionTokensPath(JSONUtil.getJsonPath("$.usage.completion_tokens"));
254        aiMessageParser.setFinishReasonPath(JSONUtil.getJsonPath("$.choices[0].finish_reason"));
255        aiMessageParser.setStopReasonPath(JSONUtil.getJsonPath("$.choices[0].stop_reason"));
256
257        aiMessageParser.setToolCallsJsonPath(JSONUtil.getJsonPath("$.choices[0].message.tool_calls"));
258        aiMessageParser.setDeltaToolCallsJsonPath(JSONUtil.getJsonPath("$.choices[0].delta.tool_calls"));
259
260        aiMessageParser.setCallsParser(toolCalls -> {
261            if (toolCalls == null || toolCalls.isEmpty()) {
262                return Collections.emptyList();
263            }
264            List<ToolCall> toolInfos = new ArrayList<>();
265            for (int i = 0; i < toolCalls.size(); i++) {
266                JSONObject jsonObject = toolCalls.getJSONObject(i);
267                JSONObject functionObject = jsonObject.getJSONObject("function");
268                if (functionObject != null) {
269                    ToolCall toolCall = new ToolCall();
270                    toolCall.setId(jsonObject.getString("id"));
271                    toolCall.setName(functionObject.getString("name"));
272                    Object arguments = functionObject.get("arguments");
273                    if (arguments instanceof Map) {
274                        toolCall.setArguments(JSON.toJSONString(arguments));
275                    } else if (arguments instanceof String) {
276                        toolCall.setArguments((String) arguments);
277                    }
278                    toolInfos.add(toolCall);
279                }
280            }
281            return toolInfos;
282        });
283
284        return aiMessageParser;
285    }
286}