test/agents-flex
1
0
Fork 0
mirror of https://gitee.com/agents-flex/agents-flex.git synced 2024-11-29 18:38:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: fix AIMessage tokens parse

Browse Source
This commit is contained in:
Michael Yang 2024-05-30 11:19:42 +08:00
parent 1afa757a00
commit 2937f5aef8
9 changed files with 63 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
agents-flex-core/src/main/java/com/agentsflex/message/AiMessage.java
View File

@ -20,6 +20,8 @@ public class AiMessage extends TextMessage {
private Integer index;
private MessageStatus status;
private Integer promptTokens;
private Integer completionTokens;
private Integer totalTokens;
private String fullContent;
@ -39,6 +41,22 @@ public class AiMessage extends TextMessage {
this.status = status;
}
public Integer getPromptTokens() {
return promptTokens;
}
public void setPromptTokens(Integer promptTokens) {
this.promptTokens = promptTokens;
}
public Integer getCompletionTokens() {
return completionTokens;
}
public void setCompletionTokens(Integer completionTokens) {
this.completionTokens = completionTokens;
}
public Integer getTotalTokens() {
return totalTokens;
}

3
agents-flex-core/src/main/java/com/agentsflex/message/HumanMessage.java
View File

@ -17,6 +17,9 @@ package com.agentsflex.message;
public class HumanMessage extends TextMessage{
public HumanMessage() {
}
public HumanMessage(String content) {
setContent(content);
}

3
agents-flex-core/src/main/java/com/agentsflex/message/SystemMessage.java
View File

@ -17,6 +17,9 @@ package com.agentsflex.message;
public class SystemMessage extends TextMessage {
public SystemMessage() {
}
public SystemMessage(String content) {
this.content = content;
}

25
agents-flex-core/src/main/java/com/agentsflex/parser/impl/DefaultAiMessageParser.java
View File

@ -30,6 +30,8 @@ public class DefaultAiMessageParser implements AiMessageParser {
private String indexPath;
private String statusPath;
private String totalTokensPath;
private String promptTokensPath;
private String completionTokensPath;
private Parser<Object, MessageStatus> statusParser;
public String getContentPath() {
@ -64,6 +66,22 @@ public class DefaultAiMessageParser implements AiMessageParser {
this.totalTokensPath = totalTokensPath;
}
public String getPromptTokensPath() {
return promptTokensPath;
}
public void setPromptTokensPath(String promptTokensPath) {
this.promptTokensPath = promptTokensPath;
}
public String getCompletionTokensPath() {
return completionTokensPath;
}
public void setCompletionTokensPath(String completionTokensPath) {
this.completionTokensPath = completionTokensPath;
}
public Parser<Object, MessageStatus> getStatusParser() {
return statusParser;
}
@ -88,6 +106,13 @@ public class DefaultAiMessageParser implements AiMessageParser {
aiMessage.setTotalTokens((Integer) JSONPath.eval(rootJson, this.totalTokensPath));
}
if (StringUtil.hasText(promptTokensPath)) {
aiMessage.setPromptTokens((Integer) JSONPath.eval(rootJson, this.promptTokensPath));
}
if (StringUtil.hasText(completionTokensPath)) {
aiMessage.setCompletionTokens((Integer) JSONPath.eval(rootJson, this.completionTokensPath));
}
if (StringUtil.hasText(this.statusPath)) {
Object statusString = JSONPath.eval(rootJson, this.statusPath);
if (this.statusParser != null) {

2
agents-flex-llm/agents-flex-llm-chatglm/src/main/java/com/agentsflex/llm/chatglm/ChatglmLlmUtil.java
View File

@ -86,6 +86,8 @@ public class ChatglmLlmUtil {
aiMessageParser.setStatusPath("$.choices[0].finish_reason");
aiMessageParser.setStatusParser(content -> parseMessageStatus((String) content));
aiMessageParser.setTotalTokensPath("$.usage.total_tokens");
aiMessageParser.setPromptTokensPath("$.usage.prompt_tokens");
aiMessageParser.setCompletionTokensPath("$.usage.completion_tokens");
return aiMessageParser;
}

2
agents-flex-llm/agents-flex-llm-llama/src/main/java/com/agentsflex/llm/llama/LlamaLlmUtil.java
View File

@ -35,6 +35,8 @@ public class LlamaLlmUtil {
aiMessageParser.setStatusPath("$.choices[0].finish_reason");
aiMessageParser.setStatusParser(content -> parseMessageStatus((String) content));
aiMessageParser.setTotalTokensPath("$.usage.total_tokens");
aiMessageParser.setPromptTokensPath("$.usage.prompt_tokens");
aiMessageParser.setCompletionTokensPath("$.usage.completion_tokens");
return aiMessageParser;
}

4
agents-flex-llm/agents-flex-llm-openai/src/main/java/com/agentsflex/llm/openai/OpenAiLLmUtil.java
View File

@ -44,6 +44,10 @@ public class OpenAiLLmUtil {
aiMessageParser.setIndexPath("$.choices[0].index");
aiMessageParser.setStatusPath("$.choices[0].finish_reason");
aiMessageParser.setStatusParser(content -> parseMessageStatus((String) content));
aiMessageParser.setTotalTokensPath("$.usage.total_tokens");
aiMessageParser.setPromptTokensPath("$.usage.prompt_tokens");
aiMessageParser.setCompletionTokensPath("$.usage.completion_tokens");
return aiMessageParser;
}

3
agents-flex-llm/agents-flex-llm-qwen/src/main/java/com/agentsflex/llm/qwen/QwenLlmUtil.java
View File

@ -36,6 +36,9 @@ public class QwenLlmUtil {
aiMessageParser.setStatusPath("$.output.finish_reason");
aiMessageParser.setTotalTokensPath("$.usage.total_tokens");
aiMessageParser.setStatusParser(content -> parseMessageStatus((String) content));
aiMessageParser.setTotalTokensPath("$.usage.total_tokens");
aiMessageParser.setPromptTokensPath("$.usage.input_tokens");
aiMessageParser.setCompletionTokensPath("$.usage.output_tokens");
return aiMessageParser;
}

3
agents-flex-llm/agents-flex-llm-spark/src/main/java/com/agentsflex/llm/spark/SparkLlmUtil.java
View File

@ -67,6 +67,9 @@ public class SparkLlmUtil {
aiMessageParser.setContentPath("$.payload.choices.text[0].content");
aiMessageParser.setIndexPath("$.payload.choices.text[0].index");
aiMessageParser.setStatusPath("$.payload.choices.status");
aiMessageParser.setCompletionTokensPath("$.payload.usage.text.completion_tokens");
aiMessageParser.setPromptTokensPath("$.payload.usage.text.prompt_tokens");
aiMessageParser.setTotalTokensPath("$.payload.usage.text.total_tokens");
aiMessageParser.setStatusParser(content -> parseMessageStatus((Integer) content));
return aiMessageParser;
}