refactor: optimize SparkLlm.java

This commit is contained in:
Michael Yang 2024-10-30 11:41:33 +08:00
parent 074d1f03ec
commit ebafb3eae4
2 changed files with 22 additions and 5 deletions

View File

@ -117,11 +117,17 @@ public class SparkLlm extends BaseLlm<SparkLlmConfig> {
waitResponse(prompt, options, messageResponse, latch, failureThrowable);
AbstractBaseMessageResponse<?> response = messageResponse[0];
if (response == null) {
return null;
}
Throwable fialureThrowable = failureThrowable[0];
if (null == response.getMessage() || fialureThrowable != null) {
if (response == null) {
if (fialureThrowable != null) {
response = new AiMessageResponse("", null);
} else {
return null;
}
}
if (fialureThrowable != null || response.getMessage() == null) {
response.setError(true);
if (fialureThrowable != null) {
response.setErrorMessage(fialureThrowable.getMessage());

View File

@ -19,6 +19,7 @@ import com.agentsflex.core.document.Document;
import com.agentsflex.core.functions.Function;
import com.agentsflex.core.functions.Parameter;
import com.agentsflex.core.llm.ChatOptions;
import com.agentsflex.core.message.AiMessage;
import com.agentsflex.core.message.MessageStatus;
import com.agentsflex.core.parser.AiMessageParser;
import com.agentsflex.core.parser.FunctionMessageParser;
@ -30,6 +31,8 @@ import com.agentsflex.core.prompt.PromptFormat;
import com.agentsflex.core.util.HashUtil;
import com.agentsflex.core.util.Maps;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
@ -64,7 +67,15 @@ public class SparkLlmUtil {
public static AiMessageParser getAiMessageParser() {
DefaultAiMessageParser aiMessageParser = new DefaultAiMessageParser();
DefaultAiMessageParser aiMessageParser = new DefaultAiMessageParser(){
@Override
public AiMessage parse(JSONObject rootJson) {
if (!rootJson.containsKey("payload")) {
throw new JSONException("json not contains payload: " + rootJson);
}
return super.parse(rootJson);
}
};
aiMessageParser.setContentPath("$.payload.choices.text[0].content");
aiMessageParser.setIndexPath("$.payload.choices.text[0].index");
aiMessageParser.setStatusPath("$.payload.choices.status");