admin/hermione
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fixed usedTokens count

Browse Source
This commit is contained in:
Emanuele Ferrelli
2024-11-15 16:45:45 +01:00
parent 5179bad78a
commit 202bb8a5e1
3 changed files with 13 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/main/java/com/olympus/hermione/services/ScenarioExecutionService.java
View File

@@ -177,17 +177,14 @@ public class ScenarioExecutionService {
executeScenarioStep(startStep, scenarioExecution);
Long usedTokens = (scenarioExecution.getUsedTokens() != null) ? scenarioExecution.getUsedTokens() : Long.valueOf(0);
scenarioExecution.setUsedTokens(usedTokens);
scenarioExecutionRepository.save(scenarioExecution);
while (scenarioExecution.getNextStepId()!=null) {
ScenarioStep step = steps.stream().filter(s -> s.getStepId().equals(scenarioExecution.getNextStepId())).findFirst().orElse(null);
executeScenarioStep(step, scenarioExecution);
if (scenarioExecution.getUsedTokens() != null && scenarioExecution.getUsedTokens() != 0) {
usedTokens += scenarioExecution.getUsedTokens();
scenarioExecution.setUsedTokens(usedTokens);
scenarioExecutionRepository.save(scenarioExecution);
scenarioExecution.setUsedTokens(Long.valueOf(0)); //resetting value for next step if is not an AI step
}
if(scenarioExecution.getLatestStepStatus() != null && scenarioExecution.getLatestStepStatus().equals("ERROR")){
@@ -196,6 +193,7 @@ public class ScenarioExecutionService {
}
}
scenarioExecution.setUsedTokens(usedTokens);
scenarioExecution.setEndDate(new java.util.Date());
scenarioExecutionRepository.save(scenarioExecution);
}

18
src/main/java/com/olympus/hermione/stepSolvers/AdvancedAIPromptSolver.java
View File

@@ -72,16 +72,6 @@ public class AdvancedAIPromptSolver extends StepSolver {
.param("chat_memory_response_size", 100))
.call();
/*Usage usage = resp.chatResponse().getMetadata().getUsage();
if (usage != null) {
Long usedTokens = usage.getTotalTokens();
this.scenarioExecution.setUsedTokens(usedTokens);
} else {
logger.info("Token usage information is not available.");
}*/
if(qai_output_entityType!=null && qai_output_entityType.equals("CiaOutputEntity")){
logger.info("Output is of type CiaOutputEntity");
@@ -99,6 +89,14 @@ public class AdvancedAIPromptSolver extends StepSolver {
String output = resp.content();
this.scenarioExecution.getExecSharedMap().put(this.qai_output_variable, output);
}
Usage usage = resp.chatResponse().getMetadata().getUsage();
if (usage != null) {
Long usedTokens = usage.getTotalTokens();
this.scenarioExecution.setUsedTokens(usedTokens);
} else {
logger.info("Token usage information is not available.");
}
this.scenarioExecution.setNextStepId(this.step.getNextStepId());

3
src/main/java/com/olympus/hermione/stepSolvers/BasicAIPromptSolver.java
View File

@@ -61,6 +61,8 @@ public class BasicAIPromptSolver extends StepSolver {
.messages(userMessage,systemMessage)
.call();
String output = resp.content();
Usage usage = resp.chatResponse().getMetadata().getUsage();
if (usage != null) {
Long usedTokens = usage.getTotalTokens();
@@ -69,7 +71,6 @@ public class BasicAIPromptSolver extends StepSolver {
logger.info("Token usage information is not available.");
}
String output = resp.content();
this.scenarioExecution.getExecSharedMap().put(this.qai_output_variable, output);
this.scenarioExecution.setNextStepId(this.step.getNextStepId());